From db4ce308522c4065e6760e7802181cba27265371 Mon Sep 17 00:00:00 2001 From: digua Date: Fri, 18 Sep 2026 10:39:59 +0800 Subject: [PATCH] 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 | 5 +- apps/cli/src/ai/chat-command.test.ts | 93 ++++++++++ apps/cli/src/ai/chat-command.ts | 21 ++- apps/desktop/main/ai/agent/index.ts | 1 + .../node-runtime/src/ai/cross-chat-agent.ts | 8 +- src/stores/aiChat.ts | 8 +- src/stores/aiChatTruncatedPersistence.test.ts | 161 ++++++++++++++++++ tests/ai-agent-output-limit.test.mts | 51 +++++- 8 files changed, 330 insertions(+), 18 deletions(-) create mode 100644 src/stores/aiChatTruncatedPersistence.test.ts diff --git a/apps/cli/src/ai/agent.ts b/apps/cli/src/ai/agent.ts index 6a1f3724..7f59eeb4 100644 --- a/apps/cli/src/ai/agent.ts +++ b/apps/cli/src/ai/agent.ts @@ -272,7 +272,10 @@ export async function runServerAgent(options: RunAgentOptions): Promise { 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, { diff --git a/apps/cli/src/ai/chat-command.test.ts b/apps/cli/src/ai/chat-command.test.ts index e79ebbb1..0f5a33f4 100644 --- a/apps/cli/src/ai/chat-command.test.ts +++ b/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() diff --git a/apps/cli/src/ai/chat-command.ts b/apps/cli/src/ai/chat-command.ts index 537a6d5f..4f176e14 100644 --- a/apps/cli/src/ai/chat-command.ts +++ b/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`) } diff --git a/apps/desktop/main/ai/agent/index.ts b/apps/desktop/main/ai/agent/index.ts index d3752a7f..0feb8838 100644 --- a/apps/desktop/main/ai/agent/index.ts +++ b/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 diff --git a/packages/node-runtime/src/ai/cross-chat-agent.ts b/packages/node-runtime/src/ai/cross-chat-agent.ts index 964596d6..783c9202 100644 --- a/packages/node-runtime/src/ai/cross-chat-agent.ts +++ b/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', { diff --git a/src/stores/aiChat.ts b/src/stores/aiChat.ts index 56d5f53b..08d88fba 100644 --- a/src/stores/aiChat.ts +++ b/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() diff --git a/src/stores/aiChatTruncatedPersistence.test.ts b/src/stores/aiChatTruncatedPersistence.test.ts new file mode 100644 index 00000000..4131bec5 --- /dev/null +++ b/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 = [] + 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) => manager.addMessagePair(...args), + replaceLatestMessageRound: async (...args: Parameters) => + 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 }) + } +}) diff --git a/tests/ai-agent-output-limit.test.mts b/tests/ai-agent-output-limit.test.mts index 18bc5f77..72b1470c 100644 --- a/tests/ai-agent-output-limit.test.mts +++ b/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) + } + } }) } }