fix(cli): fix issue updating a component while rendering a different component (#14319)

2026-03-22 20:14:58 -07:00 · 2025-12-01 17:33:03 -08:00
parent 98d7238ed6
commit 1689e9b671
9 changed files with 43 additions and 90 deletions
--- a/integration-tests/extensions-reload.test.ts
+++ b/integration-tests/extensions-reload.test.ts
@@ -76,15 +76,16 @@ describe('extension reloading', () => {
      await run.expectText(
        'test-extension (v0.0.1) - active (update available)',
      );
-      await run.sendText('/mcp list');
-      await run.type('\r');
+      // Wait a bit for the UI to settle
+      await new Promise((resolve) => setTimeout(resolve, 1000));
+      await run.sendKeys('\u0015/mcp list\r');
      await run.expectText(
        'test-server (from test-extension) - Ready (1 tool)',
      );
      await run.expectText('- hello');

      // Update the extension, expect the list to update, and mcp servers as well.
-      await run.sendKeys('/extensions update test-extension');
+      await run.sendKeys('\u0015/extensions update test-extension');
      await run.expectText('/extensions update test-extension');
      await run.sendKeys('\r');
      await new Promise((resolve) => setTimeout(resolve, 500));
@@ -96,12 +97,12 @@ describe('extension reloading', () => {
      await run.expectText(
        'Extension "test-extension" successfully updated: 0.0.1 → 0.0.2',
      );
-      await new Promise((resolve) => setTimeout(resolve, 1000));
-      await run.sendText('/extensions list');
-      await run.type('\r');
+      await new Promise((resolve) => setTimeout(resolve, 3000));
+      await run.sendKeys('\u0015/extensions list\r');
      await run.expectText('test-extension (v0.0.2) - active (updated)');
-      await run.sendText('/mcp list');
-      await run.type('\r');
+      // Wait a bit for the UI to settle
+      await new Promise((resolve) => setTimeout(resolve, 1000));
+      await run.sendKeys('\u0015/mcp list\r');
      await run.expectText(
        'test-server (from test-extension) - Ready (1 tool)',
      );
--- a/packages/cli/src/ui/AppContainer.tsx
+++ b/packages/cli/src/ui/AppContainer.tsx
@@ -832,11 +832,14 @@ Logging in with Google... Restarting Gemini CLI to continue.
  useLayoutEffect(() => {
    if (mainControlsRef.current) {
      const fullFooterMeasurement = measureElement(mainControlsRef.current);
-      if (fullFooterMeasurement.height > 0) {
+      if (
+        fullFooterMeasurement.height > 0 &&
+        fullFooterMeasurement.height !== controlsHeight
+      ) {
        setControlsHeight(fullFooterMeasurement.height);
      }
    }
-  }, [buffer, terminalWidth, terminalHeight]);
+  }, [buffer, terminalWidth, terminalHeight, controlsHeight]);

  // Compute available terminal height based on controls measurement
  const availableTerminalHeight = Math.max(
--- a/packages/cli/src/ui/components/InputPrompt.test.tsx
+++ b/packages/cli/src/ui/components/InputPrompt.test.tsx
@@ -1147,7 +1147,6 @@ describe('InputPrompt', () => {
      await waitFor(() => {
        expect(mockedUseCommandCompletion).toHaveBeenCalledWith(
          mockBuffer,
-          ['/test/project/src'],
          path.join('test', 'project', 'src'),
          mockSlashCommands,
          mockCommandContext,
@@ -2268,6 +2267,7 @@ describe('InputPrompt', () => {
  describe('queued message editing', () => {
    it('should load all queued messages when up arrow is pressed with empty input', async () => {
      const mockPopAllMessages = vi.fn();
+      mockPopAllMessages.mockReturnValue('Message 1\n\nMessage 2\n\nMessage 3');
      props.popAllMessages = mockPopAllMessages;
      props.buffer.text = '';

@@ -2279,11 +2279,7 @@ describe('InputPrompt', () => {
        stdin.write('\u001B[A');
      });
      await waitFor(() => expect(mockPopAllMessages).toHaveBeenCalled());
-      const callback = mockPopAllMessages.mock.calls[0][0];

-      await act(async () => {
-        callback('Message 1\n\nMessage 2\n\nMessage 3');
-      });
      expect(props.buffer.setText).toHaveBeenCalledWith(
        'Message 1\n\nMessage 2\n\nMessage 3',
      );
@@ -2311,6 +2307,7 @@ describe('InputPrompt', () => {

    it('should handle undefined messages from popAllMessages', async () => {
      const mockPopAllMessages = vi.fn();
+      mockPopAllMessages.mockReturnValue(undefined);
      props.popAllMessages = mockPopAllMessages;
      props.buffer.text = '';

@@ -2322,10 +2319,6 @@ describe('InputPrompt', () => {
        stdin.write('\u001B[A');
      });
      await waitFor(() => expect(mockPopAllMessages).toHaveBeenCalled());
-      const callback = mockPopAllMessages.mock.calls[0][0];
-      await act(async () => {
-        callback(undefined);
-      });

      expect(props.buffer.setText).not.toHaveBeenCalled();
      expect(mockInputHistory.navigateUp).toHaveBeenCalled();
@@ -2353,6 +2346,7 @@ describe('InputPrompt', () => {

    it('should handle single queued message', async () => {
      const mockPopAllMessages = vi.fn();
+      mockPopAllMessages.mockReturnValue('Single message');
      props.popAllMessages = mockPopAllMessages;
      props.buffer.text = '';

@@ -2365,11 +2359,6 @@ describe('InputPrompt', () => {
      });
      await waitFor(() => expect(mockPopAllMessages).toHaveBeenCalled());

-      const callback = mockPopAllMessages.mock.calls[0][0];
-      await act(async () => {
-        callback('Single message');
-      });
-
      expect(props.buffer.setText).toHaveBeenCalledWith('Single message');
      unmount();
    });
@@ -2409,6 +2398,7 @@ describe('InputPrompt', () => {

    it('should navigate input history on fresh start when no queued messages exist', async () => {
      const mockPopAllMessages = vi.fn();
+      mockPopAllMessages.mockReturnValue(undefined);
      props.popAllMessages = mockPopAllMessages;
      props.buffer.text = '';

@@ -2421,11 +2411,6 @@ describe('InputPrompt', () => {
      });
      await waitFor(() => expect(mockPopAllMessages).toHaveBeenCalled());

-      const callback = mockPopAllMessages.mock.calls[0][0];
-      await act(async () => {
-        callback(undefined);
-      });
-
      expect(mockInputHistory.navigateUp).toHaveBeenCalled();
      expect(props.buffer.setText).not.toHaveBeenCalled();

--- a/packages/cli/src/ui/components/InputPrompt.tsx
+++ b/packages/cli/src/ui/components/InputPrompt.tsx
@@ -83,7 +83,7 @@ export interface InputPromptProps {
  isEmbeddedShellFocused?: boolean;
  setQueueErrorMessage: (message: string | null) => void;
  streamingState: StreamingState;
-  popAllMessages?: (onPop: (messages: string | undefined) => void) => void;
+  popAllMessages?: () => string | undefined;
  suggestionsPosition?: 'above' | 'below';
  setBannerVisible: (visible: boolean) => void;
 }
@@ -143,15 +143,6 @@ export const InputPrompt: React.FC<InputPromptProps> = ({
  const pasteTimeoutRef = useRef<NodeJS.Timeout | null>(null);
  const innerBoxRef = useRef<DOMElement>(null);

-  const [dirs, setDirs] = useState<readonly string[]>(
-    config.getWorkspaceContext().getDirectories(),
-  );
-  const dirsChanged = config.getWorkspaceContext().getDirectories();
-  useEffect(() => {
-    if (dirs.length !== dirsChanged.length) {
-      setDirs(dirsChanged);
-    }
-  }, [dirs.length, dirsChanged]);
  const [reverseSearchActive, setReverseSearchActive] = useState(false);
  const [commandSearchActive, setCommandSearchActive] = useState(false);
  const [textBeforeReverseSearch, setTextBeforeReverseSearch] = useState('');
@@ -165,7 +156,6 @@ export const InputPrompt: React.FC<InputPromptProps> = ({

  const completion = useCommandCompletion(
    buffer,
-    dirs,
    config.getTargetDir(),
    slashCommands,
    commandContext,
@@ -310,14 +300,13 @@ export const InputPrompt: React.FC<InputPromptProps> = ({
  // Returns true if we should continue with input history navigation
  const tryLoadQueuedMessages = useCallback(() => {
    if (buffer.text.trim() === '' && popAllMessages) {
-      popAllMessages((allMessages) => {
-        if (allMessages) {
-          buffer.setText(allMessages);
-        } else {
-          // No queued messages, proceed with input history
-          inputHistory.navigateUp();
-        }
-      });
+      const allMessages = popAllMessages();
+      if (allMessages) {
+        buffer.setText(allMessages);
+      } else {
+        // No queued messages, proceed with input history
+        inputHistory.navigateUp();
+      }
      return true; // We handled the up arrow key
    }
    return false;
--- a/packages/cli/src/ui/contexts/UIActionsContext.tsx
+++ b/packages/cli/src/ui/contexts/UIActionsContext.tsx
@@ -51,7 +51,7 @@ export interface UIActions {
  handleResumeSession: (session: SessionInfo) => Promise<void>;
  handleDeleteSession: (session: SessionInfo) => Promise<void>;
  setQueueErrorMessage: (message: string | null) => void;
-  popAllMessages: (onPop: (messages: string | undefined) => void) => void;
+  popAllMessages: () => string | undefined;
  handleApiKeySubmit: (apiKey: string) => Promise<void>;
  handleApiKeyCancel: () => void;
  setBannerVisible: (visible: boolean) => void;
--- a/packages/cli/src/ui/hooks/useCommandCompletion.test.tsx
+++ b/packages/cli/src/ui/hooks/useCommandCompletion.test.tsx
@@ -94,7 +94,6 @@ describe('useCommandCompletion', () => {
    getEnablePromptCompletion: () => false,
    getGeminiClient: vi.fn(),
  } as unknown as Config;
-  const testDirs: string[] = [];
  const testRootDir = '/';

  // Helper to create real TextBuffer objects within renderHook
@@ -121,7 +120,6 @@ describe('useCommandCompletion', () => {
      const textBuffer = useTextBufferForTest(initialText, cursorOffset);
      const completion = useCommandCompletion(
        textBuffer,
-        testDirs,
        testRootDir,
        [],
        mockCommandContext,
@@ -505,7 +503,6 @@ describe('useCommandCompletion', () => {
        const textBuffer = useTextBufferForTest('// This is a line comment');
        const completion = useCommandCompletion(
          textBuffer,
-          testDirs,
          testRootDir,
          [],
          mockCommandContext,
@@ -538,7 +535,6 @@ describe('useCommandCompletion', () => {
        );
        const completion = useCommandCompletion(
          textBuffer,
-          testDirs,
          testRootDir,
          [],
          mockCommandContext,
@@ -571,7 +567,6 @@ describe('useCommandCompletion', () => {
        );
        const completion = useCommandCompletion(
          textBuffer,
-          testDirs,
          testRootDir,
          [],
          mockCommandContext,
--- a/packages/cli/src/ui/hooks/useCommandCompletion.tsx
+++ b/packages/cli/src/ui/hooks/useCommandCompletion.tsx
@@ -57,7 +57,6 @@ export interface UseCommandCompletionReturn {

 export function useCommandCompletion(
  buffer: TextBuffer,
-  dirs: readonly string[],
  cwd: string,
  slashCommands: readonly SlashCommand[],
  commandContext: CommandContext,
--- a/packages/cli/src/ui/hooks/useMessageQueue.test.tsx
+++ b/packages/cli/src/ui/hooks/useMessageQueue.test.tsx
@@ -253,9 +253,7 @@ describe('useMessageQueue', () => {
      // Pop all messages
      let poppedMessages: string | undefined;
      act(() => {
-        result.current.popAllMessages((messages) => {
-          poppedMessages = messages;
-        });
+        poppedMessages = result.current.popAllMessages();
      });

      expect(poppedMessages).toBe('Message 1\n\nMessage 2\n\nMessage 3');
@@ -271,9 +269,7 @@ describe('useMessageQueue', () => {

      let poppedMessages: string | undefined = 'not-undefined';
      act(() => {
-        result.current.popAllMessages((messages) => {
-          poppedMessages = messages;
-        });
+        poppedMessages = result.current.popAllMessages();
      });

      expect(poppedMessages).toBeUndefined();
@@ -293,9 +289,7 @@ describe('useMessageQueue', () => {

      let poppedMessages: string | undefined;
      act(() => {
-        result.current.popAllMessages((messages) => {
-          poppedMessages = messages;
-        });
+        poppedMessages = result.current.popAllMessages();
      });

      expect(poppedMessages).toBe('Single message');
@@ -315,7 +309,7 @@ describe('useMessageQueue', () => {
      });

      act(() => {
-        result.current.popAllMessages(() => {});
+        result.current.popAllMessages();
      });

      // Queue should be empty
@@ -325,9 +319,7 @@ describe('useMessageQueue', () => {
      // Popping again should return undefined
      let secondPop: string | undefined = 'not-undefined';
      act(() => {
-        result.current.popAllMessages((messages) => {
-          secondPop = messages;
-        });
+        secondPop = result.current.popAllMessages();
      });

      expect(secondPop).toBeUndefined();
@@ -349,9 +341,7 @@ describe('useMessageQueue', () => {
      // Pop all messages
      let firstPop: string | undefined;
      act(() => {
-        result.current.popAllMessages((messages) => {
-          firstPop = messages;
-        });
+        firstPop = result.current.popAllMessages();
      });

      expect(firstPop).toBe('First\n\nSecond');
@@ -365,9 +355,7 @@ describe('useMessageQueue', () => {
      // Pop again
      let secondPop: string | undefined;
      act(() => {
-        result.current.popAllMessages((messages) => {
-          secondPop = messages;
-        });
+        secondPop = result.current.popAllMessages();
      });

      expect(secondPop).toBe('Third\n\nFourth');
--- a/packages/cli/src/ui/hooks/useMessageQueue.ts
+++ b/packages/cli/src/ui/hooks/useMessageQueue.ts
@@ -18,7 +18,7 @@ export interface UseMessageQueueReturn {
  addMessage: (message: string) => void;
  clearQueue: () => void;
  getQueuedMessagesText: () => string;
-  popAllMessages: (onPop: (messages: string | undefined) => void) => void;
+  popAllMessages: () => string | undefined;
 }

 /**
@@ -53,21 +53,14 @@ export function useMessageQueue({
  }, [messageQueue]);

  // Pop all messages from the queue and return them as a single string
-  const popAllMessages = useCallback(
-    (onPop: (messages: string | undefined) => void) => {
-      setMessageQueue((prev) => {
-        if (prev.length === 0) {
-          onPop(undefined);
-          return prev;
-        }
-        // Join all messages with double newlines, same as when they're sent
-        const allMessages = prev.join('\n\n');
-        onPop(allMessages);
-        return []; // Clear the entire queue
-      });
-    },
-    [],
-  );
+  const popAllMessages = useCallback(() => {
+    if (messageQueue.length === 0) {
+      return undefined;
+    }
+    const allMessages = messageQueue.join('\n\n');
+    setMessageQueue([]);
+    return allMessages;
+  }, [messageQueue]);

  // Process queued messages when streaming becomes idle
  useEffect(() => {