feat(policy): support subagent-specific policies in TOML (#21431)

This commit is contained in:
AK
2026-03-09 12:22:46 -07:00
committed by GitHub
parent a17691f0fc
commit 527074b50a
10 changed files with 56 additions and 3 deletions

View File

@@ -219,6 +219,10 @@ Here is a breakdown of the fields available in a TOML policy rule:
# A unique name for the tool, or an array of names.
toolName = "run_shell_command"
# (Optional) The name of a subagent. If provided, the rule only applies to tool calls
# made by this specific subagent.
subagent = "generalist"
# (Optional) The name of an MCP server. Can be combined with toolName
# to form a composite name like "mcpName__toolName".
mcpName = "my-custom-server"

View File

@@ -27,6 +27,7 @@ describe('agent-scheduler', () => {
mockMessageBus = {} as Mocked<MessageBus>;
mockToolRegistry = {
getTool: vi.fn(),
getMessageBus: vi.fn().mockReturnValue(mockMessageBus),
} as unknown as Mocked<ToolRegistry>;
mockConfig = {
getMessageBus: vi.fn().mockReturnValue(mockMessageBus),

View File

@@ -57,10 +57,11 @@ export async function scheduleAgentTools(
// eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
const agentConfig: Config = Object.create(config);
agentConfig.getToolRegistry = () => toolRegistry;
agentConfig.getMessageBus = () => toolRegistry.getMessageBus();
const scheduler = new Scheduler({
config: agentConfig,
messageBus: config.getMessageBus(),
messageBus: toolRegistry.getMessageBus(),
getPreferredEditor: getPreferredEditor ?? (() => undefined),
schedulerId,
parentCallId,

View File

@@ -19,6 +19,7 @@ import { ToolRegistry } from '../tools/tool-registry.js';
import { DiscoveredMCPTool } from '../tools/mcp-tool.js';
import { CompressionStatus } from '../core/turn.js';
import { type ToolCallRequestInfo } from '../scheduler/types.js';
import { type Message } from '../confirmation-bus/types.js';
import { ChatCompressionService } from '../services/chatCompressionService.js';
import { getDirectoryContextString } from '../utils/environmentContext.js';
import { promptIdContext } from '../utils/promptIdContext.js';
@@ -113,10 +114,27 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
runtimeContext: Config,
onActivity?: ActivityCallback,
): Promise<LocalAgentExecutor<TOutput>> {
const parentMessageBus = runtimeContext.getMessageBus();
// Create an override object to inject the subagent name into tool confirmation requests
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
const subagentMessageBus = Object.create(
parentMessageBus,
) as typeof parentMessageBus;
subagentMessageBus.publish = async (message: Message) => {
if (message.type === 'tool-confirmation-request') {
return parentMessageBus.publish({
...message,
subagent: definition.name,
});
}
return parentMessageBus.publish(message);
};
// Create an isolated tool registry for this agent instance.
const agentToolRegistry = new ToolRegistry(
runtimeContext,
runtimeContext.getMessageBus(),
subagentMessageBus,
);
const parentToolRegistry = runtimeContext.getToolRegistry();
const allAgentNames = new Set(

View File

@@ -160,6 +160,7 @@ describe('MessageBus', () => {
{ name: 'test-tool', args: {} },
'test-server',
annotations,
undefined,
);
});

View File

@@ -56,6 +56,7 @@ export class MessageBus extends EventEmitter {
message.toolCall,
message.serverName,
message.toolAnnotations,
message.subagent,
);
switch (decision) {

View File

@@ -38,6 +38,10 @@ export interface ToolConfirmationRequest {
* Optional tool annotations (e.g., readOnlyHint, destructiveHint) from MCP.
*/
toolAnnotations?: Record<string, unknown>;
/**
* Optional subagent name, if this tool call was initiated by a subagent.
*/
subagent?: string;
/**
* Optional rich details for the confirmation UI (diffs, counts, etc.)
*/

View File

@@ -74,6 +74,7 @@ function ruleMatches(
serverName: string | undefined,
currentApprovalMode: ApprovalMode,
toolAnnotations?: Record<string, unknown>,
subagent?: string,
): boolean {
// Check if rule applies to current approval mode
if (rule.modes && rule.modes.length > 0) {
@@ -82,6 +83,13 @@ function ruleMatches(
}
}
// Check subagent if specified (only for PolicyRule, SafetyCheckerRule doesn't have it)
if ('subagent' in rule && rule.subagent) {
if (rule.subagent !== subagent) {
return false;
}
}
// Strictly enforce mcpName identity if the rule dictates it
if (rule.mcpName) {
if (rule.mcpName === '*') {
@@ -203,6 +211,7 @@ export class PolicyEngine {
allowRedirection?: boolean,
rule?: PolicyRule,
toolAnnotations?: Record<string, unknown>,
subagent?: string,
): Promise<CheckResult> {
if (!command) {
return {
@@ -294,6 +303,7 @@ export class PolicyEngine {
{ name: toolName, args: { command: subCmd, dir_path } },
serverName,
toolAnnotations,
subagent,
);
// subResult.decision is already filtered through applyNonInteractiveMode by this.check()
@@ -352,6 +362,7 @@ export class PolicyEngine {
toolCall: FunctionCall,
serverName: string | undefined,
toolAnnotations?: Record<string, unknown>,
subagent?: string,
): Promise<CheckResult> {
// Case 1: Metadata injection is the primary and safest way to identify an MCP server.
// If we have explicit `_serverName` metadata (usually injected by tool-registry for active tools), use it.
@@ -419,6 +430,7 @@ export class PolicyEngine {
serverName,
this.approvalMode,
toolAnnotations,
subagent,
),
);
@@ -437,6 +449,7 @@ export class PolicyEngine {
rule.allowRedirection,
rule,
toolAnnotations,
subagent,
);
decision = shellResult.decision;
if (shellResult.rule) {
@@ -463,9 +476,10 @@ export class PolicyEngine {
this.defaultDecision,
serverName,
shellDirPath,
undefined,
false,
undefined,
toolAnnotations,
subagent,
);
decision = shellResult.decision;
matchedRule = shellResult.rule;
@@ -485,6 +499,7 @@ export class PolicyEngine {
serverName,
this.approvalMode,
toolAnnotations,
subagent,
)
) {
debugLogger.debug(

View File

@@ -38,6 +38,7 @@ const MAX_TYPO_DISTANCE = 3;
*/
const PolicyRuleSchema = z.object({
toolName: z.union([z.string(), z.array(z.string())]).optional(),
subagent: z.string().optional(),
mcpName: z.string().optional(),
argsPattern: z.string().optional(),
commandPrefix: z.union([z.string(), z.array(z.string())]).optional(),
@@ -464,6 +465,7 @@ export async function loadPoliciesFromToml(
const policyRule: PolicyRule = {
toolName: effectiveToolName,
subagent: rule.subagent,
mcpName: rule.mcpName,
decision: rule.decision,
priority: transformPriority(rule.priority, tier),

View File

@@ -110,6 +110,12 @@ export interface PolicyRule {
*/
toolName?: string;
/**
* The name of the subagent this rule applies to.
* If undefined, the rule applies regardless of whether it's the main agent or a subagent.
*/
subagent?: string;
/**
* Identifies the MCP server this rule applies to.
* Enables precise rule matching against `serverName` metadata instead