Skip to content

Commit c7f1d19

Browse files
simurg79Bertan Ari
andauthored
feat(modes): per-mode MCP server restrictions (allowlist) (Zoo-Code-Org#453)
* feat(modes): add per-mode allowed MCP servers (allowlist) Adds optional `allowedMcpServers: string[]` to ModeConfig. When defined, only listed MCP servers' schemas are injected into the system prompt and exposed as native tools, preventing context bloat in specialized modes. - Schema: packages/types ModeConfig + schemas/roomodes.json - Prompt: filter mcpHub.getServers() by allowlist in system.ts and native-tools/mcp_server.ts; wired via core/task/build-tools.ts - UI: new McpServerRestriction.tsx editor (cachedState + 150ms debounce, per AGENTS.md SettingsView pattern), integrated into ModesView for both edit and create flows - Tests: schema, system-prompt filtering, native-tool filtering, component behavior; vitest config + toolkit mock updates Ports upstream RooCodeInc/Roo-Code#12004 (fix from simurg79/Roo-Code#1). * test(modes): relax over-strict Profiler assertion in McpServerRestriction Test 3 `<Profiler onRender>` fires whenever the Profiler boundary commits, which happens whenever its parent re-renders — even when every child inside bails out via React.memo. The `=== 0` assertion in Test 3 part (b) therefore measured the Profiler's own commit cadence rather than the child's render work. Relax to `<= 1` and document the caveat in the test's JSDoc. The real anti-flicker guarantee is verified by Test 2 (DOM-node identity preserved across an equivalent `mcpServers` heartbeat). No production code changed. * fix(mcp): enforce mode allowlist in access_mcp_resource availability check Restricted modes could still read MCP resources from disallowed servers because hasAnyMcpResources() inspected the full hub. Forward allowedMcpServers from build-tools into filterNativeToolsForMode so the resource-availability check only considers allowed servers. Addresses review comment from PR Zoo-Code-Org#75 (RooCodeInc/Roo-Code -> Zoo-Code-Org/Zoo-Code). * fix(modes): avoid clobbering concurrent mode edits in debounced MCP allowlist flush The 150ms debounced flush in McpServerRestriction captured the customMode from the scheduling render and spread it on commit, so an edit to another field within the debounce window was overwritten by the stale snapshot. Track the latest customMode/onCommit in refs and merge allowedMcpServers into the freshest snapshot at flush time. Adds Test 4 covering concurrent-edit safety. Addresses review comment from PR Zoo-Code-Org#75. * ci: trigger checks * fix(webview-tests): correct toolkit mock data-testid + missing exports * test(modes): align McpServerRestriction queries with updated checkbox mock shape The shared toolkit mock now forwards data-testid to the inner <input type='checkbox'>, so getByTestId resolves to the checkbox input directly. The old queries re-derived the input via .querySelector("input[type='checkbox']") on the testid element, which now returns null (the testid element IS the input). Update the 5 affected tests to target the checkbox input directly for both the restrict toggle and per-server checkboxes. Meaningful assertions (allowlist enforcement, debounced-flush-merge) are unchanged. * fix(modes): address CodeRabbit review on per-mode MCP allowlist - filter-tools-for-mode: default access_mcp_resource gating to modeConfig.allowedMcpServers when the parameter is omitted (defense in depth), so a restricted mode can never retain the tool via a caller that forgets to thread the allowlist. Adds tests covering param-omitted fallback and explicit-param precedence. - ModesView: add newModeAllowedMcpServers (and switchMode/resetFormState) to handleCreateMode useCallback deps and drop the now-unnecessary react-hooks/exhaustive-deps disable, fixing a potential stale-closure on mode creation. - system.ts: document that the capabilities MCP line is already gated by the filtered allowlist via shouldIncludeMcp. Addresses review feedback on PR Zoo-Code-Org#453. * fix(modes): address frontend PR review on per-mode MCP allowlist (F4-F7) F4: extract shared McpServerChecklist (server checkboxes + not-connected warning) and use it in both the edit panel and the create-mode dialog. F5: add slug-change reseed tests. F6: move toolkit mock registration into vitest.setup.ts via vi.mock and drop the alias from vitest.config.ts. F7: dependency array already complete (no eslint-disable). * fix(modes): address PR review on per-mode MCP allowlist (backend) - Filter MCP servers passed to the capabilities section so disallowed servers are not described in the system prompt (review: getCapabilitiesSection) - Hoist allow-list Set construction out of the server .filter() in system.ts - Add invocation-time MCP server guard (use_mcp_tool / access_mcp_resource) rejecting calls to servers not in the mode's allowedMcpServers, as a second defense layer beyond tool listing/filtering - Add unit tests for the invocation-time guard (mcpServerRestriction.spec.ts) * fix(tests): repair sections.spec MCP mock for getServers-based capability check getCapabilitiesSection now calls mcpHub.getServers() to honor the per-mode allowedMcpServers allowlist, so the prior empty {} as McpHub mock threw 'TypeError: mcpHub.getServers is not a function' in CI. Provide a proper getServers stub (matching the existing repo mock pattern) and add coverage for the new allowlist filtering paths. * fix(schema): add allowedMcpServers description to Zod schema and regenerate roomodes.json --------- Co-authored-by: Bertan Ari <bertanari@microsoft.com>
1 parent ff45fc3 commit c7f1d19

23 files changed

Lines changed: 1656 additions & 51 deletions
Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
import { modeConfigSchema } from "../mode.js"
2+
3+
describe("modeConfigSchema allowedMcpServers", () => {
4+
const baseModeConfig = {
5+
slug: "test-mode",
6+
name: "Test Mode",
7+
roleDefinition: "A test mode",
8+
groups: ["read" as const],
9+
}
10+
11+
it("should accept valid allowedMcpServers array of strings", () => {
12+
const result = modeConfigSchema.safeParse({
13+
...baseModeConfig,
14+
allowedMcpServers: ["server1", "server2"],
15+
})
16+
expect(result.success).toBe(true)
17+
if (result.success) {
18+
expect(result.data.allowedMcpServers).toEqual(["server1", "server2"])
19+
}
20+
})
21+
22+
it("should accept missing/undefined allowedMcpServers", () => {
23+
const result = modeConfigSchema.safeParse(baseModeConfig)
24+
expect(result.success).toBe(true)
25+
if (result.success) {
26+
expect(result.data.allowedMcpServers).toBeUndefined()
27+
}
28+
})
29+
30+
it("should accept empty allowedMcpServers array", () => {
31+
const result = modeConfigSchema.safeParse({
32+
...baseModeConfig,
33+
allowedMcpServers: [],
34+
})
35+
expect(result.success).toBe(true)
36+
if (result.success) {
37+
expect(result.data.allowedMcpServers).toEqual([])
38+
}
39+
})
40+
41+
it("should reject non-string array items", () => {
42+
const result = modeConfigSchema.safeParse({
43+
...baseModeConfig,
44+
allowedMcpServers: [123, 456],
45+
})
46+
expect(result.success).toBe(false)
47+
})
48+
49+
it("should reject non-array value", () => {
50+
const result = modeConfigSchema.safeParse({
51+
...baseModeConfig,
52+
allowedMcpServers: "server1",
53+
})
54+
expect(result.success).toBe(false)
55+
})
56+
})

packages/types/src/mode.ts

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -102,6 +102,12 @@ export const modeConfigSchema = z.object({
102102
customInstructions: z.string().optional(),
103103
groups: groupEntryArraySchema,
104104
source: z.enum(["global", "project"]).optional(),
105+
allowedMcpServers: z
106+
.array(z.string())
107+
.describe(
108+
"Optional list of MCP server names to include. When omitted, all servers are available. When set, only the listed servers are injected.",
109+
)
110+
.optional(),
105111
})
106112

107113
export type ModeConfig = z.infer<typeof modeConfigSchema>

schemas/roomodes.json

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,13 @@
3131
"type": "string",
3232
"enum": ["global", "project"]
3333
},
34+
"allowedMcpServers": {
35+
"type": "array",
36+
"items": {
37+
"type": "string"
38+
},
39+
"description": "Optional list of MCP server names to include. When omitted, all servers are available. When set, only the listed servers are injected."
40+
},
3441
"groups": {
3542
"type": "array",
3643
"items": {

src/core/prompts/__tests__/sections.spec.ts

Lines changed: 35 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -43,8 +43,13 @@ describe("getCapabilitiesSection", () => {
4343
expect(result).toContain("read and write files")
4444
})
4545

46-
it("includes MCP reference when mcpHub is provided", () => {
47-
const mockMcpHub = {} as McpHub
46+
const createMockMcpHub = (serverNames: string[]): McpHub =>
47+
({
48+
getServers: () => serverNames.map((name) => ({ name })),
49+
}) as unknown as McpHub
50+
51+
it("includes MCP reference when mcpHub exposes at least one server", () => {
52+
const mockMcpHub = createMockMcpHub(["test-server"])
4853
const result = getCapabilitiesSection(cwd, mockMcpHub)
4954

5055
expect(result).toContain("MCP servers")
@@ -55,6 +60,34 @@ describe("getCapabilitiesSection", () => {
5560

5661
expect(result).not.toContain("MCP servers")
5762
})
63+
64+
it("excludes MCP reference when mcpHub exposes no servers", () => {
65+
const mockMcpHub = createMockMcpHub([])
66+
const result = getCapabilitiesSection(cwd, mockMcpHub)
67+
68+
expect(result).not.toContain("MCP servers")
69+
})
70+
71+
it("includes MCP reference when allowedMcpServers matches a connected server", () => {
72+
const mockMcpHub = createMockMcpHub(["allowed-server", "other-server"])
73+
const result = getCapabilitiesSection(cwd, mockMcpHub, ["allowed-server"])
74+
75+
expect(result).toContain("MCP servers")
76+
})
77+
78+
it("excludes MCP reference when allowedMcpServers is an empty array", () => {
79+
const mockMcpHub = createMockMcpHub(["test-server"])
80+
const result = getCapabilitiesSection(cwd, mockMcpHub, [])
81+
82+
expect(result).not.toContain("MCP servers")
83+
})
84+
85+
it("excludes MCP reference when allowedMcpServers matches no connected server", () => {
86+
const mockMcpHub = createMockMcpHub(["test-server"])
87+
const result = getCapabilitiesSection(cwd, mockMcpHub, ["nonexistent-server"])
88+
89+
expect(result).not.toContain("MCP servers")
90+
})
5891
})
5992

6093
describe("getRulesSection", () => {

src/core/prompts/__tests__/system-prompt.spec.ts

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -571,6 +571,70 @@ describe("SYSTEM_PROMPT", () => {
571571
expect(prompt).toContain("OBJECTIVE")
572572
})
573573

574+
describe("allowedMcpServers filtering in system prompt", () => {
575+
it("should exclude MCP capability text when allowedMcpServers is empty array", async () => {
576+
mockMcpHub = createMockMcpHub(true)
577+
578+
const customModes: ModeConfig[] = [
579+
{
580+
slug: "filtered-mode",
581+
name: "Filtered Mode",
582+
roleDefinition: "A filtered mode",
583+
groups: ["read", "mcp"] as const,
584+
allowedMcpServers: [],
585+
},
586+
]
587+
588+
const prompt = await SYSTEM_PROMPT(
589+
mockContext,
590+
"/test/path",
591+
false,
592+
mockMcpHub, // mcpHub with servers
593+
undefined, // diffStrategy
594+
"filtered-mode", // mode
595+
undefined, // customModePrompts
596+
customModes, // customModes
597+
undefined, // globalCustomInstructions
598+
experiments,
599+
undefined, // language
600+
undefined, // rooIgnoreInstructions
601+
)
602+
603+
expect(prompt).not.toContain("MCP servers")
604+
})
605+
606+
it("should include MCP capability text when allowedMcpServers matches connected servers", async () => {
607+
mockMcpHub = createMockMcpHub(true) // has "test-server"
608+
609+
const customModes: ModeConfig[] = [
610+
{
611+
slug: "mcp-mode",
612+
name: "MCP Mode",
613+
roleDefinition: "A mode with MCP",
614+
groups: ["read", "mcp"] as const,
615+
allowedMcpServers: ["test-server"],
616+
},
617+
]
618+
619+
const prompt = await SYSTEM_PROMPT(
620+
mockContext,
621+
"/test/path",
622+
false,
623+
mockMcpHub,
624+
undefined,
625+
"mcp-mode",
626+
undefined,
627+
customModes,
628+
undefined,
629+
experiments,
630+
undefined,
631+
undefined,
632+
)
633+
634+
expect(prompt).toContain("MCP servers")
635+
})
636+
})
637+
574638
afterAll(() => {
575639
vi.restoreAllMocks()
576640
})

src/core/prompts/sections/capabilities.ts

Lines changed: 31 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,43 @@
11
import { McpHub } from "../../../services/mcp/McpHub"
22

3-
export function getCapabilitiesSection(cwd: string, mcpHub?: McpHub): string {
3+
/**
4+
* Builds the CAPABILITIES section of the system prompt.
5+
*
6+
* The MCP availability line is only emitted when at least one MCP server is actually
7+
* exposed to the current mode. When `allowedMcpServers` is provided, the hub's server
8+
* list is filtered by that allowlist BEFORE deciding whether to advertise MCP, so the
9+
* capability text matches the per-mode tool exposure:
10+
* - `undefined` allowlist → all connected servers count (backward compatible)
11+
* - empty `[]` allowlist → no servers count ⇒ MCP line omitted
12+
* - populated allowlist → only listed servers count
13+
*
14+
* @param cwd Current working directory used in the prompt text.
15+
* @param mcpHub Optional MCP hub. When omitted, the MCP line is never emitted.
16+
* @param allowedMcpServers Optional per-mode allowlist of MCP server names. When provided,
17+
* the hub's servers are filtered to this set before determining MCP availability.
18+
*/
19+
export function getCapabilitiesSection(cwd: string, mcpHub?: McpHub, allowedMcpServers?: string[]): string {
20+
// Determine whether any MCP server is actually available to the current mode.
21+
// Filtering the hub's servers by the allowlist (when provided) keeps the capability
22+
// text consistent with the tools that are exposed for the mode.
23+
let hasMcpServers = false
24+
if (mcpHub) {
25+
let servers = mcpHub.getServers()
26+
if (allowedMcpServers) {
27+
const allowSet = new Set(allowedMcpServers)
28+
servers = servers.filter((server) => allowSet.has(server.name))
29+
}
30+
hasMcpServers = servers.length > 0
31+
}
32+
433
return `====
534
635
CAPABILITIES
736
837
- You have access to tools that let you execute CLI commands on the user's computer, list files, view source code definitions, regex search, read and write files, and ask follow-up questions. These tools help you effectively accomplish a wide range of tasks, such as writing code, making edits or improvements to existing files, understanding the current state of a project, performing system operations, and much more.
938
- When the user initially gives you a task, a recursive list of all filepaths in the current workspace directory ('${cwd}') will be included in environment_details. This provides an overview of the project's file structure, offering key insights into the project from directory/file names (how developers conceptualize and organize their code) and file extensions (the language used). This can also guide decision-making on which files to explore further. If you need to further explore directories such as outside the current workspace directory, you can use the list_files tool. If you pass 'true' for the recursive parameter, it will list files recursively. Otherwise, it will list files at the top level, which is better suited for generic directories where you don't necessarily need the nested structure, like the Desktop.
1039
- You can use the execute_command tool to run commands on the user's computer whenever you feel it can help accomplish the user's task. When you need to execute a CLI command, you must provide a clear explanation of what the command does. Prefer to execute complex CLI commands over creating executable scripts, since they are more flexible and easier to run. Interactive and long-running commands are allowed, since the commands are run in the user's VSCode terminal. The user may keep commands running in the background and you will be kept updated on their status along the way. Each command you execute is run in a new terminal instance.${
11-
mcpHub
40+
hasMcpServers
1241
? `
1342
- You have access to MCP servers that may provide additional tools and resources. Each server may provide different capabilities that you can use to accomplish tasks more effectively.
1443
`

src/core/prompts/system.ts

Lines changed: 20 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,17 @@ async function generatePrompt(
6666

6767
// Check if MCP functionality should be included
6868
const hasMcpGroup = modeConfig.groups.some((groupEntry) => getGroupName(groupEntry) === "mcp")
69-
const hasMcpServers = mcpHub && mcpHub.getServers().length > 0
69+
const allowedMcpServers = modeConfig.allowedMcpServers
70+
71+
// Hoist the allowlist Set once (matches the sibling call sites, e.g. mcp_server.ts) instead
72+
// of constructing a new Set on every `.filter` iteration.
73+
const allowSet = allowedMcpServers ? new Set(allowedMcpServers) : undefined
74+
75+
let hasMcpServers = false
76+
if (mcpHub) {
77+
const servers = allowSet ? mcpHub.getServers().filter((s) => allowSet.has(s.name)) : mcpHub.getServers()
78+
hasMcpServers = servers.length > 0
79+
}
7080
const shouldIncludeMcp = hasMcpGroup && hasMcpServers
7181

7282
const codeIndexManager = CodeIndexManager.getInstance(context, cwd)
@@ -90,7 +100,15 @@ ${getSharedToolUseSection()}${toolsCatalog}
90100
91101
${getToolUseGuidelinesSection()}
92102
93-
${getCapabilitiesSection(cwd, shouldIncludeMcp ? mcpHub : undefined)}
103+
${
104+
// Forward the hub only when the mode actually exposes the MCP group, and pass the per-mode
105+
// allowlist through so the capabilities section filters servers using the SAME convention as
106+
// the tool-listing layer (a single source of truth for which servers are visible). This keeps
107+
// the capability text consistent with the tools exposed in mixed cases (e.g. one allowed +
108+
// one disallowed server), preventing the section from advertising MCP based on a disallowed
109+
// server. `shouldIncludeMcp` is still used to short-circuit when no allowed server exists.
110+
getCapabilitiesSection(cwd, hasMcpGroup ? mcpHub : undefined, allowedMcpServers)
111+
}
94112
95113
${modesSection}
96114
${skillsSection ? `\n${skillsSection}` : ""}

0 commit comments

Comments
 (0)