diff --git a/src/server/__tests__/websocket-handler.test.ts b/src/server/__tests__/websocket-handler.test.ts index 4584ec2b..34adcdba 100644 --- a/src/server/__tests__/websocket-handler.test.ts +++ b/src/server/__tests__/websocket-handler.test.ts @@ -1,6 +1,7 @@ import { afterEach, describe, expect, it, mock, spyOn } from 'bun:test' import type { ServerWebSocket } from 'bun' import { + __markPrewarmPendingForTests, __resetWebSocketHandlerStateForTests, closeSessionConnection, getActiveSessionIds, @@ -135,4 +136,30 @@ describe('WebSocket handler session isolation', () => { expect(setTimeoutSpy).toHaveBeenCalled() expect(setTimeoutSpy.mock.calls[0]?.[1]).toBeGreaterThan(30_000) }) + + it('does not forward prewarm startup status to a reconnecting client', async () => { + const sessionId = `prewarm-reconnect-${crypto.randomUUID()}` + const second = makeClientSocket(sessionId) + let outputCallback: ((cliMsg: any) => void) | null = null + + __markPrewarmPendingForTests(sessionId) + spyOn(conversationService, 'hasSession').mockReturnValue(true) + spyOn(conversationService, 'getRecentSdkMessages').mockReturnValue([]) + spyOn(conversationService, 'onOutput').mockImplementation((_sid, callback) => { + outputCallback = callback + }) + spyOn(conversationService, 'removeOutputCallback').mockImplementation(() => {}) + spyOn(conversationService, 'clearOutputCallbacks').mockImplementation(() => { + outputCallback = null + }) + + handleWebSocket.open(second) + outputCallback?.({ + type: 'stream_event', + event: { type: 'message_start' }, + }) + + const secondMessages = second.sent.map((payload) => JSON.parse(payload)) + expect(secondMessages).not.toContainEqual({ type: 'status', state: 'thinking' }) + }) }) diff --git a/src/server/ws/handler.ts b/src/server/ws/handler.ts index 2c51c768..6abb0107 100644 --- a/src/server/ws/handler.ts +++ b/src/server/ws/handler.ts @@ -150,7 +150,7 @@ export const handleWebSocket = { } addActiveClient(sessionId, ws) - if (prewarmedSessions.has(sessionId)) { + if (prewarmPendingSessions.has(sessionId) || prewarmedSessions.has(sessionId)) { bindPrewarmMetadataCapture(sessionId) } else { bindClientSessionOutput(sessionId, ws) @@ -2235,5 +2235,11 @@ export function __resetWebSocketHandlerStateForTests(): void { activeSessions.clear() clientOutputCallbacks.clear() sessionCleanupTimers.clear() + prewarmPendingSessions.clear() + prewarmedSessions.clear() prewarmIdleTimers.clear() } + +export function __markPrewarmPendingForTests(sessionId: string): void { + prewarmPendingSessions.add(sessionId) +}