mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-04-22 11:04:42 -07:00
feat(core): implement model-driven parallel tool scheduler (#21933)
This commit is contained in:
@@ -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: {
|
||||
|
||||
@@ -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);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -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.',
|
||||
},
|
||||
},
|
||||
});
|
||||
});
|
||||
|
||||
@@ -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.',
|
||||
},
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user