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) {