mirror of
https://github.com/NanmiCoder/cc-haha
synced 2026-07-18 13:23:33 +08:00
9111c0c242
#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.
307 lines
12 KiB
TypeScript
307 lines
12 KiB
TypeScript
import { afterEach, describe, expect, it, mock, spyOn } from 'bun:test'
|
|
import type { ServerWebSocket } from 'bun'
|
|
import {
|
|
__markPrewarmPendingForTests,
|
|
__markActiveTurnForTests,
|
|
__registerPendingUserTurnForTests,
|
|
__markPrewarmedForTests,
|
|
__resetWebSocketHandlerStateForTests,
|
|
closeSessionConnection,
|
|
getActiveSessionIds,
|
|
handleWebSocket,
|
|
type WebSocketData,
|
|
} from '../ws/handler.js'
|
|
import {
|
|
__resetDisconnectGraceMsForTests,
|
|
__setDisconnectGraceMsForTests,
|
|
} from '../ws/disconnectGraceConfig.js'
|
|
import { conversationService } from '../services/conversationService.js'
|
|
import { computerUseApprovalService } from '../services/computerUseApprovalService.js'
|
|
|
|
function makeClientSocket(sessionId: string) {
|
|
const sent: string[] = []
|
|
return {
|
|
data: {
|
|
sessionId,
|
|
connectedAt: Date.now(),
|
|
channel: 'client',
|
|
sdkToken: null,
|
|
serverPort: 0,
|
|
serverHost: '127.0.0.1',
|
|
},
|
|
send: mock((payload: string) => {
|
|
sent.push(payload)
|
|
}),
|
|
close: mock(() => {}),
|
|
sent,
|
|
} as unknown as ServerWebSocket<WebSocketData> & { sent: string[] }
|
|
}
|
|
|
|
describe('WebSocket handler session isolation', () => {
|
|
afterEach(() => {
|
|
__resetWebSocketHandlerStateForTests()
|
|
__resetDisconnectGraceMsForTests()
|
|
mock.restore()
|
|
})
|
|
|
|
it('ignores stale disconnects from an older socket for the same session', () => {
|
|
const sessionId = `duplicate-${crypto.randomUUID()}`
|
|
const first = makeClientSocket(sessionId)
|
|
const second = makeClientSocket(sessionId)
|
|
const clearCallbacks = spyOn(conversationService, 'clearOutputCallbacks')
|
|
const cancelComputerUse = spyOn(computerUseApprovalService, 'cancelSession')
|
|
|
|
handleWebSocket.open(first)
|
|
handleWebSocket.open(second)
|
|
clearCallbacks.mockClear()
|
|
cancelComputerUse.mockClear()
|
|
|
|
handleWebSocket.close(first, 1000, 'stale tab closed')
|
|
|
|
expect(getActiveSessionIds()).toContain(sessionId)
|
|
expect(clearCallbacks).not.toHaveBeenCalled()
|
|
expect(cancelComputerUse).not.toHaveBeenCalled()
|
|
})
|
|
|
|
it('closes and removes an active client socket when a session is deleted', () => {
|
|
const sessionId = `delete-${crypto.randomUUID()}`
|
|
const ws = makeClientSocket(sessionId)
|
|
const clearCallbacks = spyOn(conversationService, 'clearOutputCallbacks')
|
|
const cancelComputerUse = spyOn(computerUseApprovalService, 'cancelSession')
|
|
|
|
handleWebSocket.open(ws)
|
|
|
|
expect(closeSessionConnection(sessionId, 'session deleted')).toBe(true)
|
|
|
|
expect(getActiveSessionIds()).not.toContain(sessionId)
|
|
expect(ws.close).toHaveBeenCalledWith(1000, 'session deleted')
|
|
expect(clearCallbacks).toHaveBeenCalledWith(sessionId)
|
|
expect(cancelComputerUse).toHaveBeenCalledWith(sessionId)
|
|
})
|
|
|
|
it('replays pending permission requests when a client reconnects', () => {
|
|
const sessionId = `permission-reconnect-${crypto.randomUUID()}`
|
|
const ws = makeClientSocket(sessionId)
|
|
spyOn(conversationService, 'hasSession').mockReturnValue(true)
|
|
spyOn(conversationService, 'onOutput').mockImplementation(() => {})
|
|
spyOn(conversationService, 'removeOutputCallback').mockImplementation(() => {})
|
|
spyOn(conversationService, 'getPendingPermissionRequests').mockReturnValue([
|
|
{
|
|
requestId: 'request-ask-1',
|
|
toolName: 'AskUserQuestion',
|
|
toolUseId: 'tool-ask-1',
|
|
input: {
|
|
questions: [
|
|
{
|
|
header: 'Scope',
|
|
question: 'Which scope?',
|
|
options: [{ label: 'A', description: 'First' }, { label: 'B', description: 'Second' }],
|
|
},
|
|
],
|
|
},
|
|
description: 'Answer questions?',
|
|
},
|
|
])
|
|
|
|
handleWebSocket.open(ws)
|
|
|
|
expect(ws.sent.map((payload) => JSON.parse(payload))).toContainEqual({
|
|
type: 'permission_request',
|
|
requestId: 'request-ask-1',
|
|
toolName: 'AskUserQuestion',
|
|
toolUseId: 'tool-ask-1',
|
|
input: {
|
|
questions: [
|
|
{
|
|
header: 'Scope',
|
|
question: 'Which scope?',
|
|
options: [{ label: 'A', description: 'First' }, { label: 'B', description: 'Second' }],
|
|
},
|
|
],
|
|
},
|
|
description: 'Answer questions?',
|
|
})
|
|
})
|
|
|
|
it('keeps disconnected sessions alive longer while user input is pending', () => {
|
|
const sessionId = `permission-disconnect-${crypto.randomUUID()}`
|
|
const ws = makeClientSocket(sessionId)
|
|
const setTimeoutSpy = spyOn(globalThis, 'setTimeout').mockImplementation(() => 0 as any)
|
|
spyOn(conversationService, 'getPendingPermissionRequests').mockReturnValue([
|
|
{
|
|
requestId: 'request-ask-1',
|
|
toolName: 'AskUserQuestion',
|
|
toolUseId: 'tool-ask-1',
|
|
input: { questions: [] },
|
|
},
|
|
])
|
|
|
|
handleWebSocket.open(ws)
|
|
setTimeoutSpy.mockClear()
|
|
|
|
handleWebSocket.close(ws, 1006, 'renderer reconnecting')
|
|
|
|
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' })
|
|
})
|
|
|
|
it('keeps a running session alive on disconnect and cleans up only after the turn finishes (issue #764)', () => {
|
|
const sessionId = `running-disconnect-${crypto.randomUUID()}`
|
|
const ws = makeClientSocket(sessionId)
|
|
const setTimeoutSpy = spyOn(globalThis, 'setTimeout')
|
|
const stopSession = spyOn(conversationService, 'stopSession').mockImplementation(() => {})
|
|
spyOn(conversationService, 'getPendingPermissionRequests').mockReturnValue([])
|
|
|
|
let turnCompleteCallback: ((cliMsg: any) => void) | null = null
|
|
spyOn(conversationService, 'onOutput').mockImplementation((_sid, cb) => {
|
|
turnCompleteCallback = cb
|
|
})
|
|
spyOn(conversationService, 'removeOutputCallback').mockImplementation(() => {})
|
|
|
|
handleWebSocket.open(ws)
|
|
__markActiveTurnForTests(sessionId)
|
|
setTimeoutSpy.mockClear()
|
|
|
|
// Last client disconnects while the turn is still running: no kill timer,
|
|
// just a turn-completion watcher.
|
|
handleWebSocket.close(ws, 1006, 'phone locked screen')
|
|
expect(setTimeoutSpy).not.toHaveBeenCalled()
|
|
expect(stopSession).not.toHaveBeenCalled()
|
|
expect(turnCompleteCallback).not.toBeNull()
|
|
|
|
// Turn finishes while still disconnected → now the idle grace timer starts.
|
|
turnCompleteCallback?.({ type: 'result', subtype: 'success' })
|
|
expect(setTimeoutSpy).toHaveBeenCalled()
|
|
// Timer body still hasn't run, so the process is not killed yet.
|
|
expect(stopSession).not.toHaveBeenCalled()
|
|
})
|
|
|
|
it('uses the configured disconnect grace period for an idle session', () => {
|
|
const sessionId = `idle-disconnect-${crypto.randomUUID()}`
|
|
const ws = makeClientSocket(sessionId)
|
|
__setDisconnectGraceMsForTests(120_000)
|
|
const setTimeoutSpy = spyOn(globalThis, 'setTimeout').mockImplementation(() => 0 as any)
|
|
spyOn(conversationService, 'getPendingPermissionRequests').mockReturnValue([])
|
|
|
|
handleWebSocket.open(ws)
|
|
setTimeoutSpy.mockClear()
|
|
|
|
handleWebSocket.close(ws, 1006, 'tab closed')
|
|
|
|
expect(setTimeoutSpy).toHaveBeenCalled()
|
|
expect(setTimeoutSpy.mock.calls[0]?.[1]).toBe(120_000)
|
|
})
|
|
|
|
it('does not start the idle timer if the client reconnects before the turn finishes', () => {
|
|
const sessionId = `reconnect-mid-turn-${crypto.randomUUID()}`
|
|
const ws = makeClientSocket(sessionId)
|
|
const reconnected = makeClientSocket(sessionId)
|
|
const setTimeoutSpy = spyOn(globalThis, 'setTimeout')
|
|
spyOn(conversationService, 'getPendingPermissionRequests').mockReturnValue([])
|
|
spyOn(conversationService, 'hasSession').mockReturnValue(true)
|
|
|
|
let turnCompleteCallback: ((cliMsg: any) => void) | null = null
|
|
spyOn(conversationService, 'onOutput').mockImplementation((_sid, cb) => {
|
|
turnCompleteCallback = cb
|
|
})
|
|
const removeOutputCallback = spyOn(conversationService, 'removeOutputCallback').mockImplementation(() => {})
|
|
|
|
handleWebSocket.open(ws)
|
|
__markActiveTurnForTests(sessionId)
|
|
handleWebSocket.close(ws, 1006, 'phone locked screen')
|
|
expect(turnCompleteCallback).not.toBeNull()
|
|
|
|
// Reconnect tears down the watcher before the turn completes.
|
|
handleWebSocket.open(reconnected)
|
|
expect(removeOutputCallback).toHaveBeenCalled()
|
|
setTimeoutSpy.mockClear()
|
|
|
|
// A late result must not schedule cleanup now that a client is back.
|
|
turnCompleteCallback?.({ type: 'result', subtype: 'success' })
|
|
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)
|
|
})
|
|
})
|