feat: Add enableSubagents configuration and wire up subagent registration (#9988)

This commit is contained in:
Abhi
2025-10-01 16:54:00 -04:00
committed by GitHub
parent 5b16771567
commit 331ae7dbdf
14 changed files with 352 additions and 39 deletions
+53
View File
@@ -2534,6 +2534,59 @@ describe('loadCliConfig useRipgrep', () => {
expect(config.getUseModelRouter()).toBe(false); expect(config.getUseModelRouter()).toBe(false);
}); });
}); });
describe('loadCliConfig enableSubagents', () => {
it('should be false by default when enableSubagents is not set in settings', async () => {
process.argv = ['node', 'script.js'];
const argv = await parseArguments({} as Settings);
const settings: Settings = {};
const config = await loadCliConfig(
settings,
[],
new ExtensionEnablementManager(
ExtensionStorage.getUserExtensionsDir(),
argv.extensions,
),
'test-session',
argv,
);
expect(config.getEnableSubagents()).toBe(false);
});
it('should be true when enableSubagents is set to true in settings', async () => {
process.argv = ['node', 'script.js'];
const argv = await parseArguments({} as Settings);
const settings: Settings = { experimental: { enableSubagents: true } };
const config = await loadCliConfig(
settings,
[],
new ExtensionEnablementManager(
ExtensionStorage.getUserExtensionsDir(),
argv.extensions,
),
'test-session',
argv,
);
expect(config.getEnableSubagents()).toBe(true);
});
it('should be false when enableSubagents is explicitly set to false in settings', async () => {
process.argv = ['node', 'script.js'];
const argv = await parseArguments({} as Settings);
const settings: Settings = { experimental: { enableSubagents: false } };
const config = await loadCliConfig(
settings,
[],
new ExtensionEnablementManager(
ExtensionStorage.getUserExtensionsDir(),
argv.extensions,
),
'test-session',
argv,
);
expect(config.getEnableSubagents()).toBe(false);
});
});
}); });
describe('screenReader configuration', () => { describe('screenReader configuration', () => {
+1
View File
@@ -708,6 +708,7 @@ export async function loadCliConfig(
useModelRouter, useModelRouter,
enableMessageBusIntegration: enableMessageBusIntegration:
settings.tools?.enableMessageBusIntegration ?? false, settings.tools?.enableMessageBusIntegration ?? false,
enableSubagents: settings.experimental?.enableSubagents ?? false,
}); });
} }
@@ -330,5 +330,24 @@ describe('SettingsSchema', () => {
getSettingsSchema().experimental.properties.useModelRouter.default, getSettingsSchema().experimental.properties.useModelRouter.default,
).toBe(false); ).toBe(false);
}); });
it('should have enableSubagents setting in schema', () => {
expect(
getSettingsSchema().experimental.properties.enableSubagents,
).toBeDefined();
expect(
getSettingsSchema().experimental.properties.enableSubagents.type,
).toBe('boolean');
expect(
getSettingsSchema().experimental.properties.enableSubagents.category,
).toBe('Experimental');
expect(
getSettingsSchema().experimental.properties.enableSubagents.default,
).toBe(false);
expect(
getSettingsSchema().experimental.properties.enableSubagents
.requiresRestart,
).toBe(true);
});
}); });
}); });
@@ -1001,6 +1001,15 @@ const SETTINGS_SCHEMA = {
'Enable model routing to route requests to the best model based on complexity.', 'Enable model routing to route requests to the best model based on complexity.',
showInDialog: true, showInDialog: true,
}, },
enableSubagents: {
type: 'boolean',
label: 'Enable Subagents',
category: 'Experimental',
requiresRestart: true,
default: false,
description: 'Enable experimental subagents.',
showInDialog: false,
},
}, },
}, },
@@ -7,7 +7,7 @@
import type { AgentDefinition } from './types.js'; import type { AgentDefinition } from './types.js';
import { LSTool } from '../tools/ls.js'; import { LSTool } from '../tools/ls.js';
import { ReadFileTool } from '../tools/read-file.js'; import { ReadFileTool } from '../tools/read-file.js';
import { GlobTool } from '../tools/glob.js'; import { GLOB_TOOL_NAME } from '../tools/tool-names.js';
import { GrepTool } from '../tools/grep.js'; import { GrepTool } from '../tools/grep.js';
import { DEFAULT_GEMINI_MODEL } from '../config/models.js'; import { DEFAULT_GEMINI_MODEL } from '../config/models.js';
@@ -104,7 +104,7 @@ ${CODEBASE_REPORT_MARKDOWN}
toolConfig: { toolConfig: {
// Grant access only to read-only tools. // Grant access only to read-only tools.
tools: [LSTool.Name, ReadFileTool.Name, GlobTool.Name, GrepTool.Name], tools: [LSTool.Name, ReadFileTool.Name, GLOB_TOOL_NAME, GrepTool.Name],
}, },
promptConfig: { promptConfig: {
+36 -34
View File
@@ -370,44 +370,46 @@ export class AgentExecutor {
return true; return true;
}); });
const toolPromises = validatedFunctionCalls.map(async (functionCall) => { const toolPromises = validatedFunctionCalls.map(
const callId = functionCall.id ?? `${functionCall.name}-${Date.now()}`; async (functionCall, index) => {
const args = functionCall.args ?? {}; const callId = functionCall.id ?? `${promptId}-${index}`;
const args = functionCall.args ?? {};
this.emitActivity('TOOL_CALL_START', { this.emitActivity('TOOL_CALL_START', {
name: functionCall.name,
args,
});
const requestInfo: ToolCallRequestInfo = {
callId,
name: functionCall.name as string,
args: args as Record<string, unknown>,
isClientInitiated: true,
prompt_id: promptId,
};
const toolResponse = await executeToolCall(
this.runtimeContext,
requestInfo,
signal,
);
if (toolResponse.error) {
this.emitActivity('ERROR', {
context: 'tool_call',
name: functionCall.name, name: functionCall.name,
error: toolResponse.error.message, args,
}); });
} else {
this.emitActivity('TOOL_CALL_END', {
name: functionCall.name,
output: toolResponse.resultDisplay,
});
}
return toolResponse; const requestInfo: ToolCallRequestInfo = {
}); callId,
name: functionCall.name as string,
args: args as Record<string, unknown>,
isClientInitiated: true,
prompt_id: promptId,
};
const toolResponse = await executeToolCall(
this.runtimeContext,
requestInfo,
signal,
);
if (toolResponse.error) {
this.emitActivity('ERROR', {
context: 'tool_call',
name: functionCall.name,
error: toolResponse.error.message,
});
} else {
this.emitActivity('TOOL_CALL_END', {
name: functionCall.name,
output: toolResponse.resultDisplay,
});
}
return toolResponse;
},
);
const toolResponses = await Promise.all(toolPromises); const toolResponses = await Promise.all(toolPromises);
const toolResponseParts: Part[] = toolResponses const toolResponseParts: Part[] = toolResponses
@@ -16,6 +16,7 @@ import { AgentTerminateMode } from './types.js';
import { makeFakeConfig } from '../test-utils/config.js'; import { makeFakeConfig } from '../test-utils/config.js';
import { ToolErrorType } from '../tools/tool-error.js'; import { ToolErrorType } from '../tools/tool-error.js';
import type { Config } from '../config/config.js'; import type { Config } from '../config/config.js';
import type { MessageBus } from '../confirmation-bus/message-bus.js';
vi.mock('./executor.js'); vi.mock('./executor.js');
@@ -52,6 +53,21 @@ describe('SubagentInvocation', () => {
MockAgentExecutor.create.mockResolvedValue(mockExecutorInstance); MockAgentExecutor.create.mockResolvedValue(mockExecutorInstance);
}); });
it('should pass the messageBus to the parent constructor', () => {
const mockMessageBus = {} as MessageBus;
const params = { task: 'Analyze data' };
const invocation = new SubagentInvocation(
params,
testDefinition,
mockConfig,
mockMessageBus,
);
// Access the protected messageBus property by casting to any
// eslint-disable-next-line @typescript-eslint/no-explicit-any
expect((invocation as any).messageBus).toBe(mockMessageBus);
});
describe('getDescription', () => { describe('getDescription', () => {
it('should format the description with inputs', () => { it('should format the description with inputs', () => {
const params = { task: 'Analyze data', priority: 5 }; const params = { task: 'Analyze data', priority: 5 };
+4 -1
View File
@@ -14,6 +14,7 @@ import type {
AgentInputs, AgentInputs,
SubagentActivityEvent, SubagentActivityEvent,
} from './types.js'; } from './types.js';
import type { MessageBus } from '../confirmation-bus/message-bus.js';
const INPUT_PREVIEW_MAX_LENGTH = 50; const INPUT_PREVIEW_MAX_LENGTH = 50;
const DESCRIPTION_MAX_LENGTH = 200; const DESCRIPTION_MAX_LENGTH = 200;
@@ -36,13 +37,15 @@ export class SubagentInvocation extends BaseToolInvocation<
* @param params The validated input parameters for the agent. * @param params The validated input parameters for the agent.
* @param definition The definition object that configures the agent. * @param definition The definition object that configures the agent.
* @param config The global runtime configuration. * @param config The global runtime configuration.
* @param messageBus Optional message bus for policy enforcement.
*/ */
constructor( constructor(
params: AgentInputs, params: AgentInputs,
private readonly definition: AgentDefinition, private readonly definition: AgentDefinition,
private readonly config: Config, private readonly config: Config,
messageBus?: MessageBus,
) { ) {
super(params); super(params, messageBus);
} }
/** /**
@@ -12,6 +12,7 @@ import { makeFakeConfig } from '../test-utils/config.js';
import type { AgentDefinition, AgentInputs } from './types.js'; import type { AgentDefinition, AgentInputs } from './types.js';
import type { Config } from '../config/config.js'; import type { Config } from '../config/config.js';
import { Kind } from '../tools/tools.js'; import { Kind } from '../tools/tools.js';
import type { MessageBus } from '../confirmation-bus/message-bus.js';
// Mock dependencies to isolate the SubagentToolWrapper class // Mock dependencies to isolate the SubagentToolWrapper class
vi.mock('./invocation.js'); vi.mock('./invocation.js');
@@ -119,6 +120,26 @@ describe('SubagentToolWrapper', () => {
params, params,
mockDefinition, mockDefinition,
mockConfig, mockConfig,
undefined,
);
});
it('should pass the messageBus to the SubagentInvocation constructor', () => {
const mockMessageBus = {} as MessageBus;
const wrapper = new SubagentToolWrapper(
mockDefinition,
mockConfig,
mockMessageBus,
);
const params: AgentInputs = { goal: 'Test the invocation', priority: 1 };
wrapper.build(params);
expect(MockedSubagentInvocation).toHaveBeenCalledWith(
params,
mockDefinition,
mockConfig,
mockMessageBus,
); );
}); });
@@ -14,6 +14,7 @@ import type { Config } from '../config/config.js';
import type { AgentDefinition, AgentInputs } from './types.js'; import type { AgentDefinition, AgentInputs } from './types.js';
import { convertInputConfigToJsonSchema } from './schema-utils.js'; import { convertInputConfigToJsonSchema } from './schema-utils.js';
import { SubagentInvocation } from './invocation.js'; import { SubagentInvocation } from './invocation.js';
import type { MessageBus } from '../confirmation-bus/message-bus.js';
/** /**
* A tool wrapper that dynamically exposes a subagent as a standard, * A tool wrapper that dynamically exposes a subagent as a standard,
@@ -31,10 +32,12 @@ export class SubagentToolWrapper extends BaseDeclarativeTool<
* *
* @param definition The `AgentDefinition` of the subagent to wrap. * @param definition The `AgentDefinition` of the subagent to wrap.
* @param config The runtime configuration, passed down to the subagent. * @param config The runtime configuration, passed down to the subagent.
* @param messageBus Optional message bus for policy enforcement.
*/ */
constructor( constructor(
private readonly definition: AgentDefinition, private readonly definition: AgentDefinition,
private readonly config: Config, private readonly config: Config,
messageBus?: MessageBus,
) { ) {
// Dynamically generate the JSON schema required for the tool definition. // Dynamically generate the JSON schema required for the tool definition.
const parameterSchema = convertInputConfigToJsonSchema( const parameterSchema = convertInputConfigToJsonSchema(
@@ -49,6 +52,7 @@ export class SubagentToolWrapper extends BaseDeclarativeTool<
parameterSchema, parameterSchema,
/* isOutputMarkdown */ true, /* isOutputMarkdown */ true,
/* canUpdateOutput */ true, /* canUpdateOutput */ true,
messageBus,
); );
} }
@@ -64,6 +68,11 @@ export class SubagentToolWrapper extends BaseDeclarativeTool<
protected createInvocation( protected createInvocation(
params: AgentInputs, params: AgentInputs,
): ToolInvocation<AgentInputs, ToolResult> { ): ToolInvocation<AgentInputs, ToolResult> {
return new SubagentInvocation(params, this.definition, this.config); return new SubagentInvocation(
params,
this.definition,
this.config,
this.messageBus,
);
} }
} }
+108
View File
@@ -125,6 +125,17 @@ vi.mock('../ide/ide-client.js', () => ({
}, },
})); }));
vi.mock('../agents/registry.js', () => {
const AgentRegistryMock = vi.fn();
AgentRegistryMock.prototype.initialize = vi.fn();
AgentRegistryMock.prototype.getAllDefinitions = vi.fn(() => []);
return { AgentRegistry: AgentRegistryMock };
});
vi.mock('../agents/subagent-tool-wrapper.js', () => ({
SubagentToolWrapper: vi.fn(),
}));
import { BaseLlmClient } from '../core/baseLlmClient.js'; import { BaseLlmClient } from '../core/baseLlmClient.js';
import { tokenLimit } from '../core/tokenLimits.js'; import { tokenLimit } from '../core/tokenLimits.js';
import { uiTelemetryService } from '../telemetry/index.js'; import { uiTelemetryService } from '../telemetry/index.js';
@@ -583,6 +594,31 @@ describe('Server Config (config.ts)', () => {
}); });
}); });
describe('EnableSubagents Configuration', () => {
it('should default enableSubagents to false when not provided', () => {
const config = new Config(baseParams);
expect(config.getEnableSubagents()).toBe(false);
});
it('should set enableSubagents to true when provided as true', () => {
const paramsWithSubagents: ConfigParameters = {
...baseParams,
enableSubagents: true,
};
const config = new Config(paramsWithSubagents);
expect(config.getEnableSubagents()).toBe(true);
});
it('should set enableSubagents to false when explicitly provided as false', () => {
const paramsWithSubagents: ConfigParameters = {
...baseParams,
enableSubagents: false,
};
const config = new Config(paramsWithSubagents);
expect(config.getEnableSubagents()).toBe(false);
});
});
describe('createToolRegistry', () => { describe('createToolRegistry', () => {
it('should register a tool if coreTools contains an argument-specific pattern', async () => { it('should register a tool if coreTools contains an argument-specific pattern', async () => {
const params: ConfigParameters = { const params: ConfigParameters = {
@@ -612,6 +648,78 @@ describe('Server Config (config.ts)', () => {
expect(wasReadFileToolRegistered).toBe(false); expect(wasReadFileToolRegistered).toBe(false);
}); });
it('should register subagents as tools when enableSubagents is true', async () => {
const params: ConfigParameters = {
...baseParams,
enableSubagents: true,
};
const config = new Config(params);
const mockAgentDefinitions = [
{ name: 'agent1', description: 'Agent 1', instructions: 'Inst 1' },
{ name: 'agent2', description: 'Agent 2', instructions: 'Inst 2' },
];
const AgentRegistryMock = (
(await vi.importMock('../agents/registry.js')) as {
AgentRegistry: Mock;
}
).AgentRegistry;
AgentRegistryMock.prototype.getAllDefinitions.mockReturnValue(
mockAgentDefinitions,
);
const SubagentToolWrapperMock = (
(await vi.importMock('../agents/subagent-tool-wrapper.js')) as {
SubagentToolWrapper: Mock;
}
).SubagentToolWrapper;
await config.initialize();
const registerToolMock = (
(await vi.importMock('../tools/tool-registry')) as {
ToolRegistry: { prototype: { registerTool: Mock } };
}
).ToolRegistry.prototype.registerTool;
expect(SubagentToolWrapperMock).toHaveBeenCalledTimes(2);
expect(SubagentToolWrapperMock).toHaveBeenCalledWith(
mockAgentDefinitions[0],
config,
undefined,
);
expect(SubagentToolWrapperMock).toHaveBeenCalledWith(
mockAgentDefinitions[1],
config,
undefined,
);
const calls = registerToolMock.mock.calls;
const registeredWrappers = calls.filter(
(call) => call[0] instanceof SubagentToolWrapperMock,
);
expect(registeredWrappers).toHaveLength(2);
});
it('should not register subagents as tools when enableSubagents is false', async () => {
const params: ConfigParameters = {
...baseParams,
enableSubagents: false,
};
const config = new Config(params);
const SubagentToolWrapperMock = (
(await vi.importMock('../agents/subagent-tool-wrapper.js')) as {
SubagentToolWrapper: Mock;
}
).SubagentToolWrapper;
await config.initialize();
expect(SubagentToolWrapperMock).not.toHaveBeenCalled();
});
describe('with minified tool class names', () => { describe('with minified tool class names', () => {
beforeEach(() => { beforeEach(() => {
Object.defineProperty( Object.defineProperty(
+54
View File
@@ -76,6 +76,9 @@ import type { PolicyEngineConfig } from '../policy/types.js';
import type { UserTierId } from '../code_assist/types.js'; import type { UserTierId } from '../code_assist/types.js';
import { ProxyAgent, setGlobalDispatcher } from 'undici'; import { ProxyAgent, setGlobalDispatcher } from 'undici';
import { AgentRegistry } from '../agents/registry.js';
import { SubagentToolWrapper } from '../agents/subagent-tool-wrapper.js';
export enum ApprovalMode { export enum ApprovalMode {
DEFAULT = 'default', DEFAULT = 'default',
AUTO_EDIT = 'autoEdit', AUTO_EDIT = 'autoEdit',
@@ -255,11 +258,13 @@ export interface ConfigParameters {
output?: OutputSettings; output?: OutputSettings;
useModelRouter?: boolean; useModelRouter?: boolean;
enableMessageBusIntegration?: boolean; enableMessageBusIntegration?: boolean;
enableSubagents?: boolean;
} }
export class Config { export class Config {
private toolRegistry!: ToolRegistry; private toolRegistry!: ToolRegistry;
private promptRegistry!: PromptRegistry; private promptRegistry!: PromptRegistry;
private agentRegistry!: AgentRegistry;
private readonly sessionId: string; private readonly sessionId: string;
private fileSystemService: FileSystemService; private fileSystemService: FileSystemService;
private contentGeneratorConfig!: ContentGeneratorConfig; private contentGeneratorConfig!: ContentGeneratorConfig;
@@ -345,6 +350,7 @@ export class Config {
private readonly outputSettings: OutputSettings; private readonly outputSettings: OutputSettings;
private readonly useModelRouter: boolean; private readonly useModelRouter: boolean;
private readonly enableMessageBusIntegration: boolean; private readonly enableMessageBusIntegration: boolean;
private readonly enableSubagents: boolean;
constructor(params: ConfigParameters) { constructor(params: ConfigParameters) {
this.sessionId = params.sessionId; this.sessionId = params.sessionId;
@@ -433,6 +439,7 @@ export class Config {
this.useModelRouter = params.useModelRouter ?? false; this.useModelRouter = params.useModelRouter ?? false;
this.enableMessageBusIntegration = this.enableMessageBusIntegration =
params.enableMessageBusIntegration ?? false; params.enableMessageBusIntegration ?? false;
this.enableSubagents = params.enableSubagents ?? false;
this.extensionManagement = params.extensionManagement ?? true; this.extensionManagement = params.extensionManagement ?? true;
this.storage = new Storage(this.targetDir); this.storage = new Storage(this.targetDir);
this.enablePromptCompletion = params.enablePromptCompletion ?? false; this.enablePromptCompletion = params.enablePromptCompletion ?? false;
@@ -474,6 +481,10 @@ export class Config {
await this.getGitService(); await this.getGitService();
} }
this.promptRegistry = new PromptRegistry(); this.promptRegistry = new PromptRegistry();
this.agentRegistry = new AgentRegistry(this);
await this.agentRegistry.initialize();
this.toolRegistry = await this.createToolRegistry(); this.toolRegistry = await this.createToolRegistry();
await this.geminiClient.initialize(); await this.geminiClient.initialize();
@@ -620,6 +631,10 @@ export class Config {
return this.workspaceContext; return this.workspaceContext;
} }
getAgentRegistry(): AgentRegistry {
return this.agentRegistry;
}
getToolRegistry(): ToolRegistry { getToolRegistry(): ToolRegistry {
return this.toolRegistry; return this.toolRegistry;
} }
@@ -996,6 +1011,10 @@ export class Config {
return this.enableMessageBusIntegration; return this.enableMessageBusIntegration;
} }
getEnableSubagents(): boolean {
return this.enableSubagents;
}
async createToolRegistry(): Promise<ToolRegistry> { async createToolRegistry(): Promise<ToolRegistry> {
const registry = new ToolRegistry(this, this.eventEmitter); const registry = new ToolRegistry(this, this.eventEmitter);
@@ -1087,6 +1106,41 @@ export class Config {
registerCoreTool(WriteTodosTool, this); registerCoreTool(WriteTodosTool, this);
} }
// Register Subagents as Tools
if (this.getEnableSubagents()) {
const agentDefinitions = this.agentRegistry.getAllDefinitions();
for (const definition of agentDefinitions) {
// We must respect the main allowed/exclude lists for agents too.
const excludeTools = this.getExcludeTools() || [];
const allowedTools = this.getAllowedTools();
const isExcluded = excludeTools.includes(definition.name);
const isAllowed =
!allowedTools || allowedTools.includes(definition.name);
if (isAllowed && !isExcluded) {
try {
const messageBusEnabled = this.getEnableMessageBusIntegration();
const wrapper = new SubagentToolWrapper(
definition,
this,
messageBusEnabled ? this.getMessageBus() : undefined,
);
registry.registerTool(wrapper);
} catch (error) {
console.error(
`Failed to wrap agent '${definition.name}' as a tool:`,
error,
);
}
} else if (this.getDebugMode()) {
console.log(
`[Config] Skipping registration of agent '${definition.name}' due to allow/exclude configuration.`,
);
}
}
}
await registry.discoverAllTools(); await registry.discoverAllTools();
return registry; return registry;
} }
+2 -1
View File
@@ -13,6 +13,7 @@ import { shortenPath, makeRelative } from '../utils/paths.js';
import { type Config } from '../config/config.js'; import { type Config } from '../config/config.js';
import { DEFAULT_FILE_FILTERING_OPTIONS } from '../config/constants.js'; import { DEFAULT_FILE_FILTERING_OPTIONS } from '../config/constants.js';
import { ToolErrorType } from './tool-error.js'; import { ToolErrorType } from './tool-error.js';
import { GLOB_TOOL_NAME } from './tool-names.js';
// Subset of 'Path' interface provided by 'glob' that we can implement for testing // Subset of 'Path' interface provided by 'glob' that we can implement for testing
export interface GlobPath { export interface GlobPath {
@@ -259,7 +260,7 @@ class GlobToolInvocation extends BaseToolInvocation<
* Implementation of the Glob tool logic * Implementation of the Glob tool logic
*/ */
export class GlobTool extends BaseDeclarativeTool<GlobToolParams, ToolResult> { export class GlobTool extends BaseDeclarativeTool<GlobToolParams, ToolResult> {
static readonly Name = 'glob'; static readonly Name = GLOB_TOOL_NAME;
constructor(private config: Config) { constructor(private config: Config) {
super( super(
+17
View File
@@ -0,0 +1,17 @@
/**
* @license
* Copyright 2025 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
// Centralized constants for tool names.
// This prevents circular dependencies that can occur when other modules (like agents)
// need to reference a tool's name without importing the tool's implementation.
export const GLOB_TOOL_NAME = 'glob';
// TODO: Migrate other tool names here to follow this pattern and prevent future circular dependencies.
// Candidates for migration:
// - LSTool ('list_directory')
// - ReadFileTool ('read_file')
// - GrepTool ('search_file_content')