Sync upstream Gemini-CLI v0.8.2 (#838)

2025-12-20 08:47:44 +00:00 · 2025-10-23 09:27:04 +08:00
parent 096fabb5d6
commit eb95c131be
644 changed files with 70389 additions and 23709 deletions
--- a/packages/cli/src/ui/hooks/useGeminiStream.ts
+++ b/packages/cli/src/ui/hooks/useGeminiStream.ts
@@ -7,15 +7,15 @@
 import { useState, useRef, useCallback, useEffect, useMemo } from 'react';
 import type {
  Config,
-  GeminiClient,
-  ServerGeminiStreamEvent as GeminiEvent,
-  ServerGeminiContentEvent as ContentEvent,
-  ServerGeminiErrorEvent as ErrorEvent,
-  ServerGeminiChatCompressedEvent,
-  ServerGeminiFinishedEvent,
-  ToolCallRequestInfo,
  EditorType,
+  GeminiClient,
+  ServerGeminiChatCompressedEvent,
+  ServerGeminiContentEvent as ContentEvent,
+  ServerGeminiFinishedEvent,
+  ServerGeminiStreamEvent as GeminiEvent,
  ThoughtSummary,
+  ToolCallRequestInfo,
+  GeminiErrorEventValue,
 } from '@qwen-code/qwen-code-core';
 import {
  GeminiEventType as ServerGeminiEventType,
@@ -31,6 +31,8 @@ import {
  ConversationFinishedEvent,
  ApprovalMode,
  parseAndFormatApiError,
+  promptIdContext,
+  ToolConfirmationOutcome,
  logApiCancel,
  ApiCancelEvent,
 } from '@qwen-code/qwen-code-core';
@@ -50,19 +52,19 @@ import { findLastSafeSplitPoint } from '../utils/markdownUtilities.js';
 import { useStateAndRef } from './useStateAndRef.js';
 import type { UseHistoryManagerReturn } from './useHistoryManager.js';
 import { useLogger } from './useLogger.js';
-import type {
-  TrackedToolCall,
-  TrackedCompletedToolCall,
-  TrackedCancelledToolCall,
-} from './useReactToolScheduler.js';
-import { promises as fs } from 'node:fs';
-import path from 'node:path';
 import {
  useReactToolScheduler,
  mapToDisplay as mapTrackedToolCallsToDisplay,
+  type TrackedToolCall,
+  type TrackedCompletedToolCall,
+  type TrackedCancelledToolCall,
+  type TrackedWaitingToolCall,
 } from './useReactToolScheduler.js';
+import { promises as fs } from 'node:fs';
+import path from 'node:path';
 import { useSessionStats } from '../contexts/SessionContext.js';
 import { useKeypress } from './useKeypress.js';
+import type { LoadedSettings } from '../../config/settings.js';

 enum StreamProcessingStatus {
  Completed,
@@ -70,6 +72,16 @@ enum StreamProcessingStatus {
  Error,
 }

+const EDIT_TOOL_NAMES = new Set(['replace', 'write_file']);
+
+function showCitations(settings: LoadedSettings): boolean {
+  const enabled = settings?.merged?.ui?.showCitations;
+  if (enabled !== undefined) {
+    return enabled;
+  }
+  return true;
+}
+
 /**
 * Manages the Gemini stream, including user input, command processing,
 * API interaction, and tool call lifecycle.
@@ -79,24 +91,29 @@ export const useGeminiStream = (
  history: HistoryItem[],
  addItem: UseHistoryManagerReturn['addItem'],
  config: Config,
+  settings: LoadedSettings,
  onDebugMessage: (message: string) => void,
  handleSlashCommand: (
    cmd: PartListUnion,
  ) => Promise<SlashCommandProcessorResult | false>,
  shellModeActive: boolean,
  getPreferredEditor: () => EditorType | undefined,
-  onAuthError: () => void,
+  onAuthError: (error: string) => void,
  performMemoryRefresh: () => Promise<void>,
  modelSwitchedFromQuotaError: boolean,
  setModelSwitchedFromQuotaError: React.Dispatch<React.SetStateAction<boolean>>,
  onEditorClose: () => void,
  onCancelSubmit: () => void,
  visionModelPreviewEnabled: boolean,
+  setShellInputFocused: (value: boolean) => void,
+  terminalWidth: number,
+  terminalHeight: number,
  onVisionSwitchRequired?: (query: PartListUnion) => Promise<{
    modelOverride?: string;
    persistSessionModel?: string;
    showGuidance?: boolean;
  }>,
+  isShellFocused?: boolean,
 ) => {
  const [initError, setInitError] = useState<string | null>(null);
  const abortControllerRef = useRef<AbortController | null>(null);
@@ -104,7 +121,7 @@ export const useGeminiStream = (
  const isSubmittingQueryRef = useRef(false);
  const [isResponding, setIsResponding] = useState<boolean>(false);
  const [thought, setThought] = useState<ThoughtSummary | null>(null);
-  const [pendingHistoryItemRef, setPendingHistoryItem] =
+  const [pendingHistoryItem, pendingHistoryItemRef, setPendingHistoryItem] =
    useStateAndRef<HistoryItemWithoutId | null>(null);
  const processedMemoryToolsRef = useRef<Set<string>>(new Set());
  const { startNewPrompt, getPromptCount } = useSessionStats();
@@ -137,7 +154,6 @@ export const useGeminiStream = (
        }
      },
      config,
-      setPendingHistoryItem,
      getPreferredEditor,
      onEditorClose,
    );
@@ -148,20 +164,40 @@ export const useGeminiStream = (
    [toolCalls],
  );

+  const activeToolPtyId = useMemo(() => {
+    const executingShellTool = toolCalls?.find(
+      (tc) =>
+        tc.status === 'executing' && tc.request.name === 'run_shell_command',
+    );
+    if (executingShellTool) {
+      return (executingShellTool as { pid?: number }).pid;
+    }
+    return undefined;
+  }, [toolCalls]);
+
  const loopDetectedRef = useRef(false);
+  const [
+    loopDetectionConfirmationRequest,
+    setLoopDetectionConfirmationRequest,
+  ] = useState<{
+    onComplete: (result: { userSelection: 'disable' | 'keep' }) => void;
+  } | null>(null);

  const onExec = useCallback(async (done: Promise<void>) => {
    setIsResponding(true);
    await done;
    setIsResponding(false);
  }, []);
-  const { handleShellCommand } = useShellCommandProcessor(
+  const { handleShellCommand, activeShellPtyId } = useShellCommandProcessor(
    addItem,
    setPendingHistoryItem,
    onExec,
    onDebugMessage,
    config,
    geminiClient,
+    setShellInputFocused,
+    terminalWidth,
+    terminalHeight,
  );

  const { handleVisionSwitch, restoreOriginalModel } = useVisionAutoSwitch(
@@ -170,6 +206,13 @@ export const useGeminiStream = (
    visionModelPreviewEnabled,
    onVisionSwitchRequired,
  );
+  const activePtyId = activeShellPtyId || activeToolPtyId;
+
+  useEffect(() => {
+    if (!activePtyId) {
+      setShellInputFocused(false);
+    }
+  }, [activePtyId, setShellInputFocused]);

  const streamingState = useMemo(() => {
    if (toolCalls.some((tc) => tc.status === 'awaiting_approval')) {
@@ -248,19 +291,21 @@ export const useGeminiStream = (
    setPendingHistoryItem(null);
    onCancelSubmit();
    setIsResponding(false);
+    setShellInputFocused(false);
  }, [
    streamingState,
    addItem,
    setPendingHistoryItem,
    onCancelSubmit,
    pendingHistoryItemRef,
+    setShellInputFocused,
    config,
    getPromptCount,
  ]);

  useKeypress(
    (key) => {
-      if (key.name === 'escape') {
+      if (key.name === 'escape' && !isShellFocused) {
        cancelOngoingRequest();
      }
    },
@@ -288,15 +333,6 @@ export const useGeminiStream = (

      if (typeof query === 'string') {
        const trimmedQuery = query.trim();
-        logUserPrompt(
-          config,
-          new UserPromptEvent(
-            trimmedQuery.length,
-            prompt_id,
-            config.getContentGeneratorConfig()?.authType,
-            trimmedQuery,
-          ),
-        );
        onDebugMessage(`User query: '${trimmedQuery}'`);
        await logger?.logMessage(MessageSenderType.USER, trimmedQuery);

@@ -494,7 +530,7 @@ export const useGeminiStream = (
  );

  const handleErrorEvent = useCallback(
-    (eventValue: ErrorEvent['value'], userMessageTimestamp: number) => {
+    (eventValue: GeminiErrorEventValue, userMessageTimestamp: number) => {
      if (pendingHistoryItemRef.current) {
        addItem(pendingHistoryItemRef.current, userMessageTimestamp);
        setPendingHistoryItem(null);
@@ -517,9 +553,27 @@ export const useGeminiStream = (
    [addItem, pendingHistoryItemRef, setPendingHistoryItem, config, setThought],
  );

+  const handleCitationEvent = useCallback(
+    (text: string, userMessageTimestamp: number) => {
+      if (!showCitations(settings)) {
+        return;
+      }
+
+      if (pendingHistoryItemRef.current) {
+        addItem(pendingHistoryItemRef.current, userMessageTimestamp);
+        setPendingHistoryItem(null);
+      }
+      addItem({ type: MessageType.INFO, text }, userMessageTimestamp);
+    },
+    [addItem, pendingHistoryItemRef, setPendingHistoryItem, settings],
+  );
+
  const handleFinishedEvent = useCallback(
    (event: ServerGeminiFinishedEvent, userMessageTimestamp: number) => {
-      const finishReason = event.value;
+      const finishReason = event.value.reason;
+      if (!finishReason) {
+        return;
+      }

      const finishReasonMessages: Record<FinishReason, string | undefined> = {
        [FinishReason.FINISH_REASON_UNSPECIFIED]: undefined,
@@ -558,8 +612,15 @@ export const useGeminiStream = (
  );

  const handleChatCompressionEvent = useCallback(
-    (eventValue: ServerGeminiChatCompressedEvent['value']) =>
-      addItem(
+    (
+      eventValue: ServerGeminiChatCompressedEvent['value'],
+      userMessageTimestamp: number,
+    ) => {
+      if (pendingHistoryItemRef.current) {
+        addItem(pendingHistoryItemRef.current, userMessageTimestamp);
+        setPendingHistoryItem(null);
+      }
+      return addItem(
        {
          type: 'info',
          text:
@@ -569,8 +630,9 @@ export const useGeminiStream = (
            `${eventValue?.newTokenCount ?? 'unknown'} tokens).`,
        },
        Date.now(),
-      ),
-    [addItem, config],
+      );
+    },
+    [addItem, config, pendingHistoryItemRef, setPendingHistoryItem],
  );

  const handleMaxSessionTurnsEvent = useCallback(
@@ -604,15 +666,38 @@ export const useGeminiStream = (
    [addItem],
  );

+  const handleLoopDetectionConfirmation = useCallback(
+    (result: { userSelection: 'disable' | 'keep' }) => {
+      setLoopDetectionConfirmationRequest(null);
+
+      if (result.userSelection === 'disable') {
+        config.getGeminiClient().getLoopDetectionService().disableForSession();
+        addItem(
+          {
+            type: 'info',
+            text: `Loop detection has been disabled for this session. Please try your request again.`,
+          },
+          Date.now(),
+        );
+      } else {
+        addItem(
+          {
+            type: 'info',
+            text: `A potential loop was detected. This can happen due to repetitive tool calls or other model behavior. The request has been halted.`,
+          },
+          Date.now(),
+        );
+      }
+    },
+    [config, addItem],
+  );
+
  const handleLoopDetectedEvent = useCallback(() => {
-    addItem(
-      {
-        type: 'info',
-        text: `A potential loop was detected. This can happen due to repetitive tool calls or other model behavior. The request has been halted.`,
-      },
-      Date.now(),
-    );
-  }, [addItem]);
+    // Show the confirmation dialog to choose whether to disable loop detection
+    setLoopDetectionConfirmationRequest({
+      onComplete: handleLoopDetectionConfirmation,
+    });
+  }, [handleLoopDetectionConfirmation]);

  const processGeminiStreamEvents = useCallback(
    async (
@@ -644,7 +729,7 @@ export const useGeminiStream = (
            handleErrorEvent(event.value, userMessageTimestamp);
            break;
          case ServerGeminiEventType.ChatCompressed:
-            handleChatCompressionEvent(event.value);
+            handleChatCompressionEvent(event.value, userMessageTimestamp);
            break;
          case ServerGeminiEventType.ToolCallConfirmation:
          case ServerGeminiEventType.ToolCallResponse:
@@ -662,6 +747,9 @@ export const useGeminiStream = (
              userMessageTimestamp,
            );
            break;
+          case ServerGeminiEventType.Citation:
+            handleCitationEvent(event.value, userMessageTimestamp);
+            break;
          case ServerGeminiEventType.LoopDetected:
            // handle later because we want to move pending history to history
            // before we add loop detected message to history
@@ -691,6 +779,7 @@ export const useGeminiStream = (
      handleFinishedEvent,
      handleMaxSessionTurnsEvent,
      handleSessionTokenLimitExceededEvent,
+      handleCitationEvent,
    ],
  );

@@ -732,101 +821,116 @@ export const useGeminiStream = (
        prompt_id = config.getSessionId() + '########' + getPromptCount();
      }

-      const { queryToSend, shouldProceed } = await prepareQueryForGemini(
-        query,
-        userMessageTimestamp,
-        abortSignal,
-        prompt_id!,
-      );
-
-      if (!shouldProceed || queryToSend === null) {
-        isSubmittingQueryRef.current = false;
-        return;
-      }
-
-      // Handle vision switch requirement
-      const visionSwitchResult = await handleVisionSwitch(
-        queryToSend,
-        userMessageTimestamp,
-        options?.isContinuation || false,
-      );
-
-      if (!visionSwitchResult.shouldProceed) {
-        isSubmittingQueryRef.current = false;
-        return;
-      }
-
-      const finalQueryToSend = queryToSend;
-
-      if (!options?.isContinuation) {
-        startNewPrompt();
-        setThought(null); // Reset thought when starting a new prompt
-      }
-
-      setIsResponding(true);
-      setInitError(null);
-
-      try {
-        const stream = geminiClient.sendMessageStream(
-          finalQueryToSend,
+      return promptIdContext.run(prompt_id, async () => {
+        const { queryToSend, shouldProceed } = await prepareQueryForGemini(
+          query,
+          userMessageTimestamp,
          abortSignal,
          prompt_id!,
        );
-        const processingStatus = await processGeminiStreamEvents(
-          stream,
-          userMessageTimestamp,
-          abortSignal,
-        );

-        if (processingStatus === StreamProcessingStatus.UserCancelled) {
-          // Restore original model if it was temporarily overridden
-          restoreOriginalModel().catch((error) => {
-            console.error('Failed to restore original model:', error);
-          });
+        if (!shouldProceed || queryToSend === null) {
          isSubmittingQueryRef.current = false;
          return;
        }

-        if (pendingHistoryItemRef.current) {
-          addItem(pendingHistoryItemRef.current, userMessageTimestamp);
-          setPendingHistoryItem(null);
-        }
-        if (loopDetectedRef.current) {
-          loopDetectedRef.current = false;
-          handleLoopDetectedEvent();
+        // Handle vision switch requirement
+        const visionSwitchResult = await handleVisionSwitch(
+          queryToSend,
+          userMessageTimestamp,
+          options?.isContinuation || false,
+        );
+
+        if (!visionSwitchResult.shouldProceed) {
+          isSubmittingQueryRef.current = false;
+          return;
        }

-        // Restore original model if it was temporarily overridden
-        restoreOriginalModel().catch((error) => {
-          console.error('Failed to restore original model:', error);
-        });
-      } catch (error: unknown) {
-        // Restore original model if it was temporarily overridden
-        restoreOriginalModel().catch((error) => {
-          console.error('Failed to restore original model:', error);
-        });
+        const finalQueryToSend = queryToSend;

-        if (error instanceof UnauthorizedError) {
-          onAuthError();
-        } else if (!isNodeError(error) || error.name !== 'AbortError') {
-          addItem(
-            {
-              type: MessageType.ERROR,
-              text: parseAndFormatApiError(
-                getErrorMessage(error) || 'Unknown error',
+        if (!options?.isContinuation) {
+          if (typeof queryToSend === 'string') {
+            // logging the text prompts only for now
+            const promptText = queryToSend;
+            logUserPrompt(
+              config,
+              new UserPromptEvent(
+                promptText.length,
+                prompt_id,
                config.getContentGeneratorConfig()?.authType,
-                undefined,
-                config.getModel(),
-                DEFAULT_GEMINI_FLASH_MODEL,
+                promptText,
              ),
-            },
-            userMessageTimestamp,
-          );
+            );
+          }
+          startNewPrompt();
+          setThought(null); // Reset thought when starting a new prompt
        }
-      } finally {
-        setIsResponding(false);
-        isSubmittingQueryRef.current = false;
-      }
+
+        setIsResponding(true);
+        setInitError(null);
+
+        try {
+          const stream = geminiClient.sendMessageStream(
+            finalQueryToSend,
+            abortSignal,
+            prompt_id!,
+          );
+          const processingStatus = await processGeminiStreamEvents(
+            stream,
+            userMessageTimestamp,
+            abortSignal,
+          );
+
+          if (processingStatus === StreamProcessingStatus.UserCancelled) {
+            // Restore original model if it was temporarily overridden
+            restoreOriginalModel().catch((error) => {
+              console.error('Failed to restore original model:', error);
+            });
+            isSubmittingQueryRef.current = false;
+            return;
+          }
+
+          if (pendingHistoryItemRef.current) {
+            addItem(pendingHistoryItemRef.current, userMessageTimestamp);
+            setPendingHistoryItem(null);
+          }
+          if (loopDetectedRef.current) {
+            loopDetectedRef.current = false;
+            handleLoopDetectedEvent();
+          }
+
+          // Restore original model if it was temporarily overridden
+          restoreOriginalModel().catch((error) => {
+            console.error('Failed to restore original model:', error);
+          });
+        } catch (error: unknown) {
+          // Restore original model if it was temporarily overridden
+          restoreOriginalModel().catch((error) => {
+            console.error('Failed to restore original model:', error);
+          });
+
+          if (error instanceof UnauthorizedError) {
+            onAuthError('Session expired or is unauthorized.');
+          } else if (!isNodeError(error) || error.name !== 'AbortError') {
+            addItem(
+              {
+                type: MessageType.ERROR,
+                text: parseAndFormatApiError(
+                  getErrorMessage(error) || 'Unknown error',
+                  config.getContentGeneratorConfig()?.authType,
+                  undefined,
+                  config.getModel(),
+                  DEFAULT_GEMINI_FLASH_MODEL,
+                ),
+              },
+              userMessageTimestamp,
+            );
+          }
+        } finally {
+          setIsResponding(false);
+          isSubmittingQueryRef.current = false;
+        }
+      });
    },
    [
      streamingState,
@@ -848,6 +952,45 @@ export const useGeminiStream = (
    ],
  );

+  const handleApprovalModeChange = useCallback(
+    async (newApprovalMode: ApprovalMode) => {
+      // Auto-approve pending tool calls when switching to auto-approval modes
+      if (
+        newApprovalMode === ApprovalMode.YOLO ||
+        newApprovalMode === ApprovalMode.AUTO_EDIT
+      ) {
+        let awaitingApprovalCalls = toolCalls.filter(
+          (call): call is TrackedWaitingToolCall =>
+            call.status === 'awaiting_approval',
+        );
+
+        // For AUTO_EDIT mode, only approve edit tools (replace, write_file)
+        if (newApprovalMode === ApprovalMode.AUTO_EDIT) {
+          awaitingApprovalCalls = awaitingApprovalCalls.filter((call) =>
+            EDIT_TOOL_NAMES.has(call.request.name),
+          );
+        }
+
+        // Process pending tool calls sequentially to reduce UI chaos
+        for (const call of awaitingApprovalCalls) {
+          if (call.confirmationDetails?.onConfirm) {
+            try {
+              await call.confirmationDetails.onConfirm(
+                ToolConfirmationOutcome.ProceedOnce,
+              );
+            } catch (error) {
+              console.error(
+                `Failed to auto-approve tool call ${call.request.callId}:`,
+                error,
+              );
+            }
+          }
+        }
+      }
+    },
+    [toolCalls],
+  );
+
  const handleCompletedTools = useCallback(
    async (completedToolCallsFromScheduler: TrackedToolCall[]) => {
      if (isResponding) {
@@ -972,10 +1115,10 @@ export const useGeminiStream = (

  const pendingHistoryItems = useMemo(
    () =>
-      [pendingHistoryItemRef.current, pendingToolCallGroupDisplay].filter(
+      [pendingHistoryItem, pendingToolCallGroupDisplay].filter(
        (i) => i !== undefined && i !== null,
      ),
-    [pendingHistoryItemRef, pendingToolCallGroupDisplay],
+    [pendingHistoryItem, pendingToolCallGroupDisplay],
  );

  useEffect(() => {
@@ -985,8 +1128,7 @@ export const useGeminiStream = (
      }
      const restorableToolCalls = toolCalls.filter(
        (toolCall) =>
-          (toolCall.request.name === 'edit' ||
-            toolCall.request.name === 'write_file') &&
+          EDIT_TOOL_NAMES.has(toolCall.request.name) &&
          toolCall.status === 'awaiting_approval',
      );

@@ -1105,5 +1247,9 @@ export const useGeminiStream = (
    pendingHistoryItems,
    thought,
    cancelOngoingRequest,
+    pendingToolCalls: toolCalls,
+    handleApprovalModeChange,
+    activePtyId,
+    loopDetectionConfirmationRequest,
  };
 };