From bbb862a41addfc26d31313de553d388763e1bdb3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=A8=8B=E5=BA=8F=E5=91=98=E9=98=BF=E6=B1=9F=28Relakkes?= =?UTF-8?q?=29?= Date: Fri, 10 Jul 2026 04:05:19 +0800 Subject: [PATCH] fix(runtime): recover stalled long-running turns Retry idle provider streams only before tool side effects, and reconcile desktop turn state after WebSocket gaps so completed or still-running work cannot leave a stale spinner. Tested: 108 runtime/server focused tests, 128 desktop focused tests Tested: disconnected-turn WebSocket integration test Tested: bun run check:server (1489 pass) Tested: desktop typecheck and production build Not-tested: full desktop gate has one unrelated TraceSession polling assertion failure Scope-risk: broad --- desktop/src/api/websocket.test.ts | 53 +++++ desktop/src/api/websocket.ts | 68 +++++- desktop/src/stores/chatStore.test.ts | 216 ++++++++++++++++++ desktop/src/stores/chatStore.ts | 167 +++++++++++++- desktop/src/types/chat.ts | 8 +- src/entrypoints/sdk/coreSchemas.ts | 4 +- .../__tests__/conversation-service.test.ts | 53 +++++ src/server/__tests__/conversations.test.ts | 90 ++++++++ .../__tests__/websocket-handler.test.ts | 104 +++++++++ src/server/__tests__/ws-memory-events.test.ts | 162 ++++++++++++- src/server/services/conversationService.ts | 41 ++-- src/server/ws/events.ts | 6 +- src/server/ws/handler.ts | 57 ++++- src/services/api/claude.ts | 49 ++-- .../api/streamAssistantCommitBuffer.test.ts | 43 ++++ .../api/streamAssistantCommitBuffer.ts | 35 +++ src/services/api/streamRetry.test.ts | 8 + src/services/api/streamRetry.ts | 16 +- src/services/api/streamWatchdog.test.ts | 31 ++- src/services/api/streamWatchdog.ts | 32 ++- src/services/api/withRetry.test.ts | 6 + src/services/api/withRetry.ts | 5 +- src/utils/messages.ts | 12 +- 23 files changed, 1189 insertions(+), 77 deletions(-) create mode 100644 src/services/api/streamAssistantCommitBuffer.test.ts create mode 100644 src/services/api/streamAssistantCommitBuffer.ts diff --git a/desktop/src/api/websocket.test.ts b/desktop/src/api/websocket.test.ts index 19402dd5..60b5a2ab 100644 --- a/desktop/src/api/websocket.test.ts +++ b/desktop/src/api/websocket.test.ts @@ -52,6 +52,12 @@ class FakeWebSocket { this.readyState = FakeWebSocket.CLOSED ;(this.onclose as (() => void) | null)?.() } + + receive(message: unknown) { + ;(this.onmessage as ((event: { data: string }) => void) | null)?.({ + data: JSON.stringify(message), + }) + } } describe('wsManager reconnect buffering', () => { @@ -95,9 +101,56 @@ describe('wsManager reconnect buffering', () => { expect(secondSocket!.sent).toEqual([ JSON.stringify({ type: 'user_message', content: 'queued while offline' }), + JSON.stringify({ type: 'sync_state' }), ]) }) + it('closes and reconnects a half-open socket when a pong never arrives', async () => { + wsManager.connect('session-half-open') + const firstSocket = FakeWebSocket.instances[0]! + firstSocket.open() + + await vi.advanceTimersByTimeAsync(30_000) + expect(firstSocket.sent).toContain(JSON.stringify({ type: 'ping' })) + + await vi.advanceTimersByTimeAsync(10_000) + expect(firstSocket.readyState).toBe(FakeWebSocket.CLOSED) + + await vi.advanceTimersByTimeAsync(1000) + expect(FakeWebSocket.instances).toHaveLength(2) + }) + + it('keeps a healthy socket open when the server answers the heartbeat', async () => { + wsManager.connect('session-heartbeat') + const socket = FakeWebSocket.instances[0]! + socket.open() + + await vi.advanceTimersByTimeAsync(30_000) + socket.receive({ type: 'pong' }) + await vi.advanceTimersByTimeAsync(10_000) + + expect(socket.readyState).toBe(FakeWebSocket.OPEN) + expect(FakeWebSocket.instances).toHaveLength(1) + }) + + it('does not let a stale socket close stop the replacement heartbeat', async () => { + wsManager.connect('session-stale-close') + const firstSocket = FakeWebSocket.instances[0]! + firstSocket.open() + firstSocket.fail() + + await vi.advanceTimersByTimeAsync(1000) + const secondSocket = FakeWebSocket.instances[1]! + secondSocket.open() + + // A browser may deliver a duplicate or delayed close callback for the + // replaced socket. It must not clear the replacement connection's timer. + firstSocket.fail() + await vi.advanceTimersByTimeAsync(30_000) + + expect(secondSocket.sent).toContain(JSON.stringify({ type: 'ping' })) + }) + it('builds websocket URLs from http and encodes token query params', () => { clientMocks.baseUrl = 'http://10.0.0.2:3456' clientMocks.authToken = 'h5 token/with?chars' diff --git a/desktop/src/api/websocket.ts b/desktop/src/api/websocket.ts index 06f98b67..fb7a49fd 100644 --- a/desktop/src/api/websocket.ts +++ b/desktop/src/api/websocket.ts @@ -3,12 +3,16 @@ import { getAuthToken, getBaseUrl } from './client' type MessageHandler = (msg: ServerMessage) => void +const HEARTBEAT_INTERVAL_MS = 30_000 +const HEARTBEAT_TIMEOUT_MS = 10_000 + type Connection = { ws: WebSocket handlers: Set reconnectTimer: ReturnType | null reconnectAttempt: number pingInterval: ReturnType | null + pongTimeout: ReturnType | null intentionalClose: boolean pendingMessages: ClientMessage[] } @@ -47,23 +51,34 @@ class WebSocketManager { reconnectTimer: null, reconnectAttempt: existing?.reconnectAttempt ?? 0, pingInterval: null, + pongTimeout: null, intentionalClose: false, pendingMessages: existing?.pendingMessages ?? [], } this.connections.set(sessionId, conn) ws.onopen = () => { + const isReconnect = conn.reconnectAttempt > 0 conn.reconnectAttempt = 0 - this.startPingLoop(sessionId) + this.startPingLoop(sessionId, conn) while (conn.pendingMessages.length > 0) { const msg = conn.pendingMessages.shift()! ws.send(JSON.stringify(msg)) } + // Ask for authoritative turn state only on an automatic reconnect. This + // is deliberately queued after pending user messages so the server sees + // those turns before deciding whether the session is running or idle. + if (isReconnect) { + ws.send(JSON.stringify({ type: 'sync_state' } satisfies ClientMessage)) + } } ws.onmessage = (event) => { try { const msg = JSON.parse(event.data as string) as ServerMessage + if (msg.type === 'pong') { + this.clearPongTimeout(conn) + } for (const handler of conn.handlers) { handler(msg) } @@ -73,7 +88,7 @@ class WebSocketManager { } ws.onclose = () => { - this.stopPingLoop(sessionId) + this.stopPingLoopForConnection(conn) if (!conn.intentionalClose && this.connections.get(sessionId) === conn) { this.scheduleReconnect(sessionId, conn) } @@ -89,7 +104,7 @@ class WebSocketManager { if (!conn) return conn.intentionalClose = true - this.stopPingLoop(sessionId) + this.stopPingLoopForConnection(conn) if (conn.reconnectTimer) { clearTimeout(conn.reconnectTimer) conn.reconnectTimer = null @@ -143,21 +158,50 @@ class WebSocketManager { if (conn) conn.handlers.clear() } - private startPingLoop(sessionId: string) { - this.stopPingLoop(sessionId) - const conn = this.connections.get(sessionId) - if (!conn) return + private startPingLoop(sessionId: string, conn: Connection) { + this.stopPingLoopForConnection(conn) + if (this.connections.get(sessionId) !== conn) return conn.pingInterval = setInterval(() => { - this.send(sessionId, { type: 'ping' }) - }, 30_000) + if ( + this.connections.get(sessionId) !== conn || + conn.ws.readyState !== WebSocket.OPEN + ) { + return + } + + try { + conn.ws.send(JSON.stringify({ type: 'ping' } satisfies ClientMessage)) + } catch { + conn.ws.close() + return + } + + this.clearPongTimeout(conn) + conn.pongTimeout = setTimeout(() => { + conn.pongTimeout = null + if ( + this.connections.get(sessionId) === conn && + conn.ws.readyState === WebSocket.OPEN + ) { + conn.ws.close() + } + }, HEARTBEAT_TIMEOUT_MS) + }, HEARTBEAT_INTERVAL_MS) } - private stopPingLoop(sessionId: string) { - const conn = this.connections.get(sessionId) - if (conn?.pingInterval) { + private stopPingLoopForConnection(conn: Connection) { + if (conn.pingInterval) { clearInterval(conn.pingInterval) conn.pingInterval = null } + this.clearPongTimeout(conn) + } + + private clearPongTimeout(conn: Connection) { + if (conn.pongTimeout) { + clearTimeout(conn.pongTimeout) + conn.pongTimeout = null + } } private scheduleReconnect(sessionId: string, conn: Connection) { diff --git a/desktop/src/stores/chatStore.test.ts b/desktop/src/stores/chatStore.test.ts index c1709820..430d5996 100644 --- a/desktop/src/stores/chatStore.test.ts +++ b/desktop/src/stores/chatStore.test.ts @@ -3738,6 +3738,155 @@ describe('chatStore history mapping', () => { vi.useRealTimers() }) + it('reloads authoritative history when a reconnect finds the turn already idle', async () => { + vi.mocked(sessionsApi.getMessages).mockClear() + vi.mocked(sessionsApi.getMessages).mockResolvedValueOnce({ + messages: [ + { + id: 'completed-assistant', + type: 'assistant', + timestamp: '2026-07-10T00:00:00.000Z', + content: [{ type: 'text', text: 'Finished while the socket was offline.' }], + }, + ], + }) + useChatStore.setState({ + sessions: { + [TEST_SESSION_ID]: makeSession({ + chatState: 'thinking', + streamingText: 'stale partial', + messages: [{ id: 'user-1', type: 'user_text', content: 'long task', timestamp: 1 }], + }), + }, + }) + + useChatStore.getState().handleServerMessage(TEST_SESSION_ID, { + type: 'session_state', + turnState: 'idle', + }) + + await vi.waitFor(() => { + expect(useChatStore.getState().sessions[TEST_SESSION_ID]?.messages).toContainEqual( + expect.objectContaining({ + type: 'assistant_text', + content: 'Finished while the socket was offline.', + }), + ) + }) + const session = useChatStore.getState().sessions[TEST_SESSION_ID] + expect(sessionsApi.getMessages).toHaveBeenCalledWith(TEST_SESSION_ID) + expect(session?.streamingText).toBe('') + expect(session?.messages).toContainEqual(expect.objectContaining({ + type: 'assistant_text', + content: 'Finished while the socket was offline.', + })) + }) + + it('does not let delayed reconnect history overwrite a newly sent turn', async () => { + let resolveHistory!: (value: { messages: MessageEntry[] }) => void + vi.mocked(sessionsApi.getMessages).mockReturnValueOnce(new Promise((resolve) => { + resolveHistory = resolve + })) + useChatStore.setState({ + sessions: { + [TEST_SESSION_ID]: makeSession({ + chatState: 'thinking', + streamingText: 'old partial', + messages: [{ id: 'old-user', type: 'user_text', content: 'old turn', timestamp: 1 }], + }), + }, + }) + + useChatStore.getState().handleServerMessage(TEST_SESSION_ID, { + type: 'session_state', + turnState: 'idle', + }) + await vi.waitFor(() => { + expect(sessionsApi.getMessages).toHaveBeenCalledWith(TEST_SESSION_ID) + }) + useChatStore.getState().sendMessage(TEST_SESSION_ID, 'new turn') + + resolveHistory({ + messages: [{ + id: 'old-assistant', + type: 'assistant', + timestamp: '2026-07-10T00:00:00.000Z', + content: [{ type: 'text', text: 'old completed answer' }], + }], + }) + await new Promise((resolve) => setTimeout(resolve, 0)) + + const session = useChatStore.getState().sessions[TEST_SESSION_ID] + expect(session?.chatState).toBe('thinking') + expect(session?.messages).toContainEqual(expect.objectContaining({ + type: 'user_text', + content: 'new turn', + })) + expect(session?.messages).not.toContainEqual(expect.objectContaining({ + type: 'assistant_text', + content: 'old completed answer', + })) + if (session?.elapsedTimer) clearInterval(session.elapsedTimer) + }) + + it('keeps the turn running but discards stale partials when reconnect reconciliation says running', async () => { + vi.mocked(sessionsApi.getMessages).mockClear() + vi.mocked(sessionsApi.getMessages).mockResolvedValueOnce({ messages: [] }) + useChatStore.setState({ + sessions: { + [TEST_SESSION_ID]: makeSession({ + chatState: 'streaming', + streamingText: 'still arriving', + streamingToolInput: '{"stale":', + activeToolUseId: 'stale-tool', + }), + }, + }) + + useChatStore.getState().handleServerMessage(TEST_SESSION_ID, { + type: 'session_state', + turnState: 'running', + }) + await vi.waitFor(() => { + expect(sessionsApi.getMessages).toHaveBeenCalledWith(TEST_SESSION_ID) + }) + + expect(useChatStore.getState().sessions[TEST_SESSION_ID]).toMatchObject({ + chatState: 'thinking', + streamingText: '', + streamingToolInput: '', + activeToolUseId: null, + }) + }) + + it('keeps the tab running for background agents when reconnect reconciliation finds the foreground idle', () => { + useChatStore.setState({ + sessions: { + [TEST_SESSION_ID]: makeSession({ + chatState: 'thinking', + backgroundAgentTasks: { + 'agent-task-1': { + taskId: 'agent-task-1', + toolUseId: 'agent-tool-1', + status: 'running', + taskType: 'local_agent', + description: 'Review screenshots', + startedAt: 1, + updatedAt: 2, + }, + }, + }), + }, + }) + + useChatStore.getState().handleServerMessage(TEST_SESSION_ID, { + type: 'session_state', + turnState: 'idle', + }) + + expect(updateTabStatusMock).toHaveBeenLastCalledWith(TEST_SESSION_ID, 'running') + }) + it('resumes the elapsed timer when streaming continues after the timer was lost', () => { vi.useFakeTimers() @@ -4101,6 +4250,73 @@ describe('chatStore history mapping', () => { expect(useChatStore.getState().sessions[TEST_SESSION_ID]?.streamingFallback).toBeNull() }) + it('discards only the failed stream attempt before a safe retry', () => { + const completedTool = { + id: 'completed-tool', + type: 'tool_use' as const, + toolName: 'Read', + toolUseId: 'read-1', + input: { file_path: 'README.md' }, + timestamp: 1, + isPending: false, + } + const completedResult = { + id: 'completed-result', + type: 'tool_result' as const, + toolUseId: 'read-1', + content: 'ok', + isError: false, + timestamp: 2, + } + useChatStore.setState({ + sessions: { + [TEST_SESSION_ID]: makeSession({ + messages: [ + completedTool, + completedResult, + { id: 'failed-thinking', type: 'thinking', content: 'partial thought', timestamp: 3 }, + { + id: 'failed-tool', + type: 'tool_use', + toolName: 'Write', + toolUseId: 'write-partial', + input: {}, + timestamp: 4, + isPending: true, + partialInput: '{"file_path":', + }, + ], + chatState: 'tool_executing', + streamingText: 'partial answer', + streamingToolInput: '{"file_path":', + activeToolUseId: 'write-partial', + activeToolName: 'Write', + activeThinkingId: 'failed-thinking', + streamingResponseChars: 200, + streamAttemptStartIndex: 2, + streamAttemptStartResponseChars: 80, + }), + }, + }) + + useChatStore.getState().handleServerMessage(TEST_SESSION_ID, { + type: 'streaming_fallback', + cause: 'stream_retry', + }) + + expect(useChatStore.getState().sessions[TEST_SESSION_ID]).toMatchObject({ + messages: [completedTool, completedResult], + chatState: 'thinking', + streamingText: '', + streamingToolInput: '', + activeToolUseId: null, + activeToolName: null, + activeThinkingId: null, + streamingResponseChars: 80, + streamingFallback: null, + }) + }) + it('keeps the fallback notice when idle and clears it on turn completion', () => { useChatStore.setState({ sessions: { diff --git a/desktop/src/stores/chatStore.ts b/desktop/src/stores/chatStore.ts index 839d498b..c220cbff 100644 --- a/desktop/src/stores/chatStore.ts +++ b/desktop/src/stores/chatStore.ts @@ -116,10 +116,13 @@ export type PerSessionState = { * indicator — same estimation the CLI spinner uses. Reset on each send. */ streamingResponseChars: number + /** Boundary used to discard one failed, side-effect-free stream attempt. */ + streamAttemptStartIndex?: number + streamAttemptStartResponseChars?: number elapsedSeconds: number statusVerb: string apiRetry?: ApiRetryState | null - // 流式→非流式降级提示（活动回合状态，与 apiRetry 同清除时机）。 + // 流式恢复/非流式降级提示（活动回合状态，与 apiRetry 同清除时机）。 streamingFallback?: StreamingFallbackState | null slashCommands: Array<{ name: string; description: string; argumentHint?: string }> agentTaskNotifications: Record @@ -280,7 +283,13 @@ type ChatStore = { setSessionPermissionMode: (sessionId: string, mode: PermissionMode) => void stopGeneration: (sessionId: string) => void loadHistory: (sessionId: string) => Promise - reloadHistory: (sessionId: string) => Promise + reloadHistory: ( + sessionId: string, + guard?: { + messages: UIMessage[] + backgroundAgentTasks?: Record + }, + ) => Promise queueComposerPrefill: ( sessionId: string, prefill: { text: string; attachments?: UIAttachment[]; mode?: ComposerPrefillMode }, @@ -1426,7 +1435,7 @@ export const useChatStore = create((set, get) => ({ return load }, - reloadHistory: async (sessionId) => { + reloadHistory: async (sessionId, guard) => { try { const { uiMessages, @@ -1438,6 +1447,18 @@ export const useChatStore = create((set, get) => ({ tokenUsage, } = await fetchAndMapSessionHistory(sessionId) + if (guard) { + const current = get().sessions[sessionId] + if ( + !current || + current.chatState !== 'idle' || + current.messages !== guard.messages || + current.backgroundAgentTasks !== guard.backgroundAgentTasks + ) { + return + } + } + set((state) => { const session = state.sessions[sessionId] if (!session) return state @@ -1682,6 +1703,100 @@ export const useChatStore = create((set, get) => ({ case 'connected': break + case 'session_state': { + const session = get().sessions[sessionId] + if (!session) break + + if (msg.turnState === 'running') { + // Raw deltas are not replayable across a socket gap. Discard the + // uncommitted attempt instead of appending new deltas (or a missed + // stream_retry attempt) to stale text/tool JSON. Persisted completed + // messages are merged back below while the turn remains running. + consumePendingDelta(sessionId) + clearPendingToolInputDelta(sessionId) + clearPendingTaskToolUseIds(sessionId) + clearPendingToolParentUseIds(sessionId) + update((current) => { + const startIndex = Math.max( + 0, + Math.min( + current.streamAttemptStartIndex ?? current.messages.length, + current.messages.length, + ), + ) + return { + messages: [ + ...current.messages.slice(0, startIndex), + ...current.messages.slice(startIndex).filter((message) => + message.type !== 'assistant_text' && + message.type !== 'thinking' && + !(message.type === 'tool_use' && message.isPending)), + ], + chatState: 'thinking', + streamingText: '', + streamingToolInput: '', + activeThinkingId: null, + activeToolUseId: null, + activeToolName: null, + streamingResponseChars: + current.streamAttemptStartResponseChars ?? current.streamingResponseChars, + streamAttemptStartIndex: undefined, + streamAttemptStartResponseChars: undefined, + apiRetry: null, + streamingFallback: null, + statusVerb: '', + } + }) + useTabStore.getState().updateTabStatus(sessionId, 'running') + ensureElapsedTimer() + void get().loadHistory(sessionId) + break + } + + if (session.chatState === 'idle') break + + const text = `${session.streamingText}${consumePendingDelta(sessionId)}` + clearPendingToolInputDelta(sessionId) + clearPendingTaskToolUseIds(sessionId) + clearPendingToolParentUseIds(sessionId) + if (session.elapsedTimer) clearInterval(session.elapsedTimer) + const messagesWithText = text.trim() + ? appendAssistantTextMessage(session.messages, text, Date.now()) + : session.messages + update(() => ({ + messages: markPendingToolUseMessagesStopped(messagesWithText), + chatState: 'idle', + activeThinkingId: null, + activeToolUseId: null, + activeToolName: null, + pendingPermission: null, + pendingComputerUsePermission: null, + elapsedTimer: null, + statusVerb: '', + apiRetry: null, + streamingFallback: null, + streamingText: '', + streamingToolInput: '', + })) + const reconciledSession = get().sessions[sessionId] + const hasRunningBackgroundAgents = hasRunningBackgroundTasks( + reconciledSession?.backgroundAgentTasks, + ) + useTabStore.getState().updateTabStatus( + sessionId, + hasRunningBackgroundAgents ? 'running' : 'idle', + ) + // The terminal event may have arrived while this renderer was offline. + // Replace optimistic/partial state with the persisted transcript. + if (reconciledSession) { + void get().reloadHistory(sessionId, { + messages: reconciledSession.messages, + backgroundAgentTasks: reconciledSession.backgroundAgentTasks, + }) + } + break + } + case 'status': update((session) => { const pendingText = `${session.streamingText}${consumePendingDelta(sessionId)}` @@ -1718,6 +1833,10 @@ export const useChatStore = create((set, get) => ({ : '', ...(msg.state === 'idle' ? { activeThinkingId: null } : {}), ...(msg.state === 'idle' ? { apiRetry: null, streamingFallback: null } : {}), + ...(msg.attemptStart ? { + streamAttemptStartIndex: session.messages.length, + streamAttemptStartResponseChars: session.streamingResponseChars, + } : {}), ...(nextMessages !== session.messages ? { messages: nextMessages } : {}), ...(shouldFlush ? { streamingText: '', @@ -1837,6 +1956,48 @@ export const useChatStore = create((set, get) => ({ } case 'streaming_fallback': { + if (msg.cause === 'stream_retry') { + consumePendingDelta(sessionId) + clearPendingToolInputDelta(sessionId) + clearPendingTaskToolUseIds(sessionId) + clearPendingToolParentUseIds(sessionId) + update((session) => { + const startIndex = Math.max( + 0, + Math.min( + session.streamAttemptStartIndex ?? session.messages.length, + session.messages.length, + ), + ) + const messages = [ + ...session.messages.slice(0, startIndex), + ...session.messages.slice(startIndex).filter((message) => + message.type !== 'assistant_text' && + message.type !== 'thinking' && + !(message.type === 'tool_use' && message.isPending)), + ] + return { + messages, + streamingText: '', + streamingToolInput: '', + activeToolUseId: null, + activeToolName: null, + activeThinkingId: null, + streamingResponseChars: + session.streamAttemptStartResponseChars ?? session.streamingResponseChars, + streamAttemptStartIndex: undefined, + streamAttemptStartResponseChars: undefined, + streamingFallback: null, + apiRetry: null, + chatState: 'thinking', + statusVerb: '', + } + }) + ensureElapsedTimer() + useTabStore.getState().updateTabStatus(sessionId, 'running') + break + } + // 进入非流式降级阶段：旧的重试横幅（针对失败的流式请求）已过时， // 清掉换成降级提示；后续非流式重试到来的 api_retry 会重新接管显示。 update((session) => ({ diff --git a/desktop/src/types/chat.ts b/desktop/src/types/chat.ts index 3ab32bc0..2768e6ae 100644 --- a/desktop/src/types/chat.ts +++ b/desktop/src/types/chat.ts @@ -7,6 +7,7 @@ import type { RuntimeSelection } from './runtime' export type ClientMessage = | { type: 'prewarm_session' } + | { type: 'sync_state' } | { type: 'user_message'; content: string; attachments?: AttachmentRef[] } | { type: 'permission_response' @@ -75,6 +76,7 @@ export type UIAttachment = { export type ServerMessage = | { type: 'connected'; sessionId: string } + | { type: 'session_state'; turnState: 'running' | 'idle' } | { type: 'content_start'; blockType: 'text' | 'tool_use'; toolName?: string; toolUseId?: string; parentToolUseId?: string } | { type: 'content_delta'; text?: string; toolInput?: string } | { type: 'tool_use_complete'; toolName: string; toolUseId: string; input: unknown; parentToolUseId?: string } @@ -107,7 +109,7 @@ export type ServerMessage = | { type: 'user_message_replay'; content: string } | { type: 'message_complete'; usage: TokenUsage } | { type: 'thinking'; text: string } - | { type: 'status'; state: ChatState; verb?: string } + | { type: 'status'; state: ChatState; verb?: string; attemptStart?: boolean } // CLI 回传的权限模式变化（如 ExitPlanMode 退出 plan 后恢复、Shift+Tab）。 // 桌面端据此把选择器校正回 CLI 的真实权限，避免本地影子值漂移。 | { type: 'permission_mode_changed'; mode: PermissionMode } @@ -120,7 +122,7 @@ export type ServerMessage = errorType?: string errorMessage?: string } - // 流式请求失败、CLI 已降级为非流式重试：完整响应一次性返回，期间无增量输出。 + // 流式请求失败后的恢复状态：可能安全重试流，也可能降级为非流式请求。 | { type: 'streaming_fallback'; cause: StreamingFallbackCause } | { type: 'error'; message: string; code: string; retryable?: boolean; businessErrorCode?: string } | { type: 'system_notification'; subtype: string; message?: string; data?: unknown } @@ -150,7 +152,7 @@ export type ApiRetryState = { receivedAt: number } -export type StreamingFallbackCause = 'watchdog' | 'stream_error' | '404_stream_creation' | 'unknown' +export type StreamingFallbackCause = 'watchdog' | 'stream_error' | '404_stream_creation' | 'stream_retry' | 'unknown' // 活动回合状态（与 apiRetry 同生命周期），不进消息历史。 export type StreamingFallbackState = { diff --git a/src/entrypoints/sdk/coreSchemas.ts b/src/entrypoints/sdk/coreSchemas.ts index 78d4b6d0..322ae3d4 100644 --- a/src/entrypoints/sdk/coreSchemas.ts +++ b/src/entrypoints/sdk/coreSchemas.ts @@ -1592,12 +1592,12 @@ export const SDKStreamingFallbackMessageSchema = lazySchema(() => .object({ type: z.literal('system'), subtype: z.literal('streaming_fallback'), - cause: z.enum(['watchdog', 'stream_error', '404_stream_creation']), + cause: z.enum(['watchdog', 'stream_error', '404_stream_creation', 'stream_retry']), uuid: UUIDPlaceholder(), session_id: z.string(), }) .describe( - 'Emitted when a streaming request fails and the query falls back to a non-streaming request. The fallback response arrives in one piece after a potentially long wait with no incremental output.', + 'Emitted when a streaming request is recovered by either a bounded safe stream retry or a non-streaming fallback.', ), ) diff --git a/src/server/__tests__/conversation-service.test.ts b/src/server/__tests__/conversation-service.test.ts index 82dba28f..ca8987b6 100644 --- a/src/server/__tests__/conversation-service.test.ts +++ b/src/server/__tests__/conversation-service.test.ts @@ -979,6 +979,59 @@ describe('ConversationService', () => { test('default CLI shutdown wait covers the CLI graceful cleanup budget', () => { expect(DESKTOP_CLI_GRACEFUL_SHUTDOWN_TIMEOUT_MS).toBeGreaterThanOrEqual(6_000) }) + + test('isolates SDK output callbacks so one broken client cannot swallow turn completion', () => { + const service = new ConversationService() as any + let completionObserved = false + service.sessions.set('callback-isolation', { + outputCallbacks: [ + () => { throw new Error('closed client socket') }, + (message: any) => { completionObserved = message.type === 'result' }, + ], + sdkMessages: [], + initMessage: null, + pendingPermissionRequests: new Map(), + }) + + service.handleSdkPayload('callback-isolation', JSON.stringify({ + type: 'result', + subtype: 'success', + is_error: false, + })) + + expect(completionObserved).toBe(true) + }) + + test('removes an exited CLI session even when one output callback throws', async () => { + const service = new ConversationService() as any + const sessionId = 'exit-callback-isolation' + const proc = { + exited: Promise.resolve(1), + kill: () => {}, + } + let completionObserved = false + service.sessions.set(sessionId, { + proc, + startupPending: false, + startupExitCode: null, + outputDrain: Promise.resolve(), + outputCallbacks: [ + () => { throw new Error('closed client socket') }, + (message: any) => { completionObserved = message.type === 'result' }, + ], + workDir: tmpDir, + permissionMode: 'default', + stdoutLines: [], + stderrLines: [], + sdkMessages: [], + pendingPermissionRequests: new Map(), + }) + + await service.handleProcessExit(sessionId, proc, 1) + + expect(completionObserved).toBe(true) + expect(service.hasSession(sessionId)).toBe(false) + }) }) function sanitizeMemoryPath(value: string): string { diff --git a/src/server/__tests__/conversations.test.ts b/src/server/__tests__/conversations.test.ts index 5c0b1876..6e732469 100644 --- a/src/server/__tests__/conversations.test.ts +++ b/src/server/__tests__/conversations.test.ts @@ -4283,6 +4283,96 @@ describe('WebSocket Chat Integration', () => { }) }) + it('should reconcile an idle turn that completed while the client was disconnected', async () => { + await withMockStreamDelay(75, async () => { + const sessionId = `chat-reconnect-completed-${crypto.randomUUID()}` + const firstMessages: any[] = [] + const reconnectMessages: any[] = [] + let firstSocket: WebSocket | null = null + let reconnectSocket: WebSocket | null = null + + try { + await new Promise((resolve, reject) => { + const timeout = setTimeout(() => { + reject(new Error(`Timed out waiting for disconnected completion for session ${sessionId}`)) + }, 10_000) + + const fail = (error: unknown) => { + clearTimeout(timeout) + reject(error instanceof Error ? error : new Error(String(error))) + } + + firstSocket = new WebSocket(`${wsUrl}/ws/${sessionId}`) + firstSocket.onmessage = (event) => { + const message = JSON.parse(event.data as string) + firstMessages.push(message) + if (message.type === 'connected') { + firstSocket?.send(JSON.stringify({ + type: 'user_message', + content: 'finish while disconnected', + })) + return + } + if (message.type !== 'thinking') return + + firstSocket?.close() + void (async () => { + const deadline = Date.now() + 5_000 + while ( + Date.now() < deadline && + !conversationService.getRecentSdkMessages(sessionId).some( + (entry) => entry?.type === 'result', + ) + ) { + await new Promise((pollResolve) => setTimeout(pollResolve, 25)) + } + if (!conversationService.getRecentSdkMessages(sessionId).some( + (entry) => entry?.type === 'result', + )) { + fail(new Error(`CLI did not finish while disconnected for session ${sessionId}`)) + return + } + + reconnectSocket = new WebSocket(`${wsUrl}/ws/${sessionId}`) + reconnectSocket.onmessage = (reconnectEvent) => { + const reconnectMessage = JSON.parse(reconnectEvent.data as string) + reconnectMessages.push(reconnectMessage) + if (reconnectMessage.type === 'connected') { + reconnectSocket?.send(JSON.stringify({ type: 'sync_state' })) + return + } + if ( + reconnectMessage.type === 'session_state' && + reconnectMessage.turnState === 'idle' + ) { + clearTimeout(timeout) + resolve() + } + } + reconnectSocket.onerror = () => fail( + new Error(`Reconnect WebSocket error for session ${sessionId}`), + ) + })().catch(fail) + } + firstSocket.onerror = () => fail( + new Error(`Initial WebSocket error for session ${sessionId}`), + ) + }) + + expect(firstMessages.some((message) => message.type === 'thinking')).toBe(true) + expect(reconnectMessages).toContainEqual({ + type: 'session_state', + turnState: 'idle', + }) + expect(reconnectMessages.some((message) => message.type === 'message_complete')).toBe(false) + } finally { + firstSocket?.close() + reconnectSocket?.close() + conversationService.stopSession(sessionId) + } + }) + }) + it('should stream one active turn to multiple connected clients', async () => { await withMockStreamDelay(150, async () => { const sessionId = `chat-multi-client-${crypto.randomUUID()}` diff --git a/src/server/__tests__/websocket-handler.test.ts b/src/server/__tests__/websocket-handler.test.ts index 5691e816..1a354499 100644 --- a/src/server/__tests__/websocket-handler.test.ts +++ b/src/server/__tests__/websocket-handler.test.ts @@ -18,6 +18,7 @@ import { } from '../ws/disconnectGraceConfig.js' import { conversationService } from '../services/conversationService.js' import { computerUseApprovalService } from '../services/computerUseApprovalService.js' +import { sessionService } from '../services/sessionService.js' function makeClientSocket(sessionId: string) { const sent: string[] = [] @@ -443,6 +444,109 @@ describe('WebSocket handler session isolation', () => { turnCompleteCallback?.({ type: 'result', subtype: 'success' }) expect(setTimeoutSpy).not.toHaveBeenCalled() }) + + it('reports authoritative turn state when a reconnected client asks to reconcile', () => { + const runningSessionId = `sync-running-${crypto.randomUUID()}` + const runningSocket = makeClientSocket(runningSessionId) + spyOn(conversationService, 'getPendingPermissionRequests').mockReturnValue([]) + + handleWebSocket.open(runningSocket) + __markActiveTurnForTests(runningSessionId) + runningSocket.sent.length = 0 + handleWebSocket.message(runningSocket, JSON.stringify({ type: 'sync_state' })) + + expect(runningSocket.sent.map((payload) => JSON.parse(payload))).toContainEqual({ + type: 'session_state', + turnState: 'running', + }) + + const idleSessionId = `sync-idle-${crypto.randomUUID()}` + const idleSocket = makeClientSocket(idleSessionId) + handleWebSocket.open(idleSocket) + idleSocket.sent.length = 0 + handleWebSocket.message(idleSocket, JSON.stringify({ type: 'sync_state' })) + + expect(idleSocket.sent.map((payload) => JSON.parse(payload))).toContainEqual({ + type: 'session_state', + turnState: 'idle', + }) + }) + + it('terminates the desktop turn when user-message handling throws unexpectedly', async () => { + const sessionId = `user-message-failure-${crypto.randomUUID()}` + const ws = makeClientSocket(sessionId) + spyOn(conversationService, 'getPendingPermissionRequests').mockReturnValue([]) + spyOn(sessionService, 'getCustomTitle').mockRejectedValue( + new Error('metadata store unavailable'), + ) + + handleWebSocket.open(ws) + ws.sent.length = 0 + handleWebSocket.message(ws, JSON.stringify({ + type: 'user_message', + content: 'continue the long task', + })) + await new Promise((resolve) => setTimeout(resolve, 0)) + + const messages = ws.sent.map((payload) => JSON.parse(payload)) + expect(messages).toContainEqual({ + type: 'error', + message: 'The request could not be started. Please retry.', + code: 'USER_TURN_FAILED', + retryable: true, + }) + expect(messages).toContainEqual({ type: 'status', state: 'idle' }) + + ws.sent.length = 0 + handleWebSocket.message(ws, JSON.stringify({ type: 'sync_state' })) + expect(ws.sent.map((payload) => JSON.parse(payload))).toContainEqual({ + type: 'session_state', + turnState: 'idle', + }) + }) + + it('does not let an older failed handler clear a newer active turn', async () => { + const sessionId = `concurrent-user-message-${crypto.randomUUID()}` + const ws = makeClientSocket(sessionId) + spyOn(conversationService, 'getPendingPermissionRequests').mockReturnValue([]) + + let rejectFirst!: (error: Error) => void + let customTitleCalls = 0 + spyOn(sessionService, 'getCustomTitle').mockImplementation(() => { + customTitleCalls++ + if (customTitleCalls === 1) { + return new Promise((_resolve, reject) => { + rejectFirst = reject + }) + } + return new Promise(() => {}) + }) + + handleWebSocket.open(ws) + handleWebSocket.message(ws, JSON.stringify({ + type: 'user_message', + content: 'older turn', + })) + await new Promise((resolve) => setTimeout(resolve, 0)) + expect(customTitleCalls).toBe(1) + + handleWebSocket.message(ws, JSON.stringify({ + type: 'user_message', + content: 'newer turn', + })) + await new Promise((resolve) => setTimeout(resolve, 0)) + expect(customTitleCalls).toBe(2) + + rejectFirst(new Error('older metadata request failed')) + await new Promise((resolve) => setTimeout(resolve, 0)) + + ws.sent.length = 0 + handleWebSocket.message(ws, JSON.stringify({ type: 'sync_state' })) + expect(ws.sent.map((payload) => JSON.parse(payload))).toContainEqual({ + type: 'session_state', + turnState: 'running', + }) + }) }) describe('prewarm idle timer active-turn guard (issue #865 follow-up)', () => { diff --git a/src/server/__tests__/ws-memory-events.test.ts b/src/server/__tests__/ws-memory-events.test.ts index 8556bd96..7a3fc7ed 100644 --- a/src/server/__tests__/ws-memory-events.test.ts +++ b/src/server/__tests__/ws-memory-events.test.ts @@ -629,6 +629,99 @@ describe('WebSocket goal command events', () => { }) describe('WebSocket stream event translation', () => { + it('does not replay buffered assistant blocks after their raw stream events', () => { + const sessionId = `buffered-assistant-${crypto.randomUUID()}` + + translateCliMessage({ + type: 'stream_event', + event: { type: 'message_start' }, + }, sessionId) + translateCliMessage({ + type: 'stream_event', + event: { + type: 'content_block_start', + index: 0, + content_block: { type: 'thinking', thinking: '' }, + }, + }, sessionId) + translateCliMessage({ + type: 'stream_event', + event: { + type: 'content_block_delta', + index: 0, + delta: { type: 'thinking_delta', thinking: 'reasoning' }, + }, + }, sessionId) + translateCliMessage({ + type: 'stream_event', + event: { type: 'content_block_stop', index: 0 }, + }, sessionId) + translateCliMessage({ + type: 'stream_event', + event: { + type: 'content_block_start', + index: 1, + content_block: { type: 'text', text: '' }, + }, + }, sessionId) + translateCliMessage({ + type: 'stream_event', + event: { + type: 'content_block_delta', + index: 1, + delta: { type: 'text_delta', text: 'hello' }, + }, + }, sessionId) + translateCliMessage({ + type: 'stream_event', + event: { type: 'content_block_stop', index: 1 }, + }, sessionId) + translateCliMessage({ + type: 'stream_event', + event: { type: 'message_stop' }, + }, sessionId) + + expect(translateCliMessage({ + type: 'assistant', + message: { content: [{ type: 'thinking', thinking: 'reasoning' }] }, + }, sessionId)).toEqual([]) + expect(translateCliMessage({ + type: 'assistant', + message: { content: [{ type: 'text', text: 'hello' }] }, + }, sessionId)).toEqual([]) + }) + + it('accepts the complete assistant response after a non-streaming fallback', () => { + const sessionId = `non-stream-fallback-${crypto.randomUUID()}` + + translateCliMessage({ + type: 'stream_event', + event: { type: 'message_start' }, + }, sessionId) + translateCliMessage({ + type: 'stream_event', + event: { + type: 'content_block_start', + index: 0, + content_block: { type: 'text', text: '' }, + }, + }, sessionId) + + translateCliMessage({ + type: 'system', + subtype: 'streaming_fallback', + cause: 'watchdog', + }, sessionId) + + expect(translateCliMessage({ + type: 'assistant', + message: { content: [{ type: 'text', text: 'fallback answer' }] }, + }, sessionId)).toEqual([ + { type: 'content_start', blockType: 'text' }, + { type: 'content_delta', text: 'fallback answer' }, + ]) + }) + it('keeps subagent parent linkage when later stream events omit the parent id', () => { const sessionId = `subagent-parent-${crypto.randomUUID()}` @@ -699,7 +792,7 @@ describe('WebSocket stream event translation', () => { type: 'stream_event', event: { type: 'message_start' }, }, sessionId)).toEqual([ - { type: 'status', state: 'thinking' }, + { type: 'status', state: 'thinking', attemptStart: true }, ]) expect(translateCliMessage({ @@ -740,4 +833,71 @@ describe('WebSocket stream event translation', () => { { type: 'content_start', blockType: 'text' }, ]) }) + + it('resets partial block accumulation before a safe stream retry', () => { + const sessionId = `stream-retry-${crypto.randomUUID()}` + + translateCliMessage({ + type: 'stream_event', + event: { type: 'message_start' }, + }, sessionId) + translateCliMessage({ + type: 'stream_event', + event: { + type: 'content_block_start', + index: 0, + content_block: { type: 'tool_use', id: 'stale-tool', name: 'Write' }, + }, + }, sessionId) + translateCliMessage({ + type: 'stream_event', + event: { + type: 'content_block_delta', + index: 0, + delta: { type: 'input_json_delta', partial_json: '{"stale":' }, + }, + }, sessionId) + + expect(translateCliMessage({ + type: 'system', + subtype: 'streaming_fallback', + cause: 'stream_retry', + }, sessionId)).toEqual([ + { type: 'streaming_fallback', cause: 'stream_retry' }, + ]) + + translateCliMessage({ + type: 'stream_event', + event: { type: 'message_start' }, + }, sessionId) + translateCliMessage({ + type: 'stream_event', + event: { + type: 'content_block_start', + index: 0, + content_block: { type: 'tool_use', id: 'fresh-tool', name: 'Write' }, + }, + }, sessionId) + translateCliMessage({ + type: 'stream_event', + event: { + type: 'content_block_delta', + index: 0, + delta: { type: 'input_json_delta', partial_json: '{"fresh":true}' }, + }, + }, sessionId) + + expect(translateCliMessage({ + type: 'stream_event', + event: { type: 'content_block_stop', index: 0 }, + }, sessionId)).toEqual([ + { + type: 'tool_use_complete', + toolName: 'Write', + toolUseId: 'fresh-tool', + input: { fresh: true }, + parentToolUseId: undefined, + }, + ]) + }) }) diff --git a/src/server/services/conversationService.ts b/src/server/services/conversationService.ts index 4c9334ee..ec6cdfc6 100644 --- a/src/server/services/conversationService.ts +++ b/src/server/services/conversationService.ts @@ -776,9 +776,7 @@ export class ConversationService { ) { session.pendingPermissionRequests.delete(msg.response.request_id) } - for (const cb of session.outputCallbacks) { - cb(msg) - } + this.notifyOutputCallbacks(sessionId, session.outputCallbacks, msg) } catch { console.warn( `[ConversationService] Ignoring malformed SDK payload for ${sessionId}`, @@ -787,6 +785,24 @@ export class ConversationService { } } + private notifyOutputCallbacks( + sessionId: string, + callbacks: Array<(msg: any) => void>, + message: any, + ): void { + for (const callback of [...callbacks]) { + try { + callback(message) + } catch (error) { + console.warn( + `[ConversationService] Output callback failed for ${sessionId}: ${ + error instanceof Error ? error.message : String(error) + }`, + ) + } + } + } + stopSession(sessionId: string): void { const session = this.sessions.get(sessionId) if (!session) return @@ -990,17 +1006,16 @@ export class ConversationService { sdkMessages: this.summarizeSdkMessages(activeSession.sdkMessages), }, }) - for (const cb of activeSession.outputCallbacks) { - cb({ - type: 'result', - subtype: 'error', - is_error: true, - result: exitError, - usage: { input_tokens: 0, output_tokens: 0 }, - session_id: sessionId, - }) - } + const callbacks = [...activeSession.outputCallbacks] this.sessions.delete(sessionId) + this.notifyOutputCallbacks(sessionId, callbacks, { + type: 'result', + subtype: 'error', + is_error: true, + result: exitError, + usage: { input_tokens: 0, output_tokens: 0 }, + session_id: sessionId, + }) } } diff --git a/src/server/ws/events.ts b/src/server/ws/events.ts index 47cfea14..3f1a7171 100644 --- a/src/server/ws/events.ts +++ b/src/server/ws/events.ts @@ -10,6 +10,7 @@ export type ClientMessage = | { type: 'prewarm_session' } + | { type: 'sync_state' } | { type: 'user_message'; content: string; attachments?: AttachmentRef[] } | { type: 'permission_response' @@ -45,6 +46,7 @@ export type AttachmentRef = { export type ServerMessage = | { type: 'connected'; sessionId: string } + | { type: 'session_state'; turnState: 'running' | 'idle' } | { type: 'content_start'; blockType: 'text' | 'tool_use'; toolName?: string; toolUseId?: string; parentToolUseId?: string } | { type: 'content_delta'; text?: string; toolInput?: string } | { type: 'tool_use_complete'; toolName: string; toolUseId: string; input: unknown; parentToolUseId?: string } @@ -77,7 +79,7 @@ export type ServerMessage = | { type: 'user_message_replay'; content: string } | { type: 'message_complete'; usage: TokenUsage } | { type: 'thinking'; text: string } - | { type: 'status'; state: ChatState; verb?: string } + | { type: 'status'; state: ChatState; verb?: string; attemptStart?: boolean } // CLI 是权限模式的唯一真相来源。当 CLI 内部 mode 变化（如 ExitPlanMode 后 // 恢复到进入 plan 前的模式、Shift+Tab 切换）时，把新模式回传给前端，让桌面端 // 选择器与 CLI 保持同步，而不是停留在本地影子值上。 @@ -114,7 +116,7 @@ export type ChatState = 'idle' | 'thinking' | 'compacting' | 'tool_executing' | // 与 CLI 的 streaming_fallback cause 对齐；unknown 兜底未来新增的 cause 值， // 避免新 CLI + 旧 server 组合下丢消息。 -export type StreamingFallbackCause = 'watchdog' | 'stream_error' | '404_stream_creation' | 'unknown' +export type StreamingFallbackCause = 'watchdog' | 'stream_error' | '404_stream_creation' | 'stream_retry' | 'unknown' export type TeamMemberStatus = { agentId: string diff --git a/src/server/ws/handler.ts b/src/server/ws/handler.ts index bb7fc8bf..351eaf9d 100644 --- a/src/server/ws/handler.ts +++ b/src/server/ws/handler.ts @@ -233,18 +233,36 @@ export const handleWebSocket = { ) as ClientMessage switch (message.type) { - case 'user_message': - handleUserMessage(ws, message).catch((err) => { + case 'user_message': { + const activeTurn: ActiveUserTurnState = { messageSent: false } + handleUserMessage(ws, message, activeTurn).catch((err) => { + const sessionId = ws.data.sessionId void diagnosticsService.recordEvent({ type: 'ws_user_message_failed', severity: 'error', - sessionId: ws.data.sessionId, + sessionId, summary: err instanceof Error ? err.message : String(err), details: err, }) console.error(`[WS] Unhandled error in handleUserMessage:`, err) + // A queued/newer turn may have replaced this handler while an + // earlier await was pending. Only the handler that still owns the + // active-turn token may terminate the desktop state. + if (activeUserTurns.get(sessionId) === activeTurn) { + clearActiveUserTurn(sessionId, activeTurn) + const titleState = sessionTitleState.get(sessionId) + if (titleState) titleState.activeTurn = undefined + sendMessage(ws, { + type: 'error', + message: 'The request could not be started. Please retry.', + code: 'USER_TURN_FAILED', + retryable: true, + }) + sendMessage(ws, { type: 'status', state: 'idle' }) + } }) break + } case 'permission_response': handlePermissionResponse(ws, message) @@ -266,6 +284,15 @@ export const handleWebSocket = { void handlePrewarmSession(ws) break + case 'sync_state': + sendMessage(ws, { + type: 'session_state', + turnState: hasPendingOrActiveUserTurn(ws.data.sessionId) + ? 'running' + : 'idle', + }) + break + case 'stop_generation': handleStopGeneration(ws) break @@ -326,7 +353,8 @@ export const handleWebSocket = { async function handleUserMessage( ws: ServerWebSocket, - message: Extract + message: Extract, + activeTurn: ActiveUserTurnState, ) { const { sessionId } = ws.data @@ -353,7 +381,6 @@ async function handleUserMessage( // Send thinking status sendMessage(ws, { type: 'status', state: 'thinking', verb: 'Thinking' }) - const activeTurn: ActiveUserTurnState = { messageSent: false } activeUserTurns.set(sessionId, activeTurn) const initialRuntimeTransition = await waitForRuntimeTransitionBeforeUserTurn(ws, sessionId) @@ -1231,6 +1258,13 @@ function getStreamState(sessionId: string): SessionStreamState { return state } +function resetCurrentStreamAttempt(state: SessionStreamState): void { + state.hasReceivedStreamEvents = false + state.activeBlockTypes.clear() + state.activeToolBlocks.clear() + state.pendingToolBlocks.clear() +} + function cliParentToolUseId(cliMsg: any): string | undefined { return typeof cliMsg.parent_tool_use_id === 'string' && cliMsg.parent_tool_use_id.length > 0 ? cliMsg.parent_tool_use_id @@ -1536,9 +1570,6 @@ export function translateCliMessage(cliMsg: any, sessionId: string): ServerMessa } } - // Reset flags for next turn - streamState.hasReceivedStreamEvents = false - streamState.pendingToolBlocks.clear() return messages } return [] @@ -1620,7 +1651,7 @@ export function translateCliMessage(cliMsg: any, sessionId: string): ServerMessa switch (event.type) { case 'message_start': { - return [{ type: 'status', state: 'thinking' }] + return [{ type: 'status', state: 'thinking', attemptStart: true }] } case 'content_block_start': { @@ -1782,6 +1813,10 @@ export function translateCliMessage(cliMsg: any, sessionId: string): ServerMessa case 'result': { // 对话结果（成功或错误） const usage = translateCliUsage(cliMsg.usage) + // Buffered assistant blocks can arrive as a batch after all raw events + // for one provider message. Keep deduplication active across the entire + // batch, then clear it only at the terminal result boundary. + resetCurrentStreamAttempt(streamState) if (cliMsg.is_error) { // If the user requested stop, this "error" is just the interrupt @@ -1825,6 +1860,9 @@ export function translateCliMessage(cliMsg: any, sessionId: string): ServerMessa return apiRetryMessage ? [apiRetryMessage] : [] } if (subtype === 'streaming_fallback') { + // The next attempt is a new stream or a full non-streaming response; + // neither should inherit raw-event dedup/tool JSON from the failed one. + resetCurrentStreamAttempt(streamState) return [toStreamingFallbackServerMessage(cliMsg)] } if (subtype === 'init') { @@ -2069,6 +2107,7 @@ const STREAMING_FALLBACK_CAUSES: ReadonlySet = new Set([ 'watchdog', 'stream_error', '404_stream_creation', + 'stream_retry', ]) function toStreamingFallbackServerMessage(cliMsg: any): ServerMessage { diff --git a/src/services/api/claude.ts b/src/services/api/claude.ts index 4250f586..701b6365 100644 --- a/src/services/api/claude.ts +++ b/src/services/api/claude.ts @@ -212,6 +212,7 @@ import { stopSessionActivity, } from "../../utils/sessionActivity.js"; import { shouldTriggerNonStreamingFallbackForEmptyStream } from "./streamFallback.js"; +import { StreamAssistantCommitBuffer } from "./streamAssistantCommitBuffer.js"; import { StreamWatchdogTimeoutError, createStreamWatchdogState, @@ -1870,6 +1871,7 @@ async function* queryModel( } const newMessages: AssistantMessage[] = []; + const assistantCommitBuffer = new StreamAssistantCommitBuffer(); let ttftMs = 0; let partialMessage: BetaMessage | undefined = undefined; const contentBlocks: (BetaContentBlock | ConnectorTextBlock)[] = []; @@ -2420,7 +2422,12 @@ async function* queryModel( ...(advisorModel && { advisorModel }), }; newMessages.push(m); - yield m; + for (const committedMessage of assistantCommitBuffer.add( + m, + contentBlock.type, + )) { + yield committedMessage; + } break; } case "message_delta": { @@ -2460,6 +2467,18 @@ async function* queryModel( lastMsg.message.stop_reason = stopReason; } + // Max-token recovery needs the completed assistant blocks before + // its synthetic error so query.ts still sees the error as the + // terminal message and can continue from the partial response. + if ( + stopReason === "max_tokens" || + stopReason === "model_context_window_exceeded" + ) { + for (const committedMessage of assistantCommitBuffer.flush()) { + yield committedMessage; + } + } + // Update cost const costUSDForPart = calculateUSDCost(resolvedModel, usage); costUSD += addToTotalSessionCost( @@ -2473,6 +2492,9 @@ async function* queryModel( options.model, ); if (refusalMessage) { + for (const committedMessage of assistantCommitBuffer.flush()) { + yield committedMessage; + } yield refusalMessage; } @@ -2597,6 +2619,13 @@ async function* queryModel( throw new Error("Stream ended without receiving any events"); } + // No tool boundary was crossed, so completed thinking/text blocks were + // intentionally held until the response proved it could finish. This + // keeps a watchdog retry from persisting orphan partial blocks. + for (const committedMessage of assistantCommitBuffer.flush()) { + yield committedMessage; + } + // Log summary if any stalls occurred during streaming if (stallCount > 0) { logForDebugging( @@ -2699,22 +2728,16 @@ async function* queryModel( } } - // A transient, server-side error that arrived mid-stream (a local provider - // rejecting a malformed tool_call, or an upstream api_error / - // overloaded_error SSE event) is recoverable by re-establishing the - // stream. Only retry when this attempt produced NOTHING - // (newMessages.length === 0): a zero-output stream means no tool_use block - // ever completed, so query.ts never started a tool — no double-execution - // risk (cf. #766 / inc-4258), the same precondition the zero-output - // fallback below relies on. Watchdog aborts without content/tool output - // are handled first; remaining watchdog aborts keep their partial-output - // boundary and are not retried. Thrown past the outer catch — which - // re-throws it — up to withStreamRetry. + // A watchdog stall is recoverable while the failed attempt is still + // side-effect-free. Completed thinking/text and partial local tool JSON + // stay buffered, so re-establishing the stream cannot duplicate a tool. + // A completed local tool block or any server-side tool activity closes + // this retry boundary permanently for the attempt. if ( streamIdleAborted && streamingError instanceof StreamWatchdogTimeoutError && streamingError.safeToRetryStream() && - newMessages.length === 0 && + !assistantCommitBuffer.hasCrossedSideEffectBoundary() && !signal.aborted ) { logForDebugging( diff --git a/src/services/api/streamAssistantCommitBuffer.test.ts b/src/services/api/streamAssistantCommitBuffer.test.ts new file mode 100644 index 00000000..92b9c892 --- /dev/null +++ b/src/services/api/streamAssistantCommitBuffer.test.ts @@ -0,0 +1,43 @@ +import { describe, expect, test } from 'bun:test' +import { StreamAssistantCommitBuffer } from './streamAssistantCommitBuffer.js' + +function assistant(uuid: string) { + return { + type: 'assistant' as const, + uuid, + message: { role: 'assistant' as const, content: [] }, + } +} + +describe('StreamAssistantCommitBuffer', () => { + test('holds side-effect-free blocks until the stream completes', () => { + const buffer = new StreamAssistantCommitBuffer>() + const thinking = assistant('thinking') + const text = assistant('text') + + expect(buffer.add(thinking, 'thinking')).toEqual([]) + expect(buffer.add(text, 'text')).toEqual([]) + expect(buffer.flush()).toEqual([thinking, text]) + }) + + test('commits buffered blocks with the first local tool boundary', () => { + const buffer = new StreamAssistantCommitBuffer>() + const thinking = assistant('thinking') + const tool = assistant('tool') + + expect(buffer.add(thinking, 'thinking')).toEqual([]) + expect(buffer.add(tool, 'tool_use')).toEqual([thinking, tool]) + expect(buffer.add(assistant('after-tool'), 'text')).toEqual([ + expect.objectContaining({ uuid: 'after-tool' }), + ]) + expect(buffer.flush()).toEqual([]) + }) + + test('treats server tools as an irreversible boundary', () => { + const buffer = new StreamAssistantCommitBuffer>() + const serverTool = assistant('server-tool') + + expect(buffer.add(serverTool, 'server_tool_use')).toEqual([serverTool]) + expect(buffer.hasCrossedSideEffectBoundary()).toBe(true) + }) +}) diff --git a/src/services/api/streamAssistantCommitBuffer.ts b/src/services/api/streamAssistantCommitBuffer.ts new file mode 100644 index 00000000..c28f5aaa --- /dev/null +++ b/src/services/api/streamAssistantCommitBuffer.ts @@ -0,0 +1,35 @@ +/** + * Holds completed, side-effect-free assistant blocks until a stream either + * finishes or crosses a tool boundary. A watchdog retry can then discard the + * failed attempt without leaving orphan thinking/text in the transcript. + */ +export class StreamAssistantCommitBuffer { + private pending: T[] = [] + private crossedSideEffectBoundary = false + + add(value: T, blockType: string): T[] { + if (this.crossedSideEffectBoundary) return [value] + + this.pending.push(value) + if (blockType !== 'tool_use' && blockType !== 'server_tool_use') { + return [] + } + + this.crossedSideEffectBoundary = true + return this.drain() + } + + flush(): T[] { + return this.drain() + } + + hasCrossedSideEffectBoundary(): boolean { + return this.crossedSideEffectBoundary + } + + private drain(): T[] { + const values = this.pending + this.pending = [] + return values + } +} diff --git a/src/services/api/streamRetry.test.ts b/src/services/api/streamRetry.test.ts index 99c8bb52..3369974d 100644 --- a/src/services/api/streamRetry.test.ts +++ b/src/services/api/streamRetry.test.ts @@ -61,6 +61,11 @@ describe('withStreamRetry', () => { const out = await collect(withStreamRetry(attempt, 'test-model', [])) expect(calls).toBe(2) + expect(out).toContainEqual(expect.objectContaining({ + type: 'system', + subtype: 'streaming_fallback', + cause: 'stream_retry', + })) const assistants = out.filter(m => m.type === 'assistant') expect(assistants).toHaveLength(1) expect(assistants[0].uuid).toBe('ok') @@ -82,6 +87,9 @@ describe('withStreamRetry', () => { const out = await collect(withStreamRetry(attempt, 'test-model', [])) expect(calls).toBe(3) // 1 initial attempt + 2 retries + expect(out.filter( + m => m.type === 'system' && m.subtype === 'streaming_fallback' && m.cause === 'stream_retry', + )).toHaveLength(2) const last = out.at(-1) expect(last?.type).toBe('assistant') expect(last?.isApiErrorMessage).toBe(true) diff --git a/src/services/api/streamRetry.ts b/src/services/api/streamRetry.ts index 84c4adfb..06643d9a 100644 --- a/src/services/api/streamRetry.ts +++ b/src/services/api/streamRetry.ts @@ -7,6 +7,7 @@ import type { } from "../../types/message.js"; import { logForDebugging } from "../../utils/debug.js"; import { errorMessage } from "../../utils/errors.js"; +import { createSystemStreamingFallbackMessage } from "../../utils/messages.js"; import { type AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS, logEvent, @@ -36,14 +37,12 @@ type StreamQueryMessage = * every per-request value is a fresh local, so there is nothing to reset by hand. * * Safe against the double-tool-execution hazard (#766 / inc-4258): queryModel - * only throws RetriableStreamError when the failed attempt produced zero - * assistant messages (no content_block_stop completed), so query.ts never handed - * a tool_use to the StreamingToolExecutor and no tool ran. + * buffers completed thinking/text and only throws RetriableStreamError before a + * local tool block completes or server-side tool activity starts. * - * A failed attempt may already have yielded raw stream_event partials; the retry - * re-emits message_start etc. queryModelWithoutStreaming ignores stream_event - * entirely, and the streaming UI resets its in-flight partial on the next - * message_start, so a re-emit at most causes a brief redraw. + * A failed attempt may already have yielded raw stream_event partials. Before + * retrying, emit a bounded recovery signal so streaming consumers discard that + * attempt instead of appending the next attempt to stale text/tool JSON. */ export async function* withStreamRetry( attempt: () => AsyncGenerator, @@ -89,6 +88,9 @@ export async function* withStreamRetry( model: model as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS, }); + // Raw deltas from the failed attempt may already be visible. Consumers + // use this bounded retry signal to discard only that in-flight attempt. + yield createSystemStreamingFallbackMessage("stream_retry"); } } } diff --git a/src/services/api/streamWatchdog.test.ts b/src/services/api/streamWatchdog.test.ts index 15e20bbb..ea1f2833 100644 --- a/src/services/api/streamWatchdog.test.ts +++ b/src/services/api/streamWatchdog.test.ts @@ -73,13 +73,13 @@ describe('stream watchdog state', () => { expect(error.code).toBe('STREAM_IDLE_TIMEOUT') expect(error.phase).toBe('mid_stream') - expect(error.safeToRetryStream()).toBe(false) + expect(error.safeToRetryStream()).toBe(true) expect(error.message).toContain('Provider stream stalled after partial response') expect(error.message).toContain('last event: text_delta') expect(error.message).not.toContain('no chunks received') }) - test('does not retry after a tool use has started', () => { + test('retries partial local tool input but stops after the tool block completes', () => { const state = createStreamWatchdogState() state.recordEvent({ type: 'message_start' }) state.recordEvent({ @@ -88,9 +88,33 @@ describe('stream watchdog state', () => { content_block: { type: 'tool_use' }, }) + expect(state.createTimeoutError('idle', 240_000).safeToRetryStream()).toBe(true) + + state.recordEvent({ + type: 'content_block_delta', + index: 0, + delta: { type: 'input_json_delta', partial_json: '{"file_path":' }, + }) + expect(state.createTimeoutError('idle', 240_000).safeToRetryStream()).toBe(true) + + state.recordEvent({ type: 'content_block_stop', index: 0 }) + const error = state.createTimeoutError('idle', 240_000) + + expect(error.phase).toBe('mid_stream') + expect(error.safeToRetryStream()).toBe(false) + }) + + test('never retries after server-side tool activity begins', () => { + const state = createStreamWatchdogState() + state.recordEvent({ type: 'message_start' }) + state.recordEvent({ + type: 'content_block_start', + index: 0, + content_block: { type: 'server_tool_use' }, + }) + const error = state.createTimeoutError('idle', 240_000) - expect(error.phase).toBe('before_content') expect(error.safeToRetryStream()).toBe(false) }) @@ -112,6 +136,7 @@ describe('stream watchdog state', () => { expect(error.code).toBe('STREAM_MAX_DURATION') expect(error.phase).toBe('mid_stream') + expect(error.safeToRetryStream()).toBe(false) expect(error.message).toContain('Stream max duration exceeded') expect(error.message).toContain('last event: text_delta') }) diff --git a/src/services/api/streamWatchdog.ts b/src/services/api/streamWatchdog.ts index 43925086..f91130c0 100644 --- a/src/services/api/streamWatchdog.ts +++ b/src/services/api/streamWatchdog.ts @@ -17,6 +17,9 @@ export type StreamWatchdogSnapshot = { lastBlockType?: string messageStopReceived: boolean toolUseStarted: boolean + localToolUseStarted: boolean + localToolUseCompleted: boolean + serverToolUseStarted: boolean } type StreamEventLike = { @@ -148,8 +151,9 @@ export class StreamWatchdogTimeoutError extends Error { safeToRetryStream(): boolean { return ( - this.streamSnapshot.contentDeltaCount === 0 && - !this.streamSnapshot.toolUseStarted + this.reason === 'idle' && + !this.streamSnapshot.localToolUseCompleted && + !this.streamSnapshot.serverToolUseStarted ) } @@ -175,6 +179,9 @@ class StreamWatchdogState { toolInputDeltaCount: 0, messageStopReceived: false, toolUseStarted: false, + localToolUseStarted: false, + localToolUseCompleted: false, + serverToolUseStarted: false, } private readonly blockTypesByIndex = new Map() @@ -212,6 +219,10 @@ class StreamWatchdogState { this.snapshotValue.toolUseStarted || blockType === 'tool_use' || blockType === 'server_tool_use', + localToolUseStarted: + this.snapshotValue.localToolUseStarted || blockType === 'tool_use', + serverToolUseStarted: + this.snapshotValue.serverToolUseStarted || blockType === 'server_tool_use', } if (index !== undefined) { this.blockTypesByIndex.set(index, blockType) @@ -219,6 +230,23 @@ class StreamWatchdogState { } } + if (type === 'content_block_stop') { + const index = readIndex(rawEvent?.index) + const blockType = index === undefined + ? undefined + : this.blockTypesByIndex.get(index) + if (blockType === 'tool_use') { + this.snapshotValue = { + ...this.snapshotValue, + localToolUseCompleted: true, + lastBlockType: blockType, + } + } + if (index !== undefined) { + this.blockTypesByIndex.delete(index) + } + } + if (type === 'content_block_delta') { const delta = asObject(rawEvent?.delta) const deltaType = readString(delta?.type) diff --git a/src/services/api/withRetry.test.ts b/src/services/api/withRetry.test.ts index fab65b53..4b794742 100644 --- a/src/services/api/withRetry.test.ts +++ b/src/services/api/withRetry.test.ts @@ -123,6 +123,12 @@ describe('getMaxStreamTransientRetries', () => { delete process.env[ENV] }) + test('caps overrides so recovery cannot become an unbounded retry loop', () => { + process.env[ENV] = '1000' + expect(getMaxStreamTransientRetries()).toBe(5) + delete process.env[ENV] + }) + test('falls back to 2 on non-numeric input', () => { process.env[ENV] = 'abc' expect(getMaxStreamTransientRetries()).toBe(2) diff --git a/src/services/api/withRetry.ts b/src/services/api/withRetry.ts index 11494515..9461408c 100644 --- a/src/services/api/withRetry.ts +++ b/src/services/api/withRetry.ts @@ -216,11 +216,12 @@ export function isRetryableStreamError(error: unknown): boolean { * Max times withStreamRetry() re-establishes a stream after a transient * mid-stream error (see RetriableStreamError). Small by default — a malformed * tool_call or a one-off blip usually clears on the first retry; this is not a - * capacity backoff loop. Override with CLAUDE_STREAM_TRANSIENT_RETRY_MAX. + * capacity backoff loop. Override with CLAUDE_STREAM_TRANSIENT_RETRY_MAX; + * values are capped at 5 so a bad environment value cannot make it unbounded. */ export function getMaxStreamTransientRetries(): number { const raw = parseInt(process.env.CLAUDE_STREAM_TRANSIENT_RETRY_MAX || '', 10) - return Number.isFinite(raw) && raw >= 0 ? raw : 2 + return Number.isFinite(raw) && raw >= 0 ? Math.min(raw, 5) : 2 } export async function* withRetry( diff --git a/src/utils/messages.ts b/src/utils/messages.ts index a1b3b17a..075e4b51 100644 --- a/src/utils/messages.ts +++ b/src/utils/messages.ts @@ -4723,12 +4723,12 @@ export type StreamingFallbackCause = | 'watchdog' | 'stream_error' | '404_stream_creation' + | 'stream_retry' /** - * Marks the switch from a failed streaming request to the non-streaming - * fallback. The fallback response arrives in one piece after a potentially - * long wait with zero incremental output, so UIs surface this as a lightweight - * active-turn status (level info — an expected state, not an error). + * Marks recovery from a failed streaming attempt. Most causes switch to the + * non-streaming fallback; stream_retry starts a new bounded streaming attempt. + * UIs surface this as an active-turn status rather than a terminal error. */ export function createSystemStreamingFallbackMessage( cause: StreamingFallbackCause, @@ -4737,7 +4737,9 @@ export function createSystemStreamingFallbackMessage( type: 'system', subtype: 'streaming_fallback', level: 'info', - content: `Streaming request failed (${cause.replace(/_/g, ' ')}); retrying in non-streaming mode`, + content: cause === 'stream_retry' + ? 'Provider stream stalled before a tool side effect; retrying safely' + : `Streaming request failed (${cause.replace(/_/g, ' ')}); retrying in non-streaming mode`, cause, timestamp: new Date().toISOString(), uuid: randomUUID(),