From 3a7b1159ae63b7916fae5e0340c1683ffa7bc1d6 Mon Sep 17 00:00:00 2001 From: tanzhenxin Date: Fri, 5 Dec 2025 15:40:49 +0800 Subject: [PATCH 1/2] feat: add usage metadata in acp session/update event --- integration-tests/acp-integration.test.ts | 59 +++++++++++++++++++ packages/cli/src/acp-integration/schema.ts | 18 ++++++ .../src/acp-integration/service/filesystem.ts | 8 +++ .../session/HistoryReplayer.test.ts | 44 ++++++++++++++ .../session/HistoryReplayer.ts | 19 +++++- .../src/acp-integration/session/Session.ts | 12 +++- .../session/SubAgentTracker.test.ts | 2 +- .../session/emitters/MessageEmitter.test.ts | 28 +++++++++ .../session/emitters/MessageEmitter.ts | 30 +++++++++- .../session/emitters/ToolCallEmitter.test.ts | 10 ++-- .../session/emitters/ToolCallEmitter.ts | 11 +++- .../cli/src/acp-integration/session/types.ts | 2 + 12 files changed, 230 insertions(+), 13 deletions(-) diff --git a/integration-tests/acp-integration.test.ts b/integration-tests/acp-integration.test.ts index b098e025..5bffca90 100644 --- a/integration-tests/acp-integration.test.ts +++ b/integration-tests/acp-integration.test.ts @@ -25,6 +25,14 @@ type PendingRequest = { timeout: NodeJS.Timeout; }; +type UsageMetadata = { + promptTokens?: number | null; + completionTokens?: number | null; + thoughtsTokens?: number | null; + totalTokens?: number | null; + cachedTokens?: number | null; +}; + type SessionUpdateNotification = { sessionId?: string; update?: { @@ -39,6 +47,9 @@ type SessionUpdateNotification = { text?: string; }; modeId?: string; + _meta?: { + usage?: UsageMetadata; + }; }; }; @@ -587,4 +598,52 @@ function setupAcpTest( await cleanup(); } }); + + it('receives usage metadata in agent_message_chunk updates', async () => { + const rig = new TestRig(); + rig.setup('acp usage metadata'); + + const { sendRequest, cleanup, stderr, sessionUpdates } = setupAcpTest(rig); + + try { + await sendRequest('initialize', { + protocolVersion: 1, + clientCapabilities: { fs: { readTextFile: true, writeTextFile: true } }, + }); + await sendRequest('authenticate', { methodId: 'openai' }); + + const newSession = (await sendRequest('session/new', { + cwd: rig.testDir!, + mcpServers: [], + })) as { sessionId: string }; + + await sendRequest('session/prompt', { + sessionId: newSession.sessionId, + prompt: [{ type: 'text', text: 'Say "hello".' }], + }); + + await delay(500); + + // Find updates with usage metadata + const updatesWithUsage = sessionUpdates.filter( + (u) => + u.update?.sessionUpdate === 'agent_message_chunk' && + u.update?._meta?.usage, + ); + + expect(updatesWithUsage.length).toBeGreaterThan(0); + + const usage = updatesWithUsage[0].update?._meta?.usage; + expect(usage).toBeDefined(); + expect( + typeof usage?.promptTokens === 'number' || + typeof usage?.totalTokens === 'number', + ).toBe(true); + } catch (e) { + if (stderr.length) console.error('Agent stderr:', stderr.join('')); + throw e; + } finally { + await cleanup(); + } + }); }); diff --git a/packages/cli/src/acp-integration/schema.ts b/packages/cli/src/acp-integration/schema.ts index 8f21c74c..1a486818 100644 --- a/packages/cli/src/acp-integration/schema.ts +++ b/packages/cli/src/acp-integration/schema.ts @@ -316,6 +316,22 @@ export const annotationsSchema = z.object({ priority: z.number().optional().nullable(), }); +export const usageSchema = z.object({ + promptTokens: z.number().optional().nullable(), + completionTokens: z.number().optional().nullable(), + thoughtsTokens: z.number().optional().nullable(), + totalTokens: z.number().optional().nullable(), + cachedTokens: z.number().optional().nullable(), +}); + +export type Usage = z.infer; + +export const sessionUpdateMetaSchema = z.object({ + usage: usageSchema.optional().nullable(), +}); + +export type SessionUpdateMeta = z.infer; + export const requestPermissionResponseSchema = z.object({ outcome: requestPermissionOutcomeSchema, }); @@ -500,10 +516,12 @@ export const sessionUpdateSchema = z.union([ z.object({ content: contentBlockSchema, sessionUpdate: z.literal('agent_message_chunk'), + _meta: sessionUpdateMetaSchema.optional().nullable(), }), z.object({ content: contentBlockSchema, sessionUpdate: z.literal('agent_thought_chunk'), + _meta: sessionUpdateMetaSchema.optional().nullable(), }), z.object({ content: z.array(toolCallContentSchema).optional(), diff --git a/packages/cli/src/acp-integration/service/filesystem.ts b/packages/cli/src/acp-integration/service/filesystem.ts index c7db7235..af7c26ad 100644 --- a/packages/cli/src/acp-integration/service/filesystem.ts +++ b/packages/cli/src/acp-integration/service/filesystem.ts @@ -30,6 +30,14 @@ export class AcpFileSystemService implements FileSystemService { limit: null, }); + if (response.content.startsWith('ERROR: ENOENT:')) { + const err = new Error(response.content) as NodeJS.ErrnoException; + err.code = 'ENOENT'; + err.errno = -2; + err.path = filePath; + throw err; + } + return response.content; } diff --git a/packages/cli/src/acp-integration/session/HistoryReplayer.test.ts b/packages/cli/src/acp-integration/session/HistoryReplayer.test.ts index 83451592..c9cf65fb 100644 --- a/packages/cli/src/acp-integration/session/HistoryReplayer.test.ts +++ b/packages/cli/src/acp-integration/session/HistoryReplayer.test.ts @@ -411,4 +411,48 @@ describe('HistoryReplayer', () => { ]); }); }); + + describe('usage metadata replay', () => { + it('should emit usage metadata after assistant message content', async () => { + const record: ChatRecord = { + uuid: 'assistant-uuid', + parentUuid: 'user-uuid', + sessionId: 'test-session', + timestamp: new Date().toISOString(), + type: 'assistant', + cwd: '/test', + version: '1.0.0', + message: { + role: 'model', + parts: [{ text: 'Hello!' }], + }, + usageMetadata: { + promptTokenCount: 100, + candidatesTokenCount: 50, + totalTokenCount: 150, + }, + }; + + await replayer.replay([record]); + + expect(sendUpdateSpy).toHaveBeenCalledTimes(2); + expect(sendUpdateSpy).toHaveBeenNthCalledWith(1, { + sessionUpdate: 'agent_message_chunk', + content: { type: 'text', text: 'Hello!' }, + }); + expect(sendUpdateSpy).toHaveBeenNthCalledWith(2, { + sessionUpdate: 'agent_message_chunk', + content: { type: 'text', text: '' }, + _meta: { + usage: { + promptTokens: 100, + completionTokens: 50, + thoughtsTokens: undefined, + totalTokens: 150, + cachedTokens: undefined, + }, + }, + }); + }); + }); }); diff --git a/packages/cli/src/acp-integration/session/HistoryReplayer.ts b/packages/cli/src/acp-integration/session/HistoryReplayer.ts index 53a1ed8a..1bd11c79 100644 --- a/packages/cli/src/acp-integration/session/HistoryReplayer.ts +++ b/packages/cli/src/acp-integration/session/HistoryReplayer.ts @@ -5,7 +5,10 @@ */ import type { ChatRecord } from '@qwen-code/qwen-code-core'; -import type { Content } from '@google/genai'; +import type { + Content, + GenerateContentResponseUsageMetadata, +} from '@google/genai'; import type { SessionContext } from './types.js'; import { MessageEmitter } from './emitters/MessageEmitter.js'; import { ToolCallEmitter } from './emitters/ToolCallEmitter.js'; @@ -52,6 +55,9 @@ export class HistoryReplayer { if (record.message) { await this.replayContent(record.message, 'assistant'); } + if (record.usageMetadata) { + await this.replayUsageMetadata(record.usageMetadata); + } break; case 'tool_result': @@ -88,11 +94,22 @@ export class HistoryReplayer { toolName: functionName, callId, args: part.functionCall.args as Record, + status: 'in_progress', }); } } } + /** + * Replays usage metadata. + * @param usageMetadata - The usage metadata to replay + */ + private async replayUsageMetadata( + usageMetadata: GenerateContentResponseUsageMetadata, + ): Promise { + await this.messageEmitter.emitUsageMetadata(usageMetadata); + } + /** * Replays a tool result record. */ diff --git a/packages/cli/src/acp-integration/session/Session.ts b/packages/cli/src/acp-integration/session/Session.ts index b4d79433..a6466d1e 100644 --- a/packages/cli/src/acp-integration/session/Session.ts +++ b/packages/cli/src/acp-integration/session/Session.ts @@ -55,6 +55,7 @@ import type { SessionContext, ToolCallStartParams } from './types.js'; import { HistoryReplayer } from './HistoryReplayer.js'; import { ToolCallEmitter } from './emitters/ToolCallEmitter.js'; import { PlanEmitter } from './emitters/PlanEmitter.js'; +import { MessageEmitter } from './emitters/MessageEmitter.js'; import { SubAgentTracker } from './SubAgentTracker.js'; /** @@ -79,6 +80,7 @@ export class Session implements SessionContext { private readonly historyReplayer: HistoryReplayer; private readonly toolCallEmitter: ToolCallEmitter; private readonly planEmitter: PlanEmitter; + private readonly messageEmitter: MessageEmitter; // Implement SessionContext interface readonly sessionId: string; @@ -96,6 +98,7 @@ export class Session implements SessionContext { this.toolCallEmitter = new ToolCallEmitter(this); this.planEmitter = new PlanEmitter(this); this.historyReplayer = new HistoryReplayer(this); + this.messageEmitter = new MessageEmitter(this); } getId(): string { @@ -236,6 +239,10 @@ export class Session implements SessionContext { } } + if (resp.type === StreamEventType.CHUNK && resp.value.usageMetadata) { + this.messageEmitter.emitUsageMetadata(resp.value.usageMetadata); + } + if (resp.type === StreamEventType.CHUNK && resp.value.functionCalls) { functionCalls.push(...resp.value.functionCalls); } @@ -444,7 +451,9 @@ export class Session implements SessionContext { } const confirmationDetails = - await invocation.shouldConfirmExecute(abortSignal); + this.config.getApprovalMode() !== ApprovalMode.YOLO + ? await invocation.shouldConfirmExecute(abortSignal) + : false; if (confirmationDetails) { const content: acp.ToolCallContent[] = []; @@ -522,6 +531,7 @@ export class Session implements SessionContext { callId, toolName: fc.name, args, + status: 'in_progress', }; await this.toolCallEmitter.emitStart(startParams); } diff --git a/packages/cli/src/acp-integration/session/SubAgentTracker.test.ts b/packages/cli/src/acp-integration/session/SubAgentTracker.test.ts index 074c8162..f2bb7cc5 100644 --- a/packages/cli/src/acp-integration/session/SubAgentTracker.test.ts +++ b/packages/cli/src/acp-integration/session/SubAgentTracker.test.ts @@ -208,7 +208,7 @@ describe('SubAgentTracker', () => { expect.objectContaining({ sessionUpdate: 'tool_call', toolCallId: 'call-123', - status: 'in_progress', + status: 'pending', title: 'read_file', content: [], locations: [], diff --git a/packages/cli/src/acp-integration/session/emitters/MessageEmitter.test.ts b/packages/cli/src/acp-integration/session/emitters/MessageEmitter.test.ts index 52a41a48..7720d1f1 100644 --- a/packages/cli/src/acp-integration/session/emitters/MessageEmitter.test.ts +++ b/packages/cli/src/acp-integration/session/emitters/MessageEmitter.test.ts @@ -148,4 +148,32 @@ describe('MessageEmitter', () => { }); }); }); + + describe('emitUsageMetadata', () => { + it('should emit agent_message_chunk with _meta.usage containing token counts', async () => { + const usageMetadata = { + promptTokenCount: 100, + candidatesTokenCount: 50, + thoughtsTokenCount: 25, + totalTokenCount: 175, + cachedContentTokenCount: 10, + }; + + await emitter.emitUsageMetadata(usageMetadata); + + expect(sendUpdateSpy).toHaveBeenCalledWith({ + sessionUpdate: 'agent_message_chunk', + content: { type: 'text', text: '' }, + _meta: { + usage: { + promptTokens: 100, + completionTokens: 50, + thoughtsTokens: 25, + totalTokens: 175, + cachedTokens: 10, + }, + }, + }); + }); + }); }); diff --git a/packages/cli/src/acp-integration/session/emitters/MessageEmitter.ts b/packages/cli/src/acp-integration/session/emitters/MessageEmitter.ts index 9ac8943a..00770922 100644 --- a/packages/cli/src/acp-integration/session/emitters/MessageEmitter.ts +++ b/packages/cli/src/acp-integration/session/emitters/MessageEmitter.ts @@ -4,6 +4,8 @@ * SPDX-License-Identifier: Apache-2.0 */ +import type { GenerateContentResponseUsageMetadata } from '@google/genai'; +import type { Usage } from '../../schema.js'; import { BaseEmitter } from './BaseEmitter.js'; /** @@ -24,6 +26,16 @@ export class MessageEmitter extends BaseEmitter { }); } + /** + * Emits an agent thought chunk. + */ + async emitAgentThought(text: string): Promise { + await this.sendUpdate({ + sessionUpdate: 'agent_thought_chunk', + content: { type: 'text', text }, + }); + } + /** * Emits an agent message chunk. */ @@ -35,12 +47,24 @@ export class MessageEmitter extends BaseEmitter { } /** - * Emits an agent thought chunk. + * Emits usage metadata. */ - async emitAgentThought(text: string): Promise { + async emitUsageMetadata( + usageMetadata: GenerateContentResponseUsageMetadata, + text: string = '', + ): Promise { + const usage: Usage = { + promptTokens: usageMetadata.promptTokenCount, + completionTokens: usageMetadata.candidatesTokenCount, + thoughtsTokens: usageMetadata.thoughtsTokenCount, + totalTokens: usageMetadata.totalTokenCount, + cachedTokens: usageMetadata.cachedContentTokenCount, + }; + await this.sendUpdate({ - sessionUpdate: 'agent_thought_chunk', + sessionUpdate: 'agent_message_chunk', content: { type: 'text', text }, + _meta: { usage }, }); } diff --git a/packages/cli/src/acp-integration/session/emitters/ToolCallEmitter.test.ts b/packages/cli/src/acp-integration/session/emitters/ToolCallEmitter.test.ts index 52e13399..4616b859 100644 --- a/packages/cli/src/acp-integration/session/emitters/ToolCallEmitter.test.ts +++ b/packages/cli/src/acp-integration/session/emitters/ToolCallEmitter.test.ts @@ -71,7 +71,7 @@ describe('ToolCallEmitter', () => { expect(sendUpdateSpy).toHaveBeenCalledWith({ sessionUpdate: 'tool_call', toolCallId: 'call-123', - status: 'in_progress', + status: 'pending', title: 'unknown_tool', // Falls back to tool name content: [], locations: [], @@ -94,7 +94,7 @@ describe('ToolCallEmitter', () => { expect(sendUpdateSpy).toHaveBeenCalledWith({ sessionUpdate: 'tool_call', toolCallId: 'call-456', - status: 'in_progress', + status: 'pending', title: 'edit_file: Test tool description', content: [], locations: [{ path: '/test/file.ts', line: 10 }], @@ -144,7 +144,7 @@ describe('ToolCallEmitter', () => { expect(sendUpdateSpy).toHaveBeenCalledWith({ sessionUpdate: 'tool_call', toolCallId: 'call-fail', - status: 'in_progress', + status: 'pending', title: 'failing_tool', // Fallback to tool name content: [], locations: [], // Fallback to empty @@ -493,7 +493,7 @@ describe('ToolCallEmitter', () => { type: 'content', content: { type: 'text', - text: '{"output":"test output"}', + text: 'test output', }, }, ], @@ -650,7 +650,7 @@ describe('ToolCallEmitter', () => { content: [ { type: 'content', - content: { type: 'text', text: '{"output":"Function output"}' }, + content: { type: 'text', text: 'Function output' }, }, ], rawOutput: 'raw result', diff --git a/packages/cli/src/acp-integration/session/emitters/ToolCallEmitter.ts b/packages/cli/src/acp-integration/session/emitters/ToolCallEmitter.ts index 4c25570a..3ac514c0 100644 --- a/packages/cli/src/acp-integration/session/emitters/ToolCallEmitter.ts +++ b/packages/cli/src/acp-integration/session/emitters/ToolCallEmitter.ts @@ -59,7 +59,7 @@ export class ToolCallEmitter extends BaseEmitter { await this.sendUpdate({ sessionUpdate: 'tool_call', toolCallId: params.callId, - status: 'in_progress', + status: params.status || 'pending', title, content: [], locations, @@ -275,7 +275,14 @@ export class ToolCallEmitter extends BaseEmitter { // Handle functionResponse parts - stringify the response if ('functionResponse' in part && part.functionResponse) { try { - const responseText = JSON.stringify(part.functionResponse.response); + const resp = part.functionResponse.response as Record< + string, + unknown + >; + const responseText = + (resp['output'] as string) ?? + (resp['error'] as string) ?? + JSON.stringify(resp); result.push({ type: 'content', content: { type: 'text', text: responseText }, diff --git a/packages/cli/src/acp-integration/session/types.ts b/packages/cli/src/acp-integration/session/types.ts index 0c8f60a0..7812fb03 100644 --- a/packages/cli/src/acp-integration/session/types.ts +++ b/packages/cli/src/acp-integration/session/types.ts @@ -35,6 +35,8 @@ export interface ToolCallStartParams { callId: string; /** Arguments passed to the tool */ args?: Record; + /** Status of the tool call */ + status?: 'pending' | 'in_progress' | 'completed' | 'failed'; } /** From d7b946651633a44aa32a2d46ad2a215f4c08ca85 Mon Sep 17 00:00:00 2001 From: tanzhenxin Date: Tue, 9 Dec 2025 09:58:19 +0800 Subject: [PATCH 2/2] #1129, add usage update in ACP mode --- packages/cli/src/acp-integration/schema.ts | 1 + .../service/filesystem.test.ts | 59 +++++++++++++++++++ .../src/acp-integration/service/filesystem.ts | 8 ++- .../session/HistoryReplayer.ts | 50 +++++++++++++++- .../src/acp-integration/session/Session.ts | 36 +++++++---- .../session/SubAgentTracker.ts | 21 +++++++ .../session/emitters/MessageEmitter.test.ts | 27 +++++++++ .../session/emitters/MessageEmitter.ts | 6 +- .../session/emitters/ToolCallEmitter.ts | 10 +++- packages/core/src/subagents/index.ts | 1 + .../core/src/subagents/subagent-events.ts | 12 +++- .../src/subagents/subagent-statistics.test.ts | 13 +++- .../core/src/subagents/subagent-statistics.ts | 33 +++++++++-- packages/core/src/subagents/subagent.ts | 24 +++++++- packages/core/src/telemetry/uiTelemetry.ts | 6 ++ 15 files changed, 279 insertions(+), 28 deletions(-) create mode 100644 packages/cli/src/acp-integration/service/filesystem.test.ts diff --git a/packages/cli/src/acp-integration/schema.ts b/packages/cli/src/acp-integration/schema.ts index 1a486818..ac754318 100644 --- a/packages/cli/src/acp-integration/schema.ts +++ b/packages/cli/src/acp-integration/schema.ts @@ -328,6 +328,7 @@ export type Usage = z.infer; export const sessionUpdateMetaSchema = z.object({ usage: usageSchema.optional().nullable(), + durationMs: z.number().optional().nullable(), }); export type SessionUpdateMeta = z.infer; diff --git a/packages/cli/src/acp-integration/service/filesystem.test.ts b/packages/cli/src/acp-integration/service/filesystem.test.ts new file mode 100644 index 00000000..70ccfc2d --- /dev/null +++ b/packages/cli/src/acp-integration/service/filesystem.test.ts @@ -0,0 +1,59 @@ +/** + * @license + * Copyright 2025 Qwen + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, expect, it, vi } from 'vitest'; +import type { FileSystemService } from '@qwen-code/qwen-code-core'; +import { AcpFileSystemService } from './filesystem.js'; + +const createFallback = (): FileSystemService => ({ + readTextFile: vi.fn(), + writeTextFile: vi.fn(), + findFiles: vi.fn().mockReturnValue([]), +}); + +describe('AcpFileSystemService', () => { + describe('readTextFile ENOENT handling', () => { + it('parses path from ACP ENOENT message (quoted)', async () => { + const client = { + readTextFile: vi + .fn() + .mockResolvedValue({ content: 'ERROR: ENOENT: "/remote/file.txt"' }), + } as unknown as import('../acp.js').Client; + + const svc = new AcpFileSystemService( + client, + 'session-1', + { readTextFile: true, writeTextFile: true }, + createFallback(), + ); + + await expect(svc.readTextFile('/local/file.txt')).rejects.toMatchObject({ + code: 'ENOENT', + path: '/remote/file.txt', + }); + }); + + it('falls back to requested path when none provided', async () => { + const client = { + readTextFile: vi.fn().mockResolvedValue({ content: 'ERROR: ENOENT:' }), + } as unknown as import('../acp.js').Client; + + const svc = new AcpFileSystemService( + client, + 'session-2', + { readTextFile: true, writeTextFile: true }, + createFallback(), + ); + + await expect( + svc.readTextFile('/fallback/path.txt'), + ).rejects.toMatchObject({ + code: 'ENOENT', + path: '/fallback/path.txt', + }); + }); + }); +}); diff --git a/packages/cli/src/acp-integration/service/filesystem.ts b/packages/cli/src/acp-integration/service/filesystem.ts index af7c26ad..7bcaee2d 100644 --- a/packages/cli/src/acp-integration/service/filesystem.ts +++ b/packages/cli/src/acp-integration/service/filesystem.ts @@ -31,10 +31,16 @@ export class AcpFileSystemService implements FileSystemService { }); if (response.content.startsWith('ERROR: ENOENT:')) { + // Treat ACP error strings as structured ENOENT errors without + // assuming a specific platform format. + const match = /^ERROR:\s*ENOENT:\s*(?.*)$/i.exec(response.content); const err = new Error(response.content) as NodeJS.ErrnoException; err.code = 'ENOENT'; err.errno = -2; - err.path = filePath; + const rawPath = match?.groups?.['path']?.trim(); + err['path'] = rawPath + ? rawPath.replace(/^['"]|['"]$/g, '') || filePath + : filePath; throw err; } diff --git a/packages/cli/src/acp-integration/session/HistoryReplayer.ts b/packages/cli/src/acp-integration/session/HistoryReplayer.ts index 1bd11c79..0ecbccb9 100644 --- a/packages/cli/src/acp-integration/session/HistoryReplayer.ts +++ b/packages/cli/src/acp-integration/session/HistoryReplayer.ts @@ -4,7 +4,7 @@ * SPDX-License-Identifier: Apache-2.0 */ -import type { ChatRecord } from '@qwen-code/qwen-code-core'; +import type { ChatRecord, TaskResultDisplay } from '@qwen-code/qwen-code-core'; import type { Content, GenerateContentResponseUsageMetadata, @@ -135,6 +135,54 @@ export class HistoryReplayer { // Note: args aren't stored in tool_result records by default args: undefined, }); + + // Special handling: Task tool execution summary contains token usage + const { resultDisplay } = result ?? {}; + if ( + !!resultDisplay && + typeof resultDisplay === 'object' && + 'type' in resultDisplay && + (resultDisplay as { type?: unknown }).type === 'task_execution' + ) { + await this.emitTaskUsageFromResultDisplay( + resultDisplay as TaskResultDisplay, + ); + } + } + + /** + * Emits token usage from a TaskResultDisplay execution summary, if present. + */ + private async emitTaskUsageFromResultDisplay( + resultDisplay: TaskResultDisplay, + ): Promise { + const summary = resultDisplay.executionSummary; + if (!summary) { + return; + } + + const usageMetadata: GenerateContentResponseUsageMetadata = {}; + + if (Number.isFinite(summary.inputTokens)) { + usageMetadata.promptTokenCount = summary.inputTokens; + } + if (Number.isFinite(summary.outputTokens)) { + usageMetadata.candidatesTokenCount = summary.outputTokens; + } + if (Number.isFinite(summary.thoughtTokens)) { + usageMetadata.thoughtsTokenCount = summary.thoughtTokens; + } + if (Number.isFinite(summary.cachedTokens)) { + usageMetadata.cachedContentTokenCount = summary.cachedTokens; + } + if (Number.isFinite(summary.totalTokens)) { + usageMetadata.totalTokenCount = summary.totalTokens; + } + + // Only emit if we captured at least one token metric + if (Object.keys(usageMetadata).length > 0) { + await this.messageEmitter.emitUsageMetadata(usageMetadata); + } } /** diff --git a/packages/cli/src/acp-integration/session/Session.ts b/packages/cli/src/acp-integration/session/Session.ts index a6466d1e..1d90ed20 100644 --- a/packages/cli/src/acp-integration/session/Session.ts +++ b/packages/cli/src/acp-integration/session/Session.ts @@ -4,7 +4,12 @@ * SPDX-License-Identifier: Apache-2.0 */ -import type { Content, FunctionCall, Part } from '@google/genai'; +import type { + Content, + FunctionCall, + GenerateContentResponseUsageMetadata, + Part, +} from '@google/genai'; import type { Config, GeminiChat, @@ -195,6 +200,8 @@ export class Session implements SessionContext { } const functionCalls: FunctionCall[] = []; + let usageMetadata: GenerateContentResponseUsageMetadata | null = null; + const streamStartTime = Date.now(); try { const responseStream = await chat.sendMessageStream( @@ -225,22 +232,16 @@ export class Session implements SessionContext { continue; } - const content: acp.ContentBlock = { - type: 'text', - text: part.text, - }; - - this.sendUpdate({ - sessionUpdate: part.thought - ? 'agent_thought_chunk' - : 'agent_message_chunk', - content, - }); + this.messageEmitter.emitMessage( + part.text, + 'assistant', + part.thought, + ); } } if (resp.type === StreamEventType.CHUNK && resp.value.usageMetadata) { - this.messageEmitter.emitUsageMetadata(resp.value.usageMetadata); + usageMetadata = resp.value.usageMetadata; } if (resp.type === StreamEventType.CHUNK && resp.value.functionCalls) { @@ -258,6 +259,15 @@ export class Session implements SessionContext { throw error; } + if (usageMetadata) { + const durationMs = Date.now() - streamStartTime; + await this.messageEmitter.emitUsageMetadata( + usageMetadata, + '', + durationMs, + ); + } + if (functionCalls.length > 0) { const toolResponseParts: Part[] = []; diff --git a/packages/cli/src/acp-integration/session/SubAgentTracker.ts b/packages/cli/src/acp-integration/session/SubAgentTracker.ts index c6c83292..1e745b92 100644 --- a/packages/cli/src/acp-integration/session/SubAgentTracker.ts +++ b/packages/cli/src/acp-integration/session/SubAgentTracker.ts @@ -9,6 +9,7 @@ import type { SubAgentToolCallEvent, SubAgentToolResultEvent, SubAgentApprovalRequestEvent, + SubAgentUsageEvent, ToolCallConfirmationDetails, AnyDeclarativeTool, AnyToolInvocation, @@ -20,6 +21,7 @@ import { import { z } from 'zod'; import type { SessionContext } from './types.js'; import { ToolCallEmitter } from './emitters/ToolCallEmitter.js'; +import { MessageEmitter } from './emitters/MessageEmitter.js'; import type * as acp from '../acp.js'; /** @@ -62,6 +64,7 @@ const basicPermissionOptions: readonly PermissionOptionConfig[] = [ */ export class SubAgentTracker { private readonly toolCallEmitter: ToolCallEmitter; + private readonly messageEmitter: MessageEmitter; private readonly toolStates = new Map< string, { @@ -76,6 +79,7 @@ export class SubAgentTracker { private readonly client: acp.Client, ) { this.toolCallEmitter = new ToolCallEmitter(ctx); + this.messageEmitter = new MessageEmitter(ctx); } /** @@ -92,16 +96,19 @@ export class SubAgentTracker { const onToolCall = this.createToolCallHandler(abortSignal); const onToolResult = this.createToolResultHandler(abortSignal); const onApproval = this.createApprovalHandler(abortSignal); + const onUsageMetadata = this.createUsageMetadataHandler(abortSignal); eventEmitter.on(SubAgentEventType.TOOL_CALL, onToolCall); eventEmitter.on(SubAgentEventType.TOOL_RESULT, onToolResult); eventEmitter.on(SubAgentEventType.TOOL_WAITING_APPROVAL, onApproval); + eventEmitter.on(SubAgentEventType.USAGE_METADATA, onUsageMetadata); return [ () => { eventEmitter.off(SubAgentEventType.TOOL_CALL, onToolCall); eventEmitter.off(SubAgentEventType.TOOL_RESULT, onToolResult); eventEmitter.off(SubAgentEventType.TOOL_WAITING_APPROVAL, onApproval); + eventEmitter.off(SubAgentEventType.USAGE_METADATA, onUsageMetadata); // Clean up any remaining states this.toolStates.clear(); }, @@ -252,6 +259,20 @@ export class SubAgentTracker { }; } + /** + * Creates a handler for usage metadata events. + */ + private createUsageMetadataHandler( + abortSignal: AbortSignal, + ): (...args: unknown[]) => void { + return (...args: unknown[]) => { + const event = args[0] as SubAgentUsageEvent; + if (abortSignal.aborted) return; + + this.messageEmitter.emitUsageMetadata(event.usage, '', event.durationMs); + }; + } + /** * Converts confirmation details to permission options for the client. */ diff --git a/packages/cli/src/acp-integration/session/emitters/MessageEmitter.test.ts b/packages/cli/src/acp-integration/session/emitters/MessageEmitter.test.ts index 7720d1f1..d0b1ae87 100644 --- a/packages/cli/src/acp-integration/session/emitters/MessageEmitter.test.ts +++ b/packages/cli/src/acp-integration/session/emitters/MessageEmitter.test.ts @@ -175,5 +175,32 @@ describe('MessageEmitter', () => { }, }); }); + + it('should include durationMs in _meta when provided', async () => { + const usageMetadata = { + promptTokenCount: 10, + candidatesTokenCount: 5, + thoughtsTokenCount: 2, + totalTokenCount: 17, + cachedContentTokenCount: 1, + }; + + await emitter.emitUsageMetadata(usageMetadata, 'done', 1234); + + expect(sendUpdateSpy).toHaveBeenCalledWith({ + sessionUpdate: 'agent_message_chunk', + content: { type: 'text', text: 'done' }, + _meta: { + usage: { + promptTokens: 10, + completionTokens: 5, + thoughtsTokens: 2, + totalTokens: 17, + cachedTokens: 1, + }, + durationMs: 1234, + }, + }); + }); }); }); diff --git a/packages/cli/src/acp-integration/session/emitters/MessageEmitter.ts b/packages/cli/src/acp-integration/session/emitters/MessageEmitter.ts index 00770922..39cdf6a7 100644 --- a/packages/cli/src/acp-integration/session/emitters/MessageEmitter.ts +++ b/packages/cli/src/acp-integration/session/emitters/MessageEmitter.ts @@ -52,6 +52,7 @@ export class MessageEmitter extends BaseEmitter { async emitUsageMetadata( usageMetadata: GenerateContentResponseUsageMetadata, text: string = '', + durationMs?: number, ): Promise { const usage: Usage = { promptTokens: usageMetadata.promptTokenCount, @@ -61,10 +62,13 @@ export class MessageEmitter extends BaseEmitter { cachedTokens: usageMetadata.cachedContentTokenCount, }; + const meta = + typeof durationMs === 'number' ? { usage, durationMs } : { usage }; + await this.sendUpdate({ sessionUpdate: 'agent_message_chunk', content: { type: 'text', text }, - _meta: { usage }, + _meta: meta, }); } diff --git a/packages/cli/src/acp-integration/session/emitters/ToolCallEmitter.ts b/packages/cli/src/acp-integration/session/emitters/ToolCallEmitter.ts index 3ac514c0..9859ed78 100644 --- a/packages/cli/src/acp-integration/session/emitters/ToolCallEmitter.ts +++ b/packages/cli/src/acp-integration/session/emitters/ToolCallEmitter.ts @@ -279,10 +279,14 @@ export class ToolCallEmitter extends BaseEmitter { string, unknown >; + const outputField = resp['output']; + const errorField = resp['error']; const responseText = - (resp['output'] as string) ?? - (resp['error'] as string) ?? - JSON.stringify(resp); + typeof outputField === 'string' + ? outputField + : typeof errorField === 'string' + ? errorField + : JSON.stringify(resp); result.push({ type: 'content', content: { type: 'text', text: responseText }, diff --git a/packages/core/src/subagents/index.ts b/packages/core/src/subagents/index.ts index 5560b4fd..17c62a20 100644 --- a/packages/core/src/subagents/index.ts +++ b/packages/core/src/subagents/index.ts @@ -58,6 +58,7 @@ export type { SubAgentStartEvent, SubAgentRoundEvent, SubAgentStreamTextEvent, + SubAgentUsageEvent, SubAgentToolCallEvent, SubAgentToolResultEvent, SubAgentFinishEvent, diff --git a/packages/core/src/subagents/subagent-events.ts b/packages/core/src/subagents/subagent-events.ts index 3c93112d..1f793308 100644 --- a/packages/core/src/subagents/subagent-events.ts +++ b/packages/core/src/subagents/subagent-events.ts @@ -10,7 +10,7 @@ import type { ToolConfirmationOutcome, ToolResultDisplay, } from '../tools/tools.js'; -import type { Part } from '@google/genai'; +import type { Part, GenerateContentResponseUsageMetadata } from '@google/genai'; export type SubAgentEvent = | 'start' @@ -20,6 +20,7 @@ export type SubAgentEvent = | 'tool_call' | 'tool_result' | 'tool_waiting_approval' + | 'usage_metadata' | 'finish' | 'error'; @@ -31,6 +32,7 @@ export enum SubAgentEventType { TOOL_CALL = 'tool_call', TOOL_RESULT = 'tool_result', TOOL_WAITING_APPROVAL = 'tool_waiting_approval', + USAGE_METADATA = 'usage_metadata', FINISH = 'finish', ERROR = 'error', } @@ -57,6 +59,14 @@ export interface SubAgentStreamTextEvent { timestamp: number; } +export interface SubAgentUsageEvent { + subagentId: string; + round: number; + usage: GenerateContentResponseUsageMetadata; + durationMs?: number; + timestamp: number; +} + export interface SubAgentToolCallEvent { subagentId: string; round: number; diff --git a/packages/core/src/subagents/subagent-statistics.test.ts b/packages/core/src/subagents/subagent-statistics.test.ts index 5b4ae3c6..39ba70aa 100644 --- a/packages/core/src/subagents/subagent-statistics.test.ts +++ b/packages/core/src/subagents/subagent-statistics.test.ts @@ -50,6 +50,15 @@ describe('SubagentStatistics', () => { expect(summary.outputTokens).toBe(600); expect(summary.totalTokens).toBe(1800); }); + + it('should track thought and cached tokens', () => { + stats.recordTokens(100, 50, 10, 5); + + const summary = stats.getSummary(); + expect(summary.thoughtTokens).toBe(10); + expect(summary.cachedTokens).toBe(5); + expect(summary.totalTokens).toBe(165); // 100 + 50 + 10 + 5 + }); }); describe('tool usage statistics', () => { @@ -93,14 +102,14 @@ describe('SubagentStatistics', () => { stats.start(baseTime); stats.setRounds(2); stats.recordToolCall('file_read', true, 100); - stats.recordTokens(1000, 500); + stats.recordTokens(1000, 500, 20, 10); const result = stats.formatCompact('Test task', baseTime + 5000); expect(result).toContain('📋 Task Completed: Test task'); expect(result).toContain('🔧 Tool Usage: 1 calls, 100.0% success'); expect(result).toContain('⏱️ Duration: 5.0s | 🔁 Rounds: 2'); - expect(result).toContain('🔢 Tokens: 1,500 (in 1000, out 500)'); + expect(result).toContain('🔢 Tokens: 1,530 (in 1000, out 500)'); }); it('should handle zero tool calls', () => { diff --git a/packages/core/src/subagents/subagent-statistics.ts b/packages/core/src/subagents/subagent-statistics.ts index 3ef120c6..72308c63 100644 --- a/packages/core/src/subagents/subagent-statistics.ts +++ b/packages/core/src/subagents/subagent-statistics.ts @@ -23,6 +23,8 @@ export interface SubagentStatsSummary { successRate: number; inputTokens: number; outputTokens: number; + thoughtTokens: number; + cachedTokens: number; totalTokens: number; estimatedCost: number; toolUsage: ToolUsageStats[]; @@ -36,6 +38,8 @@ export class SubagentStatistics { private failedToolCalls = 0; private inputTokens = 0; private outputTokens = 0; + private thoughtTokens = 0; + private cachedTokens = 0; private toolUsage = new Map(); start(now = Date.now()) { @@ -74,9 +78,16 @@ export class SubagentStatistics { this.toolUsage.set(name, tu); } - recordTokens(input: number, output: number) { + recordTokens( + input: number, + output: number, + thought: number = 0, + cached: number = 0, + ) { this.inputTokens += Math.max(0, input || 0); this.outputTokens += Math.max(0, output || 0); + this.thoughtTokens += Math.max(0, thought || 0); + this.cachedTokens += Math.max(0, cached || 0); } getSummary(now = Date.now()): SubagentStatsSummary { @@ -86,7 +97,11 @@ export class SubagentStatistics { totalToolCalls > 0 ? (this.successfulToolCalls / totalToolCalls) * 100 : 0; - const totalTokens = this.inputTokens + this.outputTokens; + const totalTokens = + this.inputTokens + + this.outputTokens + + this.thoughtTokens + + this.cachedTokens; const estimatedCost = this.inputTokens * 3e-5 + this.outputTokens * 6e-5; return { rounds: this.rounds, @@ -97,6 +112,8 @@ export class SubagentStatistics { successRate, inputTokens: this.inputTokens, outputTokens: this.outputTokens, + thoughtTokens: this.thoughtTokens, + cachedTokens: this.cachedTokens, totalTokens, estimatedCost, toolUsage: Array.from(this.toolUsage.values()), @@ -116,8 +133,12 @@ export class SubagentStatistics { `⏱️ Duration: ${this.fmtDuration(stats.totalDurationMs)} | 🔁 Rounds: ${stats.rounds}`, ]; if (typeof stats.totalTokens === 'number') { + const parts = [ + `in ${stats.inputTokens ?? 0}`, + `out ${stats.outputTokens ?? 0}`, + ]; lines.push( - `🔢 Tokens: ${stats.totalTokens.toLocaleString()}${stats.inputTokens || stats.outputTokens ? ` (in ${stats.inputTokens ?? 0}, out ${stats.outputTokens ?? 0})` : ''}`, + `🔢 Tokens: ${stats.totalTokens.toLocaleString()}${parts.length ? ` (${parts.join(', ')})` : ''}`, ); } return lines.join('\n'); @@ -152,8 +173,12 @@ export class SubagentStatistics { `🔧 Tools: ${stats.totalToolCalls} calls, ${sr.toFixed(1)}% success (${stats.successfulToolCalls} ok, ${stats.failedToolCalls} failed)`, ); if (typeof stats.totalTokens === 'number') { + const parts = [ + `in ${stats.inputTokens ?? 0}`, + `out ${stats.outputTokens ?? 0}`, + ]; lines.push( - `🔢 Tokens: ${stats.totalTokens.toLocaleString()} (in ${stats.inputTokens ?? 0}, out ${stats.outputTokens ?? 0})`, + `🔢 Tokens: ${stats.totalTokens.toLocaleString()} (${parts.join(', ')})`, ); } if (stats.toolUsage && stats.toolUsage.length) { diff --git a/packages/core/src/subagents/subagent.ts b/packages/core/src/subagents/subagent.ts index 885e8ca6..39e43e54 100644 --- a/packages/core/src/subagents/subagent.ts +++ b/packages/core/src/subagents/subagent.ts @@ -41,6 +41,7 @@ import type { SubAgentToolResultEvent, SubAgentStreamTextEvent, SubAgentErrorEvent, + SubAgentUsageEvent, } from './subagent-events.js'; import { type SubAgentEventEmitter, @@ -369,6 +370,7 @@ export class SubAgentScope { }, }; + const roundStreamStart = Date.now(); const responseStream = await chat.sendMessageStream( this.modelConfig.model || this.runtimeContext.getModel() || @@ -439,10 +441,19 @@ export class SubAgentScope { if (lastUsage) { const inTok = Number(lastUsage.promptTokenCount || 0); const outTok = Number(lastUsage.candidatesTokenCount || 0); - if (isFinite(inTok) || isFinite(outTok)) { + const thoughtTok = Number(lastUsage.thoughtsTokenCount || 0); + const cachedTok = Number(lastUsage.cachedContentTokenCount || 0); + if ( + isFinite(inTok) || + isFinite(outTok) || + isFinite(thoughtTok) || + isFinite(cachedTok) + ) { this.stats.recordTokens( isFinite(inTok) ? inTok : 0, isFinite(outTok) ? outTok : 0, + isFinite(thoughtTok) ? thoughtTok : 0, + isFinite(cachedTok) ? cachedTok : 0, ); // mirror legacy fields for compatibility this.executionStats.inputTokens = @@ -453,11 +464,20 @@ export class SubAgentScope { (isFinite(outTok) ? outTok : 0); this.executionStats.totalTokens = (this.executionStats.inputTokens || 0) + - (this.executionStats.outputTokens || 0); + (this.executionStats.outputTokens || 0) + + (isFinite(thoughtTok) ? thoughtTok : 0) + + (isFinite(cachedTok) ? cachedTok : 0); this.executionStats.estimatedCost = (this.executionStats.inputTokens || 0) * 3e-5 + (this.executionStats.outputTokens || 0) * 6e-5; } + this.eventEmitter?.emit(SubAgentEventType.USAGE_METADATA, { + subagentId: this.subagentId, + round: turnCounter, + usage: lastUsage, + durationMs: Date.now() - roundStreamStart, + timestamp: Date.now(), + } as SubAgentUsageEvent); } if (functionCalls.length > 0) { diff --git a/packages/core/src/telemetry/uiTelemetry.ts b/packages/core/src/telemetry/uiTelemetry.ts index 9a257e5a..0f8f2146 100644 --- a/packages/core/src/telemetry/uiTelemetry.ts +++ b/packages/core/src/telemetry/uiTelemetry.ts @@ -23,6 +23,12 @@ export type UiEvent = | (ApiErrorEvent & { 'event.name': typeof EVENT_API_ERROR }) | (ToolCallEvent & { 'event.name': typeof EVENT_TOOL_CALL }); +export { + EVENT_API_ERROR, + EVENT_API_RESPONSE, + EVENT_TOOL_CALL, +} from './constants.js'; + export interface ToolCallStats { count: number; success: number;