mirror of
https://github.com/QwenLM/qwen-code.git
synced 2025-12-21 09:17:53 +00:00
feat: Change /stats to include more detailed breakdowns (#2615)
This commit is contained in:
Abhi
GitHub
@@ -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<typeof useSessionStats> | undefined
|
||||
> = { current: undefined };
|
||||
@@ -77,150 +60,60 @@ describe('SessionStatsContext', () => {
|
||||
</SessionStatsProvider>,
|
||||
);
|
||||
|
||||
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?.startNewTurn();
|
||||
uiTelemetryService.emit('update', {
|
||||
metrics: newMetrics,
|
||||
lastPromptTokenCount: 100,
|
||||
});
|
||||
});
|
||||
|
||||
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<typeof useSessionStats> | undefined
|
||||
> = { current: undefined };
|
||||
|
||||
render(
|
||||
<SessionStatsProvider>
|
||||
<TestHarness contextRef={contextRef} />
|
||||
</SessionStatsProvider>,
|
||||
);
|
||||
|
||||
act(() => {
|
||||
contextRef.current?.addUsage({ ...mockMetadata1, apiTimeMs: 123 });
|
||||
});
|
||||
|
||||
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<typeof useSessionStats> | undefined
|
||||
> = { current: undefined };
|
||||
|
||||
render(
|
||||
<SessionStatsProvider>
|
||||
<TestHarness contextRef={contextRef} />
|
||||
</SessionStatsProvider>,
|
||||
);
|
||||
|
||||
// 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<typeof useSessionStats> | undefined
|
||||
> = { current: undefined };
|
||||
|
||||
render(
|
||||
<SessionStatsProvider>
|
||||
<TestHarness contextRef={contextRef} />
|
||||
</SessionStatsProvider>,
|
||||
);
|
||||
|
||||
// 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', () => {
|
||||
|
||||
Reference in New Issue
Block a user