MediaMetz/gemini-cli
1
0
Fork 0
mirror of https://github.com/google-gemini/gemini-cli.git synced 2026-05-13 13:22:35 -07:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
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
21 changed files with 903 additions and 136 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/core/src/agents/subagent-tool.test.ts
+39
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', () => {
packages/core/src/agents/subagent-tool.ts
+21 -1
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 {
packages/core/src/core/coreToolScheduler.test.ts
+35 -10
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 () => {
packages/core/src/core/coreToolScheduler.ts
+2 -1
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) {
packages/core/src/core/loggingContentGenerator.test.ts
+165 -5
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,
});
});
});
packages/core/src/core/loggingContentGenerator.ts
+51 -17
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;
packages/core/src/scheduler/scheduler.test.ts
+27 -3
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 () => {
packages/core/src/scheduler/scheduler.ts
+10 -3
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;
},
);
}
packages/core/src/scheduler/scheduler_parallel.test.ts
+27 -3
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 () => {
packages/core/src/scheduler/tool-executor.test.ts
+76 -2
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 () => {
packages/core/src/scheduler/tool-executor.ts
+35 -17
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;
},
);
}
packages/core/src/telemetry/constants.ts
+31
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',
}
packages/core/src/telemetry/index.ts
+1
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';
packages/core/src/telemetry/trace.test.ts
+188
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();
});
});
packages/core/src/telemetry/trace.ts
+30 -19
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);