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();
   });

packages/core/src/telemetry/trace.ts

@@ -25,9 +25,42 @@ import {
 } from './constants.js';
 import { sessionId } from '../utils/session.js';
 
+import { truncateString } from '../utils/textUtils.js';
+
 const TRACER_NAME = 'gemini-cli';
 const TRACER_VERSION = 'v1';
 
+export function truncateForTelemetry(
+  value: unknown,
+  maxLength: number = 10000,
+): AttributeValue | undefined {
+  if (typeof value === 'string') {
+    return truncateString(
+      value,
+      maxLength,
+      `...[TRUNCATED: original length ${value.length}]`,
+    );
+  }
+  if (typeof value === 'object' && value !== null) {
+    const stringified = safeJsonStringify(value);
+    return truncateString(
+      stringified,
+      maxLength,
+      `...[TRUNCATED: original length ${stringified.length}]`,
+    );
+  }
+  if (typeof value === 'number' || typeof value === 'boolean') {
+    return value;
+  }
+  return undefined;
+}
+
+function isAsyncIterable<T>(value: T): value is T & AsyncIterable<unknown> {
+  return (
+    typeof value === 'object' && value !== null && Symbol.asyncIterator in value
+  );
+}
+
 /**
  * Metadata for a span.
  */
@@ -63,15 +96,10 @@ export interface SpanMetadata {
  * @returns The result of the function.
  */
 export async function runInDevTraceSpan<R>(
-  opts: SpanOptions & { operation: GeminiCliOperation; noAutoEnd?: boolean },
-  fn: ({
-    metadata,
-  }: {
-    metadata: SpanMetadata;
-    endSpan: () => void;
-  }) => Promise<R>,
+  opts: SpanOptions & { operation: GeminiCliOperation; logPrompts?: boolean },
+  fn: ({ metadata }: { metadata: SpanMetadata }) => Promise<R>,
 ): Promise<R> {
-  const { operation, noAutoEnd, ...restOfSpanOpts } = opts;
+  const { operation, logPrompts, ...restOfSpanOpts } = opts;
 
   const tracer = trace.getTracer(TRACER_NAME, TRACER_VERSION);
   return tracer.startActiveSpan(operation, restOfSpanOpts, async (span) => {
@@ -86,20 +114,25 @@ export async function runInDevTraceSpan<R>(
     };
     const endSpan = () => {
       try {
-        if (meta.input !== undefined) {
-          span.setAttribute(
-            GEN_AI_INPUT_MESSAGES,
-            safeJsonStringify(meta.input),
-          );
-        }
-        if (meta.output !== undefined) {
-          span.setAttribute(
-            GEN_AI_OUTPUT_MESSAGES,
-            safeJsonStringify(meta.output),
-          );
+        if (logPrompts !== false) {
+          if (meta.input !== undefined) {
+            const truncated = truncateForTelemetry(meta.input);
+            if (truncated !== undefined) {
+              span.setAttribute(GEN_AI_INPUT_MESSAGES, truncated);
+            }
+          }
+          if (meta.output !== undefined) {
+            const truncated = truncateForTelemetry(meta.output);
+            if (truncated !== undefined) {
+              span.setAttribute(GEN_AI_OUTPUT_MESSAGES, truncated);
+            }
+          }
         }
         for (const [key, value] of Object.entries(meta.attributes)) {
-          span.setAttribute(key, value);
+          const truncated = truncateForTelemetry(value);
+          if (truncated !== undefined) {
+            span.setAttribute(key, truncated);
+          }
         }
         if (meta.error) {
           span.setStatus({
@@ -123,20 +156,32 @@ export async function runInDevTraceSpan<R>(
         span.end();
       }
     };
+
+    let isStream = false;
     try {
-      return await fn({ metadata: meta, endSpan });
+      const result = await fn({ metadata: meta });
+
+      if (isAsyncIterable(result)) {
+        isStream = true;
+        const streamWrapper = (async function* () {
+          try {
+            yield* result;
+          } catch (e) {
+            meta.error = e;
+            throw e;
+          } finally {
+            endSpan();
+          }
+        })();
+
+        return Object.assign(streamWrapper, result);
+      }
+      return result;
     } catch (e) {
       meta.error = e;
-      if (noAutoEnd) {
-        // For streaming operations, the delegated endSpan call will not be reached
-        // on an exception, so we must end the span here to prevent a leak.
-        endSpan();
-      }
       throw e;
     } finally {
-      if (!noAutoEnd) {
-        // For non-streaming operations, this ensures the span is always closed,
-        // and if an error occurred, it will be recorded correctly by endSpan.
+      if (!isStream) {
         endSpan();
       }
     }