diff --git a/packages/a2a-server/src/agent/task.test.ts b/packages/a2a-server/src/agent/task.test.ts index 2a3a20ebf4..1d0e010c2a 100644 --- a/packages/a2a-server/src/agent/task.test.ts +++ b/packages/a2a-server/src/agent/task.test.ts @@ -294,6 +294,66 @@ describe('Task', () => { ]); }); + it('should capture usageMetadata on Finished event and include it in final status update', async () => { + const mockConfig = createMockConfig(); + const mockEventBus: ExecutionEventBus = { + publish: vi.fn(), + on: vi.fn(), + off: vi.fn(), + once: vi.fn(), + removeAllListeners: vi.fn(), + finished: vi.fn(), + }; + + // @ts-expect-error - Calling private constructor for test purposes. + const task = new Task( + 'task-id', + 'context-id', + mockConfig as Config, + mockEventBus, + ); + + const finishedEvent = { + type: GeminiEventType.Finished, + value: { + reason: 'STOP', + usageMetadata: { + promptTokenCount: 100, + candidatesTokenCount: 50, + totalTokenCount: 150, + }, + }, + }; + + await task.acceptAgentMessage(finishedEvent); + expect(task.usageMetadata).toEqual({ + promptTokenCount: 100, + candidatesTokenCount: 50, + totalTokenCount: 150, + }); + + task.setTaskStateAndPublishUpdate( + 'input-required', + { kind: CoderAgentEvent.StateChangeEvent }, + undefined, + undefined, + true, // final + ); + + expect(mockEventBus.publish).toHaveBeenCalledWith( + expect.objectContaining({ + final: true, + metadata: expect.objectContaining({ + usageMetadata: { + promptTokenCount: 100, + candidatesTokenCount: 50, + totalTokenCount: 150, + }, + }), + }), + ); + }); + it('should update modelInfo and reflect it in metadata and status updates', async () => { const mockConfig = createMockConfig(); const mockEventBus: ExecutionEventBus = { diff --git a/packages/a2a-server/src/agent/task.ts b/packages/a2a-server/src/agent/task.ts index 6ecea06c60..ee13adc8cb 100644 --- a/packages/a2a-server/src/agent/task.ts +++ b/packages/a2a-server/src/agent/task.ts @@ -89,6 +89,12 @@ export class Task { currentAgentMessageId = uuidv4(); promptCount = 0; autoExecute: boolean; + usageMetadata?: { + promptTokenCount?: number; + candidatesTokenCount?: number; + totalTokenCount?: number; + cachedContentTokenCount?: number; + }; private get isYoloMatch(): boolean { return ( this.autoExecute || this.config.getApprovalMode() === ApprovalMode.YOLO @@ -274,6 +280,7 @@ export class Task { userTier?: UserTierId; error?: string; traceId?: string; + usageMetadata?: Task['usageMetadata']; } = { coderAgent: coderAgentMessage, model: this.modelInfo || this.config.getModel(), @@ -288,6 +295,10 @@ export class Task { metadata.traceId = traceId; } + if (final && this.usageMetadata) { + metadata.usageMetadata = this.usageMetadata; + } + return { kind: 'status-update', taskId: this.id, @@ -857,8 +868,18 @@ export class Task { break; case GeminiEventType.Finished: logger.info(`[Task ${this.id}] Agent finished its turn.`); + // Capture the usage metadata when the stream finishes + if ( + event.value && + typeof event.value === 'object' && + 'usageMetadata' in event.value + ) { + this.usageMetadata = event.value + .usageMetadata as typeof this.usageMetadata; + } break; case GeminiEventType.ModelInfo: + this.usageMetadata = undefined; this.modelInfo = event.value; break; case GeminiEventType.Retry: