From 08063d7b0a75ab23716c7631b78fa0025501ac42 Mon Sep 17 00:00:00 2001
From: Sri Pasumarthi <111310667+sripasg@users.noreply.github.com>
Date: Thu, 19 Mar 2026 14:49:50 -0700
Subject: [PATCH] feat: ACP: Add token usage metadata to the `send` method's
 return value (#23148)

---
 packages/cli/src/acp/acpClient.test.ts | 16 +++----
 packages/cli/src/acp/acpClient.ts      | 64 +++++++++++++++++++++++++-
 2 files changed, 70 insertions(+), 10 deletions(-)

diff --git a/packages/cli/src/acp/acpClient.test.ts b/packages/cli/src/acp/acpClient.test.ts
index ca525182b5..0f9c4a8e5b 100644
--- a/packages/cli/src/acp/acpClient.test.ts
+++ b/packages/cli/src/acp/acpClient.test.ts
@@ -551,7 +551,7 @@ describe('GeminiAgent', () => {
     });
 
     expect(session.prompt).toHaveBeenCalled();
-    expect(result).toEqual({ stopReason: 'end_turn' });
+    expect(result).toMatchObject({ stopReason: 'end_turn' });
   });
 
   it('should delegate setMode to session', async () => {
@@ -750,7 +750,7 @@ describe('Session', () => {
         content: { type: 'text', text: 'Hello' },
       },
     });
-    expect(result).toEqual({ stopReason: 'end_turn' });
+    expect(result).toMatchObject({ stopReason: 'end_turn' });
   });
 
   it('should handle /memory command', async () => {
@@ -767,7 +767,7 @@ describe('Session', () => {
       prompt: [{ type: 'text', text: '/memory view' }],
     });
 
-    expect(result).toEqual({ stopReason: 'end_turn' });
+    expect(result).toMatchObject({ stopReason: 'end_turn' });
     expect(handleCommandSpy).toHaveBeenCalledWith(
       '/memory view',
       expect.any(Object),
@@ -789,7 +789,7 @@ describe('Session', () => {
       prompt: [{ type: 'text', text: '/extensions list' }],
     });
 
-    expect(result).toEqual({ stopReason: 'end_turn' });
+    expect(result).toMatchObject({ stopReason: 'end_turn' });
     expect(handleCommandSpy).toHaveBeenCalledWith(
       '/extensions list',
       expect.any(Object),
@@ -811,7 +811,7 @@ describe('Session', () => {
       prompt: [{ type: 'text', text: '/extensions explore' }],
     });
 
-    expect(result).toEqual({ stopReason: 'end_turn' });
+    expect(result).toMatchObject({ stopReason: 'end_turn' });
     expect(handleCommandSpy).toHaveBeenCalledWith(
       '/extensions explore',
       expect.any(Object),
@@ -833,7 +833,7 @@ describe('Session', () => {
       prompt: [{ type: 'text', text: '/restore' }],
     });
 
-    expect(result).toEqual({ stopReason: 'end_turn' });
+    expect(result).toMatchObject({ stopReason: 'end_turn' });
     expect(handleCommandSpy).toHaveBeenCalledWith(
       '/restore',
       expect.any(Object),
@@ -855,7 +855,7 @@ describe('Session', () => {
       prompt: [{ type: 'text', text: '/init' }],
     });
 
-    expect(result).toEqual({ stopReason: 'end_turn' });
+    expect(result).toMatchObject({ stopReason: 'end_turn' });
     expect(handleCommandSpy).toHaveBeenCalledWith('/init', expect.any(Object));
     expect(mockChat.sendMessageStream).not.toHaveBeenCalled();
   });
@@ -909,7 +909,7 @@ describe('Session', () => {
         }),
       }),
     );
-    expect(result).toEqual({ stopReason: 'end_turn' });
+    expect(result).toMatchObject({ stopReason: 'end_turn' });
   });
 
   it('should handle tool call permission request', async () => {
diff --git a/packages/cli/src/acp/acpClient.ts b/packages/cli/src/acp/acpClient.ts
index bd5a52f126..5e3f3666b1 100644
--- a/packages/cli/src/acp/acpClient.ts
+++ b/packages/cli/src/acp/acpClient.ts
@@ -699,10 +699,22 @@ export class Session {
       // It uses `parts` argument but effectively ignores it in current implementation
       const handled = await this.handleCommand(commandText, parts);
       if (handled) {
-        return { stopReason: 'end_turn' };
+        return {
+          stopReason: 'end_turn',
+          _meta: {
+            quota: {
+              token_count: { input_tokens: 0, output_tokens: 0 },
+              model_usage: [],
+            },
+          },
+        };
       }
     }
 
+    let totalInputTokens = 0;
+    let totalOutputTokens = 0;
+    const modelUsageMap = new Map<string, { input: number; output: number }>();
+
     let nextMessage: Content | null = { role: 'user', parts };
 
     while (nextMessage !== null) {
@@ -727,11 +739,25 @@ export class Session {
         );
         nextMessage = null;
 
+        let turnInputTokens = 0;
+        let turnOutputTokens = 0;
+        let turnModelId = model;
+
         for await (const resp of responseStream) {
           if (pendingSend.signal.aborted) {
             return { stopReason: CoreToolCallStatus.Cancelled };
           }
 
+          if (resp.type === StreamEventType.CHUNK && resp.value.usageMetadata) {
+            turnInputTokens =
+              resp.value.usageMetadata.promptTokenCount ?? turnInputTokens;
+            turnOutputTokens =
+              resp.value.usageMetadata.candidatesTokenCount ?? turnOutputTokens;
+            if (resp.value.modelVersion) {
+              turnModelId = resp.value.modelVersion;
+            }
+          }
+
           if (
             resp.type === StreamEventType.CHUNK &&
             resp.value.candidates &&
@@ -763,6 +789,19 @@ export class Session {
           }
         }
 
+        totalInputTokens += turnInputTokens;
+        totalOutputTokens += turnOutputTokens;
+
+        if (turnInputTokens > 0 || turnOutputTokens > 0) {
+          const existing = modelUsageMap.get(turnModelId) ?? {
+            input: 0,
+            output: 0,
+          };
+          existing.input += turnInputTokens;
+          existing.output += turnOutputTokens;
+          modelUsageMap.set(turnModelId, existing);
+        }
+
         if (pendingSend.signal.aborted) {
           return { stopReason: CoreToolCallStatus.Cancelled };
         }
@@ -799,7 +838,28 @@ export class Session {
       }
     }
 
-    return { stopReason: 'end_turn' };
+    const modelUsageArray = Array.from(modelUsageMap.entries()).map(
+      ([modelName, counts]) => ({
+        model: modelName,
+        token_count: {
+          input_tokens: counts.input,
+          output_tokens: counts.output,
+        },
+      }),
+    );
+
+    return {
+      stopReason: 'end_turn',
+      _meta: {
+        quota: {
+          token_count: {
+            input_tokens: totalInputTokens,
+            output_tokens: totalOutputTokens,
+          },
+          model_usage: modelUsageArray,
+        },
+      },
+    };
   }
 
   private async handleCommand(