feat(core): implement model-driven parallel tool scheduler (#21933)

Author: abhipatel12

packages/core/src/agents/subagent-tool-wrapper.test.ts

@@ -103,9 +103,19 @@ describe('SubagentToolWrapper', () => {
 
       expect(schema.name).toBe(mockDefinition.name);
       expect(schema.description).toBe(mockDefinition.description);
-      expect(schema.parametersJsonSchema).toEqual(
-        mockDefinition.inputConfig.inputSchema,
-      );
+      expect(schema.parametersJsonSchema).toEqual({
+        ...(mockDefinition.inputConfig.inputSchema as Record<string, unknown>),
+        properties: {
+          ...((
+            mockDefinition.inputConfig.inputSchema as Record<string, unknown>
+          )['properties'] as Record<string, unknown>),
+          wait_for_previous: {
+            type: 'boolean',
+            description:
+              'Set to true to wait for all previously requested tools in this turn to complete before starting. Set to false (or omit) to run in parallel. Use true when this tool depends on the output of previous tools.',
+          },
+        },
+      });
     });
   });
 

packages/core/src/core/__snapshots__/prompts.test.ts.snap

@@ -355,7 +355,8 @@ export function renderOperationalGuidelines(
 - **Security First:** Always apply security best practices. Never introduce code that exposes, logs, or commits secrets, API keys, or other sensitive information.
 
 ## Tool Usage
-- **Parallelism:** Execute multiple independent tool calls in parallel when feasible (i.e. searching the codebase).
+- **Parallelism & Sequencing:** Tools execute in parallel by default. Execute multiple independent tool calls in parallel when feasible (e.g., searching, reading files, independent shell commands, or editing *different* files). If a tool depends on the output or side-effects of a previous tool in the same turn (e.g., running a shell command that depends on the success of a previous command), you MUST set the \`wait_for_previous\` parameter to \`true\` on the dependent tool to ensure sequential execution.
+- **File Editing Collisions:** Do NOT make multiple calls to the ${formatToolName(EDIT_TOOL_NAME)} tool for the SAME file in a single turn. To make multiple edits to the same file, you MUST perform them sequentially across multiple conversational turns to prevent race conditions and ensure the file state is accurate before each edit.
 - **Command Execution:** Use the ${formatToolName(SHELL_TOOL_NAME)} tool for running shell commands, remembering the safety rule to explain modifying commands first.${toolUsageInteractive(
     options.interactive,
     options.interactiveShellEnabled,

packages/core/src/prompts/snippets.ts

@@ -134,7 +134,7 @@ describe('Scheduler (Orchestrator)', () => {
   const req2: ToolCallRequestInfo = {
     callId: 'call-2',
     name: 'test-tool',
-    args: { foo: 'baz' },
+    args: { foo: 'baz', wait_for_previous: true },
     isClientInitiated: false,
     prompt_id: 'prompt-1',
     schedulerId: ROOT_SCHEDULER_ID,

packages/core/src/scheduler/scheduler.test.ts

@@ -29,7 +29,6 @@ import { PolicyDecision, type ApprovalMode } from '../policy/types.js';
 import {
   ToolConfirmationOutcome,
   type AnyDeclarativeTool,
-  Kind,
 } from '../tools/tools.js';
 import { getToolSuggestion } from '../utils/tool-utils.js';
 import { runInDevTraceSpan } from '../telemetry/trace.js';
@@ -434,10 +433,10 @@ export class Scheduler {
       }
 
       // If the first tool is parallelizable, batch all contiguous parallelizable tools.
-      if (this._isParallelizable(next.tool)) {
+      if (this._isParallelizable(next.request)) {
         while (this.state.queueLength > 0) {
           const peeked = this.state.peekQueue();
-          if (peeked && this._isParallelizable(peeked.tool)) {
+          if (peeked && this._isParallelizable(peeked.request)) {
             this.state.dequeue();
           } else {
             break;
@@ -522,9 +521,16 @@ export class Scheduler {
     return false;
   }
 
-  private _isParallelizable(tool?: AnyDeclarativeTool): boolean {
-    if (!tool) return false;
-    return tool.isReadOnly || tool.kind === Kind.Agent;
+  private _isParallelizable(request: ToolCallRequestInfo): boolean {
+    if (request.args) {
+      const wait = request.args['wait_for_previous'];
+      if (typeof wait === 'boolean') {
+        return !wait;
+      }
+    }
+
+    // Default to parallel if the flag is omitted.
+    return true;
   }
 
   private async _processValidatingCall(

packages/core/src/scheduler/scheduler.ts

@@ -119,7 +119,7 @@ describe('Scheduler Parallel Execution', () => {
   const req3: ToolCallRequestInfo = {
     callId: 'call-3',
     name: 'write-tool',
-    args: { path: 'c.txt', content: 'hi' },
+    args: { path: 'c.txt', content: 'hi', wait_for_previous: true },
     isClientInitiated: false,
     prompt_id: 'p1',
     schedulerId: ROOT_SCHEDULER_ID,
@@ -505,4 +505,50 @@ describe('Scheduler Parallel Execution', () => {
     const start1 = executionLog.indexOf('start-call-1');
     expect(start1).toBeGreaterThan(end3);
   });
+
+  it('should execute non-read-only tools in parallel if wait_for_previous is false', async () => {
+    const executionLog: string[] = [];
+    mockExecutor.execute.mockImplementation(async ({ call }) => {
+      const id = call.request.callId;
+      executionLog.push(`start-${id}`);
+      await new Promise<void>((resolve) => setTimeout(resolve, 10));
+      executionLog.push(`end-${id}`);
+      return {
+        status: 'success',
+        response: { callId: id, responseParts: [] },
+      } as unknown as SuccessfulToolCall;
+    });
+
+    const w1 = { ...req3, callId: 'w1', args: { wait_for_previous: false } };
+    const w2 = { ...req3, callId: 'w2', args: { wait_for_previous: false } };
+
+    await scheduler.schedule([w1, w2], signal);
+
+    expect(executionLog.slice(0, 2)).toContain('start-w1');
+    expect(executionLog.slice(0, 2)).toContain('start-w2');
+  });
+
+  it('should execute read-only tools sequentially if wait_for_previous is true', async () => {
+    const executionLog: string[] = [];
+    mockExecutor.execute.mockImplementation(async ({ call }) => {
+      const id = call.request.callId;
+      executionLog.push(`start-${id}`);
+      await new Promise<void>((resolve) => setTimeout(resolve, 10));
+      executionLog.push(`end-${id}`);
+      return {
+        status: 'success',
+        response: { callId: id, responseParts: [] },
+      } as unknown as SuccessfulToolCall;
+    });
+
+    const r1 = { ...req1, callId: 'r1', args: { wait_for_previous: false } };
+    const r2 = { ...req1, callId: 'r2', args: { wait_for_previous: true } };
+
+    await scheduler.schedule([r1, r2], signal);
+
+    expect(executionLog[0]).toBe('start-r1');
+    expect(executionLog[1]).toBe('end-r1');
+    expect(executionLog[2]).toBe('start-r2');
+    expect(executionLog[3]).toBe('end-r2');
+  });
 });

packages/core/src/scheduler/scheduler_parallel.test.ts

@@ -752,6 +752,11 @@ describe('mcp-client', () => {
           param1: {
             $ref: '#/$defs/MyType',
           },
+          wait_for_previous: {
+            type: 'boolean',
+            description:
+              'Set to true to wait for all previously requested tools in this turn to complete before starting. Set to false (or omit) to run in parallel. Use true when this tool depends on the output of previous tools.',
+          },
         },
         $defs: {
           MyType: {

packages/core/src/tools/mcp-client.test.ts

@@ -150,7 +150,17 @@ describe('DiscoveredMCPTool', () => {
       );
       expect(tool.schema.description).toBe(baseDescription);
       expect(tool.schema.parameters).toBeUndefined();
-      expect(tool.schema.parametersJsonSchema).toEqual(inputSchema);
+      expect(tool.schema.parametersJsonSchema).toEqual({
+        ...inputSchema,
+        properties: {
+          ...(inputSchema['properties'] as Record<string, unknown>),
+          wait_for_previous: {
+            type: 'boolean',
+            description:
+              'Set to true to wait for all previously requested tools in this turn to complete before starting. Set to false (or omit) to run in parallel. Use true when this tool depends on the output of previous tools.',
+          },
+        },
+      });
       expect(tool.serverToolName).toBe(serverToolName);
     });
   });

packages/core/src/tools/mcp-tool.test.ts

@@ -541,6 +541,11 @@ describe('ToolRegistry', () => {
             type: 'string',
             format: 'uuid',
           },
+          wait_for_previous: {
+            type: 'boolean',
+            description:
+              'Set to true to wait for all previously requested tools in this turn to complete before starting. Set to false (or omit) to run in parallel. Use true when this tool depends on the output of previous tools.',
+          },
         },
       });
     });

packages/core/src/tools/tool-registry.test.ts

@@ -11,6 +11,7 @@ import type { ShellExecutionConfig } from '../services/shellExecutionService.js'
 import { SchemaValidator } from '../utils/schemaValidator.js';
 import type { AnsiOutput } from '../utils/terminalSerializer.js';
 import type { MessageBus } from '../confirmation-bus/message-bus.js';
+import { isRecord } from '../utils/markdownUtils.js';
 import { randomUUID } from 'node:crypto';
 import {
   MessageBusType,
@@ -394,6 +395,15 @@ export interface ToolBuilder<
   build(params: TParams): ToolInvocation<TParams, TResult>;
 }
 
+/**
+ * Represents the expected JSON Schema structure for tool parameters.
+ */
+export interface ToolParameterSchema {
+  type: string;
+  properties?: unknown;
+  [key: string]: unknown;
+}
+
 /**
  * New base class for tools that separates validation from execution.
  * New tools should extend this class.
@@ -428,7 +438,49 @@ export abstract class DeclarativeTool<
     return {
       name: this.name,
       description: this.description,
-      parametersJsonSchema: this.parameterSchema,
+      parametersJsonSchema: this.addWaitForPreviousParameter(
+        this.parameterSchema,
+      ),
+    };
+  }
+
+  /**
+   * Type guard to check if an unknown value represents a ToolParameterSchema object.
+   */
+  private isParameterSchema(obj: unknown): obj is ToolParameterSchema {
+    return isRecord(obj) && 'type' in obj;
+  }
+
+  /**
+   * Adds the `wait_for_previous` parameter to the tool's schema.
+   * This allows the model to explicitly control parallel vs sequential execution.
+   */
+  private addWaitForPreviousParameter(schema: unknown): unknown {
+    if (!this.isParameterSchema(schema) || schema.type !== 'object') {
+      return schema;
+    }
+
+    const props = schema.properties;
+    let propertiesObj: Record<string, unknown> = {};
+
+    if (props !== undefined) {
+      if (!isRecord(props)) {
+        // properties exists but is not an object, so it's a malformed schema.
+        return schema;
+      }
+      propertiesObj = props;
+    }
+
+    return {
+      ...schema,
+      properties: {
+        ...propertiesObj,
+        wait_for_previous: {
+          type: 'boolean',
+          description:
+            'Set to true to wait for all previously requested tools in this turn to complete before starting. Set to false (or omit) to run in parallel. Use true when this tool depends on the output of previous tools.',
+        },
+      },
     };
   }