diff --git a/cli/src/hooks/helpers/__tests__/send-message.test.ts b/cli/src/hooks/helpers/__tests__/send-message.test.ts index 7e23893249..8899c19b6f 100644 --- a/cli/src/hooks/helpers/__tests__/send-message.test.ts +++ b/cli/src/hooks/helpers/__tests__/send-message.test.ts @@ -37,6 +37,7 @@ const { finalizeQueueState, resetEarlyReturnState, } = await import('../send-message') +const { createRunConfig } = await import('../../../utils/create-run-config') const { createBatchedMessageUpdater } = await import('../../../utils/message-updater') import { createPaymentRequiredError } from '@codebuff/sdk' @@ -1858,4 +1859,144 @@ describe('freebuff gate errors', () => { // (which would set a userError from the message). expect(messages[0].userError).toBeUndefined() }) + + describe('session state preservation on user abort and error / session expiration', () => { + test('user abort (Esc): follow-up message inherits active snapshot with message history and tool calls', () => { + const previousRunStateRef = { current: null as RunState | null } + let committedStoreRunState: RunState | null = null + let currentChatDir = '/chat-1' + + const syncRunState = (state: RunState) => { + if (currentChatDir !== '/chat-1') return + previousRunStateRef.current = state + committedStoreRunState = state + } + + const activeSnapshot: RunState = { + sessionState: { + fileContext: { projectRoot: '/project', files: {} } as any, + mainAgentState: { + agentId: 'agent-1', + agentType: 'base2', + messageHistory: [ + { role: 'user', content: [{ type: 'text', text: 'implement authentication' }] }, + { role: 'assistant', content: [{ type: 'text', text: 'reading project files' }] }, + { + role: 'tool', + toolName: 'read_files', + content: [{ type: 'text', value: { files: ['src/auth.ts'] } }], + } as any, + ], + } as any, + }, + traceSessionId: 'trace-1', + } + + // User hits Esc mid-stream: registerActiveRun abort callback fires + syncRunState(activeSnapshot) + + // User immediately sends follow-up prompt + const runConfigB = createRunConfig({ + logger: { debug: () => {}, warn: () => {}, error: () => {}, info: () => {} } as any, + agent: 'base2', + prompt: 'also verify tests', + content: undefined, + previousRunState: previousRunStateRef.current, + agentDefinitions: [], + eventHandlerState: {} as any, + }) + + // Verify that previousRun carries full history rather than empty [] + expect(runConfigB.previousRun).toBe(activeSnapshot) + expect(runConfigB.previousRun?.sessionState?.mainAgentState.messageHistory).toHaveLength(3) + expect( + runConfigB.previousRun?.sessionState?.mainAgentState.messageHistory[0].content[0].text, + ).toBe('implement authentication') + expect( + runConfigB.previousRun?.sessionState?.mainAgentState.messageHistory[2].toolName, + ).toBe('read_files') + expect(committedStoreRunState).toBe(activeSnapshot) + }) + + test('session expiration / error: sending "continue" inherits expired session snapshot', () => { + const previousRunStateRef = { current: null as RunState | null } + let committedStoreRunState: RunState | null = null + let currentChatDir = '/chat-1' + + const syncRunState = (state: RunState) => { + if (currentChatDir !== '/chat-1') return + previousRunStateRef.current = state + committedStoreRunState = state + } + + const errorSnapshot: RunState = { + sessionState: { + fileContext: { projectRoot: '/project', files: {} } as any, + mainAgentState: { + agentId: 'agent-1', + agentType: 'base2', + messageHistory: [ + { role: 'user', content: [{ type: 'text', text: 'turn 1 prompt' }] }, + { role: 'assistant', content: [{ type: 'text', text: 'turn 1 response' }] }, + { role: 'user', content: [{ type: 'text', text: 'turn 2 prompt' }] }, + ], + } as any, + }, + traceSessionId: 'trace-2', + output: { + type: 'error', + message: 'Your free session ended', + error: 'session_expired', + } as any, + } + + // Session expiration error caught in useSendMessage catch block + syncRunState(errorSnapshot) + + // User sends "continue" after session ended banner + const runConfigB = createRunConfig({ + logger: { debug: () => {}, warn: () => {}, error: () => {}, info: () => {} } as any, + agent: 'base2', + prompt: 'continue', + content: undefined, + previousRunState: previousRunStateRef.current, + agentDefinitions: [], + eventHandlerState: {} as any, + }) + + expect(runConfigB.previousRun).toBe(errorSnapshot) + expect(runConfigB.previousRun?.sessionState?.mainAgentState.messageHistory).toHaveLength(3) + expect( + runConfigB.previousRun?.sessionState?.mainAgentState.messageHistory[2].content[0].text, + ).toBe('turn 2 prompt') + expect(committedStoreRunState).toBe(errorSnapshot) + }) + + test('chat switch isolation: syncRunState is a no-op when active chat directory has changed', () => { + const previousRunStateRef = { current: null as RunState | null } + let committedStoreRunState: RunState | null = null + let currentChatDir = '/chat-2' // User switched to chat 2 + + const syncRunState = (state: RunState) => { + if (currentChatDir !== '/chat-1') return // Stale run from chat 1 + previousRunStateRef.current = state + committedStoreRunState = state + } + + const staleSnapshot: RunState = { + sessionState: { + mainAgentState: { + messageHistory: [{ role: 'user', content: [{ type: 'text', text: 'chat 1 prompt' }] }], + } as any, + } as any, + traceSessionId: 'trace-stale', + } + + syncRunState(staleSnapshot) + + // Must NOT leak state into chat 2 + expect(previousRunStateRef.current).toBeNull() + expect(committedStoreRunState).toBeNull() + }) + }) }) diff --git a/cli/src/hooks/helpers/send-message.ts b/cli/src/hooks/helpers/send-message.ts index f5432dbc97..8602fc8fa8 100644 --- a/cli/src/hooks/helpers/send-message.ts +++ b/cli/src/hooks/helpers/send-message.ts @@ -305,11 +305,8 @@ export const setupStreamingContext = (params: { abortController.signal.addEventListener('abort', () => { // Abort means the user stopped streaming; update UI with an interruption notice. // Release the chain lock immediately so new messages can be sent directly instead - // of being queued. The minor trade-off is that if the user sends a new message - // before client.run() resolves, it may use stale previousRunStateRef. This is - // acceptable because: (1) the user explicitly cancelled, and (2) client.run() - // will update previousRunStateRef when it eventually resolves, so subsequent - // runs will have the full state. + // of being queued. registerActiveRun updates previousRunStateRef synchronously + // with the latest snapshot so immediate follow-ups retain preserved context. streamRefs.setters.setWasAbortedByUser(true) setIsRetrying(false) timerController.stop('aborted') diff --git a/cli/src/hooks/use-send-message.ts b/cli/src/hooks/use-send-message.ts index 39065317ee..7433f0fd89 100644 --- a/cli/src/hooks/use-send-message.ts +++ b/cli/src/hooks/use-send-message.ts @@ -315,6 +315,12 @@ export const useSendMessage = ({ clearActiveRun(runOwnerId) } + const syncRunState = (state: RunState) => { + if (!runChatIsCurrent()) return + previousRunStateRef.current = state + setRunState(state) + } + registerActiveRun(runOwnerId, (reason) => { if (abortController.signal.aborted) return @@ -330,6 +336,10 @@ export const useSendMessage = ({ if (isProcessingQueueRef) isProcessingQueueRef.current = false } + // Keep in-memory previousRunStateRef fresh so immediate follow-up + // messages carry the latest snapshot even before client.run settles. + syncRunState(latestRunStateSnapshot) + // Capture the old chat's array now. Context-changing callers reset the // store immediately after stopActiveRun returns. scheduleCheckpointSave( @@ -644,8 +654,7 @@ export const useSendMessage = ({ // same chat, so the interrupted turn is still saved as before.) if (runChatIsCurrent()) { // Finalize: persist state and mark complete - previousRunStateRef.current = runState - setRunState(runState) + syncRunState(runState) setIsRetrying(false) // Drop any queued/in-flight async checkpoint first so a stale write @@ -697,6 +706,7 @@ export const useSendMessage = ({ // first so a stale write can't clobber this one. Skipped after a // mid-run chat switch — the store's messages belong to the new chat. if (runChatIsCurrent()) { + syncRunState(latestRunStateSnapshot) await settleCheckpointSave() saveChatState( latestRunStateSnapshot,