summaryrefslogtreecommitdiff
path: root/packages/core/src
diff options
context:
space:
mode:
Diffstat (limited to 'packages/core/src')
-rw-r--r--packages/core/src/core/turn.test.ts7
-rw-r--r--packages/core/src/core/turn.ts6
2 files changed, 10 insertions, 3 deletions
diff --git a/packages/core/src/core/turn.test.ts b/packages/core/src/core/turn.test.ts
index 2217e5da..aeb30229 100644
--- a/packages/core/src/core/turn.test.ts
+++ b/packages/core/src/core/turn.test.ts
@@ -239,6 +239,8 @@ describe('Turn', () => {
candidates: [{ content: { parts: [{ text: 'First response' }] } }],
usageMetadata: mockMetadata1,
} as unknown as GenerateContentResponse;
+ // Add a small delay to ensure apiTimeMs is > 0
+ await new Promise((resolve) => setTimeout(resolve, 10));
yield {
functionCalls: [{ name: 'aTool' }],
usageMetadata: mockMetadata2,
@@ -262,7 +264,10 @@ describe('Turn', () => {
expect(metadataEvent.type).toBe(GeminiEventType.UsageMetadata);
// The value should be the *last* metadata object received.
- expect(metadataEvent.value).toEqual(mockMetadata2);
+ expect(metadataEvent.value).toEqual(
+ expect.objectContaining(mockMetadata2),
+ );
+ expect(metadataEvent.value.apiTimeMs).toBeGreaterThan(0);
// Also check the public getter
expect(turn.getUsageMetadata()).toEqual(mockMetadata2);
diff --git a/packages/core/src/core/turn.ts b/packages/core/src/core/turn.ts
index 34e4a494..71c02d83 100644
--- a/packages/core/src/core/turn.ts
+++ b/packages/core/src/core/turn.ts
@@ -104,7 +104,7 @@ export type ServerGeminiChatCompressedEvent = {
export type ServerGeminiUsageMetadataEvent = {
type: GeminiEventType.UsageMetadata;
- value: GenerateContentResponseUsageMetadata;
+ value: GenerateContentResponseUsageMetadata & { apiTimeMs?: number };
};
// The original union type, now composed of the individual types
@@ -137,6 +137,7 @@ export class Turn {
req: PartListUnion,
signal: AbortSignal,
): AsyncGenerator<ServerGeminiStreamEvent> {
+ const startTime = Date.now();
try {
const responseStream = await this.chat.sendMessageStream({
message: req,
@@ -174,9 +175,10 @@ export class Turn {
}
if (this.lastUsageMetadata) {
+ const durationMs = Date.now() - startTime;
yield {
type: GeminiEventType.UsageMetadata,
- value: this.lastUsageMetadata,
+ value: { ...this.lastUsageMetadata, apiTimeMs: durationMs },
};
}
} catch (error) {