feat(core): subagent local execution and tool isolation (#22718)

This commit is contained in:
AK
2026-03-17 19:34:44 -07:00
committed by GitHub
parent bd34a42ec3
commit 7bfe6ac418
7 changed files with 222 additions and 56 deletions
+77 -30
View File
@@ -4,7 +4,6 @@
* SPDX-License-Identifier: Apache-2.0
*/
import type { Config } from '../config/config.js';
import { type AgentLoopContext } from '../config/agent-loop-context.js';
import { reportError } from '../utils/errorReporting.js';
import { GeminiChat, StreamEventType } from '../core/geminiChat.js';
@@ -17,6 +16,8 @@ import {
type Schema,
} from '@google/genai';
import { ToolRegistry } from '../tools/tool-registry.js';
import { PromptRegistry } from '../prompts/prompt-registry.js';
import { ResourceRegistry } from '../resources/resource-registry.js';
import { type AnyDeclarativeTool } from '../tools/tools.js';
import {
DiscoveredMCPTool,
@@ -102,14 +103,22 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
private readonly agentId: string;
private readonly toolRegistry: ToolRegistry;
private readonly promptRegistry: PromptRegistry;
private readonly resourceRegistry: ResourceRegistry;
private readonly context: AgentLoopContext;
private readonly onActivity?: ActivityCallback;
private readonly compressionService: ChatCompressionService;
private readonly parentCallId?: string;
private hasFailedCompressionAttempt = false;
private get config(): Config {
return this.context.config;
private get executionContext(): AgentLoopContext {
return {
...this.context,
toolRegistry: this.toolRegistry,
promptRegistry: this.promptRegistry,
resourceRegistry: this.resourceRegistry,
messageBus: this.toolRegistry.getMessageBus(),
};
}
/**
@@ -133,11 +142,27 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
// Create an override object to inject the subagent name into tool confirmation requests
const subagentMessageBus = parentMessageBus.derive(definition.name);
// Create an isolated tool registry for this agent instance.
// Create isolated registries for this agent instance.
const agentToolRegistry = new ToolRegistry(
context.config,
subagentMessageBus,
);
const agentPromptRegistry = new PromptRegistry();
const agentResourceRegistry = new ResourceRegistry();
if (definition.mcpServers) {
const globalMcpManager = context.config.getMcpClientManager();
if (globalMcpManager) {
for (const [name, config] of Object.entries(definition.mcpServers)) {
await globalMcpManager.maybeDiscoverMcpServer(name, config, {
toolRegistry: agentToolRegistry,
promptRegistry: agentPromptRegistry,
resourceRegistry: agentResourceRegistry,
});
}
}
}
const parentToolRegistry = context.toolRegistry;
const allAgentNames = new Set(
context.config.getAgentRegistry().getAllAgentNames(),
@@ -153,7 +178,9 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
return;
}
agentToolRegistry.registerTool(tool);
// Clone the tool, so it gets its own state and subagent messageBus
const clonedTool = tool.clone(subagentMessageBus);
agentToolRegistry.registerTool(clonedTool);
};
const registerToolByName = (toolName: string) => {
@@ -228,10 +255,12 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
return new LocalAgentExecutor(
definition,
context,
agentToolRegistry,
parentPromptId,
parentCallId,
agentToolRegistry,
agentPromptRegistry,
agentResourceRegistry,
onActivity,
parentCallId,
);
}
@@ -244,14 +273,18 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
private constructor(
definition: LocalAgentDefinition<TOutput>,
context: AgentLoopContext,
toolRegistry: ToolRegistry,
parentPromptId: string | undefined,
parentCallId: string | undefined,
toolRegistry: ToolRegistry,
promptRegistry: PromptRegistry,
resourceRegistry: ResourceRegistry,
onActivity?: ActivityCallback,
parentCallId?: string,
) {
this.definition = definition;
this.context = context;
this.toolRegistry = toolRegistry;
this.promptRegistry = promptRegistry;
this.resourceRegistry = resourceRegistry;
this.onActivity = onActivity;
this.compressionService = new ChatCompressionService();
this.parentCallId = parentCallId;
@@ -447,7 +480,7 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
} finally {
clearTimeout(graceTimeoutId);
logRecoveryAttempt(
this.config,
this.context.config,
new RecoveryAttemptEvent(
this.agentId,
this.definition.name,
@@ -495,7 +528,7 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
const combinedSignal = AbortSignal.any([signal, deadlineTimer.signal]);
logAgentStart(
this.config,
this.context.config,
new AgentStartEvent(this.agentId, this.definition.name),
);
@@ -506,7 +539,7 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
const augmentedInputs = {
...inputs,
cliVersion: await getVersion(),
activeModel: this.config.getActiveModel(),
activeModel: this.context.config.getActiveModel(),
today: new Date().toLocaleDateString(),
};
@@ -528,14 +561,16 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
// Capture the index of the last hint before starting to avoid re-injecting old hints.
// NOTE: Hints added AFTER this point will be broadcast to all currently running
// local agents via the listener below.
const startIndex = this.config.injectionService.getLatestInjectionIndex();
this.config.injectionService.onInjection(injectionListener);
const startIndex =
this.context.config.injectionService.getLatestInjectionIndex();
this.context.config.injectionService.onInjection(injectionListener);
try {
const initialHints = this.config.injectionService.getInjectionsAfter(
startIndex,
'user_steering',
);
const initialHints =
this.context.config.injectionService.getInjectionsAfter(
startIndex,
'user_steering',
);
const formattedInitialHints = formatUserHintsForModel(initialHints);
let currentMessage: Content = formattedInitialHints
@@ -606,7 +641,16 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
}
}
} finally {
this.config.injectionService.offInjection(injectionListener);
this.context.config.injectionService.offInjection(injectionListener);
const globalMcpManager = this.context.config.getMcpClientManager();
if (globalMcpManager) {
globalMcpManager.removeRegistries({
toolRegistry: this.toolRegistry,
promptRegistry: this.promptRegistry,
resourceRegistry: this.resourceRegistry,
});
}
}
// === UNIFIED RECOVERY BLOCK ===
@@ -719,7 +763,7 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
} finally {
deadlineTimer.abort();
logAgentFinish(
this.config,
this.context.config,
new AgentFinishEvent(
this.agentId,
this.definition.name,
@@ -742,7 +786,7 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
prompt_id,
false,
model,
this.config,
this.context.config,
this.hasFailedCompressionAttempt,
);
@@ -780,10 +824,11 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
const modelConfigAlias = getModelConfigAlias(this.definition);
// Resolve the model config early to get the concrete model string (which may be `auto`).
const resolvedConfig = this.config.modelConfigService.getResolvedConfig({
model: modelConfigAlias,
overrideScope: this.definition.name,
});
const resolvedConfig =
this.context.config.modelConfigService.getResolvedConfig({
model: modelConfigAlias,
overrideScope: this.definition.name,
});
const requestedModel = resolvedConfig.model;
let modelToUse: string;
@@ -800,7 +845,7 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
signal,
requestedModel,
};
const router = this.config.getModelRouterService();
const router = this.context.config.getModelRouterService();
const decision = await router.route(routingContext);
modelToUse = decision.model;
} catch (error) {
@@ -888,7 +933,7 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
try {
return new GeminiChat(
this.config,
this.executionContext,
systemInstruction,
[{ functionDeclarations: tools }],
startHistory,
@@ -1136,13 +1181,15 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
// Execute standard tool calls using the new scheduler
if (toolRequests.length > 0) {
const completedCalls = await scheduleAgentTools(
this.config,
this.context.config,
toolRequests,
{
schedulerId: this.agentId,
schedulerId: promptId,
subagent: this.definition.name,
parentCallId: this.parentCallId,
toolRegistry: this.toolRegistry,
promptRegistry: this.promptRegistry,
resourceRegistry: this.resourceRegistry,
signal,
onWaitingForConfirmation,
},
@@ -1277,7 +1324,7 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
let finalPrompt = templateString(promptConfig.systemPrompt, inputs);
// Append environment context (CWD and folder structure).
const dirContext = await getDirectoryContextString(this.config);
const dirContext = await getDirectoryContextString(this.context.config);
finalPrompt += `\n\n# Environment Context\n${dirContext}`;
// Append standard rules for non-interactive execution.