feat: Track and return token usage metadata, including total and model-specific counts, within prompt responses.

packages/a2a-server/src/agent/task.ts

@@ -348,7 +348,8 @@ export class Task {
     } else if (isSubagentProgress(outputChunk)) {
       outputAsText = JSON.stringify(outputChunk);
     } else if (Array.isArray(outputChunk)) {
-      const ansiOutput: AnsiOutput = outputChunk;
+      // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
+      const ansiOutput = outputChunk as AnsiOutput;
       outputAsText = ansiOutput
         .map((line: AnsiLine) =>
           line.map((token: AnsiToken) => token.text).join(''),

packages/cli/src/acp/acpClient.test.ts

@@ -495,7 +495,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 () => {
@@ -687,50 +687,110 @@ describe('Session', () => {
         content: { type: 'text', text: 'Hello' },
       },
     });
-    expect(result).toEqual({ stopReason: 'end_turn' });
+    expect(result).toMatchObject({ stopReason: 'end_turn' });
+  });
+
+  it('should accumulate token counts into _meta.quota.token_count and model_usage', async () => {
+    mockConfig.getModel.mockReturnValue('gemini-1.5-pro-test');
+    const stream = createMockStream([
+      {
+        type: StreamEventType.CHUNK,
+        value: {
+          candidates: [{ content: { parts: [{ text: 'Answer 1' }] } }],
+          usageMetadata: { promptTokenCount: 10, candidatesTokenCount: 20 },
+        },
+      },
+    ]);
+    mockChat.sendMessageStream.mockResolvedValue(stream);
+
+    const result = await session.prompt({
+      sessionId: 'session-1',
+      prompt: [{ type: 'text', text: 'Hi' }],
+    });
+
+    expect(result).toEqual({
+      stopReason: 'end_turn',
+      _meta: {
+        quota: {
+          token_count: { input_tokens: 10, output_tokens: 20 },
+          model_usage: [
+            {
+              model: 'gemini-1.5-pro-test',
+              token_count: { input_tokens: 10, output_tokens: 20 },
+            },
+          ],
+        },
+      },
+    });
   });
 
   it('should handle /memory command', async () => {
-    const handleCommandSpy = vi
-      .spyOn(
-        (session as unknown as { commandHandler: CommandHandler })
-          .commandHandler,
-        'handleCommand',
-      )
-      .mockResolvedValue(true);
+    vi.spyOn(
+      (session as unknown as { commandHandler: CommandHandler }).commandHandler,
+      'handleCommand',
+    ).mockResolvedValue(true);
 
     const result = await session.prompt({
       sessionId: 'session-1',
       prompt: [{ type: 'text', text: '/memory view' }],
     });
 
-    expect(result).toEqual({ stopReason: 'end_turn' });
-    expect(handleCommandSpy).toHaveBeenCalledWith(
-      '/memory view',
-      expect.any(Object),
+    expect(mockConnection.sessionUpdate).toHaveBeenCalledWith(
+      expect.objectContaining({
+        update: expect.objectContaining({
+          sessionUpdate: 'agent_message_chunk',
+        }),
+      }),
     );
+    expect(result).toEqual({ stopReason: 'end_turn' });
+    expect(mockChat.sendMessageStream).not.toHaveBeenCalled();
+  });
+  it('should handle /status command directly', async () => {
+    mockConfig.getActiveModel.mockReturnValue('gemini-1.5-pro-test');
+    const result = await session.prompt({
+      sessionId: 'session-1',
+      prompt: [{ type: 'text', text: '/status' }],
+    });
+
+    expect(mockConnection.sessionUpdate).toHaveBeenCalledWith(
+      expect.objectContaining({
+        update: expect.objectContaining({
+          sessionUpdate: 'agent_message_chunk',
+          content: expect.objectContaining({
+            type: 'text',
+            text: expect.stringContaining('gemini-1.5-pro-test'),
+          }),
+        }),
+      }),
+    );
+    expect(result).toEqual({ stopReason: 'end_turn' });
+    // Chat should not be called
     expect(mockChat.sendMessageStream).not.toHaveBeenCalled();
   });
 
   it('should handle /extensions command', async () => {
-    const handleCommandSpy = vi
-      .spyOn(
-        (session as unknown as { commandHandler: CommandHandler })
-          .commandHandler,
-        'handleCommand',
-      )
-      .mockResolvedValue(true);
+    vi.spyOn(
+      (session as unknown as { commandHandler: CommandHandler }).commandHandler,
+      'handleCommand',
+    ).mockResolvedValue(true);
 
     const result = await session.prompt({
       sessionId: 'session-1',
       prompt: [{ type: 'text', text: '/extensions list' }],
     });
 
-    expect(result).toEqual({ stopReason: 'end_turn' });
-    expect(handleCommandSpy).toHaveBeenCalledWith(
-      '/extensions list',
-      expect.any(Object),
+    expect(mockConnection.sessionUpdate).toHaveBeenCalledWith(
+      expect.objectContaining({
+        update: expect.objectContaining({
+          sessionUpdate: 'agent_message_chunk',
+          content: expect.objectContaining({
+            type: 'text',
+            text: expect.stringContaining('`test-mcp`'),
+          }),
+        }),
+      }),
     );
+    expect(result).toEqual({ stopReason: 'end_turn' });
     expect(mockChat.sendMessageStream).not.toHaveBeenCalled();
   });
 
@@ -843,7 +903,7 @@ describe('Session', () => {
         }),
       }),
     );
-    expect(result).toEqual({ stopReason: 'end_turn' });
+    expect(result).toMatchObject({ stopReason: 'end_turn' });
   });
 
   it('should handle tool call permission request', async () => {

packages/cli/src/acp/acpClient.ts

@@ -646,6 +646,10 @@ export class Session {
 
     let nextMessage: Content | null = { role: 'user', parts };
 
+    let totalInputTokens = 0;
+    let totalOutputTokens = 0;
+    const modelUsageMap = new Map<string, { input: number; output: number }>();
+
     while (nextMessage !== null) {
       if (pendingSend.signal.aborted) {
         chat.addHistory(nextMessage);
@@ -668,11 +672,20 @@ export class Session {
         );
         nextMessage = null;
 
+        let turnInputTokens = 0;
+        let turnOutputTokens = 0;
+
         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 ?? 0;
+            turnOutputTokens =
+              resp.value.usageMetadata.candidatesTokenCount ?? 0;
+          }
+
           if (
             resp.type === StreamEventType.CHUNK &&
             resp.value.candidates &&
@@ -704,6 +717,19 @@ export class Session {
           }
         }
 
+        totalInputTokens += turnInputTokens;
+        totalOutputTokens += turnOutputTokens;
+
+        if (turnInputTokens > 0 || turnOutputTokens > 0) {
+          const existingModelUsage = modelUsageMap.get(model) || {
+            input: 0,
+            output: 0,
+          };
+          existingModelUsage.input += turnInputTokens;
+          existingModelUsage.output += turnOutputTokens;
+          modelUsageMap.set(model, existingModelUsage);
+        }
+
         if (pendingSend.signal.aborted) {
           return { stopReason: CoreToolCallStatus.Cancelled };
         }
@@ -740,7 +766,28 @@ export class Session {
       }
     }
 
-    return { stopReason: 'end_turn' };
+    const modelUsageArray = Array.from(modelUsageMap.entries()).map(
+      ([m, usage]) => ({
+        model: m,
+        token_count: {
+          input_tokens: usage.input,
+          output_tokens: usage.output,
+        },
+      }),
+    );
+
+    return {
+      stopReason: 'end_turn',
+      _meta: {
+        quota: {
+          token_count: {
+            input_tokens: totalInputTokens,
+            output_tokens: totalOutputTokens,
+          },
+          model_usage: modelUsageArray,
+        },
+      },
+    };
   }
 
   private async handleCommand(