diff --git a/src/server/__tests__/conversations.test.ts b/src/server/__tests__/conversations.test.ts index e0b82bed..e12e846c 100644 --- a/src/server/__tests__/conversations.test.ts +++ b/src/server/__tests__/conversations.test.ts @@ -2514,6 +2514,417 @@ describe('WebSocket Chat Integration', () => { } }, 20_000) + it('should clear active turn tracking when sending a user message fails after startup', async () => { + const providerService = new ProviderService() + const provider = await providerService.addProvider({ + presetId: 'custom', + name: 'Provider Send Failure Runtime', + apiKey: 'key-send-failure-runtime', + baseUrl: 'http://127.0.0.1:1/anthropic', + apiFormat: 'anthropic', + models: { + main: 'send-failure-main', + haiku: 'send-failure-haiku', + sonnet: 'send-failure-sonnet', + opus: 'send-failure-opus', + }, + }) + + const createRes = await fetch(`${baseUrl}/api/sessions`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ workDir: process.cwd() }), + }) + expect(createRes.status).toBe(201) + const { sessionId } = await createRes.json() as { sessionId: string } + + const originalStartSession = conversationService.startSession.bind(conversationService) + const originalSendMessage = conversationService.sendMessage.bind(conversationService) + const startCalls: Array<{ + sessionId: string + options: { permissionMode?: string; model?: string; effort?: string; providerId?: string | null } | undefined + }> = [] + + conversationService.startSession = (async function patchedStartSession( + sid: string, + workDir: string, + sdkUrl: string, + options?: { permissionMode?: string; model?: string; effort?: string; thinking?: 'enabled' | 'adaptive' | 'disabled'; providerId?: string | null }, + ) { + startCalls.push({ sessionId: sid, options }) + return originalStartSession(sid, workDir, sdkUrl, options) + }) as typeof conversationService.startSession + conversationService.sendMessage = (async () => false) as typeof conversationService.sendMessage + + const ws = new WebSocket(`${wsUrl}/ws/${sessionId}`) + const messages: any[] = [] + let sendFailureIdle = false + try { + await new Promise((resolve, reject) => { + const timeout = setTimeout(() => { + ws.close() + reject(new Error(`Timed out waiting for send-failure runtime switch for session ${sessionId}`)) + }, 10_000) + + ws.onmessage = (event) => { + const msg = JSON.parse(event.data as string) + messages.push(msg) + + if (msg.type === 'connected') { + ws.send(JSON.stringify({ type: 'user_message', content: 'send failure active turn cleanup' })) + return + } + + if (msg.type === 'error' && msg.code !== 'CLI_NOT_RUNNING') { + clearTimeout(timeout) + ws.close() + reject(new Error(msg.message)) + return + } + + if (msg.type === 'status' && msg.state === 'idle' && !sendFailureIdle) { + if (!messages.some((item) => item.type === 'error' && item.code === 'CLI_NOT_RUNNING')) { + return + } + sendFailureIdle = true + ws.send(JSON.stringify({ + type: 'set_runtime_config', + providerId: provider.id, + modelId: 'send-failure-sonnet', + })) + return + } + + if (msg.type === 'status' && msg.state === 'idle' && sendFailureIdle && startCalls.length > 1) { + clearTimeout(timeout) + ws.close() + resolve() + } + } + + ws.onerror = () => { + clearTimeout(timeout) + reject(new Error(`WebSocket error for send-failure runtime switch session ${sessionId}`)) + } + }) + + expect(startCalls).toHaveLength(2) + expect(startCalls[1]).toMatchObject({ + sessionId, + options: { + providerId: provider.id, + model: 'send-failure-sonnet', + }, + }) + } finally { + ws.close() + conversationService.startSession = originalStartSession + conversationService.sendMessage = originalSendMessage + conversationService.stopSession(sessionId) + } + }, 20_000) + + it('should defer runtime model switches until the active turn completes', async () => { + await withMockStreamDelay(350, async () => { + const providerService = new ProviderService() + const providerA = await providerService.addProvider({ + presetId: 'custom', + name: 'Provider Active Runtime A', + apiKey: 'key-active-runtime-a', + baseUrl: 'http://127.0.0.1:1/anthropic', + apiFormat: 'anthropic', + models: { + main: 'active-a-main', + haiku: 'active-a-haiku', + sonnet: 'active-a-sonnet', + opus: 'active-a-opus', + }, + }) + const providerB = await providerService.addProvider({ + presetId: 'custom', + name: 'Provider Active Runtime B', + apiKey: 'key-active-runtime-b', + baseUrl: 'http://127.0.0.1:1/anthropic', + apiFormat: 'anthropic', + models: { + main: 'active-b-main', + haiku: 'active-b-haiku', + sonnet: 'active-b-sonnet', + opus: 'active-b-opus', + }, + }) + + const createRes = await fetch(`${baseUrl}/api/sessions`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ workDir: process.cwd() }), + }) + expect(createRes.status).toBe(201) + const { sessionId } = await createRes.json() as { sessionId: string } + + const originalStartSession = conversationService.startSession.bind(conversationService) + const startCalls: Array<{ + sessionId: string + options: { permissionMode?: string; model?: string; effort?: string; providerId?: string | null } | undefined + }> = [] + + conversationService.startSession = (async function patchedStartSession( + sid: string, + workDir: string, + sdkUrl: string, + options?: { permissionMode?: string; model?: string; effort?: string; thinking?: 'enabled' | 'adaptive' | 'disabled'; providerId?: string | null }, + ) { + startCalls.push({ sessionId: sid, options }) + return originalStartSession(sid, workDir, sdkUrl, options) + }) as typeof conversationService.startSession + + const ws = new WebSocket(`${wsUrl}/ws/${sessionId}`) + let switchTriggered = false + let turnComplete = false + try { + await new Promise((resolve, reject) => { + const timeout = setTimeout(() => { + ws.close() + reject(new Error(`Timed out waiting for active-turn runtime switch for session ${sessionId}`)) + }, 10_000) + + ws.onmessage = (event) => { + const msg = JSON.parse(event.data as string) + + if (msg.type === 'connected') { + ws.send(JSON.stringify({ + type: 'set_runtime_config', + providerId: providerA.id, + modelId: 'active-a-sonnet', + })) + ws.send(JSON.stringify({ type: 'user_message', content: 'active turn runtime switch' })) + return + } + + if (msg.type === 'error') { + clearTimeout(timeout) + ws.close() + reject(new Error(msg.message)) + return + } + + if ( + msg.type === 'content_delta' && + typeof msg.text === 'string' && + msg.text.includes('active turn runtime switch') && + !switchTriggered + ) { + switchTriggered = true + ws.send(JSON.stringify({ + type: 'set_runtime_config', + providerId: providerB.id, + modelId: 'active-b-opus', + })) + return + } + + if ( + msg.type === 'status' && + msg.state === 'idle' && + switchTriggered && + !turnComplete && + startCalls.length > 1 + ) { + clearTimeout(timeout) + ws.close() + reject(new Error('Runtime restarted before the active turn completed')) + return + } + + if (msg.type === 'message_complete' && switchTriggered && !turnComplete) { + turnComplete = true + expect(startCalls).toHaveLength(1) + return + } + + if (msg.type === 'status' && msg.state === 'idle' && turnComplete) { + clearTimeout(timeout) + ws.close() + resolve() + } + } + + ws.onerror = () => { + clearTimeout(timeout) + reject(new Error(`WebSocket error for active-turn runtime switch session ${sessionId}`)) + } + }) + + expect(startCalls).toHaveLength(2) + expect(startCalls[0]).toMatchObject({ + sessionId, + options: { + providerId: providerA.id, + model: 'active-a-sonnet', + }, + }) + expect(startCalls[1]).toMatchObject({ + sessionId, + options: { + providerId: providerB.id, + model: 'active-b-opus', + }, + }) + } finally { + ws.close() + conversationService.startSession = originalStartSession + conversationService.stopSession(sessionId) + } + }) + }, 20_000) + + it('should surface deferred runtime restart failures after the active turn completes', async () => { + await withMockStreamDelay(350, async () => { + const providerService = new ProviderService() + const providerA = await providerService.addProvider({ + presetId: 'custom', + name: 'Provider Deferred Failure A', + apiKey: 'key-deferred-failure-a', + baseUrl: 'http://127.0.0.1:1/anthropic', + apiFormat: 'anthropic', + models: { + main: 'deferred-failure-a-main', + haiku: 'deferred-failure-a-haiku', + sonnet: 'deferred-failure-a-sonnet', + opus: 'deferred-failure-a-opus', + }, + }) + const providerB = await providerService.addProvider({ + presetId: 'custom', + name: 'Provider Deferred Failure B', + apiKey: 'key-deferred-failure-b', + baseUrl: 'http://127.0.0.1:1/anthropic', + apiFormat: 'anthropic', + models: { + main: 'deferred-failure-b-main', + haiku: 'deferred-failure-b-haiku', + sonnet: 'deferred-failure-b-sonnet', + opus: 'deferred-failure-b-opus', + }, + }) + + const createRes = await fetch(`${baseUrl}/api/sessions`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ workDir: process.cwd() }), + }) + expect(createRes.status).toBe(201) + const { sessionId } = await createRes.json() as { sessionId: string } + + const originalStartSession = conversationService.startSession.bind(conversationService) + const startCalls: Array<{ + sessionId: string + options: { permissionMode?: string; model?: string; effort?: string; providerId?: string | null } | undefined + }> = [] + + conversationService.startSession = (async function patchedStartSession( + sid: string, + workDir: string, + sdkUrl: string, + options?: { permissionMode?: string; model?: string; effort?: string; thinking?: 'enabled' | 'adaptive' | 'disabled'; providerId?: string | null }, + ) { + startCalls.push({ sessionId: sid, options }) + if (startCalls.length > 1) { + throw new Error('deferred restart failed') + } + return originalStartSession(sid, workDir, sdkUrl, options) + }) as typeof conversationService.startSession + + const ws = new WebSocket(`${wsUrl}/ws/${sessionId}`) + let switchTriggered = false + let turnComplete = false + let restartError = false + try { + await new Promise((resolve, reject) => { + const timeout = setTimeout(() => { + ws.close() + reject(new Error(`Timed out waiting for deferred restart failure for session ${sessionId}`)) + }, 10_000) + + ws.onmessage = (event) => { + const msg = JSON.parse(event.data as string) + + if (msg.type === 'connected') { + ws.send(JSON.stringify({ + type: 'set_runtime_config', + providerId: providerA.id, + modelId: 'deferred-failure-a-sonnet', + })) + ws.send(JSON.stringify({ type: 'user_message', content: 'deferred runtime restart failure' })) + return + } + + if ( + msg.type === 'content_delta' && + typeof msg.text === 'string' && + msg.text.includes('deferred runtime restart failure') && + !switchTriggered + ) { + switchTriggered = true + ws.send(JSON.stringify({ + type: 'set_runtime_config', + providerId: providerB.id, + modelId: 'deferred-failure-b-opus', + })) + return + } + + if (msg.type === 'message_complete' && switchTriggered && !turnComplete) { + turnComplete = true + expect(startCalls).toHaveLength(1) + return + } + + if (msg.type === 'error') { + if (!turnComplete) { + clearTimeout(timeout) + ws.close() + reject(new Error(`Deferred restart failed before turn completion: ${msg.message}`)) + return + } + restartError = msg.code === 'CLI_RESTART_FAILED' && + typeof msg.message === 'string' && + msg.message.includes('deferred restart failed') + return + } + + if (msg.type === 'status' && msg.state === 'idle' && restartError) { + clearTimeout(timeout) + ws.close() + resolve() + } + } + + ws.onerror = () => { + clearTimeout(timeout) + reject(new Error(`WebSocket error for deferred restart failure session ${sessionId}`)) + } + }) + + expect(switchTriggered).toBe(true) + expect(turnComplete).toBe(true) + expect(restartError).toBe(true) + expect(startCalls).toHaveLength(2) + expect(startCalls[1]).toMatchObject({ + sessionId, + options: { + providerId: providerB.id, + model: 'deferred-failure-b-opus', + }, + }) + } finally { + ws.close() + conversationService.startSession = originalStartSession + conversationService.stopSession(sessionId) + } + }) + }, 20_000) + it('should keep the session idle in the UI while restarting for a bypass permission switch', async () => { await fetch(`${baseUrl}/api/permissions/mode`, { method: 'PUT', @@ -2721,7 +3132,7 @@ describe('WebSocket Chat Integration', () => { } }, 20_000) - it('should restart when switching from bypass permissions back to default', async () => { + it('should switch from bypass permissions back to default without restarting', async () => { const createRes = await fetch(`${baseUrl}/api/sessions`, { method: 'POST', headers: { 'Content-Type': 'application/json' }, @@ -2789,24 +3200,13 @@ describe('WebSocket Chat Integration', () => { mode: 'default', })) - await waitUntil( - async () => messages.slice(switchStartIndex).some((msg) => msg.type === 'status' && msg.state === 'idle'), - `bypass-to-default permission switch completion for ${sessionId}`, - ) - - expect(startCalls).toHaveLength(2) - expect(startCalls[1]).toMatchObject({ - sessionId, - options: { - permissionMode: 'default', - }, - }) await waitUntil(async () => { const res = await fetch(`${baseUrl}/api/sessions/${sessionId}/inspection?includeContext=0`) if (!res.ok) return false const body = await res.json() as { status?: { permissionMode?: string } } return body.status?.permissionMode === 'default' }, `persisted bypass-to-default permission switch for ${sessionId}`) + expect(startCalls).toHaveLength(1) expect(messages.slice(switchStartIndex).some((msg) => msg.type === 'error')).toBe(false) } finally { ws.close() diff --git a/src/server/ws/handler.ts b/src/server/ws/handler.ts index 2ea95be7..27c64357 100644 --- a/src/server/ws/handler.ts +++ b/src/server/ws/handler.ts @@ -76,11 +76,19 @@ const sessionTitleState = new Map() -const runtimeOverrides = new Map() +} + +type ActiveUserTurnState = { + messageSent: boolean +} + +const runtimeOverrides = new Map() +const activeUserTurns = new Map() +const deferredRuntimeRestarts = new Map() const runtimeTransitionPromises = new Map>() const sessionStartupPromises = new Map>() @@ -306,8 +314,14 @@ async function handleUserMessage( // Send thinking status sendMessage(ws, { type: 'status', state: 'thinking', verb: 'Thinking' }) + const activeTurn: ActiveUserTurnState = { messageSent: false } + activeUserTurns.set(sessionId, activeTurn) + const initialRuntimeTransition = await waitForRuntimeTransitionBeforeUserTurn(ws, sessionId) - if (!initialRuntimeTransition.ok) return + if (!initialRuntimeTransition.ok) { + clearActiveUserTurn(sessionId, activeTurn) + return + } if (initialRuntimeTransition.waited) { sendMessage(ws, { type: 'status', state: 'thinking', verb: 'Thinking' }) } @@ -357,6 +371,7 @@ async function handleUserMessage( err instanceof ConversationStartupError ? err.retryable : false, }) sendMessage(ws, { type: 'status', state: 'idle' }) + clearActiveUserTurn(sessionId, activeTurn) return } @@ -366,6 +381,7 @@ async function handleUserMessage( sendMessage(ws, { type: 'status', state: 'thinking', verb: 'Thinking' }) } } else { + clearActiveUserTurn(sessionId, activeTurn) return } @@ -387,6 +403,7 @@ async function handleUserMessage( return shouldForwardCurrentTurnLocalCommand(cliMsg) }, }) + const removeActiveTurnOutputCallback = bindActiveUserTurnCompletion(ws, sessionId, activeTurn) const sent = await conversationService.sendMessage( sessionId, @@ -394,6 +411,8 @@ async function handleUserMessage( message.attachments ) if (!sent) { + removeActiveTurnOutputCallback() + clearActiveUserTurn(sessionId, activeTurn) removeTitleOutputCallback?.() discardActiveTitleTurn(sessionId, titleTurnNumber) sendMessage(ws, { @@ -406,6 +425,57 @@ async function handleUserMessage( } userMessageSent = true + activeTurn.messageSent = true +} + +function clearActiveUserTurn(sessionId: string, activeTurn: ActiveUserTurnState): void { + if (activeUserTurns.get(sessionId) === activeTurn) { + activeUserTurns.delete(sessionId) + } +} + +function bindActiveUserTurnCompletion( + ws: ServerWebSocket, + sessionId: string, + activeTurn: ActiveUserTurnState, +): () => void { + const callback = (cliMsg: any) => { + if (!activeTurn.messageSent || cliMsg?.type !== 'result') return + + conversationService.removeOutputCallback(sessionId, callback) + clearActiveUserTurn(sessionId, activeTurn) + applyDeferredRuntimeRestartAfterActiveTurn(ws, sessionId) + } + + conversationService.onOutput(sessionId, callback) + return () => conversationService.removeOutputCallback(sessionId, callback) +} + +function shouldDeferRuntimeRestartForActiveTurn(sessionId: string): boolean { + return activeUserTurns.get(sessionId)?.messageSent === true +} + +function applyDeferredRuntimeRestartAfterActiveTurn( + ws: ServerWebSocket, + sessionId: string, +): void { + const deferred = deferredRuntimeRestarts.get(sessionId) + if (!deferred) return + + deferredRuntimeRestarts.delete(sessionId) + void enqueueRuntimeTransition(sessionId, async () => { + const currentOverride = runtimeOverrides.get(sessionId) + if ( + !currentOverride || + currentOverride.providerId !== deferred.providerId || + currentOverride.modelId !== deferred.modelId || + currentOverride.effort !== deferred.effort || + !conversationService.hasSession(sessionId) + ) { + return + } + await restartSessionWithRuntimeConfig(ws, sessionId) + }) } async function handleDesktopClearCommand( @@ -621,6 +691,12 @@ async function handleSetRuntimeConfig( (runtimeOverrideVersions.get(sessionId) ?? 0) + 1, ) + if (shouldDeferRuntimeRestartForActiveTurn(sessionId)) { + deferredRuntimeRestarts.set(sessionId, nextOverride) + await persistSessionRuntimeConfig(sessionId, nextOverride) + return + } + if (conversationService.hasSession(sessionId)) { await enqueueRuntimeTransition(sessionId, async () => { await persistSessionRuntimeConfig(sessionId, nextOverride) @@ -1081,6 +1157,8 @@ function cleanupSessionRuntimeState(sessionId: string) { sessionSlashCommands.delete(sessionId) sessionTitleState.delete(sessionId) runtimeOverrides.delete(sessionId) + activeUserTurns.delete(sessionId) + deferredRuntimeRestarts.delete(sessionId) runtimeTransitionPromises.delete(sessionId) sessionStartupPromises.delete(sessionId) lastResolvedStartupWorkDirs.delete(sessionId)