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) + } + } }) } }