summaryrefslogtreecommitdiff
path: root/packages/cli/src/ui/hooks/useGeminiStream.ts
diff options
context:
space:
mode:
Diffstat (limited to 'packages/cli/src/ui/hooks/useGeminiStream.ts')
-rw-r--r--packages/cli/src/ui/hooks/useGeminiStream.ts17
1 files changed, 15 insertions, 2 deletions
diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts
index 2b47ae6f..bad9f78a 100644
--- a/packages/cli/src/ui/hooks/useGeminiStream.ts
+++ b/packages/cli/src/ui/hooks/useGeminiStream.ts
@@ -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,