diff options
Diffstat (limited to 'packages/cli/src/ui/hooks/useGeminiStream.ts')
| -rw-r--r-- | packages/cli/src/ui/hooks/useGeminiStream.ts | 17 |
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, |
