feat: add usage metadata in acp session/update event

2025-12-19 09:33:53 +00:00 · 2025-12-05 15:40:49 +08:00
parent 3e2a2255ee
commit 3a7b1159ae
12 changed files with 230 additions and 13 deletions
--- 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();
+    }
+  });
 });
--- 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<typeof usageSchema>;
+
+export const sessionUpdateMetaSchema = z.object({
+  usage: usageSchema.optional().nullable(),
+});
+
+export type SessionUpdateMeta = z.infer<typeof sessionUpdateMetaSchema>;
+
 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(),
--- 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;
  }

--- 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,
+          },
+        },
+      });
+    });
+  });
 });
--- 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<string, unknown>,
+          status: 'in_progress',
        });
      }
    }
  }

+  /**
+   * Replays usage metadata.
+   * @param usageMetadata - The usage metadata to replay
+   */
+  private async replayUsageMetadata(
+    usageMetadata: GenerateContentResponseUsageMetadata,
+  ): Promise<void> {
+    await this.messageEmitter.emitUsageMetadata(usageMetadata);
+  }
+
  /**
   * Replays a tool result record.
   */
--- 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);
      }
--- 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: [],
--- 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,
+          },
+        },
+      });
+    });
+  });
 });
--- 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<void> {
+    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<void> {
+  async emitUsageMetadata(
+    usageMetadata: GenerateContentResponseUsageMetadata,
+    text: string = '',
+  ): Promise<void> {
+    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 },
    });
  }

--- 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',
--- 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 },
--- 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<string, unknown>;
+  /** Status of the tool call */
+  status?: 'pending' | 'in_progress' | 'completed' | 'failed';
 }

 /**