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.
This commit is contained in:
digua committed 2026-09-18 10:53:06 +08:00
1 parent a2c4176f89
commit db4ce30852
8 files changed
+330 -18

No files matched your search

+4 -1
View File
@@ -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, {
+93
View File
@@ -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()
+17 -4
View File
@@ -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`)
}