feat(telemetry) Instrument traces with more attributes and make them available to OTEL users (#20237)

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: Jerop Kipruto <jerop@google.com>
Co-authored-by: MD. MOHIBUR RAHMAN <35300157+mrpmohiburrahman@users.noreply.github.com>
Co-authored-by: Jeffrey Ying <jeffrey.ying86@live.com>
Co-authored-by: Bryan Morgan <bryanmorgan@google.com>
Co-authored-by: joshualitt <joshualitt@google.com>
Co-authored-by: Dev Randalpura <devrandalpura@google.com>
Co-authored-by: Google Admin <github-admin@google.com>
Co-authored-by: Ben Knutson <benknutson@google.com>

packages/cli/src/ui/hooks/useGeminiStream.test.tsx

@@ -25,6 +25,7 @@ import type {
   Config,
   EditorType,
   AnyToolInvocation,
+  SpanMetadata,
 } from '@google/gemini-cli-core';
 import {
   CoreToolCallStatus,
@@ -39,6 +40,7 @@ import {
   coreEvents,
   CoreEvent,
   MCPDiscoveryState,
+  GeminiCliOperation,
   getPlanModeExitMessage,
 } from '@google/gemini-cli-core';
 import type { Part, PartListUnion } from '@google/genai';
@@ -101,6 +103,19 @@ const MockValidationRequiredError = vi.hoisted(
     },
 );
 
+const mockRunInDevTraceSpan = vi.hoisted(() =>
+  vi.fn(async (opts, fn) => {
+    const metadata: SpanMetadata = {
+      name: opts.operation,
+      attributes: opts.attributes || {},
+    };
+    return await fn({
+      metadata,
+      endSpan: vi.fn(),
+    });
+  }),
+);
+
 vi.mock('@google/gemini-cli-core', async (importOriginal) => {
   const actualCoreModule = (await importOriginal()) as any;
   return {
@@ -113,6 +128,7 @@ vi.mock('@google/gemini-cli-core', async (importOriginal) => {
     tokenLimit: vi.fn().mockReturnValue(100), // Mock tokenLimit
     recordToolCallInteractions: vi.fn().mockResolvedValue(undefined),
     getCodeAssistServer: vi.fn().mockReturnValue(undefined),
+    runInDevTraceSpan: mockRunInDevTraceSpan,
   };
 });
 
@@ -794,6 +810,23 @@ describe('useGeminiStream', () => {
           item.text.includes('Got it. Focusing on tests only.'),
       ),
     ).toBe(true);
+
+    expect(mockRunInDevTraceSpan).toHaveBeenCalledWith(
+      expect.objectContaining({
+        operation: GeminiCliOperation.SystemPrompt,
+      }),
+      expect.any(Function),
+    );
+
+    const spanArgs = mockRunInDevTraceSpan.mock.calls[0];
+    const fn = spanArgs[1];
+    const metadata = { attributes: {} };
+    await act(async () => {
+      await fn({ metadata, endSpan: vi.fn() });
+    });
+    expect(metadata).toMatchObject({
+      input: sentParts,
+    });
   });
 
   it('should handle all tool calls being cancelled', async () => {
@@ -2452,6 +2485,11 @@ describe('useGeminiStream', () => {
     // This is the core fix validation: Rationale comes before tools are even scheduled (awaited)
     expect(rationaleIndex).toBeLessThan(scheduleIndex);
     expect(rationaleIndex).toBeLessThan(toolGroupIndex);
+
+    // Ensure all state updates from recursive submitQuery are settled
+    await waitFor(() => {
+      expect(result.current.streamingState).toBe(StreamingState.Idle);
+    });
   });
 
   it('should process @include commands, adding user turn after processing to prevent race conditions', async () => {
@@ -3554,4 +3592,31 @@ describe('useGeminiStream', () => {
       expect(result.current.pendingHistoryItems.length).toEqual(0);
     });
   });
+
+  it('should trace UserPrompt telemetry on submitQuery', async () => {
+    const { result } = renderTestHook();
+
+    mockSendMessageStream.mockReturnValue(
+      (async function* () {
+        yield { type: ServerGeminiEventType.Content, value: 'Response' };
+      })(),
+    );
+
+    await act(async () => {
+      await result.current.submitQuery('telemetry test query');
+    });
+
+    const userPromptCall = mockRunInDevTraceSpan.mock.calls.find(
+      (call) =>
+        call[0].operation === GeminiCliOperation.UserPrompt ||
+        call[0].operation === 'UserPrompt',
+    );
+    expect(userPromptCall).toBeDefined();
+
+    const spanMetadata = {} as SpanMetadata;
+    await act(async () => {
+      await userPromptCall![1]({ metadata: spanMetadata, endSpan: vi.fn() });
+    });
+    expect(spanMetadata.input).toBe('telemetry test query');
+  });
 });

packages/cli/src/ui/hooks/useGeminiStream.ts

@@ -36,6 +36,7 @@ import {
   CoreToolCallStatus,
   buildUserSteeringHintPrompt,
   generateSteeringAckMessage,
+  GeminiCliOperation,
   getPlanModeExitMessage,
 } from '@google/gemini-cli-core';
 import type {
@@ -1262,7 +1263,11 @@ export const useGeminiStream = (
       prompt_id?: string,
     ) =>
       runInDevTraceSpan(
-        { name: 'submitQuery' },
+        {
+          operation: options?.isContinuation
+            ? GeminiCliOperation.SystemPrompt
+            : GeminiCliOperation.UserPrompt,
+        },
         async ({ metadata: spanMetadata }) => {
           spanMetadata.input = query;