mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-04-23 11:34:44 -07:00
feat(core): implement dynamic policy registration for subagents (#17838)
This commit is contained in:
@@ -26,6 +26,7 @@ import type { ConfigParameters } from '../config/config.js';
|
||||
import type { ToolRegistry } from '../tools/tool-registry.js';
|
||||
import { ThinkingLevel } from '@google/genai';
|
||||
import type { AcknowledgedAgentsService } from './acknowledgedAgents.js';
|
||||
import { PolicyDecision } from '../policy/types.js';
|
||||
|
||||
vi.mock('./agentLoader.js', () => ({
|
||||
loadAgentsFromDirectory: vi
|
||||
@@ -657,6 +658,114 @@ describe('AgentRegistry', () => {
|
||||
await Promise.all(promises);
|
||||
expect(registry.getAllDefinitions()).toHaveLength(100);
|
||||
});
|
||||
|
||||
it('should dynamically register an ALLOW policy for local agents', async () => {
|
||||
const agent: AgentDefinition = {
|
||||
...MOCK_AGENT_V1,
|
||||
name: 'PolicyTestAgent',
|
||||
};
|
||||
const policyEngine = mockConfig.getPolicyEngine();
|
||||
const addRuleSpy = vi.spyOn(policyEngine, 'addRule');
|
||||
|
||||
await registry.testRegisterAgent(agent);
|
||||
|
||||
expect(addRuleSpy).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
toolName: 'PolicyTestAgent',
|
||||
decision: PolicyDecision.ALLOW,
|
||||
priority: 1.05,
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('should dynamically register an ASK_USER policy for remote agents', async () => {
|
||||
const remoteAgent: AgentDefinition = {
|
||||
kind: 'remote',
|
||||
name: 'RemotePolicyAgent',
|
||||
description: 'A remote agent',
|
||||
agentCardUrl: 'https://example.com/card',
|
||||
inputConfig: { inputSchema: { type: 'object' } },
|
||||
};
|
||||
|
||||
vi.mocked(A2AClientManager.getInstance).mockReturnValue({
|
||||
loadAgent: vi.fn().mockResolvedValue({ name: 'RemotePolicyAgent' }),
|
||||
} as unknown as A2AClientManager);
|
||||
|
||||
const policyEngine = mockConfig.getPolicyEngine();
|
||||
const addRuleSpy = vi.spyOn(policyEngine, 'addRule');
|
||||
|
||||
await registry.testRegisterAgent(remoteAgent);
|
||||
|
||||
expect(addRuleSpy).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
toolName: 'RemotePolicyAgent',
|
||||
decision: PolicyDecision.ASK_USER,
|
||||
priority: 1.05,
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('should not register a policy if a USER policy already exists', async () => {
|
||||
const agent: AgentDefinition = {
|
||||
...MOCK_AGENT_V1,
|
||||
name: 'ExistingUserPolicyAgent',
|
||||
};
|
||||
const policyEngine = mockConfig.getPolicyEngine();
|
||||
// Mock hasRuleForTool to return true when ignoreDynamic=true (simulating a user policy)
|
||||
vi.spyOn(policyEngine, 'hasRuleForTool').mockImplementation(
|
||||
(toolName, ignoreDynamic) =>
|
||||
toolName === 'ExistingUserPolicyAgent' && ignoreDynamic === true,
|
||||
);
|
||||
const addRuleSpy = vi.spyOn(policyEngine, 'addRule');
|
||||
|
||||
await registry.testRegisterAgent(agent);
|
||||
|
||||
expect(addRuleSpy).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should replace an existing dynamic policy when an agent is overwritten', async () => {
|
||||
const localAgent: AgentDefinition = {
|
||||
...MOCK_AGENT_V1,
|
||||
name: 'OverwrittenAgent',
|
||||
};
|
||||
const remoteAgent: AgentDefinition = {
|
||||
kind: 'remote',
|
||||
name: 'OverwrittenAgent',
|
||||
description: 'A remote agent',
|
||||
agentCardUrl: 'https://example.com/card',
|
||||
inputConfig: { inputSchema: { type: 'object' } },
|
||||
};
|
||||
|
||||
vi.mocked(A2AClientManager.getInstance).mockReturnValue({
|
||||
loadAgent: vi.fn().mockResolvedValue({ name: 'OverwrittenAgent' }),
|
||||
} as unknown as A2AClientManager);
|
||||
|
||||
const policyEngine = mockConfig.getPolicyEngine();
|
||||
const removeRuleSpy = vi.spyOn(policyEngine, 'removeRulesForTool');
|
||||
const addRuleSpy = vi.spyOn(policyEngine, 'addRule');
|
||||
|
||||
// 1. Register local
|
||||
await registry.testRegisterAgent(localAgent);
|
||||
expect(addRuleSpy).toHaveBeenLastCalledWith(
|
||||
expect.objectContaining({ decision: PolicyDecision.ALLOW }),
|
||||
);
|
||||
|
||||
// 2. Overwrite with remote
|
||||
await registry.testRegisterAgent(remoteAgent);
|
||||
|
||||
// Verify old dynamic rule was removed
|
||||
expect(removeRuleSpy).toHaveBeenCalledWith(
|
||||
'OverwrittenAgent',
|
||||
'AgentRegistry (Dynamic)',
|
||||
);
|
||||
// Verify new dynamic rule (remote -> ASK_USER) was added
|
||||
expect(addRuleSpy).toHaveBeenLastCalledWith(
|
||||
expect.objectContaining({
|
||||
toolName: 'OverwrittenAgent',
|
||||
decision: PolicyDecision.ASK_USER,
|
||||
}),
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe('reload', () => {
|
||||
|
||||
@@ -21,6 +21,7 @@ import {
|
||||
type ModelConfig,
|
||||
ModelConfigService,
|
||||
} from '../services/modelConfigService.js';
|
||||
import { PolicyDecision } from '../policy/types.js';
|
||||
|
||||
/**
|
||||
* Returns the model config alias for a given agent definition.
|
||||
@@ -266,6 +267,39 @@ export class AgentRegistry {
|
||||
this.agents.set(mergedDefinition.name, mergedDefinition);
|
||||
|
||||
this.registerModelConfigs(mergedDefinition);
|
||||
this.addAgentPolicy(mergedDefinition);
|
||||
}
|
||||
|
||||
private addAgentPolicy(definition: AgentDefinition<z.ZodTypeAny>): void {
|
||||
const policyEngine = this.config.getPolicyEngine();
|
||||
if (!policyEngine) {
|
||||
return;
|
||||
}
|
||||
|
||||
// If the user has explicitly defined a policy for this tool, respect it.
|
||||
// ignoreDynamic=true means we only check for rules NOT added by this registry.
|
||||
if (policyEngine.hasRuleForTool(definition.name, true)) {
|
||||
if (this.config.getDebugMode()) {
|
||||
debugLogger.log(
|
||||
`[AgentRegistry] User policy exists for '${definition.name}', skipping dynamic registration.`,
|
||||
);
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
// Clean up any old dynamic policy for this tool (e.g. if we are overwriting an agent)
|
||||
policyEngine.removeRulesForTool(definition.name, 'AgentRegistry (Dynamic)');
|
||||
|
||||
// Add the new dynamic policy
|
||||
policyEngine.addRule({
|
||||
toolName: definition.name,
|
||||
decision:
|
||||
definition.kind === 'local'
|
||||
? PolicyDecision.ALLOW
|
||||
: PolicyDecision.ASK_USER,
|
||||
priority: 1.05,
|
||||
source: 'AgentRegistry (Dynamic)',
|
||||
});
|
||||
}
|
||||
|
||||
private isAgentEnabled<TOutput extends z.ZodTypeAny>(
|
||||
@@ -342,6 +376,7 @@ export class AgentRegistry {
|
||||
);
|
||||
}
|
||||
this.agents.set(definition.name, definition);
|
||||
this.addAgentPolicy(definition);
|
||||
} catch (e) {
|
||||
debugLogger.warn(
|
||||
`[AgentRegistry] Error loading A2A agent "${definition.name}":`,
|
||||
|
||||
Reference in New Issue
Block a user