diff --git a/packages/cli/src/ui/components/ModelStatsDisplay.test.tsx b/packages/cli/src/ui/components/ModelStatsDisplay.test.tsx index d233d3b3850..619bdd774c9 100644 --- a/packages/cli/src/ui/components/ModelStatsDisplay.test.tsx +++ b/packages/cli/src/ui/components/ModelStatsDisplay.test.tsx @@ -9,7 +9,7 @@ import { describe, it, expect, vi, beforeAll, afterAll } from 'vitest'; import { ModelStatsDisplay } from './ModelStatsDisplay.js'; import * as SessionContext from '../contexts/SessionContext.js'; import type { SessionMetrics } from '../contexts/SessionContext.js'; -import { ToolCallDecision } from '@google/gemini-cli-core'; +import { ToolCallDecision, LlmRole } from '@google/gemini-cli-core'; // Mock the context to provide controlled data for testing vi.mock('../contexts/SessionContext.js', async (importOriginal) => { @@ -95,6 +95,7 @@ describe('', () => { thoughts: 0, tool: 0, }, + roles: {}, }, }, tools: { @@ -137,6 +138,7 @@ describe('', () => { thoughts: 2, tool: 0, }, + roles: {}, }, 'gemini-2.5-flash': { api: { totalRequests: 1, totalErrors: 0, totalLatencyMs: 50 }, @@ -149,6 +151,7 @@ describe('', () => { thoughts: 0, tool: 3, }, + roles: {}, }, }, tools: { @@ -191,6 +194,7 @@ describe('', () => { thoughts: 10, tool: 5, }, + roles: {}, }, 'gemini-2.5-flash': { api: { totalRequests: 20, totalErrors: 2, totalLatencyMs: 500 }, @@ -203,6 +207,7 @@ describe('', () => { thoughts: 20, tool: 10, }, + roles: {}, }, }, tools: { @@ -248,6 +253,7 @@ describe('', () => { thoughts: 111111111, tool: 222222222, }, + roles: {}, }, }, tools: { @@ -286,6 +292,7 @@ describe('', () => { thoughts: 2, tool: 1, }, + roles: {}, }, }, tools: { @@ -328,6 +335,7 @@ describe('', () => { thoughts: 100, tool: 50, }, + roles: {}, }, 'gemini-3-flash-preview': { api: { totalRequests: 20, totalErrors: 0, totalLatencyMs: 1000 }, @@ -340,6 +348,7 @@ describe('', () => { thoughts: 200, tool: 100, }, + roles: {}, }, }, tools: { @@ -366,6 +375,120 @@ describe('', () => { const output = lastFrame(); expect(output).toContain('gemini-3-pro-'); expect(output).toContain('gemini-3-flash-'); + }); + + it('should display role breakdown correctly', () => { + const { lastFrame } = renderWithMockedStats({ + models: { + 'gemini-2.5-pro': { + api: { totalRequests: 2, totalErrors: 0, totalLatencyMs: 200 }, + tokens: { + input: 20, + prompt: 30, + candidates: 40, + total: 70, + cached: 10, + thoughts: 0, + tool: 0, + }, + roles: { + [LlmRole.MAIN]: { + totalRequests: 1, + totalErrors: 0, + totalLatencyMs: 100, + tokens: { + input: 10, + prompt: 15, + candidates: 20, + total: 35, + cached: 5, + thoughts: 0, + tool: 0, + }, + }, + }, + }, + }, + tools: { + totalCalls: 0, + totalSuccess: 0, + totalFail: 0, + totalDurationMs: 0, + totalDecisions: { + accept: 0, + reject: 0, + modify: 0, + [ToolCallDecision.AUTO_ACCEPT]: 0, + }, + byName: {}, + }, + files: { + totalLinesAdded: 0, + totalLinesRemoved: 0, + }, + }); + + const output = lastFrame(); + expect(output).toContain('main'); + expect(output).toContain('Input'); + expect(output).toContain('Output'); + expect(output).toContain('Cache Reads'); expect(output).toMatchSnapshot(); }); + + it('should handle long metric names with truncation', () => { + const longRoleName = + 'this_is_a_very_long_role_name_that_should_be_truncated' as LlmRole; + const { lastFrame } = renderWithMockedStats({ + models: { + 'gemini-2.5-pro': { + api: { totalRequests: 1, totalErrors: 0, totalLatencyMs: 100 }, + tokens: { + input: 10, + prompt: 10, + candidates: 20, + total: 30, + cached: 0, + thoughts: 0, + tool: 0, + }, + roles: { + [longRoleName]: { + totalRequests: 1, + totalErrors: 0, + totalLatencyMs: 100, + tokens: { + input: 10, + prompt: 10, + candidates: 20, + total: 30, + cached: 0, + thoughts: 0, + tool: 0, + }, + }, + }, + }, + }, + tools: { + totalCalls: 0, + totalSuccess: 0, + totalFail: 0, + totalDurationMs: 0, + totalDecisions: { + accept: 0, + reject: 0, + modify: 0, + [ToolCallDecision.AUTO_ACCEPT]: 0, + }, + byName: {}, + }, + files: { + totalLinesAdded: 0, + totalLinesRemoved: 0, + }, + }); + + expect(lastFrame()).toMatchSnapshot(); + }); }); diff --git a/packages/cli/src/ui/components/ModelStatsDisplay.tsx b/packages/cli/src/ui/components/ModelStatsDisplay.tsx index f765bcede3c..ae7e7e8f2f7 100644 --- a/packages/cli/src/ui/components/ModelStatsDisplay.tsx +++ b/packages/cli/src/ui/components/ModelStatsDisplay.tsx @@ -13,17 +13,24 @@ import { calculateCacheHitRate, calculateErrorRate, } from '../utils/computeStats.js'; -import { useSessionStats } from '../contexts/SessionContext.js'; +import { + useSessionStats, + type ModelMetrics, +} from '../contexts/SessionContext.js'; import { Table, type Column } from './Table.js'; +import { LlmRole } from '@google/gemini-cli-core'; interface StatRowData { metric: string; isSection?: boolean; - isSubtle?: boolean; + indentLevel?: number; // Dynamic keys for model values - [key: string]: string | React.ReactNode | boolean | undefined; + [key: string]: string | React.ReactNode | boolean | undefined | number; + color?: string; } +type RoleMetrics = NonNullable[LlmRole]>; + export const ModelStatsDisplay: React.FC = () => { const { stats } = useSessionStats(); const { models } = stats.metrics; @@ -56,132 +63,187 @@ export const ModelStatsDisplay: React.FC = () => { ([, metrics]) => metrics.tokens.cached > 0, ); - // Helper to create a row with values for each model - const createRow = ( + const allRoles = Array.from( + new Set( + activeModels.flatMap(([, metrics]) => Object.keys(metrics.roles || {})), + ), + ).sort((a, b) => { + if (a === LlmRole.MAIN) return -1; + if (b === LlmRole.MAIN) return 1; + return a.localeCompare(b); + }) as LlmRole[]; + + const rows: StatRowData[] = []; + + // Helper to add a row for global model metrics + const addRow = ( metric: string, - getValue: ( - metrics: (typeof activeModels)[0][1], - ) => string | React.ReactNode, - options: { isSection?: boolean; isSubtle?: boolean } = {}, - ): StatRowData => { + getValue: (metrics: ModelMetrics) => string | React.ReactNode, + options: { + isSection?: boolean; + indentLevel?: number; + } = {}, + ) => { + const { indentLevel = 0 } = options; const row: StatRowData = { metric, isSection: options.isSection, - isSubtle: options.isSubtle, + indentLevel, }; activeModels.forEach(([name, metrics]) => { row[name] = getValue(metrics); }); - return row; + rows.push(row); }; - const rows: StatRowData[] = [ - // API Section - { metric: 'API', isSection: true }, - createRow('Requests', (m) => m.api.totalRequests.toLocaleString()), - createRow('Errors', (m) => { - const errorRate = calculateErrorRate(m); - return ( - 0 ? theme.status.error : theme.text.primary - } - > - {m.api.totalErrors.toLocaleString()} ({errorRate.toFixed(1)}%) - - ); - }), - createRow('Avg Latency', (m) => formatDuration(calculateAverageLatency(m))), + // API Section + addRow('API', () => '', { isSection: true }); + addRow('Requests', (m) => m.api.totalRequests.toLocaleString()); + addRow('Errors', (m) => { + const errorRate = calculateErrorRate(m); + return ( + 0 ? theme.status.error : theme.text.primary} + > + {m.api.totalErrors.toLocaleString()} ({errorRate.toFixed(1)}%) + + ); + }); + addRow('Avg Latency', (m) => formatDuration(calculateAverageLatency(m))); - // Spacer - { metric: '' }, + // Spacer + rows.push({ metric: ' ' }); - // Tokens Section - { metric: 'Tokens', isSection: true }, - createRow('Total', (m) => ( - - {m.tokens.total.toLocaleString()} - - )), - createRow( - 'Input', - (m) => ( - - {m.tokens.input.toLocaleString()} - - ), - { isSubtle: true }, + // Tokens Section + addRow('Tokens', () => '', { isSection: true }); + addRow('Total', (m) => ( + {m.tokens.total.toLocaleString()} + )); + addRow( + 'Input', + (m) => ( + {m.tokens.input.toLocaleString()} ), - ]; + { indentLevel: 1 }, + ); if (hasCached) { - rows.push( - createRow( - 'Cache Reads', - (m) => { - const cacheHitRate = calculateCacheHitRate(m); - return ( - - {m.tokens.cached.toLocaleString()} ({cacheHitRate.toFixed(1)}%) - - ); - }, - { isSubtle: true }, - ), + addRow( + 'Cache Reads', + (m) => { + const cacheHitRate = calculateCacheHitRate(m); + return ( + + {m.tokens.cached.toLocaleString()} ({cacheHitRate.toFixed(1)}%) + + ); + }, + { indentLevel: 1 }, ); } if (hasThoughts) { - rows.push( - createRow( - 'Thoughts', - (m) => ( - - {m.tokens.thoughts.toLocaleString()} - - ), - { isSubtle: true }, + addRow( + 'Thoughts', + (m) => ( + + {m.tokens.thoughts.toLocaleString()} + ), + { indentLevel: 1 }, ); } if (hasTool) { - rows.push( - createRow( - 'Tool', - (m) => ( - - {m.tokens.tool.toLocaleString()} - - ), - { isSubtle: true }, + addRow( + 'Tool', + (m) => ( + {m.tokens.tool.toLocaleString()} ), + { indentLevel: 1 }, ); } - rows.push( - createRow( - 'Output', - (m) => ( - - {m.tokens.candidates.toLocaleString()} - - ), - { isSubtle: true }, + addRow( + 'Output', + (m) => ( + + {m.tokens.candidates.toLocaleString()} + ), + { indentLevel: 1 }, ); + // Roles Section + if (allRoles.length > 0) { + // Spacer + rows.push({ metric: ' ' }); + rows.push({ metric: 'Roles', isSection: true }); + + allRoles.forEach((role) => { + // Role Header Row + const roleHeaderRow: StatRowData = { + metric: role, + indentLevel: 1, + color: theme.text.accent, + }; + // We don't populate model values for the role header row + rows.push(roleHeaderRow); + + const addRoleMetric = ( + metric: string, + getValue: (r: RoleMetrics) => string | React.ReactNode, + ) => { + const row: StatRowData = { + metric, + indentLevel: 2, + }; + activeModels.forEach(([name, metrics]) => { + const roleMetrics = metrics.roles?.[role]; + if (roleMetrics) { + row[name] = getValue(roleMetrics); + } else { + row[name] = -; + } + }); + rows.push(row); + }; + + addRoleMetric('Requests', (r) => r.totalRequests.toLocaleString()); + addRoleMetric('Input', (r) => ( + + {r.tokens.input.toLocaleString()} + + )); + addRoleMetric('Output', (r) => ( + + {r.tokens.candidates.toLocaleString()} + + )); + addRoleMetric('Cache Reads', (r) => ( + + {r.tokens.cached.toLocaleString()} + + )); + }); + } + const columns: Array> = [ { key: 'metric', header: 'Metric', - width: 28, + width: 36, renderCell: (row) => ( - {row.isSubtle ? ` ↳ ${row.metric}` : row.metric} + {row.indentLevel + ? `${' '.repeat(row.indentLevel)}↳ ${row.metric}` + : row.metric} ), }, @@ -191,7 +253,7 @@ export const ModelStatsDisplay: React.FC = () => { flexGrow: 1, renderCell: (row: StatRowData) => { // Don't render anything for section headers in model columns - if (row.isSection) return null; + if (row.isSection && !row[name]) return null; const val = row[name]; if (val === undefined || val === null) return null; if (typeof val === 'string' || typeof val === 'number') { diff --git a/packages/cli/src/ui/components/SessionSummaryDisplay.test.tsx b/packages/cli/src/ui/components/SessionSummaryDisplay.test.tsx index 7476fa08d4e..c5ab844135a 100644 --- a/packages/cli/src/ui/components/SessionSummaryDisplay.test.tsx +++ b/packages/cli/src/ui/components/SessionSummaryDisplay.test.tsx @@ -53,6 +53,7 @@ describe('', () => { thoughts: 300, tool: 200, }, + roles: {}, }, }, tools: { diff --git a/packages/cli/src/ui/components/StatsDisplay.test.tsx b/packages/cli/src/ui/components/StatsDisplay.test.tsx index eb34fa6bd27..a002cd0b590 100644 --- a/packages/cli/src/ui/components/StatsDisplay.test.tsx +++ b/packages/cli/src/ui/components/StatsDisplay.test.tsx @@ -11,6 +11,7 @@ import * as SessionContext from '../contexts/SessionContext.js'; import type { SessionMetrics } from '../contexts/SessionContext.js'; import { ToolCallDecision, + LlmRole, type RetrieveUserQuotaResponse, } from '@google/gemini-cli-core'; @@ -93,6 +94,7 @@ describe('', () => { thoughts: 100, tool: 50, }, + roles: {}, }, 'gemini-2.5-flash': { api: { totalRequests: 5, totalErrors: 1, totalLatencyMs: 4500 }, @@ -105,6 +107,7 @@ describe('', () => { thoughts: 2000, tool: 1000, }, + roles: {}, }, }, }); @@ -133,6 +136,7 @@ describe('', () => { thoughts: 0, tool: 0, }, + roles: {}, }, }, tools: { @@ -227,6 +231,7 @@ describe('', () => { thoughts: 0, tool: 0, }, + roles: {}, }, }, }); @@ -410,6 +415,7 @@ describe('', () => { thoughts: 0, tool: 0, }, + roles: {}, }, }, }); @@ -498,4 +504,60 @@ describe('', () => { vi.useRealTimers(); }); }); + + it('renders roles correctly', () => { + const metrics = createTestMetrics({ + models: { + 'gemini-2.5-pro': { + api: { totalRequests: 3, totalErrors: 0, totalLatencyMs: 15000 }, + tokens: { + input: 500, + prompt: 1000, + candidates: 2000, + total: 43234, + cached: 500, + thoughts: 100, + tool: 50, + }, + roles: { + [LlmRole.MAIN]: { + totalRequests: 1, + totalErrors: 0, + totalLatencyMs: 1000, + tokens: { + input: 100, + prompt: 0, + candidates: 0, + total: 100, + cached: 0, + thoughts: 0, + tool: 0, + }, + }, + [LlmRole.UTILITY_TOOL]: { + totalRequests: 2, + totalErrors: 0, + totalLatencyMs: 14000, + tokens: { + input: 400, + prompt: 1000, + candidates: 2000, + total: 3400, + cached: 500, + thoughts: 100, + tool: 50, + }, + }, + }, + }, + }, + }); + + const { lastFrame } = renderWithMockedStats(metrics); + const output = lastFrame(); + + expect(output).toContain('↳ ' + LlmRole.MAIN); + expect(output).toContain('↳ ' + LlmRole.UTILITY_TOOL); + expect(output).toMatchSnapshot(); + }); }); diff --git a/packages/cli/src/ui/components/StatsDisplay.tsx b/packages/cli/src/ui/components/StatsDisplay.tsx index 8e89f54a6dd..175063347fd 100644 --- a/packages/cli/src/ui/components/StatsDisplay.tsx +++ b/packages/cli/src/ui/components/StatsDisplay.tsx @@ -95,6 +95,7 @@ const buildModelRows = ( outputTokens: metrics.tokens.candidates.toLocaleString(), bucket: quotas?.buckets?.find((b) => b.modelId === modelName), isActive: true, + roles: metrics.roles, }; }); @@ -116,6 +117,7 @@ const buildModelRows = ( outputTokens: '-', bucket, isActive: false, + roles: undefined, })) || []; return [...activeRows, ...quotaRows]; @@ -159,7 +161,7 @@ const ModelUsageTable: React.FC<{ const showQuotaColumn = !!quotas && rows.some((row) => !!row.bucket); - const nameWidth = 25; + const nameWidth = 38; const requestsWidth = 7; const uncachedWidth = 15; const cachedWidth = 14; @@ -182,7 +184,7 @@ const ModelUsageTable: React.FC<{ {/* Header */} - + Model Usage @@ -253,81 +255,135 @@ const ModelUsageTable: React.FC<{ borderRight={false} borderColor={theme.border.default} width={totalWidth} + flexShrink={1} > {rows.map((row) => ( - - - - {row.modelName} - - - - + + + + {row.modelName} + + + - {row.requests} - - - {!showQuotaColumn && ( - <> - - + + {!showQuotaColumn && ( + <> + - {row.inputTokens} - - - - {row.cachedTokens} - - - + {row.inputTokens} + + + - {row.outputTokens} - - - - )} - - {row.bucket && - row.bucket.remainingFraction != null && - row.bucket.resetTime && ( - - {(row.bucket.remainingFraction * 100).toFixed(1)}%{' '} - {formatResetTime(row.bucket.resetTime)} - - )} + {row.cachedTokens} + + + + {row.outputTokens} + + + + )} + + {row.bucket && + row.bucket.remainingFraction != null && + row.bucket.resetTime && ( + + {(row.bucket.remainingFraction * 100).toFixed(1)}%{' '} + {formatResetTime(row.bucket.resetTime)} + + )} + + {!showQuotaColumn && + row.roles && + Object.entries(row.roles).map(([role, metrics]) => ( + + + + {'↳ ' + role} + + + + + {metrics?.totalRequests} + + + + + {metrics?.tokens.input.toLocaleString()} + + + + + {metrics?.tokens.cached.toLocaleString()} + + + + + {metrics?.tokens.candidates.toLocaleString()} + + + + ))} ))} diff --git a/packages/cli/src/ui/components/__snapshots__/ModelStatsDisplay.test.tsx.snap b/packages/cli/src/ui/components/__snapshots__/ModelStatsDisplay.test.tsx.snap index f06cc3fe5af..841de3ae981 100644 --- a/packages/cli/src/ui/components/__snapshots__/ModelStatsDisplay.test.tsx.snap +++ b/packages/cli/src/ui/components/__snapshots__/ModelStatsDisplay.test.tsx.snap @@ -5,19 +5,20 @@ exports[` > should display a single model correctly 1`] = ` │ │ │ Model Stats For Nerds │ │ │ -│ Metric gemini-2.5-pro │ +│ Metric gemini-2.5-pro │ │ ────────────────────────────────────────────────────────────────────────────────────────────── │ │ API │ -│ Requests 1 │ -│ Errors 0 (0.0%) │ -│ Avg Latency 100ms │ +│ Requests 1 │ +│ Errors 0 (0.0%) │ +│ Avg Latency 100ms │ +│ │ │ Tokens │ -│ Total 30 │ -│ ↳ Input 5 │ -│ ↳ Cache Reads 5 (50.0%) │ -│ ↳ Thoughts 2 │ -│ ↳ Tool 1 │ -│ ↳ Output 20 │ +│ Total 30 │ +│ ↳ Input 5 │ +│ ↳ Cache Reads 5 (50.0%) │ +│ ↳ Thoughts 2 │ +│ ↳ Tool 1 │ +│ ↳ Output 20 │ │ │ ╰──────────────────────────────────────────────────────────────────────────────────────────────────╯" `; @@ -27,19 +28,48 @@ exports[` > should display conditional rows if at least one │ │ │ Model Stats For Nerds │ │ │ -│ Metric gemini-2.5-pro gemini-2.5-flash │ +│ Metric gemini-2.5-pro gemini-2.5-flash │ +│ ────────────────────────────────────────────────────────────────────────────────────────────── │ +│ API │ +│ Requests 1 1 │ +│ Errors 0 (0.0%) 0 (0.0%) │ +│ Avg Latency 100ms 50ms │ +│ │ +│ Tokens │ +│ Total 30 15 │ +│ ↳ Input 5 5 │ +│ ↳ Cache Reads 5 (50.0%) 0 (0.0%) │ +│ ↳ Thoughts 2 0 │ +│ ↳ Tool 0 3 │ +│ ↳ Output 20 10 │ +│ │ +╰──────────────────────────────────────────────────────────────────────────────────────────────────╯" +`; + +exports[` > should display role breakdown correctly 1`] = ` +"╭──────────────────────────────────────────────────────────────────────────────────────────────────╮ +│ │ +│ Model Stats For Nerds │ +│ │ +│ Metric gemini-2.5-pro │ │ ────────────────────────────────────────────────────────────────────────────────────────────── │ │ API │ -│ Requests 1 1 │ -│ Errors 0 (0.0%) 0 (0.0%) │ -│ Avg Latency 100ms 50ms │ +│ Requests 2 │ +│ Errors 0 (0.0%) │ +│ Avg Latency 100ms │ +│ │ │ Tokens │ -│ Total 30 15 │ -│ ↳ Input 5 5 │ -│ ↳ Cache Reads 5 (50.0%) 0 (0.0%) │ -│ ↳ Thoughts 2 0 │ -│ ↳ Tool 0 3 │ -│ ↳ Output 20 10 │ +│ Total 70 │ +│ ↳ Input 20 │ +│ ↳ Cache Reads 10 (33.3%) │ +│ ↳ Output 40 │ +│ │ +│ Roles │ +│ ↳ main │ +│ ↳ Requests 1 │ +│ ↳ Input 10 │ +│ ↳ Output 20 │ +│ ↳ Cache Reads 5 │ │ │ ╰──────────────────────────────────────────────────────────────────────────────────────────────────╯" `; @@ -49,19 +79,20 @@ exports[` > should display stats for multiple models correc │ │ │ Model Stats For Nerds │ │ │ -│ Metric gemini-2.5-pro gemini-2.5-flash │ +│ Metric gemini-2.5-pro gemini-2.5-flash │ │ ────────────────────────────────────────────────────────────────────────────────────────────── │ │ API │ -│ Requests 10 20 │ -│ Errors 1 (10.0%) 2 (10.0%) │ -│ Avg Latency 100ms 25ms │ +│ Requests 10 20 │ +│ Errors 1 (10.0%) 2 (10.0%) │ +│ Avg Latency 100ms 25ms │ +│ │ │ Tokens │ -│ Total 300 600 │ -│ ↳ Input 50 100 │ -│ ↳ Cache Reads 50 (50.0%) 100 (50.0%) │ -│ ↳ Thoughts 10 20 │ -│ ↳ Tool 5 10 │ -│ ↳ Output 200 400 │ +│ Total 300 600 │ +│ ↳ Input 50 100 │ +│ ↳ Cache Reads 50 (50.0%) 100 (50.0%) │ +│ ↳ Thoughts 10 20 │ +│ ↳ Tool 5 10 │ +│ ↳ Output 200 400 │ │ │ ╰──────────────────────────────────────────────────────────────────────────────────────────────────╯" `; @@ -71,43 +102,49 @@ exports[` > should handle large values without wrapping or │ │ │ Model Stats For Nerds │ │ │ -│ Metric gemini-2.5-pro │ +│ Metric gemini-2.5-pro │ │ ────────────────────────────────────────────────────────────────────────────────────────────── │ │ API │ -│ Requests 999,999,999 │ -│ Errors 123,456,789 (12.3%) │ -│ Avg Latency 0ms │ +│ Requests 999,999,999 │ +│ Errors 123,456,789 (12.3%) │ +│ Avg Latency 0ms │ +│ │ │ Tokens │ -│ Total 999,999,999 │ -│ ↳ Input 864,197,532 │ -│ ↳ Cache Reads 123,456,789 (12.5%) │ -│ ↳ Thoughts 111,111,111 │ -│ ↳ Tool 222,222,222 │ -│ ↳ Output 123,456,789 │ +│ Total 999,999,999 │ +│ ↳ Input 864,197,532 │ +│ ↳ Cache Reads 123,456,789 (12.5%) │ +│ ↳ Thoughts 111,111,111 │ +│ ↳ Tool 222,222,222 │ +│ ↳ Output 123,456,789 │ │ │ ╰──────────────────────────────────────────────────────────────────────────────────────────────────╯" `; -exports[` > should handle models with long names (gemini-3-*-preview) without layout breaking 1`] = ` -"╭──────────────────────────────────────────────────────────────────────────────╮ -│ │ -│ Model Stats For Nerds │ -│ │ -│ Metric gemini-3-pro-preview gemini-3-flash-preview │ -│ ────────────────────────────────────────────────────────────────────────── │ -│ API │ -│ Requests 10 20 │ -│ Errors 0 (0.0%) 0 (0.0%) │ -│ Avg Latency 200ms 50ms │ -│ Tokens │ -│ Total 6,000 12,000 │ -│ ↳ Input 1,000 2,000 │ -│ ↳ Cache Reads 500 (25.0%) 1,000 (25.0%) │ -│ ↳ Thoughts 100 200 │ -│ ↳ Tool 50 100 │ -│ ↳ Output 4,000 8,000 │ -│ │ -╰──────────────────────────────────────────────────────────────────────────────╯" +exports[` > should handle long metric names with truncation 1`] = ` +"╭──────────────────────────────────────────────────────────────────────────────────────────────────╮ +│ │ +│ Model Stats For Nerds │ +│ │ +│ Metric gemini-2.5-pro │ +│ ────────────────────────────────────────────────────────────────────────────────────────────── │ +│ API │ +│ Requests 1 │ +│ Errors 0 (0.0%) │ +│ Avg Latency 100ms │ +│ │ +│ Tokens │ +│ Total 30 │ +│ ↳ Input 10 │ +│ ↳ Output 20 │ +│ │ +│ Roles │ +│ ↳ this_is_a_very_long_role_name_… │ +│ ↳ Requests 1 │ +│ ↳ Input 10 │ +│ ↳ Output 20 │ +│ ↳ Cache Reads 0 │ +│ │ +╰──────────────────────────────────────────────────────────────────────────────────────────────────╯" `; exports[` > should not display conditional rows if no model has data for them 1`] = ` @@ -115,16 +152,17 @@ exports[` > should not display conditional rows if no model │ │ │ Model Stats For Nerds │ │ │ -│ Metric gemini-2.5-pro │ +│ Metric gemini-2.5-pro │ │ ────────────────────────────────────────────────────────────────────────────────────────────── │ │ API │ -│ Requests 1 │ -│ Errors 0 (0.0%) │ -│ Avg Latency 100ms │ +│ Requests 1 │ +│ Errors 0 (0.0%) │ +│ Avg Latency 100ms │ +│ │ │ Tokens │ -│ Total 30 │ -│ ↳ Input 10 │ -│ ↳ Output 20 │ +│ Total 30 │ +│ ↳ Input 10 │ +│ ↳ Output 20 │ │ │ ╰──────────────────────────────────────────────────────────────────────────────────────────────────╯" `; diff --git a/packages/cli/src/ui/components/__snapshots__/SessionSummaryDisplay.test.tsx.snap b/packages/cli/src/ui/components/__snapshots__/SessionSummaryDisplay.test.tsx.snap index 224758bca9d..6cb99a932d5 100644 --- a/packages/cli/src/ui/components/__snapshots__/SessionSummaryDisplay.test.tsx.snap +++ b/packages/cli/src/ui/components/__snapshots__/SessionSummaryDisplay.test.tsx.snap @@ -18,9 +18,9 @@ exports[` > renders the summary display with a title 1` │ » Tool Time: 0s (0.0%) │ │ │ │ │ -│ Model Usage Reqs Input Tokens Cache Reads Output Tokens │ -│ ──────────────────────────────────────────────────────────────────────────── │ -│ gemini-2.5-pro 10 500 500 2,000 │ +│ Model Usage Reqs Input Tokens Cache Reads Output Tokens │ +│ ───────────────────────────────────────────────────────────────────────────────────────── │ +│ gemini-2.5-pro 10 500 500 2,000 │ │ │ │ Savings Highlight: 500 (50.0%) of input tokens were served from the cache, reducing costs. │ │ │ diff --git a/packages/cli/src/ui/components/__snapshots__/StatsDisplay.test.tsx.snap b/packages/cli/src/ui/components/__snapshots__/StatsDisplay.test.tsx.snap index 7bac29f45da..a25c343ecfb 100644 --- a/packages/cli/src/ui/components/__snapshots__/StatsDisplay.test.tsx.snap +++ b/packages/cli/src/ui/components/__snapshots__/StatsDisplay.test.tsx.snap @@ -118,9 +118,9 @@ exports[` > Conditional Rendering Tests > hides Efficiency secti │ » Tool Time: 0s (0.0%) │ │ │ │ │ -│ Model Usage Reqs Input Tokens Cache Reads Output Tokens │ -│ ──────────────────────────────────────────────────────────────────────────── │ -│ gemini-2.5-pro 1 100 0 100 │ +│ Model Usage Reqs Input Tokens Cache Reads Output Tokens │ +│ ───────────────────────────────────────────────────────────────────────────────────────── │ +│ gemini-2.5-pro 1 100 0 100 │ │ │ ╰──────────────────────────────────────────────────────────────────────────────────────────────────╯" `; @@ -162,9 +162,9 @@ exports[` > Quota Display > renders quota information for unused │ » Tool Time: 0s (0.0%) │ │ │ │ │ -│ Model Usage Reqs Usage left │ -│ ──────────────────────────────────────────────────────────── │ -│ gemini-2.5-flash - 50.0% (Resets in 2h) │ +│ Model Usage Reqs Usage left │ +│ ───────────────────────────────────────────────────────────────────────── │ +│ gemini-2.5-flash - 50.0% (Resets in 2h) │ │ │ │ Usage limits span all sessions and reset daily. │ │ /auth to upgrade or switch to API key. │ @@ -192,9 +192,9 @@ exports[` > Quota Display > renders quota information when quota │ » Tool Time: 0s (0.0%) │ │ │ │ │ -│ Model Usage Reqs Usage left │ -│ ──────────────────────────────────────────────────────────── │ -│ gemini-2.5-pro 1 75.0% (Resets in 1h 30m) │ +│ Model Usage Reqs Usage left │ +│ ───────────────────────────────────────────────────────────────────────── │ +│ gemini-2.5-pro 1 75.0% (Resets in 1h 30m) │ │ │ │ Usage limits span all sessions and reset daily. │ │ /auth to upgrade or switch to API key. │ @@ -262,10 +262,10 @@ exports[` > renders a table with two models correctly 1`] = ` │ » Tool Time: 0s (0.0%) │ │ │ │ │ -│ Model Usage Reqs Input Tokens Cache Reads Output Tokens │ -│ ──────────────────────────────────────────────────────────────────────────── │ -│ gemini-2.5-pro 3 500 500 2,000 │ -│ gemini-2.5-flash 5 15,000 10,000 15,000 │ +│ Model Usage Reqs Input Tokens Cache Reads Output Tokens │ +│ ───────────────────────────────────────────────────────────────────────────────────────── │ +│ gemini-2.5-pro 3 500 500 2,000 │ +│ gemini-2.5-flash 5 15,000 10,000 15,000 │ │ │ │ Savings Highlight: 10,500 (40.4%) of input tokens were served from the cache, reducing costs. │ │ │ @@ -290,9 +290,9 @@ exports[` > renders all sections when all data is present 1`] = │ » Tool Time: 123ms (55.2%) │ │ │ │ │ -│ Model Usage Reqs Input Tokens Cache Reads Output Tokens │ -│ ──────────────────────────────────────────────────────────────────────────── │ -│ gemini-2.5-pro 1 50 50 100 │ +│ Model Usage Reqs Input Tokens Cache Reads Output Tokens │ +│ ───────────────────────────────────────────────────────────────────────────────────────── │ +│ gemini-2.5-pro 1 50 50 100 │ │ │ │ Savings Highlight: 50 (50.0%) of input tokens were served from the cache, reducing costs. │ │ │ @@ -318,3 +318,31 @@ exports[` > renders only the Performance section in its zero sta │ │ ╰──────────────────────────────────────────────────────────────────────────────────────────────────╯" `; + +exports[` > renders roles correctly 1`] = ` +"╭──────────────────────────────────────────────────────────────────────────────────────────────────╮ +│ │ +│ Session Stats │ +│ │ +│ Interaction Summary │ +│ Session ID: test-session-id │ +│ Tool Calls: 0 ( ✓ 0 x 0 ) │ +│ Success Rate: 0.0% │ +│ │ +│ Performance │ +│ Wall Time: 1s │ +│ Agent Active: 15.0s │ +│ » API Time: 15.0s (100.0%) │ +│ » Tool Time: 0s (0.0%) │ +│ │ +│ │ +│ Model Usage Reqs Input Tokens Cache Reads Output Tokens │ +│ ───────────────────────────────────────────────────────────────────────────────────────── │ +│ gemini-2.5-pro 3 500 500 2,000 │ +│ ↳ main 1 100 0 0 │ +│ ↳ utility_tool 2 400 500 2,000 │ +│ │ +│ Savings Highlight: 500 (50.0%) of input tokens were served from the cache, reducing costs. │ +│ │ +╰──────────────────────────────────────────────────────────────────────────────────────────────────╯" +`; diff --git a/packages/cli/src/ui/contexts/SessionContext.test.tsx b/packages/cli/src/ui/contexts/SessionContext.test.tsx index 5ab02042557..5ab76e45191 100644 --- a/packages/cli/src/ui/contexts/SessionContext.test.tsx +++ b/packages/cli/src/ui/contexts/SessionContext.test.tsx @@ -100,6 +100,7 @@ describe('SessionStatsContext', () => { thoughts: 20, tool: 10, }, + roles: {}, }, }, tools: { @@ -180,6 +181,7 @@ describe('SessionStatsContext', () => { thoughts: 0, tool: 0, }, + roles: {}, }, }, tools: { diff --git a/packages/cli/src/ui/hooks/usePromptCompletion.ts b/packages/cli/src/ui/hooks/usePromptCompletion.ts index c666b82900f..b17c802a553 100644 --- a/packages/cli/src/ui/hooks/usePromptCompletion.ts +++ b/packages/cli/src/ui/hooks/usePromptCompletion.ts @@ -6,7 +6,7 @@ import { useState, useCallback, useRef, useEffect, useMemo } from 'react'; import type { Config } from '@google/gemini-cli-core'; -import { debugLogger, getResponseText } from '@google/gemini-cli-core'; +import { debugLogger, getResponseText, LlmRole } from '@google/gemini-cli-core'; import type { Content } from '@google/genai'; import type { TextBuffer } from '../components/shared/text-buffer.js'; import { isSlashCommand } from '../utils/commandUtils.js'; @@ -110,6 +110,7 @@ export function usePromptCompletion({ { model: 'prompt-completion' }, contents, signal, + LlmRole.UTILITY_AUTOCOMPLETE, ); if (signal.aborted) { diff --git a/packages/cli/src/ui/utils/computeStats.test.ts b/packages/cli/src/ui/utils/computeStats.test.ts index b3677164a78..09baec304f0 100644 --- a/packages/cli/src/ui/utils/computeStats.test.ts +++ b/packages/cli/src/ui/utils/computeStats.test.ts @@ -29,6 +29,7 @@ describe('calculateErrorRate', () => { thoughts: 0, tool: 0, }, + roles: {}, }; expect(calculateErrorRate(metrics)).toBe(0); }); @@ -45,6 +46,7 @@ describe('calculateErrorRate', () => { thoughts: 0, tool: 0, }, + roles: {}, }; expect(calculateErrorRate(metrics)).toBe(20); }); @@ -63,6 +65,7 @@ describe('calculateAverageLatency', () => { thoughts: 0, tool: 0, }, + roles: {}, }; expect(calculateAverageLatency(metrics)).toBe(0); }); @@ -79,6 +82,7 @@ describe('calculateAverageLatency', () => { thoughts: 0, tool: 0, }, + roles: {}, }; expect(calculateAverageLatency(metrics)).toBe(150); }); @@ -97,6 +101,7 @@ describe('calculateCacheHitRate', () => { thoughts: 0, tool: 0, }, + roles: {}, }; expect(calculateCacheHitRate(metrics)).toBe(0); }); @@ -113,6 +118,7 @@ describe('calculateCacheHitRate', () => { thoughts: 0, tool: 0, }, + roles: {}, }; expect(calculateCacheHitRate(metrics)).toBe(25); }); @@ -170,6 +176,7 @@ describe('computeSessionStats', () => { thoughts: 0, tool: 0, }, + roles: {}, }, }, tools: { @@ -209,6 +216,7 @@ describe('computeSessionStats', () => { thoughts: 0, tool: 0, }, + roles: {}, }, }, tools: { diff --git a/packages/cli/src/zed-integration/zedIntegration.test.ts b/packages/cli/src/zed-integration/zedIntegration.test.ts index 0a6341a158b..187a5bae943 100644 --- a/packages/cli/src/zed-integration/zedIntegration.test.ts +++ b/packages/cli/src/zed-integration/zedIntegration.test.ts @@ -24,6 +24,7 @@ import { ReadManyFilesTool, type GeminiChat, type Config, + LlmRole, } from '@google/gemini-cli-core'; import { SettingScope, type LoadedSettings } from '../config/settings.js'; import { loadCliConfig, type CliArgs } from '../config/config.js'; @@ -527,7 +528,8 @@ describe('Session', () => { }), ]), expect.anything(), - expect.anything(), + expect.any(AbortSignal), + LlmRole.MAIN, ); }); diff --git a/packages/cli/src/zed-integration/zedIntegration.ts b/packages/cli/src/zed-integration/zedIntegration.ts index b87bb7e77c1..d1fd3863530 100644 --- a/packages/cli/src/zed-integration/zedIntegration.ts +++ b/packages/cli/src/zed-integration/zedIntegration.ts @@ -30,6 +30,7 @@ import { getEffectiveModel, createWorkingStdio, startupProfiler, + LlmRole, } from '@google/gemini-cli-core'; import * as acp from '@agentclientprotocol/sdk'; import { AcpFileSystemService } from './fileSystemService.js'; @@ -291,6 +292,7 @@ export class Session { nextMessage?.parts ?? [], promptId, pendingSend.signal, + LlmRole.MAIN, ); nextMessage = null; diff --git a/packages/core/src/agents/local-executor.ts b/packages/core/src/agents/local-executor.ts index 6cb78c5f113..fc3f2215de7 100644 --- a/packages/core/src/agents/local-executor.ts +++ b/packages/core/src/agents/local-executor.ts @@ -45,6 +45,7 @@ import { zodToJsonSchema } from 'zod-to-json-schema'; import { debugLogger } from '../utils/debugLogger.js'; import { getModelConfigAlias } from './registry.js'; import { ApprovalMode } from '../policy/types.js'; +import { LlmRole } from '../telemetry/types.js'; /** A callback function to report on agent activity. */ export type ActivityCallback = (activity: SubagentActivityEvent) => void; @@ -577,6 +578,11 @@ export class LocalAgentExecutor { signal: AbortSignal, promptId: string, ): Promise<{ functionCalls: FunctionCall[]; textResponse: string }> { + const role = + this.definition.name === 'codebase_investigator' + ? LlmRole.SUBAGENT_CODEBASE_INVESTIGATOR + : LlmRole.SUBAGENT; + const responseStream = await chat.sendMessageStream( { model: getModelConfigAlias(this.definition), @@ -585,6 +591,7 @@ export class LocalAgentExecutor { message.parts || [], promptId, signal, + role, ); const functionCalls: FunctionCall[] = []; diff --git a/packages/core/src/code_assist/server.test.ts b/packages/core/src/code_assist/server.test.ts index 91873a796eb..5080457efb3 100644 --- a/packages/core/src/code_assist/server.test.ts +++ b/packages/core/src/code_assist/server.test.ts @@ -9,6 +9,7 @@ import { CodeAssistServer } from './server.js'; import { OAuth2Client } from 'google-auth-library'; import { UserTierId, ActionStatus } from './types.js'; import { FinishReason } from '@google/genai'; +import { LlmRole } from '../telemetry/types.js'; vi.mock('google-auth-library'); @@ -69,6 +70,7 @@ describe('CodeAssistServer', () => { contents: [{ role: 'user', parts: [{ text: 'request' }] }], }, 'user-prompt-id', + LlmRole.MAIN, ); expect(mockRequest).toHaveBeenCalledWith({ @@ -126,6 +128,7 @@ describe('CodeAssistServer', () => { contents: [{ role: 'user', parts: [{ text: 'request' }] }], }, 'user-prompt-id', + LlmRole.MAIN, ); expect(recordConversationOfferedSpy).toHaveBeenCalledWith( @@ -170,6 +173,7 @@ describe('CodeAssistServer', () => { contents: [{ role: 'user', parts: [{ text: 'request' }] }], }, 'user-prompt-id', + LlmRole.MAIN, ); expect(server.recordCodeAssistMetrics).toHaveBeenCalledWith( @@ -205,6 +209,7 @@ describe('CodeAssistServer', () => { contents: [{ role: 'user', parts: [{ text: 'request' }] }], }, 'user-prompt-id', + LlmRole.MAIN, ); const mockResponseData = { @@ -344,6 +349,7 @@ describe('CodeAssistServer', () => { contents: [{ role: 'user', parts: [{ text: 'request' }] }], }, 'user-prompt-id', + LlmRole.MAIN, ); // Push SSE data to the stream diff --git a/packages/core/src/code_assist/server.ts b/packages/core/src/code_assist/server.ts index ccebc80977f..a3ff14663a9 100644 --- a/packages/core/src/code_assist/server.ts +++ b/packages/core/src/code_assist/server.ts @@ -51,6 +51,7 @@ import { recordConversationOffered, } from './telemetry.js'; import { getClientMetadata } from './experiments/client_metadata.js'; +import type { LlmRole } from '../telemetry/types.js'; /** HTTP options to be used in each of the requests. */ export interface HttpOptions { @@ -73,6 +74,8 @@ export class CodeAssistServer implements ContentGenerator { async generateContentStream( req: GenerateContentParameters, userPromptId: string, + // eslint-disable-next-line @typescript-eslint/no-unused-vars + role: LlmRole, ): Promise> { const responses = await this.requestStreamingPost( @@ -123,6 +126,8 @@ export class CodeAssistServer implements ContentGenerator { async generateContent( req: GenerateContentParameters, userPromptId: string, + // eslint-disable-next-line @typescript-eslint/no-unused-vars + role: LlmRole, ): Promise { const start = Date.now(); const response = await this.requestPost( diff --git a/packages/core/src/core/baseLlmClient.test.ts b/packages/core/src/core/baseLlmClient.test.ts index bcb701e739d..8e9dafa21b7 100644 --- a/packages/core/src/core/baseLlmClient.test.ts +++ b/packages/core/src/core/baseLlmClient.test.ts @@ -30,6 +30,7 @@ import { MalformedJsonResponseEvent } from '../telemetry/types.js'; import { getErrorMessage } from '../utils/errors.js'; import type { ModelConfigService } from '../services/modelConfigService.js'; import { makeResolvedModelConfig } from '../services/modelConfigServiceTestUtils.js'; +import { LlmRole } from '../telemetry/types.js'; vi.mock('../utils/errorReporting.js'); vi.mock('../telemetry/loggers.js'); @@ -129,6 +130,7 @@ describe('BaseLlmClient', () => { schema: { type: 'object', properties: { color: { type: 'string' } } }, abortSignal: abortController.signal, promptId: 'test-prompt-id', + role: LlmRole.UTILITY_TOOL, }; }); @@ -170,6 +172,7 @@ describe('BaseLlmClient', () => { }, }, 'test-prompt-id', + LlmRole.UTILITY_TOOL, ); }); @@ -192,6 +195,7 @@ describe('BaseLlmClient', () => { }), }), expect.any(String), + LlmRole.UTILITY_TOOL, ); }); @@ -210,6 +214,7 @@ describe('BaseLlmClient', () => { expect(mockGenerateContent).toHaveBeenCalledWith( expect.any(Object), customPromptId, + LlmRole.UTILITY_TOOL, ); }); @@ -529,6 +534,7 @@ describe('BaseLlmClient', () => { contents: [{ role: 'user', parts: [{ text: 'Give me content.' }] }], abortSignal: abortController.signal, promptId: 'content-prompt-id', + role: LlmRole.UTILITY_TOOL, }; const result = await client.generateContent(options); @@ -557,6 +563,7 @@ describe('BaseLlmClient', () => { }, }, 'content-prompt-id', + LlmRole.UTILITY_TOOL, ); }); @@ -569,6 +576,7 @@ describe('BaseLlmClient', () => { contents: [{ role: 'user', parts: [{ text: 'Give me content.' }] }], abortSignal: abortController.signal, promptId: 'content-prompt-id', + role: LlmRole.UTILITY_TOOL, }; await client.generateContent(options); @@ -591,6 +599,7 @@ describe('BaseLlmClient', () => { contents: [{ role: 'user', parts: [{ text: 'Give me content.' }] }], abortSignal: abortController.signal, promptId: 'content-prompt-id', + role: LlmRole.UTILITY_TOOL, }; await expect(client.generateContent(options)).rejects.toThrow( @@ -635,6 +644,7 @@ describe('BaseLlmClient', () => { contents: [{ role: 'user', parts: [{ text: 'Give me a color.' }] }], abortSignal: abortController.signal, promptId: 'content-prompt-id', + role: LlmRole.UTILITY_TOOL, }; jsonOptions = { @@ -656,6 +666,7 @@ describe('BaseLlmClient', () => { await client.generateContent({ ...contentOptions, modelConfigKey: { model: successfulModel }, + role: LlmRole.UTILITY_TOOL, }); expect(mockAvailabilityService.markHealthy).toHaveBeenCalledWith( @@ -681,6 +692,7 @@ describe('BaseLlmClient', () => { ...contentOptions, modelConfigKey: { model: firstModel }, maxAttempts: 2, + role: LlmRole.UTILITY_TOOL, }); await vi.runAllTimersAsync(); @@ -690,6 +702,7 @@ describe('BaseLlmClient', () => { ...contentOptions, modelConfigKey: { model: firstModel }, maxAttempts: 2, + role: LlmRole.UTILITY_TOOL, }); expect(mockConfig.setActiveModel).toHaveBeenCalledWith(firstModel); @@ -700,6 +713,7 @@ describe('BaseLlmClient', () => { expect(mockGenerateContent).toHaveBeenLastCalledWith( expect.objectContaining({ model: fallbackModel }), expect.any(String), + LlmRole.UTILITY_TOOL, ); }); @@ -725,6 +739,7 @@ describe('BaseLlmClient', () => { await client.generateContent({ ...contentOptions, modelConfigKey: { model: stickyModel }, + role: LlmRole.UTILITY_TOOL, }); expect(mockAvailabilityService.consumeStickyAttempt).toHaveBeenCalledWith( @@ -764,6 +779,7 @@ describe('BaseLlmClient', () => { expect(mockGenerateContent).toHaveBeenLastCalledWith( expect.objectContaining({ model: availableModel }), jsonOptions.promptId, + LlmRole.UTILITY_TOOL, ); }); @@ -815,6 +831,7 @@ describe('BaseLlmClient', () => { ...contentOptions, modelConfigKey: { model: firstModel }, maxAttempts: 2, + role: LlmRole.UTILITY_TOOL, }); expect(mockGenerateContent).toHaveBeenCalledTimes(2); diff --git a/packages/core/src/core/baseLlmClient.ts b/packages/core/src/core/baseLlmClient.ts index a508cdd038f..453d95a7c4d 100644 --- a/packages/core/src/core/baseLlmClient.ts +++ b/packages/core/src/core/baseLlmClient.ts @@ -27,6 +27,7 @@ import { applyModelSelection, createAvailabilityContextProvider, } from '../availability/policyHelpers.js'; +import type { LlmRole } from '../telemetry/types.js'; const DEFAULT_MAX_ATTEMPTS = 5; @@ -55,6 +56,10 @@ export interface GenerateJsonOptions { * The maximum number of attempts for the request. */ maxAttempts?: number; + /** + * The role of the LLM call. + */ + role: LlmRole; } /** @@ -80,6 +85,10 @@ export interface GenerateContentOptions { * The maximum number of attempts for the request. */ maxAttempts?: number; + /** + * The role of the LLM call. + */ + role: LlmRole; } interface _CommonGenerateOptions { @@ -116,6 +125,7 @@ export class BaseLlmClient { abortSignal, promptId, maxAttempts, + role, } = options; const { model } = @@ -150,6 +160,7 @@ export class BaseLlmClient { }, shouldRetryOnContent, 'generateJson', + role, ); // If we are here, the content is valid (not empty and parsable). @@ -216,6 +227,7 @@ export class BaseLlmClient { abortSignal, promptId, maxAttempts, + role, } = options; const shouldRetryOnContent = (response: GenerateContentResponse) => { @@ -234,6 +246,7 @@ export class BaseLlmClient { }, shouldRetryOnContent, 'generateContent', + role, ); } @@ -241,6 +254,7 @@ export class BaseLlmClient { options: _CommonGenerateOptions, shouldRetryOnContent: (response: GenerateContentResponse) => boolean, errorContext: 'generateJson' | 'generateContent', + role: LlmRole = 'utility_tool' as LlmRole, ): Promise { const { modelConfigKey, @@ -293,7 +307,11 @@ export class BaseLlmClient { config: finalConfig, contents, }; - return this.contentGenerator.generateContent(requestParams, promptId); + return this.contentGenerator.generateContent( + requestParams, + promptId, + role, + ); }; return await retryWithBackoff(apiCall, { diff --git a/packages/core/src/core/client.test.ts b/packages/core/src/core/client.test.ts index 3f496564ab7..063f4e95db0 100644 --- a/packages/core/src/core/client.test.ts +++ b/packages/core/src/core/client.test.ts @@ -47,6 +47,7 @@ import type { import { ClearcutLogger } from '../telemetry/clearcut-logger/clearcut-logger.js'; import { HookSystem } from '../hooks/hookSystem.js'; import * as policyCatalog from '../availability/policyCatalog.js'; +import { LlmRole } from '../telemetry/types.js'; vi.mock('../services/chatCompressionService.js'); @@ -2545,6 +2546,7 @@ ${JSON.stringify( { model: 'test-model' }, contents, abortSignal, + LlmRole.MAIN, ); expect(mockContentGenerator.generateContent).toHaveBeenCalledWith( @@ -2559,6 +2561,7 @@ ${JSON.stringify( contents, }, 'test-session-id', + LlmRole.MAIN, ); }); @@ -2570,6 +2573,7 @@ ${JSON.stringify( { model: initialModel }, contents, new AbortController().signal, + LlmRole.MAIN, ); expect(mockContentGenerator.generateContent).toHaveBeenCalledWith( @@ -2577,6 +2581,7 @@ ${JSON.stringify( model: initialModel, }), 'test-session-id', + LlmRole.MAIN, ); }); }); diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts index fba1a5bd1a0..16d3094cf22 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -58,6 +58,7 @@ import { createAvailabilityContextProvider, } from '../availability/policyHelpers.js'; import type { RetryAvailabilityContext } from '../utils/retry.js'; +import type { LlmRole } from '../telemetry/types.js'; const MAX_TURNS = 100; @@ -660,6 +661,7 @@ export class GeminiClient { modelConfigKey: ModelConfigKey, contents: Content[], abortSignal: AbortSignal, + role: LlmRole, ): Promise { const desiredModelConfig = this.config.modelConfigService.getResolvedConfig(modelConfigKey); @@ -714,6 +716,7 @@ export class GeminiClient { contents, }, this.lastPromptId, + role, ); }; const onPersistent429Callback = async ( diff --git a/packages/core/src/core/contentGenerator.ts b/packages/core/src/core/contentGenerator.ts index 7c733bce0f9..a654aedce1d 100644 --- a/packages/core/src/core/contentGenerator.ts +++ b/packages/core/src/core/contentGenerator.ts @@ -24,6 +24,7 @@ import { FakeContentGenerator } from './fakeContentGenerator.js'; import { parseCustomHeaders } from '../utils/customHeaderUtils.js'; import { RecordingContentGenerator } from './recordingContentGenerator.js'; import { getVersion, getEffectiveModel } from '../../index.js'; +import type { LlmRole } from '../telemetry/llmRole.js'; /** * Interface abstracting the core functionalities for generating content and counting tokens. @@ -32,11 +33,13 @@ export interface ContentGenerator { generateContent( request: GenerateContentParameters, userPromptId: string, + role: LlmRole, ): Promise; generateContentStream( request: GenerateContentParameters, userPromptId: string, + role: LlmRole, ): Promise>; countTokens(request: CountTokensParameters): Promise; diff --git a/packages/core/src/core/fakeContentGenerator.test.ts b/packages/core/src/core/fakeContentGenerator.test.ts index de8306e5167..673fa6b2e72 100644 --- a/packages/core/src/core/fakeContentGenerator.test.ts +++ b/packages/core/src/core/fakeContentGenerator.test.ts @@ -18,6 +18,7 @@ import { type CountTokensParameters, type EmbedContentParameters, } from '@google/genai'; +import { LlmRole } from '../telemetry/types.js'; vi.mock('node:fs', async (importOriginal) => { const actual = await importOriginal(); @@ -79,6 +80,7 @@ describe('FakeContentGenerator', () => { const response = await generator.generateContent( {} as GenerateContentParameters, 'id', + LlmRole.MAIN, ); expect(response).instanceOf(GenerateContentResponse); expect(response).toEqual(fakeGenerateContentResponse.response); @@ -91,6 +93,7 @@ describe('FakeContentGenerator', () => { const stream = await generator.generateContentStream( {} as GenerateContentParameters, 'id', + LlmRole.MAIN, ); const responses = []; for await (const response of stream) { @@ -121,7 +124,11 @@ describe('FakeContentGenerator', () => { ]; const generator = new FakeContentGenerator(fakeResponses); for (const fakeResponse of fakeResponses) { - const response = await generator[fakeResponse.method]({} as never, ''); + const response = await generator[fakeResponse.method]( + {} as never, + '', + LlmRole.MAIN, + ); if (fakeResponse.method === 'generateContentStream') { const responses = []; for await (const item of response as AsyncGenerator) { @@ -137,7 +144,11 @@ describe('FakeContentGenerator', () => { it('should throw error when no more responses', async () => { const generator = new FakeContentGenerator([fakeGenerateContentResponse]); - await generator.generateContent({} as GenerateContentParameters, 'id'); + await generator.generateContent( + {} as GenerateContentParameters, + 'id', + LlmRole.MAIN, + ); await expect( generator.embedContent({} as EmbedContentParameters), ).rejects.toThrowError('No more mock responses for embedContent'); @@ -145,10 +156,18 @@ describe('FakeContentGenerator', () => { generator.countTokens({} as CountTokensParameters), ).rejects.toThrowError('No more mock responses for countTokens'); await expect( - generator.generateContentStream({} as GenerateContentParameters, 'id'), + generator.generateContentStream( + {} as GenerateContentParameters, + 'id', + LlmRole.MAIN, + ), ).rejects.toThrow('No more mock responses for generateContentStream'); await expect( - generator.generateContent({} as GenerateContentParameters, 'id'), + generator.generateContent( + {} as GenerateContentParameters, + 'id', + LlmRole.MAIN, + ), ).rejects.toThrowError('No more mock responses for generateContent'); }); @@ -161,6 +180,7 @@ describe('FakeContentGenerator', () => { const response = await generator.generateContent( {} as GenerateContentParameters, 'id', + LlmRole.MAIN, ); expect(response).toEqual(fakeGenerateContentResponse.response); }); diff --git a/packages/core/src/core/fakeContentGenerator.ts b/packages/core/src/core/fakeContentGenerator.ts index a464c4f8faa..3ead5dcc022 100644 --- a/packages/core/src/core/fakeContentGenerator.ts +++ b/packages/core/src/core/fakeContentGenerator.ts @@ -16,6 +16,7 @@ import { promises } from 'node:fs'; import type { ContentGenerator } from './contentGenerator.js'; import type { UserTierId } from '../code_assist/types.js'; import { safeJsonStringify } from '../utils/safeJsonStringify.js'; +import type { LlmRole } from '../telemetry/types.js'; export type FakeResponse = | { @@ -76,6 +77,8 @@ export class FakeContentGenerator implements ContentGenerator { async generateContent( request: GenerateContentParameters, _userPromptId: string, + // eslint-disable-next-line @typescript-eslint/no-unused-vars + role: LlmRole, ): Promise { return Object.setPrototypeOf( this.getNextResponse('generateContent', request), @@ -86,6 +89,8 @@ export class FakeContentGenerator implements ContentGenerator { async generateContentStream( request: GenerateContentParameters, _userPromptId: string, + // eslint-disable-next-line @typescript-eslint/no-unused-vars + role: LlmRole, ): Promise> { const responses = this.getNextResponse('generateContentStream', request); async function* stream() { diff --git a/packages/core/src/core/geminiChat.test.ts b/packages/core/src/core/geminiChat.test.ts index baf8973904d..5ee20e6019f 100644 --- a/packages/core/src/core/geminiChat.test.ts +++ b/packages/core/src/core/geminiChat.test.ts @@ -28,6 +28,7 @@ import { createAvailabilityServiceMock } from '../availability/testUtils.js'; import type { ModelAvailabilityService } from '../availability/modelAvailabilityService.js'; import * as policyHelpers from '../availability/policyHelpers.js'; import { makeResolvedModelConfig } from '../services/modelConfigServiceTestUtils.js'; +import { LlmRole } from '../telemetry/types.js'; // Mock fs module to prevent actual file system operations during tests const mockFileSystem = new Map(); @@ -259,6 +260,7 @@ describe('GeminiChat', () => { 'test message', 'prompt-id-tool-call-empty-end', new AbortController().signal, + LlmRole.MAIN, ); await expect( (async () => { @@ -312,6 +314,7 @@ describe('GeminiChat', () => { 'test message', 'prompt-id-no-finish-empty-end', new AbortController().signal, + LlmRole.MAIN, ); await expect( (async () => { @@ -359,6 +362,7 @@ describe('GeminiChat', () => { 'test message', 'prompt-id-valid-then-invalid-end', new AbortController().signal, + LlmRole.MAIN, ); await expect( (async () => { @@ -407,6 +411,7 @@ describe('GeminiChat', () => { 'test message', 'prompt-id-empty-chunk-consolidation', new AbortController().signal, + LlmRole.MAIN, ); for await (const _ of stream) { // Consume the stream @@ -466,6 +471,7 @@ describe('GeminiChat', () => { 'test message', 'prompt-id-multi-chunk', new AbortController().signal, + LlmRole.MAIN, ); for await (const _ of stream) { // Consume the stream to trigger history recording. @@ -515,6 +521,7 @@ describe('GeminiChat', () => { 'test message', 'prompt-id-mixed-chunk', new AbortController().signal, + LlmRole.MAIN, ); for await (const _ of stream) { // This loop consumes the stream. @@ -584,6 +591,7 @@ describe('GeminiChat', () => { }, 'prompt-id-stream-1', new AbortController().signal, + LlmRole.MAIN, ); // 4. Assert: The stream processing should throw an InvalidStreamError. @@ -628,6 +636,7 @@ describe('GeminiChat', () => { 'test message', 'prompt-id-1', new AbortController().signal, + LlmRole.MAIN, ); // Should not throw an error @@ -665,6 +674,7 @@ describe('GeminiChat', () => { 'test message', 'prompt-id-1', new AbortController().signal, + LlmRole.MAIN, ); await expect( @@ -701,6 +711,7 @@ describe('GeminiChat', () => { 'test message', 'prompt-id-1', new AbortController().signal, + LlmRole.MAIN, ); await expect( @@ -737,6 +748,7 @@ describe('GeminiChat', () => { 'test message', 'prompt-id-1', new AbortController().signal, + LlmRole.MAIN, ); // Should not throw an error @@ -774,6 +786,7 @@ describe('GeminiChat', () => { 'test', 'prompt-id-malformed', new AbortController().signal, + LlmRole.MAIN, ); // Should throw an error @@ -821,6 +834,7 @@ describe('GeminiChat', () => { 'test retry', 'prompt-id-retry-malformed', new AbortController().signal, + LlmRole.MAIN, ); const events: StreamEvent[] = []; for await (const event of stream) { @@ -878,6 +892,7 @@ describe('GeminiChat', () => { 'hello', 'prompt-id-1', new AbortController().signal, + LlmRole.MAIN, ); for await (const _ of stream) { // consume stream @@ -903,6 +918,7 @@ describe('GeminiChat', () => { }, }, 'prompt-id-1', + LlmRole.MAIN, ); }); @@ -926,6 +942,7 @@ describe('GeminiChat', () => { 'hello', 'prompt-id-thinking-level', new AbortController().signal, + LlmRole.MAIN, ); for await (const _ of stream) { // consume stream @@ -942,6 +959,7 @@ describe('GeminiChat', () => { }), }), 'prompt-id-thinking-level', + LlmRole.MAIN, ); }); @@ -965,6 +983,7 @@ describe('GeminiChat', () => { 'hello', 'prompt-id-thinking-budget', new AbortController().signal, + LlmRole.MAIN, ); for await (const _ of stream) { // consume stream @@ -975,12 +994,13 @@ describe('GeminiChat', () => { model: 'gemini-2.0-flash', config: expect.objectContaining({ thinkingConfig: { - thinkingBudget: DEFAULT_THINKING_MODE, + thinkingBudget: 8192, thinkingLevel: undefined, }, }), }), 'prompt-id-thinking-budget', + LlmRole.MAIN, ); }); }); @@ -1032,6 +1052,7 @@ describe('GeminiChat', () => { 'test', 'prompt-id-no-retry', new AbortController().signal, + LlmRole.MAIN, ); await expect( @@ -1080,6 +1101,7 @@ describe('GeminiChat', () => { 'test message', 'prompt-id-yield-retry', new AbortController().signal, + LlmRole.MAIN, ); const events: StreamEvent[] = []; for await (const event of stream) { @@ -1122,6 +1144,7 @@ describe('GeminiChat', () => { 'test', 'prompt-id-retry-success', new AbortController().signal, + LlmRole.MAIN, ); const chunks: StreamEvent[] = []; for await (const chunk of stream) { @@ -1194,6 +1217,7 @@ describe('GeminiChat', () => { 'test message', 'prompt-id-retry-temperature', new AbortController().signal, + LlmRole.MAIN, ); for await (const _ of stream) { @@ -1215,6 +1239,7 @@ describe('GeminiChat', () => { }), }), 'prompt-id-retry-temperature', + LlmRole.MAIN, ); // Second call (retry) should have temperature 1 @@ -1228,6 +1253,7 @@ describe('GeminiChat', () => { }), }), 'prompt-id-retry-temperature', + LlmRole.MAIN, ); }); @@ -1253,6 +1279,7 @@ describe('GeminiChat', () => { 'test', 'prompt-id-retry-fail', new AbortController().signal, + LlmRole.MAIN, ); await expect(async () => { for await (const _ of stream) { @@ -1319,6 +1346,7 @@ describe('GeminiChat', () => { 'test message', 'prompt-id-400', new AbortController().signal, + LlmRole.MAIN, ); await expect( @@ -1358,9 +1386,11 @@ describe('GeminiChat', () => { 'test message', 'prompt-id-429-retry', new AbortController().signal, + LlmRole.MAIN, ); const events: StreamEvent[] = []; + for await (const event of stream) { events.push(event); } @@ -1407,9 +1437,11 @@ describe('GeminiChat', () => { 'test message', 'prompt-id-500-retry', new AbortController().signal, + LlmRole.MAIN, ); const events: StreamEvent[] = []; + for await (const event of stream) { events.push(event); } @@ -1464,9 +1496,11 @@ describe('GeminiChat', () => { 'test message', 'prompt-id-fetch-error-retry', new AbortController().signal, + LlmRole.MAIN, ); const events: StreamEvent[] = []; + for await (const event of stream) { events.push(event); } @@ -1528,6 +1562,7 @@ describe('GeminiChat', () => { 'Second question', 'prompt-id-retry-existing', new AbortController().signal, + LlmRole.MAIN, ); for await (const _ of stream) { // consume stream @@ -1600,6 +1635,7 @@ describe('GeminiChat', () => { 'test empty stream', 'prompt-id-empty-stream', new AbortController().signal, + LlmRole.MAIN, ); const chunks: StreamEvent[] = []; for await (const chunk of stream) { @@ -1681,6 +1717,7 @@ describe('GeminiChat', () => { 'first', 'prompt-1', new AbortController().signal, + LlmRole.MAIN, ); const firstStreamIterator = firstStream[Symbol.asyncIterator](); await firstStreamIterator.next(); @@ -1691,6 +1728,7 @@ describe('GeminiChat', () => { 'second', 'prompt-2', new AbortController().signal, + LlmRole.MAIN, ); // 5. Assert that only one API call has been made so far. @@ -1796,6 +1834,7 @@ describe('GeminiChat', () => { 'trigger 429', 'prompt-id-fb1', new AbortController().signal, + LlmRole.MAIN, ); // Consume stream to trigger logic @@ -1862,6 +1901,7 @@ describe('GeminiChat', () => { 'test message', 'prompt-id-discard-test', new AbortController().signal, + LlmRole.MAIN, ); const events: StreamEvent[] = []; for await (const event of stream) { @@ -2078,6 +2118,7 @@ describe('GeminiChat', () => { 'test', 'prompt-healthy', new AbortController().signal, + LlmRole.MAIN, ); for await (const _ of stream) { // consume @@ -2113,6 +2154,7 @@ describe('GeminiChat', () => { 'test', 'prompt-sticky-once', new AbortController().signal, + LlmRole.MAIN, ); for await (const _ of stream) { // consume @@ -2163,6 +2205,7 @@ describe('GeminiChat', () => { 'test', 'prompt-fallback-arg', new AbortController().signal, + LlmRole.MAIN, ); for await (const _ of stream) { // consume @@ -2241,6 +2284,7 @@ describe('GeminiChat', () => { 'test', 'prompt-config-refresh', new AbortController().signal, + LlmRole.MAIN, ); // Consume to drive both attempts for await (const _ of stream) { @@ -2253,9 +2297,12 @@ describe('GeminiChat', () => { 1, expect.objectContaining({ model: 'model-a', - config: expect.objectContaining({ temperature: 0.1 }), + config: expect.objectContaining({ + temperature: 0.1, + }), }), expect.any(String), + LlmRole.MAIN, ); expect( mockContentGenerator.generateContentStream, @@ -2263,9 +2310,12 @@ describe('GeminiChat', () => { 2, expect.objectContaining({ model: 'model-b', - config: expect.objectContaining({ temperature: 0.9 }), + config: expect.objectContaining({ + temperature: 0.9, + }), }), expect.any(String), + LlmRole.MAIN, ); }); }); diff --git a/packages/core/src/core/geminiChat.ts b/packages/core/src/core/geminiChat.ts index 3dc91e1b6c6..57be3771448 100644 --- a/packages/core/src/core/geminiChat.ts +++ b/packages/core/src/core/geminiChat.ts @@ -54,6 +54,7 @@ import { fireBeforeModelHook, fireBeforeToolSelectionHook, } from './geminiChatHookTriggers.js'; +import type { LlmRole } from '../telemetry/types.js'; export enum StreamEventType { /** A regular content chunk from the API. */ @@ -260,6 +261,7 @@ export class GeminiChat { message: PartListUnion, prompt_id: string, signal: AbortSignal, + role: LlmRole, ): Promise> { await this.sendPromise; @@ -316,6 +318,7 @@ export class GeminiChat { requestContents, prompt_id, signal, + role, ); isConnectionPhase = false; for await (const chunk of stream) { @@ -383,6 +386,7 @@ export class GeminiChat { requestContents: Content[], prompt_id: string, abortSignal: AbortSignal, + role: LlmRole, ): Promise> { const contentsForPreviewModel = this.ensureActiveLoopHasThoughtSignatures(requestContents); @@ -517,6 +521,7 @@ export class GeminiChat { config, }, prompt_id, + role, ); }; diff --git a/packages/core/src/core/geminiChat_network_retry.test.ts b/packages/core/src/core/geminiChat_network_retry.test.ts index d8bd4b726db..778585e0a90 100644 --- a/packages/core/src/core/geminiChat_network_retry.test.ts +++ b/packages/core/src/core/geminiChat_network_retry.test.ts @@ -14,6 +14,7 @@ import { setSimulate429 } from '../utils/testUtils.js'; import { HookSystem } from '../hooks/hookSystem.js'; import { createMockMessageBus } from '../test-utils/mock-message-bus.js'; import { createAvailabilityServiceMock } from '../availability/testUtils.js'; +import { LlmRole } from '../telemetry/types.js'; // Mock fs module vi.mock('node:fs', () => ({ @@ -150,6 +151,7 @@ describe('GeminiChat Network Retries', () => { 'test message', 'prompt-id-retry-network', new AbortController().signal, + LlmRole.MAIN, ); const events: StreamEvent[] = []; @@ -219,6 +221,7 @@ describe('GeminiChat Network Retries', () => { 'test message', 'prompt-id-retry-fetch', new AbortController().signal, + LlmRole.MAIN, ); const events: StreamEvent[] = []; @@ -259,6 +262,7 @@ describe('GeminiChat Network Retries', () => { 'test message', 'prompt-id-no-retry', new AbortController().signal, + LlmRole.MAIN, ); await expect(async () => { diff --git a/packages/core/src/core/loggingContentGenerator.test.ts b/packages/core/src/core/loggingContentGenerator.test.ts index e591f86be98..d884c3e5337 100644 --- a/packages/core/src/core/loggingContentGenerator.test.ts +++ b/packages/core/src/core/loggingContentGenerator.test.ts @@ -30,7 +30,7 @@ import type { import type { ContentGenerator } from './contentGenerator.js'; import { LoggingContentGenerator } from './loggingContentGenerator.js'; import type { Config } from '../config/config.js'; -import { ApiRequestEvent } from '../telemetry/types.js'; +import { ApiRequestEvent, LlmRole } from '../telemetry/types.js'; describe('LoggingContentGenerator', () => { let wrapped: ContentGenerator; @@ -87,13 +87,18 @@ describe('LoggingContentGenerator', () => { const promise = loggingContentGenerator.generateContent( req, userPromptId, + LlmRole.MAIN, ); vi.advanceTimersByTime(1000); await promise; - expect(wrapped.generateContent).toHaveBeenCalledWith(req, userPromptId); + expect(wrapped.generateContent).toHaveBeenCalledWith( + req, + userPromptId, + LlmRole.MAIN, + ); expect(logApiRequest).toHaveBeenCalledWith( config, expect.any(ApiRequestEvent), @@ -116,6 +121,7 @@ describe('LoggingContentGenerator', () => { const promise = loggingContentGenerator.generateContent( req, userPromptId, + LlmRole.MAIN, ); vi.advanceTimersByTime(1000); @@ -154,12 +160,17 @@ describe('LoggingContentGenerator', () => { vi.mocked(wrapped.generateContentStream).mockResolvedValue( createAsyncGenerator(), ); + const startTime = new Date('2025-01-01T00:00:00.000Z'); + vi.setSystemTime(startTime); const stream = await loggingContentGenerator.generateContentStream( req, + userPromptId, + + LlmRole.MAIN, ); vi.advanceTimersByTime(1000); @@ -171,6 +182,7 @@ describe('LoggingContentGenerator', () => { expect(wrapped.generateContentStream).toHaveBeenCalledWith( req, userPromptId, + LlmRole.MAIN, ); expect(logApiRequest).toHaveBeenCalledWith( config, @@ -201,6 +213,7 @@ describe('LoggingContentGenerator', () => { const stream = await loggingContentGenerator.generateContentStream( req, userPromptId, + LlmRole.MAIN, ); vi.advanceTimersByTime(1000); diff --git a/packages/core/src/core/loggingContentGenerator.ts b/packages/core/src/core/loggingContentGenerator.ts index de16b8de69e..0a7064238ed 100644 --- a/packages/core/src/core/loggingContentGenerator.ts +++ b/packages/core/src/core/loggingContentGenerator.ts @@ -22,6 +22,7 @@ import { ApiResponseEvent, ApiErrorEvent, } from '../telemetry/types.js'; +import type { LlmRole } from '../telemetry/llmRole.js'; import type { Config } from '../config/config.js'; import { logApiError, @@ -55,6 +56,7 @@ export class LoggingContentGenerator implements ContentGenerator { contents: Content[], model: string, promptId: string, + role: LlmRole, generationConfig?: GenerateContentConfig, serverDetails?: ServerDetails, ): void { @@ -70,6 +72,7 @@ export class LoggingContentGenerator implements ContentGenerator { server: serverDetails, }, requestText, + role, ), ); } @@ -112,6 +115,7 @@ export class LoggingContentGenerator implements ContentGenerator { durationMs: number, model: string, prompt_id: string, + role: LlmRole, responseId: string | undefined, responseCandidates?: Candidate[], usageMetadata?: GenerateContentResponseUsageMetadata, @@ -137,6 +141,7 @@ export class LoggingContentGenerator implements ContentGenerator { this.config.getContentGeneratorConfig()?.authType, usageMetadata, responseText, + role, ), ); } @@ -147,6 +152,7 @@ export class LoggingContentGenerator implements ContentGenerator { model: string, prompt_id: string, requestContents: Content[], + role: LlmRole, generationConfig?: GenerateContentConfig, serverDetails?: ServerDetails, ): void { @@ -170,6 +176,7 @@ export class LoggingContentGenerator implements ContentGenerator { isStructuredError(error) ? (error as StructuredError).status : undefined, + role, ), ); } @@ -177,6 +184,7 @@ export class LoggingContentGenerator implements ContentGenerator { async generateContent( req: GenerateContentParameters, userPromptId: string, + role: LlmRole, ): Promise { return runInDevTraceSpan( { @@ -192,6 +200,7 @@ export class LoggingContentGenerator implements ContentGenerator { contents, req.model, userPromptId, + role, req.config, serverDetails, ); @@ -199,6 +208,7 @@ export class LoggingContentGenerator implements ContentGenerator { const response = await this.wrapped.generateContent( req, userPromptId, + role, ); spanMetadata.output = { response, @@ -210,6 +220,7 @@ export class LoggingContentGenerator implements ContentGenerator { durationMs, response.modelVersion || req.model, userPromptId, + role, response.responseId, response.candidates, response.usageMetadata, @@ -226,6 +237,7 @@ export class LoggingContentGenerator implements ContentGenerator { req.model, userPromptId, contents, + role, req.config, serverDetails, ); @@ -238,6 +250,7 @@ export class LoggingContentGenerator implements ContentGenerator { async generateContentStream( req: GenerateContentParameters, userPromptId: string, + role: LlmRole, ): Promise> { return runInDevTraceSpan( { @@ -255,13 +268,18 @@ export class LoggingContentGenerator implements ContentGenerator { toContents(req.contents), req.model, userPromptId, + role, req.config, serverDetails, ); let stream: AsyncGenerator; try { - stream = await this.wrapped.generateContentStream(req, userPromptId); + stream = await this.wrapped.generateContentStream( + req, + userPromptId, + role, + ); } catch (error) { const durationMs = Date.now() - startTime; this._logApiError( @@ -270,6 +288,7 @@ export class LoggingContentGenerator implements ContentGenerator { req.model, userPromptId, toContents(req.contents), + role, req.config, serverDetails, ); @@ -281,6 +300,7 @@ export class LoggingContentGenerator implements ContentGenerator { stream, startTime, userPromptId, + role, spanMetadata, endSpan, ); @@ -293,6 +313,7 @@ export class LoggingContentGenerator implements ContentGenerator { stream: AsyncGenerator, startTime: number, userPromptId: string, + role: LlmRole, spanMetadata: SpanMetadata, endSpan: () => void, ): AsyncGenerator { @@ -316,6 +337,7 @@ export class LoggingContentGenerator implements ContentGenerator { durationMs, responses[0]?.modelVersion || req.model, userPromptId, + role, responses[0]?.responseId, responses.flatMap((response) => response.candidates || []), lastUsageMetadata, @@ -339,6 +361,7 @@ export class LoggingContentGenerator implements ContentGenerator { responses[0]?.modelVersion || req.model, userPromptId, requestContents, + role, req.config, serverDetails, ); diff --git a/packages/core/src/core/recordingContentGenerator.test.ts b/packages/core/src/core/recordingContentGenerator.test.ts index c69c62ebfa0..cbdb239ecfe 100644 --- a/packages/core/src/core/recordingContentGenerator.test.ts +++ b/packages/core/src/core/recordingContentGenerator.test.ts @@ -18,6 +18,7 @@ import { describe, it, expect, vi, beforeEach, type Mock } from 'vitest'; import { safeJsonStringify } from '../utils/safeJsonStringify.js'; import type { ContentGenerator } from './contentGenerator.js'; import { RecordingContentGenerator } from './recordingContentGenerator.js'; +import { LlmRole } from '../telemetry/types.js'; vi.mock('node:fs', () => ({ appendFileSync: vi.fn(), @@ -51,9 +52,14 @@ describe('RecordingContentGenerator', () => { const response = await recorder.generateContent( {} as GenerateContentParameters, 'id1', + LlmRole.MAIN, ); expect(response).toEqual(mockResponse); - expect(mockRealGenerator.generateContent).toHaveBeenCalledWith({}, 'id1'); + expect(mockRealGenerator.generateContent).toHaveBeenCalledWith( + {}, + 'id1', + LlmRole.MAIN, + ); expect(appendFileSync).toHaveBeenCalledWith( filePath, @@ -90,6 +96,7 @@ describe('RecordingContentGenerator', () => { const stream = await recorder.generateContentStream( {} as GenerateContentParameters, 'id1', + LlmRole.MAIN, ); const responses = []; for await (const response of stream) { @@ -100,6 +107,7 @@ describe('RecordingContentGenerator', () => { expect(mockRealGenerator.generateContentStream).toHaveBeenCalledWith( {}, 'id1', + LlmRole.MAIN, ); expect(appendFileSync).toHaveBeenCalledWith( diff --git a/packages/core/src/core/recordingContentGenerator.ts b/packages/core/src/core/recordingContentGenerator.ts index 27abcb418ff..121cd838f4a 100644 --- a/packages/core/src/core/recordingContentGenerator.ts +++ b/packages/core/src/core/recordingContentGenerator.ts @@ -17,6 +17,7 @@ import type { ContentGenerator } from './contentGenerator.js'; import type { FakeResponse } from './fakeContentGenerator.js'; import type { UserTierId } from '../code_assist/types.js'; import { safeJsonStringify } from '../utils/safeJsonStringify.js'; +import type { LlmRole } from '../telemetry/types.js'; // A ContentGenerator that wraps another content generator and records all the // responses, with the ability to write them out to a file. These files are @@ -35,10 +36,12 @@ export class RecordingContentGenerator implements ContentGenerator { async generateContent( request: GenerateContentParameters, userPromptId: string, + role: LlmRole, ): Promise { const response = await this.realGenerator.generateContent( request, userPromptId, + role, ); const recordedResponse: FakeResponse = { method: 'generateContent', @@ -54,6 +57,7 @@ export class RecordingContentGenerator implements ContentGenerator { async generateContentStream( request: GenerateContentParameters, userPromptId: string, + role: LlmRole, ): Promise> { const recordedResponse: FakeResponse = { method: 'generateContentStream', @@ -63,6 +67,7 @@ export class RecordingContentGenerator implements ContentGenerator { const realResponses = await this.realGenerator.generateContentStream( request, userPromptId, + role, ); async function* stream(filePath: string) { diff --git a/packages/core/src/core/turn.test.ts b/packages/core/src/core/turn.test.ts index e951d809330..20803b86651 100644 --- a/packages/core/src/core/turn.test.ts +++ b/packages/core/src/core/turn.test.ts @@ -14,6 +14,7 @@ import type { GenerateContentResponse, Part, Content } from '@google/genai'; import { reportError } from '../utils/errorReporting.js'; import type { GeminiChat } from './geminiChat.js'; import { InvalidStreamError, StreamEventType } from './geminiChat.js'; +import { LlmRole } from '../telemetry/types.js'; const mockSendMessageStream = vi.fn(); const mockGetHistory = vi.fn(); @@ -102,6 +103,7 @@ describe('Turn', () => { reqParts, 'prompt-id-1', expect.any(AbortSignal), + LlmRole.MAIN, ); expect(events).toEqual([ diff --git a/packages/core/src/core/turn.ts b/packages/core/src/core/turn.ts index 86c7fdf49c6..af0b7ee5752 100644 --- a/packages/core/src/core/turn.ts +++ b/packages/core/src/core/turn.ts @@ -32,6 +32,7 @@ import { parseThought, type ThoughtSummary } from '../utils/thoughtUtils.js'; import { createUserContent } from '@google/genai'; import type { ModelConfigKey } from '../services/modelConfigService.js'; import { getCitations } from '../utils/generateContentResponseUtilities.js'; +import { LlmRole } from '../telemetry/types.js'; // Define a structure for tools passed to the server export interface ServerTool { @@ -242,6 +243,7 @@ export class Turn { modelConfigKey: ModelConfigKey, req: PartListUnion, signal: AbortSignal, + role: LlmRole = LlmRole.MAIN, ): AsyncGenerator { try { // Note: This assumes `sendMessageStream` yields events like @@ -251,6 +253,7 @@ export class Turn { req, this.prompt_id, signal, + role, ); for await (const streamEvent of responseStream) { diff --git a/packages/core/src/output/json-formatter.test.ts b/packages/core/src/output/json-formatter.test.ts index 14d2cb47c4a..13321fae77a 100644 --- a/packages/core/src/output/json-formatter.test.ts +++ b/packages/core/src/output/json-formatter.test.ts @@ -79,6 +79,7 @@ describe('JsonFormatter', () => { thoughts: 103, tool: 0, }, + roles: {}, }, 'gemini-2.5-flash': { api: { @@ -95,6 +96,7 @@ describe('JsonFormatter', () => { thoughts: 138, tool: 0, }, + roles: {}, }, }, tools: { diff --git a/packages/core/src/output/stream-json-formatter.test.ts b/packages/core/src/output/stream-json-formatter.test.ts index 557b72a0a9d..69dbaac23bb 100644 --- a/packages/core/src/output/stream-json-formatter.test.ts +++ b/packages/core/src/output/stream-json-formatter.test.ts @@ -289,6 +289,7 @@ describe('StreamJsonFormatter', () => { thoughts: 0, tool: 0, }, + roles: {}, }; metrics.tools.totalCalls = 2; metrics.tools.totalDecisions[ToolCallDecision.AUTO_ACCEPT] = 2; @@ -319,6 +320,7 @@ describe('StreamJsonFormatter', () => { thoughts: 0, tool: 0, }, + roles: {}, }; metrics.models['gemini-ultra'] = { api: { totalRequests: 1, totalErrors: 0, totalLatencyMs: 2000 }, @@ -331,6 +333,7 @@ describe('StreamJsonFormatter', () => { thoughts: 0, tool: 0, }, + roles: {}, }; metrics.tools.totalCalls = 5; @@ -360,6 +363,7 @@ describe('StreamJsonFormatter', () => { thoughts: 0, tool: 0, }, + roles: {}, }; const result = formatter.convertToStreamStats(metrics, 1200); diff --git a/packages/core/src/routing/strategies/classifierStrategy.ts b/packages/core/src/routing/strategies/classifierStrategy.ts index 4747bc5352a..944efc6041c 100644 --- a/packages/core/src/routing/strategies/classifierStrategy.ts +++ b/packages/core/src/routing/strategies/classifierStrategy.ts @@ -20,6 +20,7 @@ import { isFunctionResponse, } from '../../utils/messageInspectors.js'; import { debugLogger } from '../../utils/debugLogger.js'; +import { LlmRole } from '../../telemetry/types.js'; // The number of recent history turns to provide to the router for context. const HISTORY_TURNS_FOR_CONTEXT = 4; @@ -161,6 +162,7 @@ export class ClassifierStrategy implements RoutingStrategy { systemInstruction: CLASSIFIER_SYSTEM_PROMPT, abortSignal: context.signal, promptId, + role: LlmRole.UTILITY_ROUTER, }); const routerResponse = ClassifierResponseSchema.parse(jsonResponse); diff --git a/packages/core/src/services/chatCompressionService.ts b/packages/core/src/services/chatCompressionService.ts index 2336b75f551..d71fb81cedc 100644 --- a/packages/core/src/services/chatCompressionService.ts +++ b/packages/core/src/services/chatCompressionService.ts @@ -24,6 +24,7 @@ import { } from '../config/models.js'; import { firePreCompressHook } from '../core/sessionHookTriggers.js'; import { PreCompressTrigger } from '../hooks/types.js'; +import { LlmRole } from '../telemetry/types.js'; /** * Default threshold for compression token count as a fraction of the model's @@ -194,6 +195,7 @@ export class ChatCompressionService { promptId, // TODO(joshualitt): wire up a sensible abort signal, abortSignal: new AbortController().signal, + role: LlmRole.UTILITY_COMPRESSOR, }); const summary = getResponseText(summaryResponse) ?? ''; diff --git a/packages/core/src/services/loopDetectionService.ts b/packages/core/src/services/loopDetectionService.ts index d8049dbdb68..41c5566a425 100644 --- a/packages/core/src/services/loopDetectionService.ts +++ b/packages/core/src/services/loopDetectionService.ts @@ -25,6 +25,7 @@ import { isFunctionResponse, } from '../utils/messageInspectors.js'; import { debugLogger } from '../utils/debugLogger.js'; +import { LlmRole } from '../telemetry/types.js'; const TOOL_CALL_LOOP_THRESHOLD = 5; const CONTENT_LOOP_THRESHOLD = 10; @@ -529,6 +530,7 @@ export class LoopDetectionService { abortSignal: signal, promptId: this.promptId, maxAttempts: 2, + role: LlmRole.UTILITY_LOOP_DETECTOR, }); if ( diff --git a/packages/core/src/services/sessionSummaryService.ts b/packages/core/src/services/sessionSummaryService.ts index 98ffd66fcac..09c60a2e310 100644 --- a/packages/core/src/services/sessionSummaryService.ts +++ b/packages/core/src/services/sessionSummaryService.ts @@ -10,6 +10,7 @@ import { partListUnionToString } from '../core/geminiRequest.js'; import { debugLogger } from '../utils/debugLogger.js'; import type { Content } from '@google/genai'; import { getResponseText } from '../utils/partUtils.js'; +import { LlmRole } from '../telemetry/types.js'; const DEFAULT_MAX_MESSAGES = 20; const DEFAULT_TIMEOUT_MS = 5000; @@ -124,6 +125,7 @@ export class SessionSummaryService { contents, abortSignal: abortController.signal, promptId: 'session-summary-generation', + role: LlmRole.UTILITY_SUMMARIZER, }); const summary = getResponseText(response); diff --git a/packages/core/src/telemetry/index.ts b/packages/core/src/telemetry/index.ts index 11bc00773f9..661dce3c0bd 100644 --- a/packages/core/src/telemetry/index.ts +++ b/packages/core/src/telemetry/index.ts @@ -63,6 +63,7 @@ export { WebFetchFallbackAttemptEvent, ToolCallDecision, } from './types.js'; +export { LlmRole } from './llmRole.js'; export { makeSlashCommandEvent, makeChatCompressionEvent } from './types.js'; export type { TelemetryEvent } from './types.js'; export { SpanStatusCode, ValueType } from '@opentelemetry/api'; diff --git a/packages/core/src/telemetry/llmRole.ts b/packages/core/src/telemetry/llmRole.ts new file mode 100644 index 00000000000..6f8474dbb73 --- /dev/null +++ b/packages/core/src/telemetry/llmRole.ts @@ -0,0 +1,19 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +export enum LlmRole { + MAIN = 'main', + SUBAGENT = 'subagent', + SUBAGENT_CODEBASE_INVESTIGATOR = 'subagent_codebase_investigator', + UTILITY_TOOL = 'utility_tool', + UTILITY_COMPRESSOR = 'utility_compressor', + UTILITY_SUMMARIZER = 'utility_summarizer', + UTILITY_ROUTER = 'utility_router', + UTILITY_LOOP_DETECTOR = 'utility_loop_detector', + UTILITY_NEXT_SPEAKER = 'utility_next_speaker', + UTILITY_EDIT_CORRECTOR = 'utility_edit_corrector', + UTILITY_AUTOCOMPLETE = 'utility_autocomplete', +} diff --git a/packages/core/src/telemetry/loggers.ts b/packages/core/src/telemetry/loggers.ts index 8a2c08f2a5e..411def04e09 100644 --- a/packages/core/src/telemetry/loggers.ts +++ b/packages/core/src/telemetry/loggers.ts @@ -184,7 +184,12 @@ export function logApiRequest(config: Config, event: ApiRequestEvent): void { ClearcutLogger.getInstance(config)?.logApiRequestEvent(event); bufferTelemetryEvent(() => { const logger = logs.getLogger(SERVICE_NAME); - logger.emit(event.toLogRecord(config)); + const logRecord: LogRecord = event.toLogRecord(config); + if (event.role) { + if (!logRecord.attributes) logRecord.attributes = {}; + logRecord.attributes['role'] = event.role; + } + logger.emit(logRecord); logger.emit(event.toSemanticLogRecord(config)); }); } diff --git a/packages/core/src/telemetry/types.ts b/packages/core/src/telemetry/types.ts index 95893f34d02..79df62e4c44 100644 --- a/packages/core/src/telemetry/types.ts +++ b/packages/core/src/telemetry/types.ts @@ -39,6 +39,8 @@ import { toSystemInstruction, } from './semantic.js'; import { sanitizeHookName } from './sanitize.js'; +import { LlmRole } from './llmRole.js'; +export { LlmRole }; export interface BaseTelemetryEvent { 'event.name': string; @@ -363,17 +365,20 @@ export class ApiRequestEvent implements BaseTelemetryEvent { model: string; prompt: GenAIPromptDetails; request_text?: string; + role?: LlmRole; constructor( model: string, prompt_details: GenAIPromptDetails, request_text?: string, + role?: LlmRole, ) { this['event.name'] = 'api_request'; this['event.timestamp'] = new Date().toISOString(); this.model = model; this.prompt = prompt_details; this.request_text = request_text; + this.role = role; } toLogRecord(config: Config): LogRecord { @@ -385,6 +390,9 @@ export class ApiRequestEvent implements BaseTelemetryEvent { prompt_id: this.prompt.prompt_id, request_text: this.request_text, }; + if (this.role) { + attributes['role'] = this.role; + } return { body: `API request to ${this.model}.`, attributes }; } @@ -433,6 +441,7 @@ export class ApiErrorEvent implements BaseTelemetryEvent { status_code?: number | string; duration_ms: number; auth_type?: string; + role?: LlmRole; constructor( model: string, @@ -442,6 +451,7 @@ export class ApiErrorEvent implements BaseTelemetryEvent { auth_type?: string, error_type?: string, status_code?: number | string, + role?: LlmRole, ) { this['event.name'] = 'api_error'; this['event.timestamp'] = new Date().toISOString(); @@ -452,6 +462,7 @@ export class ApiErrorEvent implements BaseTelemetryEvent { this.duration_ms = duration_ms; this.prompt = prompt_details; this.auth_type = auth_type; + this.role = role; } toLogRecord(config: Config): LogRecord { @@ -470,6 +481,10 @@ export class ApiErrorEvent implements BaseTelemetryEvent { auth_type: this.auth_type, }; + if (this.role) { + attributes['role'] = this.role; + } + if (this.error_type) { attributes['error.type'] = this.error_type; } @@ -578,6 +593,7 @@ export class ApiResponseEvent implements BaseTelemetryEvent { response: GenAIResponseDetails; usage: GenAIUsageDetails; finish_reasons: OTelFinishReason[]; + role?: LlmRole; constructor( model: string, @@ -587,6 +603,7 @@ export class ApiResponseEvent implements BaseTelemetryEvent { auth_type?: string, usage_data?: GenerateContentResponseUsageMetadata, response_text?: string, + role?: LlmRole, ) { this['event.name'] = 'api_response'; this['event.timestamp'] = new Date().toISOString(); @@ -607,6 +624,7 @@ export class ApiResponseEvent implements BaseTelemetryEvent { total_token_count: usage_data?.totalTokenCount ?? 0, }; this.finish_reasons = toFinishReasons(this.response.candidates); + this.role = role; } toLogRecord(config: Config): LogRecord { @@ -627,6 +645,9 @@ export class ApiResponseEvent implements BaseTelemetryEvent { status_code: this.status_code, finish_reasons: this.finish_reasons, }; + if (this.role) { + attributes['role'] = this.role; + } if (this.response_text) { attributes['response_text'] = this.response_text; } diff --git a/packages/core/src/telemetry/uiTelemetry.test.ts b/packages/core/src/telemetry/uiTelemetry.test.ts index 825852f5070..52f0911730c 100644 --- a/packages/core/src/telemetry/uiTelemetry.test.ts +++ b/packages/core/src/telemetry/uiTelemetry.test.ts @@ -181,6 +181,7 @@ describe('UiTelemetryService', () => { thoughts: 2, tool: 3, }, + roles: {}, }); expect(service.getLastPromptTokenCount()).toBe(0); }); @@ -236,6 +237,7 @@ describe('UiTelemetryService', () => { thoughts: 6, tool: 9, }, + roles: {}, }); expect(service.getLastPromptTokenCount()).toBe(0); }); @@ -311,6 +313,7 @@ describe('UiTelemetryService', () => { thoughts: 0, tool: 0, }, + roles: {}, }); }); @@ -356,6 +359,35 @@ describe('UiTelemetryService', () => { thoughts: 2, tool: 3, }, + roles: {}, + }); + }); + + it('should update role metrics when processing an ApiErrorEvent with a role', () => { + const event = { + 'event.name': EVENT_API_ERROR, + model: 'gemini-2.5-pro', + duration_ms: 300, + error: 'Something went wrong', + role: 'utility_tool', + } as unknown as ApiErrorEvent & { 'event.name': typeof EVENT_API_ERROR }; + + service.addEvent(event); + + const metrics = service.getMetrics(); + expect(metrics.models['gemini-2.5-pro'].roles['utility_tool']).toEqual({ + totalRequests: 1, + totalErrors: 1, + totalLatencyMs: 300, + tokens: { + input: 0, + prompt: 0, + candidates: 0, + total: 0, + cached: 0, + thoughts: 0, + tool: 0, + }, }); }); }); diff --git a/packages/core/src/telemetry/uiTelemetry.ts b/packages/core/src/telemetry/uiTelemetry.ts index 6caf2a86063..8c9f2adb83d 100644 --- a/packages/core/src/telemetry/uiTelemetry.ts +++ b/packages/core/src/telemetry/uiTelemetry.ts @@ -18,6 +18,8 @@ import type { ToolCallEvent, } from './types.js'; +import type { LlmRole } from './types.js'; + export type UiEvent = | (ApiResponseEvent & { 'event.name': typeof EVENT_API_RESPONSE }) | (ApiErrorEvent & { 'event.name': typeof EVENT_API_ERROR }) @@ -36,6 +38,21 @@ export interface ToolCallStats { }; } +export interface RoleMetrics { + totalRequests: number; + totalErrors: number; + totalLatencyMs: number; + tokens: { + input: number; + prompt: number; + candidates: number; + total: number; + cached: number; + thoughts: number; + tool: number; + }; +} + export interface ModelMetrics { api: { totalRequests: number; @@ -51,6 +68,7 @@ export interface ModelMetrics { thoughts: number; tool: number; }; + roles: Partial>; } export interface SessionMetrics { @@ -74,6 +92,21 @@ export interface SessionMetrics { }; } +const createInitialRoleMetrics = (): RoleMetrics => ({ + totalRequests: 0, + totalErrors: 0, + totalLatencyMs: 0, + tokens: { + input: 0, + prompt: 0, + candidates: 0, + total: 0, + cached: 0, + thoughts: 0, + tool: 0, + }, +}); + const createInitialModelMetrics = (): ModelMetrics => ({ api: { totalRequests: 0, @@ -89,6 +122,7 @@ const createInitialModelMetrics = (): ModelMetrics => ({ thoughts: 0, tool: 0, }, + roles: {}, }); const createInitialMetrics = (): SessionMetrics => ({ @@ -177,6 +211,25 @@ export class UiTelemetryService extends EventEmitter { 0, modelMetrics.tokens.prompt - modelMetrics.tokens.cached, ); + + if (event.role) { + if (!modelMetrics.roles[event.role]) { + modelMetrics.roles[event.role] = createInitialRoleMetrics(); + } + const roleMetrics = modelMetrics.roles[event.role]!; + roleMetrics.totalRequests++; + roleMetrics.totalLatencyMs += event.duration_ms; + roleMetrics.tokens.prompt += event.usage.input_token_count; + roleMetrics.tokens.candidates += event.usage.output_token_count; + roleMetrics.tokens.total += event.usage.total_token_count; + roleMetrics.tokens.cached += event.usage.cached_content_token_count; + roleMetrics.tokens.thoughts += event.usage.thoughts_token_count; + roleMetrics.tokens.tool += event.usage.tool_token_count; + roleMetrics.tokens.input = Math.max( + 0, + roleMetrics.tokens.prompt - roleMetrics.tokens.cached, + ); + } } private processApiError(event: ApiErrorEvent) { @@ -184,6 +237,16 @@ export class UiTelemetryService extends EventEmitter { modelMetrics.api.totalRequests++; modelMetrics.api.totalErrors++; modelMetrics.api.totalLatencyMs += event.duration_ms; + + if (event.role) { + if (!modelMetrics.roles[event.role]) { + modelMetrics.roles[event.role] = createInitialRoleMetrics(); + } + const roleMetrics = modelMetrics.roles[event.role]!; + roleMetrics.totalRequests++; + roleMetrics.totalErrors++; + roleMetrics.totalLatencyMs += event.duration_ms; + } } private processToolCall(event: ToolCallEvent) { diff --git a/packages/core/src/tools/web-fetch.ts b/packages/core/src/tools/web-fetch.ts index 57591343f3f..d7ce113dde8 100644 --- a/packages/core/src/tools/web-fetch.ts +++ b/packages/core/src/tools/web-fetch.ts @@ -27,6 +27,7 @@ import { logWebFetchFallbackAttempt, WebFetchFallbackAttemptEvent, } from '../telemetry/index.js'; +import { LlmRole } from '../telemetry/llmRole.js'; import { WEB_FETCH_TOOL_NAME } from './tool-names.js'; import { debugLogger } from '../utils/debugLogger.js'; import { retryWithBackoff } from '../utils/retry.js'; @@ -187,6 +188,7 @@ ${textContent} { model: 'web-fetch-fallback' }, [{ role: 'user', parts: [{ text: fallbackPrompt }] }], signal, + LlmRole.UTILITY_TOOL, ); const resultText = getResponseText(result) || ''; return { @@ -271,6 +273,7 @@ ${textContent} { model: 'web-fetch' }, [{ role: 'user', parts: [{ text: userPrompt }] }], signal, // Pass signal + LlmRole.UTILITY_TOOL, ); debugLogger.debug( diff --git a/packages/core/src/tools/web-search.ts b/packages/core/src/tools/web-search.ts index c85ca02a6cd..42891ea1261 100644 --- a/packages/core/src/tools/web-search.ts +++ b/packages/core/src/tools/web-search.ts @@ -14,6 +14,7 @@ import { ToolErrorType } from './tool-error.js'; import { getErrorMessage } from '../utils/errors.js'; import { type Config } from '../config/config.js'; import { getResponseText } from '../utils/partUtils.js'; +import { LlmRole } from '../telemetry/llmRole.js'; interface GroundingChunkWeb { uri?: string; @@ -83,6 +84,7 @@ class WebSearchToolInvocation extends BaseToolInvocation< { model: 'web-search' }, [{ role: 'user', parts: [{ text: this.params.query }] }], signal, + LlmRole.UTILITY_TOOL, ); const responseText = getResponseText(response); diff --git a/packages/core/src/utils/editCorrector.ts b/packages/core/src/utils/editCorrector.ts index 9c74d81f5a9..50f91c93c5a 100644 --- a/packages/core/src/utils/editCorrector.ts +++ b/packages/core/src/utils/editCorrector.ts @@ -22,6 +22,7 @@ import { } from '../utils/messageInspectors.js'; import * as fs from 'node:fs'; import { promptIdContext } from './promptIdContext.js'; +import { LlmRole } from '../telemetry/types.js'; const CODE_CORRECTION_SYSTEM_PROMPT = ` You are an expert code-editing assistant. Your task is to analyze a failed edit attempt and provide a corrected version of the text snippets. @@ -418,6 +419,7 @@ Return ONLY the corrected target snippet in the specified JSON format with the k abortSignal, systemInstruction: CODE_CORRECTION_SYSTEM_PROMPT, promptId: getPromptId(), + role: LlmRole.UTILITY_EDIT_CORRECTOR, }); if ( @@ -507,6 +509,7 @@ Return ONLY the corrected string in the specified JSON format with the key 'corr abortSignal, systemInstruction: CODE_CORRECTION_SYSTEM_PROMPT, promptId: getPromptId(), + role: LlmRole.UTILITY_EDIT_CORRECTOR, }); if ( @@ -577,6 +580,7 @@ Return ONLY the corrected string in the specified JSON format with the key 'corr abortSignal, systemInstruction: CODE_CORRECTION_SYSTEM_PROMPT, promptId: getPromptId(), + role: LlmRole.UTILITY_EDIT_CORRECTOR, }); if ( @@ -644,6 +648,7 @@ Return ONLY the corrected string in the specified JSON format with the key 'corr abortSignal, systemInstruction: CODE_CORRECTION_SYSTEM_PROMPT, promptId: getPromptId(), + role: LlmRole.UTILITY_EDIT_CORRECTOR, }); if ( diff --git a/packages/core/src/utils/llm-edit-fixer.ts b/packages/core/src/utils/llm-edit-fixer.ts index 363c973cd7e..d83bcb2c5fd 100644 --- a/packages/core/src/utils/llm-edit-fixer.ts +++ b/packages/core/src/utils/llm-edit-fixer.ts @@ -10,6 +10,7 @@ import { type BaseLlmClient } from '../core/baseLlmClient.js'; import { LruCache } from './LruCache.js'; import { promptIdContext } from './promptIdContext.js'; import { debugLogger } from './debugLogger.js'; +import { LlmRole } from '../telemetry/types.js'; const MAX_CACHE_SIZE = 50; const GENERATE_JSON_TIMEOUT_MS = 40000; // 40 seconds @@ -182,6 +183,7 @@ export async function FixLLMEditWithInstruction( systemInstruction: EDIT_SYS_PROMPT, promptId, maxAttempts: 1, + role: LlmRole.UTILITY_EDIT_CORRECTOR, }, GENERATE_JSON_TIMEOUT_MS, ); diff --git a/packages/core/src/utils/nextSpeakerChecker.ts b/packages/core/src/utils/nextSpeakerChecker.ts index 76b1c6a440a..e5b54340f6d 100644 --- a/packages/core/src/utils/nextSpeakerChecker.ts +++ b/packages/core/src/utils/nextSpeakerChecker.ts @@ -9,6 +9,7 @@ import type { BaseLlmClient } from '../core/baseLlmClient.js'; import type { GeminiChat } from '../core/geminiChat.js'; import { isFunctionResponse } from './messageInspectors.js'; import { debugLogger } from './debugLogger.js'; +import { LlmRole } from '../telemetry/types.js'; const CHECK_PROMPT = `Analyze *only* the content and structure of your immediately preceding response (your last turn in the conversation history). Based *strictly* on that response, determine who should logically speak next: the 'user' or the 'model' (you). **Decision Rules (apply in order):** @@ -115,6 +116,7 @@ export async function checkNextSpeaker( schema: RESPONSE_SCHEMA, abortSignal, promptId, + role: LlmRole.UTILITY_NEXT_SPEAKER, })) as unknown as NextSpeakerResponse; if ( diff --git a/packages/core/src/utils/summarizer.ts b/packages/core/src/utils/summarizer.ts index b25961e1492..99653d4c598 100644 --- a/packages/core/src/utils/summarizer.ts +++ b/packages/core/src/utils/summarizer.ts @@ -11,6 +11,7 @@ import { getResponseText, partToString } from './partUtils.js'; import { debugLogger } from './debugLogger.js'; import type { ModelConfigKey } from '../services/modelConfigService.js'; import type { Config } from '../config/config.js'; +import { LlmRole } from '../telemetry/llmRole.js'; /** * A function that summarizes the result of a tool execution. @@ -94,6 +95,7 @@ export async function summarizeToolOutput( modelConfigKey, contents, abortSignal, + LlmRole.UTILITY_SUMMARIZER, ); return getResponseText(parsedResponse) || textToSummarize; } catch (error) {