fix(browser-agent): enable "Allow all server tools" session policy (#22343)

This commit is contained in:
cynthialong0-0
2026-03-19 09:32:35 -07:00
committed by GitHub
parent 39d3b0e28c
commit 7de0616229
7 changed files with 297 additions and 6 deletions

View File

@@ -31,6 +31,8 @@ import type { MessageBus } from '../../confirmation-bus/message-bus.js';
import type { BrowserManager, McpToolCallResult } from './browserManager.js';
import { debugLogger } from '../../utils/debugLogger.js';
import { suspendInputBlocker, resumeInputBlocker } from './inputBlocker.js';
import { MCP_TOOL_PREFIX } from '../../tools/mcp-tool.js';
import { BROWSER_AGENT_NAME } from './browserAgentDefinition.js';
/**
* Tools that interact with page elements and require the input blocker
@@ -62,7 +64,13 @@ class McpToolInvocation extends BaseToolInvocation<
messageBus: MessageBus,
private readonly shouldDisableInput: boolean,
) {
super(params, messageBus, toolName, toolName);
super(
params,
messageBus,
`${MCP_TOOL_PREFIX}${BROWSER_AGENT_NAME}_${toolName}`,
toolName,
BROWSER_AGENT_NAME,
);
}
getDescription(): string {
@@ -79,7 +87,7 @@ class McpToolInvocation extends BaseToolInvocation<
return {
type: 'mcp',
title: `Confirm MCP Tool: ${this.toolName}`,
serverName: 'browser-agent',
serverName: BROWSER_AGENT_NAME,
toolName: this.toolName,
toolDisplayName: this.toolName,
onConfirm: async (outcome: ToolConfirmationOutcome) => {
@@ -92,7 +100,7 @@ class McpToolInvocation extends BaseToolInvocation<
_outcome: ToolConfirmationOutcome,
): PolicyUpdateOptions | undefined {
return {
mcpName: 'browser-agent',
mcpName: BROWSER_AGENT_NAME,
};
}
@@ -202,6 +210,14 @@ class McpDeclarativeTool extends DeclarativeTool<
);
}
// Used for determining tool identity in the policy engine to check if a tool
// call is allowed based on policy.
override get toolAnnotations(): Record<string, unknown> {
return {
_serverName: BROWSER_AGENT_NAME,
};
}
build(
params: Record<string, unknown>,
): ToolInvocation<Record<string, unknown>, ToolResult> {

View File

@@ -61,7 +61,7 @@ describe('mcpToolWrapper Confirmation', () => {
expect(details).toEqual(
expect.objectContaining({
type: 'mcp',
serverName: 'browser-agent',
serverName: 'browser_agent',
toolName: 'test_tool',
}),
);
@@ -76,7 +76,7 @@ describe('mcpToolWrapper Confirmation', () => {
expect(mockMessageBus.publish).toHaveBeenCalledWith(
expect.objectContaining({
type: MessageBusType.UPDATE_POLICY,
mcpName: 'browser-agent',
mcpName: 'browser_agent',
persist: false,
}),
);
@@ -94,7 +94,7 @@ describe('mcpToolWrapper Confirmation', () => {
);
expect(options).toEqual({
mcpName: 'browser-agent',
mcpName: 'browser_agent',
});
});
});

View File

@@ -630,6 +630,34 @@ name = "invalid-name"
).toBeUndefined();
});
it('should support mcpName in policy rules from TOML', async () => {
mockPolicyFile(
nodePath.join(MOCK_DEFAULT_DIR, 'mcp.toml'),
`
[[rule]]
toolName = "my-tool"
mcpName = "my-server"
decision = "allow"
priority = 150
`,
);
const config = await createPolicyEngineConfig(
{},
ApprovalMode.DEFAULT,
MOCK_DEFAULT_DIR,
);
const rule = config.rules?.find(
(r) =>
r.toolName === 'mcp_my-server_my-tool' &&
r.mcpName === 'my-server' &&
r.decision === PolicyDecision.ALLOW,
);
expect(rule).toBeDefined();
expect(rule?.priority).toBeCloseTo(1.15, 5);
});
it('should have default ASK_USER rule for discovered tools', async () => {
const config = await createPolicyEngineConfig({}, ApprovalMode.DEFAULT);
const discoveredRule = config.rules?.find(

View File

@@ -576,6 +576,7 @@ export function createPolicyUpdater(
decision: PolicyDecision.ALLOW,
priority,
argsPattern: new RegExp(pattern),
mcpName: message.mcpName,
source: 'Dynamic (Confirmed)',
});
}
@@ -611,6 +612,7 @@ export function createPolicyUpdater(
decision: PolicyDecision.ALLOW,
priority,
argsPattern,
mcpName: message.mcpName,
source: 'Dynamic (Confirmed)',
});
}

View File

@@ -30,6 +30,8 @@ vi.mock('../utils/shell-utils.js', () => ({
interface ParsedPolicy {
rule?: Array<{
commandPrefix?: string | string[];
mcpName?: string;
toolName?: string;
}>;
}
@@ -67,6 +69,7 @@ describe('createPolicyUpdater', () => {
type: MessageBusType.UPDATE_POLICY,
toolName: 'run_shell_command',
commandPrefix: ['echo', 'ls'],
mcpName: 'test-mcp',
persist: false,
});
@@ -76,6 +79,7 @@ describe('createPolicyUpdater', () => {
expect.objectContaining({
toolName: 'run_shell_command',
priority: ALWAYS_ALLOW_PRIORITY,
mcpName: 'test-mcp',
argsPattern: new RegExp(
escapeRegex('"command":"echo') + '(?:[\\s"]|\\\\")',
),
@@ -86,6 +90,7 @@ describe('createPolicyUpdater', () => {
expect.objectContaining({
toolName: 'run_shell_command',
priority: ALWAYS_ALLOW_PRIORITY,
mcpName: 'test-mcp',
argsPattern: new RegExp(
escapeRegex('"command":"ls') + '(?:[\\s"]|\\\\")',
),
@@ -93,6 +98,63 @@ describe('createPolicyUpdater', () => {
);
});
it('should pass mcpName to policyEngine.addRule for argsPattern updates', async () => {
createPolicyUpdater(policyEngine, messageBus, mockStorage);
await messageBus.publish({
type: MessageBusType.UPDATE_POLICY,
toolName: 'test_tool',
argsPattern: '"foo":"bar"',
mcpName: 'test-mcp',
persist: false,
});
expect(policyEngine.addRule).toHaveBeenCalledWith(
expect.objectContaining({
toolName: 'test_tool',
mcpName: 'test-mcp',
argsPattern: /"foo":"bar"/,
}),
);
});
it('should persist mcpName to TOML', async () => {
createPolicyUpdater(policyEngine, messageBus, mockStorage);
vi.mocked(fs.readFile).mockRejectedValue({ code: 'ENOENT' });
vi.mocked(fs.mkdir).mockResolvedValue(undefined);
const mockFileHandle = {
writeFile: vi.fn().mockResolvedValue(undefined),
close: vi.fn().mockResolvedValue(undefined),
};
vi.mocked(fs.open).mockResolvedValue(
mockFileHandle as unknown as fs.FileHandle,
);
vi.mocked(fs.rename).mockResolvedValue(undefined);
await messageBus.publish({
type: MessageBusType.UPDATE_POLICY,
toolName: 'mcp_test-mcp_tool',
mcpName: 'test-mcp',
commandPrefix: 'ls',
persist: true,
});
// Wait for the async listener to complete
await new Promise((resolve) => setTimeout(resolve, 0));
expect(fs.open).toHaveBeenCalled();
const [content] = mockFileHandle.writeFile.mock.calls[0] as [
string,
string,
];
const parsed = toml.parse(content) as unknown as ParsedPolicy;
expect(parsed.rule).toHaveLength(1);
expect(parsed.rule![0].mcpName).toBe('test-mcp');
expect(parsed.rule![0].toolName).toBe('tool'); // toolName should be stripped of MCP prefix
});
it('should add a single rule when commandPrefix is a string', async () => {
createPolicyUpdater(policyEngine, messageBus, mockStorage);