fix(core): preserve session state on abort and error / session expiration (replaces #1055)

cli/src/hooks/__tests__/use-send-message.test.tsx

@@ -186,3 +186,127 @@ describe('useSendMessage continuation state', () => {
     }
   })
 })
+
+// Regression tests for the syncRunState fix (#1054).
+//
+// Both scenarios previously caused a follow-up message to lose all conversation
+// context because previousRunStateRef was not updated before client.run()
+// settled. The tests below drive the real hook and assert that the snapshot
+// passed to onStateSnapshot() is the one carried into the next run's
+// previousRun, verifying the actual wiring — not a reimplemented proxy.
+describe('useSendMessage syncRunState regression (#1054)', () => {
+  test('abort path: latestRunStateSnapshot is committed to previousRunStateRef before client.run() settles', async () => {
+    // This exercises use-send-message.ts line ~355:
+    //   syncRunState(latestRunStateSnapshot)  ← inside registerActiveRun callback
+    // Calling stopActiveRun fires the real abort callback synchronously, before
+    // client.run()'s promise resolves. The follow-up run must receive the
+    // snapshot that was live at abort time, not an empty/null state.
+    const { setup, root } = await mountHost()
+    const runs: Promise<void>[] = []
+
+    try {
+      runs.push(
+        sendMessageFromHost!({ content: 'first message', agentMode: 'DEFAULT' }),
+      )
+      await waitFor('first run registered', () => runCalls.length === 1)
+
+      const liveSnapshot = makeRunState('mid-stream')
+      // Simulate a partial streaming state update arriving before Esc.
+      runCalls[0].runConfig.onStateSnapshot(liveSnapshot)
+      // User presses Esc — fires the real registerActiveRun abort callback.
+      stopActiveRun('user-interrupt')
+
+      runs.push(
+        sendMessageFromHost!({ content: 'follow-up after abort', agentMode: 'DEFAULT' }),
+      )
+      await waitFor('second run registered', () => runCalls.length === 2)
+
+      // The real hook must have assigned liveSnapshot into previousRunStateRef
+      // via syncRunState before we got here — not the blank sentinel.
+      expect(runCalls[1].runConfig.previousRun).toBe(liveSnapshot)
+    } finally {
+      settlePendingRuns()
+      await Promise.all(runs)
+      flushSync(() => root.unmount())
+      setup.renderer.destroy()
+    }
+  })
+
+  test('error path: latestRunStateSnapshot is committed to previousRunStateRef when client.run() rejects', async () => {
+    // This exercises use-send-message.ts line ~752:
+    //   syncRunState(latestRunStateSnapshot)  ← inside catch (error) block
+    // When client.run() throws (network error, session expiry, gate error),
+    // the catch block must persist the last received snapshot so the user's
+    // conversation context survives the failure.
+    const { setup, root } = await mountHost()
+    const runs: Promise<void>[] = []
+
+    try {
+      runs.push(
+        sendMessageFromHost!({ content: 'first message', agentMode: 'DEFAULT' }),
+      )
+      await waitFor('first run registered', () => runCalls.length === 1)
+
+      const lastSnapshot = makeRunState('before-error')
+      // Simulate a snapshot arriving mid-stream, then a network / gate error.
+      runCalls[0].runConfig.onStateSnapshot(lastSnapshot)
+      runCalls[0].reject(new Error('session expired'))
+      await runs[0]
+
+      runs.push(
+        sendMessageFromHost!({ content: 'continue', agentMode: 'DEFAULT' }),
+      )
+      await waitFor('second run registered', () => runCalls.length === 2)
+
+      // The catch block must have called syncRunState(latestRunStateSnapshot),
+      // making lastSnapshot available to the next run.
+      expect(runCalls[1].runConfig.previousRun).toBe(lastSnapshot)
+    } finally {
+      settlePendingRuns()
+      await Promise.all(runs)
+      flushSync(() => root.unmount())
+      setup.renderer.destroy()
+    }
+  })
+  test('abort path: falls back to prior run state when client.run() is aborted before any snapshot arrives', async () => {
+    // latestRunStateSnapshot is initialized from previousRunStateRef.current (line ~313).
+    // If the user presses Esc immediately — before the SDK emits any onStateSnapshot —
+    // syncRunState is called with that initial value, which is the prior completed run's
+    // state. This directly answers the question: "is latestRunStateSnapshot guaranteed
+    // to be populated at the abort callsite?" Yes — it is never null.
+    const { setup, root } = await mountHost()
+    const runs: Promise<void>[] = []
+
+    try {
+      // Run 1: complete successfully so there IS a known prior state.
+      runs.push(
+        sendMessageFromHost!({ content: 'first message', agentMode: 'DEFAULT' }),
+      )
+      await waitFor('first run registered', () => runCalls.length === 1)
+      const priorState = makeRunState('completed')
+      runCalls[0].resolve(priorState)
+      await runs[0]
+
+      // Run 2: abort immediately, before any onStateSnapshot arrives.
+      runs.push(
+        sendMessageFromHost!({ content: 'second message', agentMode: 'DEFAULT' }),
+      )
+      await waitFor('second run registered', () => runCalls.length === 2)
+      // Deliberately NO onStateSnapshot call — simulates Esc before any streaming progress.
+      stopActiveRun('user-interrupt')
+
+      // Run 3 should carry priorState (from run 1), not a blank/null sentinel.
+      runs.push(
+        sendMessageFromHost!({ content: 'follow-up', agentMode: 'DEFAULT' }),
+      )
+      await waitFor('third run registered', () => runCalls.length === 3)
+
+      expect(runCalls[2].runConfig.previousRun).toBe(priorState)
+    } finally {
+      settlePendingRuns()
+      await Promise.all(runs)
+      flushSync(() => root.unmount())
+      setup.renderer.destroy()
+    }
+  })
+})

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

@@ -305,7 +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.
+    // 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')

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

@@ -329,6 +329,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
 
@@ -344,6 +350,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(
@@ -687,8 +697,7 @@ export const useSendMessage = ({
         // same chat, so the interrupted turn is still saved as before.)
         if (!abortController.signal.aborted && 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
@@ -740,6 +749,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,