Fix blank-session restarts after interrupt or failed run

Fixes #1054. Esc-aborting a session and sending a
follow-up could start a brand new, blank conversation because the
continuation state (previousRunStateRef) was only synced when run()
settled, while the abort listener released the input lock immediately.

- SetupStreamingContext gains an onAbort callback invoked synchronously
  at the top of the abort listener, before the lock is released, so the
  run owner can checkpoint its latest SDK snapshot.
- useSendMessage passes syncRunState(latestRunStateSnapshot) as onAbort,
  and syncRunState now guards with a generation token so a superseded
  run settling late can never adopt state, persist a checkpoint, or
  touch shared queue state over the run that replaced it.
- The catch path (failed/expired runs) now also syncs the ref, not just
  disk, so the next prompt resumes from the last snapshot instead of
  stale or null history.
- loadMostRecentChatState stops adopting/persisting sessionState-less
  run states, which made the SDK build a blank session on restart.

Adds hook-level regression tests through the real createRunConfig /
client.run wiring: abort then follow-up carries full history, late
superseded runs cannot clobber, the rejected-run path resumes from the
last snapshot, and sessionState-less states are never adopted. Helper
and storage suites updated.

Author: bunnysayzz

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

@@ -0,0 +1,398 @@
+import { describe, expect, test, beforeEach, afterEach } from 'bun:test'
+import * as fs from 'fs'
+import * as os from 'os'
+import * as path from 'path'
+import { createTestRenderer } from '@opentui/core/testing'
+import { createRoot, flushSync } from '@opentui/react'
+import React, { useEffect, useRef } from 'react'
+
+import { useSendMessage } from '../use-send-message'
+import { setClientFactoryOverrideForTesting } from '../../utils/codebuff-client'
+import { stopActiveRun } from '../../utils/active-run'
+import { setChatDirOverrideForTesting } from '../../utils/run-state-storage'
+import { setProjectRoot } from '../../project-files'
+import { useChatStore } from '../../state/chat-store'
+
+import type { RunState } from '@codebuff/sdk'
+import type { SendMessageFn } from '../../types/contracts/send-message'
+import type { ElapsedTimeTracker } from '../use-elapsed-time'
+
+// Ensure required env vars exist so logger/env parsing succeeds in tests.
+const ensureEnv = () => {
+  process.env.NEXT_PUBLIC_CB_ENVIRONMENT =
+    process.env.NEXT_PUBLIC_CB_ENVIRONMENT || 'test'
+  process.env.NEXT_PUBLIC_CODEBUFF_APP_URL =
+    process.env.NEXT_PUBLIC_CODEBUFF_APP_URL || 'https://app.codebuff.test'
+  process.env.NEXT_PUBLIC_SUPPORT_EMAIL =
+    process.env.NEXT_PUBLIC_SUPPORT_EMAIL || 'support@codebuff.test'
+  process.env.NEXT_PUBLIC_POSTHOG_API_KEY =
+    process.env.NEXT_PUBLIC_POSTHOG_API_KEY || 'phc_test_key'
+  process.env.NEXT_PUBLIC_POSTHOG_HOST_URL =
+    process.env.NEXT_PUBLIC_POSTHOG_HOST_URL || 'https://posthog.codebuff.test'
+  process.env.NEXT_PUBLIC_STRIPE_PUBLISHABLE_KEY =
+    process.env.NEXT_PUBLIC_STRIPE_PUBLISHABLE_KEY || 'pk_test_123'
+  process.env.NEXT_PUBLIC_STRIPE_CUSTOMER_PORTAL =
+    process.env.NEXT_PUBLIC_STRIPE_CUSTOMER_PORTAL ||
+    'https://stripe.codebuff.test'
+  process.env.NEXT_PUBLIC_WEB_PORT = process.env.NEXT_PUBLIC_WEB_PORT || '3000'
+}
+ensureEnv()
+
+let api: { sendMessage: SendMessageFn; clearMessages: () => void } | null = null
+
+const makeTimer = (): ElapsedTimeTracker => ({
+  start: () => {},
+  stop: () => {},
+  pause: () => {},
+  resume: () => {},
+  elapsedSeconds: 0,
+  startTime: null,
+  isPaused: false,
+})
+
+const Host = () => {
+  const inputRef = useRef<any>(null)
+  const activeSubagentsRef = useRef<Set<string>>(new Set())
+  const isChainInProgressRef = useRef(false)
+  const isQueuePausedRef = useRef(false)
+  const isProcessingQueueRef = useRef(false)
+  const mainAgentTimer = useRef<ElapsedTimeTracker>(makeTimer()).current
+
+  const { sendMessage, clearMessages } = useSendMessage({
+    inputRef,
+    activeSubagentsRef,
+    isChainInProgressRef,
+    setStreamStatus: () => {},
+    setCanProcessQueue: () => {},
+    agentId: undefined,
+    onBeforeMessageSend: async () => ({ success: true, errors: [] }),
+    mainAgentTimer,
+    scrollToLatest: () => {},
+    onTimerEvent: () => {},
+    isQueuePausedRef,
+    isProcessingQueueRef,
+    resumeQueue: () => {},
+    requeueMessageAtFront: () => {},
+    continueChat: false,
+    subscriptionData: null,
+  })
+
+  useEffect(() => {
+    api = { sendMessage, clearMessages }
+  }, [sendMessage, clearMessages])
+
+  return <text>host</text>
+}
+
+type RunCall = {
+  runConfig: any
+  resolve: (state: RunState) => void
+  reject: (error: unknown) => void
+  settled: boolean
+}
+
+let runCalls: RunCall[] = []
+
+const fakeClient = {
+  run: (runConfig: any) =>
+    new Promise<RunState>((resolve, reject) => {
+      runCalls.push({ runConfig, resolve, reject, settled: false })
+    }),
+}
+
+const TEST_ROOT = fs.mkdtempSync(path.join(os.tmpdir(), 'use-send-message-'))
+
+// Timer-based poll, not setImmediate: the sendMessage preflight does real
+// async work (checkpoint drain, filesystem, yieldToEventLoop), and a tight
+// setImmediate loop starves the timer phase before the next run registers.
+const waitFor = async (label: string, fn: () => boolean, timeoutMs = 5000) => {
+  const deadline = Date.now() + timeoutMs
+  while (Date.now() < deadline) {
+    if (fn()) return
+    await new Promise((r) => setTimeout(r, 10))
+  }
+  throw new Error(`Timed out waiting for ${label}`)
+}
+
+/** History shape a real in-flight SDK snapshot carries: a user prompt, an
+ *  assistant response, and a completed tool call. */
+const threeMessageHistory = () => [
+  {
+    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'] } }],
+  },
+]
+
+const makeSnapshot = (marker: string, history: any[]): RunState =>
+  ({
+    traceSessionId: `trace-${marker}`,
+    sessionState: {
+      mainAgentState: {
+        messageHistory: history,
+      },
+    },
+    output: {
+      type: 'error',
+      message: 'Session ended before this response completed.',
+    },
+  }) as unknown as RunState
+
+const makeCompleted = (marker: string, history: any[]): RunState => ({
+  traceSessionId: `trace-${marker}-done`,
+  sessionState: {
+    mainAgentState: {
+      messageHistory: history,
+    },
+  } as any,
+  output: { type: 'lastMessage', value: [] },
+})
+
+const resolveAllPending = () => {
+  for (const call of runCalls) {
+    if (!call.settled) {
+      call.settled = true
+      call.resolve(
+        makeCompleted(call.runConfig.previousRun?.traceSessionId ?? 'x', []),
+      )
+    }
+  }
+}
+
+beforeEach(() => {
+  runCalls = []
+  api = null
+  setProjectRoot(process.cwd())
+  setChatDirOverrideForTesting(TEST_ROOT)
+  setClientFactoryOverrideForTesting(async () => fakeClient as any)
+  useChatStore.getState().reset()
+})
+
+afterEach(() => {
+  stopActiveRun('process-exit')
+  setClientFactoryOverrideForTesting(undefined)
+  setChatDirOverrideForTesting(undefined)
+})
+
+describe('useSendMessage session-state preservation (freebuff #1054)', () => {
+  test('Esc abort: an immediate follow-up run inherits the interrupted snapshot', async () => {
+    const setup = await createTestRenderer({ width: 80, height: 3 })
+    const root = createRoot(setup.renderer)
+    flushSync(() => {
+      root.render(<Host />)
+    })
+    await setup.renderOnce()
+
+    let run1: Promise<void> | null = null
+    let run2: Promise<void> | null = null
+    try {
+      flushSync(() => {
+        run1 = api!.sendMessage({ content: 'first', agentMode: 'DEFAULT' })
+      })
+      await waitFor('first run to start', () => runCalls.length === 1)
+
+      // The SDK emits an in-flight snapshot that includes real history.
+      const richSnapshot = makeSnapshot('A', threeMessageHistory())
+      runCalls[0].runConfig.onStateSnapshot(richSnapshot)
+
+      // User hits Esc. The abort listener must synchronously adopt the
+      // latest snapshot into continuation state before releasing input.
+      flushSync(() => {
+        stopActiveRun('user-interrupt')
+      })
+
+      flushSync(() => {
+        run2 = api!.sendMessage({
+          content: 'also verify tests',
+          agentMode: 'DEFAULT',
+        })
+      })
+      await waitFor('second run to start', () => runCalls.length === 2)
+
+      // Regression: the follow-up must carry the interrupted run's full
+      // history instead of an empty previousRun that makes the SDK start a
+      // blank session.
+      const previousRun = runCalls[1].runConfig.previousRun as RunState
+      expect(previousRun).toBeDefined()
+      expect(
+        (previousRun.sessionState as any).mainAgentState.messageHistory,
+      ).toHaveLength(3)
+      expect(
+        (previousRun.sessionState as any).mainAgentState.messageHistory[0]
+          .content[0].text,
+      ).toBe('implement authentication')
+      expect(
+        (previousRun.sessionState as any).mainAgentState.messageHistory[2]
+          .toolName,
+      ).toBe('read_files')
+    } finally {
+      resolveAllPending()
+      await Promise.all([run1, run2])
+      flushSync(() => root.unmount())
+      setup.renderer.destroy()
+    }
+  })
+
+  test('rejected run: the next prompt resumes from the last snapshot', async () => {
+    const setup = await createTestRenderer({ width: 80, height: 3 })
+    const root = createRoot(setup.renderer)
+    flushSync(() => {
+      root.render(<Host />)
+    })
+    await setup.renderOnce()
+
+    let run1: Promise<void> | null = null
+    let run2: Promise<void> | null = null
+    try {
+      flushSync(() => {
+        run1 = api!.sendMessage({ content: 'first', agentMode: 'DEFAULT' })
+      })
+      await waitFor('first run to start', () => runCalls.length === 1)
+
+      const richSnapshot = makeSnapshot('A', threeMessageHistory())
+      runCalls[0].runConfig.onStateSnapshot(richSnapshot)
+
+      // The SDK run throws (network/auth/session gate). The catch block
+      // must keep the latest snapshot available in memory for the next
+      // message, not only on disk.
+      runCalls[0].reject(new Error('boom'))
+      await run1
+
+      flushSync(() => {
+        run2 = api!.sendMessage({ content: 'continue', agentMode: 'DEFAULT' })
+      })
+      await waitFor('second run to start', () => runCalls.length === 2)
+
+      const previousRun = runCalls[1].runConfig.previousRun as RunState
+      expect(previousRun).toBeDefined()
+      expect(previousRun.traceSessionId).toBe('trace-A')
+      expect(
+        (previousRun.sessionState as any).mainAgentState.messageHistory,
+      ).toHaveLength(3)
+    } finally {
+      resolveAllPending()
+      await Promise.all([run1, run2])
+      flushSync(() => root.unmount())
+      setup.renderer.destroy()
+    }
+  })
+
+  test('superseded run settling late cannot clobber the newer run state', async () => {
+    const setup = await createTestRenderer({ width: 80, height: 3 })
+    const root = createRoot(setup.renderer)
+    flushSync(() => {
+      root.render(<Host />)
+    })
+    await setup.renderOnce()
+
+    let run1: Promise<void> | null = null
+    let run2: Promise<void> | null = null
+    let run3: Promise<void> | null = null
+    try {
+      flushSync(() => {
+        run1 = api!.sendMessage({ content: 'first', agentMode: 'DEFAULT' })
+      })
+      await waitFor('first run to start', () => runCalls.length === 1)
+      const snapshotA = makeSnapshot('A', threeMessageHistory())
+      runCalls[0].runConfig.onStateSnapshot(snapshotA)
+
+      // Abort run A, then start run B before A's SDK promise settles.
+      flushSync(() => {
+        stopActiveRun('user-interrupt')
+        run2 = api!.sendMessage({ content: 'second', agentMode: 'DEFAULT' })
+      })
+      await waitFor('second run to start', () => runCalls.length === 2)
+
+      // B settles first with its own final state.
+      runCalls[1].resolve(
+        makeCompleted('B', [
+          { role: 'user', content: [{ type: 'text', text: 'B final' }] },
+        ]),
+      )
+      await run2
+
+      // Now A settles late. Its state must NOT replace B's: B is the
+      // current run.
+      runCalls[0].resolve(
+        makeCompleted('A-late', [
+          { role: 'user', content: [{ type: 'text', text: 'A late result' }] },
+        ]),
+      )
+      await run1
+
+      flushSync(() => {
+        run3 = api!.sendMessage({ content: 'third', agentMode: 'DEFAULT' })
+      })
+      await waitFor('third run to start', () => runCalls.length === 3)
+
+      const previousRun = runCalls[2].runConfig.previousRun as RunState
+      expect(previousRun.traceSessionId).toBe('trace-B-done')
+      expect(
+        (previousRun.sessionState as any).mainAgentState.messageHistory[0]
+          .content[0].text,
+      ).toBe('B final')
+    } finally {
+      resolveAllPending()
+      await Promise.all([run1, run2, run3])
+      flushSync(() => root.unmount())
+      setup.renderer.destroy()
+    }
+  })
+
+  test('sessionState-less snapshot is never adopted as continuation state', async () => {
+    const setup = await createTestRenderer({ width: 80, height: 3 })
+    const root = createRoot(setup.renderer)
+    flushSync(() => {
+      root.render(<Host />)
+    })
+    await setup.renderOnce()
+
+    let run1: Promise<void> | null = null
+    let run2: Promise<void> | null = null
+    try {
+      flushSync(() => {
+        run1 = api!.sendMessage({ content: 'first', agentMode: 'DEFAULT' })
+      })
+      await waitFor('first run to start', () => runCalls.length === 1)
+
+      // Abort before the SDK ever emitted a sessionState-bearing snapshot.
+      // The placeholder snapshot has no session state and must not be
+      // adopted (it would blank the next run).
+      flushSync(() => {
+        stopActiveRun('user-interrupt')
+      })
+
+      // The abort resolver returns a state without any sessionState (no
+      // snapshot was ever emitted). It must not be adopted as continuation
+      // state: handing it to the SDK would silently blank the next session.
+      runCalls[0].resolve({
+        traceSessionId: 'trace-A-no-state',
+        output: {
+          type: 'error',
+          message: 'Session ended before this response completed.',
+        },
+      } as RunState)
+      await run1
+
+      flushSync(() => {
+        run2 = api!.sendMessage({ content: 'second', agentMode: 'DEFAULT' })
+      })
+      await waitFor('second run to start', () => runCalls.length === 2)
+
+      expect(runCalls[1].runConfig.previousRun).toBeUndefined()
+    } finally {
+      resolveAllPending()
+      await Promise.all([run1, run2])
+      flushSync(() => root.unmount())
+      setup.renderer.destroy()
+    }
+  })
+})