From 770f862832dfef477705bee69bd2a84397d105a8 Mon Sep 17 00:00:00 2001 From: Abhi <43648792+abhipatel12@users.noreply.github.com> Date: Sun, 29 Jun 2025 20:44:33 -0400 Subject: feat: Change /stats to include more detailed breakdowns (#2615) --- .../cli/src/ui/contexts/SessionContext.test.tsx | 223 ++++++--------------- 1 file changed, 58 insertions(+), 165 deletions(-) (limited to 'packages/cli/src/ui/contexts/SessionContext.test.tsx') diff --git a/packages/cli/src/ui/contexts/SessionContext.test.tsx b/packages/cli/src/ui/contexts/SessionContext.test.tsx index fedb5341..5b05c284 100644 --- a/packages/cli/src/ui/contexts/SessionContext.test.tsx +++ b/packages/cli/src/ui/contexts/SessionContext.test.tsx @@ -8,28 +8,13 @@ import { type MutableRefObject } from 'react'; import { render } from 'ink-testing-library'; import { renderHook } from '@testing-library/react'; import { act } from 'react-dom/test-utils'; -import { SessionStatsProvider, useSessionStats } from './SessionContext.js'; +import { + SessionStatsProvider, + useSessionStats, + SessionMetrics, +} from './SessionContext.js'; import { describe, it, expect, vi } from 'vitest'; -import { GenerateContentResponseUsageMetadata } from '@google/genai'; - -// Mock data that simulates what the Gemini API would return. -const mockMetadata1: GenerateContentResponseUsageMetadata = { - promptTokenCount: 100, - candidatesTokenCount: 200, - totalTokenCount: 300, - cachedContentTokenCount: 50, - toolUsePromptTokenCount: 10, - thoughtsTokenCount: 20, -}; - -const mockMetadata2: GenerateContentResponseUsageMetadata = { - promptTokenCount: 10, - candidatesTokenCount: 20, - totalTokenCount: 30, - cachedContentTokenCount: 5, - toolUsePromptTokenCount: 1, - thoughtsTokenCount: 2, -}; +import { uiTelemetryService } from '@google/gemini-cli-core'; /** * A test harness component that uses the hook and exposes the context value @@ -60,13 +45,11 @@ describe('SessionStatsContext', () => { const stats = contextRef.current?.stats; expect(stats?.sessionStartTime).toBeInstanceOf(Date); - expect(stats?.currentTurn).toBeDefined(); - expect(stats?.cumulative.turnCount).toBe(0); - expect(stats?.cumulative.totalTokenCount).toBe(0); - expect(stats?.cumulative.promptTokenCount).toBe(0); + expect(stats?.metrics).toBeDefined(); + expect(stats?.metrics.models).toEqual({}); }); - it('should increment turnCount when startNewTurn is called', () => { + it('should update metrics when the uiTelemetryService emits an update', () => { const contextRef: MutableRefObject< ReturnType | undefined > = { current: undefined }; @@ -77,150 +60,60 @@ describe('SessionStatsContext', () => { , ); - act(() => { - contextRef.current?.startNewTurn(); - }); - - const stats = contextRef.current?.stats; - expect(stats?.currentTurn.totalTokenCount).toBe(0); - expect(stats?.cumulative.turnCount).toBe(1); - // Ensure token counts are unaffected - expect(stats?.cumulative.totalTokenCount).toBe(0); - }); - - it('should aggregate token usage correctly when addUsage is called', () => { - const contextRef: MutableRefObject< - ReturnType | undefined - > = { current: undefined }; - - render( - - - , - ); + const newMetrics: SessionMetrics = { + models: { + 'gemini-pro': { + api: { + totalRequests: 1, + totalErrors: 0, + totalLatencyMs: 123, + }, + tokens: { + prompt: 100, + candidates: 200, + total: 300, + cached: 50, + thoughts: 20, + tool: 10, + }, + }, + }, + tools: { + totalCalls: 1, + totalSuccess: 1, + totalFail: 0, + totalDurationMs: 456, + totalDecisions: { + accept: 1, + reject: 0, + modify: 0, + }, + byName: { + 'test-tool': { + count: 1, + success: 1, + fail: 0, + durationMs: 456, + decisions: { + accept: 1, + reject: 0, + modify: 0, + }, + }, + }, + }, + }; act(() => { - contextRef.current?.addUsage({ ...mockMetadata1, apiTimeMs: 123 }); + uiTelemetryService.emit('update', { + metrics: newMetrics, + lastPromptTokenCount: 100, + }); }); const stats = contextRef.current?.stats; - - // Check that token counts are updated - expect(stats?.cumulative.totalTokenCount).toBe( - mockMetadata1.totalTokenCount ?? 0, - ); - expect(stats?.cumulative.promptTokenCount).toBe( - mockMetadata1.promptTokenCount ?? 0, - ); - expect(stats?.cumulative.apiTimeMs).toBe(123); - - // Check that turn count is NOT incremented - expect(stats?.cumulative.turnCount).toBe(0); - - // Check that currentTurn is updated - expect(stats?.currentTurn?.totalTokenCount).toEqual( - mockMetadata1.totalTokenCount, - ); - expect(stats?.currentTurn?.apiTimeMs).toBe(123); - }); - - it('should correctly track a full logical turn with multiple API calls', () => { - const contextRef: MutableRefObject< - ReturnType | undefined - > = { current: undefined }; - - render( - - - , - ); - - // 1. User starts a new turn - act(() => { - contextRef.current?.startNewTurn(); - }); - - // 2. First API call (e.g., prompt with a tool request) - act(() => { - contextRef.current?.addUsage({ ...mockMetadata1, apiTimeMs: 100 }); - }); - - // 3. Second API call (e.g., sending tool response back) - act(() => { - contextRef.current?.addUsage({ ...mockMetadata2, apiTimeMs: 50 }); - }); - - const stats = contextRef.current?.stats; - - // Turn count should only be 1 - expect(stats?.cumulative.turnCount).toBe(1); - - // --- Check Cumulative Stats --- - // These fields should be the SUM of both calls - expect(stats?.cumulative.totalTokenCount).toBe(300 + 30); - expect(stats?.cumulative.candidatesTokenCount).toBe(200 + 20); - expect(stats?.cumulative.thoughtsTokenCount).toBe(20 + 2); - expect(stats?.cumulative.apiTimeMs).toBe(100 + 50); - - // These fields should be the SUM of both calls - expect(stats?.cumulative.promptTokenCount).toBe(100 + 10); - expect(stats?.cumulative.cachedContentTokenCount).toBe(50 + 5); - expect(stats?.cumulative.toolUsePromptTokenCount).toBe(10 + 1); - - // --- Check Current Turn Stats --- - // All fields should be the SUM of both calls for the turn - expect(stats?.currentTurn.totalTokenCount).toBe(300 + 30); - expect(stats?.currentTurn.candidatesTokenCount).toBe(200 + 20); - expect(stats?.currentTurn.thoughtsTokenCount).toBe(20 + 2); - expect(stats?.currentTurn.promptTokenCount).toBe(100 + 10); - expect(stats?.currentTurn.cachedContentTokenCount).toBe(50 + 5); - expect(stats?.currentTurn.toolUsePromptTokenCount).toBe(10 + 1); - expect(stats?.currentTurn.apiTimeMs).toBe(100 + 50); - }); - - it('should overwrite currentResponse with each API call', () => { - const contextRef: MutableRefObject< - ReturnType | undefined - > = { current: undefined }; - - render( - - - , - ); - - // 1. First API call - act(() => { - contextRef.current?.addUsage({ ...mockMetadata1, apiTimeMs: 100 }); - }); - - let stats = contextRef.current?.stats; - - // currentResponse should match the first call - expect(stats?.currentResponse.totalTokenCount).toBe(300); - expect(stats?.currentResponse.apiTimeMs).toBe(100); - - // 2. Second API call - act(() => { - contextRef.current?.addUsage({ ...mockMetadata2, apiTimeMs: 50 }); - }); - - stats = contextRef.current?.stats; - - // currentResponse should now match the second call - expect(stats?.currentResponse.totalTokenCount).toBe(30); - expect(stats?.currentResponse.apiTimeMs).toBe(50); - - // 3. Start a new turn - act(() => { - contextRef.current?.startNewTurn(); - }); - - stats = contextRef.current?.stats; - - // currentResponse should be reset - expect(stats?.currentResponse.totalTokenCount).toBe(0); - expect(stats?.currentResponse.apiTimeMs).toBe(0); + expect(stats?.metrics).toEqual(newMetrics); + expect(stats?.lastPromptTokenCount).toBe(100); }); it('should throw an error when useSessionStats is used outside of a provider', () => { -- cgit v1.2.3