Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 13 additions & 3 deletions packages/core/src/agents/subagent-tool-wrapper.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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.',
},
},
});
});
});

Expand Down
57 changes: 38 additions & 19 deletions packages/core/src/core/__snapshots__/prompts.test.ts.snap

Large diffs are not rendered by default.

3 changes: 2 additions & 1 deletion packages/core/src/prompts/snippets.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
2 changes: 1 addition & 1 deletion packages/core/src/scheduler/scheduler.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
18 changes: 12 additions & 6 deletions packages/core/src/scheduler/scheduler.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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(
Expand Down
48 changes: 47 additions & 1 deletion packages/core/src/scheduler/scheduler_parallel.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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');
});
});
5 changes: 5 additions & 0 deletions packages/core/src/tools/mcp-client.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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: {
Expand Down
12 changes: 11 additions & 1 deletion packages/core/src/tools/mcp-tool.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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);
});
});
Expand Down
5 changes: 5 additions & 0 deletions packages/core/src/tools/tool-registry.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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.',
},
},
});
});
Expand Down
54 changes: 53 additions & 1 deletion packages/core/src/tools/tools.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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.',
},
},
};
}

Expand Down
Loading