Feat/a2a expose usage metadata (#27288)

This commit is contained in:
jvargassanchez-dot
2026-05-22 11:51:39 -06:00
committed by GitHub
parent ba04e99bea
commit d1fa323cfb
2 changed files with 81 additions and 0 deletions
@@ -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 = {