diff options
Diffstat (limited to 'packages/cli/src/ui/contexts')
| -rw-r--r-- | packages/cli/src/ui/contexts/SessionContext.test.tsx | 45 | ||||
| -rw-r--r-- | packages/cli/src/ui/contexts/SessionContext.tsx | 33 |
2 files changed, 78 insertions, 0 deletions
diff --git a/packages/cli/src/ui/contexts/SessionContext.test.tsx b/packages/cli/src/ui/contexts/SessionContext.test.tsx index b00a5d75..e9fc33e6 100644 --- a/packages/cli/src/ui/contexts/SessionContext.test.tsx +++ b/packages/cli/src/ui/contexts/SessionContext.test.tsx @@ -177,6 +177,51 @@ describe('SessionStatsContext', () => { 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); + }); + 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(() => {}); diff --git a/packages/cli/src/ui/contexts/SessionContext.tsx b/packages/cli/src/ui/contexts/SessionContext.tsx index 0d574e75..f59e17e1 100644 --- a/packages/cli/src/ui/contexts/SessionContext.tsx +++ b/packages/cli/src/ui/contexts/SessionContext.tsx @@ -31,6 +31,7 @@ interface SessionStatsState { sessionStartTime: Date; cumulative: CumulativeStats; currentTurn: CumulativeStats; + currentResponse: CumulativeStats; } // Defines the final "value" of our context, including the state @@ -97,6 +98,16 @@ export const SessionStatsProvider: React.FC<{ children: React.ReactNode }> = ({ thoughtsTokenCount: 0, apiTimeMs: 0, }, + currentResponse: { + turnCount: 0, + promptTokenCount: 0, + candidatesTokenCount: 0, + totalTokenCount: 0, + cachedContentTokenCount: 0, + toolUsePromptTokenCount: 0, + thoughtsTokenCount: 0, + apiTimeMs: 0, + }, }); // A single, internal worker function to handle all metadata aggregation. @@ -107,15 +118,27 @@ export const SessionStatsProvider: React.FC<{ children: React.ReactNode }> = ({ setStats((prevState) => { const newCumulative = { ...prevState.cumulative }; const newCurrentTurn = { ...prevState.currentTurn }; + const newCurrentResponse = { + turnCount: 0, + promptTokenCount: 0, + candidatesTokenCount: 0, + totalTokenCount: 0, + cachedContentTokenCount: 0, + toolUsePromptTokenCount: 0, + thoughtsTokenCount: 0, + apiTimeMs: 0, + }; // Add all tokens to the current turn's stats as well as cumulative stats. addTokens(newCurrentTurn, metadata); addTokens(newCumulative, metadata); + addTokens(newCurrentResponse, metadata); return { ...prevState, cumulative: newCumulative, currentTurn: newCurrentTurn, + currentResponse: newCurrentResponse, }; }); }, @@ -139,6 +162,16 @@ export const SessionStatsProvider: React.FC<{ children: React.ReactNode }> = ({ thoughtsTokenCount: 0, apiTimeMs: 0, }, + currentResponse: { + turnCount: 0, + promptTokenCount: 0, + candidatesTokenCount: 0, + totalTokenCount: 0, + cachedContentTokenCount: 0, + toolUsePromptTokenCount: 0, + thoughtsTokenCount: 0, + apiTimeMs: 0, + }, })); }, []); |
