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.ts36
1 files changed, 29 insertions, 7 deletions
diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts
index 234652db..fcfa1c57 100644
--- a/packages/cli/src/ui/hooks/useGeminiStream.ts
+++ b/packages/cli/src/ui/hooks/useGeminiStream.ts
@@ -92,6 +92,7 @@ export const useGeminiStream = (
) => {
const [initError, setInitError] = useState<string | null>(null);
const abortControllerRef = useRef<AbortController | null>(null);
+ const turnCancelledRef = useRef(false);
const [isResponding, setIsResponding] = useState<boolean>(false);
const [thought, setThought] = useState<ThoughtSummary | null>(null);
const [pendingHistoryItemRef, setPendingHistoryItem] =
@@ -168,15 +169,25 @@ export const useGeminiStream = (
return StreamingState.Idle;
}, [isResponding, toolCalls]);
- useEffect(() => {
- if (streamingState === StreamingState.Idle) {
- abortControllerRef.current = null;
- }
- }, [streamingState]);
-
useInput((_input, key) => {
- if (streamingState !== StreamingState.Idle && key.escape) {
+ if (streamingState === StreamingState.Responding && key.escape) {
+ if (turnCancelledRef.current) {
+ return;
+ }
+ turnCancelledRef.current = true;
abortControllerRef.current?.abort();
+ if (pendingHistoryItemRef.current) {
+ addItem(pendingHistoryItemRef.current, Date.now());
+ }
+ addItem(
+ {
+ type: MessageType.INFO,
+ text: 'Request cancelled.',
+ },
+ Date.now(),
+ );
+ setPendingHistoryItem(null);
+ setIsResponding(false);
}
});
@@ -189,6 +200,9 @@ export const useGeminiStream = (
queryToSend: PartListUnion | null;
shouldProceed: boolean;
}> => {
+ if (turnCancelledRef.current) {
+ return { queryToSend: null, shouldProceed: false };
+ }
if (typeof query === 'string' && query.trim().length === 0) {
return { queryToSend: null, shouldProceed: false };
}
@@ -285,6 +299,10 @@ export const useGeminiStream = (
currentGeminiMessageBuffer: string,
userMessageTimestamp: number,
): string => {
+ if (turnCancelledRef.current) {
+ // Prevents additional output after a user initiated cancel.
+ return '';
+ }
let newGeminiMessageBuffer = currentGeminiMessageBuffer + eventValue;
if (
pendingHistoryItemRef.current?.type !== 'gemini' &&
@@ -335,6 +353,9 @@ export const useGeminiStream = (
const handleUserCancelledEvent = useCallback(
(userMessageTimestamp: number) => {
+ if (turnCancelledRef.current) {
+ return;
+ }
if (pendingHistoryItemRef.current) {
if (pendingHistoryItemRef.current.type === 'tool_group') {
const updatedTools = pendingHistoryItemRef.current.tools.map(
@@ -469,6 +490,7 @@ export const useGeminiStream = (
abortControllerRef.current = new AbortController();
const abortSignal = abortControllerRef.current.signal;
+ turnCancelledRef.current = false;
const { queryToSend, shouldProceed } = await prepareQueryForGemini(
query,