Skip to content

Commit d474779

Browse files
qdaxbdaniel-lxs
qdaxb
authored and
daniel-lxs
committed
Support specifying available and unavailable MCP servers in custom mode
1 parent 1237eb8 commit d474779

File tree

23 files changed

+693
-68
lines changed

23 files changed

+693
-68
lines changed

‎packages/types/src/mode.ts‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,23 @@ export const groupOptionsSchema = z.object({
2626
{ message: "Invalid regular expression pattern" },
2727
),
2828
description: z.string().optional(),
29+
mcp: z
30+
.object({
31+
included: z.array(
32+
z.union([
33+
z.string(),
34+
z.record(
35+
z.string(),
36+
z.object({
37+
allowedTools: z.array(z.string()).optional(), // not used yet
38+
disallowedTools: z.array(z.string()).optional(), // not used yet
39+
}),
40+
),
41+
]),
42+
),
43+
description: z.string().optional(),
44+
})
45+
.optional(),
2946
})
3047

3148
export type GroupOptions = z.infer<typeof groupOptionsSchema>

‎src/core/prompts/sections/mcp-servers.ts‎

Lines changed: 63 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,79 @@
11
import { DiffStrategy } from "../../../shared/tools"
22
import { McpHub } from "../../../services/mcp/McpHub"
3+
import { GroupEntry, ModeConfig } from "@roo-code/types"
4+
import { getGroupName } from "../../../shared/modes"
5+
import { McpServer } from "../../../shared/mcp"
6+
7+
let lastMcpHub: McpHub | undefined
8+
let lastMcpIncludedList: string[] | undefined
9+
let lastFilteredServers: McpServer[] = []
10+
11+
function memoizeFilteredServers(mcpHub: McpHub, mcpIncludedList?: string[]): McpServer[] {
12+
const mcpHubChanged = mcpHub !== lastMcpHub
13+
const listChanged = !areArraysEqual(mcpIncludedList, lastMcpIncludedList)
14+
15+
if (!mcpHubChanged && !listChanged) {
16+
return lastFilteredServers
17+
}
18+
19+
lastMcpHub = mcpHub
20+
lastMcpIncludedList = mcpIncludedList
21+
22+
lastFilteredServers = (
23+
mcpIncludedList && mcpIncludedList.length > 0 ? mcpHub.getAllServers() : mcpHub.getServers()
24+
).filter((server) => {
25+
if (mcpIncludedList && mcpIncludedList.length > 0) {
26+
return mcpIncludedList.includes(server.name) && server.status === "connected"
27+
}
28+
return server.status === "connected"
29+
})
30+
31+
return lastFilteredServers
32+
}
33+
function areArraysEqual(arr1?: string[], arr2?: string[]): boolean {
34+
if (!arr1 && !arr2) return true
35+
if (!arr1 || !arr2) return false
36+
if (arr1.length !== arr2.length) return false
37+
38+
return arr1.every((item, index) => item === arr2[index])
39+
}
340

441
export async function getMcpServersSection(
542
mcpHub?: McpHub,
643
diffStrategy?: DiffStrategy,
744
enableMcpServerCreation?: boolean,
45+
currentMode?: ModeConfig,
846
): Promise<string> {
947
if (!mcpHub) {
1048
return ""
1149
}
1250

51+
// Get MCP configuration for current mode
52+
let mcpIncludedList: string[] | undefined
53+
54+
if (currentMode) {
55+
// Find MCP group configuration
56+
const mcpGroup = currentMode.groups.find((group: GroupEntry) => {
57+
if (Array.isArray(group) && group.length === 2 && group[0] === "mcp") {
58+
return true
59+
}
60+
return getGroupName(group) === "mcp"
61+
})
62+
63+
// If MCP group configuration is found, get mcpIncludedList from mcp.included
64+
if (mcpGroup && Array.isArray(mcpGroup) && mcpGroup.length === 2) {
65+
const options = mcpGroup[1] as { mcp?: { included?: unknown[] } }
66+
mcpIncludedList = Array.isArray(options.mcp?.included)
67+
? options.mcp.included.filter((item: unknown): item is string => typeof item === "string")
68+
: undefined
69+
}
70+
}
71+
72+
const filteredServers = memoizeFilteredServers(mcpHub, mcpIncludedList)
73+
1374
const connectedServers =
14-
mcpHub.getServers().length > 0
15-
? `${mcpHub
16-
.getServers()
17-
.filter((server) => server.status === "connected")
75+
filteredServers.length > 0
76+
? `${filteredServers
1877
.map((server) => {
1978
const tools = server.tools
2079
?.filter((tool) => tool.enabledForPrompt !== false)

‎src/core/prompts/system.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -81,7 +81,7 @@ async function generatePrompt(
8181
const [modesSection, mcpServersSection] = await Promise.all([
8282
getModesSection(context),
8383
shouldIncludeMcp
84-
? getMcpServersSection(mcpHub, effectiveDiffStrategy, enableMcpServerCreation)
84+
? getMcpServersSection(mcpHub, effectiveDiffStrategy, enableMcpServerCreation, modeConfig)
8585
: Promise.resolve(""),
8686
])
8787

‎webview-ui/src/components/modes/McpSelector.tsx‎

Lines changed: 256 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,256 @@
1+
import React, { useState, useRef, useEffect } from "react"
2+
import { ChevronsUpDown, X } from "lucide-react"
3+
import { VSCodeCheckbox } from "@vscode/webview-ui-toolkit/react"
4+
import {
5+
Popover,
6+
PopoverContent,
7+
PopoverTrigger,
8+
Button,
9+
Command,
10+
CommandInput,
11+
CommandList,
12+
CommandEmpty,
13+
CommandGroup,
14+
CommandItem,
15+
} from "@src/components/ui"
16+
import { useAppTranslation } from "@src/i18n/TranslationContext"
17+
import { ModeConfig, GroupEntry, GroupOptions } from "@roo-code/types"
18+
import { McpServer } from "@roo/mcp"
19+
20+
interface McpSelectorProps {
21+
group: string
22+
isEnabled: boolean
23+
isCustomMode: boolean
24+
mcpServers: McpServer[]
25+
currentMode?: ModeConfig
26+
visualMode: string
27+
customModes: ModeConfig[]
28+
findModeBySlug: (slug: string, modes: ModeConfig[]) => ModeConfig | undefined
29+
updateCustomMode: (slug: string, config: ModeConfig) => void
30+
}
31+
32+
const McpSelector: React.FC<McpSelectorProps> = ({
33+
group,
34+
isEnabled,
35+
isCustomMode,
36+
mcpServers,
37+
currentMode,
38+
visualMode,
39+
customModes,
40+
findModeBySlug,
41+
updateCustomMode,
42+
}) => {
43+
const { t } = useAppTranslation()
44+
45+
// State
46+
const [isDialogOpen, setIsDialogOpen] = useState(false)
47+
const [mcpIncludedList, setMcpIncludedList] = useState<string[]>([])
48+
const [searchValue, setSearchValue] = useState("")
49+
const searchInputRef = useRef<HTMLInputElement | null>(null)
50+
51+
// Sync MCP settings
52+
useEffect(() => {
53+
if (!currentMode) {
54+
setMcpIncludedList([])
55+
return
56+
}
57+
58+
const mcpGroupArr = currentMode.groups?.find(
59+
(g: GroupEntry): g is ["mcp", GroupOptions] => Array.isArray(g) && g.length === 2 && g[0] === "mcp",
60+
)
61+
62+
const rawGroupOptions: GroupOptions | undefined = mcpGroupArr ? mcpGroupArr[1] : undefined
63+
64+
const included = Array.isArray(rawGroupOptions?.mcp?.included)
65+
? (rawGroupOptions.mcp.included.filter((item) => typeof item === "string") as string[])
66+
: []
67+
68+
// Sync MCP settings when mode changes
69+
setMcpIncludedList(included)
70+
}, [currentMode])
71+
// Handle save
72+
function updateMcpGroupOptions(groups: GroupEntry[] = [], group: string, mcpIncludedList: string[]): GroupEntry[] {
73+
let mcpGroupFound = false
74+
const newGroups = groups
75+
.map((g) => {
76+
if (Array.isArray(g) && g[0] === group) {
77+
mcpGroupFound = true
78+
return [
79+
group,
80+
{
81+
...(g[1] || {}),
82+
mcp: mcpIncludedList.length > 0 ? { included: mcpIncludedList } : undefined,
83+
},
84+
] as GroupEntry
85+
}
86+
if (typeof g === "string" && g === group) {
87+
mcpGroupFound = true
88+
return [
89+
group,
90+
{
91+
mcp: mcpIncludedList.length > 0 ? { included: mcpIncludedList } : undefined,
92+
},
93+
] as GroupEntry
94+
}
95+
return g
96+
})
97+
.filter((g) => g !== undefined)
98+
99+
if (!mcpGroupFound && group === "mcp") {
100+
const groupsWithoutSimpleMcp = newGroups.filter((g) => g !== "mcp")
101+
groupsWithoutSimpleMcp.push([
102+
"mcp",
103+
{
104+
mcp: mcpIncludedList.length > 0 ? { included: mcpIncludedList } : undefined,
105+
},
106+
])
107+
return groupsWithoutSimpleMcp as GroupEntry[]
108+
} else {
109+
return newGroups as GroupEntry[]
110+
}
111+
}
112+
113+
// Handle save
114+
const handleSave = () => {
115+
const customMode = findModeBySlug(visualMode, customModes)
116+
if (!customMode) {
117+
setIsDialogOpen(false)
118+
return
119+
}
120+
121+
const updatedGroups = updateMcpGroupOptions(customMode.groups, group, mcpIncludedList)
122+
123+
updateCustomMode(customMode.slug, {
124+
...customMode,
125+
groups: updatedGroups,
126+
source: customMode.source || "global",
127+
})
128+
129+
setIsDialogOpen(false)
130+
}
131+
if (!isCustomMode || !isEnabled) {
132+
return null
133+
}
134+
135+
return (
136+
<Popover
137+
open={isDialogOpen}
138+
onOpenChange={(open) => {
139+
setIsDialogOpen(open)
140+
// Reset search box
141+
if (!open) {
142+
setTimeout(() => {
143+
setSearchValue("")
144+
}, 100)
145+
}
146+
}}>
147+
<PopoverTrigger asChild>
148+
<Button variant="secondary" size="sm" style={{ marginLeft: 4 }} className="flex items-center gap-1">
149+
{/* Dynamically display button text */}
150+
{mcpIncludedList.length === 0
151+
? t("prompts:tools.mcpAll")
152+
: t("prompts:tools.mcpSelectedCount", {
153+
included: mcpIncludedList.length,
154+
})}
155+
<ChevronsUpDown className="opacity-50 size-3" />
156+
</Button>
157+
</PopoverTrigger>
158+
<PopoverContent className="p-0 w-[400px] bg-vscode-editor-background">
159+
<Command>
160+
<div className="flex items-center border-b border-vscode-input-border p-2">
161+
<div className="font-medium text-sm flex-1">{t("prompts:tools.selectMcpServers")}</div>
162+
<div className="flex gap-2">
163+
<Button variant="secondary" size="sm" onClick={() => setMcpIncludedList([])}>
164+
{t("prompts:tools.buttons.clearAll")}
165+
</Button>
166+
<Button variant="default" size="sm" onClick={handleSave}>
167+
{t("prompts:tools.buttons.save")}
168+
</Button>
169+
</div>
170+
</div>
171+
<div className="relative">
172+
<CommandInput
173+
ref={searchInputRef}
174+
value={searchValue}
175+
onValueChange={setSearchValue}
176+
placeholder={t("prompts:tools.searchMcpServers")}
177+
className="h-9 mr-4"
178+
/>
179+
{searchValue.length > 0 && (
180+
<div className="absolute right-2 top-0 bottom-0 flex items-center justify-center">
181+
<X
182+
className="text-vscode-input-foreground opacity-50 hover:opacity-100 size-4 p-0.5 cursor-pointer"
183+
onClick={() => {
184+
setSearchValue("")
185+
searchInputRef.current?.focus()
186+
}}
187+
/>
188+
</div>
189+
)}
190+
</div>
191+
<div className="border-b border-vscode-input-border p-2">
192+
<div className="text-sm font-medium text-vscode-foreground mb-2">
193+
{t("prompts:tools.requiredMcpList")}
194+
</div>
195+
<div className="text-sm text-vscode-descriptionForeground mb-2">
196+
{t("prompts:tools.mcpDefaultDescription")}
197+
</div>
198+
<CommandList className="max-h-[150px] overflow-auto bg-vscode-editorWidget-background">
199+
<CommandEmpty>
200+
{mcpServers.length === 0 ? (
201+
<div className="py-2 px-2 text-sm text-vscode-descriptionForeground">
202+
{t("prompts:tools.noMcpServers")}
203+
</div>
204+
) : (
205+
<div className="py-2 px-2 text-sm">{t("prompts:tools.noMatchFound")}</div>
206+
)}
207+
</CommandEmpty>
208+
<CommandGroup>
209+
{mcpServers
210+
.filter(
211+
(server) =>
212+
!searchValue ||
213+
server.name.toLowerCase().includes(searchValue.toLowerCase()),
214+
)
215+
.map((server) => (
216+
<CommandItem
217+
key={`included-${server.name}`}
218+
value={`included-${server.name}`}
219+
onSelect={() => {
220+
const isIncluded = mcpIncludedList.includes(server.name)
221+
if (isIncluded) {
222+
setMcpIncludedList(mcpIncludedList.filter((n) => n !== server.name))
223+
} else {
224+
setMcpIncludedList([...mcpIncludedList, server.name])
225+
}
226+
}}
227+
className="flex items-center px-2 py-1">
228+
<div className="flex items-center flex-1 gap-2">
229+
<VSCodeCheckbox
230+
checked={mcpIncludedList.includes(server.name)}
231+
onClick={(e) => {
232+
e.stopPropagation()
233+
const isIncluded = mcpIncludedList.includes(server.name)
234+
if (isIncluded) {
235+
setMcpIncludedList(
236+
mcpIncludedList.filter((n) => n !== server.name),
237+
)
238+
} else {
239+
setMcpIncludedList([...mcpIncludedList, server.name])
240+
}
241+
}}
242+
/>
243+
<span>{server.name}</span>
244+
</div>
245+
</CommandItem>
246+
))}
247+
</CommandGroup>
248+
</CommandList>
249+
</div>
250+
</Command>
251+
</PopoverContent>
252+
</Popover>
253+
)
254+
}
255+
256+
export default McpSelector

0 commit comments

Comments
 (0)