MediaMetz/gemini-cli
1
0
Fork 0
mirror of https://github.com/google-gemini/gemini-cli.git synced 2026-04-20 10:10:56 -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/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;
},
);
}