From 8b75e095e9ac0bc5d3270fa823621a1317f3be8f Mon Sep 17 00:00:00 2001 From: viettranx Date: Sun, 29 Mar 2026 09:00:06 +0700 Subject: [PATCH] feat(cancel): add run.cancelled event and fix stop across web UI, desktop, and channels - Add distinct `run.cancelled` agent event to differentiate user-initiated cancellation from real errors (previously both emitted `run.failed`) - Fix desktop stop button: wire onStop callback to InputBar, add abort() function in useChat hook, restore isRunning state on session switch - Fix web UI: handle run.cancelled without showing error message, preserve partial streamed content on cancel - Fix hanging chat.send promise: send {cancelled: true} OK response instead of silently dropping (was hanging 600s until timeout) - Fix channel stream cleanup: handle run.cancelled for stream stop + run context deletion (prevents memory leak) + reaction status clear - Fix child trace status: cancelled subagent runs now correctly get TraceStatusCancelled instead of TraceStatusError - Add run.cancelled to all event consumers: query invalidation, session detail refresh, trace dialog refresh, gateway activity clearing --- cmd/gateway.go | 2 +- internal/agent/loop_run.go | 33 +++- internal/agent/loop_types.go | 2 +- internal/agent/router_abort_test.go | 164 ++++++++++++++++++ internal/channels/events.go | 10 +- internal/gateway/methods/chat.go | 6 +- pkg/protocol/events.go | 1 + .../src/components/chat/ChatCanvas.tsx | 7 +- ui/desktop/frontend/src/hooks/use-chat.ts | 35 +++- ui/desktop/frontend/src/stores/chat-store.ts | 8 + ui/web/src/api/protocol.ts | 1 + ui/web/src/hooks/use-query-invalidation.ts | 2 +- .../src/pages/chat/hooks/use-chat-messages.ts | 31 ++++ .../pages/sessions/session-detail-page.tsx | 2 +- .../src/pages/traces/trace-detail-dialog.tsx | 2 +- ui/web/src/types/chat.ts | 2 +- 16 files changed, 287 insertions(+), 21 deletions(-) create mode 100644 internal/agent/router_abort_test.go diff --git a/cmd/gateway.go b/cmd/gateway.go index 5a6789ac..27bbe4dc 100644 --- a/cmd/gateway.go +++ b/cmd/gateway.go @@ -944,7 +944,7 @@ func runGateway() { } // Clear activity on terminal events - if agentEvent.Type == protocol.AgentEventRunCompleted || agentEvent.Type == protocol.AgentEventRunFailed { + if agentEvent.Type == protocol.AgentEventRunCompleted || agentEvent.Type == protocol.AgentEventRunFailed || agentEvent.Type == protocol.AgentEventRunCancelled { if sessionKey := agentRouter.SessionKeyForRun(agentEvent.RunID); sessionKey != "" { agentRouter.ClearActivity(sessionKey) } diff --git a/internal/agent/loop_run.go b/internal/agent/loop_run.go index 03e53756..619f3f80 100644 --- a/internal/agent/loop_run.go +++ b/internal/agent/loop_run.go @@ -143,18 +143,35 @@ func (l *Loop) Run(ctx context.Context, req RunRequest) (*RunResult, error) { if isChildTrace && l.traceCollector != nil && traceID != uuid.Nil { status := store.TraceStatusCompleted if err != nil { - status = store.TraceStatusError + if ctx.Err() != nil { + status = store.TraceStatusCancelled + } else { + status = store.TraceStatusError + } } - l.traceCollector.SetTraceStatus(ctx, traceID, status) + traceCtx := ctx + if ctx.Err() != nil { + traceCtx = context.WithoutCancel(ctx) + } + l.traceCollector.SetTraceStatus(traceCtx, traceID, status) } if err != nil { - emitRun(AgentEvent{ - Type: protocol.AgentEventRunFailed, - AgentID: l.id, - RunID: req.RunID, - Payload: map[string]string{"error": err.Error()}, - }) + // Distinguish user-initiated cancellation from real errors. + if ctx.Err() != nil { + emitRun(AgentEvent{ + Type: protocol.AgentEventRunCancelled, + AgentID: l.id, + RunID: req.RunID, + }) + } else { + emitRun(AgentEvent{ + Type: protocol.AgentEventRunFailed, + AgentID: l.id, + RunID: req.RunID, + Payload: map[string]string{"error": err.Error()}, + }) + } // Only finish trace for root runs; child traces don't own the trace lifecycle. // Use background context when the run context is cancelled (/stop command) // so the DB update still succeeds. diff --git a/internal/agent/loop_types.go b/internal/agent/loop_types.go index f77d833c..117a15a2 100644 --- a/internal/agent/loop_types.go +++ b/internal/agent/loop_types.go @@ -177,7 +177,7 @@ type Loop struct { // AgentEvent is emitted during agent execution for WS broadcasting. type AgentEvent struct { - Type string `json:"type"` // "run.started", "run.completed", "run.failed", "chunk", "tool.call", "tool.result" + Type string `json:"type"` // "run.started", "run.completed", "run.failed", "run.cancelled", "chunk", "tool.call", "tool.result" AgentID string `json:"agentId"` RunID string `json:"runId"` RunKind string `json:"runKind,omitempty"` // "delegation", "announce" — omitted for user-initiated runs diff --git a/internal/agent/router_abort_test.go b/internal/agent/router_abort_test.go new file mode 100644 index 00000000..cdbcf7ec --- /dev/null +++ b/internal/agent/router_abort_test.go @@ -0,0 +1,164 @@ +package agent + +import ( + "context" + "testing" +) + +// TestAbortRun_CancelsContext verifies that AbortRun calls the cancel func +// and removes the run from tracking. +func TestAbortRun_CancelsContext(t *testing.T) { + r := NewRouter() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + runID := "run-1" + sessionKey := "session-1" + r.RegisterRun(runID, sessionKey, "agent-1", cancel) + + if !r.IsSessionBusy(sessionKey) { + t.Fatal("session should be busy after RegisterRun") + } + + ok := r.AbortRun(runID, sessionKey) + if !ok { + t.Fatal("AbortRun should return true for active run") + } + if ctx.Err() == nil { + t.Fatal("context should be cancelled after AbortRun") + } + if r.IsSessionBusy(sessionKey) { + t.Fatal("session should not be busy after AbortRun") + } +} + +// TestAbortRun_WrongSessionKey verifies authorization check. +func TestAbortRun_WrongSessionKey(t *testing.T) { + r := NewRouter() + _, cancel := context.WithCancel(context.Background()) + defer cancel() + + r.RegisterRun("run-1", "session-1", "agent-1", cancel) + + ok := r.AbortRun("run-1", "wrong-session") + if ok { + t.Fatal("AbortRun should return false for mismatched sessionKey") + } + if !r.IsSessionBusy("session-1") { + t.Fatal("session should still be busy after failed abort") + } +} + +// TestAbortRun_NotFound returns false for unknown run IDs. +func TestAbortRun_NotFound(t *testing.T) { + r := NewRouter() + if r.AbortRun("nonexistent", "") { + t.Fatal("AbortRun should return false for unknown runID") + } +} + +// TestAbortRunsForSession_CancelsAllRuns verifies that all runs in a session +// are cancelled when AbortRunsForSession is called. +func TestAbortRunsForSession_CancelsAllRuns(t *testing.T) { + r := NewRouter() + ctx1, cancel1 := context.WithCancel(context.Background()) + defer cancel1() + ctx2, cancel2 := context.WithCancel(context.Background()) + defer cancel2() + + sessionKey := "session-multi" + r.RegisterRun("run-1", sessionKey, "agent-1", cancel1) + r.RegisterRun("run-2", sessionKey, "agent-1", cancel2) + + aborted := r.AbortRunsForSession(sessionKey) + if len(aborted) != 2 { + t.Fatalf("expected 2 aborted runs, got %d", len(aborted)) + } + if ctx1.Err() == nil || ctx2.Err() == nil { + t.Fatal("both contexts should be cancelled") + } + if r.IsSessionBusy(sessionKey) { + t.Fatal("session should not be busy after AbortRunsForSession") + } +} + +// TestAbortRunsForSession_EmptySession returns empty slice. +func TestAbortRunsForSession_EmptySession(t *testing.T) { + r := NewRouter() + aborted := r.AbortRunsForSession("no-such-session") + if len(aborted) != 0 { + t.Fatalf("expected 0 aborted runs, got %d", len(aborted)) + } +} + +// TestAbortRun_AfterUnregister verifies the race scenario: +// if run is unregistered before abort, abort returns false (idempotent). +func TestAbortRun_AfterUnregister(t *testing.T) { + r := NewRouter() + _, cancel := context.WithCancel(context.Background()) + defer cancel() + + r.RegisterRun("run-1", "session-1", "agent-1", cancel) + r.UnregisterRun("run-1") + + ok := r.AbortRun("run-1", "session-1") + if ok { + t.Fatal("AbortRun should return false after UnregisterRun") + } +} + +// TestUnregisterRun_ClearsSessionIndex verifies that UnregisterRun cleans up +// both activeRuns and sessionRuns. +func TestUnregisterRun_ClearsSessionIndex(t *testing.T) { + r := NewRouter() + _, cancel := context.WithCancel(context.Background()) + defer cancel() + + r.RegisterRun("run-1", "session-1", "agent-1", cancel) + if !r.IsSessionBusy("session-1") { + t.Fatal("session should be busy") + } + + r.UnregisterRun("run-1") + if r.IsSessionBusy("session-1") { + t.Fatal("session should not be busy after unregister") + } +} + +// TestSessionKeyForRun returns correct session key or empty. +func TestSessionKeyForRun(t *testing.T) { + r := NewRouter() + _, cancel := context.WithCancel(context.Background()) + defer cancel() + + r.RegisterRun("run-1", "session-1", "agent-1", cancel) + + if got := r.SessionKeyForRun("run-1"); got != "session-1" { + t.Fatalf("expected session-1, got %q", got) + } + if got := r.SessionKeyForRun("nonexistent"); got != "" { + t.Fatalf("expected empty string, got %q", got) + } +} + +// TestActivityClearedOnAbort verifies that activity is cleared when the gateway +// subscriber processes a terminal event. We test the Router methods directly. +func TestActivityClearedOnAbort(t *testing.T) { + r := NewRouter() + _, cancel := context.WithCancel(context.Background()) + defer cancel() + + sessionKey := "session-act" + r.RegisterRun("run-act", sessionKey, "agent-1", cancel) + r.UpdateActivity(sessionKey, "run-act", "thinking", "", 0) + + if r.GetActivity(sessionKey) == nil { + t.Fatal("activity should be set") + } + + // Simulate what gateway.go does on terminal events + r.ClearActivity(sessionKey) + if r.GetActivity(sessionKey) != nil { + t.Fatal("activity should be cleared after ClearActivity") + } +} diff --git a/internal/channels/events.go b/internal/channels/events.go index f791ee1c..0b10c144 100644 --- a/internal/channels/events.go +++ b/internal/channels/events.go @@ -15,7 +15,7 @@ import ( // HandleAgentEvent routes agent lifecycle events to streaming/reaction channels. // Called from the bus event subscriber — must be non-blocking. -// eventType: "run.started", "chunk", "tool.call", "tool.result", "run.completed", "run.failed" +// eventType: "run.started", "chunk", "tool.call", "tool.result", "run.completed", "run.failed", "run.cancelled" func (m *Manager) HandleAgentEvent(eventType, runID string, payload any) { val, ok := m.runs.Load(runID) if !ok { @@ -230,8 +230,8 @@ func (m *Manager) HandleAgentEvent(eventType, runID string, payload any) { } sc.FinalizeStream(ctx, rc.ChatID, currentStream) } - case protocol.AgentEventRunFailed: - // Clean up streaming state on failure + case protocol.AgentEventRunFailed, protocol.AgentEventRunCancelled: + // Clean up streaming state on failure or cancellation rc.mu.Lock() currentStream := rc.stream rc.stream = nil @@ -315,6 +315,8 @@ func (m *Manager) HandleAgentEvent(eventType, runID string, payload any) { status = "done" case protocol.AgentEventRunFailed: status = "error" + case protocol.AgentEventRunCancelled: + status = "done" } if status != "" { if err := reactionCh.OnReactionEvent(ctx, rc.ChatID, rc.MessageID, status); err != nil { @@ -324,7 +326,7 @@ func (m *Manager) HandleAgentEvent(eventType, runID string, payload any) { } // Clean up on terminal events - if eventType == protocol.AgentEventRunCompleted || eventType == protocol.AgentEventRunFailed { + if eventType == protocol.AgentEventRunCompleted || eventType == protocol.AgentEventRunFailed || eventType == protocol.AgentEventRunCancelled { m.runs.Delete(runID) } } diff --git a/internal/gateway/methods/chat.go b/internal/gateway/methods/chat.go index 26d63ec9..42891196 100644 --- a/internal/gateway/methods/chat.go +++ b/internal/gateway/methods/chat.go @@ -242,8 +242,12 @@ func (m *ChatMethods) handleSend(ctx context.Context, client *gateway.Client, re }) if err != nil { - // Don't send error if context was cancelled (abort) + // Send cancelled response so the frontend's chat.send promise resolves + // instead of hanging until the 600s timeout. if runCtx.Err() != nil { + client.SendResponse(protocol.NewOKResponse(req.ID, map[string]any{ + "cancelled": true, + })) return } client.SendResponse(protocol.NewErrorResponse(req.ID, protocol.ErrInternal, err.Error())) diff --git a/pkg/protocol/events.go b/pkg/protocol/events.go index d4cb89ad..c9ed2e68 100644 --- a/pkg/protocol/events.go +++ b/pkg/protocol/events.go @@ -108,6 +108,7 @@ const ( AgentEventRunStarted = "run.started" AgentEventRunCompleted = "run.completed" AgentEventRunFailed = "run.failed" + AgentEventRunCancelled = "run.cancelled" AgentEventRunRetrying = "run.retrying" AgentEventToolCall = "tool.call" AgentEventToolResult = "tool.result" diff --git a/ui/desktop/frontend/src/components/chat/ChatCanvas.tsx b/ui/desktop/frontend/src/components/chat/ChatCanvas.tsx index 9f57f47b..860609f9 100644 --- a/ui/desktop/frontend/src/components/chat/ChatCanvas.tsx +++ b/ui/desktop/frontend/src/components/chat/ChatCanvas.tsx @@ -14,7 +14,7 @@ import type { TeamTaskData } from '../../types/team' export function ChatCanvas() { const { t } = useTranslation('common') - const { messages, isRunning, activity, sendMessage } = useChat() + const { messages, isRunning, activity, sendMessage, abort } = useChat() const { selectedAgent } = useAgents() const { members, fetchTaskDetail } = useTeamTasks() const activeSessionKey = useSessionStore((s) => s.activeSessionKey) @@ -50,6 +50,10 @@ export function ChatCanvas() { sendMessage(text, selectedAgent.id, files) }, [selectedAgent, sendMessage]) + const handleStop = useCallback(() => { + abort() + }, [abort]) + const hasMessages = messages.length > 0 return ( @@ -95,6 +99,7 @@ export function ChatCanvas() { {/* Input bar */} { + const sk = sessionKeyRef.current + if (!ws || !sk) return + try { + await ws.call('chat.abort', { sessionKey: sk }) + } catch { + // ignore abort errors + } + }, [ws]) + // Reset streaming state + load history when session changes. // Messages are NOT cleared here — loadHistory() replaces them atomically. // Clearing would race with sendMessage's auto-session-creation flow. @@ -373,8 +392,21 @@ export function useChat() { loadHistory(activeSessionKey).then(() => { if (cancelled) return }) + + // Restore session running state (mirrors web UI pattern). + // If user switches to a session with an active run, show the stop button. + ws?.call('chat.session_status', { sessionKey: activeSessionKey }) + .then((raw: unknown) => { + if (cancelled) return + const res = raw as { isRunning?: boolean; activity?: { phase: string; tool?: string; iteration?: number } } + if (res?.isRunning) { + useChatStore.getState().restoreRunning(res.activity ?? null) + } + }) + .catch(() => {}) + return () => { cancelled = true } - }, [activeSessionKey, loadHistory, chunkBatcher, thinkingBatcher]) + }, [activeSessionKey, ws, loadHistory, chunkBatcher, thinkingBatcher]) return { messages, @@ -382,5 +414,6 @@ export function useChat() { activity, sendMessage, loadHistory, + abort, } } diff --git a/ui/desktop/frontend/src/stores/chat-store.ts b/ui/desktop/frontend/src/stores/chat-store.ts index b55cd7df..b48b3b7f 100644 --- a/ui/desktop/frontend/src/stores/chat-store.ts +++ b/ui/desktop/frontend/src/stores/chat-store.ts @@ -42,6 +42,8 @@ interface ChatState { setActivity: (activity: Activity | null) => void completeRun: (content: string, usage?: Message['usage'], media?: Message['media']) => void failRun: (error: string) => void + cancelRun: () => void + restoreRunning: (activity?: Activity | null) => void setMessages: (messages: Message[]) => void clear: () => void } @@ -160,6 +162,12 @@ export const useChatStore = create((set) => ({ }) }, + // Cancellation: preserve partial streamed content, just clear running state. + cancelRun: () => set({ isRunning: false, activity: null, currentRunId: null }), + + // Restore running state on session switch (without creating a new assistant message). + restoreRunning: (activity) => set({ isRunning: true, activity: activity ?? null }), + setMessages: (messages) => set({ messages }), clear: () => set({ messages: [], isRunning: false, activity: null, currentRunId: null }), })) diff --git a/ui/web/src/api/protocol.ts b/ui/web/src/api/protocol.ts index e23bce40..919479a4 100644 --- a/ui/web/src/api/protocol.ts +++ b/ui/web/src/api/protocol.ts @@ -279,6 +279,7 @@ export const AgentEventTypes = { RUN_STARTED: "run.started", RUN_COMPLETED: "run.completed", RUN_FAILED: "run.failed", + RUN_CANCELLED: "run.cancelled", TOOL_CALL: "tool.call", TOOL_RESULT: "tool.result", BLOCK_REPLY: "block.reply", diff --git a/ui/web/src/hooks/use-query-invalidation.ts b/ui/web/src/hooks/use-query-invalidation.ts index 2bccc64b..a6d47da7 100644 --- a/ui/web/src/hooks/use-query-invalidation.ts +++ b/ui/web/src/hooks/use-query-invalidation.ts @@ -21,7 +21,7 @@ export function useWsQueryInvalidation() { if (event.type === "run.started") { queryClient.invalidateQueries({ queryKey: queryKeys.traces.all }); } - if (event.type === "run.completed" || event.type === "run.failed") { + if (event.type === "run.completed" || event.type === "run.failed" || event.type === "run.cancelled") { queryClient.invalidateQueries({ queryKey: queryKeys.sessions.all }); queryClient.invalidateQueries({ queryKey: queryKeys.traces.all }); queryClient.invalidateQueries({ queryKey: queryKeys.usage.all }); diff --git a/ui/web/src/pages/chat/hooks/use-chat-messages.ts b/ui/web/src/pages/chat/hooks/use-chat-messages.ts index fdf21bd9..c97e431a 100644 --- a/ui/web/src/pages/chat/hooks/use-chat-messages.ts +++ b/ui/web/src/pages/chat/hooks/use-chat-messages.ts @@ -385,6 +385,37 @@ export function useChatMessages(sessionKey: string, agentId: string) { ]); break; } + // User-initiated cancellation — clear state, preserve partial content. + case "run.cancelled": { + cancelAnimationFrame(rafHandleRef.current); + rafPendingRef.current = false; + + setIsRunning(false); + runIdRef.current = null; + + const streamed = streamRef.current; + setStreamText(null); + setThinkingText(null); + setToolStream([]); + streamRef.current = ""; + thinkingRef.current = ""; + toolStreamRef.current = []; + activityRef.current = null; + setActivity(null); + blockRepliesRef.current = []; + setBlockReplies([]); + + // Promote partial streamed text or reload history + if (streamed) { + setMessages((prev) => [ + ...prev, + { role: "assistant", content: streamed, timestamp: Date.now() }, + ]); + } else { + loadHistory(); + } + break; + } } }, [loadHistory], diff --git a/ui/web/src/pages/sessions/session-detail-page.tsx b/ui/web/src/pages/sessions/session-detail-page.tsx index a142afb8..3bc38032 100644 --- a/ui/web/src/pages/sessions/session-detail-page.tsx +++ b/ui/web/src/pages/sessions/session-detail-page.tsx @@ -118,7 +118,7 @@ export function SessionDetailPage({ const event = payload as AgentEventPayload; if (!event) return; if ( - (event.type === "run.completed" || event.type === "run.failed") && + (event.type === "run.completed" || event.type === "run.failed" || event.type === "run.cancelled") && event.agentId === parsed.agentId ) { debouncedRefresh(); diff --git a/ui/web/src/pages/traces/trace-detail-dialog.tsx b/ui/web/src/pages/traces/trace-detail-dialog.tsx index b58775bf..b456830e 100644 --- a/ui/web/src/pages/traces/trace-detail-dialog.tsx +++ b/ui/web/src/pages/traces/trace-detail-dialog.tsx @@ -124,7 +124,7 @@ export function TraceDetailDialog({ traceId, onClose, getTrace, onNavigateTrace, useCallback( (payload: unknown) => { const event = payload as AgentEventPayload; - if (event?.type === "run.completed" || event?.type === "run.failed") { + if (event?.type === "run.completed" || event?.type === "run.failed" || event?.type === "run.cancelled") { fetchTrace(); } }, diff --git a/ui/web/src/types/chat.ts b/ui/web/src/types/chat.ts index fb82abbb..9873aecb 100644 --- a/ui/web/src/types/chat.ts +++ b/ui/web/src/types/chat.ts @@ -48,7 +48,7 @@ export interface ChatMessage extends Message { /** Agent event payload from WS event "agent" */ export interface AgentEventPayload { - type: string; // "run.started" | "run.completed" | "run.failed" | "chunk" | "tool.call" | "tool.result" | "activity" | "block.reply" | "run.retrying" + type: string; // "run.started" | "run.completed" | "run.failed" | "run.cancelled" | "chunk" | "tool.call" | "tool.result" | "activity" | "block.reply" | "run.retrying" agentId: string; runId: string; runKind?: string; // "delegation" | "announce" — omitted for user-initiated runs