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>
This commit is contained in:
heaventourist 2026-02-26 18:26:16 -08:00 committed by GitHub
parent 4b7ce1fe67
commit b1befee8fb
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
21 changed files with 903 additions and 136 deletions

View file

@ -176,11 +176,12 @@ Sends telemetry directly to Google Cloud services. No collector needed.
}
```
2. Run Gemini CLI and send prompts.
3. View logs and metrics:
3. View logs, metrics, and traces:
- Open the Google Cloud Console in your browser after sending prompts:
- Logs: https://console.cloud.google.com/logs/
- Metrics: https://console.cloud.google.com/monitoring/metrics-explorer
- Traces: https://console.cloud.google.com/traces/list
- Logs (Logs Explorer): https://console.cloud.google.com/logs/
- Metrics (Metrics Explorer):
https://console.cloud.google.com/monitoring/metrics-explorer
- Traces (Trace Explorer): https://console.cloud.google.com/traces/list
### Collector-based export (advanced)
@ -208,11 +209,12 @@ forward data to Google Cloud.
- Save collector logs to `~/.gemini/tmp/<projectHash>/otel/collector-gcp.log`
- Stop collector on exit (e.g. `Ctrl+C`)
3. Run Gemini CLI and send prompts.
4. View logs and metrics:
4. View logs, metrics, and traces:
- Open the Google Cloud Console in your browser after sending prompts:
- Logs: https://console.cloud.google.com/logs/
- Metrics: https://console.cloud.google.com/monitoring/metrics-explorer
- Traces: https://console.cloud.google.com/traces/list
- Logs (Logs Explorer): https://console.cloud.google.com/logs/
- Metrics (Metrics Explorer):
https://console.cloud.google.com/monitoring/metrics-explorer
- Traces (Trace Explorer): https://console.cloud.google.com/traces/list
- Open `~/.gemini/tmp/<projectHash>/otel/collector-gcp.log` to view local
collector logs.
@ -270,10 +272,10 @@ For local development and debugging, you can capture telemetry data locally:
3. View traces at http://localhost:16686 and logs/metrics in the collector log
file.
## Logs and metrics
## Logs, metrics, and traces
The following section describes the structure of logs and metrics generated for
Gemini CLI.
The following section describes the structure of logs, metrics, and traces
generated for Gemini CLI.
The `session.id`, `installation.id`, `active_approval_mode`, and `user.email`
(available only when authenticated with a Google account) are included as common
@ -824,6 +826,32 @@ Optional performance monitoring for startup, CPU/memory, and phase timing.
- `current_value` (number)
- `baseline_value` (number)
### Traces
Traces offer a granular, "under-the-hood" view of every agent and backend
operation. By providing a high-fidelity execution map, they enable precise
debugging of complex tool interactions and deep performance optimization. Each
trace captures rich, consistent metadata via custom span attributes:
- `gen_ai.operation.name` (string): The high-level operation kind (e.g.
"tool_call", "llm_call").
- `gen_ai.agent.name` (string): The service agent identifier ("gemini-cli").
- `gen_ai.agent.description` (string): The service agent description.
- `gen_ai.input.messages` (string): Input messages or metadata specific to the
operation.
- `gen_ai.output.messages` (string): Output messages or metadata generated from
the operation.
- `gen_ai.request.model` (string): The request model name.
- `gen_ai.response.model` (string): The response model name.
- `gen_ai.system_instructions` (json string): The system instructions.
- `gen_ai.prompt.name` (string): The prompt name.
- `gen_ai.tool.name` (string): The executed tool's name.
- `gen_ai.tool.call_id` (string): The generated specific ID of the tool call.
- `gen_ai.tool.description` (string): The executed tool's description.
- `gen_ai.tool.definitions` (json string): The executed tool's description.
- `gen_ai.conversation.id` (string): The current CLI session ID.
- Additional user-defined Custom Attributes passed via the span's configuration.
#### GenAI semantic convention
The following metrics comply with [OpenTelemetry GenAI semantic conventions] for

View file

@ -1,23 +1,21 @@
# Local development guide
This guide provides instructions for setting up and using local development
features, such as development tracing.
features, such as tracing.
## Development tracing
## Tracing
Development traces (dev traces) are OpenTelemetry (OTel) traces that help you
debug your code by instrumenting interesting events like model calls, tool
scheduler, tool calls, etc.
Traces are OpenTelemetry (OTel) records that help you debug your code by
instrumenting key events like model calls, tool scheduler operations, and tool
calls.
Dev traces are verbose and are specifically meant for understanding agent
behavior and debugging issues. They are disabled by default.
Traces provide deep visibility into agent behavior and are invaluable for
debugging complex issues. They are captured automatically when telemetry is
enabled.
To enable dev traces, set the `GEMINI_DEV_TRACING=true` environment variable
when running Gemini CLI.
### Viewing traces
### Viewing dev traces
You can view dev traces using either Jaeger or the Genkit Developer UI.
You can view traces using either Jaeger or the Genkit Developer UI.
#### Using Genkit
@ -37,13 +35,12 @@ Genkit provides a web-based UI for viewing traces and other telemetry data.
Genkit Developer UI: http://localhost:4000
```
2. **Run Gemini CLI with dev tracing:**
2. **Run Gemini CLI:**
In a separate terminal, run your Gemini CLI command with the
`GEMINI_DEV_TRACING` environment variable:
In a separate terminal, run your Gemini CLI command:
```bash
GEMINI_DEV_TRACING=true gemini
gemini
```
3. **View the traces:**
@ -53,7 +50,7 @@ Genkit provides a web-based UI for viewing traces and other telemetry data.
#### Using Jaeger
You can view dev traces in the Jaeger UI. To get started, follow these steps:
You can view traces in the Jaeger UI. To get started, follow these steps:
1. **Start the telemetry collector:**
@ -67,13 +64,12 @@ You can view dev traces in the Jaeger UI. To get started, follow these steps:
This command also configures your workspace for local telemetry and provides
a link to the Jaeger UI (usually `http://localhost:16686`).
2. **Run Gemini CLI with dev tracing:**
2. **Run Gemini CLI:**
In a separate terminal, run your Gemini CLI command with the
`GEMINI_DEV_TRACING` environment variable:
In a separate terminal, run your Gemini CLI command:
```bash
GEMINI_DEV_TRACING=true gemini
gemini
```
3. **View the traces:**
@ -84,10 +80,10 @@ You can view dev traces in the Jaeger UI. To get started, follow these steps:
For more detailed information on telemetry, see the
[telemetry documentation](./cli/telemetry.md).
### Instrumenting code with dev traces
### Instrumenting code with traces
You can add dev traces to your own code for more detailed instrumentation. This
is useful for debugging and understanding the flow of execution.
You can add traces to your own code for more detailed instrumentation. This is
useful for debugging and understanding the flow of execution.
Use the `runInDevTraceSpan` function to wrap any section of code in a trace
span.
@ -96,29 +92,39 @@ Here is a basic example:
```typescript
import { runInDevTraceSpan } from '@google/gemini-cli-core';
import { GeminiCliOperation } from '@google/gemini-cli-core/lib/telemetry/constants.js';
await runInDevTraceSpan({ name: 'my-custom-span' }, async ({ metadata }) => {
// The `metadata` object allows you to record the input and output of the
// operation as well as other attributes.
metadata.input = { key: 'value' };
// Set custom attributes.
metadata.attributes['gen_ai.request.model'] = 'gemini-4.0-mega';
await runInDevTraceSpan(
{
operation: GeminiCliOperation.ToolCall,
attributes: {
[GEN_AI_AGENT_NAME]: 'gemini-cli',
},
},
async ({ metadata }) => {
// The `metadata` object allows you to record the input and output of the
// operation as well as other attributes.
metadata.input = { key: 'value' };
// Set custom attributes.
metadata.attributes['custom.attribute'] = 'custom.value';
// Your code to be traced goes here
try {
const output = await somethingRisky();
metadata.output = output;
return output;
} catch (e) {
metadata.error = e;
throw e;
}
});
// Your code to be traced goes here
try {
const output = await somethingRisky();
metadata.output = output;
return output;
} catch (e) {
metadata.error = e;
throw e;
}
},
);
```
In this example:
- `name`: The name of the span, which will be displayed in the trace.
- `operation`: The operation type of the span, represented by the
`GeminiCliOperation` enum.
- `metadata.input`: (Optional) An object containing the input data for the
traced operation.
- `metadata.output`: (Optional) An object containing the output data from the

View file

@ -72,7 +72,6 @@ describe('ACP telemetry', () => {
GEMINI_TELEMETRY_ENABLED: 'true',
GEMINI_TELEMETRY_TARGET: 'local',
GEMINI_TELEMETRY_OUTFILE: telemetryPath,
// GEMINI_DEV_TRACING not set: fake responses aren't instrumented for spans
},
},
);

View file

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

View file

@ -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;

View file

@ -23,10 +23,30 @@ import type {
ToolInvocation,
ToolResult,
} from '../tools/tools.js';
import {
GeminiCliOperation,
GEN_AI_AGENT_DESCRIPTION,
GEN_AI_AGENT_NAME,
} from '../telemetry/constants.js';
import type { ToolRegistry } from 'src/tools/tool-registry.js';
vi.mock('./subagent-tool-wrapper.js');
// Mock runInDevTraceSpan
const runInDevTraceSpan = vi.hoisted(() =>
vi.fn(async (opts, fn) => {
const metadata = { attributes: opts.attributes || {} };
return fn({
metadata,
endSpan: vi.fn(),
});
}),
);
vi.mock('../telemetry/trace.js', () => ({
runInDevTraceSpan,
}));
const MockSubagentToolWrapper = vi.mocked(SubagentToolWrapper);
const testDefinition: LocalAgentDefinition = {
@ -155,6 +175,25 @@ describe('SubAgentInvocation', () => {
abortSignal,
updateOutput,
);
expect(runInDevTraceSpan).toHaveBeenCalledWith(
expect.objectContaining({
operation: GeminiCliOperation.AgentCall,
attributes: expect.objectContaining({
[GEN_AI_AGENT_NAME]: testDefinition.name,
[GEN_AI_AGENT_DESCRIPTION]: testDefinition.description,
}),
}),
expect.any(Function),
);
// Verify metadata was set on the span
const spanCallback = vi.mocked(runInDevTraceSpan).mock.calls[0][1];
const mockMetadata = { input: undefined, output: undefined };
const mockSpan = { metadata: mockMetadata, endSpan: vi.fn() };
await spanCallback(mockSpan as Parameters<typeof spanCallback>[0]);
expect(mockMetadata.input).toBe(params);
expect(mockMetadata.output).toBe(mockResult);
});
describe('withUserHints', () => {

View file

@ -20,6 +20,12 @@ import type { AgentDefinition, AgentInputs } from './types.js';
import { SubagentToolWrapper } from './subagent-tool-wrapper.js';
import { SchemaValidator } from '../utils/schemaValidator.js';
import { formatUserHintsForModel } from '../utils/fastAckHelper.js';
import { runInDevTraceSpan } from '../telemetry/trace.js';
import {
GeminiCliOperation,
GEN_AI_AGENT_DESCRIPTION,
GEN_AI_AGENT_NAME,
} from '../telemetry/constants.js';
export class SubagentTool extends BaseDeclarativeTool<AgentInputs, ToolResult> {
constructor(
@ -167,7 +173,21 @@ class SubAgentInvocation extends BaseToolInvocation<AgentInputs, ToolResult> {
this.withUserHints(this.params),
);
return invocation.execute(signal, updateOutput);
return runInDevTraceSpan(
{
operation: GeminiCliOperation.AgentCall,
attributes: {
[GEN_AI_AGENT_NAME]: this.definition.name,
[GEN_AI_AGENT_DESCRIPTION]: this.definition.description,
},
},
async ({ metadata }) => {
metadata.input = this.params;
const result = await invocation.execute(signal, updateOutput);
metadata.output = result;
return result;
},
);
}
private withUserHints(agentArgs: AgentInputs): AgentInputs {

View file

@ -14,16 +14,14 @@ import {
type ErroredToolCall,
CoreToolCallStatus,
} from '../scheduler/types.js';
import type {
ToolCallConfirmationDetails,
ToolConfirmationPayload,
ToolInvocation,
ToolResult,
Config,
ToolRegistry,
MessageBus,
} from '../index.js';
import {
type ToolCallConfirmationDetails,
type ToolConfirmationPayload,
type ToolInvocation,
type ToolResult,
type Config,
type ToolRegistry,
type MessageBus,
DEFAULT_TRUNCATE_TOOL_OUTPUT_THRESHOLD,
BaseDeclarativeTool,
BaseToolInvocation,
@ -33,6 +31,8 @@ import {
HookSystem,
PolicyDecision,
ToolErrorType,
DiscoveredMCPTool,
GeminiCliOperation,
} from '../index.js';
import { createMockMessageBus } from '../test-utils/mock-message-bus.js';
import {
@ -43,12 +43,22 @@ import {
import * as modifiableToolModule from '../tools/modifiable-tool.js';
import { DEFAULT_GEMINI_MODEL } from '../config/models.js';
import type { PolicyEngine } from '../policy/policy-engine.js';
import { DiscoveredMCPTool } from '../tools/mcp-tool.js';
import { runInDevTraceSpan, type SpanMetadata } from '../telemetry/trace.js';
vi.mock('fs/promises', () => ({
writeFile: vi.fn(),
}));
vi.mock('../telemetry/trace.js', () => ({
runInDevTraceSpan: vi.fn(async (opts, fn) => {
const metadata = { attributes: opts.attributes || {} };
return fn({
metadata,
endSpan: vi.fn(),
});
}),
}));
class TestApprovalTool extends BaseDeclarativeTool<{ id: string }, ToolResult> {
static readonly Name = 'testApprovalTool';
@ -362,6 +372,21 @@ describe('CoreToolScheduler', () => {
const completedCalls = onAllToolCallsComplete.mock
.calls[0][0] as ToolCall[];
expect(completedCalls[0].status).toBe(CoreToolCallStatus.Cancelled);
expect(runInDevTraceSpan).toHaveBeenCalledWith(
expect.objectContaining({
operation: GeminiCliOperation.ScheduleToolCalls,
}),
expect.any(Function),
);
const spanArgs = vi.mocked(runInDevTraceSpan).mock.calls[0];
const fn = spanArgs[1];
const metadata: SpanMetadata = { name: '', attributes: {} };
await fn({ metadata, endSpan: vi.fn() });
expect(metadata).toMatchObject({
input: [request],
});
});
it('should cancel all tools when cancelAll is called', async () => {

View file

@ -46,6 +46,7 @@ import { CoreToolCallStatus } from '../scheduler/types.js';
import { ToolExecutor } from '../scheduler/tool-executor.js';
import { DiscoveredMCPTool } from '../tools/mcp-tool.js';
import { getPolicyDenialError } from '../scheduler/policy.js';
import { GeminiCliOperation } from '../telemetry/constants.js';
export type {
ToolCall,
@ -424,7 +425,7 @@ export class CoreToolScheduler {
signal: AbortSignal,
): Promise<void> {
return runInDevTraceSpan(
{ name: 'schedule' },
{ operation: GeminiCliOperation.ScheduleToolCalls },
async ({ metadata: spanMetadata }) => {
spanMetadata.input = request;
if (this.isRunning() || this.isScheduling) {

View file

@ -15,7 +15,13 @@ vi.mock('../telemetry/loggers.js', () => ({
}));
const runInDevTraceSpan = vi.hoisted(() =>
vi.fn(async (meta, fn) => fn({ metadata: {}, endSpan: vi.fn() })),
vi.fn(async (opts, fn) => {
const metadata = { attributes: opts.attributes || {} };
return fn({
metadata,
endSpan: vi.fn(),
});
}),
);
vi.mock('../telemetry/trace.js', () => ({
@ -38,6 +44,16 @@ import type { Config } from '../config/config.js';
import { UserTierId } from '../code_assist/types.js';
import { ApiRequestEvent, LlmRole } from '../telemetry/types.js';
import { FatalAuthenticationError } from '../utils/errors.js';
import {
GeminiCliOperation,
GEN_AI_PROMPT_NAME,
GEN_AI_REQUEST_MODEL,
GEN_AI_SYSTEM_INSTRUCTIONS,
GEN_AI_TOOL_DEFINITIONS,
GEN_AI_USAGE_INPUT_TOKENS,
GEN_AI_USAGE_OUTPUT_TOKENS,
} from '../telemetry/constants.js';
import { type SpanMetadata } from '../telemetry/trace.js';
describe('LoggingContentGenerator', () => {
let wrapped: ContentGenerator;
@ -73,10 +89,20 @@ describe('LoggingContentGenerator', () => {
const req = {
contents: [{ role: 'user', parts: [{ text: 'hello' }] }],
model: 'gemini-pro',
config: {
systemInstruction: { parts: [{ text: 'system instructions' }] },
tools: [{ functionDeclarations: [{ name: 'myTool' }] }],
},
};
const userPromptId = 'prompt-123';
const response: GenerateContentResponse = {
candidates: [],
candidates: [
{
content: {
parts: [{ text: 'hello' }],
},
},
],
usageMetadata: {
promptTokenCount: 1,
candidatesTokenCount: 2,
@ -113,12 +139,47 @@ describe('LoggingContentGenerator', () => {
);
const responseEvent = vi.mocked(logApiResponse).mock.calls[0][1];
expect(responseEvent.duration_ms).toBe(1000);
expect(runInDevTraceSpan).toHaveBeenCalledWith(
expect.objectContaining({
operation: GeminiCliOperation.LLMCall,
attributes: expect.objectContaining({
[GEN_AI_REQUEST_MODEL]: 'gemini-pro',
[GEN_AI_PROMPT_NAME]: userPromptId,
[GEN_AI_SYSTEM_INSTRUCTIONS]: JSON.stringify(
req.config.systemInstruction,
),
[GEN_AI_TOOL_DEFINITIONS]: JSON.stringify(req.config.tools),
}),
}),
expect.any(Function),
);
const spanArgs = vi.mocked(runInDevTraceSpan).mock.calls[0];
const fn = spanArgs[1];
const metadata: SpanMetadata = { name: '', attributes: {} };
await fn({ metadata, endSpan: vi.fn() });
expect(metadata).toMatchObject({
input: req.contents,
output: response.candidates?.[0]?.content,
attributes: {
[GEN_AI_USAGE_INPUT_TOKENS]: 1,
[GEN_AI_USAGE_OUTPUT_TOKENS]: 2,
},
});
});
it('should log error on failure', async () => {
const req = {
contents: [{ role: 'user', parts: [{ text: 'hello' }] }],
model: 'gemini-pro',
config: {
systemInstruction: {
parts: [{ text: 'stream system instructions' }],
},
tools: [{ functionDeclarations: [{ name: 'streamTool' }] }],
},
};
const userPromptId = 'prompt-123';
const error = new Error('test error');
@ -126,7 +187,7 @@ describe('LoggingContentGenerator', () => {
const startTime = new Date('2025-01-01T00:00:00.000Z');
vi.setSystemTime(startTime);
const promise = loggingContentGenerator.generateContent(
let promise = loggingContentGenerator.generateContent(
req,
userPromptId,
LlmRole.MAIN,
@ -142,6 +203,32 @@ describe('LoggingContentGenerator', () => {
);
const errorEvent = vi.mocked(logApiError).mock.calls[0][1];
expect(errorEvent.duration_ms).toBe(1000);
expect(runInDevTraceSpan).toHaveBeenCalledWith(
expect.objectContaining({
operation: GeminiCliOperation.LLMCall,
attributes: expect.objectContaining({
[GEN_AI_REQUEST_MODEL]: 'gemini-pro',
[GEN_AI_PROMPT_NAME]: userPromptId,
[GEN_AI_SYSTEM_INSTRUCTIONS]: JSON.stringify(
req.config.systemInstruction,
),
[GEN_AI_TOOL_DEFINITIONS]: JSON.stringify(req.config.tools),
}),
}),
expect.any(Function),
);
const spanArgs = vi.mocked(runInDevTraceSpan).mock.calls[0];
const fn = spanArgs[1];
const metadata: SpanMetadata = { name: '', attributes: {} };
promise = fn({ metadata, endSpan: vi.fn() });
await expect(promise).rejects.toThrow(error);
expect(metadata).toMatchObject({
error,
});
});
describe('error type extraction', () => {
@ -163,10 +250,22 @@ describe('LoggingContentGenerator', () => {
const req = {
contents: [{ role: 'user', parts: [{ text: 'hello' }] }],
model: 'gemini-pro',
config: {
systemInstruction: {
parts: [{ text: 'stream system instructions' }],
},
tools: [{ functionDeclarations: [{ name: 'streamTool' }] }],
},
};
const userPromptId = 'prompt-123';
const response = {
candidates: [],
candidates: [
{
content: {
parts: [{ text: 'hello' }],
},
},
],
usageMetadata: {
promptTokenCount: 1,
candidatesTokenCount: 2,
@ -186,7 +285,7 @@ describe('LoggingContentGenerator', () => {
vi.setSystemTime(startTime);
const stream = await loggingContentGenerator.generateContentStream(
let stream = await loggingContentGenerator.generateContentStream(
req,
userPromptId,
@ -211,6 +310,44 @@ describe('LoggingContentGenerator', () => {
);
const responseEvent = vi.mocked(logApiResponse).mock.calls[0][1];
expect(responseEvent.duration_ms).toBe(1000);
expect(runInDevTraceSpan).toHaveBeenCalledWith(
expect.objectContaining({
operation: GeminiCliOperation.LLMCall,
noAutoEnd: true,
attributes: expect.objectContaining({
[GEN_AI_REQUEST_MODEL]: 'gemini-pro',
[GEN_AI_PROMPT_NAME]: userPromptId,
[GEN_AI_SYSTEM_INSTRUCTIONS]: JSON.stringify(
req.config.systemInstruction,
),
[GEN_AI_TOOL_DEFINITIONS]: JSON.stringify(req.config.tools),
}),
}),
expect.any(Function),
);
const spanArgs = vi.mocked(runInDevTraceSpan).mock.calls[0];
const fn = spanArgs[1];
const metadata: SpanMetadata = { name: '', attributes: {} };
vi.mocked(wrapped.generateContentStream).mockResolvedValue(
createAsyncGenerator(),
);
stream = await fn({ metadata, endSpan: vi.fn() });
for await (const _ of stream) {
// consume stream
}
expect(metadata).toMatchObject({
input: req.contents,
output: [response.candidates?.[0]?.content],
attributes: {
[GEN_AI_USAGE_INPUT_TOKENS]: 1,
[GEN_AI_USAGE_OUTPUT_TOKENS]: 2,
},
});
});
it('should log error on failure', async () => {
@ -328,6 +465,9 @@ describe('LoggingContentGenerator', () => {
const req = {
contents: [{ role: 'user', parts: [] }],
model: 'gemini-pro',
config: {
mimeType: 'text/plain',
},
};
const response: EmbedContentResponse = { embeddings: [{ values: [] }] };
vi.mocked(wrapped.embedContent).mockResolvedValue(response);
@ -336,6 +476,26 @@ describe('LoggingContentGenerator', () => {
expect(wrapped.embedContent).toHaveBeenCalledWith(req);
expect(result).toBe(response);
expect(runInDevTraceSpan).toHaveBeenCalledWith(
expect.objectContaining({
operation: GeminiCliOperation.LLMCall,
attributes: expect.objectContaining({
[GEN_AI_REQUEST_MODEL]: req.model,
}),
}),
expect.any(Function),
);
const spanArgs = vi.mocked(runInDevTraceSpan).mock.calls[0];
const fn = spanArgs[1];
const metadata: SpanMetadata = { name: '', attributes: {} };
await fn({ metadata, endSpan: vi.fn() });
expect(metadata).toMatchObject({
input: req.contents,
output: response,
});
});
});

View file

@ -37,6 +37,16 @@ import { isStructuredError } from '../utils/quotaErrorDetection.js';
import { runInDevTraceSpan, type SpanMetadata } from '../telemetry/trace.js';
import { debugLogger } from '../utils/debugLogger.js';
import { getErrorType } from '../utils/errors.js';
import {
GeminiCliOperation,
GEN_AI_PROMPT_NAME,
GEN_AI_REQUEST_MODEL,
GEN_AI_SYSTEM_INSTRUCTIONS,
GEN_AI_TOOL_DEFINITIONS,
GEN_AI_USAGE_INPUT_TOKENS,
GEN_AI_USAGE_OUTPUT_TOKENS,
} from '../telemetry/constants.js';
import { safeJsonStringify } from '../utils/safeJsonStringify.js';
import { isMcpToolName } from '../tools/mcp-tool.js';
import { estimateTokenCountSync } from '../utils/tokenCalculation.js';
@ -303,10 +313,18 @@ export class LoggingContentGenerator implements ContentGenerator {
): Promise<GenerateContentResponse> {
return runInDevTraceSpan(
{
name: 'generateContent',
operation: GeminiCliOperation.LLMCall,
attributes: {
[GEN_AI_REQUEST_MODEL]: req.model,
[GEN_AI_PROMPT_NAME]: userPromptId,
[GEN_AI_SYSTEM_INSTRUCTIONS]: safeJsonStringify(
req.config?.systemInstruction ?? [],
),
[GEN_AI_TOOL_DEFINITIONS]: safeJsonStringify(req.config?.tools ?? []),
},
},
async ({ metadata: spanMetadata }) => {
spanMetadata.input = { request: req, userPromptId, model: req.model };
spanMetadata.input = req.contents;
const startTime = Date.now();
const contents: Content[] = toContents(req.contents);
@ -326,10 +344,11 @@ export class LoggingContentGenerator implements ContentGenerator {
userPromptId,
role,
);
spanMetadata.output = {
response,
usageMetadata: response.usageMetadata,
};
spanMetadata.output = response.candidates?.[0]?.content ?? null;
spanMetadata.attributes[GEN_AI_USAGE_INPUT_TOKENS] =
response.usageMetadata?.promptTokenCount ?? 0;
spanMetadata.attributes[GEN_AI_USAGE_OUTPUT_TOKENS] =
response.usageMetadata?.candidatesTokenCount ?? 0;
const durationMs = Date.now() - startTime;
this._logApiResponse(
contents,
@ -355,6 +374,7 @@ export class LoggingContentGenerator implements ContentGenerator {
.catch((e) => debugLogger.debug('quota refresh failed', e));
return response;
} catch (error) {
spanMetadata.error = error;
const durationMs = Date.now() - startTime;
this._logApiError(
durationMs,
@ -379,11 +399,20 @@ export class LoggingContentGenerator implements ContentGenerator {
): Promise<AsyncGenerator<GenerateContentResponse>> {
return runInDevTraceSpan(
{
name: 'generateContentStream',
operation: GeminiCliOperation.LLMCall,
noAutoEnd: true,
attributes: {
[GEN_AI_REQUEST_MODEL]: req.model,
[GEN_AI_PROMPT_NAME]: userPromptId,
[GEN_AI_SYSTEM_INSTRUCTIONS]: safeJsonStringify(
req.config?.systemInstruction ?? [],
),
[GEN_AI_TOOL_DEFINITIONS]: safeJsonStringify(req.config?.tools ?? []),
},
},
async ({ metadata: spanMetadata, endSpan }) => {
spanMetadata.input = { request: req, userPromptId, model: req.model };
spanMetadata.input = req.contents;
const startTime = Date.now();
const serverDetails = this._getEndpointUrl(
req,
@ -488,13 +517,15 @@ export class LoggingContentGenerator implements ContentGenerator {
this.config
.refreshUserQuotaIfStale()
.catch((e) => debugLogger.debug('quota refresh failed', e));
spanMetadata.output = {
streamChunks: responses.map((r) => ({
content: r.candidates?.[0]?.content ?? null,
})),
usageMetadata: lastUsageMetadata,
durationMs,
};
spanMetadata.output = responses.map(
(response) => response.candidates?.[0]?.content ?? null,
);
if (lastUsageMetadata) {
spanMetadata.attributes[GEN_AI_USAGE_INPUT_TOKENS] =
lastUsageMetadata.promptTokenCount ?? 0;
spanMetadata.attributes[GEN_AI_USAGE_OUTPUT_TOKENS] =
lastUsageMetadata.candidatesTokenCount ?? 0;
}
} catch (error) {
spanMetadata.error = error;
const durationMs = Date.now() - startTime;
@ -523,10 +554,13 @@ export class LoggingContentGenerator implements ContentGenerator {
): Promise<EmbedContentResponse> {
return runInDevTraceSpan(
{
name: 'embedContent',
operation: GeminiCliOperation.LLMCall,
attributes: {
[GEN_AI_REQUEST_MODEL]: req.model,
},
},
async ({ metadata: spanMetadata }) => {
spanMetadata.input = { request: req };
spanMetadata.input = req.contents;
const output = await this.wrapped.embedContent(req);
spanMetadata.output = output;
return output;

View file

@ -20,10 +20,18 @@ vi.mock('node:crypto', () => ({
randomUUID: vi.fn(),
}));
const runInDevTraceSpan = vi.hoisted(() =>
vi.fn(async (opts, fn) => {
const metadata = { attributes: opts.attributes || {} };
return fn({
metadata,
endSpan: vi.fn(),
});
}),
);
vi.mock('../telemetry/trace.js', () => ({
runInDevTraceSpan: vi.fn(async (_opts, fn) =>
fn({ metadata: { input: {}, output: {} } }),
),
runInDevTraceSpan,
}));
import { logToolCall } from '../telemetry/loggers.js';
@ -81,6 +89,7 @@ import type {
} from './types.js';
import { CoreToolCallStatus, ROOT_SCHEDULER_ID } from './types.js';
import { ToolErrorType } from '../tools/tool-error.js';
import { GeminiCliOperation } from '../telemetry/constants.js';
import * as ToolUtils from '../utils/tool-utils.js';
import type { EditorType } from '../utils/editor.js';
import {
@ -366,6 +375,21 @@ describe('Scheduler (Orchestrator)', () => {
}),
]),
);
expect(runInDevTraceSpan).toHaveBeenCalledWith(
expect.objectContaining({
operation: GeminiCliOperation.ScheduleToolCalls,
}),
expect.any(Function),
);
const spanArgs = vi.mocked(runInDevTraceSpan).mock.calls[0];
const fn = spanArgs[1];
const metadata = { attributes: {} };
await fn({ metadata, endSpan: vi.fn() });
expect(metadata).toMatchObject({
input: [req1],
});
});
it('should set approvalMode to PLAN when config returns PLAN', async () => {

View file

@ -46,6 +46,7 @@ import {
CoreEvent,
type McpProgressPayload,
} from '../utils/events.js';
import { GeminiCliOperation } from '../telemetry/constants.js';
interface SchedulerQueueItem {
requests: ToolCallRequestInfo[];
@ -186,16 +187,22 @@ export class Scheduler {
signal: AbortSignal,
): Promise<CompletedToolCall[]> {
return runInDevTraceSpan(
{ name: 'schedule' },
{ operation: GeminiCliOperation.ScheduleToolCalls },
async ({ metadata: spanMetadata }) => {
const requests = Array.isArray(request) ? request : [request];
spanMetadata.input = requests;
let toolCallResponse: CompletedToolCall[] = [];
if (this.isProcessing || this.state.isActive) {
return this._enqueueRequest(requests, signal);
toolCallResponse = await this._enqueueRequest(requests, signal);
} else {
toolCallResponse = await this._startBatch(requests, signal);
}
return this._startBatch(requests, signal);
spanMetadata.output = toolCallResponse;
return toolCallResponse;
},
);
}

View file

@ -20,10 +20,18 @@ vi.mock('node:crypto', () => ({
randomUUID: vi.fn(),
}));
const runInDevTraceSpan = vi.hoisted(() =>
vi.fn(async (opts, fn) => {
const metadata = { name: '', attributes: opts.attributes || {} };
return fn({
metadata,
endSpan: vi.fn(),
});
}),
);
vi.mock('../telemetry/trace.js', () => ({
runInDevTraceSpan: vi.fn(async (_opts, fn) =>
fn({ metadata: { input: {}, output: {} } }),
),
runInDevTraceSpan,
}));
vi.mock('../telemetry/loggers.js', () => ({
logToolCall: vi.fn(),
@ -71,6 +79,7 @@ import type {
ToolCall,
} from './types.js';
import { ROOT_SCHEDULER_ID } from './types.js';
import { GeminiCliOperation } from '../telemetry/constants.js';
import type { EditorType } from '../utils/editor.js';
describe('Scheduler Parallel Execution', () => {
@ -306,6 +315,21 @@ describe('Scheduler Parallel Execution', () => {
);
expect(executionLog).toContain('end-call-3');
expect(runInDevTraceSpan).toHaveBeenCalledWith(
expect.objectContaining({
operation: GeminiCliOperation.ScheduleToolCalls,
}),
expect.any(Function),
);
const spanArgs = vi.mocked(runInDevTraceSpan).mock.calls[0];
const fn = spanArgs[1];
const metadata = { name: '', attributes: {} };
await fn({ metadata, endSpan: vi.fn() });
expect(metadata).toMatchObject({
input: [req1, req2, req3],
});
});
it('should execute non-read-only tools sequentially', async () => {

View file

@ -6,8 +6,11 @@
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
import { ToolExecutor } from './tool-executor.js';
import type { Config, AnyToolInvocation } from '../index.js';
import type { ToolResult } from '../tools/tools.js';
import {
type Config,
type ToolResult,
type AnyToolInvocation,
} from '../index.js';
import { makeFakeConfig } from '../test-utils/config.js';
import { MockTool } from '../test-utils/mock-tool.js';
import type { ScheduledToolCall } from './types.js';
@ -17,6 +20,12 @@ import * as fileUtils from '../utils/fileUtils.js';
import * as coreToolHookTriggers from '../core/coreToolHookTriggers.js';
import { ShellToolInvocation } from '../tools/shell.js';
import { createMockMessageBus } from '../test-utils/mock-message-bus.js';
import {
GeminiCliOperation,
GEN_AI_TOOL_CALL_ID,
GEN_AI_TOOL_DESCRIPTION,
GEN_AI_TOOL_NAME,
} from '../telemetry/constants.js';
// Mock file utils
vi.mock('../utils/fileUtils.js', () => ({
@ -28,6 +37,24 @@ vi.mock('../utils/fileUtils.js', () => ({
vi.mock('../core/coreToolHookTriggers.js', () => ({
executeToolWithHooks: vi.fn(),
}));
// Mock runInDevTraceSpan
const runInDevTraceSpan = vi.hoisted(() =>
vi.fn(async (opts, fn) => {
const metadata = { attributes: opts.attributes || {} };
return fn({
metadata,
endSpan: vi.fn(),
});
}),
);
vi.mock('../index.js', async (importOriginal) => {
const actual = await importOriginal<Record<string, unknown>>();
return {
...actual,
runInDevTraceSpan,
};
});
describe('ToolExecutor', () => {
let config: Config;
@ -57,6 +84,7 @@ describe('ToolExecutor', () => {
it('should execute a tool successfully', async () => {
const mockTool = new MockTool({
name: 'testTool',
description: 'Mock description',
execute: async () => ({
llmContent: 'Tool output',
returnDisplay: 'Tool output',
@ -97,11 +125,37 @@ describe('ToolExecutor', () => {
?.response as Record<string, unknown>;
expect(response).toEqual({ output: 'Tool output' });
}
expect(runInDevTraceSpan).toHaveBeenCalledWith(
expect.objectContaining({
operation: GeminiCliOperation.ToolCall,
attributes: expect.objectContaining({
[GEN_AI_TOOL_NAME]: 'testTool',
[GEN_AI_TOOL_CALL_ID]: 'call-1',
[GEN_AI_TOOL_DESCRIPTION]: 'Mock description',
}),
}),
expect.any(Function),
);
const spanArgs = vi.mocked(runInDevTraceSpan).mock.calls[0];
const fn = spanArgs[1];
const metadata = { attributes: {} };
await fn({ metadata, endSpan: vi.fn() });
expect(metadata).toMatchObject({
input: scheduledCall.request,
output: {
...result,
durationMs: expect.any(Number),
endTime: expect.any(Number),
},
});
});
it('should handle execution errors', async () => {
const mockTool = new MockTool({
name: 'failTool',
description: 'Mock description',
});
const invocation = mockTool.build({});
@ -134,6 +188,26 @@ describe('ToolExecutor', () => {
if (result.status === CoreToolCallStatus.Error) {
expect(result.response.error?.message).toBe('Tool Failed');
}
expect(runInDevTraceSpan).toHaveBeenCalledWith(
expect.objectContaining({
operation: GeminiCliOperation.ToolCall,
attributes: expect.objectContaining({
[GEN_AI_TOOL_NAME]: 'failTool',
[GEN_AI_TOOL_CALL_ID]: 'call-2',
[GEN_AI_TOOL_DESCRIPTION]: 'Mock description',
}),
}),
expect.any(Function),
);
const spanArgs = vi.mocked(runInDevTraceSpan).mock.calls[0];
const fn = spanArgs[1];
const metadata = { attributes: {} };
await fn({ metadata, endSpan: vi.fn() });
expect(metadata).toMatchObject({
error: new Error('Tool Failed'),
});
});
it('should return cancelled result when signal is aborted', async () => {

View file

@ -34,6 +34,12 @@ import type {
CancelledToolCall,
} from './types.js';
import { CoreToolCallStatus } from './types.js';
import {
GeminiCliOperation,
GEN_AI_TOOL_CALL_ID,
GEN_AI_TOOL_DESCRIPTION,
GEN_AI_TOOL_NAME,
} from '../telemetry/constants.js';
export interface ToolExecutionContext {
call: ToolCall;
@ -70,11 +76,17 @@ export class ToolExecutor {
return runInDevTraceSpan(
{
name: tool.name,
attributes: { type: 'tool-call' },
operation: GeminiCliOperation.ToolCall,
attributes: {
[GEN_AI_TOOL_NAME]: toolName,
[GEN_AI_TOOL_CALL_ID]: callId,
[GEN_AI_TOOL_DESCRIPTION]: tool.description,
},
},
async ({ metadata: spanMetadata }) => {
spanMetadata.input = { request };
spanMetadata.input = request;
let completedToolCall: CompletedToolCall;
try {
let promise: Promise<ToolResult>;
@ -116,21 +128,23 @@ export class ToolExecutor {
}
const toolResult: ToolResult = await promise;
spanMetadata.output = toolResult;
if (signal.aborted) {
return this.createCancelledResult(
completedToolCall = this.createCancelledResult(
call,
'User cancelled tool execution.',
);
} else if (toolResult.error === undefined) {
return await this.createSuccessResult(call, toolResult);
completedToolCall = await this.createSuccessResult(
call,
toolResult,
);
} else {
const displayText =
typeof toolResult.returnDisplay === 'string'
? toolResult.returnDisplay
: undefined;
return this.createErrorResult(
completedToolCall = this.createErrorResult(
call,
new Error(toolResult.error.message),
toolResult.error.type,
@ -141,21 +155,25 @@ export class ToolExecutor {
} catch (executionError: unknown) {
spanMetadata.error = executionError;
if (signal.aborted) {
return this.createCancelledResult(
completedToolCall = this.createCancelledResult(
call,
'User cancelled tool execution.',
);
} else {
const error =
executionError instanceof Error
? executionError
: new Error(String(executionError));
completedToolCall = this.createErrorResult(
call,
error,
ToolErrorType.UNHANDLED_EXCEPTION,
);
}
const error =
executionError instanceof Error
? executionError
: new Error(String(executionError));
return this.createErrorResult(
call,
error,
ToolErrorType.UNHANDLED_EXCEPTION,
);
}
spanMetadata.output = completedToolCall;
return completedToolCall;
},
);
}

View file

@ -5,3 +5,34 @@
*/
export const SERVICE_NAME = 'gemini-cli';
export const SERVICE_DESCRIPTION =
'Gemini CLI is an open-source AI agent that brings the power of Gemini directly into your terminal. It is designed to be a terminal-first, extensible, and powerful tool for developers, engineers, SREs, and beyond.';
// Gemini CLI specific semantic conventions
// https://opentelemetry.io/docs/specs/semconv/registry/attributes/gen-ai/#genai-attributes
export const GEN_AI_OPERATION_NAME = 'gen_ai.operation.name';
export const GEN_AI_AGENT_NAME = 'gen_ai.agent.name';
export const GEN_AI_AGENT_DESCRIPTION = 'gen_ai.agent.description';
export const GEN_AI_INPUT_MESSAGES = 'gen_ai.input.messages';
export const GEN_AI_OUTPUT_MESSAGES = 'gen_ai.output.messages';
export const GEN_AI_REQUEST_MODEL = 'gen_ai.request.model';
export const GEN_AI_RESPONSE_MODEL = 'gen_ai.response.model';
export const GEN_AI_PROMPT_NAME = 'gen_ai.prompt.name';
export const GEN_AI_TOOL_NAME = 'gen_ai.tool.name';
export const GEN_AI_TOOL_CALL_ID = 'gen_ai.tool.call_id';
export const GEN_AI_TOOL_DESCRIPTION = 'gen_ai.tool.description';
export const GEN_AI_USAGE_INPUT_TOKENS = 'gen_ai.usage.input_tokens';
export const GEN_AI_USAGE_OUTPUT_TOKENS = 'gen_ai.usage.output_tokens';
export const GEN_AI_SYSTEM_INSTRUCTIONS = 'gen_ai.system_instructions';
export const GEN_AI_TOOL_DEFINITIONS = 'gen_ai.tool.definitions';
export const GEN_AI_CONVERSATION_ID = 'gen_ai.conversation.id';
// Gemini CLI specific operations
export enum GeminiCliOperation {
ToolCall = 'tool_call',
LLMCall = 'llm_call',
UserPrompt = 'user_prompt',
SystemPrompt = 'system_prompt',
AgentCall = 'agent_call',
ScheduleToolCalls = 'schedule_tool_calls',
}

View file

@ -148,3 +148,4 @@ export {
} from './metrics.js';
export { runInDevTraceSpan, type SpanMetadata } from './trace.js';
export { startupProfiler, StartupProfiler } from './startupProfiler.js';
export * from './constants.js';

View file

@ -0,0 +1,188 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
import { trace, SpanStatusCode, diag, type Tracer } from '@opentelemetry/api';
import { runInDevTraceSpan } from './trace.js';
import {
GeminiCliOperation,
GEN_AI_CONVERSATION_ID,
GEN_AI_AGENT_DESCRIPTION,
GEN_AI_AGENT_NAME,
GEN_AI_INPUT_MESSAGES,
GEN_AI_OPERATION_NAME,
GEN_AI_OUTPUT_MESSAGES,
SERVICE_DESCRIPTION,
SERVICE_NAME,
} from './constants.js';
vi.mock('@opentelemetry/api', async (importOriginal) => {
const original = await importOriginal<typeof import('@opentelemetry/api')>();
return {
...original,
trace: {
getTracer: vi.fn(),
},
diag: {
error: vi.fn(),
},
};
});
vi.mock('../utils/session.js', () => ({
sessionId: 'test-session-id',
}));
describe('runInDevTraceSpan', () => {
const mockSpan = {
setAttribute: vi.fn(),
setStatus: vi.fn(),
recordException: vi.fn(),
end: vi.fn(),
};
const mockTracer = {
startActiveSpan: vi.fn((name, options, callback) => callback(mockSpan)),
} as unknown as Tracer;
beforeEach(() => {
vi.clearAllMocks();
vi.mocked(trace.getTracer).mockReturnValue(mockTracer);
});
afterEach(() => {
vi.unstubAllEnvs();
});
it('should start an active span', async () => {
const fn = vi.fn(async () => 'result');
const result = await runInDevTraceSpan(
{ operation: GeminiCliOperation.LLMCall },
fn,
);
expect(result).toBe('result');
expect(trace.getTracer).toHaveBeenCalled();
expect(mockTracer.startActiveSpan).toHaveBeenCalledWith(
GeminiCliOperation.LLMCall,
{},
expect.any(Function),
);
});
it('should set default attributes on the span metadata', async () => {
await runInDevTraceSpan(
{ operation: GeminiCliOperation.LLMCall },
async ({ metadata }) => {
expect(metadata.attributes[GEN_AI_OPERATION_NAME]).toBe(
GeminiCliOperation.LLMCall,
);
expect(metadata.attributes[GEN_AI_AGENT_NAME]).toBe(SERVICE_NAME);
expect(metadata.attributes[GEN_AI_AGENT_DESCRIPTION]).toBe(
SERVICE_DESCRIPTION,
);
expect(metadata.attributes[GEN_AI_CONVERSATION_ID]).toBe(
'test-session-id',
);
},
);
});
it('should set span attributes from metadata on completion', async () => {
await runInDevTraceSpan(
{ operation: GeminiCliOperation.LLMCall },
async ({ metadata }) => {
metadata.input = { query: 'hello' };
metadata.output = { response: 'world' };
metadata.attributes['custom.attr'] = 'value';
},
);
expect(mockSpan.setAttribute).toHaveBeenCalledWith(
GEN_AI_INPUT_MESSAGES,
JSON.stringify({ query: 'hello' }),
);
expect(mockSpan.setAttribute).toHaveBeenCalledWith(
GEN_AI_OUTPUT_MESSAGES,
JSON.stringify({ response: 'world' }),
);
expect(mockSpan.setAttribute).toHaveBeenCalledWith('custom.attr', 'value');
expect(mockSpan.setStatus).toHaveBeenCalledWith({
code: SpanStatusCode.OK,
});
expect(mockSpan.end).toHaveBeenCalled();
});
it('should handle errors in the wrapped function', async () => {
const error = new Error('test error');
await expect(
runInDevTraceSpan({ operation: GeminiCliOperation.LLMCall }, async () => {
throw error;
}),
).rejects.toThrow(error);
expect(mockSpan.setStatus).toHaveBeenCalledWith({
code: SpanStatusCode.ERROR,
message: 'test error',
});
expect(mockSpan.recordException).toHaveBeenCalledWith(error);
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';
},
);
expect(result).toBe('streaming');
expect(mockSpan.end).not.toHaveBeenCalled();
capturedEndSpan();
expect(mockSpan.end).toHaveBeenCalled();
});
it('should automatically end span on error even if noAutoEnd is true', async () => {
const error = new Error('streaming error');
await expect(
runInDevTraceSpan(
{ operation: GeminiCliOperation.LLMCall, noAutoEnd: true },
async () => {
throw error;
},
),
).rejects.toThrow(error);
expect(mockSpan.end).toHaveBeenCalled();
});
it('should handle exceptions in endSpan gracefully', async () => {
mockSpan.setAttribute.mockImplementation(() => {
throw new Error('attribute error');
});
await runInDevTraceSpan(
{ operation: GeminiCliOperation.LLMCall },
async ({ metadata }) => {
metadata.input = 'trigger error';
},
);
expect(diag.error).toHaveBeenCalled();
expect(mockSpan.setStatus).toHaveBeenCalledWith(
expect.objectContaining({
code: SpanStatusCode.ERROR,
message: expect.stringContaining('attribute error'),
}),
);
expect(mockSpan.end).toHaveBeenCalled();
});
});

View file

@ -12,6 +12,18 @@ import {
type SpanOptions,
} from '@opentelemetry/api';
import { safeJsonStringify } from '../utils/safeJsonStringify.js';
import {
type GeminiCliOperation,
GEN_AI_AGENT_DESCRIPTION,
GEN_AI_AGENT_NAME,
GEN_AI_CONVERSATION_ID,
GEN_AI_INPUT_MESSAGES,
GEN_AI_OPERATION_NAME,
GEN_AI_OUTPUT_MESSAGES,
SERVICE_DESCRIPTION,
SERVICE_NAME,
} from './constants.js';
import { sessionId } from '../utils/session.js';
const TRACER_NAME = 'gemini-cli';
const TRACER_VERSION = 'v1';
@ -51,7 +63,7 @@ export interface SpanMetadata {
* @returns The result of the function.
*/
export async function runInDevTraceSpan<R>(
opts: SpanOptions & { name: string; noAutoEnd?: boolean },
opts: SpanOptions & { operation: GeminiCliOperation; noAutoEnd?: boolean },
fn: ({
metadata,
}: {
@ -59,33 +71,32 @@ export async function runInDevTraceSpan<R>(
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 fn({
metadata: {
name: spanName,
attributes: {},
},
endSpan: () => {
// noop
},
});
}
const { operation, noAutoEnd, ...restOfSpanOpts } = opts;
const tracer = trace.getTracer(TRACER_NAME, TRACER_VERSION);
return tracer.startActiveSpan(opts.name, restOfSpanOpts, async (span) => {
return tracer.startActiveSpan(operation, restOfSpanOpts, async (span) => {
const meta: SpanMetadata = {
name: spanName,
attributes: {},
name: operation,
attributes: {
[GEN_AI_OPERATION_NAME]: operation,
[GEN_AI_AGENT_NAME]: SERVICE_NAME,
[GEN_AI_AGENT_DESCRIPTION]: SERVICE_DESCRIPTION,
[GEN_AI_CONVERSATION_ID]: sessionId,
},
};
const endSpan = () => {
try {
if (meta.input !== undefined) {
span.setAttribute('input-json', safeJsonStringify(meta.input));
span.setAttribute(
GEN_AI_INPUT_MESSAGES,
safeJsonStringify(meta.input),
);
}
if (meta.output !== undefined) {
span.setAttribute('output-json', safeJsonStringify(meta.output));
span.setAttribute(
GEN_AI_OUTPUT_MESSAGES,
safeJsonStringify(meta.output),
);
}
for (const [key, value] of Object.entries(meta.attributes)) {
span.setAttribute(key, value);

View file

@ -4,7 +4,7 @@
* SPDX-License-Identifier: Apache-2.0
*/
import { describe, it, expect } from 'vitest';
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
import { GeminiCliAgent } from './agent.js';
import * as path from 'node:path';
import { fileURLToPath } from 'node:url';
@ -20,6 +20,13 @@ const getGoldenPath = (name: string) =>
path.resolve(__dirname, '../test-data', `${name}.json`);
describe('GeminiCliAgent Integration', () => {
beforeEach(() => {
vi.stubEnv('GEMINI_API_KEY', 'test-api-key');
});
afterEach(() => {
vi.unstubAllEnvs();
});
it('handles static instructions', async () => {
const goldenFile = getGoldenPath('agent-static-instructions');