From 9111c0c2422103c0d72838069e7050bd8dfa3629 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: Sat, 20 Jun 2026 16:17:08 +0800 Subject: [PATCH] fix(server): close prewarm idle timer race in CLI-startup blind window MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit #865 added guards against the prewarm_session/user_message concurrency race, but both keyed on isSessionTurnActive (messageSent===true). messageSent only flips true after CLI startup completes, so during the startup window a user turn is registered yet invisible to the guards — the prewarm idle timer could still arm on, and later fire against, an active turn, killing the CLI and leaving the UI stuck on "执行中". - Add hasPendingOrActiveUserTurn() (turn registered regardless of messageSent) and use it for both prewarm guards, covering the blind window. - Structurally disarm the timer when a turn completes (clearPrewarmState in bindActiveUserTurnCompletion), so a timer armed by the race never outlives the turn regardless of ordering. - Add test hooks + 3 regression tests: blind-window turn not killed, active turn not killed, truly-idle session still reclaimed. --- .../__tests__/websocket-handler.test.ts | 59 +++++++++++++++++++ src/server/ws/handler.ts | 52 +++++++++++++--- 2 files changed, 103 insertions(+), 8 deletions(-) diff --git a/src/server/__tests__/websocket-handler.test.ts b/src/server/__tests__/websocket-handler.test.ts index 1241bc99..dcd20d0f 100644 --- a/src/server/__tests__/websocket-handler.test.ts +++ b/src/server/__tests__/websocket-handler.test.ts @@ -3,6 +3,8 @@ import type { ServerWebSocket } from 'bun' import { __markPrewarmPendingForTests, __markActiveTurnForTests, + __registerPendingUserTurnForTests, + __markPrewarmedForTests, __resetWebSocketHandlerStateForTests, closeSessionConnection, getActiveSessionIds, @@ -245,3 +247,60 @@ describe('WebSocket handler session isolation', () => { expect(setTimeoutSpy).not.toHaveBeenCalled() }) }) + +describe('prewarm idle timer active-turn guard (issue #865 follow-up)', () => { + afterEach(() => { + __resetWebSocketHandlerStateForTests() + mock.restore() + }) + + // Arm the prewarm idle timer the way markPrewarmed does, and return its fire + // callback so a test can trigger it deterministically without waiting 5 min. + function armPrewarmIdleTimer(sessionId: string): () => void { + const setTimeoutSpy = spyOn(globalThis, 'setTimeout').mockImplementation( + (() => 0) as unknown as typeof setTimeout, + ) + __markPrewarmedForTests(sessionId) + const fire = setTimeoutSpy.mock.calls.at(-1)?.[0] as (() => void) | undefined + if (!fire) throw new Error('prewarm idle timer was not armed') + return fire + } + + it('does not kill a prewarmed session once a user turn is registered, even before messageSent flips (CLI-startup blind window)', () => { + const sessionId = `prewarm-blind-window-${crypto.randomUUID()}` + const stopSession = spyOn(conversationService, 'stopSession').mockImplementation(() => {}) + const fire = armPrewarmIdleTimer(sessionId) + + // The concurrent prewarm_session/user_message race: the turn is registered + // (activeUserTurns has it) but messageSent is still false during CLI startup + // when the idle timer fires. The old isSessionTurnActive guard was blind to + // this window — the turn-registered guard must catch it. + __registerPendingUserTurnForTests(sessionId) + fire() + + expect(stopSession).not.toHaveBeenCalled() + }) + + it('does not kill a prewarmed session with a fully active (messageSent) turn', () => { + const sessionId = `prewarm-active-turn-${crypto.randomUUID()}` + const stopSession = spyOn(conversationService, 'stopSession').mockImplementation(() => {}) + const fire = armPrewarmIdleTimer(sessionId) + + __markActiveTurnForTests(sessionId) + fire() + + expect(stopSession).not.toHaveBeenCalled() + }) + + it('still reclaims a truly idle prewarmed session with no turn and no clients', () => { + const sessionId = `prewarm-truly-idle-${crypto.randomUUID()}` + const stopSession = spyOn(conversationService, 'stopSession').mockImplementation(() => {}) + const fire = armPrewarmIdleTimer(sessionId) + + // No registered turn and no connected client → the reaper must still fire, + // otherwise the timer's whole purpose (reclaiming idle prewarmed CLIs) is lost. + fire() + + expect(stopSession).toHaveBeenCalledWith(sessionId) + }) +}) diff --git a/src/server/ws/handler.ts b/src/server/ws/handler.ts index 5f785c04..e1e37b6d 100644 --- a/src/server/ws/handler.ts +++ b/src/server/ws/handler.ts @@ -455,6 +455,11 @@ function bindActiveUserTurnCompletion( conversationService.removeOutputCallback(sessionId, callback) clearActiveUserTurn(sessionId, activeTurn) + // Structurally disarm any prewarm idle timer that a concurrent + // prewarm_session/user_message flush may have armed on this session: once a + // turn completes the session is firmly user-owned, so no prewarm reaper + // should survive — regardless of the order in which the two raced. + clearPrewarmState(sessionId) applyDeferredPermissionModeAfterActiveTurn(ws, sessionId) applyDeferredRuntimeRestartAfterActiveTurn(ws, sessionId) } @@ -566,10 +571,13 @@ async function handlePrewarmSession(ws: ServerWebSocket) { .then(() => { const stillPending = prewarmPendingSessions.delete(sessionId) if (!stillPending) return - // Safety: if a user message arrived and activated a turn while we were - // waiting for startup, do NOT arm the prewarm idle timer — the session - // is now owned by the user conversation, not prewarm. - if (isSessionTurnActive(sessionId)) { + // Safety: if a user message arrived and claimed this session while we + // were waiting for startup, do NOT arm the prewarm idle timer — the + // session is now owned by the user conversation, not prewarm. Use the + // turn-registered check (not messageSent) so the CLI-startup window is + // covered: in the concurrent race the turn is registered but messageSent + // is still false when this .then runs, which made the old guard dead code. + if (hasPendingOrActiveUserTurn(sessionId)) { return } bindPrewarmMetadataCapture(sessionId) @@ -1246,11 +1254,13 @@ function markPrewarmed(sessionId: string) { const timer = setTimeout(() => { prewarmIdleTimers.delete(sessionId) if (!prewarmedSessions.has(sessionId)) return - const turnActive = isSessionTurnActive(sessionId) + const turnActive = hasPendingOrActiveUserTurn(sessionId) const hasClients = hasActiveClients(sessionId) - // Safety guard: never kill a session that has an active user turn or - // connected clients. The prewarm idle timer is only meant to reclaim - // truly idle prewarmed sessions — not to interrupt an active conversation. + // Safety guard: never kill a session that has a registered user turn or + // connected clients. The turn-registered check (not messageSent) covers the + // CLI-startup window, so a turn racing through startup is protected even if + // the client has briefly disconnected. The prewarm idle timer is only meant + // to reclaim truly idle prewarmed sessions — not to interrupt a conversation. if (turnActive || hasClients) { prewarmedSessions.delete(sessionId) return @@ -2004,6 +2014,19 @@ function isSessionTurnActive(sessionId: string): boolean { return activeUserTurns.get(sessionId)?.messageSent === true } +/** + * Whether a user turn has been registered for this session and not yet settled, + * INCLUDING the CLI-startup window before messageSent flips true. handleUserMessage + * registers the turn in its synchronous prefix (activeUserTurns.set), well before + * the message is actually sent. Unlike isSessionTurnActive, this is not blind to + * that window, so the prewarm idle timer can neither arm on nor fire against a + * session a user turn has already claimed — even when a concurrent + * prewarm_session/user_message flush inverts their ordering. + */ +function hasPendingOrActiveUserTurn(sessionId: string): boolean { + return activeUserTurns.has(sessionId) +} + /** * Start the idle grace timer for a disconnected, idle session. If no client * reconnects before it fires, the CLI subprocess is stopped. @@ -2771,3 +2794,16 @@ export function __markPrewarmPendingForTests(sessionId: string): void { export function __markActiveTurnForTests(sessionId: string): void { activeUserTurns.set(sessionId, { messageSent: true }) } + +/** + * Test hook: register a user turn still in the pre-send (messageSent:false) + * window — i.e. the CLI-startup window that isSessionTurnActive is blind to. + */ +export function __registerPendingUserTurnForTests(sessionId: string): void { + activeUserTurns.set(sessionId, { messageSent: false }) +} + +/** Test hook: arm the prewarm idle timer for a session, as markPrewarmed does. */ +export function __markPrewarmedForTests(sessionId: string): void { + markPrewarmed(sessionId) +}