diff --git a/desktop/src/api/websocket.test.ts b/desktop/src/api/websocket.test.ts new file mode 100644 index 00000000..7481319e --- /dev/null +++ b/desktop/src/api/websocket.test.ts @@ -0,0 +1,92 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' + +vi.mock('./client', () => ({ + getBaseUrl: () => 'http://127.0.0.1:3456', +})) + +import { wsManager } from './websocket' + +type SocketHandler = (() => void) | ((event: { data: string }) => void) + +class FakeWebSocket { + static readonly CONNECTING = 0 + static readonly OPEN = 1 + static readonly CLOSING = 2 + static readonly CLOSED = 3 + static instances: FakeWebSocket[] = [] + + readonly url: string + readyState = FakeWebSocket.CONNECTING + onopen: SocketHandler | null = null + onmessage: SocketHandler | null = null + onclose: SocketHandler | null = null + onerror: SocketHandler | null = null + sent: string[] = [] + + constructor(url: string) { + this.url = url + FakeWebSocket.instances.push(this) + } + + send(data: string) { + this.sent.push(data) + } + + close() { + this.readyState = FakeWebSocket.CLOSED + ;(this.onclose as (() => void) | null)?.() + } + + open() { + this.readyState = FakeWebSocket.OPEN + ;(this.onopen as (() => void) | null)?.() + } + + fail() { + this.readyState = FakeWebSocket.CLOSED + ;(this.onclose as (() => void) | null)?.() + } +} + +describe('wsManager reconnect buffering', () => { + const originalWebSocket = globalThis.WebSocket + + beforeEach(() => { + vi.useFakeTimers() + FakeWebSocket.instances = [] + globalThis.WebSocket = FakeWebSocket as unknown as typeof WebSocket + wsManager.disconnectAll() + }) + + afterEach(() => { + wsManager.disconnectAll() + globalThis.WebSocket = originalWebSocket + vi.useRealTimers() + }) + + it('replays queued messages after an unexpected reconnect', async () => { + wsManager.connect('session-reconnect') + + const firstSocket = FakeWebSocket.instances[0] + expect(firstSocket?.url).toContain('/ws/session-reconnect') + + firstSocket!.open() + wsManager.send('session-reconnect', { type: 'user_message', content: 'first' }) + expect(firstSocket!.sent).toEqual([ + JSON.stringify({ type: 'user_message', content: 'first' }), + ]) + + firstSocket!.fail() + wsManager.send('session-reconnect', { type: 'user_message', content: 'queued while offline' }) + + await vi.advanceTimersByTimeAsync(1000) + + const secondSocket = FakeWebSocket.instances[1] + expect(secondSocket).toBeDefined() + secondSocket!.open() + + expect(secondSocket!.sent).toEqual([ + JSON.stringify({ type: 'user_message', content: 'queued while offline' }), + ]) + }) +}) diff --git a/desktop/src/api/websocket.ts b/desktop/src/api/websocket.ts index 44d9130b..2e10622b 100644 --- a/desktop/src/api/websocket.ts +++ b/desktop/src/api/websocket.ts @@ -27,7 +27,17 @@ class WebSocketManager { connect(sessionId: string) { const existing = this.connections.get(sessionId) - if (existing && !existing.intentionalClose) return + if ( + existing && + !existing.intentionalClose && + ( + existing.ws.readyState === WebSocket.OPEN || + existing.ws.readyState === WebSocket.CONNECTING || + existing.reconnectTimer !== null + ) + ) { + return + } const wsUrl = getBaseUrl().replace(/^http/, 'ws') const ws = new WebSocket(`${wsUrl}/ws/${sessionId}`) @@ -36,10 +46,10 @@ class WebSocketManager { ws, handlers: existing?.handlers ?? new Set(), reconnectTimer: null, - reconnectAttempt: 0, + reconnectAttempt: existing?.reconnectAttempt ?? 0, pingInterval: null, intentionalClose: false, - pendingMessages: [], + pendingMessages: existing?.pendingMessages ?? [], } this.connections.set(sessionId, conn) @@ -98,13 +108,27 @@ class WebSocketManager { } send(sessionId: string, message: ClientMessage) { - const conn = this.connections.get(sessionId) - if (!conn) return + let conn = this.connections.get(sessionId) + if (!conn) { + this.connect(sessionId) + conn = this.connections.get(sessionId) + if (!conn) return + } if (conn.ws.readyState === WebSocket.OPEN) { conn.ws.send(JSON.stringify(message)) - } else if (conn.ws.readyState === WebSocket.CONNECTING) { - conn.pendingMessages.push(message) + return + } + + conn.pendingMessages.push(message) + + if ( + conn.ws.readyState === WebSocket.CLOSED || + conn.ws.readyState === WebSocket.CLOSING + ) { + if (!conn.intentionalClose && !conn.reconnectTimer) { + this.scheduleReconnect(sessionId, conn) + } } } @@ -147,15 +171,8 @@ class WebSocketManager { conn.reconnectTimer = setTimeout(() => { if (this.connections.get(sessionId) === conn && !conn.intentionalClose) { - this.connections.delete(sessionId) + conn.reconnectTimer = null this.connect(sessionId) - // Migrate handlers to new connection - const newConn = this.connections.get(sessionId) - if (newConn) { - for (const handler of conn.handlers) { - newConn.handlers.add(handler) - } - } } }, delay) } diff --git a/src/server/__tests__/conversations.test.ts b/src/server/__tests__/conversations.test.ts index fa759077..18f143ec 100644 --- a/src/server/__tests__/conversations.test.ts +++ b/src/server/__tests__/conversations.test.ts @@ -211,6 +211,29 @@ describe('WebSocket Chat Integration', () => { } } + async function withMockStreamDelay( + delayMs: number | undefined, + callback: () => Promise, + ): Promise { + const previousDelay = process.env.MOCK_SDK_STREAM_DELAY_MS + + if (delayMs && delayMs > 0) { + process.env.MOCK_SDK_STREAM_DELAY_MS = String(delayMs) + } else { + delete process.env.MOCK_SDK_STREAM_DELAY_MS + } + + try { + return await callback() + } finally { + if (previousDelay === undefined) { + delete process.env.MOCK_SDK_STREAM_DELAY_MS + } else { + process.env.MOCK_SDK_STREAM_DELAY_MS = previousDelay + } + } + } + async function runTurn(sessionId: string, content: string): Promise { const messages: any[] = [] const ws = new WebSocket(`${wsUrl}/ws/${sessionId}`) @@ -480,4 +503,72 @@ describe('WebSocket Chat Integration', () => { expect(secondTurn.some((m) => m.type === 'message_complete')).toBe(true) expect(secondTurn.some((m) => m.type === 'error')).toBe(false) }) + + it('should resume streaming to a reconnected client during an active turn', async () => { + await withMockStreamDelay(150, async () => { + const sessionId = `chat-reconnect-${crypto.randomUUID()}` + const firstMessages: any[] = [] + const secondMessages: any[] = [] + + await new Promise((resolve, reject) => { + let reconnected = false + let ws2: WebSocket | null = null + const timeout = setTimeout(() => { + ws2?.close() + reject(new Error(`Timed out waiting for reconnect completion for session ${sessionId}`)) + }, 10_000) + + const cleanup = () => { + clearTimeout(timeout) + ws2?.close() + resolve() + } + + const handleFailure = (message: string) => { + clearTimeout(timeout) + ws2?.close() + reject(new Error(message)) + } + + const ws1 = new WebSocket(`${wsUrl}/ws/${sessionId}`) + ws1.onmessage = (event) => { + const msg = JSON.parse(event.data as string) + firstMessages.push(msg) + + if (msg.type === 'connected') { + ws1.send(JSON.stringify({ type: 'user_message', content: 'resume after reconnect' })) + return + } + + if (msg.type === 'thinking' && !reconnected) { + reconnected = true + ws1.close() + + setTimeout(() => { + ws2 = new WebSocket(`${wsUrl}/ws/${sessionId}`) + ws2.onmessage = (reconnectEvent) => { + const reconnectMsg = JSON.parse(reconnectEvent.data as string) + secondMessages.push(reconnectMsg) + if (reconnectMsg.type === 'error') { + handleFailure(reconnectMsg.message) + return + } + if (reconnectMsg.type === 'message_complete') { + cleanup() + } + } + ws2.onerror = () => handleFailure(`WebSocket reconnect error for session ${sessionId}`) + }, 50) + } + } + + ws1.onerror = () => handleFailure(`Initial WebSocket error for session ${sessionId}`) + }) + + expect(firstMessages.some((msg) => msg.type === 'thinking')).toBe(true) + expect(secondMessages.some((msg) => msg.type === 'connected')).toBe(true) + expect(secondMessages.some((msg) => msg.type === 'content_delta')).toBe(true) + expect(secondMessages.some((msg) => msg.type === 'message_complete')).toBe(true) + }) + }) }) diff --git a/src/server/__tests__/fixtures/mock-sdk-cli.ts b/src/server/__tests__/fixtures/mock-sdk-cli.ts index 54119c1f..b0e624c3 100644 --- a/src/server/__tests__/fixtures/mock-sdk-cli.ts +++ b/src/server/__tests__/fixtures/mock-sdk-cli.ts @@ -9,6 +9,10 @@ function emit(ws: WebSocket, payload: Record) { ws.send(JSON.stringify(payload) + '\n') } +function delay(ms: number) { + return new Promise((resolve) => setTimeout(resolve, ms)) +} + function extractUserText(message: any): string { const content = message?.message?.content if (!Array.isArray(content)) return '' @@ -21,6 +25,7 @@ function extractUserText(message: any): string { const sdkUrl = getArg('--sdk-url') const sessionId = getArg('--session-id') || crypto.randomUUID() const initMode = process.env.MOCK_SDK_INIT_MODE || 'on_open' +const streamDelayMs = Number(process.env.MOCK_SDK_STREAM_DELAY_MS || '0') let initSent = false if (!sdkUrl) { @@ -52,70 +57,74 @@ ws.addEventListener('message', (event) => { const payload = typeof event.data === 'string' ? event.data : String(event.data) const lines = payload.split('\n').map(line => line.trim()).filter(Boolean) - for (const line of lines) { - const parsed = JSON.parse(line) + void (async () => { + for (const line of lines) { + const parsed = JSON.parse(line) - if (parsed.type === 'user') { - sendInit() - const text = extractUserText(parsed) - emit(ws, { - type: 'stream_event', - event: { type: 'message_start' }, - session_id: sessionId, - }) - emit(ws, { - type: 'stream_event', - event: { - type: 'content_block_start', - index: 0, - content_block: { type: 'text', text: '' }, - }, - session_id: sessionId, - }) - emit(ws, { - type: 'stream_event', - event: { - type: 'content_block_delta', - index: 0, - delta: { type: 'thinking_delta', thinking: 'Mock thinking...' }, - }, - session_id: sessionId, - }) - emit(ws, { - type: 'stream_event', - event: { - type: 'content_block_delta', - index: 0, - delta: { type: 'text_delta', text: `Echo: ${text}` }, - }, - session_id: sessionId, - }) - emit(ws, { - type: 'stream_event', - event: { type: 'content_block_stop', index: 0 }, - session_id: sessionId, - }) - emit(ws, { - type: 'result', - subtype: 'success', - is_error: false, - result: `Echo: ${text}`, - usage: { input_tokens: 3, output_tokens: 2 }, - session_id: sessionId, - }) - } + if (parsed.type === 'user') { + sendInit() + const text = extractUserText(parsed) + emit(ws, { + type: 'stream_event', + event: { type: 'message_start' }, + session_id: sessionId, + }) + emit(ws, { + type: 'stream_event', + event: { + type: 'content_block_start', + index: 0, + content_block: { type: 'text', text: '' }, + }, + session_id: sessionId, + }) + emit(ws, { + type: 'stream_event', + event: { + type: 'content_block_delta', + index: 0, + delta: { type: 'thinking_delta', thinking: 'Mock thinking...' }, + }, + session_id: sessionId, + }) + if (streamDelayMs > 0) await delay(streamDelayMs) + emit(ws, { + type: 'stream_event', + event: { + type: 'content_block_delta', + index: 0, + delta: { type: 'text_delta', text: `Echo: ${text}` }, + }, + session_id: sessionId, + }) + if (streamDelayMs > 0) await delay(streamDelayMs) + emit(ws, { + type: 'stream_event', + event: { type: 'content_block_stop', index: 0 }, + session_id: sessionId, + }) + emit(ws, { + type: 'result', + subtype: 'success', + is_error: false, + result: `Echo: ${text}`, + usage: { input_tokens: 3, output_tokens: 2 }, + session_id: sessionId, + }) + } - if (parsed.type === 'control_request' && parsed.request?.subtype === 'interrupt') { - emit(ws, { - type: 'result', - subtype: 'success', - is_error: false, - result: 'Interrupted', - usage: { input_tokens: 0, output_tokens: 0 }, - session_id: sessionId, - }) + if (parsed.type === 'control_request' && parsed.request?.subtype === 'interrupt') { + emit(ws, { + type: 'result', + subtype: 'success', + is_error: false, + result: 'Interrupted', + usage: { input_tokens: 0, output_tokens: 0 }, + session_id: sessionId, + }) + } } - } + })() }) ws.addEventListener('close', () => { diff --git a/src/server/__tests__/sessions.test.ts b/src/server/__tests__/sessions.test.ts index 680bcfcb..ece8e50e 100644 --- a/src/server/__tests__/sessions.test.ts +++ b/src/server/__tests__/sessions.test.ts @@ -263,6 +263,52 @@ describe('SessionService', () => { expect(messages).toHaveLength(2) }) + it('should hide synthetic interruption, no-response, and command breadcrumb transcript entries', async () => { + const sessionId = 'aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee' + await writeSessionFile('-tmp-project', sessionId, [ + makeSnapshotEntry(), + makeUserEntry('正常用户消息', crypto.randomUUID()), + { + type: 'user', + message: { + role: 'user', + content: [{ type: 'text', text: '[Request interrupted by user]' }], + }, + uuid: crypto.randomUUID(), + timestamp: '2026-01-01T00:00:02.000Z', + }, + { + type: 'assistant', + message: { + role: 'assistant', + content: [{ type: 'text', text: 'No response requested.' }], + model: '', + }, + uuid: crypto.randomUUID(), + timestamp: '2026-01-01T00:00:03.000Z', + }, + { + type: 'user', + message: { + role: 'user', + content: '/exit\nexit\n', + }, + uuid: crypto.randomUUID(), + timestamp: '2026-01-01T00:00:04.000Z', + }, + makeAssistantEntry('正常助手消息', crypto.randomUUID()), + ]) + + const messages = await service.getSessionMessages(sessionId) + + expect(messages).toHaveLength(2) + expect(messages[0]).toMatchObject({ type: 'user', content: '正常用户消息' }) + expect(messages[1]).toMatchObject({ + type: 'assistant', + content: [{ type: 'text', text: '正常助手消息' }], + }) + }) + it('should reconstruct parent agent tool linkage from parentUuid chains', async () => { const sessionId = 'aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee' const userUuid = crypto.randomUUID() diff --git a/src/server/services/sessionService.ts b/src/server/services/sessionService.ts index 0787bb74..73bcab5d 100644 --- a/src/server/services/sessionService.ts +++ b/src/server/services/sessionService.ts @@ -71,6 +71,13 @@ type RawEntry = { [key: string]: unknown } +const USER_INTERRUPTION_TEXTS = new Set([ + '[Request interrupted by user]', + '[Request interrupted by user for tool use]', +]) + +const NO_RESPONSE_REQUESTED_TEXT = 'No response requested.' + // ============================================================================ // Service // ============================================================================ @@ -204,6 +211,67 @@ export class SessionService { } } + private extractTextBlocks(content: unknown): string[] { + if (typeof content === 'string') return [content] + if (!Array.isArray(content)) return [] + + return content + .flatMap((block) => { + if (!block || typeof block !== 'object') return [] + const record = block as Record + return record.type === 'text' && typeof record.text === 'string' + ? [record.text] + : [] + }) + .map((text) => text.trim()) + .filter(Boolean) + } + + private isInternalCommandBreadcrumb(content: unknown): boolean { + if (typeof content !== 'string') return false + + return ( + content.includes('') || + content.includes('') || + content.includes('') || + content.includes('') + ) + } + + private isSyntheticUserInterruption(content: unknown): boolean { + const textBlocks = this.extractTextBlocks(content) + return ( + textBlocks.length > 0 && + textBlocks.every((text) => USER_INTERRUPTION_TEXTS.has(text)) + ) + } + + private isSyntheticNoResponseAssistant(content: unknown): boolean { + const textBlocks = this.extractTextBlocks(content) + return ( + textBlocks.length > 0 && + textBlocks.every((text) => text === NO_RESPONSE_REQUESTED_TEXT) + ) + } + + private shouldHideTranscriptEntry(entry: RawEntry): boolean { + const role = entry.message?.role + const content = entry.message?.content + + if (role === 'user') { + return ( + this.isInternalCommandBreadcrumb(content) || + this.isSyntheticUserInterruption(content) + ) + } + + if (role === 'assistant') { + return this.isSyntheticNoResponseAssistant(content) + } + + return false + } + private extractAgentToolUseId(entry: RawEntry): string | undefined { const content = entry.message?.content if (!Array.isArray(content)) return undefined @@ -754,6 +822,8 @@ export class SessionService { // Skip meta entries (CLI internal bookkeeping) if (entry.isMeta) continue + if (this.shouldHideTranscriptEntry(entry)) continue + // Skip non-transcript entry types const entryType = entry.type if ( diff --git a/src/server/ws/handler.ts b/src/server/ws/handler.ts index 14d0b91d..86ae831a 100644 --- a/src/server/ws/handler.ts +++ b/src/server/ws/handler.ts @@ -91,6 +91,7 @@ export const handleWebSocket = { } activeSessions.set(sessionId, ws) + rebindSessionOutput(sessionId, ws) const msg: ServerMessage = { type: 'connected', sessionId } ws.send(JSON.stringify(msg)) @@ -155,9 +156,7 @@ export const handleWebSocket = { console.log(`[WS] Client disconnected from session: ${sessionId} (${code}: ${reason})`) computerUseApprovalService.cancelSession(sessionId) activeSessions.delete(sessionId) - cleanupStreamState(sessionId) - sessionSlashCommands.delete(sessionId) - sessionTitleState.delete(sessionId) + conversationService.clearOutputCallbacks(sessionId) // Schedule delayed cleanup: if the client doesn't reconnect within 30 seconds, // stop the CLI subprocess to avoid leaking resources. @@ -166,6 +165,7 @@ export const handleWebSocket = { if (!activeSessions.has(sessionId)) { console.log(`[WS] Session ${sessionId} not reconnected after 30s, stopping CLI subprocess`) conversationService.stopSession(sessionId) + cleanupSessionRuntimeState(sessionId) } }, 30_000) sessionCleanupTimers.set(sessionId, cleanupTimer) @@ -252,26 +252,8 @@ async function handleUserMessage( // any pre-turn SDK chatter as fresh chat history. let userMessageSent = false - conversationService.clearOutputCallbacks(sessionId) - conversationService.onOutput(sessionId, (cliMsg) => { - // Before the current turn is sent, only surface startup errors. - if (!userMessageSent) { - if (cliMsg.type === 'result' && cliMsg.is_error) { - const serverMsgs = translateCliMessage(cliMsg, sessionId) - for (const msg of serverMsgs) { - sendMessage(ws, msg) - } - } - return - } - const serverMsgs = translateCliMessage(cliMsg, sessionId) - for (const msg of serverMsgs) { - sendMessage(ws, msg) - } - // Trigger title generation on message_complete - if (cliMsg.type === 'result') { - triggerTitleGeneration(ws, sessionId) - } + rebindSessionOutput(sessionId, ws, { + shouldForward: (cliMsg) => userMessageSent || (cliMsg.type === 'result' && cliMsg.is_error), }) const sent = conversationService.sendMessage( @@ -485,6 +467,12 @@ function cleanupStreamState(sessionId: string) { sessionStreamStates.delete(sessionId) } +function cleanupSessionRuntimeState(sessionId: string) { + cleanupStreamState(sessionId) + sessionSlashCommands.delete(sessionId) + sessionTitleState.delete(sessionId) +} + function translateCliMessage(cliMsg: any, sessionId: string): ServerMessage[] { const streamState = getStreamState(sessionId) switch (cliMsg.type) { @@ -831,6 +819,32 @@ function sendError(ws: ServerWebSocket, message: string, code: st sendMessage(ws, { type: 'error', message, code }) } +function rebindSessionOutput( + sessionId: string, + ws: ServerWebSocket, + options?: { + shouldForward?: (cliMsg: any) => boolean + }, +) { + if (!conversationService.hasSession(sessionId)) return + + conversationService.clearOutputCallbacks(sessionId) + conversationService.onOutput(sessionId, (cliMsg) => { + if (options?.shouldForward && !options.shouldForward(cliMsg)) { + return + } + + const serverMsgs = translateCliMessage(cliMsg, sessionId) + for (const msg of serverMsgs) { + sendMessage(ws, msg) + } + + if (cliMsg.type === 'result') { + triggerTitleGeneration(ws, sessionId) + } + }) +} + async function getRuntimeSettings(): Promise<{ permissionMode?: string model?: string