Add a request queue to the tool scheduler (#5845)

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

@@ -63,7 +63,7 @@ export type TrackedToolCall =
   | TrackedCancelledToolCall;
 
 export function useReactToolScheduler(
-  onComplete: (tools: CompletedToolCall[]) => void,
+  onComplete: (tools: CompletedToolCall[]) => Promise<void>,
   config: Config,
   setPendingHistoryItem: React.Dispatch<
     React.SetStateAction<HistoryItemWithoutId | null>
@@ -106,8 +106,8 @@ export function useReactToolScheduler(
   );
 
   const allToolCallsCompleteHandler: AllToolCallsCompleteHandler = useCallback(
-    (completedToolCalls) => {
-      onComplete(completedToolCalls);
+    async (completedToolCalls) => {
+      await onComplete(completedToolCalls);
     },
     [onComplete],
   );
@@ -157,7 +157,7 @@ export function useReactToolScheduler(
       request: ToolCallRequestInfo | ToolCallRequestInfo[],
       signal: AbortSignal,
     ) => {
-      scheduler.schedule(request, signal);
+      void scheduler.schedule(request, signal);
     },
     [scheduler],
   );

packages/core/src/core/coreToolScheduler.test.ts

@@ -592,3 +592,195 @@ describe('CoreToolScheduler YOLO mode', () => {
     }
   });
 });
+
+describe('CoreToolScheduler request queueing', () => {
+  it('should queue a request if another is running', async () => {
+    let resolveFirstCall: (result: ToolResult) => void;
+    const firstCallPromise = new Promise<ToolResult>((resolve) => {
+      resolveFirstCall = resolve;
+    });
+
+    const mockTool = new MockTool();
+    mockTool.executeFn.mockImplementation(() => firstCallPromise);
+    const declarativeTool = mockTool;
+
+    const toolRegistry = {
+      getTool: () => declarativeTool,
+      getToolByName: () => declarativeTool,
+      getFunctionDeclarations: () => [],
+      tools: new Map(),
+      discovery: {} as any,
+      registerTool: () => {},
+      getToolByDisplayName: () => declarativeTool,
+      getTools: () => [],
+      discoverTools: async () => {},
+      getAllTools: () => [],
+      getToolsByServer: () => [],
+    };
+
+    const onAllToolCallsComplete = vi.fn();
+    const onToolCallsUpdate = vi.fn();
+
+    const mockConfig = {
+      getSessionId: () => 'test-session-id',
+      getUsageStatisticsEnabled: () => true,
+      getDebugMode: () => false,
+      getApprovalMode: () => ApprovalMode.YOLO, // Use YOLO to avoid confirmation prompts
+    } as unknown as Config;
+
+    const scheduler = new CoreToolScheduler({
+      config: mockConfig,
+      toolRegistry: Promise.resolve(toolRegistry as any),
+      onAllToolCallsComplete,
+      onToolCallsUpdate,
+      getPreferredEditor: () => 'vscode',
+      onEditorClose: vi.fn(),
+    });
+
+    const abortController = new AbortController();
+    const request1 = {
+      callId: '1',
+      name: 'mockTool',
+      args: { a: 1 },
+      isClientInitiated: false,
+      prompt_id: 'prompt-1',
+    };
+    const request2 = {
+      callId: '2',
+      name: 'mockTool',
+      args: { b: 2 },
+      isClientInitiated: false,
+      prompt_id: 'prompt-2',
+    };
+
+    // Schedule the first call, which will pause execution.
+    scheduler.schedule([request1], abortController.signal);
+
+    // Wait for the first call to be in the 'executing' state.
+    await vi.waitFor(() => {
+      const calls = onToolCallsUpdate.mock.calls.at(-1)?.[0] as ToolCall[];
+      expect(calls?.[0]?.status).toBe('executing');
+    });
+
+    // Schedule the second call while the first is "running".
+    const schedulePromise2 = scheduler.schedule(
+      [request2],
+      abortController.signal,
+    );
+
+    // Ensure the second tool call hasn't been executed yet.
+    expect(mockTool.executeFn).toHaveBeenCalledTimes(1);
+    expect(mockTool.executeFn).toHaveBeenCalledWith({ a: 1 });
+
+    // Complete the first tool call.
+    resolveFirstCall!({
+      llmContent: 'First call complete',
+      returnDisplay: 'First call complete',
+    });
+
+    // Wait for the second schedule promise to resolve.
+    await schedulePromise2;
+
+    // Wait for the second call to be in the 'executing' state.
+    await vi.waitFor(() => {
+      const calls = onToolCallsUpdate.mock.calls.at(-1)?.[0] as ToolCall[];
+      expect(calls?.[0]?.status).toBe('executing');
+    });
+
+    // Now the second tool call should have been executed.
+    expect(mockTool.executeFn).toHaveBeenCalledTimes(2);
+    expect(mockTool.executeFn).toHaveBeenCalledWith({ b: 2 });
+
+    // Let the second call finish.
+    const secondCallResult = {
+      llmContent: 'Second call complete',
+      returnDisplay: 'Second call complete',
+    };
+    // Since the mock is shared, we need to resolve the current promise.
+    // In a real scenario, a new promise would be created for the second call.
+    resolveFirstCall!(secondCallResult);
+
+    // Wait for the second completion.
+    await vi.waitFor(() => {
+      expect(onAllToolCallsComplete).toHaveBeenCalledTimes(2);
+    });
+
+    // Verify the completion callbacks were called correctly.
+    expect(onAllToolCallsComplete.mock.calls[0][0][0].status).toBe('success');
+    expect(onAllToolCallsComplete.mock.calls[1][0][0].status).toBe('success');
+  });
+
+  it('should handle two synchronous calls to schedule', async () => {
+    const mockTool = new MockTool();
+    const declarativeTool = mockTool;
+    const toolRegistry = {
+      getTool: () => declarativeTool,
+      getToolByName: () => declarativeTool,
+      getFunctionDeclarations: () => [],
+      tools: new Map(),
+      discovery: {} as any,
+      registerTool: () => {},
+      getToolByDisplayName: () => declarativeTool,
+      getTools: () => [],
+      discoverTools: async () => {},
+      getAllTools: () => [],
+      getToolsByServer: () => [],
+    };
+
+    const onAllToolCallsComplete = vi.fn();
+    const onToolCallsUpdate = vi.fn();
+
+    const mockConfig = {
+      getSessionId: () => 'test-session-id',
+      getUsageStatisticsEnabled: () => true,
+      getDebugMode: () => false,
+      getApprovalMode: () => ApprovalMode.YOLO,
+    } as unknown as Config;
+
+    const scheduler = new CoreToolScheduler({
+      config: mockConfig,
+      toolRegistry: Promise.resolve(toolRegistry as any),
+      onAllToolCallsComplete,
+      onToolCallsUpdate,
+      getPreferredEditor: () => 'vscode',
+      onEditorClose: vi.fn(),
+    });
+
+    const abortController = new AbortController();
+    const request1 = {
+      callId: '1',
+      name: 'mockTool',
+      args: { a: 1 },
+      isClientInitiated: false,
+      prompt_id: 'prompt-1',
+    };
+    const request2 = {
+      callId: '2',
+      name: 'mockTool',
+      args: { b: 2 },
+      isClientInitiated: false,
+      prompt_id: 'prompt-2',
+    };
+
+    // Schedule two calls synchronously.
+    const schedulePromise1 = scheduler.schedule(
+      [request1],
+      abortController.signal,
+    );
+    const schedulePromise2 = scheduler.schedule(
+      [request2],
+      abortController.signal,
+    );
+
+    // Wait for both promises to resolve.
+    await Promise.all([schedulePromise1, schedulePromise2]);
+
+    // Ensure the tool was called twice with the correct arguments.
+    expect(mockTool.executeFn).toHaveBeenCalledTimes(2);
+    expect(mockTool.executeFn).toHaveBeenCalledWith({ a: 1 });
+    expect(mockTool.executeFn).toHaveBeenCalledWith({ b: 2 });
+
+    // Ensure completion callbacks were called twice.
+    expect(onAllToolCallsComplete).toHaveBeenCalledTimes(2);
+  });
+});

packages/core/src/core/coreToolScheduler.ts

@@ -125,7 +125,7 @@ export type OutputUpdateHandler = (
 
 export type AllToolCallsCompleteHandler = (
   completedToolCalls: CompletedToolCall[],
-) => void;
+) => Promise<void>;
 
 export type ToolCallsUpdateHandler = (toolCalls: ToolCall[]) => void;
 
@@ -244,6 +244,14 @@ export class CoreToolScheduler {
   private getPreferredEditor: () => EditorType | undefined;
   private config: Config;
   private onEditorClose: () => void;
+  private isFinalizingToolCalls = false;
+  private isScheduling = false;
+  private requestQueue: Array<{
+    request: ToolCallRequestInfo | ToolCallRequestInfo[];
+    signal: AbortSignal;
+    resolve: () => void;
+    reject: (reason?: Error) => void;
+  }> = [];
 
   constructor(options: CoreToolSchedulerOptions) {
     this.config = options.config;
@@ -455,9 +463,12 @@ export class CoreToolScheduler {
   }
 
   private isRunning(): boolean {
-    return this.toolCalls.some(
+    return (
+      this.isFinalizingToolCalls ||
+      this.toolCalls.some(
         (call) =>
           call.status === 'executing' || call.status === 'awaiting_approval',
+      )
     );
   }
 
@@ -475,10 +486,48 @@ export class CoreToolScheduler {
     }
   }
 
-  async schedule(
+  schedule(
     request: ToolCallRequestInfo | ToolCallRequestInfo[],
     signal: AbortSignal,
   ): Promise<void> {
+    if (this.isRunning() || this.isScheduling) {
+      return new Promise((resolve, reject) => {
+        const abortHandler = () => {
+          // Find and remove the request from the queue
+          const index = this.requestQueue.findIndex(
+            (item) => item.request === request,
+          );
+          if (index > -1) {
+            this.requestQueue.splice(index, 1);
+            reject(new Error('Tool call cancelled while in queue.'));
+          }
+        };
+
+        signal.addEventListener('abort', abortHandler, { once: true });
+
+        this.requestQueue.push({
+          request,
+          signal,
+          resolve: () => {
+            signal.removeEventListener('abort', abortHandler);
+            resolve();
+          },
+          reject: (reason?: Error) => {
+            signal.removeEventListener('abort', abortHandler);
+            reject(reason);
+          },
+        });
+      });
+    }
+    return this._schedule(request, signal);
+  }
+
+  private async _schedule(
+    request: ToolCallRequestInfo | ToolCallRequestInfo[],
+    signal: AbortSignal,
+  ): Promise<void> {
+    this.isScheduling = true;
+    try {
       if (this.isRunning()) {
         throw new Error(
           'Cannot schedule new tool calls while other tool calls are actively running (executing or awaiting approval).',
@@ -618,7 +667,10 @@ export class CoreToolScheduler {
         }
       }
       this.attemptExecutionOfScheduledCalls(signal);
-    this.checkAndNotifyCompletion();
+      void this.checkAndNotifyCompletion();
+    } finally {
+      this.isScheduling = false;
+    }
   }
 
   async handleConfirmationResponse(
@@ -822,7 +874,7 @@ export class CoreToolScheduler {
     }
   }
 
-  private checkAndNotifyCompletion(): void {
+  private async checkAndNotifyCompletion(): Promise<void> {
     const allCallsAreTerminal = this.toolCalls.every(
       (call) =>
         call.status === 'success' ||
@@ -839,9 +891,18 @@ export class CoreToolScheduler {
       }
 
       if (this.onAllToolCallsComplete) {
-        this.onAllToolCallsComplete(completedCalls);
+        this.isFinalizingToolCalls = true;
+        await this.onAllToolCallsComplete(completedCalls);
+        this.isFinalizingToolCalls = false;
       }
       this.notifyToolCallsUpdate();
+      // After completion, process the next item in the queue.
+      if (this.requestQueue.length > 0) {
+        const next = this.requestQueue.shift()!;
+        this._schedule(next.request, next.signal)
+          .then(next.resolve)
+          .catch(next.reject);
+      }
     }
   }