From 1d2eae586c628bd62c682061caf6b3b4142e08e0 Mon Sep 17 00:00:00 2001 From: digua Date: Fri, 18 Sep 2026 09:57:45 +0800 Subject: [PATCH] fix(electron): preserve usage on truncated AI responses Retain cumulative Agent usage from status events and emit it in the terminal done event after an error so truncated Desktop replies keep their token totals. Cover the real Desktop Agent and stream runner alongside CLI and global entries, including partial output, completion state, and persisted usage. --- apps/desktop/main/ai/agent-stream-runner.ts | 5 ++ tests/ai-agent-output-limit.test.mts | 83 ++++++++++++++++++++- 2 files changed, 86 insertions(+), 2 deletions(-) diff --git a/apps/desktop/main/ai/agent-stream-runner.ts b/apps/desktop/main/ai/agent-stream-runner.ts index e3838fd8..8896e898 100644 --- a/apps/desktop/main/ai/agent-stream-runner.ts +++ b/apps/desktop/main/ai/agent-stream-runner.ts @@ -353,9 +353,13 @@ export function createElectronRunAgentStream( skillCtx ) + let lastUsage: AgentStreamChunk['usage'] try { await agent.executeStream(userMessage, (chunk: AgentStreamChunk) => { if (abortSignal.aborted) return + if (chunk.type === 'status' && chunk.status) { + lastUsage = chunk.status.totalUsage + } onEvent(chunk as SharedAgentStreamChunk) }) } catch (error) { @@ -368,6 +372,7 @@ export function createElectronRunAgentStream( if (!serializedError.url && activeAIConfig.baseUrl) serializedError.url = activeAIConfig.baseUrl aiLogger.error('AgentStream', `Agent execution error: ${requestId}`, serializedError) onEvent({ type: 'error', error: serializedError, isFinished: true }) + onEvent({ type: 'done', isFinished: true, usage: lastUsage }) } } } diff --git a/tests/ai-agent-output-limit.test.mts b/tests/ai-agent-output-limit.test.mts index a15520db..e3dfc9c0 100644 --- a/tests/ai-agent-output-limit.test.mts +++ b/tests/ai-agent-output-limit.test.mts @@ -3,13 +3,16 @@ import { mkdtempSync, rmSync } from 'node:fs' import { tmpdir } from 'node:os' import { join } from 'node:path' import { test } from 'node:test' +import * as nodeRuntime from '@openchatlab/node-runtime' import { runServerAgent } from '../apps/cli/src/ai/agent' 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' import { buildPiModel } from '../packages/node-runtime/src/ai/llm-builder' +import { createAiTranslate } from '../packages/node-runtime/src/ai/i18n' +import type { LLMConfigStore } from '../packages/node-runtime/src/ai/llm-config-store' -// Exercise the real Pi SSE parser and both Node entry points, without a real provider or user data. +// Exercise the real Pi SSE parser and all Node entry points, without a real provider or user data. test('session and global agents surface truncated output instead of reporting successful completion', async (t) => { const dir = mkdtempSync(join(tmpdir(), 'chatlab-agent-output-limit-')) const manager = new AIChatManager(dir, { nativeBinding: process.env.CHATLAB_TEST_SQLITE_NATIVE_BINDING }) @@ -20,7 +23,52 @@ test('session and global agents surface truncated output instead of reporting su apiKey: 'test-placeholder', } try { - for (const kind of ['session', 'global'] as const) { + // Desktop sources load as CommonJS in tsx; keep the shared ESM runtime unchanged. + await t.mock.module('@openchatlab/node-runtime', { namedExports: nodeRuntime }) + // Isolate Desktop host services while keeping its Agent and stream runner real. + await t.mock.module('../apps/desktop/main/ai/chats', { + namedExports: { + getManager: () => manager, + getHistoryForAgent: manager.getHistoryForAgent.bind(manager), + setPendingDebugContext: manager.setPendingDebugContext.bind(manager), + }, + }) + await t.mock.module('../apps/desktop/main/ai/llm', { + namedExports: { buildPiModel, findModelDefinition: () => null, getProviderInfo: () => null }, + }) + await t.mock.module('../apps/desktop/main/ai/logger', { + namedExports: { + aiLogger: { + info: () => undefined, + warn: () => undefined, + error: () => undefined, + debug: () => undefined, + }, + isDebugMode: () => false, + }, + }) + await t.mock.module('../apps/desktop/main/i18n', { namedExports: { t: createAiTranslate('en-US') } }) + await t.mock.module('../apps/desktop/main/ai/tools', { + namedExports: { getAllTools: async () => [], createActivateSkillTool: () => undefined }, + }) + await t.mock.module('../apps/desktop/main/ai/assistant/manager', { + namedExports: { getAssistantConfig: () => null }, + }) + await t.mock.module('../apps/desktop/main/ai/skills/manager', { + namedExports: { getSkillConfig: () => null, getSkillMenu: () => '' }, + }) + await t.mock.module('../apps/desktop/main/worker/workerManager', { + namedExports: { getChatOverview: async () => null }, + }) + await t.mock.module('../apps/desktop/main/ai/cross-chat-tool-adapter', { + namedExports: { createElectronCrossChatTools: () => [] }, + }) + const { createElectronRunAgentStream } = await import('../apps/desktop/main/ai/agent-stream-runner') + const runDesktopAgent = createElectronRunAgentStream({ + getDefaultAssistantConfig: () => config, + } as LLMConfigStore) + + for (const kind of ['session', 'global', 'desktop session'] as const) { for (const scenario of [ { name: 'thinking only', text: '', finishReason: 'length' }, { name: 'partial answer', text: 'The conclusion is', finishReason: 'length' }, @@ -38,6 +86,16 @@ test('session and global agents surface truncated output instead of reporting su created: 1, model: config.model, choices: [{ index: 0, ...choice }], + ...(choice.finish_reason + ? { + usage: { + prompt_tokens: 100, + completion_tokens: 25, + total_tokens: 125, + prompt_tokens_details: { cached_tokens: 20 }, + }, + } + : {}), })) return new Response( chunks.map((chunk) => `data: ${JSON.stringify(chunk)}\n\n`).join('') + 'data: [DONE]\n\n', @@ -67,6 +125,12 @@ test('session and global agents surface truncated output instead of reporting su tools: [], memoryService: { list: () => [] }, }) + } else if (kind === 'desktop session') { + await runDesktopAgent( + { ...common, sessionId: 'session-1', chatType: 'private' }, + common.onEvent, + new AbortController().signal + ) } else { await runServerAgent({ ...common, llmConfig: config }) } @@ -92,6 +156,21 @@ test('session and global agents surface truncated output instead of reporting su ) assert.equal(events.filter((event) => event.type === 'done').length, 1) assert.equal(events.at(-1)?.type, 'done') + const doneUsage = events.at(-1)?.usage + const expectedUsage = { + promptTokens: 80, + completionTokens: 25, + totalTokens: 125, + cacheReadTokens: 20, + 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) }) } }