From 3ff5cfaaf66b89b1b5d14a743482f89c1c97348d Mon Sep 17 00:00:00 2001 From: Sandy Tao Date: Tue, 24 Feb 2026 15:26:28 -0800 Subject: [PATCH] feat(telemetry): Add context breakdown to API response event (#19699) --- .../scripts/fetch-pr-info.js | 3 +- .../src/core/loggingContentGenerator.test.ts | 284 +++++++++++++++++- .../core/src/core/loggingContentGenerator.ts | 147 +++++++-- packages/core/src/hooks/hookRegistry.test.ts | 2 +- .../clearcut-logger/clearcut-logger.ts | 34 +++ .../clearcut-logger/event-metadata-key.ts | 17 +- packages/core/src/telemetry/types.ts | 9 + packages/core/src/tools/mcp-tool.ts | 10 + 8 files changed, 479 insertions(+), 27 deletions(-) diff --git a/.gemini/skills/pr-address-comments/scripts/fetch-pr-info.js b/.gemini/skills/pr-address-comments/scripts/fetch-pr-info.js index de99def0ce..be465d17f2 100755 --- a/.gemini/skills/pr-address-comments/scripts/fetch-pr-info.js +++ b/.gemini/skills/pr-address-comments/scripts/fetch-pr-info.js @@ -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; } } diff --git a/packages/core/src/core/loggingContentGenerator.test.ts b/packages/core/src/core/loggingContentGenerator.test.ts index 45e5028553..01a7162f1d 100644 --- a/packages/core/src/core/loggingContentGenerator.test.ts +++ b/packages/core/src/core/loggingContentGenerator.test.ts @@ -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); + }); +}); diff --git a/packages/core/src/core/loggingContentGenerator.ts b/packages/core/src/core/loggingContentGenerator.ts index 1544087ae0..013600a0b9 100644 --- a/packages/core/src/core/loggingContentGenerator.ts +++ b/packages/core/src/core/loggingContentGenerator.ts @@ -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 = {}; + + 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( diff --git a/packages/core/src/hooks/hookRegistry.test.ts b/packages/core/src/hooks/hookRegistry.test.ts index 078990608c..d8157f4ef5 100644 --- a/packages/core/src/hooks/hookRegistry.test.ts +++ b/packages/core/src/hooks/hookRegistry.test.ts @@ -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', () => ({ diff --git a/packages/core/src/telemetry/clearcut-logger/clearcut-logger.ts b/packages/core/src/telemetry/clearcut-logger/clearcut-logger.ts index 7838d985f1..8646a3f6d4 100644 --- a/packages/core/src/telemetry/clearcut-logger/clearcut-logger.ts +++ b/packages/core/src/telemetry/clearcut-logger/clearcut-logger.ts @@ -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)); diff --git a/packages/core/src/telemetry/clearcut-logger/event-metadata-key.ts b/packages/core/src/telemetry/clearcut-logger/event-metadata-key.ts index 4d3bc27d27..d799ca1caf 100644 --- a/packages/core/src/telemetry/clearcut-logger/event-metadata-key.ts +++ b/packages/core/src/telemetry/clearcut-logger/event-metadata-key.ts @@ -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 // =========================================================================== diff --git a/packages/core/src/telemetry/types.ts b/packages/core/src/telemetry/types.ts index ccebef242b..867a47a9c7 100644 --- a/packages/core/src/telemetry/types.ts +++ b/packages/core/src/telemetry/types.ts @@ -569,6 +569,14 @@ export interface GenAIResponseDetails { candidates?: Candidate[]; } +export interface ContextBreakdown { + system_instructions: number; + tool_definitions: number; + history: number; + tool_calls: Record; + 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'; diff --git a/packages/core/src/tools/mcp-tool.ts b/packages/core/src/tools/mcp-tool.ts index e0bd04dd27..18c1638c95 100644 --- a/packages/core/src/tools/mcp-tool.ts +++ b/packages/core/src/tools/mcp-tool.ts @@ -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; // Discriminated union for MCP Content Blocks to ensure type safety.