refactor(telemetry): pass config object to telemetry functions

This commit refactors the telemetry system to pass a  object to various logging and metrics functions. This change centralizes configuration management within the telemetry system, making it more modular and easier to maintain.

The  constructor and various tool execution functions have been updated to accept the  object, which is then passed down to the telemetry functions. This eliminates the need to pass individual configuration values, such as , through multiple layers of the application.

docs/core/telemetry.md

@@ -252,6 +252,8 @@ Use this method if you prefer not to use Docker.
 
 ## Data Reference: Logs & Metrics
 
+A `sessionId` is included as a common attribute on all logs and metrics.
+
 ### Logs
 
 These are timestamped records of specific events.

packages/cli/src/nonInteractiveCli.test.ts

@@ -134,6 +134,7 @@ describe('runNonInteractive', () => {
 
     expect(mockChat.sendMessageStream).toHaveBeenCalledTimes(2);
     expect(mockCoreExecuteToolCall).toHaveBeenCalledWith(
+      mockConfig,
       expect.objectContaining({ callId: 'fc1', name: 'testTool' }),
       mockToolRegistry,
       expect.any(AbortSignal),

packages/cli/src/nonInteractiveCli.ts

@@ -85,6 +85,7 @@ export async function runNonInteractive(
           };
 
           const toolResponse = await executeToolCall(
+            config,
             requestInfo,
             toolRegistry,
             abortController.signal,

packages/cli/src/ui/hooks/useReactToolScheduler.ts

@@ -122,7 +122,7 @@ export function useReactToolScheduler(
         }
         duration = call.durationMs || 0;
 
-        logToolCall({
+        logToolCall(config, {
           function_name: call.request.name,
           function_args: call.request.args,
           duration_ms: duration,

packages/core/src/core/client.ts

@@ -160,9 +160,9 @@ export class GeminiClient {
       const systemInstruction = getCoreSystemPrompt(userMemory);
 
       return new GeminiChat(
+        this.config,
         await this.contentGenerator,
         this.model,
-        this.config.getSessionId(),
         {
           systemInstruction,
           ...this.generateContentConfig,
@@ -214,7 +214,7 @@ export class GeminiClient {
   }
 
   private _logApiRequest(model: string, inputTokenCount: number): void {
-    logApiRequest({
+    logApiRequest(this.config, {
       model,
       input_token_count: inputTokenCount,
       duration_ms: 0, // Duration is not known at request time
@@ -239,7 +239,7 @@ export class GeminiClient {
       responseError = `Finished with reason: ${finishReason}`;
     }
 
-    logApiResponse({
+    logApiResponse(this.config, {
       model,
       duration_ms: durationMs,
       attempt,
@@ -277,7 +277,7 @@ export class GeminiClient {
       }
     }
 
-    logApiError({
+    logApiError(this.config, {
       model,
       error: errorMessage,
       status_code: statusCode,

packages/core/src/core/geminiChat.test.ts

@@ -7,6 +7,7 @@
 import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
 import { Content, Models, GenerateContentConfig, Part } from '@google/genai';
 import { GeminiChat } from './geminiChat.js';
+import { Config } from '../config/config.js';
 
 // Mocks
 const mockModelsModule = {
@@ -17,16 +18,19 @@ const mockModelsModule = {
   batchEmbedContents: vi.fn(),
 } as unknown as Models;
 
+const mockConfig = {
+  getSessionId: () => 'test-session-id',
+} as unknown as Config;
+
 describe('GeminiChat', () => {
   let chat: GeminiChat;
   const model = 'gemini-pro';
   const config: GenerateContentConfig = {};
-  const sessionId = 'test-session-id';
 
   beforeEach(() => {
     vi.clearAllMocks();
     // Reset history for each test by creating a new instance
-    chat = new GeminiChat(mockModelsModule, model, sessionId, config, []);
+    chat = new GeminiChat(mockConfig, mockModelsModule, model, config, []);
   });
 
   afterEach(() => {
@@ -121,7 +125,7 @@ describe('GeminiChat', () => {
       chat.recordHistory(userInput, newModelOutput); // userInput here is for the *next* turn, but history is already primed
 
       // Reset and set up a more realistic scenario for merging with existing history
-      chat = new GeminiChat(mockModelsModule, model, sessionId, config, []);
+      chat = new GeminiChat(mockConfig, mockModelsModule, model, config, []);
       const firstUserInput: Content = {
         role: 'user',
         parts: [{ text: 'First user input' }],
@@ -164,7 +168,7 @@ describe('GeminiChat', () => {
         role: 'model',
         parts: [{ text: 'Initial model answer.' }],
       };
-      chat = new GeminiChat(mockModelsModule, model, sessionId, config, [
+      chat = new GeminiChat(mockConfig, mockModelsModule, model, config, [
         initialUser,
         initialModel,
       ]);

packages/core/src/core/geminiChat.ts

@@ -18,7 +18,7 @@ import {
 import { retryWithBackoff } from '../utils/retry.js';
 import { isFunctionResponse } from '../utils/messageInspectors.js';
 import { ContentGenerator } from './contentGenerator.js';
-import { Logger } from './logger.js';
+import { Config } from '../config/config.js';
 
 /**
  * Returns true if the response is valid, false otherwise.
@@ -118,17 +118,15 @@ export class GeminiChat {
   // A promise to represent the current state of the message being sent to the
   // model.
   private sendPromise: Promise<void> = Promise.resolve();
-  private logger: Logger;
 
   constructor(
+    private readonly config: Config,
     private readonly contentGenerator: ContentGenerator,
     private readonly model: string,
-    sessionId: string,
-    private readonly config: GenerateContentConfig = {},
+    private readonly generationConfig: GenerateContentConfig = {},
     private history: Content[] = [],
   ) {
     validateHistory(history);
-    this.logger = new Logger(sessionId);
   }
 
   /**
@@ -161,7 +159,7 @@ export class GeminiChat {
       this.contentGenerator.generateContent({
         model: this.model,
         contents: this.getHistory(true).concat(userContent),
-        config: { ...this.config, ...params.config },
+        config: { ...this.generationConfig, ...params.config },
       });
 
     const responsePromise = retryWithBackoff(apiCall);
@@ -230,7 +228,7 @@ export class GeminiChat {
       this.contentGenerator.generateContentStream({
         model: this.model,
         contents: this.getHistory(true).concat(userContent),
-        config: { ...this.config, ...params.config },
+        config: { ...this.generationConfig, ...params.config },
       });
 
     // Note: Retrying streams can be complex. If generateContentStream itself doesn't handle retries

packages/core/src/core/nonInteractiveToolExecutor.test.ts

@@ -12,9 +12,12 @@ import {
   ToolResult,
   Tool,
   ToolCallConfirmationDetails,
+  Config,
 } from '../index.js';
 import { Part, Type } from '@google/genai';
 
+const mockConfig = {} as unknown as Config;
+
 describe('executeToolCall', () => {
   let mockToolRegistry: ToolRegistry;
   let mockTool: Tool;
@@ -68,6 +71,7 @@ describe('executeToolCall', () => {
     vi.mocked(mockTool.execute).mockResolvedValue(toolResult);
 
     const response = await executeToolCall(
+      mockConfig,
       request,
       mockToolRegistry,
       abortController.signal,
@@ -99,6 +103,7 @@ describe('executeToolCall', () => {
     vi.mocked(mockToolRegistry.getTool).mockReturnValue(undefined);
 
     const response = await executeToolCall(
+      mockConfig,
       request,
       mockToolRegistry,
       abortController.signal,
@@ -134,6 +139,7 @@ describe('executeToolCall', () => {
     vi.mocked(mockTool.execute).mockRejectedValue(executionError);
 
     const response = await executeToolCall(
+      mockConfig,
       request,
       mockToolRegistry,
       abortController.signal,
@@ -184,6 +190,7 @@ describe('executeToolCall', () => {
 
     abortController.abort(); // Abort before calling
     const response = await executeToolCall(
+      mockConfig,
       request,
       mockToolRegistry,
       abortController.signal,
@@ -211,6 +218,7 @@ describe('executeToolCall', () => {
     vi.mocked(mockTool.execute).mockResolvedValue(toolResult);
 
     const response = await executeToolCall(
+      mockConfig,
       request,
       mockToolRegistry,
       abortController.signal,

packages/core/src/core/nonInteractiveToolExecutor.ts

@@ -10,6 +10,7 @@ import {
   ToolRegistry,
   ToolResult,
 } from '../index.js';
+import { Config } from '../config/config.js';
 import { convertToFunctionResponse } from './coreToolScheduler.js';
 
 /**
@@ -17,6 +18,7 @@ import { convertToFunctionResponse } from './coreToolScheduler.js';
  * It does not handle confirmations, multiple calls, or live updates.
  */
 export async function executeToolCall(
+  config: Config,
   toolCallRequest: ToolCallRequestInfo,
   toolRegistry: ToolRegistry,
   abortSignal?: AbortSignal,

packages/core/src/telemetry/loggers.ts

@@ -34,10 +34,17 @@ import { isTelemetrySdkInitialized } from './sdk.js';
 const shouldLogUserPrompts = (config: Config): boolean =>
   config.getTelemetryLogUserPromptsEnabled() ?? false;
 
+function getCommonAttributes(config: Config): LogAttributes {
+  return {
+    'session.id': config.getSessionId(),
+  };
+}
+
 export function logCliConfiguration(config: Config): void {
   if (!isTelemetrySdkInitialized()) return;
 
   const attributes: LogAttributes = {
+    ...getCommonAttributes(config),
     'event.name': EVENT_CLI_CONFIG,
     'event.timestamp': new Date().toISOString(),
     model: config.getModel(),
@@ -69,6 +76,7 @@ export function logUserPrompt(
   if (!isTelemetrySdkInitialized()) return;
   const { prompt, ...restOfEventArgs } = event;
   const attributes: LogAttributes = {
+    ...getCommonAttributes(config),
     ...restOfEventArgs,
     'event.name': EVENT_USER_PROMPT,
     'event.timestamp': new Date().toISOString(),
@@ -85,10 +93,12 @@ export function logUserPrompt(
 }
 
 export function logToolCall(
+  config: Config,
   event: Omit<ToolCallEvent, 'event.name' | 'event.timestamp'>,
 ): void {
   if (!isTelemetrySdkInitialized()) return;
   const attributes: LogAttributes = {
+    ...getCommonAttributes(config),
     ...event,
     'event.name': EVENT_TOOL_CALL,
     'event.timestamp': new Date().toISOString(),
@@ -106,14 +116,21 @@ export function logToolCall(
     attributes,
   };
   logger.emit(logRecord);
-  recordToolCallMetrics(event.function_name, event.duration_ms, event.success);
+  recordToolCallMetrics(
+    config,
+    event.function_name,
+    event.duration_ms,
+    event.success,
+  );
 }
 
 export function logApiRequest(
+  config: Config,
   event: Omit<ApiRequestEvent, 'event.name' | 'event.timestamp'>,
 ): void {
   if (!isTelemetrySdkInitialized()) return;
   const attributes: LogAttributes = {
+    ...getCommonAttributes(config),
     ...event,
     'event.name': EVENT_API_REQUEST,
     'event.timestamp': new Date().toISOString(),
@@ -124,14 +141,21 @@ export function logApiRequest(
     attributes,
   };
   logger.emit(logRecord);
-  recordTokenUsageMetrics(event.model, event.input_token_count, 'input');
+  recordTokenUsageMetrics(
+    config,
+    event.model,
+    event.input_token_count,
+    'input',
+  );
 }
 
 export function logApiError(
+  config: Config,
   event: Omit<ApiErrorEvent, 'event.name' | 'event.timestamp'>,
 ): void {
   if (!isTelemetrySdkInitialized()) return;
   const attributes: LogAttributes = {
+    ...getCommonAttributes(config),
     ...event,
     'event.name': EVENT_API_ERROR,
     'event.timestamp': new Date().toISOString(),
@@ -152,6 +176,7 @@ export function logApiError(
   };
   logger.emit(logRecord);
   recordApiErrorMetrics(
+    config,
     event.model,
     event.duration_ms,
     event.status_code,
@@ -160,10 +185,12 @@ export function logApiError(
 }
 
 export function logApiResponse(
+  config: Config,
   event: Omit<ApiResponseEvent, 'event.name' | 'event.timestamp'>,
 ): void {
   if (!isTelemetrySdkInitialized()) return;
   const attributes: LogAttributes = {
+    ...getCommonAttributes(config),
     ...event,
     'event.name': EVENT_API_RESPONSE,
     'event.timestamp': new Date().toISOString(),
@@ -183,17 +210,29 @@ export function logApiResponse(
   };
   logger.emit(logRecord);
   recordApiResponseMetrics(
+    config,
     event.model,
     event.duration_ms,
     event.status_code,
     event.error,
   );
-  recordTokenUsageMetrics(event.model, event.output_token_count, 'output');
   recordTokenUsageMetrics(
+    config,
+    event.model,
+    event.output_token_count,
+    'output',
+  );
+  recordTokenUsageMetrics(
+    config,
     event.model,
     event.cached_content_token_count,
     'cache',
   );
-  recordTokenUsageMetrics(event.model, event.thoughts_token_count, 'thought');
-  recordTokenUsageMetrics(event.model, event.tool_token_count, 'tool');
+  recordTokenUsageMetrics(
+    config,
+    event.model,
+    event.thoughts_token_count,
+    'thought',
+  );
+  recordTokenUsageMetrics(config, event.model, event.tool_token_count, 'tool');
 }

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

@@ -7,6 +7,7 @@
 import { describe, it, expect, vi, beforeEach, type Mock } from 'vitest';
 import { Counter, Meter, metrics } from '@opentelemetry/api';
 import { initializeMetrics, recordTokenUsageMetrics } from './metrics.js';
+import { Config } from '../config/config.js';
 
 const mockCounter = {
   add: vi.fn(),
@@ -33,51 +34,61 @@ describe('Telemetry Metrics', () => {
   });
 
   describe('recordTokenUsageMetrics', () => {
+    const mockConfig = {
+      getSessionId: () => 'test-session-id',
+    } as unknown as Config;
+
     it('should not record metrics if not initialized', () => {
-      recordTokenUsageMetrics('gemini-pro', 100, 'input');
+      recordTokenUsageMetrics(mockConfig, 'gemini-pro', 100, 'input');
       expect(mockCounter.add).not.toHaveBeenCalled();
     });
 
     it('should record token usage with the correct attributes', () => {
-      initializeMetrics();
-      recordTokenUsageMetrics('gemini-pro', 100, 'input');
+      initializeMetrics(mockConfig);
+      recordTokenUsageMetrics(mockConfig, 'gemini-pro', 100, 'input');
       expect(mockCounter.add).toHaveBeenCalledWith(100, {
+        'session.id': 'test-session-id',
         model: 'gemini-pro',
         type: 'input',
       });
     });
 
     it('should record token usage for different types', () => {
-      initializeMetrics();
-      recordTokenUsageMetrics('gemini-pro', 50, 'output');
+      initializeMetrics(mockConfig);
+      recordTokenUsageMetrics(mockConfig, 'gemini-pro', 50, 'output');
       expect(mockCounter.add).toHaveBeenCalledWith(50, {
+        'session.id': 'test-session-id',
         model: 'gemini-pro',
         type: 'output',
       });
 
-      recordTokenUsageMetrics('gemini-pro', 25, 'thought');
+      recordTokenUsageMetrics(mockConfig, 'gemini-pro', 25, 'thought');
       expect(mockCounter.add).toHaveBeenCalledWith(25, {
+        'session.id': 'test-session-id',
         model: 'gemini-pro',
         type: 'thought',
       });
 
-      recordTokenUsageMetrics('gemini-pro', 75, 'cache');
+      recordTokenUsageMetrics(mockConfig, 'gemini-pro', 75, 'cache');
       expect(mockCounter.add).toHaveBeenCalledWith(75, {
+        'session.id': 'test-session-id',
         model: 'gemini-pro',
         type: 'cache',
       });
 
-      recordTokenUsageMetrics('gemini-pro', 125, 'tool');
+      recordTokenUsageMetrics(mockConfig, 'gemini-pro', 125, 'tool');
       expect(mockCounter.add).toHaveBeenCalledWith(125, {
+        'session.id': 'test-session-id',
         model: 'gemini-pro',
         type: 'tool',
       });
     });
 
     it('should handle different models', () => {
-      initializeMetrics();
-      recordTokenUsageMetrics('gemini-ultra', 200, 'input');
+      initializeMetrics(mockConfig);
+      recordTokenUsageMetrics(mockConfig, 'gemini-ultra', 200, 'input');
       expect(mockCounter.add).toHaveBeenCalledWith(200, {
+        'session.id': 'test-session-id',
         model: 'gemini-ultra',
         type: 'input',
       });

packages/core/src/telemetry/metrics.ts

@@ -21,6 +21,7 @@ import {
   METRIC_TOKEN_USAGE,
   METRIC_SESSION_COUNT,
 } from './constants.js';
+import { Config } from '../config/config.js';
 
 let cliMeter: Meter | undefined;
 let toolCallCounter: Counter | undefined;
@@ -30,6 +31,12 @@ let apiRequestLatencyHistogram: Histogram | undefined;
 let tokenUsageCounter: Counter | undefined;
 let isMetricsInitialized = false;
 
+function getCommonAttributes(config: Config): Attributes {
+  return {
+    'session.id': config.getSessionId(),
+  };
+}
+
 export function getMeter(): Meter | undefined {
   if (!cliMeter) {
     cliMeter = metrics.getMeter(SERVICE_NAME);
@@ -37,7 +44,7 @@ export function getMeter(): Meter | undefined {
   return cliMeter;
 }
 
-export function initializeMetrics(): void {
+export function initializeMetrics(config: Config): void {
   if (isMetricsInitialized) return;
 
   const meter = getMeter();
@@ -73,11 +80,12 @@ export function initializeMetrics(): void {
     description: 'Count of CLI sessions started.',
     valueType: ValueType.INT,
   });
-  sessionCounter.add(1);
+  sessionCounter.add(1, getCommonAttributes(config));
   isMetricsInitialized = true;
 }
 
 export function recordToolCallMetrics(
+  config: Config,
   functionName: string,
   durationMs: number,
   success: boolean,
@@ -86,25 +94,33 @@ export function recordToolCallMetrics(
     return;
 
   const metricAttributes: Attributes = {
+    ...getCommonAttributes(config),
     function_name: functionName,
     success,
   };
   toolCallCounter.add(1, metricAttributes);
   toolCallLatencyHistogram.record(durationMs, {
+    ...getCommonAttributes(config),
     function_name: functionName,
   });
 }
 
 export function recordTokenUsageMetrics(
+  config: Config,
   model: string,
   tokenCount: number,
   type: 'input' | 'output' | 'thought' | 'cache' | 'tool',
 ): void {
   if (!tokenUsageCounter || !isMetricsInitialized) return;
-  tokenUsageCounter.add(tokenCount, { model, type });
+  tokenUsageCounter.add(tokenCount, {
+    ...getCommonAttributes(config),
+    model,
+    type,
+  });
 }
 
 export function recordApiResponseMetrics(
+  config: Config,
   model: string,
   durationMs: number,
   statusCode?: number | string,
@@ -117,14 +133,19 @@ export function recordApiResponseMetrics(
   )
     return;
   const metricAttributes: Attributes = {
+    ...getCommonAttributes(config),
     model,
     status_code: statusCode ?? (error ? 'error' : 'ok'),
   };
   apiRequestCounter.add(1, metricAttributes);
-  apiRequestLatencyHistogram.record(durationMs, { model });
+  apiRequestLatencyHistogram.record(durationMs, {
+    ...getCommonAttributes(config),
+    model,
+  });
 }
 
 export function recordApiErrorMetrics(
+  config: Config,
   model: string,
   durationMs: number,
   statusCode?: number | string,
@@ -137,10 +158,14 @@ export function recordApiErrorMetrics(
   )
     return;
   const metricAttributes: Attributes = {
+    ...getCommonAttributes(config),
     model,
     status_code: statusCode ?? 'error',
     error_type: errorType ?? 'unknown',
   };
   apiRequestCounter.add(1, metricAttributes);
-  apiRequestLatencyHistogram.record(durationMs, { model });
+  apiRequestLatencyHistogram.record(durationMs, {
+    ...getCommonAttributes(config),
+    model,
+  });
 }

packages/core/src/telemetry/sdk.ts

@@ -112,7 +112,7 @@ export function initializeTelemetry(config: Config): void {
     sdk.start();
     console.log('OpenTelemetry SDK started successfully.');
     telemetryInitialized = true;
-    initializeMetrics();
+    initializeMetrics(config);
     logCliConfiguration(config);
   } catch (error) {
     console.error('Error starting OpenTelemetry SDK:', error);

packages/core/src/utils/nextSpeakerChecker.test.ts

@@ -69,9 +69,9 @@ describe('checkNextSpeaker', () => {
 
     // GeminiChat will receive the mocked instances via the mocked GoogleGenAI constructor
     chatInstance = new GeminiChat(
+      mockConfigInstance,
       mockModelsInstance, // This is the instance returned by mockGoogleGenAIInstance.getGenerativeModel
       'gemini-pro', // model name
-      'test-session-id',
       {},
       [], // initial history
     );