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
2025-06-13 03:44:17 -04:00 · 2025-06-13 03:44:17 -04:00 · b20c8389f3
parent 8bb6eca915
commit b20c8389f3
10 changed files with 126 additions and 89 deletions
--- a/docs/core/telemetry.md
+++ b/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.
--- a/packages/cli/src/gemini.tsx
+++ b/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,
--- a/packages/cli/src/nonInteractiveCli.ts
+++ b/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();
+    }
  }
 }
--- a/packages/core/src/core/geminiChat.ts
+++ b/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,
-    fullStreamedText?: string,
+    usageMetadata?: GenerateContentResponseUsageMetadata,
+    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,
-      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,
+      input_token_count: usageMetadata?.promptTokenCount ?? 0,
+      output_token_count: usageMetadata?.candidatesTokenCount ?? 0,
+      cached_content_token_count: usageMetadata?.cachedContentTokenCount ?? 0,
+      thoughts_token_count: usageMetadata?.thoughtsTokenCount ?? 0,
+      tool_token_count: 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);
  }
--- a/packages/core/src/core/nonInteractiveToolExecutor.ts
+++ b/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: [
--- a/packages/core/src/telemetry/index.ts
+++ b/packages/core/src/telemetry/index.ts
@ -16,6 +16,7 @@ export {
  logApiRequest,
  logApiError,
  logApiResponse,
+  combinedUsageMetadata,
 } from './loggers.js';
 export {
  UserPromptEvent,
--- a/packages/core/src/telemetry/loggers.test.ts
+++ b/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',
-      );
    });
  });

--- a/packages/core/src/telemetry/loggers.ts
+++ b/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;
+}
--- a/packages/core/src/telemetry/types.ts
+++ b/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;
 }

--- a/packages/core/src/utils/requestUtils.ts
+++ b/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('');
-}