Make console message support more robust to logging in the middle of rendering. (#521)

6 files changed, 313 insertions, 16 deletions

diff --git a/packages/cli/src/ui/App.tsx b/packages/cli/src/ui/App.tsx
index 79b6b9ab..1453487f 100644
--- a/packages/cli/src/ui/App.tsx
+++ b/packages/cli/src/ui/App.tsx
@@ -14,18 +14,14 @@ import {
   useInput,
   type Key as InkKeyType,
 } from 'ink';
-import {
-  StreamingState,
-  type HistoryItem,
-  ConsoleMessageItem,
-  MessageType,
-} from './types.js';
+import { StreamingState, type HistoryItem, MessageType } from './types.js';
 import { useTerminalSize } from './hooks/useTerminalSize.js';
 import { useGeminiStream } from './hooks/useGeminiStream.js';
 import { useLoadingIndicator } from './hooks/useLoadingIndicator.js';
 import { useThemeCommand } from './hooks/useThemeCommand.js';
 import { useSlashCommandProcessor } from './hooks/slashCommandProcessor.js';
 import { useAutoAcceptIndicator } from './hooks/useAutoAcceptIndicator.js';
+import { useConsoleMessages } from './hooks/useConsoleMessages.js';
 import { Header } from './components/Header.js';
 import { LoadingIndicator } from './components/LoadingIndicator.js';
 import { AutoAcceptIndicator } from './components/AutoAcceptIndicator.js';
@@ -60,6 +56,11 @@ export const App = ({
   startupWarnings = [],
 }: AppProps) => {
   const { history, addItem, clearItems } = useHistory();
+  const {
+    consoleMessages,
+    handleNewMessage,
+    clearConsoleMessages: clearConsoleMessagesState,
+  } = useConsoleMessages();
   const [staticNeedsRefresh, setStaticNeedsRefresh] = useState(false);
   const [staticKey, setStaticKey] = useState(0);
   const refreshStatic = useCallback(() => {
@@ -73,10 +74,6 @@ export const App = ({
   const [footerHeight, setFooterHeight] = useState<number>(0);
   const [corgiMode, setCorgiMode] = useState(false);
   const [shellModeActive, setShellModeActive] = useState(false);
-
-  const [consoleMessages, setConsoleMessages] = useState<ConsoleMessageItem[]>(
-    [],
-  );
   const [showErrorDetails, setShowErrorDetails] = useState<boolean>(false);
 
   const errorCount = useMemo(
@@ -90,10 +87,6 @@ export const App = ({
     }
   });
 
-  const handleNewMessage = useCallback((message: ConsoleMessageItem) => {
-    setConsoleMessages((prevMessages) => [...prevMessages, message]);
-  }, []);
-
   useConsolePatcher({
     onNewMessage: handleNewMessage,
     debugMode: config.getDebugMode(),
@@ -232,10 +225,10 @@ export const App = ({
 
   const handleClearScreen = useCallback(() => {
     clearItems();
-    setConsoleMessages([]);
+    clearConsoleMessagesState();
     console.clear();
     refreshStatic();
-  }, [clearItems, refreshStatic]);
+  }, [clearItems, clearConsoleMessagesState, refreshStatic]);
 
   const { rows: terminalHeight } = useTerminalSize();
   const mainControlsRef = useRef<DOMElement>(null);
diff --git a/packages/cli/src/ui/components/ConsolePatcher.tsx b/packages/cli/src/ui/components/ConsolePatcher.tsx
index 240a32e5..843c6320 100644
--- a/packages/cli/src/ui/components/ConsolePatcher.tsx
+++ b/packages/cli/src/ui/components/ConsolePatcher.tsx
@@ -40,6 +40,7 @@ export const useConsolePatcher = ({
           onNewMessage({
             type,
             content: formatArgs(args),
+            count: 1,
           });
         }
       };
diff --git a/packages/cli/src/ui/components/DetailedMessagesDisplay.tsx b/packages/cli/src/ui/components/DetailedMessagesDisplay.tsx
index c1d65a9a..1ab78426 100644
--- a/packages/cli/src/ui/components/DetailedMessagesDisplay.tsx
+++ b/packages/cli/src/ui/components/DetailedMessagesDisplay.tsx
@@ -64,6 +64,9 @@ export const DetailedMessagesDisplay: React.FC<
             <Text color={textColor}>{icon} </Text>
             <Text color={textColor} wrap="wrap">
               {msg.content}
+              {msg.count && msg.count > 1 && (
+                <Text color={Colors.SubtleComment}> (x{msg.count})</Text>
+              )}
             </Text>
           </Box>
         );
diff --git a/packages/cli/src/ui/hooks/useConsoleMessages.test.ts b/packages/cli/src/ui/hooks/useConsoleMessages.test.ts
new file mode 100644
index 00000000..3b225ecf
--- /dev/null
+++ b/packages/cli/src/ui/hooks/useConsoleMessages.test.ts
@@ -0,0 +1,212 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { act, renderHook } from '@testing-library/react';
+import { useConsoleMessages } from './useConsoleMessages.js';
+import { ConsoleMessageItem } from '../types.js';
+
+// Mock setTimeout and clearTimeout
+vi.useFakeTimers();
+
+describe('useConsoleMessages', () => {
+  it('should initialize with an empty array of console messages', () => {
+    const { result } = renderHook(() => useConsoleMessages());
+    expect(result.current.consoleMessages).toEqual([]);
+  });
+
+  it('should add a new message', () => {
+    const { result } = renderHook(() => useConsoleMessages());
+    const message: ConsoleMessageItem = {
+      type: 'log',
+      content: 'Test message',
+      count: 1,
+    };
+
+    act(() => {
+      result.current.handleNewMessage(message);
+    });
+
+    act(() => {
+      vi.runAllTimers(); // Process the queue
+    });
+
+    expect(result.current.consoleMessages).toEqual([{ ...message, count: 1 }]);
+  });
+
+  it('should consolidate identical consecutive messages', () => {
+    const { result } = renderHook(() => useConsoleMessages());
+    const message: ConsoleMessageItem = {
+      type: 'log',
+      content: 'Test message',
+      count: 1,
+    };
+
+    act(() => {
+      result.current.handleNewMessage(message);
+      result.current.handleNewMessage(message);
+    });
+
+    act(() => {
+      vi.runAllTimers();
+    });
+
+    expect(result.current.consoleMessages).toEqual([{ ...message, count: 2 }]);
+  });
+
+  it('should not consolidate different messages', () => {
+    const { result } = renderHook(() => useConsoleMessages());
+    const message1: ConsoleMessageItem = {
+      type: 'log',
+      content: 'Test message 1',
+      count: 1,
+    };
+    const message2: ConsoleMessageItem = {
+      type: 'error',
+      content: 'Test message 2',
+      count: 1,
+    };
+
+    act(() => {
+      result.current.handleNewMessage(message1);
+      result.current.handleNewMessage(message2);
+    });
+
+    act(() => {
+      vi.runAllTimers();
+    });
+
+    expect(result.current.consoleMessages).toEqual([
+      { ...message1, count: 1 },
+      { ...message2, count: 1 },
+    ]);
+  });
+
+  it('should not consolidate messages if type is different', () => {
+    const { result } = renderHook(() => useConsoleMessages());
+    const message1: ConsoleMessageItem = {
+      type: 'log',
+      content: 'Test message',
+      count: 1,
+    };
+    const message2: ConsoleMessageItem = {
+      type: 'error',
+      content: 'Test message',
+      count: 1,
+    };
+
+    act(() => {
+      result.current.handleNewMessage(message1);
+      result.current.handleNewMessage(message2);
+    });
+
+    act(() => {
+      vi.runAllTimers();
+    });
+
+    expect(result.current.consoleMessages).toEqual([
+      { ...message1, count: 1 },
+      { ...message2, count: 1 },
+    ]);
+  });
+
+  it('should clear console messages', () => {
+    const { result } = renderHook(() => useConsoleMessages());
+    const message: ConsoleMessageItem = {
+      type: 'log',
+      content: 'Test message',
+      count: 1,
+    };
+
+    act(() => {
+      result.current.handleNewMessage(message);
+    });
+
+    act(() => {
+      vi.runAllTimers();
+    });
+
+    expect(result.current.consoleMessages).toHaveLength(1);
+
+    act(() => {
+      result.current.clearConsoleMessages();
+    });
+
+    expect(result.current.consoleMessages).toEqual([]);
+  });
+
+  it('should clear pending timeout on clearConsoleMessages', () => {
+    const { result } = renderHook(() => useConsoleMessages());
+    const message: ConsoleMessageItem = {
+      type: 'log',
+      content: 'Test message',
+      count: 1,
+    };
+
+    act(() => {
+      result.current.handleNewMessage(message); // This schedules a timeout
+    });
+
+    act(() => {
+      result.current.clearConsoleMessages();
+    });
+
+    // Ensure the queue is empty and no more messages are processed
+    act(() => {
+      vi.runAllTimers(); // If timeout wasn't cleared, this would process the queue
+    });
+
+    expect(result.current.consoleMessages).toEqual([]);
+  });
+
+  it('should clear message queue on clearConsoleMessages', () => {
+    const { result } = renderHook(() => useConsoleMessages());
+    const message: ConsoleMessageItem = {
+      type: 'log',
+      content: 'Test message',
+      count: 1,
+    };
+
+    act(() => {
+      // Add a message but don't process the queue yet
+      result.current.handleNewMessage(message);
+    });
+
+    act(() => {
+      result.current.clearConsoleMessages();
+    });
+
+    // Process any pending timeouts (should be none related to message queue)
+    act(() => {
+      vi.runAllTimers();
+    });
+
+    // The consoleMessages should be empty because the queue was cleared before processing
+    expect(result.current.consoleMessages).toEqual([]);
+  });
+
+  it('should cleanup timeout on unmount', () => {
+    const { result, unmount } = renderHook(() => useConsoleMessages());
+    const message: ConsoleMessageItem = {
+      type: 'log',
+      content: 'Test message',
+      count: 1,
+    };
+
+    act(() => {
+      result.current.handleNewMessage(message);
+    });
+
+    unmount();
+
+    // This is a bit indirect. We check that clearTimeout was called.
+    // If clearTimeout was not called, and we run timers, an error might occur
+    // or the state might change, which it shouldn't after unmount.
+    // Vitest's vi.clearAllTimers() or specific checks for clearTimeout calls
+    // would be more direct if available and easy to set up here.
+    // For now, we rely on the useEffect cleanup pattern.
+    expect(vi.getTimerCount()).toBe(0); // Check if all timers are cleared
+  });
+});
diff --git a/packages/cli/src/ui/hooks/useConsoleMessages.ts b/packages/cli/src/ui/hooks/useConsoleMessages.ts
new file mode 100644
index 00000000..90dc6f81
--- /dev/null
+++ b/packages/cli/src/ui/hooks/useConsoleMessages.ts
@@ -0,0 +1,87 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { useCallback, useEffect, useRef, useState } from 'react';
+import { ConsoleMessageItem } from '../types.js';
+
+export interface UseConsoleMessagesReturn {
+  consoleMessages: ConsoleMessageItem[];
+  handleNewMessage: (message: ConsoleMessageItem) => void;
+  clearConsoleMessages: () => void;
+}
+
+export function useConsoleMessages(): UseConsoleMessagesReturn {
+  const [consoleMessages, setConsoleMessages] = useState<ConsoleMessageItem[]>(
+    [],
+  );
+  const messageQueueRef = useRef<ConsoleMessageItem[]>([]);
+  const messageQueueTimeoutRef = useRef<number | null>(null);
+
+  const processMessageQueue = useCallback(() => {
+    if (messageQueueRef.current.length === 0) {
+      return;
+    }
+
+    setConsoleMessages((prevMessages) => {
+      const newMessages = [...prevMessages];
+      messageQueueRef.current.forEach((queuedMessage) => {
+        if (
+          newMessages.length > 0 &&
+          newMessages[newMessages.length - 1].type === queuedMessage.type &&
+          newMessages[newMessages.length - 1].content === queuedMessage.content
+        ) {
+          newMessages[newMessages.length - 1].count =
+            (newMessages[newMessages.length - 1].count || 1) + 1;
+        } else {
+          newMessages.push({ ...queuedMessage, count: 1 });
+        }
+      });
+      return newMessages;
+    });
+
+    messageQueueRef.current = [];
+    messageQueueTimeoutRef.current = null; // Allow next scheduling
+  }, []);
+
+  const scheduleQueueProcessing = useCallback(() => {
+    if (messageQueueTimeoutRef.current === null) {
+      messageQueueTimeoutRef.current = setTimeout(
+        processMessageQueue,
+        0,
+      ) as unknown as number;
+    }
+  }, [processMessageQueue]);
+
+  const handleNewMessage = useCallback(
+    (message: ConsoleMessageItem) => {
+      messageQueueRef.current.push(message);
+      scheduleQueueProcessing();
+    },
+    [scheduleQueueProcessing],
+  );
+
+  const clearConsoleMessages = useCallback(() => {
+    setConsoleMessages([]);
+    if (messageQueueTimeoutRef.current !== null) {
+      clearTimeout(messageQueueTimeoutRef.current);
+      messageQueueTimeoutRef.current = null;
+    }
+    messageQueueRef.current = [];
+  }, []);
+
+  useEffect(
+    () =>
+      // Cleanup on unmount
+      () => {
+        if (messageQueueTimeoutRef.current !== null) {
+          clearTimeout(messageQueueTimeoutRef.current);
+        }
+      },
+    [],
+  );
+
+  return { consoleMessages, handleNewMessage, clearConsoleMessages };
+}
diff --git a/packages/cli/src/ui/types.ts b/packages/cli/src/ui/types.ts
index 03b9287f..2826bb1e 100644
--- a/packages/cli/src/ui/types.ts
+++ b/packages/cli/src/ui/types.ts
@@ -143,4 +143,5 @@ export type Message =
 export interface ConsoleMessageItem {
   type: 'log' | 'warn' | 'error' | 'debug';
   content: string;
+  count: number;
 }