fix(server): close prewarm idle timer race in CLI-startup blind window

#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.
This commit is contained in:
程序员阿江(Relakkes) 2026-06-20 16:17:08 +08:00
parent 0b93792e78
commit 9111c0c242
2 changed files with 103 additions and 8 deletions

View File

@ -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)
})
})

View File

@ -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<WebSocketData>) {
.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)
}