Fix race condition by awaiting scheduleToolCalls (#16759)

Co-authored-by: Tommaso Sciortino <sciortino@gmail.com>

Author: 2 people

packages/cli/src/ui/hooks/useGeminiStream.ts

@@ -444,7 +444,7 @@ export const useGeminiStream = (
                   isClientInitiated: true,
                   prompt_id,
                 };
-                scheduleToolCalls([toolCallRequest], abortSignal);
+                await scheduleToolCalls([toolCallRequest], abortSignal);
                 return { queryToSend: null, shouldProceed: false };
               }
               case 'submit_prompt': {
@@ -911,7 +911,7 @@ export const useGeminiStream = (
         }
       }
       if (toolCallRequests.length > 0) {
-        scheduleToolCalls(toolCallRequests, signal);
+        await scheduleToolCalls(toolCallRequests, signal);
       }
       return StreamProcessingStatus.Completed;
     },

packages/cli/src/ui/hooks/useReactToolScheduler.ts

@@ -31,7 +31,7 @@ import { ToolCallStatus } from '../types.js';
 export type ScheduleFn = (
   request: ToolCallRequestInfo | ToolCallRequestInfo[],
   signal: AbortSignal,
-) => void;
+) => Promise<void>;
 export type MarkToolsAsSubmittedFn = (callIds: string[]) => void;
 
 export type TrackedScheduledToolCall = ScheduledToolCall & {
@@ -180,7 +180,7 @@ export function useReactToolScheduler(
       signal: AbortSignal,
     ) => {
       setToolCallsForDisplay([]);
-      void scheduler.schedule(request, signal);
+      return scheduler.schedule(request, signal);
     },
     [scheduler, setToolCallsForDisplay],
   );

packages/cli/src/ui/hooks/useToolScheduler.test.ts

@@ -163,8 +163,8 @@ describe('useReactToolScheduler in YOLO Mode', () => {
       args: { data: 'any data' },
     } as any;
 
-    act(() => {
-      schedule(request, new AbortController().signal);
+    await act(async () => {
+      await schedule(request, new AbortController().signal);
     });
 
     await act(async () => {
@@ -220,11 +220,11 @@ describe('useReactToolScheduler', () => {
     schedule: (
       req: ToolCallRequestInfo | ToolCallRequestInfo[],
       signal: AbortSignal,
-    ) => void,
+    ) => Promise<void>,
     request: ToolCallRequestInfo | ToolCallRequestInfo[],
   ) => {
-    act(() => {
-      schedule(request, new AbortController().signal);
+    await act(async () => {
+      await schedule(request, new AbortController().signal);
     });
 
     await advanceAndSettle();
@@ -313,10 +313,13 @@ describe('useReactToolScheduler', () => {
 
   it('should clear previous tool calls when scheduling new ones', async () => {
     mockToolRegistry.getTool.mockReturnValue(mockTool);
-    (mockTool.execute as Mock).mockResolvedValue({
-      llmContent: 'Tool output',
-      returnDisplay: 'Formatted tool output',
-    } as ToolResult);
+    (mockTool.execute as Mock).mockImplementation(async () => {
+      await new Promise((r) => setTimeout(r, 10));
+      return {
+        llmContent: 'Tool output',
+        returnDisplay: 'Formatted tool output',
+      };
+    });
 
     const { result } = renderScheduler();
     const schedule = result.current[1];
@@ -337,10 +340,13 @@ describe('useReactToolScheduler', () => {
       name: 'mockTool',
       args: {},
     } as any;
-    act(() => {
-      schedule(newRequest, new AbortController().signal);
+    let schedulePromise: Promise<void>;
+    await act(async () => {
+      schedulePromise = schedule(newRequest, new AbortController().signal);
     });
 
+    await advanceAndSettle();
+
     // After scheduling, the old call should be gone,
     // and the new one should be in the display in its initial state.
     expect(result.current[0].length).toBe(1);
@@ -349,14 +355,13 @@ describe('useReactToolScheduler', () => {
 
     // Let the new call finish.
     await act(async () => {
-      await vi.advanceTimersByTimeAsync(0);
+      await vi.advanceTimersByTimeAsync(20);
     });
+
     await act(async () => {
-      await vi.advanceTimersByTimeAsync(0);
-    });
-    await act(async () => {
-      await vi.advanceTimersByTimeAsync(0);
+      await schedulePromise;
     });
+
     expect(onComplete).toHaveBeenCalled();
   });
 
@@ -379,16 +384,14 @@ describe('useReactToolScheduler', () => {
       args: {},
     } as any;
 
-    act(() => {
-      schedule(request, new AbortController().signal);
-    });
+    let schedulePromise: Promise<void>;
     await act(async () => {
-      await vi.advanceTimersByTimeAsync(0);
-    }); // validation
-    await act(async () => {
-      await vi.advanceTimersByTimeAsync(0); // Process scheduling
+      schedulePromise = schedule(request, new AbortController().signal);
     });
 
+    await advanceAndSettle(); // validation
+    await advanceAndSettle(); // Process scheduling
+
     // At this point, the tool is 'executing' and waiting on the promise.
     expect(result.current[0][0].status).toBe('executing');
 
@@ -397,9 +400,7 @@ describe('useReactToolScheduler', () => {
       cancelAllToolCalls(cancelController.signal);
     });
 
-    await act(async () => {
-      await vi.advanceTimersByTimeAsync(0);
-    });
+    await advanceAndSettle();
 
     expect(onComplete).toHaveBeenCalledWith([
       expect.objectContaining({
@@ -412,6 +413,11 @@ describe('useReactToolScheduler', () => {
     await act(async () => {
       resolveExecute({ llmContent: 'output', returnDisplay: 'display' });
     });
+
+    // Now await the schedule promise
+    await act(async () => {
+      await schedulePromise;
+    });
   });
 
   it.each([
@@ -511,8 +517,9 @@ describe('useReactToolScheduler', () => {
       args: { data: 'sensitive' },
     } as any;
 
-    act(() => {
-      schedule(request, new AbortController().signal);
+    let schedulePromise: Promise<void>;
+    await act(async () => {
+      schedulePromise = schedule(request, new AbortController().signal);
     });
     await advanceAndSettle();
 
@@ -525,10 +532,13 @@ describe('useReactToolScheduler', () => {
       await capturedOnConfirmForTest?.(ToolConfirmationOutcome.ProceedOnce);
     });
 
-    await advanceAndSettle();
-    await advanceAndSettle();
     await advanceAndSettle();
 
+    // Now await the schedule promise as it should complete
+    await act(async () => {
+      await schedulePromise;
+    });
+
     expect(mockOnUserConfirmForToolConfirmation).toHaveBeenCalledWith(
       ToolConfirmationOutcome.ProceedOnce,
     );
@@ -558,8 +568,9 @@ describe('useReactToolScheduler', () => {
       args: {},
     } as any;
 
-    act(() => {
-      schedule(request, new AbortController().signal);
+    let schedulePromise: Promise<void>;
+    await act(async () => {
+      schedulePromise = schedule(request, new AbortController().signal);
     });
     await advanceAndSettle();
 
@@ -571,8 +582,13 @@ describe('useReactToolScheduler', () => {
     await act(async () => {
       await capturedOnConfirmForTest?.(ToolConfirmationOutcome.Cancel);
     });
+
     await advanceAndSettle();
-    await advanceAndSettle();
+
+    // Now await the schedule promise
+    await act(async () => {
+      await schedulePromise;
+    });
 
     expect(mockOnUserConfirmForToolConfirmation).toHaveBeenCalledWith(
       ToolConfirmationOutcome.Cancel,
@@ -619,8 +635,12 @@ describe('useReactToolScheduler', () => {
       args: {},
     } as any;
 
-    act(() => {
-      result.current[1](request, new AbortController().signal);
+    let schedulePromise: Promise<void>;
+    await act(async () => {
+      schedulePromise = result.current[1](
+        request,
+        new AbortController().signal,
+      );
     });
     await advanceAndSettle();
 
@@ -644,7 +664,11 @@ describe('useReactToolScheduler', () => {
       } as ToolResult);
     });
     await advanceAndSettle();
-    await advanceAndSettle();
+
+    // Now await schedule
+    await act(async () => {
+      await schedulePromise;
+    });
 
     const completedCalls = onComplete.mock.calls[0][0] as ToolCall[];
     expect(completedCalls[0].status).toBe('success');
@@ -690,8 +714,8 @@ describe('useReactToolScheduler', () => {
       { callId: 'multi2', name: 'tool2', args: { p: 2 } } as any,
     ];
 
-    act(() => {
-      schedule(requests, new AbortController().signal);
+    await act(async () => {
+      await schedule(requests, new AbortController().signal);
     });
     await act(async () => {
       await vi.advanceTimersByTimeAsync(0);
@@ -782,38 +806,48 @@ describe('useReactToolScheduler', () => {
       args: {},
     } as any;
 
-    act(() => {
-      schedule(request1, new AbortController().signal);
+    let schedulePromise1: Promise<void>;
+    let schedulePromise2: Promise<void>;
+
+    await act(async () => {
+      schedulePromise1 = schedule(request1, new AbortController().signal);
     });
     await act(async () => {
       await vi.advanceTimersByTimeAsync(0);
     });
 
-    act(() => {
-      schedule(request2, new AbortController().signal);
+    await act(async () => {
+      schedulePromise2 = schedule(request2, new AbortController().signal);
     });
 
     await act(async () => {
       await vi.advanceTimersByTimeAsync(50);
       await vi.advanceTimersByTimeAsync(0);
-      await act(async () => {
-        await vi.advanceTimersByTimeAsync(0);
-      });
     });
+
+    // Wait for first to complete
+    await act(async () => {
+      await schedulePromise1;
+    });
+
     expect(onComplete).toHaveBeenCalledWith([
       expect.objectContaining({
         status: 'success',
         request: request1,
         response: expect.objectContaining({ resultDisplay: 'done display' }),
       }),
     ]);
+
     await act(async () => {
       await vi.advanceTimersByTimeAsync(50);
       await vi.advanceTimersByTimeAsync(0);
-      await act(async () => {
-        await vi.advanceTimersByTimeAsync(0);
-      });
     });
+
+    // Wait for second to complete
+    await act(async () => {
+      await schedulePromise2;
+    });
+
     expect(onComplete).toHaveBeenCalledWith([
       expect.objectContaining({
         status: 'success',