feat: added basic dev otel trace instrumentation (#11690)

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

packages/core/src/core/coreToolScheduler.ts

@@ -25,6 +25,7 @@ import {
   ToolCallEvent,
   logToolOutputTruncated,
   ToolOutputTruncatedEvent,
+  runInDevTraceSpan,
 } from '../index.js';
 import { READ_FILE_TOOL_NAME, SHELL_TOOL_NAME } from '../tools/tool-names.js';
 import type { Part, PartListUnion } from '@google/genai';
@@ -668,36 +669,42 @@ export class CoreToolScheduler {
     request: ToolCallRequestInfo | ToolCallRequestInfo[],
     signal: AbortSignal,
   ): Promise<void> {
-    if (this.isRunning() || this.isScheduling) {
-      return new Promise((resolve, reject) => {
-        const abortHandler = () => {
-          // Find and remove the request from the queue
-          const index = this.requestQueue.findIndex(
-            (item) => item.request === request,
-          );
-          if (index > -1) {
-            this.requestQueue.splice(index, 1);
-            reject(new Error('Tool call cancelled while in queue.'));
-          }
-        };
+    return runInDevTraceSpan(
+      { name: 'schedule' },
+      async ({ metadata: spanMetadata }) => {
+        spanMetadata.input = request;
+        if (this.isRunning() || this.isScheduling) {
+          return new Promise((resolve, reject) => {
+            const abortHandler = () => {
+              // Find and remove the request from the queue
+              const index = this.requestQueue.findIndex(
+                (item) => item.request === request,
+              );
+              if (index > -1) {
+                this.requestQueue.splice(index, 1);
+                reject(new Error('Tool call cancelled while in queue.'));
+              }
+            };
 
-        signal.addEventListener('abort', abortHandler, { once: true });
+            signal.addEventListener('abort', abortHandler, { once: true });
 
-        this.requestQueue.push({
-          request,
-          signal,
-          resolve: () => {
-            signal.removeEventListener('abort', abortHandler);
-            resolve();
-          },
-          reject: (reason?: Error) => {
-            signal.removeEventListener('abort', abortHandler);
-            reject(reason);
-          },
-        });
-      });
-    }
-    return this._schedule(request, signal);
+            this.requestQueue.push({
+              request,
+              signal,
+              resolve: () => {
+                signal.removeEventListener('abort', abortHandler);
+                resolve();
+              },
+              reject: (reason?: Error) => {
+                signal.removeEventListener('abort', abortHandler);
+                reject(reason);
+              },
+            });
+          });
+        }
+        return this._schedule(request, signal);
+      },
+    );
   }
 
   cancelAll(signal: AbortSignal): void {
@@ -1091,134 +1098,153 @@ export class CoreToolScheduler {
 
         const shellExecutionConfig = this.config.getShellExecutionConfig();
 
-        // TODO: Refactor to remove special casing for ShellToolInvocation.
-        // Introduce a generic callbacks object for the execute method to handle
-        // things like `onPid` and `onLiveOutput`. This will make the scheduler
-        // agnostic to the invocation type.
-        let promise: Promise<ToolResult>;
-        if (invocation instanceof ShellToolInvocation) {
-          const setPidCallback = (pid: number) => {
-            this.toolCalls = this.toolCalls.map((tc) =>
-              tc.request.callId === callId && tc.status === 'executing'
-                ? { ...tc, pid }
-                : tc,
-            );
-            this.notifyToolCallsUpdate();
-          };
-          promise = invocation.execute(
-            signal,
-            liveOutputCallback,
-            shellExecutionConfig,
-            setPidCallback,
-          );
-        } else {
-          promise = invocation.execute(
-            signal,
-            liveOutputCallback,
-            shellExecutionConfig,
-          );
-        }
-
-        try {
-          const toolResult: ToolResult = await promise;
-          if (signal.aborted) {
-            this.setStatusInternal(
-              callId,
-              'cancelled',
-              signal,
-              'User cancelled tool execution.',
-            );
-          } else if (toolResult.error === undefined) {
-            let content = toolResult.llmContent;
-            let outputFile: string | undefined = undefined;
-            const contentLength =
-              typeof content === 'string' ? content.length : undefined;
-            if (
-              typeof content === 'string' &&
-              toolName === SHELL_TOOL_NAME &&
-              this.config.getEnableToolOutputTruncation() &&
-              this.config.getTruncateToolOutputThreshold() > 0 &&
-              this.config.getTruncateToolOutputLines() > 0
-            ) {
-              const originalContentLength = content.length;
-              const threshold = this.config.getTruncateToolOutputThreshold();
-              const lines = this.config.getTruncateToolOutputLines();
-              const truncatedResult = await truncateAndSaveToFile(
-                content,
-                callId,
-                this.config.storage.getProjectTempDir(),
-                threshold,
-                lines,
+        await runInDevTraceSpan(
+          {
+            name: toolCall.tool.name,
+            attributes: { type: 'tool-call' },
+          },
+          async ({ metadata: spanMetadata }) => {
+            spanMetadata.input = {
+              request: toolCall.request,
+            };
+            // TODO: Refactor to remove special casing for ShellToolInvocation.
+            // Introduce a generic callbacks object for the execute method to handle
+            // things like `onPid` and `onLiveOutput`. This will make the scheduler
+            // agnostic to the invocation type.
+            let promise: Promise<ToolResult>;
+            if (invocation instanceof ShellToolInvocation) {
+              const setPidCallback = (pid: number) => {
+                this.toolCalls = this.toolCalls.map((tc) =>
+                  tc.request.callId === callId && tc.status === 'executing'
+                    ? { ...tc, pid }
+                    : tc,
+                );
+                this.notifyToolCallsUpdate();
+              };
+              promise = invocation.execute(
+                signal,
+                liveOutputCallback,
+                shellExecutionConfig,
+                setPidCallback,
               );
-              content = truncatedResult.content;
-              outputFile = truncatedResult.outputFile;
+            } else {
+              promise = invocation.execute(
+                signal,
+                liveOutputCallback,
+                shellExecutionConfig,
+              );
+            }
 
-              if (outputFile) {
-                logToolOutputTruncated(
-                  this.config,
-                  new ToolOutputTruncatedEvent(
-                    scheduledCall.request.prompt_id,
-                    {
-                      toolName,
-                      originalContentLength,
-                      truncatedContentLength: content.length,
-                      threshold,
-                      lines,
-                    },
+            try {
+              const toolResult: ToolResult = await promise;
+              spanMetadata.output = toolResult;
+              if (signal.aborted) {
+                this.setStatusInternal(
+                  callId,
+                  'cancelled',
+                  signal,
+                  'User cancelled tool execution.',
+                );
+              } else if (toolResult.error === undefined) {
+                let content = toolResult.llmContent;
+                let outputFile: string | undefined = undefined;
+                const contentLength =
+                  typeof content === 'string' ? content.length : undefined;
+                if (
+                  typeof content === 'string' &&
+                  toolName === SHELL_TOOL_NAME &&
+                  this.config.getEnableToolOutputTruncation() &&
+                  this.config.getTruncateToolOutputThreshold() > 0 &&
+                  this.config.getTruncateToolOutputLines() > 0
+                ) {
+                  const originalContentLength = content.length;
+                  const threshold =
+                    this.config.getTruncateToolOutputThreshold();
+                  const lines = this.config.getTruncateToolOutputLines();
+                  const truncatedResult = await truncateAndSaveToFile(
+                    content,
+                    callId,
+                    this.config.storage.getProjectTempDir(),
+                    threshold,
+                    lines,
+                  );
+                  content = truncatedResult.content;
+                  outputFile = truncatedResult.outputFile;
+
+                  if (outputFile) {
+                    logToolOutputTruncated(
+                      this.config,
+                      new ToolOutputTruncatedEvent(
+                        scheduledCall.request.prompt_id,
+                        {
+                          toolName,
+                          originalContentLength,
+                          truncatedContentLength: content.length,
+                          threshold,
+                          lines,
+                        },
+                      ),
+                    );
+                  }
+                }
+
+                const response = convertToFunctionResponse(
+                  toolName,
+                  callId,
+                  content,
+                );
+                const successResponse: ToolCallResponseInfo = {
+                  callId,
+                  responseParts: response,
+                  resultDisplay: toolResult.returnDisplay,
+                  error: undefined,
+                  errorType: undefined,
+                  outputFile,
+                  contentLength,
+                };
+                this.setStatusInternal(
+                  callId,
+                  'success',
+                  signal,
+                  successResponse,
+                );
+              } else {
+                // It is a failure
+                const error = new Error(toolResult.error.message);
+                const errorResponse = createErrorResponse(
+                  scheduledCall.request,
+                  error,
+                  toolResult.error.type,
+                );
+                this.setStatusInternal(callId, 'error', signal, errorResponse);
+              }
+            } catch (executionError: unknown) {
+              spanMetadata.error = executionError;
+              if (signal.aborted) {
+                this.setStatusInternal(
+                  callId,
+                  'cancelled',
+                  signal,
+                  'User cancelled tool execution.',
+                );
+              } else {
+                this.setStatusInternal(
+                  callId,
+                  'error',
+                  signal,
+                  createErrorResponse(
+                    scheduledCall.request,
+                    executionError instanceof Error
+                      ? executionError
+                      : new Error(String(executionError)),
+                    ToolErrorType.UNHANDLED_EXCEPTION,
                   ),
                 );
               }
             }
-
-            const response = convertToFunctionResponse(
-              toolName,
-              callId,
-              content,
-            );
-            const successResponse: ToolCallResponseInfo = {
-              callId,
-              responseParts: response,
-              resultDisplay: toolResult.returnDisplay,
-              error: undefined,
-              errorType: undefined,
-              outputFile,
-              contentLength,
-            };
-            this.setStatusInternal(callId, 'success', signal, successResponse);
-          } else {
-            // It is a failure
-            const error = new Error(toolResult.error.message);
-            const errorResponse = createErrorResponse(
-              scheduledCall.request,
-              error,
-              toolResult.error.type,
-            );
-            this.setStatusInternal(callId, 'error', signal, errorResponse);
-          }
-        } catch (executionError: unknown) {
-          if (signal.aborted) {
-            this.setStatusInternal(
-              callId,
-              'cancelled',
-              signal,
-              'User cancelled tool execution.',
-            );
-          } else {
-            this.setStatusInternal(
-              callId,
-              'error',
-              signal,
-              createErrorResponse(
-                scheduledCall.request,
-                executionError instanceof Error
-                  ? executionError
-                  : new Error(String(executionError)),
-                ToolErrorType.UNHANDLED_EXCEPTION,
-              ),
-            );
-          }
-        }
-        await this.checkAndNotifyCompletion(signal);
+            await this.checkAndNotifyCompletion(signal);
+          },
+        );
       }
     }
   }

packages/core/src/core/loggingContentGenerator.ts

@@ -28,6 +28,7 @@ import {
 import type { ContentGenerator } from './contentGenerator.js';
 import { toContents } from '../code_assist/converter.js';
 import { isStructuredError } from '../utils/quotaErrorDetection.js';
+import { runInDevTraceSpan, type SpanMetadata } from '../telemetry/trace.js';
 
 interface StructuredError {
   status: number;
@@ -107,47 +108,74 @@ export class LoggingContentGenerator implements ContentGenerator {
     req: GenerateContentParameters,
     userPromptId: string,
   ): Promise<GenerateContentResponse> {
-    const startTime = Date.now();
-    this.logApiRequest(toContents(req.contents), req.model, userPromptId);
-    try {
-      const response = await this.wrapped.generateContent(req, userPromptId);
-      const durationMs = Date.now() - startTime;
-      this._logApiResponse(
-        durationMs,
-        response.modelVersion || req.model,
-        userPromptId,
-        response.usageMetadata,
-        JSON.stringify(response),
-      );
-      return response;
-    } catch (error) {
-      const durationMs = Date.now() - startTime;
-      this._logApiError(durationMs, error, req.model, userPromptId);
-      throw error;
-    }
+    return runInDevTraceSpan(
+      {
+        name: 'generateContent',
+      },
+      async ({ metadata: spanMetadata }) => {
+        spanMetadata.input = { request: req, userPromptId, model: req.model };
+
+        const startTime = Date.now();
+        this.logApiRequest(toContents(req.contents), req.model, userPromptId);
+        try {
+          const response = await this.wrapped.generateContent(
+            req,
+            userPromptId,
+          );
+          spanMetadata.output = {
+            response,
+            usageMetadata: response.usageMetadata,
+          };
+          const durationMs = Date.now() - startTime;
+          this._logApiResponse(
+            durationMs,
+            response.modelVersion || req.model,
+            userPromptId,
+            response.usageMetadata,
+            JSON.stringify(response),
+          );
+          return response;
+        } catch (error) {
+          const durationMs = Date.now() - startTime;
+          this._logApiError(durationMs, error, req.model, userPromptId);
+          throw error;
+        }
+      },
+    );
   }
 
   async generateContentStream(
     req: GenerateContentParameters,
     userPromptId: string,
   ): Promise<AsyncGenerator<GenerateContentResponse>> {
-    const startTime = Date.now();
-    this.logApiRequest(toContents(req.contents), req.model, userPromptId);
+    return runInDevTraceSpan(
+      {
+        name: 'generateContentStream',
+        noAutoEnd: true,
+      },
+      async ({ metadata: spanMetadata, endSpan }) => {
+        spanMetadata.input = { request: req, userPromptId, model: req.model };
+        const startTime = Date.now();
+        this.logApiRequest(toContents(req.contents), req.model, userPromptId);
 
-    let stream: AsyncGenerator<GenerateContentResponse>;
-    try {
-      stream = await this.wrapped.generateContentStream(req, userPromptId);
-    } catch (error) {
-      const durationMs = Date.now() - startTime;
-      this._logApiError(durationMs, error, req.model, userPromptId);
-      throw error;
-    }
+        let stream: AsyncGenerator<GenerateContentResponse>;
+        try {
+          stream = await this.wrapped.generateContentStream(req, userPromptId);
+        } catch (error) {
+          const durationMs = Date.now() - startTime;
+          this._logApiError(durationMs, error, req.model, userPromptId);
+          throw error;
+        }
 
-    return this.loggingStreamWrapper(
-      stream,
-      startTime,
-      userPromptId,
-      req.model,
+        return this.loggingStreamWrapper(
+          stream,
+          startTime,
+          userPromptId,
+          req.model,
+          spanMetadata,
+          endSpan,
+        );
+      },
     );
   }
 
@@ -156,6 +184,8 @@ export class LoggingContentGenerator implements ContentGenerator {
     startTime: number,
     userPromptId: string,
     model: string,
+    spanMetadata: SpanMetadata,
+    endSpan: () => void,
   ): AsyncGenerator<GenerateContentResponse> {
     const responses: GenerateContentResponse[] = [];
 
@@ -177,7 +207,15 @@ export class LoggingContentGenerator implements ContentGenerator {
         lastUsageMetadata,
         JSON.stringify(responses),
       );
+      spanMetadata.output = {
+        streamChunks: responses.map((r) => ({
+          content: r.candidates?.[0]?.content ?? null,
+        })),
+        usageMetadata: lastUsageMetadata,
+        durationMs,
+      };
     } catch (error) {
+      spanMetadata.error = error;
       const durationMs = Date.now() - startTime;
       this._logApiError(
         durationMs,
@@ -186,6 +224,8 @@ export class LoggingContentGenerator implements ContentGenerator {
         userPromptId,
       );
       throw error;
+    } finally {
+      endSpan();
     }
   }
 
@@ -196,6 +236,16 @@ export class LoggingContentGenerator implements ContentGenerator {
   async embedContent(
     req: EmbedContentParameters,
   ): Promise<EmbedContentResponse> {
-    return this.wrapped.embedContent(req);
+    return runInDevTraceSpan(
+      {
+        name: 'embedContent',
+      },
+      async ({ metadata: spanMetadata }) => {
+        spanMetadata.input = { request: req };
+        const output = await this.wrapped.embedContent(req);
+        spanMetadata.output = output;
+        return output;
+      },
+    );
   }
 }

packages/core/src/telemetry/index.ts

@@ -137,3 +137,4 @@ export {
   GenAiProviderName,
   GenAiTokenType,
 } from './metrics.js';
+export { runInDevTraceSpan, type SpanMetadata } from './trace.js';

packages/core/src/telemetry/trace.ts

@@ -0,0 +1,153 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import {
+  diag,
+  SpanStatusCode,
+  trace,
+  type AttributeValue,
+  type SpanOptions,
+} from '@opentelemetry/api';
+import { safeJsonStringify } from '../utils/safeJsonStringify.js';
+
+const TRACER_NAME = 'gemini-cli';
+const TRACER_VERSION = 'v1';
+
+/**
+ * Metadata for a span.
+ */
+export interface SpanMetadata {
+  /** The name of the span. */
+  name: string;
+  /** The input to the span. */
+  input?: unknown;
+  /** The output of the span. */
+  output?: unknown;
+  error?: unknown;
+  /** Additional attributes for the span. */
+  attributes: Record<string, AttributeValue>;
+}
+
+/**
+ * Runs a function in a new OpenTelemetry span.
+ *
+ * The `meta` object will be automatically used to set the span's status and attributes upon completion.
+ *
+ * @example
+ * ```typescript
+ * runInDevTraceSpan({ name: 'my-operation' }, ({ metadata }) => {
+ *   metadata.input = { foo: 'bar' };
+ *   // ... do work ...
+ *   metadata.output = { result: 'baz' };
+ *   metadata.attributes['my.custom.attribute'] = 'some-value';
+ * });
+ * ```
+ *
+ * @param opts The options for the span.
+ * @param fn The function to run in the span.
+ * @returns The result of the function.
+ */
+export async function runInDevTraceSpan<R>(
+  opts: SpanOptions & { name: string; noAutoEnd?: boolean },
+  fn: ({
+    metadata,
+  }: {
+    metadata: SpanMetadata;
+    endSpan: () => void;
+  }) => Promise<R>,
+): Promise<R> {
+  const { name: spanName, noAutoEnd, ...restOfSpanOpts } = opts;
+  if (process.env['GEMINI_DEV_TRACING'] !== 'true') {
+    // If GEMINI_DEV_TRACING env var not set, we do not trace.
+    return await fn({
+      metadata: {
+        name: spanName,
+        attributes: {},
+      },
+      endSpan: () => {
+        // noop
+      },
+    });
+  }
+
+  const tracer = trace.getTracer(TRACER_NAME, TRACER_VERSION);
+  return await tracer.startActiveSpan(
+    opts.name,
+    restOfSpanOpts,
+    async (span) => {
+      const meta: SpanMetadata = {
+        name: spanName,
+        attributes: {},
+      };
+      const endSpan = () => {
+        try {
+          if (meta.input !== undefined) {
+            span.setAttribute('input-json', safeJsonStringify(meta.input));
+          }
+          if (meta.output !== undefined) {
+            span.setAttribute('output-json', safeJsonStringify(meta.output));
+          }
+          for (const [key, value] of Object.entries(meta.attributes)) {
+            span.setAttribute(key, value as AttributeValue);
+          }
+          if (meta.error) {
+            span.setStatus({
+              code: SpanStatusCode.ERROR,
+              message: getErrorMessage(meta.error),
+            });
+            if (meta.error instanceof Error) {
+              span.recordException(meta.error);
+            }
+          } else {
+            span.setStatus({ code: SpanStatusCode.OK });
+          }
+        } catch (e) {
+          // Log the error but don't rethrow, to ensure span.end() is called.
+          diag.error('Error setting span attributes in endSpan', e);
+          span.setStatus({
+            code: SpanStatusCode.ERROR,
+            message: `Error in endSpan: ${getErrorMessage(e)}`,
+          });
+        } finally {
+          span.end();
+        }
+      };
+      try {
+        return await fn({ metadata: meta, endSpan });
+      } 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.
+          endSpan();
+        }
+      }
+    },
+  );
+}
+
+/**
+ * Gets the error message from an error object.
+ *
+ * @param e The error object.
+ * @returns The error message.
+ */
+function getErrorMessage(e: unknown): string {
+  if (e instanceof Error) {
+    return e.message;
+  }
+  if (typeof e === 'string') {
+    return e;
+  }
+  return safeJsonStringify(e);
+}