Skip to content

Commit f1418c3

Browse files
committed
fix(core): [CRITICAL] Fix silent conversation memory wipe on user interruption and compaction
1 parent 92112b2 commit f1418c3

3 files changed

Lines changed: 155 additions & 7 deletions

File tree

cli/src/hooks/helpers/__tests__/send-message.test.ts

Lines changed: 141 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ const {
3737
finalizeQueueState,
3838
resetEarlyReturnState,
3939
} = await import('../send-message')
40+
const { createRunConfig } = await import('../../../utils/create-run-config')
4041
const { createBatchedMessageUpdater } =
4142
await import('../../../utils/message-updater')
4243
import { createPaymentRequiredError } from '@codebuff/sdk'
@@ -1858,4 +1859,144 @@ describe('freebuff gate errors', () => {
18581859
// (which would set a userError from the message).
18591860
expect(messages[0].userError).toBeUndefined()
18601861
})
1862+
1863+
describe('session state preservation on user abort and error / session expiration', () => {
1864+
test('user abort (Esc): follow-up message inherits active snapshot with message history and tool calls', () => {
1865+
const previousRunStateRef = { current: null as RunState | null }
1866+
let committedStoreRunState: RunState | null = null
1867+
let currentChatDir = '/chat-1'
1868+
1869+
const syncRunState = (state: RunState) => {
1870+
if (currentChatDir !== '/chat-1') return
1871+
previousRunStateRef.current = state
1872+
committedStoreRunState = state
1873+
}
1874+
1875+
const activeSnapshot: RunState = {
1876+
sessionState: {
1877+
fileContext: { projectRoot: '/project', files: {} } as any,
1878+
mainAgentState: {
1879+
agentId: 'agent-1',
1880+
agentType: 'base2',
1881+
messageHistory: [
1882+
{ role: 'user', content: [{ type: 'text', text: 'implement authentication' }] },
1883+
{ role: 'assistant', content: [{ type: 'text', text: 'reading project files' }] },
1884+
{
1885+
role: 'tool',
1886+
toolName: 'read_files',
1887+
content: [{ type: 'text', value: { files: ['src/auth.ts'] } }],
1888+
} as any,
1889+
],
1890+
} as any,
1891+
},
1892+
traceSessionId: 'trace-1',
1893+
}
1894+
1895+
// User hits Esc mid-stream: registerActiveRun abort callback fires
1896+
syncRunState(activeSnapshot)
1897+
1898+
// User immediately sends follow-up prompt
1899+
const runConfigB = createRunConfig({
1900+
logger: { debug: () => {}, warn: () => {}, error: () => {}, info: () => {} } as any,
1901+
agent: 'base2',
1902+
prompt: 'also verify tests',
1903+
content: undefined,
1904+
previousRunState: previousRunStateRef.current,
1905+
agentDefinitions: [],
1906+
eventHandlerState: {} as any,
1907+
})
1908+
1909+
// Verify that previousRun carries full history rather than empty []
1910+
expect(runConfigB.previousRun).toBe(activeSnapshot)
1911+
expect(runConfigB.previousRun?.sessionState?.mainAgentState.messageHistory).toHaveLength(3)
1912+
expect(
1913+
runConfigB.previousRun?.sessionState?.mainAgentState.messageHistory[0].content[0].text,
1914+
).toBe('implement authentication')
1915+
expect(
1916+
runConfigB.previousRun?.sessionState?.mainAgentState.messageHistory[2].toolName,
1917+
).toBe('read_files')
1918+
expect(committedStoreRunState).toBe(activeSnapshot)
1919+
})
1920+
1921+
test('session expiration / error: sending "continue" inherits expired session snapshot', () => {
1922+
const previousRunStateRef = { current: null as RunState | null }
1923+
let committedStoreRunState: RunState | null = null
1924+
let currentChatDir = '/chat-1'
1925+
1926+
const syncRunState = (state: RunState) => {
1927+
if (currentChatDir !== '/chat-1') return
1928+
previousRunStateRef.current = state
1929+
committedStoreRunState = state
1930+
}
1931+
1932+
const errorSnapshot: RunState = {
1933+
sessionState: {
1934+
fileContext: { projectRoot: '/project', files: {} } as any,
1935+
mainAgentState: {
1936+
agentId: 'agent-1',
1937+
agentType: 'base2',
1938+
messageHistory: [
1939+
{ role: 'user', content: [{ type: 'text', text: 'turn 1 prompt' }] },
1940+
{ role: 'assistant', content: [{ type: 'text', text: 'turn 1 response' }] },
1941+
{ role: 'user', content: [{ type: 'text', text: 'turn 2 prompt' }] },
1942+
],
1943+
} as any,
1944+
},
1945+
traceSessionId: 'trace-2',
1946+
output: {
1947+
type: 'error',
1948+
message: 'Your free session ended',
1949+
error: 'session_expired',
1950+
} as any,
1951+
}
1952+
1953+
// Session expiration error caught in useSendMessage catch block
1954+
syncRunState(errorSnapshot)
1955+
1956+
// User sends "continue" after session ended banner
1957+
const runConfigB = createRunConfig({
1958+
logger: { debug: () => {}, warn: () => {}, error: () => {}, info: () => {} } as any,
1959+
agent: 'base2',
1960+
prompt: 'continue',
1961+
content: undefined,
1962+
previousRunState: previousRunStateRef.current,
1963+
agentDefinitions: [],
1964+
eventHandlerState: {} as any,
1965+
})
1966+
1967+
expect(runConfigB.previousRun).toBe(errorSnapshot)
1968+
expect(runConfigB.previousRun?.sessionState?.mainAgentState.messageHistory).toHaveLength(3)
1969+
expect(
1970+
runConfigB.previousRun?.sessionState?.mainAgentState.messageHistory[2].content[0].text,
1971+
).toBe('turn 2 prompt')
1972+
expect(committedStoreRunState).toBe(errorSnapshot)
1973+
})
1974+
1975+
test('chat switch isolation: syncRunState is a no-op when active chat directory has changed', () => {
1976+
const previousRunStateRef = { current: null as RunState | null }
1977+
let committedStoreRunState: RunState | null = null
1978+
let currentChatDir = '/chat-2' // User switched to chat 2
1979+
1980+
const syncRunState = (state: RunState) => {
1981+
if (currentChatDir !== '/chat-1') return // Stale run from chat 1
1982+
previousRunStateRef.current = state
1983+
committedStoreRunState = state
1984+
}
1985+
1986+
const staleSnapshot: RunState = {
1987+
sessionState: {
1988+
mainAgentState: {
1989+
messageHistory: [{ role: 'user', content: [{ type: 'text', text: 'chat 1 prompt' }] }],
1990+
} as any,
1991+
} as any,
1992+
traceSessionId: 'trace-stale',
1993+
}
1994+
1995+
syncRunState(staleSnapshot)
1996+
1997+
// Must NOT leak state into chat 2
1998+
expect(previousRunStateRef.current).toBeNull()
1999+
expect(committedStoreRunState).toBeNull()
2000+
})
2001+
})
18612002
})

cli/src/hooks/helpers/send-message.ts

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -305,11 +305,8 @@ export const setupStreamingContext = (params: {
305305
abortController.signal.addEventListener('abort', () => {
306306
// Abort means the user stopped streaming; update UI with an interruption notice.
307307
// Release the chain lock immediately so new messages can be sent directly instead
308-
// of being queued. The minor trade-off is that if the user sends a new message
309-
// before client.run() resolves, it may use stale previousRunStateRef. This is
310-
// acceptable because: (1) the user explicitly cancelled, and (2) client.run()
311-
// will update previousRunStateRef when it eventually resolves, so subsequent
312-
// runs will have the full state.
308+
// of being queued. registerActiveRun updates previousRunStateRef synchronously
309+
// with the latest snapshot so immediate follow-ups retain preserved context.
313310
streamRefs.setters.setWasAbortedByUser(true)
314311
setIsRetrying(false)
315312
timerController.stop('aborted')

cli/src/hooks/use-send-message.ts

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -315,6 +315,12 @@ export const useSendMessage = ({
315315
clearActiveRun(runOwnerId)
316316
}
317317

318+
const syncRunState = (state: RunState) => {
319+
if (!runChatIsCurrent()) return
320+
previousRunStateRef.current = state
321+
setRunState(state)
322+
}
323+
318324
registerActiveRun(runOwnerId, (reason) => {
319325
if (abortController.signal.aborted) return
320326

@@ -330,6 +336,10 @@ export const useSendMessage = ({
330336
if (isProcessingQueueRef) isProcessingQueueRef.current = false
331337
}
332338

339+
// Keep in-memory previousRunStateRef fresh so immediate follow-up
340+
// messages carry the latest snapshot even before client.run settles.
341+
syncRunState(latestRunStateSnapshot)
342+
333343
// Capture the old chat's array now. Context-changing callers reset the
334344
// store immediately after stopActiveRun returns.
335345
scheduleCheckpointSave(
@@ -644,8 +654,7 @@ export const useSendMessage = ({
644654
// same chat, so the interrupted turn is still saved as before.)
645655
if (runChatIsCurrent()) {
646656
// Finalize: persist state and mark complete
647-
previousRunStateRef.current = runState
648-
setRunState(runState)
657+
syncRunState(runState)
649658
setIsRetrying(false)
650659

651660
// Drop any queued/in-flight async checkpoint first so a stale write
@@ -697,6 +706,7 @@ export const useSendMessage = ({
697706
// first so a stale write can't clobber this one. Skipped after a
698707
// mid-run chat switch — the store's messages belong to the new chat.
699708
if (runChatIsCurrent()) {
709+
syncRunState(latestRunStateSnapshot)
700710
await settleCheckpointSave()
701711
saveChatState(
702712
latestRunStateSnapshot,

0 commit comments

Comments
 (0)