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

4 files changed, 57 insertions, 41 deletions

diff --git a/packages/core/src/telemetry/index.ts b/packages/core/src/telemetry/index.ts
index cbb7b4d2..e8248bf9 100644
--- 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,
diff --git a/packages/core/src/telemetry/loggers.test.ts b/packages/core/src/telemetry/loggers.test.ts
index 092d70f0..58ea3a10 100644
--- 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',
-      );
     });
   });
 
diff --git a/packages/core/src/telemetry/loggers.ts b/packages/core/src/telemetry/loggers.ts
index 66f584e7..f6896def 100644
--- 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(
@@ -261,6 +259,12 @@ export function logApiResponse(
   recordTokenUsageMetrics(
     config,
     event.model,
+    event.input_token_count,
+    'input',
+  );
+  recordTokenUsageMetrics(
+    config,
+    event.model,
     event.output_token_count,
     'output',
   );
@@ -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;
+}
diff --git a/packages/core/src/telemetry/types.ts b/packages/core/src/telemetry/types.ts
index e9d397a8..bed9b16d 100644
--- 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;
 }