1 files changed, 21 insertions, 10 deletions

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',
       });