fix(ai): persist truncated turns across chat entry points

Preserve partial output and token usage in CLI turns and edited UI rounds while leaving ordinary failures and cancellation unchanged. Use OutputLimitError across agent runners so incomplete turns remain visibly unsuccessful but can be reloaded and continued.

Cover real stream handling, SQLite persistence, reload, CLI continuation, and unsuccessful one-shot JSON output.

apps/cli/src/ai/agent.ts

@@ -272,7 +272,10 @@ export async function runServerAgent(options: RunAgentOptions): Promise<void> {
 
     if (result.error && !abortSignal?.aborted) {
       const friendlyMessage = formatAIError(result.error)
-      onEvent({ type: 'error', error: { name: 'AgentError', message: friendlyMessage } })
+      onEvent({
+        type: 'error',
+        error: { name: result.stopReason === 'length' ? 'OutputLimitError' : 'AgentError', message: friendlyMessage },
+      })
     }
 
     handler.emitStatus(abortSignal?.aborted ? 'aborted' : result.error ? 'error' : 'completed', cachedMessages, {

apps/cli/src/ai/chat-command.test.ts

@@ -498,6 +498,99 @@ describe('runChatTurn', () => {
 })
 
 describe('runChatCommand', () => {
+  it('keeps the first truncated turn in the same interactive conversation for continuation', async () => {
+    const stdout = new MemoryWritable()
+    const stderr = new MemoryWritable()
+    const stdin = new PromptDrivenReadable(['first', 'continue', 'exit'])
+    const aiChatManager = createAIChatManager()
+    const chatIds: string[] = []
+    stdout.onChunk = (text) => {
+      if (text.includes('chatlab> ')) stdin.pushNext()
+    }
+    await runChatCommand(
+      { sessionId: 'session-1' },
+      {
+        dbManager: createDbManager(['session-1']),
+        pathProvider: {} as never,
+        aiChatManager,
+        stdout,
+        stderr,
+        stdin,
+        createRunAgentStream: () => async (params, onEvent) => {
+          chatIds.push(params.aiChatId)
+          onEvent({ type: 'content', content: params.userMessage === 'first' ? 'Partial answer' : 'Continued answer' })
+          if (params.userMessage === 'first') {
+            onEvent({ type: 'error', error: { name: 'OutputLimitError', message: 'Output limit reached' } })
+          }
+          onEvent({ type: 'done', isFinished: true })
+        },
+      }
+    )
+    assert.deepEqual(chatIds, ['ai_chat_1', 'ai_chat_1'])
+    assert.match(stderr.text(), /Output limit reached/)
+    assert.deepEqual(
+      (aiChatManager as unknown as { __messages: Array<{ content: string }> }).__messages.map(
+        (message) => message.content
+      ),
+      ['first', 'Partial answer', 'continue', 'Continued answer']
+    )
+  })
+
+  it('outputs a persisted partial JSON result but still fails a one-shot truncated command', async () => {
+    const stdout = new MemoryWritable()
+    const aiChatManager = createAIChatManager()
+    await assert.rejects(
+      () =>
+        runChatCommand(
+          { sessionId: 'session-1', question: 'hello', json: true },
+          {
+            dbManager: createDbManager(['session-1']),
+            pathProvider: {} as never,
+            aiChatManager,
+            stdout,
+            createRunAgentStream: () => async (_params, onEvent) => {
+              onEvent({
+                type: 'plan',
+                plan: {
+                  type: 'plan',
+                  version: 1,
+                  status: 'executing',
+                  plan: {
+                    version: 1,
+                    title: 'Review the evidence',
+                    route: 'planned_execution',
+                    intent: 'summary',
+                    steps: [{ goal: 'Summarize the evidence', suggestedTools: [], evidenceNeeded: 'Recent messages' }],
+                    successCriteria: ['Provide a complete summary'],
+                  },
+                },
+              })
+              onEvent({ type: 'content', content: 'Partial answer' })
+              onEvent({ type: 'error', error: { name: 'OutputLimitError', message: 'Output limit reached' } })
+              onEvent({
+                type: 'done',
+                isFinished: true,
+                usage: {
+                  promptTokens: 10,
+                  completionTokens: 10,
+                  totalTokens: 20,
+                  cacheReadTokens: 0,
+                  cacheWriteTokens: 0,
+                },
+              })
+            },
+          }
+        ),
+      { name: 'OutputLimitError' }
+    )
+    const result = JSON.parse(stdout.text())
+    assert.equal(result.aiChatId, 'ai_chat_1')
+    assert.equal(result.answer, 'Partial answer')
+    assert.equal(result.error.name, 'OutputLimitError')
+    assert.equal(result.contentBlocks.find((block: ContentBlock) => block.type === 'plan')?.status, 'skipped')
+    assert.equal((aiChatManager as unknown as { __messages: unknown[] }).__messages.length, 2)
+  })
+
   it('keeps interactive mode alive after a single failed turn', async () => {
     const stdout = new MemoryWritable()
     const stderr = new MemoryWritable()

apps/cli/src/ai/chat-command.ts

@@ -45,6 +45,8 @@ export interface ChatTurnResult {
   aiChatId: string
   question: string
   answer: string
+  /** Present for a persisted but incomplete turn, never for a successful completion. */
+  error?: { name: string; message: string }
   events?: AgentStreamChunk[]
   contentBlocks?: ContentBlock[]
   usage: {
@@ -161,7 +163,7 @@ export async function runChatTurn(
   const startedAt = Date.now()
   let answer = ''
   let tokenUsage: TokenUsageData | null = null
-  let streamError: Error | null = null
+  const streamState: { error: Error | null } = { error: null }
   const events: AgentStreamChunk[] = []
   const contentBlocks: ContentBlock[] = []
   let hasReplayContentBlocks = false
@@ -269,7 +271,7 @@ export async function runChatTurn(
       }
       if (chunk.type === 'error') {
         removePlanDraftBlocks()
-        streamError = createAgentStreamError(chunk.error)
+        streamState.error = createAgentStreamError(chunk.error)
         return
       }
       if (chunk.type === 'plan_delta' && chunk.planDelta) {
@@ -302,7 +304,7 @@ export async function runChatTurn(
         for (let index = contentBlocks.length - 1; index >= 0; index--) {
           const block = contentBlocks[index]
           if (block.type === 'plan') {
-            block.status = 'done'
+            block.status = streamState.error ? 'skipped' : 'done'
             break
           }
         }
@@ -311,7 +313,15 @@ export async function runChatTurn(
     new AbortController().signal
   )
 
-  if (streamError) throw streamError
+  const streamError = streamState.error
+  if (streamError && streamError.name !== 'OutputLimitError') throw streamError
+  if (streamError) {
+    contentBlocks.push({
+      type: 'error',
+      error: { name: streamError.name, message: streamError.message, stack: null },
+    })
+    hasReplayContentBlocks = true
+  }
 
   if (!options.json && options.stream !== false && answer && !answer.endsWith('\n')) {
     write(stdout, '\n')
@@ -332,6 +342,7 @@ export async function runChatTurn(
     aiChatId: target.aiChatId,
     question: options.question,
     answer,
+    ...(streamError ? { error: { name: streamError.name, message: streamError.message } } : {}),
     ...(options.includeEvents ? { events } : {}),
     ...(hasReplayContentBlocks ? { contentBlocks } : {}),
     usage: {
@@ -352,6 +363,7 @@ export async function runChatCommand(options: ChatCommandOptions, deps: ChatComm
     } else if (options.stream === false) {
       write(stdout, `${result.answer}\n`)
     }
+    if (result.error) throw createAgentStreamError(result.error)
     return
   }
 
@@ -368,6 +380,7 @@ export async function runChatCommand(options: ChatCommandOptions, deps: ChatComm
       try {
         const result = await runChatTurn({ ...options, aiChatId, question, json: false, stream: true }, deps)
         aiChatId = result.aiChatId
+        if (result.error) write(stderr, `${result.error.message}\n`)
       } catch (error) {
         write(stderr, `${error instanceof Error ? error.message : String(error)}\n`)
       }

apps/desktop/main/ai/agent/index.ts

@@ -301,6 +301,7 @@ export class Agent {
             requestBody?: string
           }
         }
+        if (result.stopReason === 'length') agentError.name = 'OutputLimitError'
         const lastMsg = [...result.finalMessages].reverse().find((m) => m.role === 'assistant') as
           | (PiAssistantMessage & { provider?: string; model?: string; api?: string })
           | undefined

packages/node-runtime/src/ai/cross-chat-agent.ts

@@ -196,7 +196,13 @@ export async function runCrossChatAgent(options: RunCrossChatAgentOptions): Prom
       return
     }
     if (result.error) {
-      onEvent({ type: 'error', error: { name: 'AgentError', message: formatAIError(result.error) } })
+      onEvent({
+        type: 'error',
+        error: {
+          name: result.stopReason === 'length' ? 'OutputLimitError' : 'AgentError',
+          message: formatAIError(result.error),
+        },
+      })
     }
     handler.emitStatus(result.error ? 'error' : 'completed', cachedMessages, { force: true })
     logger?.info('CrossChatAgent', 'Cross-chat agent execution ended', {

src/stores/aiChat.ts

@@ -2024,7 +2024,8 @@ export const useAIChatStore = defineStore('aiChatRuntime', () => {
         return { success: false, reason: 'busy', activeTask: activeTask.value }
       }
 
-      if (result.success && result.result) {
+      const wasTruncated = result.error?.name === 'OutputLimitError'
+      if ((result.success || wasTruncated) && result.result) {
         updateAIMessage({
           dataSource: { toolsUsed: result.result.toolsUsed, toolRounds: result.result.toolRounds },
           isStreaming: false,
@@ -2036,7 +2037,8 @@ export const useAIChatStore = defineStore('aiChatRuntime', () => {
         updateAIMessage({ contentBlocks: [...blocks], isStreaming: false })
       }
 
-      if (!result.success) {
+      // A completed provider stream can be incomplete without invalidating its generated content.
+      if (!result.success && !wasTruncated) {
         restoreOriginal()
         return { success: false, reason: 'error' }
       }
@@ -2065,7 +2067,7 @@ export const useAIChatStore = defineStore('aiChatRuntime', () => {
 
       targetBuffer.sessionTokenUsage = toTokenUsage(await useAIService().getAIChatTokenUsage(state.currentAIChatId!))
       state.sessionTokenUsage = { ...targetBuffer.sessionTokenUsage }
-      return { success: true }
+      return result.success ? { success: true } : { success: false, reason: 'error' }
     } catch (error) {
       if (state.isAborted) {
         restoreOriginal()

src/stores/aiChatTruncatedPersistence.test.ts

@@ -0,0 +1,161 @@
+import assert from 'node:assert/strict'
+import { mkdtempSync, rmSync } from 'node:fs'
+import { tmpdir } from 'node:os'
+import { join } from 'node:path'
+import test from 'node:test'
+import { createPinia, setActivePinia } from 'pinia'
+import { ref } from 'vue'
+import { AIChatManager } from '../../packages/node-runtime/src/ai/chats'
+import type { AgentStreamChunk } from '@/services/ai-stream/types'
+
+test('edited turns retain truncated output and usage for reload and continuation, but roll back real failures', async (t) => {
+  const dir = mkdtempSync(join(tmpdir(), 'chatlab-truncated-edit-'))
+  const manager = new AIChatManager(dir, { nativeBinding: process.env.CHATLAB_TEST_SQLITE_NATIVE_BINDING })
+  const usage = { promptTokens: 100, completionTokens: 25, totalTokens: 125, cacheReadTokens: 0, cacheWriteTokens: 0 }
+  let events: AgentStreamChunk[] = []
+  let nextHistory: ReturnType<AIChatManager['getHistoryForAgent']> = []
+  const aiService = {
+    getAIChat: async (id: string) => manager.getAIChat(id),
+    getMessages: async (id: string) => manager.getMessages(id),
+    getAIChatTokenUsage: async (id: string) => manager.getAIChatTokenUsage(id),
+    addMessagePair: async (...args: Parameters<AIChatManager['addMessagePair']>) => manager.addMessagePair(...args),
+    replaceLatestMessageRound: async (...args: Parameters<AIChatManager['replaceLatestMessageRound']>) =>
+      manager.replaceLatestMessageRound(...args),
+  }
+  try {
+    await t.mock.module('@/stores/session', { namedExports: { useSessionStore: () => ({ sessions: [] }) } })
+    await t.mock.module('@/stores/settings', {
+      namedExports: {
+        useSettingsStore: () => ({
+          aiPreprocessConfig: {
+            dataCleaning: false,
+            mergeConsecutive: false,
+            mergeWindowSeconds: 180,
+            blacklistKeywords: [],
+            denoise: false,
+            desensitize: false,
+            desensitizeRules: [],
+            anonymizeNames: false,
+          },
+        }),
+      },
+    })
+    await t.mock.module('@/stores/assistant', {
+      namedExports: {
+        useAssistantStore: () => ({
+          isLoaded: true,
+          loadAssistants: async () => undefined,
+          selectAssistant: () => undefined,
+          clearSelection: () => undefined,
+        }),
+      },
+    })
+    await t.mock.module('@/stores/skill', {
+      namedExports: { useSkillStore: () => ({ activeSkillId: null, activeSkill: ref(null) }) },
+    })
+    await t.mock.module('@/stores/llm', { namedExports: { useLLMStore: () => ({}) } })
+    await t.mock.module('@/services', {
+      namedExports: {
+        useAIService: () => aiService,
+        useDataService: () => ({}),
+        useLLMService: () => ({ hasConfig: async () => true }),
+      },
+    })
+    await t.mock.module('@/services/product-analytics', { namedExports: { trackProductEvent: () => undefined } })
+    // Keep the real stream service: an error followed by done must still yield success=false.
+    await t.mock.module('@/services/utils/sse', {
+      namedExports: {
+        fetchSSE: async (options: {
+          body: { aiChatId: string }
+          onEvent: (event: { event: string; data: string }) => void
+        }) => {
+          nextHistory = manager.getHistoryForAgent(options.body.aiChatId)
+          for (const event of events) options.onEvent({ event: event.type, data: JSON.stringify(event) })
+        },
+      },
+    })
+
+    const { useAIChatStore } = await import('./aiChat')
+    for (const scenario of [
+      { name: 'partial answer', text: 'Partial answer about Project Cedar', errorName: 'OutputLimitError' },
+      { name: 'thinking only', text: '', errorName: 'OutputLimitError' },
+      { name: 'configuration failure', text: '', errorName: 'ConfigError' },
+    ]) {
+      await t.test(scenario.name, async () => {
+        setActivePinia(createPinia())
+        const store = useAIChatStore()
+        const chat = manager.createAIChat('session-1', scenario.name, 'general_cn')
+        const original = manager.addMessagePair(
+          chat.id,
+          { content: 'Original question' },
+          { content: 'Original answer' }
+        )
+        const { chatKey, state } = store.ensureSessionState({
+          sessionId: 'session-1',
+          sessionName: 'Fixture',
+          chatType: 'private',
+          locale: 'zh-CN',
+        })
+        await store.loadAIChat(chatKey, chat.id)
+        const truncated = scenario.errorName === 'OutputLimitError'
+        events = []
+        if (truncated) {
+          events.push(
+            { type: 'think', content: 'Checking the evidence.', thinkTag: 'thinking' },
+            { type: 'content', content: scenario.text }
+          )
+        }
+        events.push(
+          { type: 'error', error: { name: scenario.errorName, message: 'Generation was interrupted.' } },
+          { type: 'done', isFinished: true, ...(truncated ? { usage } : {}) }
+        )
+        const result = await store.editMessageAndRegenerate(chatKey, original.userMessage.id, 'Edited question')
+        assert.deepEqual(result, { success: false, reason: 'error' })
+        const expectedContents = truncated
+          ? ['Edited question', scenario.text]
+          : ['Original question', 'Original answer']
+        assert.deepEqual(
+          state.messages.map((message) => message.content),
+          expectedContents
+        )
+        assert.deepEqual(
+          manager.getMessages(chat.id).map((message) => message.content),
+          expectedContents
+        )
+        assert.equal(manager.getAIChatTokenUsage(chat.id).totalTokens, truncated ? usage.totalTokens : 0)
+        if (!truncated) return
+        const savedAnswer = manager.getMessages(chat.id)[1]
+        assert.ok(savedAnswer.contentBlocks?.some((block) => block.type === 'think'))
+        assert.ok(
+          savedAnswer.contentBlocks?.some((block) => block.type === 'error' && block.error.name === 'OutputLimitError')
+        )
+        // Re-open the database, then load the stored turn into a fresh frontend store.
+        manager.close()
+        setActivePinia(createPinia())
+        const reloadedStore = useAIChatStore()
+        const reloaded = reloadedStore.ensureSessionState({
+          sessionId: 'session-1',
+          sessionName: 'Fixture',
+          chatType: 'private',
+          locale: 'zh-CN',
+        })
+        await reloadedStore.loadAIChat(reloaded.chatKey, chat.id)
+        assert.deepEqual(
+          reloaded.state.messages.map((message) => message.content),
+          expectedContents
+        )
+        events = [
+          { type: 'content', content: 'Continued answer' },
+          { type: 'done', isFinished: true },
+        ]
+        assert.deepEqual(await reloadedStore.sendMessage(reloaded.chatKey, 'Continue'), { success: true })
+        assert.ok(nextHistory.some((message) => message.role === 'user' && message.content === 'Edited question'))
+        if (scenario.text) assert.equal(nextHistory.at(-1)?.content, scenario.text)
+        assert.equal(manager.getMessages(chat.id).at(-1)?.content, 'Continued answer')
+      })
+    }
+  } finally {
+    manager.close()
+    rmSync(dir, { recursive: true, force: true })
+  }
+})

tests/ai-agent-output-limit.test.mts

@@ -5,6 +5,7 @@ import { join } from 'node:path'
 import { test } from 'node:test'
 import * as nodeRuntime from '../packages/node-runtime/src/index'
 import { runServerAgent } from '../apps/cli/src/ai/agent'
+import { runChatTurn } from '../apps/cli/src/ai/chat-command'
 import { AIChatManager } from '../packages/node-runtime/src/ai/chats'
 import { runCrossChatAgent } from '../packages/node-runtime/src/ai/cross-chat-agent'
 import type { AgentStreamChunk } from '../packages/node-runtime/src/ai/agent/event-handler'
@@ -68,7 +69,7 @@ test('session and global agents surface truncated output instead of reporting su
       getDefaultAssistantConfig: () => config,
     } as LLMConfigStore)
 
-    for (const kind of ['session', 'global', 'desktop session'] as const) {
+    for (const kind of ['session', 'global', 'desktop session', 'CLI command'] as const) {
       for (const scenario of [
         { name: 'thinking only', text: '', finishReason: 'length' },
         { name: 'partial answer', text: 'The conclusion is', finishReason: 'length' },
@@ -117,7 +118,26 @@ test('session and global agents surface truncated output instead of reporting su
             aiChatManager: manager,
             onEvent: (event: AgentStreamChunk) => events.push(event),
           }
-          if (kind === 'global') {
+          if (kind === 'CLI command') {
+            await runChatTurn(
+              { aiChatId: chat.id, question: common.userMessage, json: true, locale: common.locale },
+              {
+                dbManager: { open: () => ({}) } as never,
+                pathProvider: {} as never,
+                aiChatManager: manager,
+                createRunAgentStream: () => async (_params, onEvent) => {
+                  await runServerAgent({
+                    ...common,
+                    llmConfig: config,
+                    onEvent: (event) => {
+                      common.onEvent(event)
+                      onEvent(event)
+                    },
+                  })
+                },
+              }
+            )
+          } else if (kind === 'global') {
             await runCrossChatAgent({
               ...common,
               piModel: buildPiModel(config),
@@ -146,7 +166,10 @@ test('session and global agents surface truncated output instead of reporting su
           const errors = events.filter((event) => event.type === 'error')
           const truncated = scenario.finishReason === 'length'
           assert.equal(errors.length, truncated ? 1 : 0)
-          if (truncated) assert.match((errors[0].error as { message: string }).message, /output limit/i)
+          if (truncated) {
+            assert.match((errors[0].error as { message: string }).message, /output limit/i)
+            assert.equal((errors[0].error as { name: string }).name, 'OutputLimitError')
+          }
           assert.deepEqual(
             events
               .filter((event) => event.type === 'status')
@@ -165,12 +188,22 @@ test('session and global agents surface truncated output instead of reporting su
             cacheWriteTokens: 0,
           }
           assert.deepEqual(doneUsage, expectedUsage)
-          manager.addMessagePair(
-            chat.id,
-            { content: common.userMessage },
-            { content: scenario.text, tokenUsage: doneUsage }
-          )
-          assert.deepEqual(manager.getAIChatTokenUsage(chat.id), expectedUsage)
+          if (kind === 'CLI command') {
+            const messages = manager.getMessages(chat.id)
+            assert.deepEqual(
+              messages.map((message) => message.content),
+              [common.userMessage, scenario.text]
+            )
+            assert.ok(messages[1].contentBlocks?.some((block) => block.type === 'think'))
+            assert.equal(
+              messages[1].contentBlocks?.some((block) => block.type === 'error'),
+              truncated
+            )
+            assert.deepEqual(manager.getAIChatTokenUsage(chat.id), expectedUsage)
+            if (scenario.text) {
+              assert.equal(manager.getHistoryForAgent(chat.id).at(-1)?.content, scenario.text)
+            }
+          }
         })
       }
     }