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.

4 files changed, 96 insertions, 21 deletions

diff --git a/packages/core/src/telemetry/loggers.ts b/packages/core/src/telemetry/loggers.ts
index 66be0fca..d2b01f65 100644
--- a/packages/core/src/telemetry/loggers.ts
+++ b/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');
 }
diff --git a/packages/core/src/telemetry/metrics.test.ts b/packages/core/src/telemetry/metrics.test.ts
index b7864131..7e24b9ad 100644
--- a/packages/core/src/telemetry/metrics.test.ts
+++ b/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',
       });
diff --git a/packages/core/src/telemetry/metrics.ts b/packages/core/src/telemetry/metrics.ts
index 69ac91e9..93aa2189 100644
--- a/packages/core/src/telemetry/metrics.ts
+++ b/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,
+  });
 }
diff --git a/packages/core/src/telemetry/sdk.ts b/packages/core/src/telemetry/sdk.ts
index 704661c5..61f501a6 100644
--- a/packages/core/src/telemetry/sdk.ts
+++ b/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);