Merge pull request #1176 from QwenLM/feat/acp-usage-metadata

Feat/acp usage metadata

Author: tanzhenxin

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();
+    }
+  });
 });

packages/cli/src/acp-integration/schema.ts

@@ -316,6 +316,23 @@ 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(),
+  durationMs: z.number().optional().nullable(),
+});
+
+export type SessionUpdateMeta = z.infer<typeof sessionUpdateMetaSchema>;
+
 export const requestPermissionResponseSchema = z.object({
   outcome: requestPermissionOutcomeSchema,
 });
@@ -500,10 +517,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(),

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',
+      });
+    });
+  });
+});

packages/cli/src/acp-integration/service/filesystem.ts

@@ -30,6 +30,20 @@ export class AcpFileSystemService implements FileSystemService {
       limit: null,
     });
 
+    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*(?<path>.*)$/i.exec(response.content);
+      const err = new Error(response.content) as NodeJS.ErrnoException;
+      err.code = 'ENOENT';
+      err.errno = -2;
+      const rawPath = match?.groups?.['path']?.trim();
+      err['path'] = rawPath
+        ? rawPath.replace(/^['"]|['"]$/g, '') || filePath
+        : filePath;
+      throw err;
+    }
+
     return response.content;
   }
 

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,
+          },
+        },
+      });
+    });
+  });
 });

packages/cli/src/acp-integration/session/HistoryReplayer.ts

@@ -4,8 +4,11 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import type { ChatRecord } from '@qwen-code/qwen-code-core';
-import type { Content } from '@google/genai';
+import type { ChatRecord, TaskResultDisplay } from '@qwen-code/qwen-code-core';
+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.
    */
@@ -118,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<void> {
+    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);
+    }
   }
 
   /**