fix: manually drain queued chat messages after cancel (#1166)

* fix: drain queued chat messages after cancel

* fix: address queue drain review feedback

* fix: keep queued messages parked after cancel

* fix: stabilize queue drain errors

* fix: keep queue parked during cancel events

* fix: wait for lifecycle stream goroutines

* fix: address lifecycle stream review feedback
This commit is contained in:
José Almeida
2026-05-31 11:58:01 +01:00
committed by GitHub
parent 8e1f55a0ac
commit 137336de86
26 changed files with 570 additions and 78 deletions
+1 -1
View File
@@ -141,7 +141,7 @@ func provideGateway(
orchestratorHandlers.RegisterHandlers(gateway.Dispatcher)
// Register message queue handlers
queueHandlers := orchestratorhandlers.NewQueueHandlers(orchestratorSvc.GetMessageQueue(), orchestratorSvc.GetEventBus(), log)
queueHandlers := orchestratorhandlers.NewQueueHandlers(orchestratorSvc.GetMessageQueue(), orchestratorSvc.GetEventBus(), log, orchestratorSvc)
queueHandlers.RegisterHandlers(gateway.Dispatcher)
if lifecycleMgr != nil && agentRegistry != nil {
@@ -218,7 +218,7 @@ func (m *Manager) ensureWorkspaceExecutionLocked(ctx context.Context, taskID, se
if m.streamManager != nil {
m.logger.Info("connecting workspace stream for workspace-only execution",
zap.String("execution_id", execution.ID))
go m.streamManager.connectWorkspaceStream(execution, nil)
m.streamManager.ConnectWorkspaceStream(execution, nil)
}
}()
@@ -671,7 +671,7 @@ func newEnvironmentExecutionTestManager(t *testing.T, provider WorkspaceInfoProv
ExecutorFallbackWarn, "", log,
)
mgr.workspaceInfoProvider = provider
t.Cleanup(func() { close(mgr.stopCh) })
cleanupManagerStopCh(t, mgr)
return mgr, backend
}
@@ -237,7 +237,7 @@ func TestLaunch_PublishesPrepareCompletedAfterRuntimeProgress(t *testing.T) {
)
mgr.preparerRegistry = NewPreparerRegistry(log)
mgr.preparerRegistry.Register(models.ExecutorTypeLocalDocker, &progressPreparer{})
t.Cleanup(func() { close(mgr.stopCh) })
cleanupManagerStopCh(t, mgr)
errCh := make(chan error, 1)
go func() {
@@ -506,7 +506,7 @@ func TestLaunch_RaceRollback(t *testing.T) {
ExecutorFallbackWarn, "", log,
)
mgr.dataDir = t.TempDir()
t.Cleanup(func() { close(mgr.stopCh) })
cleanupManagerStopCh(t, mgr)
req := &LaunchRequest{
TaskID: "task-1",
@@ -571,7 +571,7 @@ func TestLaunch_PersistsDockerRuntimeSecrets(t *testing.T) {
)
mgr.SetSecretStore(store)
mgr.dataDir = t.TempDir()
t.Cleanup(func() { close(mgr.stopCh) })
cleanupManagerStopCh(t, mgr)
execution, err := mgr.Launch(context.Background(), &LaunchRequest{
TaskID: "task-1",
@@ -170,6 +170,9 @@ func (m *Manager) Stop() error {
m.logger.Info("stopping lifecycle manager")
close(m.stopCh)
if m.streamManager != nil {
m.streamManager.Wait()
}
m.wg.Wait()
// Close executor backends that hold resources (e.g., Docker SDK client).
@@ -468,7 +468,7 @@ func (m *Manager) startPassthroughSession(ctx context.Context, execution *AgentE
m.startPassthroughShell(ctx, execution, "failed to start shell for passthrough session")
if m.streamManager != nil && execution.agentctl != nil {
go m.streamManager.connectWorkspaceStream(execution, nil)
m.streamManager.ConnectWorkspaceStream(execution, nil)
}
go m.autoInjectInitialPrompt(execution, pt)
@@ -687,7 +687,7 @@ func (m *Manager) ResumePassthroughSession(ctx context.Context, sessionID string
// Connect to workspace stream for shell/git/file features.
// Only connect if not already connected (process restart reuses the same agentctl).
if m.streamManager != nil && execution.agentctl != nil && execution.GetWorkspaceStream() == nil {
go m.streamManager.connectWorkspaceStream(execution, nil)
m.streamManager.ConnectWorkspaceStream(execution, nil)
}
return nil
@@ -959,7 +959,7 @@ func (m *Manager) attemptResumeFallback(execution *AgentExecution, runner *proce
// works but the user's shell session and workspace stream stay torn down.
m.startPassthroughShell(ctx, execution, "failed to start shell after passthrough resume fallback")
if m.streamManager != nil && execution.agentctl != nil && execution.GetWorkspaceStream() == nil {
go m.streamManager.connectWorkspaceStream(execution, nil)
m.streamManager.ConnectWorkspaceStream(execution, nil)
}
// Fallback path is a fresh session (no --resume) — re-inject the prompt.
@@ -4,6 +4,7 @@ import (
"context"
"io"
"strings"
"sync"
"testing"
"time"
@@ -177,6 +178,31 @@ func newTestRegistry() *registry.Registry {
return reg
}
var testStopChClosers sync.Map
func closeStopChOnce(stopCh chan struct{}) {
closer, _ := testStopChClosers.LoadOrStore(stopCh, &sync.Once{})
closer.(*sync.Once).Do(func() { close(stopCh) })
}
func cleanupManagerStopCh(t *testing.T, mgr *Manager) {
t.Helper()
t.Cleanup(func() {
closeStopChOnce(mgr.stopCh)
if mgr.streamManager != nil {
mgr.streamManager.Wait()
}
})
}
func cleanupStreamManager(t *testing.T, stopCh chan struct{}, streamMgr *StreamManager) {
t.Helper()
t.Cleanup(func() {
closeStopChOnce(stopCh)
streamMgr.Wait()
})
}
// MockCredentialsManager implements CredentialsManager for testing
type MockCredentialsManager struct{}
@@ -650,7 +650,7 @@ func (sm *SessionManager) retryPromptAfterReconnect(
attachments []v1.MessageAttachment,
) error {
ready := make(chan struct{})
go sm.streamManager.connectUpdatesStream(execution, ready)
sm.streamManager.connectUpdatesStreamAsync(execution, ready)
select {
case <-ready:
@@ -36,13 +36,7 @@ func newSessionTestLogger() *logger.Logger {
func newTestStopCh(t *testing.T) chan struct{} {
t.Helper()
stopCh := make(chan struct{})
t.Cleanup(func() {
select {
case <-stopCh:
default:
close(stopCh)
}
})
t.Cleanup(func() { closeStopChOnce(stopCh) })
return stopCh
}
@@ -221,6 +215,7 @@ func TestInitializeAndPrompt_StreamBeforeInitialize(t *testing.T) {
streamMgr := NewStreamManager(log, StreamCallbacks{
OnAgentEvent: func(execution *AgentExecution, event agentctl.AgentEvent) {},
}, nil, stopCh)
cleanupStreamManager(t, stopCh, streamMgr)
sm.SetDependencies(nil, streamMgr, nil, nil)
client := createTestClient(t, mock.server.URL)
@@ -295,6 +290,7 @@ func TestInitializeAndPrompt_StreamTimeout(t *testing.T) {
streamMgr := NewStreamManager(log, StreamCallbacks{
OnAgentEvent: func(execution *AgentExecution, event agentctl.AgentEvent) {},
}, nil, stopCh)
cleanupStreamManager(t, stopCh, streamMgr)
sm.SetDependencies(nil, streamMgr, nil, nil)
// Bind to a random port and immediately close it so the port is guaranteed
@@ -357,6 +353,7 @@ func TestInitializeAndPrompt_WithTaskDescription(t *testing.T) {
streamMgr := NewStreamManager(log, StreamCallbacks{
OnAgentEvent: func(execution *AgentExecution, event agentctl.AgentEvent) {},
}, nil, stopCh)
cleanupStreamManager(t, stopCh, streamMgr)
sm.SetDependencies(nil, streamMgr, nil, nil)
client := createTestClient(t, mock.server.URL)
@@ -2,6 +2,7 @@ package lifecycle
import (
"context"
"sync"
"time"
"go.uber.org/zap"
@@ -35,7 +36,10 @@ type StreamManager struct {
// reconnect/backoff loops below select on stopCh to drain on Manager.Stop.
// nil is treated as "no shutdown signal" and falls back to the prior
// uncancellable behaviour (compat with constructors used by isolated tests).
stopCh <-chan struct{}
stopCh <-chan struct{}
wg sync.WaitGroup
wgMu sync.Mutex
stopped bool
}
type stopChannelContext struct {
@@ -78,8 +82,48 @@ func NewStreamManager(log *logger.Logger, callbacks StreamCallbacks, mcpHandler
// completes (success or failure). Agent operations require the updates stream;
// workspace stream readiness is handled independently.
func (sm *StreamManager) ConnectAll(execution *AgentExecution, ready chan<- struct{}) {
go sm.connectUpdatesStream(execution, ready)
go sm.connectWorkspaceStream(execution, nil)
sm.connectUpdatesStreamAsync(execution, ready)
sm.ConnectWorkspaceStream(execution, nil)
}
func (sm *StreamManager) connectUpdatesStreamAsync(execution *AgentExecution, ready chan<- struct{}) {
if !sm.start(func() {
sm.connectUpdatesStream(execution, ready)
}) && ready != nil {
close(ready)
}
}
// ConnectWorkspaceStream starts the workspace stream and tracks the goroutine
// so shutdown and tests can wait for it to drain after stopCh closes.
func (sm *StreamManager) ConnectWorkspaceStream(execution *AgentExecution, ready chan<- struct{}) {
if !sm.start(func() {
sm.connectWorkspaceStream(execution, ready)
}) && ready != nil {
close(ready)
}
}
// Wait blocks until all StreamManager-owned stream goroutines have exited.
func (sm *StreamManager) Wait() {
sm.wgMu.Lock()
sm.stopped = true
sm.wgMu.Unlock()
sm.wg.Wait()
}
func (sm *StreamManager) start(fn func()) bool {
sm.wgMu.Lock()
defer sm.wgMu.Unlock()
if sm.stopped {
return false
}
sm.wg.Add(1)
go func() {
defer sm.wg.Done()
fn()
}()
return true
}
// ReconnectAll reconnects to all streams (used after backend restart).
@@ -133,6 +133,12 @@ func (s *Service) handleAgentBootReady(ctx context.Context, data watcher.AgentEv
zap.String("session_id", data.SessionID))
return
}
if s.isCancelInFlight(data.SessionID) {
s.logger.Debug("ignoring agent.boot_ready while cancel is in progress",
zap.String("task_id", data.TaskID),
zap.String("session_id", data.SessionID))
return
}
session, err := s.repo.GetTaskSession(ctx, data.SessionID)
if err != nil {
@@ -180,7 +186,7 @@ func (s *Service) handleAgentBootReady(ctx context.Context, data watcher.AgentEv
// on the queue until the user manually sends another message. After the
// agent has booted and the session is back to WAITING_FOR_INPUT it's safe
// to dispatch any pending message.
s.drainQueuedMessageAfterTransition(ctx, data.SessionID)
s.drainQueuedMessageForPromptableSession(ctx, data.SessionID)
}
// handleAgentReady handles turn-end ready events: the agent finished processing
@@ -202,6 +208,13 @@ func (s *Service) handleAgentReady(ctx context.Context, data watcher.AgentEventD
zap.String("agent_execution_id", data.AgentExecutionID))
return
}
if s.isCancelInFlight(data.SessionID) {
s.logger.Debug("ignoring agent.ready while cancel is in progress",
zap.String("task_id", data.TaskID),
zap.String("session_id", data.SessionID),
zap.String("agent_execution_id", data.AgentExecutionID))
return
}
session, err := s.repo.GetTaskSession(ctx, data.SessionID)
if err != nil {
@@ -232,7 +232,7 @@ func TestAutoStartTransientError_BootReadyDrainsOrphanedQueue(t *testing.T) {
SessionID: sessionID,
})
// drainQueuedMessageAfterTransition pops the queue synchronously and fires
// drainQueuedMessageForPromptableSession pops the queue synchronously and fires
// `go executeQueuedMessage(...)`. With the user_message_recorded flag set
// at queue time (see autoStartStepPrompt's retry branch), the goroutine
// SKIPS its CreateUserMessage and just calls PromptTask → PromptAgent.
@@ -676,7 +676,7 @@ func TestHandleAgentBootReady_DrainsOrphanedQueuedMessage(t *testing.T) {
agentManager: agentMgr,
messageQueue: messagequeue.NewServiceMemory(log),
// Wire a real executor so the executeQueuedMessage goroutine
// spawned by drainQueuedMessageAfterTransition can safely call
// spawned by drainQueuedMessageForPromptableSession can safely call
// PromptTask -> executor.GetExecutionBySession without nil-derefing.
executor: executor.NewExecutor(agentMgr, repo, log, executor.ExecutorConfig{}),
}
@@ -780,3 +780,44 @@ func TestHandleAgentBootReady_DoesNotDrainForTerminalSession(t *testing.T) {
t.Errorf("queue count after boot ready on terminal session = %d, want 1 (must not drain)", got)
}
}
func TestHandleAgentBootReady_DoesNotDrainWhileCancelInFlight(t *testing.T) {
ctx := context.Background()
repo := setupTestRepo(t)
seedSession(t, repo, "t1", "s1", "step1")
seedExecutorRunning(t, repo, "s1", "t1", "exec-1")
log := testLogger()
agentMgr := &mockAgentManager{
isAgentRunning: true,
repoForExecutionLookup: repo,
promptDone: make(chan struct{}),
}
svc := &Service{
logger: log,
repo: repo,
taskRepo: newMockTaskRepo(),
agentManager: agentMgr,
messageQueue: messagequeue.NewServiceMemory(log),
executor: executor.NewExecutor(agentMgr, repo, log, executor.ExecutorConfig{}),
}
if _, err := svc.messageQueue.QueueMessage(
ctx, "s1", "t1", "queued after cancel", "",
messagequeue.QueuedByUser, false, nil,
); err != nil {
t.Fatalf("queue prompt: %v", err)
}
svc.cancelInFlight.Store("s1", struct{}{})
defer svc.cancelInFlight.Delete("s1")
svc.handleAgentBootReady(ctx, watcher.AgentEventData{TaskID: "t1", SessionID: "s1"})
status := svc.messageQueue.GetStatus(ctx, "s1")
if status.Count != 1 {
t.Fatalf("queue count after boot ready during cancel = %d, want 1", status.Count)
}
if len(agentMgr.capturedPrompts) != 0 {
t.Fatalf("expected no queued prompt dispatch during cancel, got %d prompts", len(agentMgr.capturedPrompts))
}
}
@@ -616,6 +616,34 @@ func TestHandleAgentReadyGuards(t *testing.T) {
}
})
t.Run("ignores ready while cancel is in flight", func(t *testing.T) {
repo := setupTestRepo(t)
seedSession(t, repo, "t1", "s1", "step1")
stepGetter := newMockStepGetter()
stepGetter.steps["step1"] = &wfmodels.WorkflowStep{
ID: "step1", WorkflowID: "wf1", Name: "Step 1", Position: 0,
}
agentMgr := &mockAgentManager{isAgentRunning: true}
svc := createTestServiceWithAgent(repo, stepGetter, newMockTaskRepo(), agentMgr)
if _, err := svc.messageQueue.QueueMessage(ctx, "s1", "t1", "queued", "", messagequeue.QueuedByUser, false, nil); err != nil {
t.Fatalf("failed to queue message: %v", err)
}
svc.cancelInFlight.Store("s1", struct{}{})
defer svc.cancelInFlight.Delete("s1")
svc.handleAgentReady(ctx, watcher.AgentEventData{TaskID: "t1", SessionID: "s1"})
status := svc.messageQueue.GetStatus(ctx, "s1")
if status.Count != 1 {
t.Fatalf("expected queued message to remain parked during cancel, count=%d entries=%+v", status.Count, status.Entries)
}
if len(agentMgr.capturedPrompts) != 0 {
t.Fatalf("expected no queued prompt dispatch during cancel, got %d prompts", len(agentMgr.capturedPrompts))
}
})
t.Run("moves STARTING session to waiting for input", func(t *testing.T) {
repo := setupTestRepo(t)
seedSession(t, repo, "t1", "s1", "step1")
@@ -838,7 +838,7 @@ func (s *Service) processOnEnter(ctx context.Context, taskID string, session *mo
}
s.setSessionWaitingForInput(ctx, taskID, sessionID, session)
s.publishSessionWaitingEvent(ctx, taskID, sessionID, step.ID, session)
s.drainQueuedMessageAfterTransition(ctx, sessionID)
s.drainQueuedMessageForPromptableSession(ctx, sessionID)
return
}
@@ -924,7 +924,7 @@ func (s *Service) processOnEnter(ctx context.Context, taskID string, session *mo
zap.Error(err))
s.setSessionWaitingForInput(asyncCtx, taskID, sessionID, session)
s.publishSessionWaitingEvent(asyncCtx, taskID, sessionID, stepID, session)
s.drainQueuedMessageAfterTransition(asyncCtx, sessionID)
s.drainQueuedMessageForPromptableSession(asyncCtx, sessionID)
}
}()
return
@@ -942,7 +942,7 @@ func (s *Service) processOnEnter(ctx context.Context, taskID string, session *mo
// steps without auto_start_agent (e.g. Review). Drain here to match the
// pre-#677 behavior where handleAgentReady always drained after returning
// from inline processOnEnter.
s.drainQueuedMessageAfterTransition(ctx, sessionID)
s.drainQueuedMessageForPromptableSession(ctx, sessionID)
}
}
@@ -953,7 +953,7 @@ func (s *Service) processOnEnter(ctx context.Context, taskID string, session *mo
// task.moved handler doesn't run a second processStepExitAndEnter for the same
// transition. The message queue is left intact — any user-supplied prompt
// already queued by handleMoveTask is delivered by the on_enter path or by
// drainQueuedMessageAfterTransition.
// drainQueuedMessageForPromptableSession.
func (s *Service) applyPendingMove(ctx context.Context, taskID, sessionID string, session *models.TaskSession, move *messagequeue.PendingMove) {
// reinsertPendingMove restores the move so a future agent.ready can retry.
// Used on early failure paths (load errors, config issues) where the state
@@ -992,7 +992,7 @@ func (s *Service) applyPendingMove(ctx context.Context, taskID, sessionID string
s.logger.Info("pending move target equals current step; skipping transition",
zap.String("task_id", taskID),
zap.String("step_id", fromStepID))
s.drainQueuedMessageAfterTransition(ctx, sessionID)
s.drainQueuedMessageForPromptableSession(ctx, sessionID)
return
}
@@ -1048,26 +1048,29 @@ func (s *Service) applyPendingMove(ctx context.Context, taskID, sessionID string
go s.processStepExitAndEnter(context.WithoutCancel(ctx), taskID, session, fromStepID, move.WorkflowStepID, taskDescription)
}
// drainQueuedMessageAfterTransition takes any user-queued message and dispatches
// it for execution. Used by processOnEnter branches that don't auto-start the
// agent — without this, the message is orphaned because handleAgentReady skips
// the queue when a workflow transition occurred.
func (s *Service) drainQueuedMessageAfterTransition(ctx context.Context, sessionID string) {
// drainQueuedMessageForPromptableSession takes the next queued message and dispatches
// it for execution. Callers must ensure the session is ready for input first.
//
// Used by processOnEnter branches that don't auto-start the agent, manual
// drain requests, and cancel recovery. Without this, the message is orphaned
// because handleAgentReady only drains from a normal turn-end event.
func (s *Service) drainQueuedMessageForPromptableSession(ctx context.Context, sessionID string) bool {
if s.messageQueue == nil {
return
return false
}
queuedMsg, ok := s.messageQueue.TakeQueued(ctx, sessionID)
if !ok || queuedMsg == nil {
return
return false
}
s.publishQueueStatusEvent(ctx, sessionID)
if queuedMsg.Content == "" && len(queuedMsg.Attachments) == 0 {
s.logger.Warn("skipping empty queued message after transition",
zap.String("session_id", sessionID),
zap.String("queue_id", queuedMsg.ID))
return
return false
}
go s.executeQueuedMessage(sessionID, queuedMsg)
return true
}
// deliverPassthroughPrompt writes a prompt to PTY stdin and marks the session as running.
@@ -1359,7 +1362,7 @@ func (s *Service) fallbackFreshLaunchOnMissingExecution(
}
// takeAndMergeHandoffMessage drains any queued hand-off message for the session
// (set by handleMoveTask via move_task_kandev or by drainQueuedMessageAfterTransition)
// (set by handleMoveTask via move_task_kandev or by drainQueuedMessageForPromptableSession)
// and merges its content + attachments into the auto-start prompt. Returns the
// original queued message (so terminal failure paths can re-queue it via
// requeueMessage), the merged prompt, and the converted attachments. Empty
@@ -7,6 +7,7 @@ import (
"github.com/kandev/kandev/internal/common/logger"
"github.com/kandev/kandev/internal/events"
"github.com/kandev/kandev/internal/events/bus"
"github.com/kandev/kandev/internal/orchestrator"
"github.com/kandev/kandev/internal/orchestrator/messagequeue"
ws "github.com/kandev/kandev/pkg/websocket"
"go.uber.org/zap"
@@ -16,6 +17,8 @@ const (
// queueErrorCodeEntryNotFound is surfaced when an edit/remove targets an entry
// that has already been drained (atomic-take won the race).
queueErrorCodeEntryNotFound = "entry_not_found"
queueErrorCodeSessionBusy = "session_busy"
queueErrorCodeNotPromptable = "session_not_promptable"
// Payload field names — extracted to satisfy goconst (≥3 occurrences).
fieldSessionID = "session_id"
@@ -35,17 +38,27 @@ type QueueService interface {
GetStatus(ctx context.Context, sessionID string) *messagequeue.QueueStatus
}
type QueueDrainer interface {
DrainQueuedMessage(ctx context.Context, sessionID string) (bool, error)
}
// QueueHandlers handles WebSocket message-queue operations.
type QueueHandlers struct {
queueService QueueService
queueDrainer QueueDrainer
eventBus bus.EventBus
logger *logger.Logger
}
// NewQueueHandlers creates a new QueueHandlers instance.
func NewQueueHandlers(queueService QueueService, eventBus bus.EventBus, log *logger.Logger) *QueueHandlers {
func NewQueueHandlers(queueService QueueService, eventBus bus.EventBus, log *logger.Logger, drainer ...QueueDrainer) *QueueHandlers {
var queueDrainer QueueDrainer
if len(drainer) > 0 {
queueDrainer = drainer[0]
}
return &QueueHandlers{
queueService: queueService,
queueDrainer: queueDrainer,
eventBus: eventBus,
logger: log.WithFields(zap.String("component", "queue-handlers")),
}
@@ -58,6 +71,7 @@ func (h *QueueHandlers) RegisterHandlers(d *ws.Dispatcher) {
d.RegisterFunc(ws.ActionMessageQueueGet, h.wsGetQueueStatus)
d.RegisterFunc(ws.ActionMessageQueueUpdate, h.wsUpdateMessage)
d.RegisterFunc(ws.ActionMessageQueueAppend, h.wsAppendToQueue)
d.RegisterFunc(ws.ActionMessageQueueDrain, h.wsDrainQueue)
d.RegisterFunc(ws.ActionMessageQueueRemove, h.wsRemoveEntry)
}
@@ -141,6 +155,42 @@ func (h *QueueHandlers) wsCancelAll(ctx context.Context, msg *ws.Message) (*ws.M
})
}
type wsDrainQueueRequest struct {
SessionID string `json:"session_id"`
}
func (h *QueueHandlers) wsDrainQueue(ctx context.Context, msg *ws.Message) (*ws.Message, error) {
var req wsDrainQueueRequest
if err := msg.ParsePayload(&req); err != nil {
return ws.NewError(msg.ID, msg.Action, ws.ErrorCodeBadRequest, "Invalid payload: "+err.Error(), nil)
}
if req.SessionID == "" {
return ws.NewError(msg.ID, msg.Action, ws.ErrorCodeValidation, "session_id is required", nil)
}
if h.queueDrainer == nil {
return ws.NewError(msg.ID, msg.Action, ws.ErrorCodeInternalError, "Queue drain is unavailable", nil)
}
drained, err := h.queueDrainer.DrainQueuedMessage(ctx, req.SessionID)
if err != nil {
switch {
case errors.Is(err, orchestrator.ErrAgentPromptInProgress):
return ws.NewError(msg.ID, msg.Action, queueErrorCodeSessionBusy, "Session is busy", nil)
case errors.Is(err, orchestrator.ErrSessionNotPromptable):
return ws.NewError(msg.ID, msg.Action, queueErrorCodeNotPromptable, "Session is not ready for input", nil)
default:
h.logger.Error("failed to drain queued message", zap.String(fieldSessionID, req.SessionID), zap.Error(err))
return ws.NewError(msg.ID, msg.Action, ws.ErrorCodeInternalError, "Failed to drain queued message", nil)
}
}
h.publishStatus(ctx, req.SessionID)
return ws.NewResponse(msg.ID, msg.Action, map[string]interface{}{
fieldSessionID: req.SessionID,
"drained": drained,
})
}
type wsGetQueueStatusRequest struct {
SessionID string `json:"session_id"`
}
@@ -3,11 +3,13 @@ package handlers
import (
"context"
"encoding/json"
"errors"
"testing"
"time"
"github.com/kandev/kandev/internal/common/logger"
"github.com/kandev/kandev/internal/events/bus"
"github.com/kandev/kandev/internal/orchestrator"
"github.com/kandev/kandev/internal/orchestrator/messagequeue"
ws "github.com/kandev/kandev/pkg/websocket"
"github.com/stretchr/testify/assert"
@@ -32,7 +34,24 @@ func (m *mockEventBus) Request(_ context.Context, _ string, _ *bus.Event, _ time
func (m *mockEventBus) Close() {}
func (m *mockEventBus) IsConnected() bool { return true }
type mockQueueDrainer struct {
calls int
sessionID string
drained bool
err error
}
func (m *mockQueueDrainer) DrainQueuedMessage(_ context.Context, sessionID string) (bool, error) {
m.calls++
m.sessionID = sessionID
return m.drained, m.err
}
func setupQueueHandlers(t *testing.T) (*QueueHandlers, *messagequeue.Service) {
return setupQueueHandlersWithDrainer(t, nil)
}
func setupQueueHandlersWithDrainer(t *testing.T, drainer QueueDrainer) (*QueueHandlers, *messagequeue.Service) {
t.Helper()
log, err := logger.NewLogger(logger.LoggingConfig{
Level: "error",
@@ -41,7 +60,7 @@ func setupQueueHandlers(t *testing.T) (*QueueHandlers, *messagequeue.Service) {
})
require.NoError(t, err)
svc := messagequeue.NewServiceMemory(log)
return NewQueueHandlers(svc, &mockEventBus{}, log), svc
return NewQueueHandlers(svc, &mockEventBus{}, log, drainer), svc
}
func createTestMessage(t *testing.T, action string, payload interface{}) *ws.Message {
@@ -164,6 +183,85 @@ func TestWsCancelAll(t *testing.T) {
})
}
func TestWsDrainQueue(t *testing.T) {
t.Run("drains next queued message", func(t *testing.T) {
drainer := &mockQueueDrainer{drained: true}
handlers, _ := setupQueueHandlersWithDrainer(t, drainer)
response, err := handlers.wsDrainQueue(context.Background(),
createTestMessage(t, ws.ActionMessageQueueDrain, map[string]interface{}{"session_id": "s"}))
require.NoError(t, err)
assert.Equal(t, ws.MessageTypeResponse, response.Type)
assert.Equal(t, 1, drainer.calls)
assert.Equal(t, "s", drainer.sessionID)
var payload map[string]interface{}
require.NoError(t, json.Unmarshal(response.Payload, &payload))
assert.Equal(t, true, payload["drained"])
})
t.Run("reports no queued message", func(t *testing.T) {
drainer := &mockQueueDrainer{drained: false}
handlers, _ := setupQueueHandlersWithDrainer(t, drainer)
response, err := handlers.wsDrainQueue(context.Background(),
createTestMessage(t, ws.ActionMessageQueueDrain, map[string]interface{}{"session_id": "s"}))
require.NoError(t, err)
assert.Equal(t, ws.MessageTypeResponse, response.Type)
assert.Equal(t, 1, drainer.calls)
var payload map[string]interface{}
require.NoError(t, json.Unmarshal(response.Payload, &payload))
assert.Equal(t, false, payload["drained"])
})
t.Run("rejects missing session_id", func(t *testing.T) {
handlers, _ := setupQueueHandlersWithDrainer(t, &mockQueueDrainer{})
response, err := handlers.wsDrainQueue(context.Background(),
createTestMessage(t, ws.ActionMessageQueueDrain, map[string]interface{}{}))
require.NoError(t, err)
assert.Equal(t, ws.MessageTypeError, response.Type)
assert.Contains(t, parseError(t, response).Message, "session_id is required")
})
t.Run("rejects unavailable drainer", func(t *testing.T) {
handlers, _ := setupQueueHandlers(t)
response, err := handlers.wsDrainQueue(context.Background(),
createTestMessage(t, ws.ActionMessageQueueDrain, map[string]interface{}{"session_id": "s"}))
require.NoError(t, err)
assert.Equal(t, ws.MessageTypeError, response.Type)
assert.Equal(t, ws.ErrorCodeInternalError, parseError(t, response).Code)
})
t.Run("reports busy session", func(t *testing.T) {
handlers, _ := setupQueueHandlersWithDrainer(t, &mockQueueDrainer{err: orchestrator.ErrAgentPromptInProgress})
response, err := handlers.wsDrainQueue(context.Background(),
createTestMessage(t, ws.ActionMessageQueueDrain, map[string]interface{}{"session_id": "s"}))
require.NoError(t, err)
assert.Equal(t, ws.MessageTypeError, response.Type)
assert.Equal(t, queueErrorCodeSessionBusy, parseError(t, response).Code)
})
t.Run("reports not promptable session", func(t *testing.T) {
handlers, _ := setupQueueHandlersWithDrainer(t, &mockQueueDrainer{err: orchestrator.ErrSessionNotPromptable})
response, err := handlers.wsDrainQueue(context.Background(),
createTestMessage(t, ws.ActionMessageQueueDrain, map[string]interface{}{"session_id": "s"}))
require.NoError(t, err)
assert.Equal(t, ws.MessageTypeError, response.Type)
assert.Equal(t, queueErrorCodeNotPromptable, parseError(t, response).Code)
assert.Equal(t, "Session is not ready for input", parseError(t, response).Message)
})
t.Run("reports internal drain errors", func(t *testing.T) {
handlers, _ := setupQueueHandlersWithDrainer(t, &mockQueueDrainer{err: errors.New("boom")})
response, err := handlers.wsDrainQueue(context.Background(),
createTestMessage(t, ws.ActionMessageQueueDrain, map[string]interface{}{"session_id": "s"}))
require.NoError(t, err)
assert.Equal(t, ws.MessageTypeError, response.Type)
assert.Equal(t, ws.ErrorCodeInternalError, parseError(t, response).Code)
})
}
func TestWsGetQueueStatus(t *testing.T) {
t.Run("returns ordered list", func(t *testing.T) {
handlers, svc := setupQueueHandlers(t)
@@ -340,12 +340,9 @@ type Service struct {
clarificationWatchdogTimeout time.Duration
// cancelInFlight tracks sessionIDs whose CancelAgent call is currently in
// progress. Used to deduplicate impatient retries from the UI: while the
// first cancel is still propagating through the agent (which can take several
// seconds when a long-running tool like Claude's Monitor is being torn down),
// the user often clicks the button repeatedly. Without this guard each click
// would create another "Turn cancelled by user" message and lazily start a
// phantom turn just to host it.
// progress. It deduplicates impatient retries from the UI and lets late
// agent.ready/boot_ready events from the cancelled turn return before they
// evaluate workflow transitions or drain queued messages.
cancelInFlight sync.Map
// transientRetries tracks in-progress transient-provider-error (529
@@ -2199,6 +2199,28 @@ func (s *Service) RespondToPermission(ctx context.Context, sessionID, pendingID,
return nil
}
// DrainQueuedMessage dispatches one queued message for a session that is ready
// for input. It is intentionally one-at-a-time: each successful prompt will
// complete its own turn and then drain the next entry through handleAgentReady.
func (s *Service) DrainQueuedMessage(ctx context.Context, sessionID string) (bool, error) {
if sessionID == "" {
return false, fmt.Errorf("session_id is required")
}
session, err := s.repo.GetTaskSession(ctx, sessionID)
if err != nil {
return false, fmt.Errorf("failed to get session: %w", err)
}
if err := s.checkSessionPromptable(session.TaskID, sessionID, session.State); err != nil {
return false, err
}
return s.drainQueuedMessageForPromptableSession(ctx, sessionID), nil
}
func (s *Service) isCancelInFlight(sessionID string) bool {
_, ok := s.cancelInFlight.Load(sessionID)
return ok
}
// CancelAgent interrupts the current agent turn without terminating the process,
// allowing the user to send a new prompt.
//
@@ -15,6 +15,7 @@ import (
"github.com/kandev/kandev/internal/events"
"github.com/kandev/kandev/internal/orchestrator/dto"
"github.com/kandev/kandev/internal/orchestrator/executor"
"github.com/kandev/kandev/internal/orchestrator/messagequeue"
"github.com/kandev/kandev/internal/orchestrator/queue"
"github.com/kandev/kandev/internal/orchestrator/scheduler"
"github.com/kandev/kandev/internal/sysprompt"
@@ -371,6 +372,45 @@ func TestCancelAgent_DeduplicatesConcurrentCalls(t *testing.T) {
}
}
func TestCancelAgent_LeavesQueuedMessageForManualDrain(t *testing.T) {
ctx := context.Background()
repo := setupTestRepo(t)
agentMgr := &mockAgentManager{isAgentRunning: true}
svc := createTestServiceWithAgent(repo, newMockStepGetter(), newMockTaskRepo(), agentMgr)
svc.executor = executor.NewExecutor(agentMgr, repo, testLogger(), executor.ExecutorConfig{})
seedTaskAndSession(t, repo, "task1", "session1", models.TaskSessionStateRunning)
seedExecutorRunning(t, repo, "session1", "task1", "exec-1")
if _, err := svc.messageQueue.QueueMessage(
ctx, "session1", "task1", "queued after cancel", "", messagequeue.QueuedByUser, false, nil,
); err != nil {
t.Fatalf("queue message: %v", err)
}
if err := svc.CancelAgent(ctx, "session1"); err != nil {
t.Fatalf("cancel agent: %v", err)
}
updated, err := repo.GetTaskSession(ctx, "session1")
if err != nil {
t.Fatalf("get updated session: %v", err)
}
if updated.State != models.TaskSessionStateWaitingForInput {
t.Fatalf("expected session state %q, got %q", models.TaskSessionStateWaitingForInput, updated.State)
}
status := svc.messageQueue.GetStatus(ctx, "session1")
if status.Count != 1 {
t.Fatalf("expected cancel to leave queued message for manual drain, count=%d entries=%+v", status.Count, status.Entries)
}
if got := status.Entries[0].Content; got != "queued after cancel" {
t.Fatalf("queued prompt = %q, want %q", got, "queued after cancel")
}
if len(agentMgr.capturedPrompts) != 0 {
t.Fatalf("expected cancel not to prompt queued message, got %d prompts", len(agentMgr.capturedPrompts))
}
}
// --- StartCreatedSession ---
func TestStartCreatedSession_WrongTask(t *testing.T) {
+1
View File
@@ -113,6 +113,7 @@ const (
ActionMessageQueueGet = "message.queue.get"
ActionMessageQueueUpdate = "message.queue.update"
ActionMessageQueueAppend = "message.queue.append"
ActionMessageQueueDrain = "message.queue.drain" // Dispatch one queued entry now when the session is promptable
ActionMessageQueueRemove = "message.queue.remove" // Delete a single entry by id
ActionMessageQueueStatusChanged = "message.queue.status_changed" // Notification: queue status changed
@@ -523,6 +523,8 @@ export function ChatInputArea({
}: ChatInputAreaProps) {
const { resolvedSessionId, taskId, isAgentBusy, needsRecovery, planModeEnabled, todoItems } =
panelState;
const sessionState = panelState.session?.state ?? null;
const canDrainQueue = sessionState === "WAITING_FOR_INPUT" || sessionState === "IDLE";
const { planActions, executor, placeholder } = useChatInputDerived(
panelState,
chatInputRef,
@@ -533,12 +535,13 @@ export function ChatInputArea({
<div className="bg-card flex-shrink-0 px-2 pb-2 pt-1">
<QueueAffordance
sessionId={resolvedSessionId}
canDrain={canDrainQueue}
renderStatusBar={(queueChip) => (
<ChatStatusBar
todoItems={todoItems}
taskId={taskId}
sessionId={resolvedSessionId}
sessionState={panelState.session?.state ?? null}
sessionState={sessionState}
nextStepName={proceedStepName}
onProceed={proceed}
isAgentBusy={isAgentBusy}
@@ -62,6 +62,7 @@ function baseState(entries: QueuedMessage[]) {
isLoading: false,
queue: vi.fn(async () => {}),
clearAll: vi.fn(async () => {}),
drainNext: vi.fn(async () => {}),
editEntry: vi.fn(async () => {}),
removeEntry: vi.fn(async () => {}),
refetch: vi.fn(async () => {}),
@@ -165,6 +166,30 @@ describe("QueueAffordance", () => {
fireEvent.click(screen.getByTestId("queue-clear-all"));
expect(state.clearAll).toHaveBeenCalledTimes(1);
});
it("shows a run-next action when manual drain is available", () => {
const state = queueState([entry()]);
useQueueMock.mockReturnValue(state);
render(
<QueueAffordance sessionId={SESSION_ID} canDrain>
{CHILD}
</QueueAffordance>,
);
fireEvent.click(screen.getByTestId(CHIP_ID));
fireEvent.click(screen.getByTestId("queue-drain-next"));
expect(state.drainNext).toHaveBeenCalledTimes(1);
});
it("hides the run-next action while the agent is busy", () => {
useQueueMock.mockReturnValue(queueState([entry()]));
render(
<QueueAffordance sessionId={SESSION_ID} canDrain={false}>
{CHILD}
</QueueAffordance>,
);
fireEvent.click(screen.getByTestId(CHIP_ID));
expect(screen.queryByTestId("queue-drain-next")).toBeNull();
});
});
describe("QueueAffordance — renderStatusBar prop", () => {
@@ -2,7 +2,7 @@
import { useCallback, useEffect, useState } from "react";
import type { ReactNode } from "react";
import { IconLayoutList, IconTrash, IconX } from "@tabler/icons-react";
import { IconLayoutList, IconPlayerPlay, IconTrash, IconX } from "@tabler/icons-react";
import { toast } from "sonner";
import { Button } from "@kandev/ui";
import { Collapsible, CollapsibleContent } from "@kandev/ui/collapsible";
@@ -57,6 +57,7 @@ function useEscToClose(open: boolean, onClose: () => void): void {
type QueueAffordanceProps = {
sessionId: string | null;
children: ReactNode;
canDrain?: boolean;
/**
* Optional render slot for placing the chip inside an external row (e.g. the
* chat status bar). The callback receives the chip node (or `null` when no
@@ -66,6 +67,52 @@ type QueueAffordanceProps = {
renderStatusBar?: (queueChip: ReactNode) => ReactNode;
};
type QueuePanelHandlerArgs = {
clearAll: () => Promise<void>;
drainNext: () => Promise<void>;
editEntry: (entryId: string, content: string) => Promise<void>;
removeEntry: (entryId: string) => Promise<void>;
};
function useQueuePanelHandlers({
clearAll,
drainNext,
editEntry,
removeEntry,
}: QueuePanelHandlerArgs) {
const handleSave = useCallback(
async (entryId: string, content: string) => {
await editEntry(entryId, content);
},
[editEntry],
);
const handleRemove = useCallback(
async (entryId: string) => {
try {
await removeEntry(entryId);
} catch (err) {
console.error("Failed to remove queued entry:", err);
toast.error("Failed to remove queued message.");
}
},
[removeEntry],
);
const handleClear = useCallback(() => {
clearAll().catch((err) => {
console.error("Failed to clear queued messages:", err);
toast.error("Failed to clear queued messages.");
});
}, [clearAll]);
const handleDrain = useCallback(() => {
drainNext().catch((err) => {
console.error("Failed to run queued message:", err);
toast.error("Failed to run queued message.");
});
}, [drainNext]);
return { handleSave, handleRemove, handleClear, handleDrain };
}
/**
* Wraps the chat input with the per-session queue affordance:
* - When there are no queued entries, just renders `children` (the input).
@@ -74,9 +121,21 @@ type QueueAffordanceProps = {
* it expands a panel above the input. Drained or session-switched queues
* auto-collapse.
*/
export function QueueAffordance({ sessionId, children, renderStatusBar }: QueueAffordanceProps) {
const { entries, count, max, isFull, clearAll, editEntry, removeEntry } = useQueue(sessionId);
export function QueueAffordance({
sessionId,
children,
canDrain = false,
renderStatusBar,
}: QueueAffordanceProps) {
const { entries, count, max, isFull, isLoading, clearAll, drainNext, editEntry, removeEntry } =
useQueue(sessionId);
const [isOpen, setIsOpen] = useState(false);
const { handleSave, handleRemove, handleClear, handleDrain } = useQueuePanelHandlers({
clearAll,
drainNext,
editEntry,
removeEntry,
});
// Reset disclosure on session switch or full drain using render-phase state
// adjustment (React docs: "Adjusting some state when a prop changes"). This
@@ -96,32 +155,6 @@ export function QueueAffordance({ sessionId, children, renderStatusBar }: QueueA
const close = useCallback(() => setIsOpen(false), []);
useEscToClose(isOpen, close);
const handleSave = useCallback(
async (entryId: string, content: string) => {
await editEntry(entryId, content);
},
[editEntry],
);
const handleRemove = useCallback(
async (entryId: string) => {
try {
await removeEntry(entryId);
} catch (err) {
console.error("Failed to remove queued entry:", err);
toast.error("Failed to remove queued message.");
}
},
[removeEntry],
);
const handleClear = useCallback(() => {
clearAll().catch((err) => {
console.error("Failed to clear queued messages:", err);
toast.error("Failed to clear queued messages.");
});
}, [clearAll]);
const hasEntries = !!sessionId && entryCount > 0;
const chipNode =
hasEntries && !isOpen ? (
@@ -158,8 +191,11 @@ export function QueueAffordance({ sessionId, children, renderStatusBar }: QueueA
count={count}
max={max}
isFull={isFull}
canDrain={canDrain}
isLoading={isLoading}
onClose={close}
onClear={handleClear}
onDrain={handleDrain}
onSave={handleSave}
onRemove={handleRemove}
/>
@@ -229,8 +265,11 @@ type QueuePanelProps = {
count: number;
max: number;
isFull: boolean;
canDrain: boolean;
isLoading: boolean;
onClose: () => void;
onClear: () => void;
onDrain: () => void;
onSave: (entryId: string, content: string) => Promise<void>;
onRemove: (entryId: string) => Promise<void>;
};
@@ -240,8 +279,11 @@ function QueuePanel({
count,
max,
isFull,
canDrain,
isLoading,
onClose,
onClear,
onDrain,
onSave,
onRemove,
}: QueuePanelProps) {
@@ -260,7 +302,10 @@ function QueuePanel({
count={count}
max={max}
isFull={isFull}
canDrain={canDrain}
isLoading={isLoading}
onClear={onClear}
onDrain={onDrain}
onClose={onClose}
/>
<div className="space-y-1.5">
@@ -283,11 +328,23 @@ type QueuePanelHeaderProps = {
count: number;
max: number;
isFull: boolean;
canDrain: boolean;
isLoading: boolean;
onClear: () => void;
onDrain: () => void;
onClose: () => void;
};
function QueuePanelHeader({ count, max, isFull, onClear, onClose }: QueuePanelHeaderProps) {
function QueuePanelHeader({
count,
max,
isFull,
canDrain,
isLoading,
onClear,
onDrain,
onClose,
}: QueuePanelHeaderProps) {
const capacityText = max > 0 ? `${count} of ${max}` : `${count}`;
return (
<div className="flex items-center justify-between gap-3 py-1">
@@ -300,6 +357,20 @@ function QueuePanelHeader({ count, max, isFull, onClear, onClose }: QueuePanelHe
</span>
</div>
<div className="flex items-center gap-1">
{canDrain && (
<Button
variant="ghost"
size="sm"
className="h-6 px-1.5 text-xs text-muted-foreground hover:text-foreground cursor-pointer"
onClick={onDrain}
disabled={isLoading}
title="Run next queued message"
data-testid="queue-drain-next"
>
<IconPlayerPlay className="mr-1 h-3 w-3" />
Run next
</Button>
)}
<Button
variant="ghost"
size="sm"
+23 -2
View File
@@ -3,6 +3,7 @@ import { useAppStore } from "@/components/state-provider";
import {
queueMessage,
clearQueue,
drainQueuedMessage,
getQueueStatus,
updateQueuedMessage,
removeQueuedEntry,
@@ -43,6 +44,23 @@ type QueueActionsArgs = {
metaMax: number | undefined;
};
function useDrainNextAction(
sessionId: string | null,
setQueueLoading: ReturnType<typeof useQueueState>["setQueueLoading"],
refetch: (sid: string) => Promise<void>,
) {
return useCallback(async () => {
if (!sessionId) return;
setQueueLoading(sessionId, true);
try {
await drainQueuedMessage(sessionId);
await refetch(sessionId);
} finally {
setQueueLoading(sessionId, false);
}
}, [sessionId, refetch, setQueueLoading]);
}
/** Build an action set bound to the supplied session + slice setters. */
function useQueueActions({
sessionId,
@@ -106,6 +124,8 @@ function useQueueActions({
}
}, [sessionId, setQueueEntries, setQueueLoading, metaMax]);
const drainNext = useDrainNextAction(sessionId, setQueueLoading, refetch);
const editEntry = useCallback(
async (entryId: string, content: string, attachments?: MessageAttachment[]) => {
if (!sessionId) return;
@@ -142,7 +162,7 @@ function useQueueActions({
[sessionId, refetch, removeQueueEntry],
);
return { refetch, queue, clearAll, editEntry, removeEntry };
return { refetch, queue, clearAll, drainNext, editEntry, removeEntry };
}
/**
@@ -156,7 +176,7 @@ function useQueueActions({
export function useQueue(sessionId: string | null) {
const state = useQueueState(sessionId);
const { entries, meta, isLoading } = state;
const { refetch, queue, clearAll, editEntry, removeEntry } = useQueueActions({
const { refetch, queue, clearAll, drainNext, editEntry, removeEntry } = useQueueActions({
sessionId,
setQueueEntries: state.setQueueEntries,
removeQueueEntry: state.removeQueueEntry,
@@ -184,6 +204,7 @@ export function useQueue(sessionId: string | null) {
isLoading,
queue,
clearAll,
drainNext,
editEntry,
removeEntry,
refetch: refetchBound,
+9
View File
@@ -94,6 +94,15 @@ export async function clearQueue(sessionId: string): Promise<{ removed: number }
return client.request<{ removed: number }>("message.queue.cancel", { session_id: sessionId });
}
/** Dispatch one queued entry now when the session is ready for input. */
export async function drainQueuedMessage(sessionId: string): Promise<{ drained: boolean }> {
const client = getWebSocketClient();
if (!client) {
throw new Error(WS_CLIENT_UNAVAILABLE);
}
return client.request<{ drained: boolean }>("message.queue.drain", { session_id: sessionId });
}
/** Fetch the full queue snapshot (entries + capacity). */
export async function getQueueStatus(sessionId: string): Promise<QueueStatus> {
const client = getWebSocketClient();