fix(telemetry): patch memory leak and enforce logPrompts privacy (#23281)

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

@@ -6,7 +6,7 @@
 
 import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
 import { trace, SpanStatusCode, diag, type Tracer } from '@opentelemetry/api';
-import { runInDevTraceSpan } from './trace.js';
+import { runInDevTraceSpan, truncateForTelemetry } from './trace.js';
 import {
   GeminiCliOperation,
   GEN_AI_CONVERSATION_ID,
@@ -36,6 +36,55 @@ vi.mock('../utils/session.js', () => ({
   sessionId: 'test-session-id',
 }));
 
+describe('truncateForTelemetry', () => {
+  it('should return string unchanged if within maxLength', () => {
+    expect(truncateForTelemetry('hello', 10)).toBe('hello');
+  });
+
+  it('should truncate string if exceeding maxLength', () => {
+    const result = truncateForTelemetry('hello world', 5);
+    expect(result).toBe('hello...[TRUNCATED: original length 11]');
+  });
+
+  it('should correctly truncate strings with multi-byte unicode characters (emojis)', () => {
+    // 5 emojis, each is multiple bytes in UTF-16
+    const emojis = '👋🌍🚀🔥🎉';
+
+    // Truncating to length 5 (which is 2.5 emojis in UTF-16 length terms)
+    // truncateString will stop after the full grapheme clusters that fit within 5
+    const result = truncateForTelemetry(emojis, 5);
+
+    expect(result).toBe('👋🌍...[TRUNCATED: original length 10]');
+  });
+
+  it('should stringify and truncate objects if exceeding maxLength', () => {
+    const obj = { message: 'hello world', nested: { a: 1 } };
+    const stringified = JSON.stringify(obj);
+    const result = truncateForTelemetry(obj, 10);
+    expect(result).toBe(
+      stringified.substring(0, 10) +
+        `...[TRUNCATED: original length ${stringified.length}]`,
+    );
+  });
+
+  it('should stringify objects unchanged if within maxLength', () => {
+    const obj = { a: 1 };
+    expect(truncateForTelemetry(obj, 100)).toBe(JSON.stringify(obj));
+  });
+
+  it('should return booleans and numbers unchanged', () => {
+    expect(truncateForTelemetry(100)).toBe(100);
+    expect(truncateForTelemetry(true)).toBe(true);
+    expect(truncateForTelemetry(false)).toBe(false);
+  });
+
+  it('should return undefined for unsupported types', () => {
+    expect(truncateForTelemetry(undefined)).toBeUndefined();
+    expect(truncateForTelemetry(() => {})).toBeUndefined();
+    expect(truncateForTelemetry(Symbol('test'))).toBeUndefined();
+  });
+});
+
 describe('runInDevTraceSpan', () => {
   const mockSpan = {
     setAttribute: vi.fn(),
@@ -133,33 +182,45 @@ describe('runInDevTraceSpan', () => {
     expect(mockSpan.end).toHaveBeenCalled();
   });
 
-  it('should respect noAutoEnd option', async () => {
-    let capturedEndSpan: () => void = () => {};
-    const result = await runInDevTraceSpan(
-      { operation: GeminiCliOperation.LLMCall, noAutoEnd: true },
-      async ({ endSpan }) => {
-        capturedEndSpan = endSpan;
-        return 'streaming';
-      },
+  it('should auto-wrap async iterators and end span when iterator completes', async () => {
+    async function* testStream() {
+      yield 1;
+      yield 2;
+    }
+
+    const resultStream = await runInDevTraceSpan(
+      { operation: GeminiCliOperation.LLMCall },
+      async () => testStream(),
     );
 
-    expect(result).toBe('streaming');
     expect(mockSpan.end).not.toHaveBeenCalled();
 
-    capturedEndSpan();
+    const results = [];
+    for await (const val of resultStream) {
+      results.push(val);
+    }
+
+    expect(results).toEqual([1, 2]);
     expect(mockSpan.end).toHaveBeenCalled();
   });
 
-  it('should automatically end span on error even if noAutoEnd is true', async () => {
+  it('should end span automatically on error in async iterators', async () => {
     const error = new Error('streaming error');
-    await expect(
-      runInDevTraceSpan(
-        { operation: GeminiCliOperation.LLMCall, noAutoEnd: true },
-        async () => {
-          throw error;
-        },
-      ),
-    ).rejects.toThrow(error);
+    async function* errorStream() {
+      yield 1;
+      throw error;
+    }
+
+    const resultStream = await runInDevTraceSpan(
+      { operation: GeminiCliOperation.LLMCall },
+      async () => errorStream(),
+    );
+
+    await expect(async () => {
+      for await (const _ of resultStream) {
+        // iterate
+      }
+    }).rejects.toThrow(error);
 
     expect(mockSpan.end).toHaveBeenCalled();
   });