Handle unhandled rejections more gracefully. (#4417)

Co-authored-by: Tommaso Sciortino <[email protected]>

2 files changed, 140 insertions, 184 deletions

diff --git a/packages/cli/src/ui/hooks/useConsoleMessages.test.ts b/packages/cli/src/ui/hooks/useConsoleMessages.test.ts
index 3b225ecf..b1d1acd6 100644
--- a/packages/cli/src/ui/hooks/useConsoleMessages.test.ts
+++ b/packages/cli/src/ui/hooks/useConsoleMessages.test.ts
@@ -5,127 +5,105 @@
  */
 
 import { act, renderHook } from '@testing-library/react';
-import { useConsoleMessages } from './useConsoleMessages.js';
-import { ConsoleMessageItem } from '../types.js';
-
-// Mock setTimeout and clearTimeout
-vi.useFakeTimers();
+import { vi } from 'vitest';
+import { useConsoleMessages } from './useConsoleMessages';
+import { useCallback } from 'react';
 
 describe('useConsoleMessages', () => {
-  it('should initialize with an empty array of console messages', () => {
-    const { result } = renderHook(() => useConsoleMessages());
-    expect(result.current.consoleMessages).toEqual([]);
+  beforeEach(() => {
+    vi.useFakeTimers();
   });
 
-  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);
-    });
+  afterEach(() => {
+    vi.runOnlyPendingTimers();
+    vi.useRealTimers();
+  });
 
-    act(() => {
-      vi.runAllTimers(); // Process the queue
-    });
+  const useTestableConsoleMessages = () => {
+    const { handleNewMessage, ...rest } = useConsoleMessages();
+    const log = useCallback(
+      (content: string) => handleNewMessage({ type: 'log', content, count: 1 }),
+      [handleNewMessage],
+    );
+    const error = useCallback(
+      (content: string) =>
+        handleNewMessage({ type: 'error', content, count: 1 }),
+      [handleNewMessage],
+    );
+    return {
+      ...rest,
+      log,
+      error,
+      clearConsoleMessages: rest.clearConsoleMessages,
+    };
+  };
 
-    expect(result.current.consoleMessages).toEqual([{ ...message, count: 1 }]);
+  it('should initialize with an empty array of console messages', () => {
+    const { result } = renderHook(() => useTestableConsoleMessages());
+    expect(result.current.consoleMessages).toEqual([]);
   });
 
-  it('should consolidate identical consecutive messages', () => {
-    const { result } = renderHook(() => useConsoleMessages());
-    const message: ConsoleMessageItem = {
-      type: 'log',
-      content: 'Test message',
-      count: 1,
-    };
+  it('should add a new message when log is called', async () => {
+    const { result } = renderHook(() => useTestableConsoleMessages());
 
     act(() => {
-      result.current.handleNewMessage(message);
-      result.current.handleNewMessage(message);
+      result.current.log('Test message');
     });
 
-    act(() => {
-      vi.runAllTimers();
+    await act(async () => {
+      await vi.advanceTimersByTimeAsync(20);
     });
 
-    expect(result.current.consoleMessages).toEqual([{ ...message, count: 2 }]);
+    expect(result.current.consoleMessages).toEqual([
+      { type: 'log', content: 'Test message', count: 1 },
+    ]);
   });
 
-  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,
-    };
+  it('should batch and count identical consecutive messages', async () => {
+    const { result } = renderHook(() => useTestableConsoleMessages());
 
     act(() => {
-      result.current.handleNewMessage(message1);
-      result.current.handleNewMessage(message2);
+      result.current.log('Test message');
+      result.current.log('Test message');
+      result.current.log('Test message');
     });
 
-    act(() => {
-      vi.runAllTimers();
+    await act(async () => {
+      await vi.advanceTimersByTimeAsync(20);
     });
 
     expect(result.current.consoleMessages).toEqual([
-      { ...message1, count: 1 },
-      { ...message2, count: 1 },
+      { type: 'log', content: 'Test message', count: 3 },
     ]);
   });
 
-  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,
-    };
+  it('should not batch different messages', async () => {
+    const { result } = renderHook(() => useTestableConsoleMessages());
 
     act(() => {
-      result.current.handleNewMessage(message1);
-      result.current.handleNewMessage(message2);
+      result.current.log('First message');
+      result.current.error('Second message');
     });
 
-    act(() => {
-      vi.runAllTimers();
+    await act(async () => {
+      await vi.advanceTimersByTimeAsync(20);
     });
 
     expect(result.current.consoleMessages).toEqual([
-      { ...message1, count: 1 },
-      { ...message2, count: 1 },
+      { type: 'log', content: 'First message', count: 1 },
+      { type: 'error', content: 'Second message', count: 1 },
     ]);
   });
 
-  it('should clear console messages', () => {
-    const { result } = renderHook(() => useConsoleMessages());
-    const message: ConsoleMessageItem = {
-      type: 'log',
-      content: 'Test message',
-      count: 1,
-    };
+  it('should clear all messages when clearConsoleMessages is called', async () => {
+    const { result } = renderHook(() => useTestableConsoleMessages());
 
     act(() => {
-      result.current.handleNewMessage(message);
+      result.current.log('A message');
     });
 
-    act(() => {
-      vi.runAllTimers();
+    await act(async () => {
+      await vi.advanceTimersByTimeAsync(20);
     });
 
     expect(result.current.consoleMessages).toHaveLength(1);
@@ -134,79 +112,36 @@ describe('useConsoleMessages', () => {
       result.current.clearConsoleMessages();
     });
 
-    expect(result.current.consoleMessages).toEqual([]);
+    expect(result.current.consoleMessages).toHaveLength(0);
   });
 
-  it('should clear pending timeout on clearConsoleMessages', () => {
-    const { result } = renderHook(() => useConsoleMessages());
-    const message: ConsoleMessageItem = {
-      type: 'log',
-      content: 'Test message',
-      count: 1,
-    };
+  it('should clear the pending timeout when clearConsoleMessages is called', () => {
+    const { result } = renderHook(() => useTestableConsoleMessages());
+    const clearTimeoutSpy = vi.spyOn(global, 'clearTimeout');
 
     act(() => {
-      result.current.handleNewMessage(message); // This schedules a timeout
+      result.current.log('A message');
     });
 
     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([]);
+    expect(clearTimeoutSpy).toHaveBeenCalled();
+    clearTimeoutSpy.mockRestore();
   });
 
-  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,
-    };
+  it('should clean up the timeout on unmount', () => {
+    const { result, unmount } = renderHook(() => useTestableConsoleMessages());
+    const clearTimeoutSpy = vi.spyOn(global, 'clearTimeout');
 
     act(() => {
-      result.current.handleNewMessage(message);
+      result.current.log('A 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
+    expect(clearTimeoutSpy).toHaveBeenCalled();
+    clearTimeoutSpy.mockRestore();
   });
 });
diff --git a/packages/cli/src/ui/hooks/useConsoleMessages.ts b/packages/cli/src/ui/hooks/useConsoleMessages.ts
index 52ffbd39..3b71560e 100644
--- a/packages/cli/src/ui/hooks/useConsoleMessages.ts
+++ b/packages/cli/src/ui/hooks/useConsoleMessages.ts
@@ -4,7 +4,13 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { useCallback, useEffect, useRef, useState } from 'react';
+import {
+  useCallback,
+  useEffect,
+  useReducer,
+  useRef,
+  useTransition,
+} from 'react';
 import { ConsoleMessageItem } from '../types.js';
 
 export interface UseConsoleMessagesReturn {
@@ -13,75 +19,90 @@ export interface UseConsoleMessagesReturn {
   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;
-    }
-
-    const newMessagesToAdd = messageQueueRef.current;
-    messageQueueRef.current = [];
+type Action =
+  | { type: 'ADD_MESSAGES'; payload: ConsoleMessageItem[] }
+  | { type: 'CLEAR' };
 
-    setConsoleMessages((prevMessages) => {
-      const newMessages = [...prevMessages];
-      newMessagesToAdd.forEach((queuedMessage) => {
+function consoleMessagesReducer(
+  state: ConsoleMessageItem[],
+  action: Action,
+): ConsoleMessageItem[] {
+  switch (action.type) {
+    case 'ADD_MESSAGES': {
+      const newMessages = [...state];
+      for (const queuedMessage of action.payload) {
+        const lastMessage = newMessages[newMessages.length - 1];
         if (
-          newMessages.length > 0 &&
-          newMessages[newMessages.length - 1].type === queuedMessage.type &&
-          newMessages[newMessages.length - 1].content === queuedMessage.content
+          lastMessage &&
+          lastMessage.type === queuedMessage.type &&
+          lastMessage.content === queuedMessage.content
         ) {
-          newMessages[newMessages.length - 1].count =
-            (newMessages[newMessages.length - 1].count || 1) + 1;
+          // Create a new object for the last message to ensure React detects
+          // the change, preventing mutation of the existing state object.
+          newMessages[newMessages.length - 1] = {
+            ...lastMessage,
+            count: lastMessage.count + 1,
+          };
         } else {
           newMessages.push({ ...queuedMessage, count: 1 });
         }
-      });
+      }
       return newMessages;
-    });
+    }
+    case 'CLEAR':
+      return [];
+    default:
+      return state;
+  }
+}
 
-    messageQueueTimeoutRef.current = null; // Allow next scheduling
-  }, []);
+export function useConsoleMessages(): UseConsoleMessagesReturn {
+  const [consoleMessages, dispatch] = useReducer(consoleMessagesReducer, []);
+  const messageQueueRef = useRef<ConsoleMessageItem[]>([]);
+  const timeoutRef = useRef<NodeJS.Timeout | null>(null);
+  const [, startTransition] = useTransition();
 
-  const scheduleQueueProcessing = useCallback(() => {
-    if (messageQueueTimeoutRef.current === null) {
-      messageQueueTimeoutRef.current = setTimeout(
-        processMessageQueue,
-        0,
-      ) as unknown as number;
+  const processQueue = useCallback(() => {
+    if (messageQueueRef.current.length > 0) {
+      const messagesToProcess = messageQueueRef.current;
+      messageQueueRef.current = [];
+      startTransition(() => {
+        dispatch({ type: 'ADD_MESSAGES', payload: messagesToProcess });
+      });
     }
-  }, [processMessageQueue]);
+    timeoutRef.current = null;
+  }, []);
 
   const handleNewMessage = useCallback(
     (message: ConsoleMessageItem) => {
       messageQueueRef.current.push(message);
-      scheduleQueueProcessing();
+      if (!timeoutRef.current) {
+        // Batch updates using a timeout. 16ms is a reasonable delay to batch
+        // rapid-fire messages without noticeable lag.
+        timeoutRef.current = setTimeout(processQueue, 16);
+      }
     },
-    [scheduleQueueProcessing],
+    [processQueue],
   );
 
   const clearConsoleMessages = useCallback(() => {
-    setConsoleMessages([]);
-    if (messageQueueTimeoutRef.current !== null) {
-      clearTimeout(messageQueueTimeoutRef.current);
-      messageQueueTimeoutRef.current = null;
+    if (timeoutRef.current) {
+      clearTimeout(timeoutRef.current);
+      timeoutRef.current = null;
     }
     messageQueueRef.current = [];
+    startTransition(() => {
+      dispatch({ type: 'CLEAR' });
+    });
   }, []);
 
+  // Cleanup on unmount
   useEffect(
-    () =>
-      // Cleanup on unmount
-      () => {
-        if (messageQueueTimeoutRef.current !== null) {
-          clearTimeout(messageQueueTimeoutRef.current);
-        }
-      },
+    () => () => {
+      if (timeoutRef.current) {
+        clearTimeout(timeoutRef.current);
+      }
+    },
     [],
   );