Fix blank-session restarts after interrupt or failed run

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

@@ -1859,3 +1859,144 @@ describe('freebuff gate errors', () => {
     expect(messages[0].userError).toBeUndefined()
   })
 })
+
+describe('onAbort checkpoint callback (freebuff #1054)', () => {
+  test('runs synchronously at the top of the abort listener, before the input lock is released', () => {
+    let messages = createBaseMessages()
+    const streamRefs = createStreamController()
+    const timerController = createMockTimerController()
+    // Cast the initializers: TS narrows `let` vars from literal
+    // initializers and the abort listener's closure assignments never
+    // reset that narrowing, which would mis-type the assertions below.
+    let streamStatus = 'streaming' as StreamStatus
+    let statusSeenInOnAbort = null as StreamStatus | null
+    let chainInProgress = true
+    let onAbortCalls = 0
+
+    const { abortController } = setupStreamingContext({
+      aiMessageId: 'ai-1',
+      timerController,
+      setMessages: (fn: any) => {
+        messages = fn(messages)
+      },
+      streamRefs,
+      onAbort: () => {
+        onAbortCalls++
+        // The listener below swaps this to 'idle'; seeing 'streaming'
+        // proves onAbort ran BEFORE the lock release / UI reset.
+        statusSeenInOnAbort = streamStatus
+      },
+      setStreamStatus: (status: StreamStatus) => {
+        streamStatus = status
+      },
+      setCanProcessQueue: () => {
+        chainInProgress = false
+      },
+      updateChainInProgress: () => {},
+      setIsRetrying: () => {},
+      setStreamingAgents: () => {},
+    })
+
+    abortController.abort()
+
+    expect(onAbortCalls).toBe(1)
+    expect(statusSeenInOnAbort).toBe('streaming')
+    expect(streamStatus).toBe('idle')
+  })
+
+  test('a throwing onAbort still runs the interruption cleanup', () => {
+    let messages = createBaseMessages()
+    const streamRefs = createStreamController()
+    const timerController = createMockTimerController()
+    let chainInProgress = true
+
+    const { abortController } = setupStreamingContext({
+      aiMessageId: 'ai-1',
+      timerController,
+      setMessages: (fn: any) => {
+        messages = fn(messages)
+      },
+      streamRefs,
+      onAbort: () => {
+        throw new Error('checkpoint boom')
+      },
+      setStreamStatus: () => {},
+      setCanProcessQueue: () => {
+        chainInProgress = false
+      },
+      updateChainInProgress: () => {},
+      setIsRetrying: () => {},
+      setStreamingAgents: () => {},
+    })
+
+    expect(() => abortController.abort()).not.toThrow()
+    expect(chainInProgress).toBe(false)
+    expect(streamRefs.state.wasAbortedByUser).toBe(true)
+  })
+
+  test('abort leaves the full message history in the continuation snapshot', () => {
+    // Mirrors the hook wiring: onAbort stores the latest in-flight snapshot
+    // into the continuation ref used by the next sendMessage. This is the
+    // exact path that used to hand the SDK an empty previousRun after an
+    // interrupt, resetting the conversation to a blank session.
+    const previousRunStateRef = { current: null as RunState | null }
+    const snapshot = {
+      traceSessionId: 'trace-1',
+      sessionState: {
+        mainAgentState: {
+          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,
+      output: {
+        type: 'error',
+        message: 'Session ended before this response completed.',
+      },
+    } as unknown as RunState
+
+    const { abortController } = setupStreamingContext({
+      aiMessageId: 'ai-1',
+      timerController: createMockTimerController(),
+      setMessages: (fn: any) => {},
+      streamRefs: createStreamController(),
+      onAbort: () => {
+        previousRunStateRef.current = snapshot
+      },
+      setStreamStatus: () => {},
+      setCanProcessQueue: () => {},
+      updateChainInProgress: () => {},
+      setIsRetrying: () => {},
+      setStreamingAgents: () => {},
+    })
+
+    abortController.abort()
+
+    // The continuation snapshot now carries the full history instead of
+    // null (which would make the SDK build a blank session and re-explore
+    // from scratch on the next prompt). The hook-level test in
+    // cli/src/hooks/__tests__/use-send-message.test.tsx asserts the same
+    // through the real createRunConfig/client.run path.
+    expect(previousRunStateRef.current).toBe(snapshot)
+    expect(
+      (previousRunStateRef.current!.sessionState as any).mainAgentState
+        .messageHistory,
+    ).toHaveLength(3)
+    expect(
+      (previousRunStateRef.current!.sessionState as any).mainAgentState
+        .messageHistory[2].toolName,
+    ).toBe('read_files')
+  })
+})

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

@@ -272,6 +272,11 @@ export const setupStreamingContext = (params: {
   setMessages: (updater: (messages: ChatMessage[]) => ChatMessage[]) => void
   streamRefs: StreamController
   abortController?: AbortController
+  /** Invoked synchronously at the top of the abort listener, before the
+   *  input lock is released. Lets the run owner checkpoint its latest SDK
+   *  snapshot so a follow-up message sent immediately after an interrupt
+   *  resumes from fresh state instead of stale or null continuation state. */
+  onAbort?: () => void
   setStreamStatus: (status: StreamStatus) => void
   setCanProcessQueue: (can: boolean) => void
   isQueuePausedRef?: MutableRefObject<boolean>
@@ -303,9 +308,19 @@ export const setupStreamingContext = (params: {
   const abortController = params.abortController ?? new AbortController()
 
   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.
+    // Let the run owner checkpoint its latest SDK snapshot synchronously
+    // BEFORE the input lock below is released, so a follow-up message sent
+    // the moment the user hits Esc inherits the latest state instead of
+    // stale (or null) continuation state. Best-effort: an onAbort failure
+    // must never block interruption cleanup.
+    try {
+      params.onAbort?.()
+    } catch {
+      // Checkpoint callbacks are best-effort; never skip abort cleanup.
+    }
+    // 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.
     streamRefs.setters.setWasAbortedByUser(true)
     setIsRetrying(false)
     timerController.stop('aborted')

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

@@ -172,6 +172,11 @@ export const useSendMessage = ({
   const previousRunStateRef = useRef<RunState | null>(
     useChatStore.getState().runState,
   )
+  // Incremented for every run that is admitted as a real SDK run. Late
+  // results from a superseded run (one replaced by a newer run after the
+  // input lock was released) must never adopt state, persist a checkpoint,
+  // or touch shared queue state over the run that replaced it.
+  const runGenerationRef = useRef(0)
   // Memoize stream controller to maintain referential stability across renders
   const streamRefsRef = useRef<ReturnType<
     typeof createStreamController
@@ -310,6 +315,23 @@ export const useSendMessage = ({
       const abortController = new AbortController()
       const runChatDir = resolveCurrentChatDir()
       const runChatIsCurrent = () => resolveCurrentChatDir() === runChatDir
+      // Bump only after the run-start guard admits the message: a
+      // session-ended message that gets requeued must not supersede an
+      // active run.
+      const runGeneration = ++runGenerationRef.current
+      const runIsCurrent = () =>
+        runGenerationRef.current === runGeneration && runChatIsCurrent()
+      // Adopt a snapshot as the continuation state for the next message, in
+      // memory (ref) and React state. Skipped when the run has been
+      // superseded or the chat switched away, and for snapshots that carry
+      // no session state at all: adopting one would silently make the SDK
+      // start a blank session on the next prompt.
+      const syncRunState = (state: RunState) => {
+        if (!runIsCurrent()) return
+        if (!state.sessionState) return
+        previousRunStateRef.current = state
+        setRunState(state)
+      }
       let latestRunStateSnapshot: RunState = previousRunStateRef.current ?? {
         traceSessionId: randomUUID(),
         output: {
@@ -520,6 +542,14 @@ export const useSendMessage = ({
         setMessages,
         streamRefs,
         abortController,
+        onAbort: () => {
+          // Sync before the abort listener releases the input lock, so a
+          // message sent the moment the user hits Esc resumes from the
+          // latest snapshot instead of stale (or null) state. The
+          // generation/chat guards in syncRunState keep this a no-op for
+          // superseded runs and context-changing stops.
+          syncRunState(latestRunStateSnapshot)
+        },
         setStreamStatus,
         setCanProcessQueue,
         isQueuePausedRef,
@@ -565,7 +595,7 @@ export const useSendMessage = ({
         )
 
         const eventHandlerState = createEventHandlerState({
-          isActive: () => !abortController.signal.aborted && runChatIsCurrent(),
+          isActive: () => !abortController.signal.aborted && runIsCurrent(),
           streamRefs,
           setStreamingAgents,
           setStreamStatus,
@@ -625,7 +655,7 @@ export const useSendMessage = ({
             // conversation, and checkpointing them into this run's directory
             // would overwrite that chat's transcript with foreign (possibly
             // empty) state — the chat would then be hidden from /history.
-            if (abortController.signal.aborted || !runChatIsCurrent()) {
+            if (abortController.signal.aborted || !runIsCurrent()) {
               return
             }
             previousRunStateRef.current = snapshot
@@ -685,10 +715,9 @@ export const useSendMessage = ({
         // context, and previousRunStateRef/setRunState would leak this run's
         // agent state into the other chat. (A plain Esc interrupt keeps the
         // same chat, so the interrupted turn is still saved as before.)
-        if (!abortController.signal.aborted && runChatIsCurrent()) {
+        if (runIsCurrent()) {
           // 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
@@ -700,29 +729,31 @@ export const useSendMessage = ({
           // traps is several times slower.
           saveChatState(runState, useChatStore.getState().messages, runChatDir)
         }
-        handleRunCompletion({
-          runState,
-          actualCredits,
-          agentMode,
-          timerController,
-          updater,
-          aiMessageId,
-          wasAbortedByUser: abortController.signal.aborted,
-          hasReceivedContent: hasReceivedContentRef.current,
-          setStreamStatus,
-          setCanProcessQueue,
-          updateChainInProgress,
-          setHasReceivedPlanResponse,
-          resumeQueue,
-          isProcessingQueueRef,
-          isQueuePausedRef,
-        })
+        if (runIsCurrent()) {
+          handleRunCompletion({
+            runState,
+            actualCredits,
+            agentMode,
+            timerController,
+            updater,
+            aiMessageId,
+            wasAbortedByUser: abortController.signal.aborted,
+            hasReceivedContent: hasReceivedContentRef.current,
+            setStreamStatus,
+            setCanProcessQueue,
+            updateChainInProgress,
+            setHasReceivedPlanResponse,
+            resumeQueue,
+            isProcessingQueueRef,
+            isQueuePausedRef,
+          })
+        }
       } catch (error) {
         // If this run was aborted, the abort handler already handled cleanup.
         // Don't run error handling to avoid interfering with any new run that
         // may have started. Uses per-run abortController.signal (not shared
         // streamRefs) so a newer run's reset() can't clear this flag.
-        if (!abortController.signal.aborted) {
+        if (!abortController.signal.aborted && runIsCurrent()) {
           handleRunError({
             error,
             timerController,
@@ -735,20 +766,26 @@ export const useSendMessage = ({
             isQueuePausedRef,
             hasReceivedContent: hasReceivedContentRef.current,
           })
+          // Keep the latest snapshot available to the next message in this
+          // process, not only on disk: without this, a failed or expired
+          // turn is followed by a run with stale (or null) history.
+          syncRunState(latestRunStateSnapshot)
           // Persist the last checkpoint plus the error banner so a restart
-          // after a failed run still shows this turn. Settle async checkpoints
-          // 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()) {
-            await settleCheckpointSave()
-            saveChatState(
-              latestRunStateSnapshot,
-              useChatStore.getState().messages,
-              runChatDir,
-            )
-          }
-        } else {
+          // after a failed run still shows this turn. Settle async
+          // checkpoints first so a stale write can't clobber this one.
+          await settleCheckpointSave()
+          saveChatState(
+            latestRunStateSnapshot,
+            useChatStore.getState().messages,
+            runChatDir,
+          )
+        } else if (abortController.signal.aborted) {
           logger.debug({ error }, '[send-message] Ignoring error after abort')
+        } else {
+          logger.debug(
+            { error },
+            '[send-message] Ignoring error after run superseded',
+          )
         }
       } finally {
         // Close the steering mailbox. Anything the run never drained was
@@ -792,7 +829,7 @@ export const useSendMessage = ({
         // interfering with any new run that may have started after the abort.
         // Uses per-run abortController.signal (not shared streamRefs) so a newer
         // run's reset() can't clear this flag.
-        if (!abortController.signal.aborted) {
+        if (!abortController.signal.aborted && runIsCurrent()) {
           if (isChainInProgressRef.current) {
             logger.warn(
               {},