From cfdc4cfca8162e99cdfd169b8765031847ef5296 Mon Sep 17 00:00:00 2001 From: christine betts Date: Fri, 16 Jan 2026 11:55:15 -0500 Subject: [PATCH] Fix race condition by awaiting scheduleToolCalls (#16759) Co-authored-by: Tommaso Sciortino --- packages/cli/src/ui/hooks/useGeminiStream.ts | 4 +- .../cli/src/ui/hooks/useReactToolScheduler.ts | 4 +- .../cli/src/ui/hooks/useToolScheduler.test.ts | 130 +++++++++++------- 3 files changed, 86 insertions(+), 52 deletions(-) diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index 2d3152bba8..02295dbb88 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -460,7 +460,7 @@ export const useGeminiStream = ( isClientInitiated: true, prompt_id, }; - scheduleToolCalls([toolCallRequest], abortSignal); + await scheduleToolCalls([toolCallRequest], abortSignal); return { queryToSend: null, shouldProceed: false }; } case 'submit_prompt': { @@ -923,7 +923,7 @@ export const useGeminiStream = ( } } if (toolCallRequests.length > 0) { - scheduleToolCalls(toolCallRequests, signal); + await scheduleToolCalls(toolCallRequests, signal); } return StreamProcessingStatus.Completed; }, diff --git a/packages/cli/src/ui/hooks/useReactToolScheduler.ts b/packages/cli/src/ui/hooks/useReactToolScheduler.ts index 8939f5aa80..53f345948a 100644 --- a/packages/cli/src/ui/hooks/useReactToolScheduler.ts +++ b/packages/cli/src/ui/hooks/useReactToolScheduler.ts @@ -32,7 +32,7 @@ import { ToolCallStatus } from '../types.js'; export type ScheduleFn = ( request: ToolCallRequestInfo | ToolCallRequestInfo[], signal: AbortSignal, -) => void; +) => Promise; export type MarkToolsAsSubmittedFn = (callIds: string[]) => void; export type TrackedScheduledToolCall = ScheduledToolCall & { @@ -181,7 +181,7 @@ export function useReactToolScheduler( signal: AbortSignal, ) => { setToolCallsForDisplay([]); - void scheduler.schedule(request, signal); + return scheduler.schedule(request, signal); }, [scheduler, setToolCallsForDisplay], ); diff --git a/packages/cli/src/ui/hooks/useToolScheduler.test.ts b/packages/cli/src/ui/hooks/useToolScheduler.test.ts index 1ffaa61cc7..a1fbe21dd3 100644 --- a/packages/cli/src/ui/hooks/useToolScheduler.test.ts +++ b/packages/cli/src/ui/hooks/useToolScheduler.test.ts @@ -172,8 +172,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 () => { @@ -229,11 +229,11 @@ describe('useReactToolScheduler', () => { schedule: ( req: ToolCallRequestInfo | ToolCallRequestInfo[], signal: AbortSignal, - ) => void, + ) => Promise, request: ToolCallRequestInfo | ToolCallRequestInfo[], ) => { - act(() => { - schedule(request, new AbortController().signal); + await act(async () => { + await schedule(request, new AbortController().signal); }); await advanceAndSettle(); @@ -322,10 +322,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]; @@ -346,10 +349,13 @@ describe('useReactToolScheduler', () => { name: 'mockTool', args: {}, } as any; - act(() => { - schedule(newRequest, new AbortController().signal); + let schedulePromise: Promise; + 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); @@ -358,14 +364,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(); }); @@ -388,16 +393,14 @@ describe('useReactToolScheduler', () => { args: {}, } as any; - act(() => { - schedule(request, new AbortController().signal); - }); + let schedulePromise: Promise; 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'); @@ -406,9 +409,7 @@ describe('useReactToolScheduler', () => { cancelAllToolCalls(cancelController.signal); }); - await act(async () => { - await vi.advanceTimersByTimeAsync(0); - }); + await advanceAndSettle(); expect(onComplete).toHaveBeenCalledWith([ expect.objectContaining({ @@ -421,6 +422,11 @@ describe('useReactToolScheduler', () => { await act(async () => { resolveExecute({ llmContent: 'output', returnDisplay: 'display' }); }); + + // Now await the schedule promise + await act(async () => { + await schedulePromise; + }); }); it.each([ @@ -520,8 +526,9 @@ describe('useReactToolScheduler', () => { args: { data: 'sensitive' }, } as any; - act(() => { - schedule(request, new AbortController().signal); + let schedulePromise: Promise; + await act(async () => { + schedulePromise = schedule(request, new AbortController().signal); }); await advanceAndSettle(); @@ -535,8 +542,11 @@ describe('useReactToolScheduler', () => { }); 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, @@ -567,8 +577,9 @@ describe('useReactToolScheduler', () => { args: {}, } as any; - act(() => { - schedule(request, new AbortController().signal); + let schedulePromise: Promise; + await act(async () => { + schedulePromise = schedule(request, new AbortController().signal); }); await advanceAndSettle(); @@ -580,8 +591,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, @@ -628,8 +644,12 @@ describe('useReactToolScheduler', () => { args: {}, } as any; - act(() => { - result.current[1](request, new AbortController().signal); + let schedulePromise: Promise; + await act(async () => { + schedulePromise = result.current[1]( + request, + new AbortController().signal, + ); }); await advanceAndSettle(); @@ -653,7 +673,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'); @@ -699,8 +723,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); @@ -791,24 +815,30 @@ describe('useReactToolScheduler', () => { args: {}, } as any; - act(() => { - schedule(request1, new AbortController().signal); + let schedulePromise1: Promise; + let schedulePromise2: Promise; + + 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', @@ -816,13 +846,17 @@ describe('useReactToolScheduler', () => { 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',