Handle telemetry in non-interactive mode (#1002)

Changes:
- Ensure proper shutdown in non-interactive mode
- Ensures the initial user prompt is logged in non-interactive mode
- Improve telemetry for streaming - handle chunks and input token count is now alongside other token counts in response

To test:
- Follow instructions in https://github.com/google-gemini/gemini-cli/blob/main/docs/core/telemetry.md#google-cloud
- Run CLI in non-interactive mode and observe logs/metrics in GCP Logs Explorer and Metrics Explorer

#750

docs/core/telemetry.md

@@ -296,7 +296,6 @@ These are timestamped records of specific events.
 
   - **Attributes**:
     - `model`
-    - `input_token_count`
     - `request_text` (optional)
 
 - `gemini_cli.api_error`: Fired if the API request fails.

packages/cli/src/gemini.tsx

@@ -33,6 +33,7 @@ import {
   WebSearchTool,
   WriteFileTool,
   sessionId,
+  logUserPrompt,
 } from '@gemini-cli/core';
 
 export async function main() {
@@ -120,6 +121,11 @@ export async function main() {
     process.exit(1);
   }
 
+  logUserPrompt(config, {
+    prompt: input,
+    prompt_length: input.length,
+  });
+
   // Non-interactive mode handled by runNonInteractive
   const nonInteractiveConfig = await loadNonInteractiveConfig(
     config,

packages/cli/src/nonInteractiveCli.ts

@@ -9,6 +9,8 @@ import {
   ToolCallRequestInfo,
   executeToolCall,
   ToolRegistry,
+  shutdownTelemetry,
+  isTelemetrySdkInitialized,
 } from '@gemini-cli/core';
 import {
   Content,
@@ -119,5 +121,9 @@ export async function runNonInteractive(
   } catch (error) {
     console.error('Error processing input:', error);
     process.exit(1);
+  } finally {
+    if (isTelemetrySdkInitialized()) {
+      await shutdownTelemetry();
+    }
   }
 }

packages/core/src/core/geminiChat.ts

@@ -14,6 +14,7 @@ import {
   SendMessageParameters,
   createUserContent,
   Part,
+  GenerateContentResponseUsageMetadata,
 } from '@google/genai';
 import { retryWithBackoff } from '../utils/retry.js';
 import { isFunctionResponse } from '../utils/messageInspectors.js';
@@ -23,13 +24,12 @@ import {
   logApiRequest,
   logApiResponse,
   logApiError,
+  combinedUsageMetadata,
 } from '../telemetry/loggers.js';
 import {
   getResponseText,
   getResponseTextFromParts,
 } from '../utils/generateContentResponseUtilities.js';
-import { getErrorMessage } from '../utils/errors.js';
-import { getRequestTextFromContents } from '../utils/requestUtils.js';
 
 /**
  * Returns true if the response is valid, false otherwise.
@@ -140,51 +140,42 @@ export class GeminiChat {
     validateHistory(history);
   }
 
+  private _getRequestTextFromContents(contents: Content[]): string {
+    return contents
+      .flatMap((content) => content.parts ?? [])
+      .map((part) => part.text)
+      .filter(Boolean)
+      .join('');
+  }
+
   private async _logApiRequest(
     contents: Content[],
     model: string,
   ): Promise<void> {
-    let inputTokenCount = 0;
-    try {
-      const { totalTokens } = await this.contentGenerator.countTokens({
-        model,
-        contents,
-      });
-      inputTokenCount = totalTokens || 0;
-    } catch (_e) {
-      console.warn(
-        `Failed to count tokens for model ${model}. Proceeding with inputTokenCount = 0. Error: ${getErrorMessage(_e)}`,
-      );
-      inputTokenCount = 0;
-    }
-
     const shouldLogUserPrompts = (config: Config): boolean =>
       config.getTelemetryLogUserPromptsEnabled() ?? false;
 
-    const requestText = getRequestTextFromContents(contents);
+    const requestText = this._getRequestTextFromContents(contents);
     logApiRequest(this.config, {
       model,
-      input_token_count: inputTokenCount,
       request_text: shouldLogUserPrompts(this.config) ? requestText : undefined,
     });
   }
 
   private async _logApiResponse(
     durationMs: number,
-    response: GenerateContentResponse,
+    usageMetadata?: GenerateContentResponseUsageMetadata,
-    fullStreamedText?: string,
+    responseText?: string,
   ): Promise<void> {
-    const responseText = fullStreamedText ?? getResponseText(response);
     logApiResponse(this.config, {
       model: this.model,
       duration_ms: durationMs,
       status_code: 200, // Assuming 200 for success
-      input_token_count: response.usageMetadata?.promptTokenCount ?? 0,
+      input_token_count: usageMetadata?.promptTokenCount ?? 0,
-      output_token_count: response.usageMetadata?.candidatesTokenCount ?? 0,
+      output_token_count: usageMetadata?.candidatesTokenCount ?? 0,
-      cached_content_token_count:
+      cached_content_token_count: usageMetadata?.cachedContentTokenCount ?? 0,
-        response.usageMetadata?.cachedContentTokenCount ?? 0,
+      thoughts_token_count: usageMetadata?.thoughtsTokenCount ?? 0,
-      thoughts_token_count: response.usageMetadata?.thoughtsTokenCount ?? 0,
+      tool_token_count: usageMetadata?.toolUsePromptTokenCount ?? 0,
-      tool_token_count: response.usageMetadata?.toolUsePromptTokenCount ?? 0,
       response_text: responseText,
     });
   }
@@ -245,7 +236,11 @@ export class GeminiChat {
 
       response = await retryWithBackoff(apiCall);
       const durationMs = Date.now() - startTime;
-      await this._logApiResponse(durationMs, response);
+      await this._logApiResponse(
+        durationMs,
+        response.usageMetadata,
+        getResponseText(response),
+      );
 
       this.sendPromise = (async () => {
         const outputContent = response.candidates?.[0]?.content;
@@ -413,18 +408,18 @@ export class GeminiChat {
     startTime: number,
   ) {
     const outputContent: Content[] = [];
-    let lastChunk: GenerateContentResponse | undefined;
+    const chunks: GenerateContentResponse[] = [];
     let errorOccurred = false;
 
     try {
       for await (const chunk of streamResponse) {
         if (isValidResponse(chunk)) {
+          chunks.push(chunk);
           const content = chunk.candidates?.[0]?.content;
           if (content !== undefined) {
             outputContent.push(content);
           }
         }
-        lastChunk = chunk;
         yield chunk;
       }
     } catch (error) {
@@ -434,7 +429,7 @@ export class GeminiChat {
       throw error;
     }
 
-    if (!errorOccurred && lastChunk) {
+    if (!errorOccurred) {
       const durationMs = Date.now() - startTime;
       const allParts: Part[] = [];
       for (const content of outputContent) {
@@ -443,7 +438,11 @@ export class GeminiChat {
         }
       }
       const fullText = getResponseTextFromParts(allParts);
-      await this._logApiResponse(durationMs, lastChunk, fullText);
+      await this._logApiResponse(
+        durationMs,
+        combinedUsageMetadata(chunks),
+        fullText,
+      );
     }
     this.recordHistory(inputContent, outputContent);
   }

packages/core/src/core/nonInteractiveToolExecutor.ts

@@ -5,6 +5,7 @@
  */
 
 import {
+  logToolCall,
   ToolCallRequestInfo,
   ToolCallResponseInfo,
   ToolRegistry,
@@ -25,10 +26,19 @@ export async function executeToolCall(
 ): Promise<ToolCallResponseInfo> {
   const tool = toolRegistry.getTool(toolCallRequest.name);
 
+  const startTime = Date.now();
   if (!tool) {
     const error = new Error(
       `Tool "${toolCallRequest.name}" not found in registry.`,
     );
+    const durationMs = Date.now() - startTime;
+    logToolCall(config, {
+      function_name: toolCallRequest.name,
+      function_args: toolCallRequest.args,
+      duration_ms: durationMs,
+      success: false,
+      error: error.message,
+    });
     // Ensure the response structure matches what the API expects for an error
     return {
       callId: toolCallRequest.callId,
@@ -55,6 +65,14 @@ export async function executeToolCall(
       // No live output callback for non-interactive mode
     );
 
+    const durationMs = Date.now() - startTime;
+    logToolCall(config, {
+      function_name: toolCallRequest.name,
+      function_args: toolCallRequest.args,
+      duration_ms: durationMs,
+      success: true,
+    });
+
     const response = convertToFunctionResponse(
       toolCallRequest.name,
       toolCallRequest.callId,
@@ -69,6 +87,14 @@ export async function executeToolCall(
     };
   } catch (e) {
     const error = e instanceof Error ? e : new Error(String(e));
+    const durationMs = Date.now() - startTime;
+    logToolCall(config, {
+      function_name: toolCallRequest.name,
+      function_args: toolCallRequest.args,
+      duration_ms: durationMs,
+      success: false,
+      error: error.message,
+    });
     return {
       callId: toolCallRequest.callId,
       responseParts: [

packages/core/src/telemetry/index.ts

@@ -16,6 +16,7 @@ export {
   logApiRequest,
   logApiError,
   logApiResponse,
+  combinedUsageMetadata,
 } from './loggers.js';
 export {
   UserPromptEvent,

packages/core/src/telemetry/loggers.test.ts

@@ -113,7 +113,7 @@ describe('loggers', () => {
       logUserPrompt(mockConfig, event);
 
       expect(mockLogger.emit).toHaveBeenCalledWith({
-        body: 'User prompt. Length: 11',
+        body: 'User prompt. Length: 11.',
         attributes: {
           'session.id': 'test-session-id',
           'event.name': EVENT_USER_PROMPT,
@@ -137,7 +137,7 @@ describe('loggers', () => {
       logUserPrompt(mockConfig, event);
 
       expect(mockLogger.emit).toHaveBeenCalledWith({
-        body: 'User prompt. Length: 11',
+        body: 'User prompt. Length: 11.',
         attributes: {
           'session.id': 'test-session-id',
           'event.name': EVENT_USER_PROMPT,
@@ -250,70 +250,42 @@ describe('loggers', () => {
       getSessionId: () => 'test-session-id',
     } as Config;
 
-    const mockMetrics = {
-      recordTokenUsageMetrics: vi.fn(),
-    };
-
-    beforeEach(() => {
-      vi.spyOn(metrics, 'recordTokenUsageMetrics').mockImplementation(
-        mockMetrics.recordTokenUsageMetrics,
-      );
-    });
-
     it('should log an API request with request_text', () => {
       const event = {
         model: 'test-model',
-        input_token_count: 123,
         request_text: 'This is a test request',
       };
 
       logApiRequest(mockConfig, event);
 
       expect(mockLogger.emit).toHaveBeenCalledWith({
-        body: 'API request to test-model. Tokens: 123.',
+        body: 'API request to test-model.',
         attributes: {
           'session.id': 'test-session-id',
           'event.name': EVENT_API_REQUEST,
           'event.timestamp': '2025-01-01T00:00:00.000Z',
           model: 'test-model',
-          input_token_count: 123,
           request_text: 'This is a test request',
         },
       });
-
-      expect(mockMetrics.recordTokenUsageMetrics).toHaveBeenCalledWith(
-        mockConfig,
-        'test-model',
-        123,
-        'input',
-      );
     });
 
     it('should log an API request without request_text', () => {
       const event = {
         model: 'test-model',
-        input_token_count: 456,
       };
 
       logApiRequest(mockConfig, event);
 
       expect(mockLogger.emit).toHaveBeenCalledWith({
-        body: 'API request to test-model. Tokens: 456.',
+        body: 'API request to test-model.',
         attributes: {
           'session.id': 'test-session-id',
           'event.name': EVENT_API_REQUEST,
           'event.timestamp': '2025-01-01T00:00:00.000Z',
           model: 'test-model',
-          input_token_count: 456,
         },
       });
-
-      expect(mockMetrics.recordTokenUsageMetrics).toHaveBeenCalledWith(
-        mockConfig,
-        'test-model',
-        456,
-        'input',
-      );
     });
   });
 

packages/core/src/telemetry/loggers.ts

@@ -31,6 +31,10 @@ import {
 } from './metrics.js';
 import { isTelemetrySdkInitialized } from './sdk.js';
 import { ToolConfirmationOutcome } from '../index.js';
+import {
+  GenerateContentResponse,
+  GenerateContentResponseUsageMetadata,
+} from '@google/genai';
 
 const shouldLogUserPrompts = (config: Config): boolean =>
   config.getTelemetryLogUserPromptsEnabled() ?? false;
@@ -119,7 +123,7 @@ export function logUserPrompt(
 
   const logger = logs.getLogger(SERVICE_NAME);
   const logRecord: LogRecord = {
-    body: `User prompt. Length: ${event.prompt_length}`,
+    body: `User prompt. Length: ${event.prompt_length}.`,
     attributes,
   };
   logger.emit(logRecord);
@@ -176,16 +180,10 @@ export function logApiRequest(
   };
   const logger = logs.getLogger(SERVICE_NAME);
   const logRecord: LogRecord = {
-    body: `API request to ${event.model}. Tokens: ${event.input_token_count}.`,
+    body: `API request to ${event.model}.`,
     attributes,
   };
   logger.emit(logRecord);
-  recordTokenUsageMetrics(
-    config,
-    event.model,
-    event.input_token_count,
-    'input',
-  );
 }
 
 export function logApiError(
@@ -258,6 +256,12 @@ export function logApiResponse(
     event.status_code,
     event.error,
   );
+  recordTokenUsageMetrics(
+    config,
+    event.model,
+    event.input_token_count,
+    'input',
+  );
   recordTokenUsageMetrics(
     config,
     event.model,
@@ -278,3 +282,43 @@ export function logApiResponse(
   );
   recordTokenUsageMetrics(config, event.model, event.tool_token_count, 'tool');
 }
+
+export function combinedUsageMetadata(
+  chunks: GenerateContentResponse[],
+): GenerateContentResponseUsageMetadata {
+  const metadataKeys: Array<keyof GenerateContentResponseUsageMetadata> = [
+    'promptTokenCount',
+    'candidatesTokenCount',
+    'cachedContentTokenCount',
+    'thoughtsTokenCount',
+    'toolUsePromptTokenCount',
+    'totalTokenCount',
+  ];
+
+  const totals: Record<keyof GenerateContentResponseUsageMetadata, number> = {
+    promptTokenCount: 0,
+    candidatesTokenCount: 0,
+    cachedContentTokenCount: 0,
+    thoughtsTokenCount: 0,
+    toolUsePromptTokenCount: 0,
+    totalTokenCount: 0,
+    cacheTokensDetails: 0,
+    candidatesTokensDetails: 0,
+    promptTokensDetails: 0,
+    toolUsePromptTokensDetails: 0,
+    trafficType: 0,
+  };
+
+  for (const chunk of chunks) {
+    if (chunk.usageMetadata) {
+      for (const key of metadataKeys) {
+        const chunkValue = chunk.usageMetadata[key];
+        if (typeof chunkValue === 'number') {
+          totals[key] += chunkValue;
+        }
+      }
+    }
+  }
+
+  return totals as unknown as GenerateContentResponseUsageMetadata;
+}

packages/core/src/telemetry/types.ts

@@ -29,7 +29,6 @@ export interface ApiRequestEvent {
   'event.name': 'api_request';
   'event.timestamp': string; // ISO 8601
   model: string;
-  input_token_count: number;
   request_text?: string;
 }
 

packages/core/src/utils/requestUtils.ts

@@ -1,15 +0,0 @@
-/**
- * @license
- * Copyright 2025 Google LLC
- * SPDX-License-Identifier: Apache-2.0
- */
-
-import { Content } from '@google/genai';
-
-export function getRequestTextFromContents(contents: Content[]): string {
-  return contents
-    .flatMap((content) => content.parts ?? [])
-    .map((part) => part.text)
-    .filter(Boolean)
-    .join('');
-}