update token usage with types

7 files changed, 120 insertions, 15 deletions

diff --git a/docs/core/telemetry.md b/docs/core/telemetry.md
index ed1dbeaf..bd448776 100644
--- a/docs/core/telemetry.md
+++ b/docs/core/telemetry.md
@@ -338,6 +338,7 @@ These are numerical measurements of behavior over time.
   - **Attributes**:
     - `model`
 
-- `gemini_cli.token.input.count` (Counter, Int): Counts the total number of input tokens sent to the API.
+- `gemini_cli.token.usage` (Counter, Int): Counts the number of tokens used.
   - **Attributes**:
     - `model`
+    - `type` (string: "input", "output", "thought", "cache", or "tool")
diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts
index d69d571c..349caf1f 100644
--- a/packages/core/src/core/client.ts
+++ b/packages/core/src/core/client.ts
@@ -216,7 +216,7 @@ export class GeminiClient {
   private _logApiRequest(model: string, inputTokenCount: number): void {
     logApiRequest({
       model,
-      prompt_token_count: inputTokenCount,
+      input_token_count: inputTokenCount,
       duration_ms: 0, // Duration is not known at request time
     });
   }
@@ -245,6 +245,11 @@ export class GeminiClient {
       attempt,
       status_code: undefined,
       error: responseError,
+      output_token_count: response.usageMetadata?.candidatesTokenCount ?? 0,
+      cached_content_token_count:
+        response.usageMetadata?.cachedContentTokenCount ?? 0,
+      thoughts_token_count: response.usageMetadata?.thoughtsTokenCount ?? 0,
+      tool_token_count: response.usageMetadata?.toolUsePromptTokenCount ?? 0,
     });
   }
 
diff --git a/packages/core/src/telemetry/constants.ts b/packages/core/src/telemetry/constants.ts
index 7b4aad49..bbcdbada 100644
--- a/packages/core/src/telemetry/constants.ts
+++ b/packages/core/src/telemetry/constants.ts
@@ -17,5 +17,5 @@ export const METRIC_TOOL_CALL_COUNT = 'gemini_cli.tool.call.count';
 export const METRIC_TOOL_CALL_LATENCY = 'gemini_cli.tool.call.latency';
 export const METRIC_API_REQUEST_COUNT = 'gemini_cli.api.request.count';
 export const METRIC_API_REQUEST_LATENCY = 'gemini_cli.api.request.latency';
-export const METRIC_TOKEN_INPUT_COUNT = 'gemini_cli.token.input.count';
+export const METRIC_TOKEN_USAGE = 'gemini_cli.token.usage';
 export const METRIC_SESSION_COUNT = 'gemini_cli.session.count';
diff --git a/packages/core/src/telemetry/loggers.ts b/packages/core/src/telemetry/loggers.ts
index e8a6156c..66be0fca 100644
--- a/packages/core/src/telemetry/loggers.ts
+++ b/packages/core/src/telemetry/loggers.ts
@@ -25,7 +25,7 @@ import {
 } from './types.js';
 import {
   recordApiErrorMetrics,
-  recordApiRequestMetrics,
+  recordTokenUsageMetrics,
   recordApiResponseMetrics,
   recordToolCallMetrics,
 } from './metrics.js';
@@ -120,11 +120,11 @@ export function logApiRequest(
   };
   const logger = logs.getLogger(SERVICE_NAME);
   const logRecord: LogRecord = {
-    body: `API request to ${event.model}. Tokens: ${event.prompt_token_count}.`,
+    body: `API request to ${event.model}. Tokens: ${event.input_token_count}.`,
     attributes,
   };
   logger.emit(logRecord);
-  recordApiRequestMetrics(event.model, event.prompt_token_count);
+  recordTokenUsageMetrics(event.model, event.input_token_count, 'input');
 }
 
 export function logApiError(
@@ -188,4 +188,12 @@ export function logApiResponse(
     event.status_code,
     event.error,
   );
+  recordTokenUsageMetrics(event.model, event.output_token_count, 'output');
+  recordTokenUsageMetrics(
+    event.model,
+    event.cached_content_token_count,
+    'cache',
+  );
+  recordTokenUsageMetrics(event.model, event.thoughts_token_count, 'thought');
+  recordTokenUsageMetrics(event.model, event.tool_token_count, 'tool');
 }
diff --git a/packages/core/src/telemetry/metrics.test.ts b/packages/core/src/telemetry/metrics.test.ts
new file mode 100644
index 00000000..b7864131
--- /dev/null
+++ b/packages/core/src/telemetry/metrics.test.ts
@@ -0,0 +1,86 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { describe, it, expect, vi, beforeEach, type Mock } from 'vitest';
+import { Counter, Meter, metrics } from '@opentelemetry/api';
+import { initializeMetrics, recordTokenUsageMetrics } from './metrics.js';
+
+const mockCounter = {
+  add: vi.fn(),
+} as unknown as Counter;
+
+const mockMeter = {
+  createCounter: vi.fn().mockReturnValue(mockCounter),
+  createHistogram: vi.fn().mockReturnValue({ record: vi.fn() }),
+} as unknown as Meter;
+
+vi.mock('@opentelemetry/api', () => ({
+  metrics: {
+    getMeter: vi.fn(),
+  },
+  ValueType: {
+    INT: 1,
+  },
+}));
+
+describe('Telemetry Metrics', () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+    (metrics.getMeter as Mock).mockReturnValue(mockMeter);
+  });
+
+  describe('recordTokenUsageMetrics', () => {
+    it('should not record metrics if not initialized', () => {
+      recordTokenUsageMetrics('gemini-pro', 100, 'input');
+      expect(mockCounter.add).not.toHaveBeenCalled();
+    });
+
+    it('should record token usage with the correct attributes', () => {
+      initializeMetrics();
+      recordTokenUsageMetrics('gemini-pro', 100, 'input');
+      expect(mockCounter.add).toHaveBeenCalledWith(100, {
+        model: 'gemini-pro',
+        type: 'input',
+      });
+    });
+
+    it('should record token usage for different types', () => {
+      initializeMetrics();
+      recordTokenUsageMetrics('gemini-pro', 50, 'output');
+      expect(mockCounter.add).toHaveBeenCalledWith(50, {
+        model: 'gemini-pro',
+        type: 'output',
+      });
+
+      recordTokenUsageMetrics('gemini-pro', 25, 'thought');
+      expect(mockCounter.add).toHaveBeenCalledWith(25, {
+        model: 'gemini-pro',
+        type: 'thought',
+      });
+
+      recordTokenUsageMetrics('gemini-pro', 75, 'cache');
+      expect(mockCounter.add).toHaveBeenCalledWith(75, {
+        model: 'gemini-pro',
+        type: 'cache',
+      });
+
+      recordTokenUsageMetrics('gemini-pro', 125, 'tool');
+      expect(mockCounter.add).toHaveBeenCalledWith(125, {
+        model: 'gemini-pro',
+        type: 'tool',
+      });
+    });
+
+    it('should handle different models', () => {
+      initializeMetrics();
+      recordTokenUsageMetrics('gemini-ultra', 200, 'input');
+      expect(mockCounter.add).toHaveBeenCalledWith(200, {
+        model: 'gemini-ultra',
+        type: 'input',
+      });
+    });
+  });
+});
diff --git a/packages/core/src/telemetry/metrics.ts b/packages/core/src/telemetry/metrics.ts
index 2e6bd909..69ac91e9 100644
--- a/packages/core/src/telemetry/metrics.ts
+++ b/packages/core/src/telemetry/metrics.ts
@@ -18,7 +18,7 @@ import {
   METRIC_TOOL_CALL_LATENCY,
   METRIC_API_REQUEST_COUNT,
   METRIC_API_REQUEST_LATENCY,
-  METRIC_TOKEN_INPUT_COUNT,
+  METRIC_TOKEN_USAGE,
   METRIC_SESSION_COUNT,
 } from './constants.js';
 
@@ -27,7 +27,7 @@ let toolCallCounter: Counter | undefined;
 let toolCallLatencyHistogram: Histogram | undefined;
 let apiRequestCounter: Counter | undefined;
 let apiRequestLatencyHistogram: Histogram | undefined;
-let tokenInputCounter: Counter | undefined;
+let tokenUsageCounter: Counter | undefined;
 let isMetricsInitialized = false;
 
 export function getMeter(): Meter | undefined {
@@ -64,8 +64,8 @@ export function initializeMetrics(): void {
       valueType: ValueType.INT,
     },
   );
-  tokenInputCounter = meter.createCounter(METRIC_TOKEN_INPUT_COUNT, {
-    description: 'Counts the total number of input tokens sent to the API.',
+  tokenUsageCounter = meter.createCounter(METRIC_TOKEN_USAGE, {
+    description: 'Counts the total number of tokens used.',
     valueType: ValueType.INT,
   });
 
@@ -95,12 +95,13 @@ export function recordToolCallMetrics(
   });
 }
 
-export function recordApiRequestMetrics(
+export function recordTokenUsageMetrics(
   model: string,
-  inputTokenCount: number,
+  tokenCount: number,
+  type: 'input' | 'output' | 'thought' | 'cache' | 'tool',
 ): void {
-  if (!tokenInputCounter || !isMetricsInitialized) return;
-  tokenInputCounter.add(inputTokenCount, { model });
+  if (!tokenUsageCounter || !isMetricsInitialized) return;
+  tokenUsageCounter.add(tokenCount, { model, type });
 }
 
 export function recordApiResponseMetrics(
diff --git a/packages/core/src/telemetry/types.ts b/packages/core/src/telemetry/types.ts
index ea65d6de..4e2933a0 100644
--- a/packages/core/src/telemetry/types.ts
+++ b/packages/core/src/telemetry/types.ts
@@ -27,7 +27,7 @@ export interface ApiRequestEvent {
   'event.timestamp': string; // ISO 8601
   model: string;
   duration_ms: number;
-  prompt_token_count: number;
+  input_token_count: number;
 }
 
 export interface ApiErrorEvent {
@@ -49,6 +49,10 @@ export interface ApiResponseEvent {
   duration_ms: number;
   error?: string;
   attempt: number;
+  output_token_count: number;
+  cached_content_token_count: number;
+  thoughts_token_count: number;
+  tool_token_count: number;
 }
 
 export interface CliConfigEvent {