diff --git a/packages/types/src/__tests__/mode-allowedMcpServers.spec.ts b/packages/types/src/__tests__/mode-allowedMcpServers.spec.ts new file mode 100644 index 0000000000..9aeece79a1 --- /dev/null +++ b/packages/types/src/__tests__/mode-allowedMcpServers.spec.ts @@ -0,0 +1,56 @@ +import { modeConfigSchema } from "../mode.js" + +describe("modeConfigSchema allowedMcpServers", () => { + const baseModeConfig = { + slug: "test-mode", + name: "Test Mode", + roleDefinition: "A test mode", + groups: ["read" as const], + } + + it("should accept valid allowedMcpServers array of strings", () => { + const result = modeConfigSchema.safeParse({ + ...baseModeConfig, + allowedMcpServers: ["server1", "server2"], + }) + expect(result.success).toBe(true) + if (result.success) { + expect(result.data.allowedMcpServers).toEqual(["server1", "server2"]) + } + }) + + it("should accept missing/undefined allowedMcpServers", () => { + const result = modeConfigSchema.safeParse(baseModeConfig) + expect(result.success).toBe(true) + if (result.success) { + expect(result.data.allowedMcpServers).toBeUndefined() + } + }) + + it("should accept empty allowedMcpServers array", () => { + const result = modeConfigSchema.safeParse({ + ...baseModeConfig, + allowedMcpServers: [], + }) + expect(result.success).toBe(true) + if (result.success) { + expect(result.data.allowedMcpServers).toEqual([]) + } + }) + + it("should reject non-string array items", () => { + const result = modeConfigSchema.safeParse({ + ...baseModeConfig, + allowedMcpServers: [123, 456], + }) + expect(result.success).toBe(false) + }) + + it("should reject non-array value", () => { + const result = modeConfigSchema.safeParse({ + ...baseModeConfig, + allowedMcpServers: "server1", + }) + expect(result.success).toBe(false) + }) +}) diff --git a/packages/types/src/mode.ts b/packages/types/src/mode.ts index d0cd984fda..e86999688b 100644 --- a/packages/types/src/mode.ts +++ b/packages/types/src/mode.ts @@ -102,6 +102,7 @@ export const modeConfigSchema = z.object({ customInstructions: z.string().optional(), groups: groupEntryArraySchema, source: z.enum(["global", "project"]).optional(), + allowedMcpServers: z.array(z.string()).optional(), }) export type ModeConfig = z.infer diff --git a/schemas/roomodes.json b/schemas/roomodes.json index 2b4349afcb..25d122b176 100644 --- a/schemas/roomodes.json +++ b/schemas/roomodes.json @@ -80,6 +80,13 @@ "required": ["relativePath"], "additionalProperties": false } + }, + "allowedMcpServers": { + "type": "array", + "items": { + "type": "string" + }, + "description": "Optional list of MCP server names to include. When omitted, all servers are available. When set, only the listed servers are injected." } }, "required": ["slug", "name", "roleDefinition", "groups"], diff --git a/src/core/prompts/__tests__/system-prompt.spec.ts b/src/core/prompts/__tests__/system-prompt.spec.ts index f555daba06..8731473a2b 100644 --- a/src/core/prompts/__tests__/system-prompt.spec.ts +++ b/src/core/prompts/__tests__/system-prompt.spec.ts @@ -575,6 +575,70 @@ describe("SYSTEM_PROMPT", () => { expect(prompt).toContain("OBJECTIVE") }) + describe("allowedMcpServers filtering in system prompt", () => { + it("should exclude MCP capability text when allowedMcpServers is empty array", async () => { + mockMcpHub = createMockMcpHub(true) + + const customModes: ModeConfig[] = [ + { + slug: "filtered-mode", + name: "Filtered Mode", + roleDefinition: "A filtered mode", + groups: ["read", "mcp"] as const, + allowedMcpServers: [], + }, + ] + + const prompt = await SYSTEM_PROMPT( + mockContext, + "/test/path", + false, + mockMcpHub, // mcpHub with servers + undefined, // diffStrategy + "filtered-mode", // mode + undefined, // customModePrompts + customModes, // customModes + undefined, // globalCustomInstructions + experiments, + undefined, // language + undefined, // rooIgnoreInstructions + ) + + expect(prompt).not.toContain("MCP servers") + }) + + it("should include MCP capability text when allowedMcpServers matches connected servers", async () => { + mockMcpHub = createMockMcpHub(true) // has "test-server" + + const customModes: ModeConfig[] = [ + { + slug: "mcp-mode", + name: "MCP Mode", + roleDefinition: "A mode with MCP", + groups: ["read", "mcp"] as const, + allowedMcpServers: ["test-server"], + }, + ] + + const prompt = await SYSTEM_PROMPT( + mockContext, + "/test/path", + false, + mockMcpHub, + undefined, + "mcp-mode", + undefined, + customModes, + undefined, + experiments, + undefined, + undefined, + ) + + expect(prompt).toContain("MCP servers") + }) + }) + afterAll(() => { vi.restoreAllMocks() }) diff --git a/src/core/prompts/system.ts b/src/core/prompts/system.ts index 0d6071644a..16ce87a1a6 100644 --- a/src/core/prompts/system.ts +++ b/src/core/prompts/system.ts @@ -66,7 +66,15 @@ async function generatePrompt( // Check if MCP functionality should be included const hasMcpGroup = modeConfig.groups.some((groupEntry) => getGroupName(groupEntry) === "mcp") - const hasMcpServers = mcpHub && mcpHub.getServers().length > 0 + const allowedMcpServers = modeConfig.allowedMcpServers + + let hasMcpServers = false + if (mcpHub) { + const servers = allowedMcpServers + ? mcpHub.getServers().filter((s) => new Set(allowedMcpServers).has(s.name)) + : mcpHub.getServers() + hasMcpServers = servers.length > 0 + } const shouldIncludeMcp = hasMcpGroup && hasMcpServers const codeIndexManager = CodeIndexManager.getInstance(context, cwd) @@ -90,7 +98,15 @@ ${getSharedToolUseSection()}${toolsCatalog} ${getToolUseGuidelinesSection()} -${getCapabilitiesSection(cwd, shouldIncludeMcp ? mcpHub : undefined)} +${ + // `shouldIncludeMcp` already accounts for the mode's allowedMcpServers allowlist + // (see hasMcpServers above), so a mode that restricts MCP servers down to none will + // pass `undefined` here and omit the MCP capabilities line entirely. The capabilities + // section only emits a generic MCP availability statement (it does not enumerate + // individual servers), so forwarding the hub when at least one allowed server exists + // is consistent with the per-mode tool exposure. + getCapabilitiesSection(cwd, shouldIncludeMcp ? mcpHub : undefined) +} ${modesSection} ${skillsSection ? `\n${skillsSection}` : ""} diff --git a/src/core/prompts/tools/__tests__/filter-tools-for-mode.spec.ts b/src/core/prompts/tools/__tests__/filter-tools-for-mode.spec.ts index 0b776a2bad..bc3cd0a360 100644 --- a/src/core/prompts/tools/__tests__/filter-tools-for-mode.spec.ts +++ b/src/core/prompts/tools/__tests__/filter-tools-for-mode.spec.ts @@ -89,3 +89,141 @@ describe("filterNativeToolsForMode - disabledTools", () => { expect(resultNames).not.toContain("edit") }) }) + +describe("filterNativeToolsForMode - access_mcp_resource allowlist", () => { + const nativeTools: OpenAI.Chat.ChatCompletionTool[] = [makeTool("read_file"), makeTool("access_mcp_resource")] + + // Minimal McpHub stub exposing only getServers(), which is all the resource + // availability check uses. + function makeMcpHub(servers: Array<{ name: string; resources?: unknown[] }>): any { + return { + getServers: () => servers, + } + } + + it("keeps access_mcp_resource when an allowed server has resources", () => { + const mcpHub = makeMcpHub([{ name: "allowed-server", resources: [{ uri: "res://x" }] }]) + + const result = filterNativeToolsForMode(nativeTools, "code", undefined, undefined, undefined, {}, mcpHub, [ + "allowed-server", + ]) + + const resultNames = result.map((t) => (t as any).function.name) + expect(resultNames).toContain("access_mcp_resource") + }) + + it("removes access_mcp_resource when only a disallowed server has resources", () => { + // The server with resources is NOT in the allowlist, so the restricted + // mode must not retain access_mcp_resource. + const mcpHub = makeMcpHub([ + { name: "allowed-server", resources: [] }, + { name: "blocked-server", resources: [{ uri: "res://secret" }] }, + ]) + + const result = filterNativeToolsForMode(nativeTools, "code", undefined, undefined, undefined, {}, mcpHub, [ + "allowed-server", + ]) + + const resultNames = result.map((t) => (t as any).function.name) + expect(resultNames).not.toContain("access_mcp_resource") + expect(resultNames).toContain("read_file") + }) + + it("considers all servers when no allowlist is provided (unrestricted mode)", () => { + const mcpHub = makeMcpHub([{ name: "any-server", resources: [{ uri: "res://y" }] }]) + + const result = filterNativeToolsForMode(nativeTools, "code", undefined, undefined, undefined, {}, mcpHub) + + const resultNames = result.map((t) => (t as any).function.name) + expect(resultNames).toContain("access_mcp_resource") + }) + + it("removes access_mcp_resource when the allowlist is empty", () => { + const mcpHub = makeMcpHub([{ name: "some-server", resources: [{ uri: "res://z" }] }]) + + const result = filterNativeToolsForMode(nativeTools, "code", undefined, undefined, undefined, {}, mcpHub, []) + + const resultNames = result.map((t) => (t as any).function.name) + expect(resultNames).not.toContain("access_mcp_resource") + }) + + // Defense in depth: even if a caller forgets to thread `allowedMcpServers`, the + // function must fall back to the mode config's own allowlist so a restricted mode + // can never retain access_mcp_resource based on resources from disallowed servers. + describe("falls back to modeConfig.allowedMcpServers when the parameter is omitted", () => { + const restrictedMode = { + slug: "restricted", + name: "Restricted", + roleDefinition: "restricted role", + groups: ["read", "mcp"], + allowedMcpServers: ["allowed-server"], + } as any + + it("removes access_mcp_resource when only a disallowed server has resources (param omitted)", () => { + const mcpHub = makeMcpHub([ + { name: "allowed-server", resources: [] }, + { name: "blocked-server", resources: [{ uri: "res://secret" }] }, + ]) + + // Note: the 8th argument (allowedMcpServers) is intentionally omitted to + // simulate a caller that does not thread the allowlist through. + const result = filterNativeToolsForMode( + nativeTools, + "restricted", + [restrictedMode], + undefined, + undefined, + {}, + mcpHub, + ) + + const resultNames = result.map((t) => (t as any).function.name) + expect(resultNames).not.toContain("access_mcp_resource") + expect(resultNames).toContain("read_file") + }) + + it("keeps access_mcp_resource when an allowed server has resources (param omitted)", () => { + const mcpHub = makeMcpHub([ + { name: "allowed-server", resources: [{ uri: "res://x" }] }, + { name: "blocked-server", resources: [{ uri: "res://secret" }] }, + ]) + + const result = filterNativeToolsForMode( + nativeTools, + "restricted", + [restrictedMode], + undefined, + undefined, + {}, + mcpHub, + ) + + const resultNames = result.map((t) => (t as any).function.name) + expect(resultNames).toContain("access_mcp_resource") + }) + + it("prefers the explicit parameter over the mode config allowlist when both are provided", () => { + // The mode config allows "allowed-server", but the explicit parameter + // allows only "blocked-server" (which has the resources), so the explicit + // parameter must win and access_mcp_resource is retained. + const mcpHub = makeMcpHub([ + { name: "allowed-server", resources: [] }, + { name: "blocked-server", resources: [{ uri: "res://secret" }] }, + ]) + + const result = filterNativeToolsForMode( + nativeTools, + "restricted", + [restrictedMode], + undefined, + undefined, + {}, + mcpHub, + ["blocked-server"], + ) + + const resultNames = result.map((t) => (t as any).function.name) + expect(resultNames).toContain("access_mcp_resource") + }) + }) +}) diff --git a/src/core/prompts/tools/filter-tools-for-mode.ts b/src/core/prompts/tools/filter-tools-for-mode.ts index fdd41e7e33..2b31714a4c 100644 --- a/src/core/prompts/tools/filter-tools-for-mode.ts +++ b/src/core/prompts/tools/filter-tools-for-mode.ts @@ -220,6 +220,9 @@ export function applyModelToolCustomization( * @param codeIndexManager - Code index manager for codebase_search feature check * @param settings - Additional settings for tool filtering (includes modelInfo for model-specific customization) * @param mcpHub - MCP hub for checking available resources + * @param allowedMcpServers - Optional allowlist of MCP server names for the current mode. When + * provided, the resource-availability check only considers servers in this list, so a mode that + * restricts MCP servers cannot retain `access_mcp_resource` based on resources from disallowed servers. * @returns Filtered array of tools allowed for the mode */ export function filterNativeToolsForMode( @@ -230,6 +233,7 @@ export function filterNativeToolsForMode( codeIndexManager?: CodeIndexManager, settings?: Record, mcpHub?: McpHub, + allowedMcpServers?: string[], ): OpenAI.Chat.ChatCompletionTool[] { // Get mode configuration and all tools for this mode const modeSlug = mode ?? defaultModeSlug @@ -301,8 +305,13 @@ export function filterNativeToolsForMode( } } - // Conditionally exclude access_mcp_resource if MCP is not enabled or there are no resources - if (!mcpHub || !hasAnyMcpResources(mcpHub)) { + // Conditionally exclude access_mcp_resource if MCP is not enabled or there are no resources. + // When the mode restricts MCP servers via allowedMcpServers, only resources from allowed + // servers count — otherwise a restricted mode could still read resources from disallowed servers. + // Fall back to the mode config's own allowlist when the caller omits the parameter, so the + // restriction is enforced regardless of call site (defense in depth). + const effectiveAllowedMcpServers = allowedMcpServers ?? modeConfig.allowedMcpServers + if (!mcpHub || !hasAnyMcpResources(mcpHub, effectiveAllowedMcpServers)) { allowedToolNames.delete("access_mcp_resource") } @@ -330,10 +339,18 @@ export function filterNativeToolsForMode( } /** - * Helper function to check if any MCP server has resources available + * Helper function to check if any MCP server has resources available. + * + * When `allowedServers` is provided, only servers whose name is in the allowlist are considered. + * This keeps the `access_mcp_resource` availability check consistent with the mode's MCP server + * allowlist so a restricted mode cannot retain the tool based on resources from disallowed servers. */ -function hasAnyMcpResources(mcpHub: McpHub): boolean { - const servers = mcpHub.getServers() +function hasAnyMcpResources(mcpHub: McpHub, allowedServers?: string[]): boolean { + let servers = mcpHub.getServers() + if (allowedServers) { + const allowSet = new Set(allowedServers) + servers = servers.filter((server) => allowSet.has(server.name)) + } return servers.some((server) => server.resources && server.resources.length > 0) } diff --git a/src/core/prompts/tools/native-tools/__tests__/mcp_server.spec.ts b/src/core/prompts/tools/native-tools/__tests__/mcp_server.spec.ts index ddd7caaccf..acdc2724fa 100644 --- a/src/core/prompts/tools/native-tools/__tests__/mcp_server.spec.ts +++ b/src/core/prompts/tools/native-tools/__tests__/mcp_server.spec.ts @@ -197,4 +197,53 @@ describe("getMcpServerTools", () => { }) expect(getFunction(result[0]).parameters).not.toHaveProperty("required") }) + describe("allowedServers filtering", () => { + it("should return all tools when allowedServers is undefined", () => { + const server1 = createMockServer("server1", [createMockTool("tool1")]) + const server2 = createMockServer("server2", [createMockTool("tool2")]) + const mockHub = createMockMcpHub([server1, server2]) + + const result = getMcpServerTools(mockHub as McpHub, undefined) + + expect(result).toHaveLength(2) + const toolNames = result.map((t) => getFunction(t).name) + expect(toolNames).toContain("mcp--server1--tool1") + expect(toolNames).toContain("mcp--server2--tool2") + }) + + it("should return only tools from allowed servers when allowedServers is provided", () => { + const server1 = createMockServer("server1", [createMockTool("tool1")]) + const server2 = createMockServer("server2", [createMockTool("tool2")]) + const server3 = createMockServer("server3", [createMockTool("tool3")]) + const mockHub = createMockMcpHub([server1, server2, server3]) + + const result = getMcpServerTools(mockHub as McpHub, ["server1", "server3"]) + + expect(result).toHaveLength(2) + const toolNames = result.map((t) => getFunction(t).name) + expect(toolNames).toContain("mcp--server1--tool1") + expect(toolNames).toContain("mcp--server3--tool3") + expect(toolNames).not.toContain("mcp--server2--tool2") + }) + + it("should return empty array when allowedServers is empty array", () => { + const server1 = createMockServer("server1", [createMockTool("tool1")]) + const server2 = createMockServer("server2", [createMockTool("tool2")]) + const mockHub = createMockMcpHub([server1, server2]) + + const result = getMcpServerTools(mockHub as McpHub, []) + + expect(result).toEqual([]) + }) + + it("should ignore server names in allowedServers not found in hub", () => { + const server1 = createMockServer("server1", [createMockTool("tool1")]) + const mockHub = createMockMcpHub([server1]) + + const result = getMcpServerTools(mockHub as McpHub, ["server1", "nonexistent-server"]) + + expect(result).toHaveLength(1) + expect(getFunction(result[0]).name).toBe("mcp--server1--tool1") + }) + }) }) diff --git a/src/core/prompts/tools/native-tools/mcp_server.ts b/src/core/prompts/tools/native-tools/mcp_server.ts index 3fbd1fbcf4..11e3906934 100644 --- a/src/core/prompts/tools/native-tools/mcp_server.ts +++ b/src/core/prompts/tools/native-tools/mcp_server.ts @@ -11,12 +11,18 @@ import { normalizeToolSchema, type JsonSchema } from "../../../../utils/json-sch * @param mcpHub The McpHub instance containing connected servers. * @returns An array of OpenAI.Chat.ChatCompletionTool definitions. */ -export function getMcpServerTools(mcpHub?: McpHub): OpenAI.Chat.ChatCompletionTool[] { +export function getMcpServerTools(mcpHub?: McpHub, allowedServers?: string[]): OpenAI.Chat.ChatCompletionTool[] { if (!mcpHub) { return [] } - const servers = mcpHub.getServers() + let servers = mcpHub.getServers() + + // Filter servers by allowlist if provided + if (allowedServers) { + const allowSet = new Set(allowedServers) + servers = servers.filter((s) => allowSet.has(s.name)) + } const tools: OpenAI.Chat.ChatCompletionTool[] = [] // Track seen tool names to prevent duplicates (e.g., when same server exists in both global and project configs) const seenToolNames = new Set() diff --git a/src/core/task/build-tools.ts b/src/core/task/build-tools.ts index c32d8f6f9b..ebbdc050dc 100644 --- a/src/core/task/build-tools.ts +++ b/src/core/task/build-tools.ts @@ -7,6 +7,7 @@ import { customToolRegistry, formatNative } from "@roo-code/core" import type { ClineProvider } from "../webview/ClineProvider" import { getRooDirectoriesForCwd } from "../../services/roo-config/index.js" +import { getModeBySlug, defaultModeSlug } from "../../shared/modes" import { getNativeTools, getMcpServerTools } from "../prompts/tools/native-tools" import { @@ -113,7 +114,13 @@ export async function buildNativeToolsArrayWithRestrictions(options: BuildToolsO supportsImages, }) - // Filter native tools based on mode restrictions. + // Resolve mode config to get allowedMcpServers for MCP server filtering. + const modeConfig = getModeBySlug(mode ?? defaultModeSlug, customModes) + const allowedMcpServers = modeConfig?.allowedMcpServers + + // Filter native tools based on mode restrictions. The allowlist is forwarded so the + // access_mcp_resource availability check only considers resources from allowed servers; + // otherwise a restricted mode could still read resources from disallowed servers. const filteredNativeTools = filterNativeToolsForMode( nativeTools, mode, @@ -122,10 +129,11 @@ export async function buildNativeToolsArrayWithRestrictions(options: BuildToolsO codeIndexManager, filterSettings, mcpHub, + allowedMcpServers, ) // Filter MCP tools based on mode restrictions. - const mcpTools = getMcpServerTools(mcpHub) + const mcpTools = getMcpServerTools(mcpHub, allowedMcpServers) const filteredMcpTools = filterMcpToolsForMode(mcpTools, mode, customModes, experiments) // Add custom tools if they are available and the experiment is enabled. diff --git a/webview-ui/src/__mocks__/@vscode/webview-ui-toolkit/react.tsx b/webview-ui/src/__mocks__/@vscode/webview-ui-toolkit/react.tsx new file mode 100644 index 0000000000..cc6d089a1a --- /dev/null +++ b/webview-ui/src/__mocks__/@vscode/webview-ui-toolkit/react.tsx @@ -0,0 +1,88 @@ +import React from "react" + +export const VSCodeCheckbox = ({ children, onChange, checked, "data-testid": dataTestId, ...props }: any) => ( + +) + +export const VSCodeRadioGroup = ({ children, ...props }: any) =>
{children}
+ +export const VSCodeRadio = ({ children, value, ...props }: any) => ( + +) + +export const VSCodeTextArea = ({ value, onChange, "data-testid": dataTestId, ...props }: any) => ( +