@@ -37,6 +37,7 @@ const {
3737 finalizeQueueState,
3838 resetEarlyReturnState,
3939} = await import ( '../send-message' )
40+ const { createRunConfig } = await import ( '../../../utils/create-run-config' )
4041const { createBatchedMessageUpdater } =
4142 await import ( '../../../utils/message-updater' )
4243import { 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} )
0 commit comments