From ee425228fe88b48908802b95625674d06ff05e95 Mon Sep 17 00:00:00 2001 From: Akhilesh Kumar Date: Fri, 13 Mar 2026 17:30:32 +0000 Subject: [PATCH] fix(core): ensure policy engine compatibility with isolated MCP servers This commit addresses PR feedback regarding the prefixing of isolated subagent MCP servers and its potential to break existing security policies relying on standard FQNs. 1. Added `originalName` to `MCPServerConfig` and `originalServerName` to `DiscoveredMCPTool`. 2. Updated `CoreToolScheduler` to reconstruct the original FQN (without the `__agent__` prefix) when performing policy checks via the Policy Engine. This ensures policies mapping to standard `mcp_{server}_{tool}` formats still apply correctly to isolated agents. 3. Added a remote agent back to `NewAgentsNotification.test.tsx` to maintain coverage for both local and remote agents. --- .../ui/components/NewAgentsNotification.test.tsx | 7 +++++++ .../NewAgentsNotification.test.tsx.snap | 1 + packages/core/src/agents/agentLoader.ts | 6 ++++++ packages/core/src/config/config.ts | 5 +++++ packages/core/src/core/coreToolScheduler.ts | 14 +++++++++++--- packages/core/src/tools/mcp-client.ts | 1 + packages/core/src/tools/mcp-tool.ts | 1 + 7 files changed, 32 insertions(+), 3 deletions(-) diff --git a/packages/cli/src/ui/components/NewAgentsNotification.test.tsx b/packages/cli/src/ui/components/NewAgentsNotification.test.tsx index 5962d4d6da..d234b70c4d 100644 --- a/packages/cli/src/ui/components/NewAgentsNotification.test.tsx +++ b/packages/cli/src/ui/components/NewAgentsNotification.test.tsx @@ -38,6 +38,13 @@ describe('NewAgentsNotification', () => { }, }, }, + { + name: 'Agent C', + description: 'Description C', + kind: 'remote' as const, + agentCardUrl: '', + inputConfig: { inputSchema: {} }, + }, ]; const onSelect = vi.fn(); diff --git a/packages/cli/src/ui/components/__snapshots__/NewAgentsNotification.test.tsx.snap b/packages/cli/src/ui/components/__snapshots__/NewAgentsNotification.test.tsx.snap index 0e90c6a9e7..74dcb8a914 100644 --- a/packages/cli/src/ui/components/__snapshots__/NewAgentsNotification.test.tsx.snap +++ b/packages/cli/src/ui/components/__snapshots__/NewAgentsNotification.test.tsx.snap @@ -11,6 +11,7 @@ exports[`NewAgentsNotification > renders agent list 1`] = ` │ │ - Agent A: Description A │ │ │ │ - Agent B: Description B │ │ │ │ (Includes MCP servers: github, postgres) │ │ + │ │ - Agent C: Description C │ │ │ │ │ │ │ └────────────────────────────────────────────────────────────────────────────────────────────┘ │ │ │ diff --git a/packages/core/src/agents/agentLoader.ts b/packages/core/src/agents/agentLoader.ts index 2cb7b3c439..2141f6d915 100644 --- a/packages/core/src/agents/agentLoader.ts +++ b/packages/core/src/agents/agentLoader.ts @@ -550,6 +550,12 @@ export function markdownToAgentDefinition( config.description, config.include_tools, config.exclude_tools, + undefined, // extension + undefined, // oauth + undefined, // authProviderType + undefined, // targetAudience + undefined, // targetServiceAccount + name, // originalName ); } } diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index d12270ec77..55104e91a9 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -445,6 +445,11 @@ export class MCPServerConfig { readonly targetAudience?: string, /* targetServiceAccount format: @.iam.gserviceaccount.com */ readonly targetServiceAccount?: string, + /** + * The original name of the server before any prefixing (e.g. for subagents). + * This is used by the policy engine to match rules. + */ + readonly originalName?: string, ) {} } diff --git a/packages/core/src/core/coreToolScheduler.ts b/packages/core/src/core/coreToolScheduler.ts index 23473e199d..60cfeb930e 100644 --- a/packages/core/src/core/coreToolScheduler.ts +++ b/packages/core/src/core/coreToolScheduler.ts @@ -47,7 +47,11 @@ import { type ToolCallResponseInfo, } from '../scheduler/types.js'; import { ToolExecutor } from '../scheduler/tool-executor.js'; -import { DiscoveredMCPTool } from '../tools/mcp-tool.js'; +import { + DiscoveredMCPTool, + MCP_TOOL_PREFIX, + MCP_QUALIFIED_NAME_SEPARATOR, +} from '../tools/mcp-tool.js'; import { getPolicyDenialError } from '../scheduler/policy.js'; import { GeminiCliOperation } from '../telemetry/constants.js'; @@ -638,12 +642,16 @@ export class CoreToolScheduler { // Policy Check using PolicyEngine // We must reconstruct the FunctionCall format expected by PolicyEngine const toolCallForPolicy = { - name: toolCall.request.name, + name: + toolCall.tool instanceof DiscoveredMCPTool && + toolCall.tool.originalServerName + ? `${MCP_TOOL_PREFIX}${toolCall.tool.originalServerName}${MCP_QUALIFIED_NAME_SEPARATOR}${toolCall.tool.serverToolName}` + : toolCall.request.name, args: toolCall.request.args, }; const serverName = toolCall.tool instanceof DiscoveredMCPTool - ? toolCall.tool.serverName + ? (toolCall.tool.originalServerName ?? toolCall.tool.serverName) : undefined; const toolAnnotations = toolCall.tool.toolAnnotations; diff --git a/packages/core/src/tools/mcp-client.ts b/packages/core/src/tools/mcp-client.ts index 7932e35f38..9b501e8387 100644 --- a/packages/core/src/tools/mcp-client.ts +++ b/packages/core/src/tools/mcp-client.ts @@ -1266,6 +1266,7 @@ export async function discoverTools( mcpServerConfig.extension?.name, mcpServerConfig.extension?.id, annotations as Record | undefined, + mcpServerConfig.originalName, ); discoveredTools.push(tool); diff --git a/packages/core/src/tools/mcp-tool.ts b/packages/core/src/tools/mcp-tool.ts index 5702f88a52..0d8717b378 100644 --- a/packages/core/src/tools/mcp-tool.ts +++ b/packages/core/src/tools/mcp-tool.ts @@ -348,6 +348,7 @@ export class DiscoveredMCPTool extends BaseDeclarativeTool< override readonly extensionName?: string, override readonly extensionId?: string, private readonly _toolAnnotations?: Record, + readonly originalServerName?: string, ) { super( nameOverride ??