abbyswag/qwen-code
1
0
Fork 0
mirror of https://github.com/QwenLM/qwen-code.git synced 2025-12-20 08:47:44 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: Display initial token usage metrics in /stats (#879)

Browse Source
This commit is contained in:
Abhi
2025-06-09 20:25:37 -04:00
committed by GitHub
parent 6484dc9008
commit 7f1252d364
11 changed files with 608 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files

75
packages/cli/src/ui/hooks/slashCommandProcessor.test.ts
View File

@@ -61,13 +61,13 @@ import {
MCPServerStatus,
getMCPServerStatus,
} from '@gemini-cli/core';
import { useSession } from '../contexts/SessionContext.js';
import { useSessionStats } from '../contexts/SessionContext.js';
import * as ShowMemoryCommandModule from './useShowMemoryCommand.js';
import { GIT_COMMIT_INFO } from '../../generated/git-commit.js';
vi.mock('../contexts/SessionContext.js', () => ({
useSession: vi.fn(),
useSessionStats: vi.fn(),
}));
vi.mock('./useShowMemoryCommand.js', () => ({
@@ -89,7 +89,7 @@ describe('useSlashCommandProcessor', () => {
let mockPerformMemoryRefresh: ReturnType<typeof vi.fn>;
let mockConfig: Config;
let mockCorgiMode: ReturnType<typeof vi.fn>;
const mockUseSession = useSession as Mock;
const mockUseSessionStats = useSessionStats as Mock;
beforeEach(() => {
mockAddItem = vi.fn();
@@ -105,8 +105,19 @@ describe('useSlashCommandProcessor', () => {
getModel: vi.fn(() => 'test-model'),
} as unknown as Config;
mockCorgiMode = vi.fn();
mockUseSession.mockReturnValue({
startTime: new Date('2025-01-01T00:00:00.000Z'),
mockUseSessionStats.mockReturnValue({
stats: {
sessionStartTime: new Date('2025-01-01T00:00:00.000Z'),
cumulative: {
turnCount: 0,
promptTokenCount: 0,
candidatesTokenCount: 0,
totalTokenCount: 0,
cachedContentTokenCount: 0,
toolUsePromptTokenCount: 0,
thoughtsTokenCount: 0,
},
},
});
(open as Mock).mockClear();
@@ -240,29 +251,55 @@ describe('useSlashCommandProcessor', () => {
});
describe('/stats command', () => {
it('should show the session duration', async () => {
const { handleSlashCommand } = getProcessor();
let commandResult: SlashCommandActionReturn | boolean = false;
// Mock current time
const mockDate = new Date('2025-01-01T00:01:05.000Z');
vi.setSystemTime(mockDate);
await act(async () => {
commandResult = handleSlashCommand('/stats');
it('should show detailed session statistics', async () => {
// Arrange
mockUseSessionStats.mockReturnValue({
stats: {
sessionStartTime: new Date('2025-01-01T00:00:00.000Z'),
cumulative: {
totalTokenCount: 900,
promptTokenCount: 200,
candidatesTokenCount: 400,
cachedContentTokenCount: 100,
turnCount: 1,
toolUsePromptTokenCount: 50,
thoughtsTokenCount: 150,
},
},
});
const { handleSlashCommand } = getProcessor();
const mockDate = new Date('2025-01-01T01:02:03.000Z'); // 1h 2m 3s duration
vi.setSystemTime(mockDate);
// Act
await act(async () => {
handleSlashCommand('/stats');
});
// Assert
const expectedContent = [
` ⎿ Total duration (wall): 1h 2m 3s`,
` Total Token usage:`,
` Turns: 1`,
` Total: 900`,
` ├─ Input: 200`,
` ├─ Output: 400`,
` ├─ Cached: 100`,
` └─ Overhead: 200`,
` ├─ Model thoughts: 150`,
` └─ Tool-use prompts: 50`,
].join('\n');
expect(mockAddItem).toHaveBeenNthCalledWith(
2,
2, // Called after the user message
expect.objectContaining({
type: MessageType.INFO,
text: 'Session duration: 1m 5s',
text: expectedContent,
}),
expect.any(Number),
);
expect(commandResult).toBe(true);
// Restore system time
vi.useRealTimers();
});
});

29
packages/cli/src/ui/hooks/slashCommandProcessor.ts
View File

@@ -11,7 +11,7 @@ import process from 'node:process';
import { UseHistoryManagerReturn } from './useHistoryManager.js';
import { Config, MCPServerStatus, getMCPServerStatus } from '@gemini-cli/core';
import { Message, MessageType, HistoryItemWithoutId } from '../types.js';
import { useSession } from '../contexts/SessionContext.js';
import { useSessionStats } from '../contexts/SessionContext.js';
import { createShowMemoryAction } from './useShowMemoryCommand.js';
import { GIT_COMMIT_INFO } from '../../generated/git-commit.js';
import { formatMemoryUsage } from '../utils/formatters.js';
@@ -50,8 +50,7 @@ export const useSlashCommandProcessor = (
toggleCorgiMode: () => void,
showToolDescriptions: boolean = false,
) => {
const session = useSession();
const session = useSessionStats();
const addMessage = useCallback(
(message: Message) => {
// Convert Message to HistoryItemWithoutId
@@ -147,7 +146,9 @@ export const useSlashCommandProcessor = (
description: 'check session stats',
action: (_mainCommand, _subCommand, _args) => {
const now = new Date();
const duration = now.getTime() - session.startTime.getTime();
const { sessionStartTime, cumulative } = session.stats;
const duration = now.getTime() - sessionStartTime.getTime();
const durationInSeconds = Math.floor(duration / 1000);
const hours = Math.floor(durationInSeconds / 3600);
const minutes = Math.floor((durationInSeconds % 3600) / 60);
@@ -161,9 +162,25 @@ export const useSlashCommandProcessor = (
.filter(Boolean)
.join(' ');
const overheadTotal =
cumulative.thoughtsTokenCount + cumulative.toolUsePromptTokenCount;
const statsContent = [
` ⎿ Total duration (wall): ${durationString}`,
` Total Token usage:`,
` Turns: ${cumulative.turnCount.toLocaleString()}`,
` Total: ${cumulative.totalTokenCount.toLocaleString()}`,
` ├─ Input: ${cumulative.promptTokenCount.toLocaleString()}`,
` ├─ Output: ${cumulative.candidatesTokenCount.toLocaleString()}`,
` ├─ Cached: ${cumulative.cachedContentTokenCount.toLocaleString()}`,
` └─ Overhead: ${overheadTotal.toLocaleString()}`,
` ├─ Model thoughts: ${cumulative.thoughtsTokenCount.toLocaleString()}`,
` └─ Tool-use prompts: ${cumulative.toolUsePromptTokenCount.toLocaleString()}`,
].join('\n');
addMessage({
type: MessageType.INFO,
content: `Session duration: ${durationString}`,
content: statsContent,
timestamp: new Date(),
});
},
@@ -477,7 +494,7 @@ Add any other context about the problem here.
toggleCorgiMode,
config,
showToolDescriptions,
session.startTime,
session,
],
);

68
packages/cli/src/ui/hooks/useGeminiStream.test.tsx
View File

@@ -96,6 +96,15 @@ vi.mock('./useLogger.js', () => ({
}),
}));
const mockStartNewTurn = vi.fn();
const mockAddUsage = vi.fn();
vi.mock('../contexts/SessionContext.js', () => ({
useSessionStats: vi.fn(() => ({
startNewTurn: mockStartNewTurn,
addUsage: mockAddUsage,
})),
}));
vi.mock('./slashCommandProcessor.js', () => ({
handleSlashCommand: vi.fn().mockReturnValue(false),
}));
@@ -531,4 +540,63 @@ describe('useGeminiStream', () => {
});
});
});
describe('Session Stats Integration', () => {
it('should call startNewTurn and addUsage for a simple prompt', async () => {
const mockMetadata = { totalTokenCount: 123 };
const mockStream = (async function* () {
yield { type: 'content', value: 'Response' };
yield { type: 'usage_metadata', value: mockMetadata };
})();
mockSendMessageStream.mockReturnValue(mockStream);
const { result } = renderTestHook();
await act(async () => {
await result.current.submitQuery('Hello, world!');
});
expect(mockStartNewTurn).toHaveBeenCalledTimes(1);
expect(mockAddUsage).toHaveBeenCalledTimes(1);
expect(mockAddUsage).toHaveBeenCalledWith(mockMetadata);
});
it('should only call addUsage for a tool continuation prompt', async () => {
const mockMetadata = { totalTokenCount: 456 };
const mockStream = (async function* () {
yield { type: 'content', value: 'Final Answer' };
yield { type: 'usage_metadata', value: mockMetadata };
})();
mockSendMessageStream.mockReturnValue(mockStream);
const { result } = renderTestHook();
await act(async () => {
await result.current.submitQuery([{ text: 'tool response' }], {
isContinuation: true,
});
});
expect(mockStartNewTurn).not.toHaveBeenCalled();
expect(mockAddUsage).toHaveBeenCalledTimes(1);
expect(mockAddUsage).toHaveBeenCalledWith(mockMetadata);
});
it('should not call addUsage if the stream contains no usage metadata', async () => {
// Arrange: A stream that yields content but never a usage_metadata event
const mockStream = (async function* () {
yield { type: 'content', value: 'Some response text' };
})();
mockSendMessageStream.mockReturnValue(mockStream);
const { result } = renderTestHook();
await act(async () => {
await result.current.submitQuery('Query with no usage data');
});
expect(mockStartNewTurn).toHaveBeenCalledTimes(1);
expect(mockAddUsage).not.toHaveBeenCalled();
});
});
});

17
packages/cli/src/ui/hooks/useGeminiStream.ts
View File

@@ -42,6 +42,7 @@ import {
TrackedCompletedToolCall,
TrackedCancelledToolCall,
} from './useReactToolScheduler.js';
import { useSessionStats } from '../contexts/SessionContext.js';
export function mergePartListUnions(list: PartListUnion[]): PartListUnion {
const resultParts: PartListUnion = [];
@@ -82,6 +83,7 @@ export const useGeminiStream = (
const [pendingHistoryItemRef, setPendingHistoryItem] =
useStateAndRef<HistoryItemWithoutId | null>(null);
const logger = useLogger();
const { startNewTurn, addUsage } = useSessionStats();
const [toolCalls, scheduleToolCalls, markToolsAsSubmitted] =
useReactToolScheduler(
@@ -390,6 +392,9 @@ export const useGeminiStream = (
case ServerGeminiEventType.ChatCompressed:
handleChatCompressionEvent();
break;
case ServerGeminiEventType.UsageMetadata:
addUsage(event.value);
break;
case ServerGeminiEventType.ToolCallConfirmation:
case ServerGeminiEventType.ToolCallResponse:
// do nothing
@@ -412,11 +417,12 @@ export const useGeminiStream = (
handleErrorEvent,
scheduleToolCalls,
handleChatCompressionEvent,
addUsage,
],
);
const submitQuery = useCallback(
async (query: PartListUnion) => {
async (query: PartListUnion, options?: { isContinuation: boolean }) => {
if (
streamingState === StreamingState.Responding ||
streamingState === StreamingState.WaitingForConfirmation
@@ -426,6 +432,10 @@ export const useGeminiStream = (
const userMessageTimestamp = Date.now();
setShowHelp(false);
if (!options?.isContinuation) {
startNewTurn();
}
abortControllerRef.current = new AbortController();
const abortSignal = abortControllerRef.current.signal;
@@ -491,6 +501,7 @@ export const useGeminiStream = (
setPendingHistoryItem,
setInitError,
geminiClient,
startNewTurn,
],
);
@@ -576,7 +587,9 @@ export const useGeminiStream = (
);
markToolsAsSubmitted(callIdsToMarkAsSubmitted);
submitQuery(mergePartListUnions(responsesToSend));
submitQuery(mergePartListUnions(responsesToSend), {
isContinuation: true,
});
}
}, [
toolCalls,