mirror of
https://github.com/QwenLM/qwen-code.git
synced 2025-12-20 16:57:46 +00:00
feat: Display initial token usage metrics in /stats (#879)
This commit is contained in:
@@ -48,7 +48,7 @@ import {
|
||||
} from '@gemini-cli/core';
|
||||
import { useLogger } from './hooks/useLogger.js';
|
||||
import { StreamingContext } from './contexts/StreamingContext.js';
|
||||
import { SessionProvider } from './contexts/SessionContext.js';
|
||||
import { SessionStatsProvider } from './contexts/SessionContext.js';
|
||||
import { useGitBranchName } from './hooks/useGitBranchName.js';
|
||||
|
||||
const CTRL_C_PROMPT_DURATION_MS = 1000;
|
||||
@@ -60,9 +60,9 @@ interface AppProps {
|
||||
}
|
||||
|
||||
export const AppWrapper = (props: AppProps) => (
|
||||
<SessionProvider>
|
||||
<SessionStatsProvider>
|
||||
<App {...props} />
|
||||
</SessionProvider>
|
||||
</SessionStatsProvider>
|
||||
);
|
||||
|
||||
const App = ({ config, settings, startupWarnings = [] }: AppProps) => {
|
||||
|
||||
@@ -4,26 +4,181 @@
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { type MutableRefObject } from 'react';
|
||||
import { render } from 'ink-testing-library';
|
||||
import { Text } from 'ink';
|
||||
import { SessionProvider, useSession } from './SessionContext.js';
|
||||
import { describe, it, expect } from 'vitest';
|
||||
import { act } from 'react-dom/test-utils';
|
||||
import { SessionStatsProvider, useSessionStats } from './SessionContext.js';
|
||||
import { describe, it, expect, vi } from 'vitest';
|
||||
import { GenerateContentResponseUsageMetadata } from '@google/genai';
|
||||
|
||||
const TestComponent = () => {
|
||||
const { startTime } = useSession();
|
||||
return <Text>{startTime.toISOString()}</Text>;
|
||||
// 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,
|
||||
};
|
||||
|
||||
describe('SessionContext', () => {
|
||||
it('should provide a start time', () => {
|
||||
const { lastFrame } = render(
|
||||
<SessionProvider>
|
||||
<TestComponent />
|
||||
</SessionProvider>,
|
||||
const mockMetadata2: GenerateContentResponseUsageMetadata = {
|
||||
promptTokenCount: 10,
|
||||
candidatesTokenCount: 20,
|
||||
totalTokenCount: 30,
|
||||
cachedContentTokenCount: 5,
|
||||
toolUsePromptTokenCount: 1,
|
||||
thoughtsTokenCount: 2,
|
||||
};
|
||||
|
||||
/**
|
||||
* A test harness component that uses the hook and exposes the context value
|
||||
* via a mutable ref. This allows us to interact with the context's functions
|
||||
* and assert against its state directly in our tests.
|
||||
*/
|
||||
const TestHarness = ({
|
||||
contextRef,
|
||||
}: {
|
||||
contextRef: MutableRefObject<ReturnType<typeof useSessionStats> | undefined>;
|
||||
}) => {
|
||||
contextRef.current = useSessionStats();
|
||||
return null;
|
||||
};
|
||||
|
||||
describe('SessionStatsContext', () => {
|
||||
it('should provide the correct initial state', () => {
|
||||
const contextRef: MutableRefObject<
|
||||
ReturnType<typeof useSessionStats> | undefined
|
||||
> = { current: undefined };
|
||||
|
||||
render(
|
||||
<SessionStatsProvider>
|
||||
<TestHarness contextRef={contextRef} />
|
||||
</SessionStatsProvider>,
|
||||
);
|
||||
|
||||
const frameText = lastFrame();
|
||||
// Check if the output is a valid ISO string, which confirms it's a Date object.
|
||||
expect(new Date(frameText!).toString()).not.toBe('Invalid Date');
|
||||
const stats = contextRef.current?.stats;
|
||||
|
||||
expect(stats?.sessionStartTime).toBeInstanceOf(Date);
|
||||
expect(stats?.lastTurn).toBeNull();
|
||||
expect(stats?.cumulative.turnCount).toBe(0);
|
||||
expect(stats?.cumulative.totalTokenCount).toBe(0);
|
||||
expect(stats?.cumulative.promptTokenCount).toBe(0);
|
||||
});
|
||||
|
||||
it('should increment turnCount when startNewTurn is called', () => {
|
||||
const contextRef: MutableRefObject<
|
||||
ReturnType<typeof useSessionStats> | undefined
|
||||
> = { current: undefined };
|
||||
|
||||
render(
|
||||
<SessionStatsProvider>
|
||||
<TestHarness contextRef={contextRef} />
|
||||
</SessionStatsProvider>,
|
||||
);
|
||||
|
||||
act(() => {
|
||||
contextRef.current?.startNewTurn();
|
||||
});
|
||||
|
||||
const stats = contextRef.current?.stats;
|
||||
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);
|
||||
});
|
||||
|
||||
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,
|
||||
);
|
||||
|
||||
// Check that turn count is NOT incremented
|
||||
expect(stats?.cumulative.turnCount).toBe(0);
|
||||
|
||||
// Check that lastTurn is updated
|
||||
expect(stats?.lastTurn?.metadata).toEqual(mockMetadata1);
|
||||
});
|
||||
|
||||
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);
|
||||
});
|
||||
|
||||
// 3. Second API call (e.g., sending tool response back)
|
||||
act(() => {
|
||||
contextRef.current?.addUsage(mockMetadata2);
|
||||
});
|
||||
|
||||
const stats = contextRef.current?.stats;
|
||||
|
||||
// Turn count should only be 1
|
||||
expect(stats?.cumulative.turnCount).toBe(1);
|
||||
|
||||
// These fields should be the SUM of both calls
|
||||
expect(stats?.cumulative.totalTokenCount).toBe(330); // 300 + 30
|
||||
expect(stats?.cumulative.candidatesTokenCount).toBe(220); // 200 + 20
|
||||
expect(stats?.cumulative.thoughtsTokenCount).toBe(22); // 20 + 2
|
||||
|
||||
// These fields should ONLY be from the FIRST call, because isNewTurnForAggregation was true
|
||||
expect(stats?.cumulative.promptTokenCount).toBe(100);
|
||||
expect(stats?.cumulative.cachedContentTokenCount).toBe(50);
|
||||
expect(stats?.cumulative.toolUsePromptTokenCount).toBe(10);
|
||||
|
||||
// Last turn should hold the metadata from the most recent call
|
||||
expect(stats?.lastTurn?.metadata).toEqual(mockMetadata2);
|
||||
});
|
||||
|
||||
it('should throw an error when useSessionStats is used outside of a provider', () => {
|
||||
// Suppress the expected console error during this test.
|
||||
const errorSpy = vi.spyOn(console, 'error').mockImplementation(() => {});
|
||||
|
||||
const contextRef = { current: undefined };
|
||||
|
||||
// We expect rendering to fail, which React will catch and log as an error.
|
||||
render(<TestHarness contextRef={contextRef} />);
|
||||
|
||||
// Assert that the first argument of the first call to console.error
|
||||
// contains the expected message. This is more robust than checking
|
||||
// the exact arguments, which can be affected by React/JSDOM internals.
|
||||
expect(errorSpy.mock.calls[0][0]).toContain(
|
||||
'useSessionStats must be used within a SessionStatsProvider',
|
||||
);
|
||||
|
||||
errorSpy.mockRestore();
|
||||
});
|
||||
});
|
||||
|
||||
@@ -4,35 +4,140 @@
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import React, { createContext, useContext, useState, useMemo } from 'react';
|
||||
import React, {
|
||||
createContext,
|
||||
useContext,
|
||||
useState,
|
||||
useMemo,
|
||||
useCallback,
|
||||
} from 'react';
|
||||
|
||||
interface SessionContextType {
|
||||
startTime: Date;
|
||||
import { type GenerateContentResponseUsageMetadata } from '@google/genai';
|
||||
|
||||
// --- Interface Definitions ---
|
||||
|
||||
interface CumulativeStats {
|
||||
turnCount: number;
|
||||
promptTokenCount: number;
|
||||
candidatesTokenCount: number;
|
||||
totalTokenCount: number;
|
||||
cachedContentTokenCount: number;
|
||||
toolUsePromptTokenCount: number;
|
||||
thoughtsTokenCount: number;
|
||||
}
|
||||
|
||||
const SessionContext = createContext<SessionContextType | null>(null);
|
||||
interface LastTurnStats {
|
||||
metadata: GenerateContentResponseUsageMetadata;
|
||||
// TODO(abhipatel12): Add apiTime, etc. here in a future step.
|
||||
}
|
||||
|
||||
export const SessionProvider: React.FC<{ children: React.ReactNode }> = ({
|
||||
interface SessionStatsState {
|
||||
sessionStartTime: Date;
|
||||
cumulative: CumulativeStats;
|
||||
lastTurn: LastTurnStats | null;
|
||||
isNewTurnForAggregation: boolean;
|
||||
}
|
||||
|
||||
// Defines the final "value" of our context, including the state
|
||||
// and the functions to update it.
|
||||
interface SessionStatsContextValue {
|
||||
stats: SessionStatsState;
|
||||
startNewTurn: () => void;
|
||||
addUsage: (metadata: GenerateContentResponseUsageMetadata) => void;
|
||||
}
|
||||
|
||||
// --- Context Definition ---
|
||||
|
||||
const SessionStatsContext = createContext<SessionStatsContextValue | undefined>(
|
||||
undefined,
|
||||
);
|
||||
|
||||
// --- Provider Component ---
|
||||
|
||||
export const SessionStatsProvider: React.FC<{ children: React.ReactNode }> = ({
|
||||
children,
|
||||
}) => {
|
||||
const [startTime] = useState(new Date());
|
||||
const [stats, setStats] = useState<SessionStatsState>({
|
||||
sessionStartTime: new Date(),
|
||||
cumulative: {
|
||||
turnCount: 0,
|
||||
promptTokenCount: 0,
|
||||
candidatesTokenCount: 0,
|
||||
totalTokenCount: 0,
|
||||
cachedContentTokenCount: 0,
|
||||
toolUsePromptTokenCount: 0,
|
||||
thoughtsTokenCount: 0,
|
||||
},
|
||||
lastTurn: null,
|
||||
isNewTurnForAggregation: true,
|
||||
});
|
||||
|
||||
// A single, internal worker function to handle all metadata aggregation.
|
||||
const aggregateTokens = useCallback(
|
||||
(metadata: GenerateContentResponseUsageMetadata) => {
|
||||
setStats((prevState) => {
|
||||
const { isNewTurnForAggregation } = prevState;
|
||||
const newCumulative = { ...prevState.cumulative };
|
||||
|
||||
newCumulative.candidatesTokenCount +=
|
||||
metadata.candidatesTokenCount ?? 0;
|
||||
newCumulative.thoughtsTokenCount += metadata.thoughtsTokenCount ?? 0;
|
||||
newCumulative.totalTokenCount += metadata.totalTokenCount ?? 0;
|
||||
|
||||
if (isNewTurnForAggregation) {
|
||||
newCumulative.promptTokenCount += metadata.promptTokenCount ?? 0;
|
||||
newCumulative.cachedContentTokenCount +=
|
||||
metadata.cachedContentTokenCount ?? 0;
|
||||
newCumulative.toolUsePromptTokenCount +=
|
||||
metadata.toolUsePromptTokenCount ?? 0;
|
||||
}
|
||||
|
||||
return {
|
||||
...prevState,
|
||||
cumulative: newCumulative,
|
||||
lastTurn: { metadata },
|
||||
isNewTurnForAggregation: false,
|
||||
};
|
||||
});
|
||||
},
|
||||
[],
|
||||
);
|
||||
|
||||
const startNewTurn = useCallback(() => {
|
||||
setStats((prevState) => ({
|
||||
...prevState,
|
||||
cumulative: {
|
||||
...prevState.cumulative,
|
||||
turnCount: prevState.cumulative.turnCount + 1,
|
||||
},
|
||||
isNewTurnForAggregation: true,
|
||||
}));
|
||||
}, []);
|
||||
|
||||
const value = useMemo(
|
||||
() => ({
|
||||
startTime,
|
||||
stats,
|
||||
startNewTurn,
|
||||
addUsage: aggregateTokens,
|
||||
}),
|
||||
[startTime],
|
||||
[stats, startNewTurn, aggregateTokens],
|
||||
);
|
||||
|
||||
return (
|
||||
<SessionContext.Provider value={value}>{children}</SessionContext.Provider>
|
||||
<SessionStatsContext.Provider value={value}>
|
||||
{children}
|
||||
</SessionStatsContext.Provider>
|
||||
);
|
||||
};
|
||||
|
||||
export const useSession = () => {
|
||||
const context = useContext(SessionContext);
|
||||
if (!context) {
|
||||
throw new Error('useSession must be used within a SessionProvider');
|
||||
// --- Consumer Hook ---
|
||||
|
||||
export const useSessionStats = () => {
|
||||
const context = useContext(SessionStatsContext);
|
||||
if (context === undefined) {
|
||||
throw new Error(
|
||||
'useSessionStats must be used within a SessionStatsProvider',
|
||||
);
|
||||
}
|
||||
return context;
|
||||
};
|
||||
|
||||
@@ -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();
|
||||
});
|
||||
});
|
||||
|
||||
@@ -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,
|
||||
],
|
||||
);
|
||||
|
||||
|
||||
@@ -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();
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user