Upgrade to Ink 6 and React 19 (#2096)

Co-authored-by: jacob314 <jacob314@gmail.com>
2025-12-20 08:47:44 +00:00 · 2025-06-27 16:39:54 -07:00
parent 19d2a0fb35
commit 150df382f8
18 changed files with 1129 additions and 1571 deletions
--- 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
  }, []);

--- 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(() => {
--- 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,
--- 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,
    ) => {