Feat/a2a expose usage metadata

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 = {

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: