7 files changed, 191 insertions, 186 deletions

diff --git a/packages/cli/src/ui/App.tsx b/packages/cli/src/ui/App.tsx
index bccba9e6..6bbe47d4 100644
--- a/packages/cli/src/ui/App.tsx
+++ b/packages/cli/src/ui/App.tsx
@@ -351,18 +351,6 @@ const App = ({ config, settings, startupWarnings = [] }: AppProps) => {
       // the user starts interacting with the app.
       enteringConstrainHeightMode = true;
       setConstrainHeight(true);
-
-      // If our pending history item happens to exceed the terminal height we will most likely need to refresh
-      // our static collection to ensure no duplication or tearing. This is currently working around a core bug
-      // in Ink which we have a PR out to fix: https://github.com/vadimdemedes/ink/pull/717
-      if (pendingHistoryItemRef.current && pendingHistoryItems.length > 0) {
-        const pendingItemDimensions = measureElement(
-          pendingHistoryItemRef.current,
-        );
-        if (pendingItemDimensions.height > availableTerminalHeight) {
-          refreshStatic();
-        }
-      }
     }
 
     if (key.ctrl && input === 'o') {
@@ -531,23 +519,6 @@ const App = ({ config, settings, startupWarnings = [] }: AppProps) => {
   }, [terminalWidth, terminalHeight, refreshStatic]);
 
   useEffect(() => {
-    if (!pendingHistoryItems.length) {
-      return;
-    }
-
-    const pendingItemDimensions = measureElement(
-      pendingHistoryItemRef.current!,
-    );
-
-    // If our pending history item happens to exceed the terminal height we will most likely need to refresh
-    // our static collection to ensure no duplication or tearing. This is currently working around a core bug
-    // in Ink which we have a PR out to fix: https://github.com/vadimdemedes/ink/pull/717
-    if (pendingItemDimensions.height > availableTerminalHeight) {
-      setStaticNeedsRefresh(true);
-    }
-  }, [pendingHistoryItems.length, availableTerminalHeight, streamingState]);
-
-  useEffect(() => {
     if (streamingState === StreamingState.Idle && staticNeedsRefresh) {
       setStaticNeedsRefresh(false);
       refreshStatic();
diff --git a/packages/cli/src/ui/components/shared/MaxSizedBox.tsx b/packages/cli/src/ui/components/shared/MaxSizedBox.tsx
index cd31ff04..8880e894 100644
--- a/packages/cli/src/ui/components/shared/MaxSizedBox.tsx
+++ b/packages/cli/src/ui/components/shared/MaxSizedBox.tsx
@@ -116,13 +116,15 @@ export const MaxSizedBox: React.FC<MaxSizedBoxProps> = ({
     throw new Error('maxWidth must be defined when maxHeight is set.');
   }
   function visitRows(element: React.ReactNode) {
-    if (!React.isValidElement(element)) {
+    if (!React.isValidElement<{ children?: React.ReactNode }>(element)) {
       return;
     }
+
     if (element.type === Fragment) {
       React.Children.forEach(element.props.children, visitRows);
       return;
     }
+
     if (element.type === Box) {
       layoutInkElementAsStyledText(element, maxWidth!, laidOutStyledText);
       return;
@@ -246,7 +248,10 @@ interface Row {
  * @returns An array of `Row` objects.
  */
 function visitBoxRow(element: React.ReactNode): Row {
-  if (!React.isValidElement(element) || element.type !== Box) {
+  if (
+    !React.isValidElement<{ children?: React.ReactNode }>(element) ||
+    element.type !== Box
+  ) {
     debugReportError(
       `All children of MaxSizedBox must be <Box> elements`,
       element,
@@ -258,7 +263,15 @@ function visitBoxRow(element: React.ReactNode): Row {
   }
 
   if (enableDebugLog) {
-    const boxProps = element.props;
+    const boxProps = element.props as {
+      children?: React.ReactNode | undefined;
+      readonly flexDirection?:
+        | 'row'
+        | 'column'
+        | 'row-reverse'
+        | 'column-reverse'
+        | undefined;
+    };
     // Ensure the Box has no props other than the default ones and key.
     let maxExpectedProps = 4;
     if (boxProps.children !== undefined) {
@@ -323,14 +336,13 @@ function visitBoxRow(element: React.ReactNode): Row {
       return;
     }
 
-    if (!React.isValidElement(element)) {
+    if (!React.isValidElement<{ children?: React.ReactNode }>(element)) {
       debugReportError('Invalid element.', element);
       return;
     }
 
     if (element.type === Fragment) {
-      const fragmentChildren = element.props.children;
-      React.Children.forEach(fragmentChildren, (child) =>
+      React.Children.forEach(element.props.children, (child) =>
         visitRowChild(child, parentProps),
       );
       return;
diff --git a/packages/cli/src/ui/contexts/SessionContext.test.tsx b/packages/cli/src/ui/contexts/SessionContext.test.tsx
index e9fc33e6..fedb5341 100644
--- a/packages/cli/src/ui/contexts/SessionContext.test.tsx
+++ b/packages/cli/src/ui/contexts/SessionContext.test.tsx
@@ -6,6 +6,7 @@
 
 import { type MutableRefObject } from 'react';
 import { render } from 'ink-testing-library';
+import { renderHook } from '@testing-library/react';
 import { act } from 'react-dom/test-utils';
 import { SessionStatsProvider, useSessionStats } from './SessionContext.js';
 import { describe, it, expect, vi } from 'vitest';
@@ -223,21 +224,16 @@ describe('SessionStatsContext', () => {
   });
 
   it('should throw an error when useSessionStats is used outside of a provider', () => {
-    // Suppress the expected console error during this test.
-    const errorSpy = vi.spyOn(console, 'error').mockImplementation(() => {});
+    // Suppress console.error for this test since we expect an error
+    const consoleSpy = vi.spyOn(console, 'error').mockImplementation(() => {});
 
-    const contextRef = { current: undefined };
-
-    // We expect rendering to fail, which React will catch and log as an error.
-    render(<TestHarness contextRef={contextRef} />);
-
-    // Assert that the first argument of the first call to console.error
-    // contains the expected message. This is more robust than checking
-    // the exact arguments, which can be affected by React/JSDOM internals.
-    expect(errorSpy.mock.calls[0][0]).toContain(
-      'useSessionStats must be used within a SessionStatsProvider',
-    );
-
-    errorSpy.mockRestore();
+    try {
+      // Expect renderHook itself to throw when the hook is used outside a provider
+      expect(() => {
+        renderHook(() => useSessionStats());
+      }).toThrow('useSessionStats must be used within a SessionStatsProvider');
+    } finally {
+      consoleSpy.mockRestore();
+    }
   });
 });
diff --git a/packages/cli/src/ui/hooks/useConsoleMessages.ts b/packages/cli/src/ui/hooks/useConsoleMessages.ts
index 90dc6f81..52ffbd39 100644
--- a/packages/cli/src/ui/hooks/useConsoleMessages.ts
+++ b/packages/cli/src/ui/hooks/useConsoleMessages.ts
@@ -25,9 +25,12 @@ export function useConsoleMessages(): UseConsoleMessagesReturn {
       return;
     }
 
+    const newMessagesToAdd = messageQueueRef.current;
+    messageQueueRef.current = [];
+
     setConsoleMessages((prevMessages) => {
       const newMessages = [...prevMessages];
-      messageQueueRef.current.forEach((queuedMessage) => {
+      newMessagesToAdd.forEach((queuedMessage) => {
         if (
           newMessages.length > 0 &&
           newMessages[newMessages.length - 1].type === queuedMessage.type &&
@@ -42,7 +45,6 @@ export function useConsoleMessages(): UseConsoleMessagesReturn {
       return newMessages;
     });
 
-    messageQueueRef.current = [];
     messageQueueTimeoutRef.current = null; // Allow next scheduling
   }, []);
 
diff --git a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx
index 5ac9aaa7..0c8b261e 100644
--- a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx
+++ b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx
@@ -496,13 +496,17 @@ describe('useGeminiStream', () => {
       } as TrackedCompletedToolCall, // Treat error as a form of completion for submission
     ];
 
-    // 1. On the first render, there are no tool calls.
-    mockUseReactToolScheduler.mockReturnValue([
-      [],
-      mockScheduleToolCalls,
-      mockMarkToolsAsSubmitted,
-    ]);
-    const { rerender } = renderHook(() =>
+    // Capture the onComplete callback
+    let capturedOnComplete:
+      | ((completedTools: TrackedToolCall[]) => Promise<void>)
+      | null = null;
+
+    mockUseReactToolScheduler.mockImplementation((onComplete) => {
+      capturedOnComplete = onComplete;
+      return [[], mockScheduleToolCalls, mockMarkToolsAsSubmitted];
+    });
+
+    renderHook(() =>
       useGeminiStream(
         new MockedGeminiClientClass(mockConfig),
         [],
@@ -518,16 +522,11 @@ describe('useGeminiStream', () => {
       ),
     );
 
-    // 2. Before the second render, change the mock to return the completed tools.
-    mockUseReactToolScheduler.mockReturnValue([
-      completedToolCalls,
-      mockScheduleToolCalls,
-      mockMarkToolsAsSubmitted,
-    ]);
-
-    // 3. Trigger a re-render. The hook will now receive the completed tools, causing the effect to run.
-    act(() => {
-      rerender();
+    // Trigger the onComplete callback with completed tools
+    await act(async () => {
+      if (capturedOnComplete) {
+        await capturedOnComplete(completedToolCalls);
+      }
     });
 
     await waitFor(() => {
@@ -561,13 +560,17 @@ describe('useGeminiStream', () => {
     ];
     const client = new MockedGeminiClientClass(mockConfig);
 
-    // 1. First render: no tool calls.
-    mockUseReactToolScheduler.mockReturnValue([
-      [],
-      mockScheduleToolCalls,
-      mockMarkToolsAsSubmitted,
-    ]);
-    const { rerender } = renderHook(() =>
+    // Capture the onComplete callback
+    let capturedOnComplete:
+      | ((completedTools: TrackedToolCall[]) => Promise<void>)
+      | null = null;
+
+    mockUseReactToolScheduler.mockImplementation((onComplete) => {
+      capturedOnComplete = onComplete;
+      return [[], mockScheduleToolCalls, mockMarkToolsAsSubmitted];
+    });
+
+    renderHook(() =>
       useGeminiStream(
         client,
         [],
@@ -583,16 +586,11 @@ describe('useGeminiStream', () => {
       ),
     );
 
-    // 2. Second render: tool calls are now cancelled.
-    mockUseReactToolScheduler.mockReturnValue([
-      cancelledToolCalls,
-      mockScheduleToolCalls,
-      mockMarkToolsAsSubmitted,
-    ]);
-
-    // 3. Trigger the re-render.
-    act(() => {
-      rerender();
+    // Trigger the onComplete callback with cancelled tools
+    await act(async () => {
+      if (capturedOnComplete) {
+        await capturedOnComplete(cancelledToolCalls);
+      }
     });
 
     await waitFor(() => {
@@ -685,7 +683,12 @@ describe('useGeminiStream', () => {
 
     const initialToolCalls: TrackedToolCall[] = [
       {
-        request: { callId: 'call1', name: 'tool1', args: {} },
+        request: {
+          callId: 'call1',
+          name: 'tool1',
+          args: {},
+          isClientInitiated: false,
+        },
         status: 'executing',
         responseSubmittedToGemini: false,
         tool: {
@@ -711,36 +714,67 @@ describe('useGeminiStream', () => {
       } as TrackedCompletedToolCall,
     ];
 
-    const { result, rerender, client } = renderTestHook(initialToolCalls);
+    // Capture the onComplete callback
+    let capturedOnComplete:
+      | ((completedTools: TrackedToolCall[]) => Promise<void>)
+      | null = null;
+    let currentToolCalls = initialToolCalls;
+
+    mockUseReactToolScheduler.mockImplementation((onComplete) => {
+      capturedOnComplete = onComplete;
+      return [
+        currentToolCalls,
+        mockScheduleToolCalls,
+        mockMarkToolsAsSubmitted,
+      ];
+    });
+
+    const { result, rerender } = renderHook(() =>
+      useGeminiStream(
+        new MockedGeminiClientClass(mockConfig),
+        [],
+        mockAddItem,
+        mockSetShowHelp,
+        mockConfig,
+        mockOnDebugMessage,
+        mockHandleSlashCommand,
+        false,
+        () => 'vscode' as EditorType,
+        () => {},
+        () => Promise.resolve(),
+      ),
+    );
 
     // 1. Initial state should be Responding because a tool is executing.
     expect(result.current.streamingState).toBe(StreamingState.Responding);
 
-    // 2. Rerender with the completed tool call.
-    // The useEffect should pick this up but hasn't called submitQuery yet.
+    // 2. Update the tool calls to completed state and rerender
+    currentToolCalls = completedToolCalls;
+    mockUseReactToolScheduler.mockImplementation((onComplete) => {
+      capturedOnComplete = onComplete;
+      return [
+        completedToolCalls,
+        mockScheduleToolCalls,
+        mockMarkToolsAsSubmitted,
+      ];
+    });
+
     act(() => {
-      rerender({
-        client,
-        history: [],
-        addItem: mockAddItem,
-        setShowHelp: mockSetShowHelp,
-        config: mockConfig,
-        onDebugMessage: mockOnDebugMessage,
-        handleSlashCommand:
-          mockHandleSlashCommand as unknown as typeof mockHandleSlashCommand,
-        shellModeActive: false,
-        loadedSettings: mockLoadedSettings,
-        // This is the key part of the test: update the toolCalls array
-        // to simulate the tool finishing.
-        toolCalls: completedToolCalls,
-      });
+      rerender();
     });
 
     // 3. The state should *still* be Responding, not Idle.
     // This is because the completed tool's response has not been submitted yet.
     expect(result.current.streamingState).toBe(StreamingState.Responding);
 
-    // 4. Wait for the useEffect to call submitQuery.
+    // 4. Trigger the onComplete callback to simulate tool completion
+    await act(async () => {
+      if (capturedOnComplete) {
+        await capturedOnComplete(completedToolCalls);
+      }
+    });
+
+    // 5. Wait for submitQuery to be called
     await waitFor(() => {
       expect(mockSendMessageStream).toHaveBeenCalledWith(
         toolCallResponseParts,
@@ -748,7 +782,7 @@ describe('useGeminiStream', () => {
       );
     });
 
-    // 5. After submission, the state should remain Responding.
+    // 6. After submission, the state should remain Responding until the stream completes.
     expect(result.current.streamingState).toBe(StreamingState.Responding);
   });
 
@@ -929,14 +963,17 @@ describe('useGeminiStream', () => {
         } as any,
       };
 
-      // 1. Initial render state: no tool calls
-      mockUseReactToolScheduler.mockReturnValue([
-        [],
-        mockScheduleToolCalls,
-        mockMarkToolsAsSubmitted,
-      ]);
+      // Capture the onComplete callback
+      let capturedOnComplete:
+        | ((completedTools: TrackedToolCall[]) => Promise<void>)
+        | null = null;
 
-      const { result, rerender } = renderHook(() =>
+      mockUseReactToolScheduler.mockImplementation((onComplete) => {
+        capturedOnComplete = onComplete;
+        return [[], mockScheduleToolCalls, mockMarkToolsAsSubmitted];
+      });
+
+      const { result } = renderHook(() =>
         useGeminiStream(
           new MockedGeminiClientClass(mockConfig),
           [],
@@ -957,17 +994,11 @@ describe('useGeminiStream', () => {
         await result.current.submitQuery('/memory add "test fact"');
       });
 
-      // The command handler schedules the tool. Now we simulate the tool completing.
-      // 2. Before the next render, set the mock to return the completed tool.
-      mockUseReactToolScheduler.mockReturnValue([
-        [completedToolCall],
-        mockScheduleToolCalls,
-        mockMarkToolsAsSubmitted,
-      ]);
-
-      // 3. Trigger a re-render to process the completed tool.
-      act(() => {
-        rerender();
+      // Trigger the onComplete callback with the completed client-initiated tool
+      await act(async () => {
+        if (capturedOnComplete) {
+          await capturedOnComplete([completedToolCall]);
+        }
       });
 
       // --- Assert the outcome ---
@@ -1007,13 +1038,17 @@ describe('useGeminiStream', () => {
         } as any,
       };
 
-      mockUseReactToolScheduler.mockReturnValue([
-        [completedToolCall],
-        mockScheduleToolCalls,
-        mockMarkToolsAsSubmitted,
-      ]);
+      // Capture the onComplete callback
+      let capturedOnComplete:
+        | ((completedTools: TrackedToolCall[]) => Promise<void>)
+        | null = null;
 
-      const { rerender } = renderHook(() =>
+      mockUseReactToolScheduler.mockImplementation((onComplete) => {
+        capturedOnComplete = onComplete;
+        return [[], mockScheduleToolCalls, mockMarkToolsAsSubmitted];
+      });
+
+      renderHook(() =>
         useGeminiStream(
           new MockedGeminiClientClass(mockConfig),
           [],
@@ -1029,8 +1064,11 @@ describe('useGeminiStream', () => {
         ),
       );
 
-      act(() => {
-        rerender();
+      // Trigger the onComplete callback with the completed save_memory tool
+      await act(async () => {
+        if (capturedOnComplete) {
+          await capturedOnComplete([completedToolCall]);
+        }
       });
 
       await waitFor(() => {
diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts
index ddd7c185..3d24ede7 100644
--- a/packages/cli/src/ui/hooks/useGeminiStream.ts
+++ b/packages/cli/src/ui/hooks/useGeminiStream.ts
@@ -111,17 +111,21 @@ export const useGeminiStream = (
 
   const [toolCalls, scheduleToolCalls, markToolsAsSubmitted] =
     useReactToolScheduler(
-      (completedToolCallsFromScheduler) => {
+      async (completedToolCallsFromScheduler) => {
         // This onComplete is called when ALL scheduled tools for a given batch are done.
         if (completedToolCallsFromScheduler.length > 0) {
           // Add the final state of these tools to the history for display.
-          // The new useEffect will handle submitting their responses.
           addItem(
             mapTrackedToolCallsToDisplay(
               completedToolCallsFromScheduler as TrackedToolCall[],
             ),
             Date.now(),
           );
+
+          // Handle tool response submission immediately when tools complete
+          await handleCompletedTools(
+            completedToolCallsFromScheduler as TrackedToolCall[],
+          );
         }
       },
       config,
@@ -570,40 +574,33 @@ export const useGeminiStream = (
     ],
   );
 
-  /**
-   * Automatically submits responses for completed tool calls.
-   * This effect runs when `toolCalls` or `isResponding` changes.
-   * It ensures that tool responses are sent back to Gemini only when
-   * all processing for a given set of tools is finished and Gemini
-   * is not already generating a response.
-   */
-  useEffect(() => {
-    const run = async () => {
+  const handleCompletedTools = useCallback(
+    async (completedToolCallsFromScheduler: TrackedToolCall[]) => {
       if (isResponding) {
         return;
       }
 
-      const completedAndReadyToSubmitTools = toolCalls.filter(
-        (
-          tc: TrackedToolCall,
-        ): tc is TrackedCompletedToolCall | TrackedCancelledToolCall => {
-          const isTerminalState =
-            tc.status === 'success' ||
-            tc.status === 'error' ||
-            tc.status === 'cancelled';
+      const completedAndReadyToSubmitTools =
+        completedToolCallsFromScheduler.filter(
+          (
+            tc: TrackedToolCall,
+          ): tc is TrackedCompletedToolCall | TrackedCancelledToolCall => {
+            const isTerminalState =
+              tc.status === 'success' ||
+              tc.status === 'error' ||
+              tc.status === 'cancelled';
 
-          if (isTerminalState) {
-            const completedOrCancelledCall = tc as
-              | TrackedCompletedToolCall
-              | TrackedCancelledToolCall;
-            return (
-              !completedOrCancelledCall.responseSubmittedToGemini &&
-              completedOrCancelledCall.response?.responseParts !== undefined
-            );
-          }
-          return false;
-        },
-      );
+            if (isTerminalState) {
+              const completedOrCancelledCall = tc as
+                | TrackedCompletedToolCall
+                | TrackedCancelledToolCall;
+              return (
+                completedOrCancelledCall.response?.responseParts !== undefined
+              );
+            }
+            return false;
+          },
+        );
 
       // Finalize any client-initiated tools as soon as they are done.
       const clientTools = completedAndReadyToSubmitTools.filter(
@@ -630,15 +627,6 @@ export const useGeminiStream = (
         );
       }
 
-      // Only proceed with submitting to Gemini if ALL tools are complete.
-      const allToolsAreComplete =
-        toolCalls.length > 0 &&
-        toolCalls.length === completedAndReadyToSubmitTools.length;
-
-      if (!allToolsAreComplete) {
-        return;
-      }
-
       const geminiTools = completedAndReadyToSubmitTools.filter(
         (t) => !t.request.isClientInitiated,
       );
@@ -693,17 +681,15 @@ export const useGeminiStream = (
       submitQuery(mergePartListUnions(responsesToSend), {
         isContinuation: true,
       });
-    };
-    void run();
-  }, [
-    toolCalls,
-    isResponding,
-    submitQuery,
-    markToolsAsSubmitted,
-    addItem,
-    geminiClient,
-    performMemoryRefresh,
-  ]);
+    },
+    [
+      isResponding,
+      submitQuery,
+      markToolsAsSubmitted,
+      geminiClient,
+      performMemoryRefresh,
+    ],
+  );
 
   const pendingHistoryItems = [
     pendingHistoryItemRef.current,
diff --git a/packages/cli/src/ui/hooks/useReactToolScheduler.ts b/packages/cli/src/ui/hooks/useReactToolScheduler.ts
index 5e1bbd25..22988fef 100644
--- a/packages/cli/src/ui/hooks/useReactToolScheduler.ts
+++ b/packages/cli/src/ui/hooks/useReactToolScheduler.ts
@@ -128,7 +128,7 @@ export function useReactToolScheduler(
         }),
       );
     },
-    [],
+    [setToolCallsForDisplay],
   );
 
   const scheduler = useMemo(
@@ -152,7 +152,7 @@ export function useReactToolScheduler(
   );
 
   const schedule: ScheduleFn = useCallback(
-    async (
+    (
       request: ToolCallRequestInfo | ToolCallRequestInfo[],
       signal: AbortSignal,
     ) => {