mirror of
https://github.com/google-gemini/gemini-cli
synced 2026-04-21 13:37:17 +00:00
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:
parent
4b7ce1fe67
commit
b1befee8fb
21 changed files with 903 additions and 136 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
},
|
||||
},
|
||||
);
|
||||
|
|
|
|||
|
|
@ -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 });
|
||||
});
|
||||
expect(spanMetadata.input).toBe('telemetry test query');
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
||||
|
|
|
|||
|
|
@ -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', () => {
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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 () => {
|
||||
|
|
|
|||
|
|
@ -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) {
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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 () => {
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
},
|
||||
);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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 () => {
|
||||
|
|
|
|||
|
|
@ -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 () => {
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
},
|
||||
);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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',
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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';
|
||||
|
|
|
|||
188
packages/core/src/telemetry/trace.test.ts
Normal file
188
packages/core/src/telemetry/trace.test.ts
Normal 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();
|
||||
});
|
||||
});
|
||||
|
|
@ -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);
|
||||
|
|
|
|||
|
|
@ -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');
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue