feat(telemetry): Add context breakdown to API response event (#19699)

This commit is contained in:
Sandy Tao
2026-02-24 15:26:28 -08:00
committed by GitHub
parent 70b650122f
commit 3ff5cfaaf6
8 changed files with 479 additions and 27 deletions

View File

@@ -20,8 +20,7 @@ async function run(cmd) {
stdio: ['pipe', 'pipe', 'ignore'],
});
return stdout.trim();
} catch (_e) {
// eslint-disable-line @typescript-eslint/no-unused-vars
} catch {
return null;
}
}

View File

@@ -24,11 +24,16 @@ vi.mock('../telemetry/trace.js', () => ({
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
import type {
Content,
GenerateContentConfig,
GenerateContentResponse,
EmbedContentResponse,
} from '@google/genai';
import type { ContentGenerator } from './contentGenerator.js';
import { LoggingContentGenerator } from './loggingContentGenerator.js';
import {
LoggingContentGenerator,
estimateContextBreakdown,
} from './loggingContentGenerator.js';
import type { Config } from '../config/config.js';
import { UserTierId } from '../code_assist/types.js';
import { ApiRequestEvent, LlmRole } from '../telemetry/types.js';
@@ -346,3 +351,280 @@ describe('LoggingContentGenerator', () => {
});
});
});
describe('estimateContextBreakdown', () => {
it('should return zeros for empty contents and no config', () => {
const result = estimateContextBreakdown([], undefined);
expect(result).toEqual({
system_instructions: 0,
tool_definitions: 0,
history: 0,
tool_calls: {},
mcp_servers: 0,
});
});
it('should estimate system instruction tokens', () => {
const config = {
systemInstruction: 'You are a helpful assistant.',
} as GenerateContentConfig;
const result = estimateContextBreakdown([], config);
expect(result.system_instructions).toBeGreaterThan(0);
expect(result.tool_definitions).toBe(0);
expect(result.history).toBe(0);
});
it('should estimate non-MCP tool definition tokens', () => {
const config = {
tools: [
{
functionDeclarations: [
{ name: 'read_file', description: 'Reads a file', parameters: {} },
],
},
],
} as unknown as GenerateContentConfig;
const result = estimateContextBreakdown([], config);
expect(result.tool_definitions).toBeGreaterThan(0);
expect(result.mcp_servers).toBe(0);
});
it('should classify MCP tool definitions into mcp_servers, not tool_definitions', () => {
const config = {
tools: [
{
functionDeclarations: [
{
name: 'myserver__search',
description: 'Search via MCP',
parameters: {},
},
{
name: 'read_file',
description: 'Reads a file',
parameters: {},
},
],
},
],
} as unknown as GenerateContentConfig;
const result = estimateContextBreakdown([], config);
expect(result.mcp_servers).toBeGreaterThan(0);
expect(result.tool_definitions).toBeGreaterThan(0);
// MCP tokens should not be in tool_definitions
const configOnlyBuiltin = {
tools: [
{
functionDeclarations: [
{
name: 'read_file',
description: 'Reads a file',
parameters: {},
},
],
},
],
} as unknown as GenerateContentConfig;
const builtinOnly = estimateContextBreakdown([], configOnlyBuiltin);
// tool_definitions should be smaller when MCP tools are separated out
expect(result.tool_definitions).toBeLessThan(
result.tool_definitions + result.mcp_servers,
);
expect(builtinOnly.mcp_servers).toBe(0);
});
it('should not classify tools with __ in the middle of a segment as MCP', () => {
// "__" at start or end (not a valid server__tool pattern) should not be MCP
const config = {
tools: [
{
functionDeclarations: [
{ name: '__leading', description: 'test', parameters: {} },
{ name: 'trailing__', description: 'test', parameters: {} },
{
name: 'a__b__c',
description: 'three parts - not valid MCP',
parameters: {},
},
],
},
],
} as unknown as GenerateContentConfig;
const result = estimateContextBreakdown([], config);
expect(result.mcp_servers).toBe(0);
});
it('should estimate history tokens excluding tool call/response parts', () => {
const contents: Content[] = [
{ role: 'user', parts: [{ text: 'Hello world' }] },
{ role: 'model', parts: [{ text: 'Hi there!' }] },
];
const result = estimateContextBreakdown(contents);
expect(result.history).toBeGreaterThan(0);
expect(result.tool_calls).toEqual({});
});
it('should separate tool call tokens from history', () => {
const contents: Content[] = [
{
role: 'model',
parts: [
{
functionCall: {
name: 'read_file',
args: { path: '/tmp/test.txt' },
},
},
],
},
{
role: 'function',
parts: [
{
functionResponse: {
name: 'read_file',
response: { content: 'file contents here' },
},
},
],
},
];
const result = estimateContextBreakdown(contents);
expect(result.tool_calls['read_file']).toBeGreaterThan(0);
// history should be zero since all parts are tool calls
expect(result.history).toBe(0);
});
it('should produce additive (non-overlapping) fields', () => {
const contents: Content[] = [
{ role: 'user', parts: [{ text: 'Hello' }] },
{
role: 'model',
parts: [
{
functionCall: {
name: 'read_file',
args: { path: '/tmp/test.txt' },
},
},
],
},
{
role: 'function',
parts: [
{
functionResponse: {
name: 'read_file',
response: { content: 'data' },
},
},
],
},
];
const config = {
systemInstruction: 'Be helpful.',
tools: [
{
functionDeclarations: [
{ name: 'read_file', description: 'Read', parameters: {} },
{
name: 'myserver__search',
description: 'MCP search',
parameters: {},
},
],
},
],
} as unknown as GenerateContentConfig;
const result = estimateContextBreakdown(contents, config);
// All fields should be non-overlapping
expect(result.system_instructions).toBeGreaterThan(0);
expect(result.tool_definitions).toBeGreaterThan(0);
expect(result.history).toBeGreaterThan(0);
// tool_calls should only contain non-MCP tools
expect(result.tool_calls['read_file']).toBeGreaterThan(0);
expect(result.tool_calls['myserver__search']).toBeUndefined();
// MCP tokens are only in mcp_servers
expect(result.mcp_servers).toBeGreaterThan(0);
});
it('should classify MCP tool calls into mcp_servers only, not tool_calls', () => {
const contents: Content[] = [
{
role: 'model',
parts: [
{
functionCall: {
name: 'myserver__search',
args: { query: 'test' },
},
},
],
},
{
role: 'function',
parts: [
{
functionResponse: {
name: 'myserver__search',
response: { results: [] },
},
},
],
},
];
const result = estimateContextBreakdown(contents);
// MCP tool calls should NOT appear in tool_calls
expect(result.tool_calls['myserver__search']).toBeUndefined();
// MCP call tokens should only be counted in mcp_servers
expect(result.mcp_servers).toBeGreaterThan(0);
});
it('should handle mixed MCP and non-MCP tool calls', () => {
const contents: Content[] = [
{
role: 'model',
parts: [
{
functionCall: {
name: 'read_file',
args: { path: '/test' },
},
},
{
functionCall: {
name: 'myserver__search',
args: { q: 'hello' },
},
},
],
},
];
const result = estimateContextBreakdown(contents);
// Non-MCP tools should be in tool_calls
expect(result.tool_calls['read_file']).toBeGreaterThan(0);
// MCP tools should NOT be in tool_calls
expect(result.tool_calls['myserver__search']).toBeUndefined();
// MCP tool calls should only be in mcp_servers
expect(result.mcp_servers).toBeGreaterThan(0);
});
it('should use "unknown" for tool calls without a name', () => {
const contents: Content[] = [
{
role: 'model',
parts: [
{
functionCall: {
name: undefined as unknown as string,
args: { x: 1 },
},
},
],
},
];
const result = estimateContextBreakdown(contents);
expect(result.tool_calls['unknown']).toBeGreaterThan(0);
});
});

View File

@@ -16,7 +16,7 @@ import type {
GenerateContentResponseUsageMetadata,
GenerateContentResponse,
} from '@google/genai';
import type { ServerDetails } from '../telemetry/types.js';
import type { ServerDetails, ContextBreakdown } from '../telemetry/types.js';
import {
ApiRequestEvent,
ApiResponseEvent,
@@ -37,14 +37,104 @@ 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 { isMcpToolName } from '../tools/mcp-tool.js';
import { estimateTokenCountSync } from '../utils/tokenCalculation.js';
interface StructuredError {
status: number;
}
/**
* A decorator that wraps a ContentGenerator to add logging to API calls.
* Rough token estimate for non-Part config objects (tool definitions, etc.)
* where estimateTokenCountSync cannot be used directly.
*/
function estimateConfigTokens(value: unknown): number {
return Math.floor(JSON.stringify(value).length / 4);
}
/**
* Estimates the context breakdown for telemetry. All returned fields are
* additive (non-overlapping), so their sum approximates the total context size.
*
* - system_instructions: tokens from system instruction config
* - tool_definitions: tokens from non-MCP tool definitions
* - history: tokens from conversation history, excluding tool call/response parts
* - tool_calls: per-tool token counts for non-MCP function call + response parts
* - mcp_servers: tokens from MCP tool definitions + MCP tool call/response parts
*
* MCP tool calls are excluded from tool_calls and counted only in mcp_servers
* to keep fields non-overlapping and avoid leaking MCP server names in telemetry.
*/
export function estimateContextBreakdown(
contents: Content[],
config?: GenerateContentConfig,
): ContextBreakdown {
let systemInstructions = 0;
let toolDefinitions = 0;
let history = 0;
let mcpServers = 0;
const toolCalls: Record<string, number> = {};
if (config?.systemInstruction) {
systemInstructions += estimateConfigTokens(config.systemInstruction);
}
if (config?.tools) {
for (const tool of config.tools) {
const toolTokens = estimateConfigTokens(tool);
if (
tool &&
typeof tool === 'object' &&
'functionDeclarations' in tool &&
tool.functionDeclarations
) {
let mcpTokensInTool = 0;
for (const func of tool.functionDeclarations) {
if (func.name && isMcpToolName(func.name)) {
mcpTokensInTool += estimateConfigTokens(func);
}
}
mcpServers += mcpTokensInTool;
toolDefinitions += toolTokens - mcpTokensInTool;
} else {
toolDefinitions += toolTokens;
}
}
}
for (const content of contents) {
for (const part of content.parts || []) {
if (part.functionCall) {
const name = part.functionCall.name || 'unknown';
const tokens = estimateTokenCountSync([part]);
if (isMcpToolName(name)) {
mcpServers += tokens;
} else {
toolCalls[name] = (toolCalls[name] || 0) + tokens;
}
} else if (part.functionResponse) {
const name = part.functionResponse.name || 'unknown';
const tokens = estimateTokenCountSync([part]);
if (isMcpToolName(name)) {
mcpServers += tokens;
} else {
toolCalls[name] = (toolCalls[name] || 0) + tokens;
}
} else {
history += estimateTokenCountSync([part]);
}
}
}
return {
system_instructions: systemInstructions,
tool_definitions: toolDefinitions,
history,
tool_calls: toolCalls,
mcp_servers: mcpServers,
};
}
export class LoggingContentGenerator implements ContentGenerator {
constructor(
private readonly wrapped: ContentGenerator,
@@ -134,27 +224,40 @@ export class LoggingContentGenerator implements ContentGenerator {
generationConfig?: GenerateContentConfig,
serverDetails?: ServerDetails,
): void {
logApiResponse(
this.config,
new ApiResponseEvent(
model,
durationMs,
{
prompt_id,
contents: requestContents,
generate_content_config: generationConfig,
server: serverDetails,
},
{
candidates: responseCandidates,
response_id: responseId,
},
this.config.getContentGeneratorConfig()?.authType,
usageMetadata,
responseText,
role,
),
const event = new ApiResponseEvent(
model,
durationMs,
{
prompt_id,
contents: requestContents,
generate_content_config: generationConfig,
server: serverDetails,
},
{
candidates: responseCandidates,
response_id: responseId,
},
this.config.getContentGeneratorConfig()?.authType,
usageMetadata,
responseText,
role,
);
// Only compute context breakdown for turn-ending responses (when the user
// gets back control to type). If the response contains function calls, the
// model is in a tool-use loop and will make more API calls — skip to avoid
// emitting redundant cumulative snapshots for every intermediate step.
const hasToolCalls = responseCandidates?.some((c) =>
c.content?.parts?.some((p) => p.functionCall),
);
if (!hasToolCalls) {
event.usage.context_breakdown = estimateContextBreakdown(
requestContents,
generationConfig,
);
}
logApiResponse(this.config, event);
}
private _logApiError(

View File

@@ -14,9 +14,9 @@ import {
HookType,
HOOKS_CONFIG_FIELDS,
type CommandHookConfig,
type HookDefinition,
} from './types.js';
import type { Config } from '../config/config.js';
import type { HookDefinition } from './types.js';
// Mock fs
vi.mock('fs', () => ({

View File

@@ -846,6 +846,40 @@ export class ClearcutLogger {
EventMetadataKey.GEMINI_CLI_API_RESPONSE_TOOL_TOKEN_COUNT,
value: JSON.stringify(event.usage.tool_token_count),
},
// Context breakdown fields are only populated on turn-ending responses
// (when the user gets back control), not during intermediate tool-use
// loops. Values still grow across turns as conversation history
// accumulates, so downstream consumers should use the last event per
// session (MAX) rather than summing across events.
{
gemini_cli_key:
EventMetadataKey.GEMINI_CLI_API_RESPONSE_CONTEXT_BREAKDOWN_SYSTEM_INSTRUCTIONS,
value: JSON.stringify(
event.usage.context_breakdown?.system_instructions ?? 0,
),
},
{
gemini_cli_key:
EventMetadataKey.GEMINI_CLI_API_RESPONSE_CONTEXT_BREAKDOWN_TOOL_DEFINITIONS,
value: JSON.stringify(
event.usage.context_breakdown?.tool_definitions ?? 0,
),
},
{
gemini_cli_key:
EventMetadataKey.GEMINI_CLI_API_RESPONSE_CONTEXT_BREAKDOWN_HISTORY,
value: JSON.stringify(event.usage.context_breakdown?.history ?? 0),
},
{
gemini_cli_key:
EventMetadataKey.GEMINI_CLI_API_RESPONSE_CONTEXT_BREAKDOWN_TOOL_CALLS,
value: JSON.stringify(event.usage.context_breakdown?.tool_calls ?? {}),
},
{
gemini_cli_key:
EventMetadataKey.GEMINI_CLI_API_RESPONSE_CONTEXT_BREAKDOWN_MCP_SERVERS,
value: JSON.stringify(event.usage.context_breakdown?.mcp_servers ?? 0),
},
];
this.enqueueLogEvent(this.createLogEvent(EventNames.API_RESPONSE, data));

View File

@@ -7,7 +7,7 @@
// Defines valid event metadata keys for Clearcut logging.
export enum EventMetadataKey {
// Deleted enums: 24
// Next ID: 167
// Next ID: 172
GEMINI_CLI_KEY_UNKNOWN = 0,
@@ -137,6 +137,21 @@ export enum EventMetadataKey {
// Logs the tool use token count of the API call.
GEMINI_CLI_API_RESPONSE_TOOL_TOKEN_COUNT = 29,
// Logs the token count for system instructions.
GEMINI_CLI_API_RESPONSE_CONTEXT_BREAKDOWN_SYSTEM_INSTRUCTIONS = 167,
// Logs the token count for tool definitions.
GEMINI_CLI_API_RESPONSE_CONTEXT_BREAKDOWN_TOOL_DEFINITIONS = 168,
// Logs the token count for conversation history.
GEMINI_CLI_API_RESPONSE_CONTEXT_BREAKDOWN_HISTORY = 169,
// Logs the token count for tool calls (JSON map of tool name to tokens).
GEMINI_CLI_API_RESPONSE_CONTEXT_BREAKDOWN_TOOL_CALLS = 170,
// Logs the token count from MCP servers (tool definitions + tool inputs/outputs).
GEMINI_CLI_API_RESPONSE_CONTEXT_BREAKDOWN_MCP_SERVERS = 171,
// ==========================================================================
// GenAI API Error Event Keys
// ===========================================================================

View File

@@ -569,6 +569,14 @@ export interface GenAIResponseDetails {
candidates?: Candidate[];
}
export interface ContextBreakdown {
system_instructions: number;
tool_definitions: number;
history: number;
tool_calls: Record<string, number>;
mcp_servers: number;
}
export interface GenAIUsageDetails {
input_token_count: number;
output_token_count: number;
@@ -576,6 +584,7 @@ export interface GenAIUsageDetails {
thoughts_token_count: number;
tool_token_count: number;
total_token_count: number;
context_breakdown?: ContextBreakdown;
}
export const EVENT_API_RESPONSE = 'gemini_cli.api_response';

View File

@@ -29,6 +29,16 @@ import type { MessageBus } from '../confirmation-bus/message-bus.js';
*/
export const MCP_QUALIFIED_NAME_SEPARATOR = '__';
/**
* Returns true if `name` matches the MCP qualified name format: "server__tool",
* i.e. exactly two non-empty parts separated by the MCP_QUALIFIED_NAME_SEPARATOR.
*/
export function isMcpToolName(name: string): boolean {
if (!name.includes(MCP_QUALIFIED_NAME_SEPARATOR)) return false;
const parts = name.split(MCP_QUALIFIED_NAME_SEPARATOR);
return parts.length === 2 && parts[0].length > 0 && parts[1].length > 0;
}
type ToolParams = Record<string, unknown>;
// Discriminated union for MCP Content Blocks to ensure type safety.