mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-03-10 14:10:37 -07:00
feat(policy): support subagent-specific policies in TOML (#21431)
This commit is contained in:
@@ -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"
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -160,6 +160,7 @@ describe('MessageBus', () => {
|
||||
{ name: 'test-tool', args: {} },
|
||||
'test-server',
|
||||
annotations,
|
||||
undefined,
|
||||
);
|
||||
});
|
||||
|
||||
|
||||
@@ -56,6 +56,7 @@ export class MessageBus extends EventEmitter {
|
||||
message.toolCall,
|
||||
message.serverName,
|
||||
message.toolAnnotations,
|
||||
message.subagent,
|
||||
);
|
||||
|
||||
switch (decision) {
|
||||
|
||||
@@ -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.)
|
||||
*/
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user