From c0455b7351a0f9b731517f1fede7365af0e2ea6d Mon Sep 17 00:00:00 2001 From: Sri Pasumarthi Date: Thu, 5 Mar 2026 13:38:15 -0800 Subject: [PATCH] feat: Track and return token usage metadata, including total and model-specific counts, within prompt responses. --- packages/a2a-server/src/agent/task.ts | 3 +- packages/cli/src/acp/acpClient.test.ts | 110 +++++++++++++++++++------ packages/cli/src/acp/acpClient.ts | 49 ++++++++++- 3 files changed, 135 insertions(+), 27 deletions(-) diff --git a/packages/a2a-server/src/agent/task.ts b/packages/a2a-server/src/agent/task.ts index fe15aed37b..562f3ead97 100644 --- a/packages/a2a-server/src/agent/task.ts +++ b/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(''), diff --git a/packages/cli/src/acp/acpClient.test.ts b/packages/cli/src/acp/acpClient.test.ts index 399b365c46..2682c2993e 100644 --- a/packages/cli/src/acp/acpClient.test.ts +++ b/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 () => { diff --git a/packages/cli/src/acp/acpClient.ts b/packages/cli/src/acp/acpClient.ts index a4afb1ade2..fe3531720c 100644 --- a/packages/cli/src/acp/acpClient.ts +++ b/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(); + 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(