summaryrefslogtreecommitdiff
path: root/packages/cli/src/ui/contexts
diff options
context:
space:
mode:
Diffstat (limited to 'packages/cli/src/ui/contexts')
-rw-r--r--packages/cli/src/ui/contexts/SessionContext.test.tsx45
-rw-r--r--packages/cli/src/ui/contexts/SessionContext.tsx33
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,
+ },
}));
}, []);