Bug fix telemetry token count (#1250)

Co-authored-by: N. Taylor Mullen <[email protected]>

1 files changed, 74 insertions, 0 deletions

diff --git a/packages/core/src/telemetry/loggers.test.ts b/packages/core/src/telemetry/loggers.test.ts
index 5b1cb707..3f61530f 100644
--- a/packages/core/src/telemetry/loggers.test.ts
+++ b/packages/core/src/telemetry/loggers.test.ts
@@ -20,10 +20,12 @@ import {
   logUserPrompt,
   logToolCall,
   ToolCallDecision,
+  getFinalUsageMetadata,
 } from './loggers.js';
 import * as metrics from './metrics.js';
 import * as sdk from './sdk.js';
 import { vi, describe, beforeEach, it, expect } from 'vitest';
+import { GenerateContentResponse } from '@google/genai';
 
 vi.mock('@gemini-cli/cli/dist/src/utils/version', () => ({
   getCliVersion: () => 'test-version',
@@ -520,3 +522,75 @@ describe('loggers', () => {
     });
   });
 });
+
+describe('getFinalUsageMetadata', () => {
+  const createMockResponse = (
+    usageMetadata?: GenerateContentResponse['usageMetadata'],
+  ): GenerateContentResponse =>
+    ({
+      text: () => '',
+      data: () => ({}) as Record<string, unknown>,
+      functionCalls: () => [],
+      executableCode: () => [],
+      codeExecutionResult: () => [],
+      usageMetadata,
+    }) as unknown as GenerateContentResponse;
+
+  it('should return the usageMetadata from the last chunk that has it', () => {
+    const chunks: GenerateContentResponse[] = [
+      createMockResponse({
+        promptTokenCount: 10,
+        candidatesTokenCount: 20,
+        totalTokenCount: 30,
+      }),
+      createMockResponse(),
+      createMockResponse({
+        promptTokenCount: 15,
+        candidatesTokenCount: 25,
+        totalTokenCount: 40,
+      }),
+      createMockResponse(),
+    ];
+
+    const result = getFinalUsageMetadata(chunks);
+    expect(result).toEqual({
+      promptTokenCount: 15,
+      candidatesTokenCount: 25,
+      totalTokenCount: 40,
+    });
+  });
+
+  it('should return undefined if no chunks have usageMetadata', () => {
+    const chunks: GenerateContentResponse[] = [
+      createMockResponse(),
+      createMockResponse(),
+      createMockResponse(),
+    ];
+
+    const result = getFinalUsageMetadata(chunks);
+    expect(result).toBeUndefined();
+  });
+
+  it('should return the metadata from the only chunk if it has it', () => {
+    const chunks: GenerateContentResponse[] = [
+      createMockResponse({
+        promptTokenCount: 1,
+        candidatesTokenCount: 2,
+        totalTokenCount: 3,
+      }),
+    ];
+
+    const result = getFinalUsageMetadata(chunks);
+    expect(result).toEqual({
+      promptTokenCount: 1,
+      candidatesTokenCount: 2,
+      totalTokenCount: 3,
+    });
+  });
+
+  it('should return undefined for an empty array of chunks', () => {
+    const chunks: GenerateContentResponse[] = [];
+    const result = getFinalUsageMetadata(chunks);
+    expect(result).toBeUndefined();
+  });
+});