From 3a7b1159ae63b7916fae5e0340c1683ffa7bc1d6 Mon Sep 17 00:00:00 2001 From: tanzhenxin Date: Fri, 5 Dec 2025 15:40:49 +0800 Subject: [PATCH] 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'; } /**