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:
@@ -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) {
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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();
|
||||
|
||||
Reference in New Issue
Block a user