From cc09cfd7a4e09773523f683ab8da222afdc7dfe1 Mon Sep 17 00:00:00 2001 From: Mahima Shanware Date: Mon, 30 Mar 2026 19:22:22 +0000 Subject: [PATCH] feat(telemetry): add usage metric for /btw command --- packages/core/src/core/client.test.ts | 46 ++++++++++++++++++++- packages/core/src/core/client.ts | 3 ++ packages/core/src/telemetry/index.ts | 2 + packages/core/src/telemetry/metrics.test.ts | 24 +++++++++++ packages/core/src/telemetry/metrics.ts | 16 +++++++ 5 files changed, 90 insertions(+), 1 deletion(-) diff --git a/packages/core/src/core/client.test.ts b/packages/core/src/core/client.test.ts index f8178488bd..f7350bf7f1 100644 --- a/packages/core/src/core/client.test.ts +++ b/packages/core/src/core/client.test.ts @@ -21,7 +21,7 @@ import { type ContentGenerator, type ContentGeneratorConfig, } from './contentGenerator.js'; -import { GeminiChat } from './geminiChat.js'; +import { GeminiChat, StreamEventType, type StreamEvent } from './geminiChat.js'; import type { Config } from '../config/config.js'; import type { AgentLoopContext } from '../config/agent-loop-context.js'; import { @@ -50,6 +50,7 @@ import type { import { ClearcutLogger } from '../telemetry/clearcut-logger/clearcut-logger.js'; import * as policyCatalog from '../availability/policyCatalog.js'; import { LlmRole, LoopType } from '../telemetry/types.js'; +import { recordBtwUsageMetrics } from '../telemetry/index.js'; import { partToString } from '../utils/partUtils.js'; import { coreEvents, CoreEvent } from '../utils/events.js'; import type { MessageBus } from '../confirmation-bus/message-bus.js'; @@ -130,6 +131,7 @@ vi.mock('../telemetry/index.js', () => ({ logApiRequest: vi.fn(), logApiResponse: vi.fn(), logApiError: vi.fn(), + recordBtwUsageMetrics: vi.fn(), })); vi.mock('../ide/ideContext.js'); vi.mock('../telemetry/uiTelemetry.js', () => ({ @@ -3294,6 +3296,48 @@ ${JSON.stringify( }); }); + describe('sendBtwStream', () => { + it('should record telemetry and call chat.sendBtwStream', async () => { + // Arrange + const request = [{ text: 'btw, how does this work?' }]; + const abortSignal = new AbortController().signal; + const promptId = 'test-prompt-id'; + + const mockBtwStream = (async function* (): AsyncGenerator< + StreamEvent, + void, + unknown + > { + yield { + type: StreamEventType.CHUNK, + value: { + text: () => 'this works fine', + } as unknown as GenerateContentResponse, + }; + })(); + + const sendBtwStreamSpy = vi + .spyOn(client.getChat(), 'sendBtwStream') + .mockReturnValue(mockBtwStream); + + // Act + const stream = client.sendBtwStream(request, abortSignal, promptId); + for await (const _ of stream) { + // Consume stream + } + + // Assert + expect(recordBtwUsageMetrics).toHaveBeenCalledWith(mockConfig); + expect(sendBtwStreamSpy).toHaveBeenCalledWith( + expect.any(Object), // modelConfigKey + request, + promptId, + abortSignal, + LlmRole.MAIN, + ); + }); + }); + describe('generateContent', () => { it('should call generateContent with the correct parameters', async () => { const contents = [{ role: 'user', parts: [{ text: 'hello' }] }]; diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts index d9ec90cb58..c2f160736e 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -50,6 +50,7 @@ import { logContentRetryFailure, logNextSpeakerCheck, } from '../telemetry/loggers.js'; +import { recordBtwUsageMetrics } from '../telemetry/index.js'; import type { DefaultHookOutput, AfterAgentHookOutput, @@ -1247,6 +1248,8 @@ export class GeminiClient { yield { type: GeminiEventType.ModelInfo, value: modelToUse }; + recordBtwUsageMetrics(this.config); + // Use a custom role for BTW to avoid side-effects in telemetry if needed, // but for now LlmRole.MAIN is fine as it's the primary model talking. const btwStream = this.getChat().sendBtwStream( diff --git a/packages/core/src/telemetry/index.ts b/packages/core/src/telemetry/index.ts index ea65941e06..050b74c270 100644 --- a/packages/core/src/telemetry/index.ts +++ b/packages/core/src/telemetry/index.ts @@ -129,6 +129,8 @@ export { recordGenAiClientTokenUsage, recordGenAiClientOperationDuration, getConventionAttributes, + // Btw Metrics + recordBtwUsageMetrics, // Performance monitoring functions recordStartupPerformance, recordMemoryUsage, diff --git a/packages/core/src/telemetry/metrics.test.ts b/packages/core/src/telemetry/metrics.test.ts index 0bca699b16..197d750d5a 100644 --- a/packages/core/src/telemetry/metrics.test.ts +++ b/packages/core/src/telemetry/metrics.test.ts @@ -104,6 +104,7 @@ describe('Telemetry Metrics', () => { let recordLinesChangedModule: typeof import('./metrics.js').recordLinesChanged; let recordSlowRenderModule: typeof import('./metrics.js').recordSlowRender; let recordPlanExecutionModule: typeof import('./metrics.js').recordPlanExecution; + let recordBtwUsageMetricsModule: typeof import('./metrics.js').recordBtwUsageMetrics; let recordKeychainAvailabilityModule: typeof import('./metrics.js').recordKeychainAvailability; let recordTokenStorageInitializationModule: typeof import('./metrics.js').recordTokenStorageInitialization; let recordInvalidChunkModule: typeof import('./metrics.js').recordInvalidChunk; @@ -158,6 +159,7 @@ describe('Telemetry Metrics', () => { recordLinesChangedModule = metricsJsModule.recordLinesChanged; recordSlowRenderModule = metricsJsModule.recordSlowRender; recordPlanExecutionModule = metricsJsModule.recordPlanExecution; + recordBtwUsageMetricsModule = metricsJsModule.recordBtwUsageMetrics; recordKeychainAvailabilityModule = metricsJsModule.recordKeychainAvailability; recordTokenStorageInitializationModule = @@ -273,6 +275,28 @@ describe('Telemetry Metrics', () => { }); }); + describe('recordBtwUsageMetrics', () => { + it('does not record metrics if not initialized', () => { + const config = makeFakeConfig({}); + recordBtwUsageMetricsModule(config); + expect(mockCounterAddFn).not.toHaveBeenCalled(); + }); + + it('records a btw usage event when initialized', () => { + const config = makeFakeConfig({}); + initializeMetricsModule(config); + recordBtwUsageMetricsModule(config); + + // Called for session, then for btw usage + expect(mockCounterAddFn).toHaveBeenCalledTimes(2); + expect(mockCounterAddFn).toHaveBeenNthCalledWith(2, 1, { + 'session.id': 'test-session-id', + 'installation.id': 'test-installation-id', + 'user.email': 'test@example.com', + }); + }); + }); + describe('initializeMetrics', () => { const mockConfig = { getSessionId: () => 'test-session-id', diff --git a/packages/core/src/telemetry/metrics.ts b/packages/core/src/telemetry/metrics.ts index 422f0222a5..3850efa7e9 100644 --- a/packages/core/src/telemetry/metrics.ts +++ b/packages/core/src/telemetry/metrics.ts @@ -101,6 +101,7 @@ const FLICKER_FRAME_COUNT = 'gemini_cli.ui.flicker.count'; const SLOW_RENDER_LATENCY = 'gemini_cli.ui.slow_render.latency'; const EXIT_FAIL_COUNT = 'gemini_cli.exit.fail.count'; const PLAN_EXECUTION_COUNT = 'gemini_cli.plan.execution.count'; +const BTW_USAGE_COUNT = 'gemini_cli.btw.usage.count'; const baseMetricDefinition = { getCommonAttributes, @@ -269,6 +270,12 @@ const COUNTER_DEFINITIONS = { approval_mode: string; }, }, + [BTW_USAGE_COUNT]: { + description: 'Counts the usage of the /btw side inquiry command.', + valueType: ValueType.INT, + assign: (c: Counter) => (btwUsageCounter = c), + attributes: {} as Record, + }, [EVENT_HOOK_CALL_COUNT]: { description: 'Counts hook calls, tagged by hook event name and success.', valueType: ValueType.INT, @@ -777,6 +784,7 @@ let agentRecoveryAttemptDurationHistogram: Histogram | undefined; let flickerFrameCounter: Counter | undefined; let exitFailCounter: Counter | undefined; let planExecutionCounter: Counter | undefined; +let btwUsageCounter: Counter | undefined; let slowRenderHistogram: Histogram | undefined; let hookCallCounter: Counter | undefined; let hookCallLatencyHistogram: Histogram | undefined; @@ -1034,6 +1042,14 @@ export function recordPlanExecution( }); } +/** + * Records a metric for when the /btw side inquiry command is used + */ +export function recordBtwUsageMetrics(config: Config): void { + if (!btwUsageCounter || !isMetricsInitialized) return; + btwUsageCounter.add(1, baseMetricDefinition.getCommonAttributes(config)); +} + /** * Records a metric for when a UI frame is slow in rendering */