Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
60 changes: 51 additions & 9 deletions src/services/mcp/McpHub.ts
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,12 @@ const BaseConfigSchema = z.object({
alwaysAllow: z.array(z.string()).default([]),
watchPaths: z.array(z.string()).optional(), // paths to watch for changes and restart server
disabledTools: z.array(z.string()).default([]),
/**
* Allowlist of tool names. When non-empty, ONLY these tools will be enabled
* (enabledForPrompt=true). All other tools on the server are disabled.
* Takes precedence over disabledTools.
*/
onlyAllow: z.array(z.string()).optional(),
})

// Custom error messages for better user feedback
Expand Down Expand Up @@ -993,6 +999,7 @@ export class McpHub {
let configPath: string
let alwaysAllowConfig: string[] = []
let disabledToolsList: string[] = []
let onlyAllowList: string[] | undefined = undefined

// Read from the appropriate config file based on the actual source
try {
Expand All @@ -1014,6 +1021,8 @@ export class McpHub {
if (serverConfigData) {
alwaysAllowConfig = serverConfigData.mcpServers?.[serverName]?.alwaysAllow || []
disabledToolsList = serverConfigData.mcpServers?.[serverName]?.disabledTools || []
// onlyAllow is undefined when not set (distinct from empty array)
onlyAllowList = serverConfigData.mcpServers?.[serverName]?.onlyAllow
}
} catch (error) {
console.error(`Failed to read tool configuration for ${serverName}:`, error)
Expand All @@ -1023,11 +1032,18 @@ export class McpHub {
// Check if wildcard "*" is in the alwaysAllow config
const hasWildcard = alwaysAllowConfig.includes("*")

// Mark tools as always allowed and enabled for prompt based on settings
// Determine if onlyAllow mode is active (list is defined and non-empty)
const hasOnlyAllow = Array.isArray(onlyAllowList) && onlyAllowList.length > 0

// Mark tools as always allowed and enabled for prompt based on settings.
// onlyAllow takes precedence over disabledTools: when onlyAllow is active,
// a tool is enabled only if its name is in the onlyAllow list.
const tools = (response?.tools || []).map((tool) => ({
...tool,
alwaysAllow: hasWildcard || alwaysAllowConfig.includes(tool.name),
enabledForPrompt: !disabledToolsList.includes(tool.name),
enabledForPrompt: hasOnlyAllow
? onlyAllowList!.includes(tool.name)
: !disabledToolsList.includes(tool.name),
}))

return tools
Expand Down Expand Up @@ -1769,19 +1785,19 @@ export class McpHub {
}

/**
* Helper method to update a specific tool list (alwaysAllow or disabledTools)
* Helper method to update a specific tool list (alwaysAllow, disabledTools, or onlyAllow)
* in the appropriate settings file.
* @param serverName The name of the server to update
* @param source Whether to update the global or project config
* @param toolName The name of the tool to add or remove
* @param listName The name of the list to modify ("alwaysAllow" or "disabledTools")
* @param listName The name of the list to modify ("alwaysAllow", "disabledTools", or "onlyAllow")
* @param addTool Whether to add (true) or remove (false) the tool from the list
*/
private async updateServerToolList(
serverName: string,
source: "global" | "project",
toolName: string,
listName: "alwaysAllow" | "disabledTools",
listName: "alwaysAllow" | "disabledTools" | "onlyAllow",
addTool: boolean,
): Promise<void> {
// Find the connection with matching name and source
Expand Down Expand Up @@ -1883,10 +1899,36 @@ export class McpHub {
isEnabled: boolean,
): Promise<void> {
try {
// When isEnabled is true, we want to remove the tool from the disabledTools list.
// When isEnabled is false, we want to add the tool to the disabledTools list.
const addToolToDisabledList = !isEnabled
await this.updateServerToolList(serverName, source, toolName, "disabledTools", addToolToDisabledList)
// Determine the correct config path based on the source
let configPath: string
if (source === "project") {
const projectMcpPath = await this.getProjectMcpPath()
if (!projectMcpPath) {
throw new Error("Project MCP configuration file not found")
}
configPath = projectMcpPath
} else {
configPath = await this.getMcpSettingsFilePath()
}

// Read config to check if onlyAllow exists
const content = await fs.readFile(configPath, "utf-8")
const config = JSON.parse(content)
const onlyAllowList = config.mcpServers?.[serverName]?.onlyAllow

// Determine which list to modify based on onlyAllow presence
const hasOnlyAllow = Array.isArray(onlyAllowList) && onlyAllowList.length > 0

if (hasOnlyAllow) {
// When onlyAllow is active, toggle membership in onlyAllow list:
// isEnabled true = add to onlyAllow, isEnabled false = remove from onlyAllow
await this.updateServerToolList(serverName, source, toolName, "onlyAllow", isEnabled)
} else {
// Fall back to disabledTools behavior:
// isEnabled true = remove from disabledTools, isEnabled false = add to disabledTools
const addToolToDisabledList = !isEnabled
await this.updateServerToolList(serverName, source, toolName, "disabledTools", addToolToDisabledList)
}
} catch (error) {
this.showErrorMessage(`Failed to update settings for tool ${toolName}`, error)
throw error // Re-throw to ensure the error is properly handled
Expand Down
264 changes: 264 additions & 0 deletions src/services/mcp/__tests__/McpHub.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -1051,6 +1051,187 @@ describe("McpHub", () => {
expect(tools[0].alwaysAllow).toBe(true) // allowed-tool
expect(tools[1].alwaysAllow).toBe(false) // not-allowed-tool
})

it("should enable only tools listed in onlyAllow and disable all others", async () => {
const mockConfig = {
mcpServers: {
"test-server": {
type: "stdio",
command: "node",
args: ["test.js"],
onlyAllow: ["tool1", "tool3"],
},
},
}

vi.mocked(fs.readFile).mockResolvedValue(JSON.stringify(mockConfig))

const mockConnection: ConnectedMcpConnection = {
type: "connected",
server: {
name: "test-server",
type: "stdio",
command: "node",
args: ["test.js"],
source: "global",
} as any,
client: {
request: vi.fn().mockResolvedValue({
tools: [
{ name: "tool1", description: "Tool 1" },
{ name: "tool2", description: "Tool 2" },
{ name: "tool3", description: "Tool 3" },
{ name: "tool4", description: "Tool 4" },
],
}),
} as any,
transport: {} as any,
}
mcpHub.connections = [mockConnection]

const tools = await mcpHub["fetchToolsList"]("test-server", "global")

expect(tools.length).toBe(4)
expect(tools[0].enabledForPrompt).toBe(true) // tool1 – in onlyAllow
expect(tools[1].enabledForPrompt).toBe(false) // tool2 – not in onlyAllow
expect(tools[2].enabledForPrompt).toBe(true) // tool3 – in onlyAllow
expect(tools[3].enabledForPrompt).toBe(false) // tool4 – not in onlyAllow
})

it("should take precedence over disabledTools when onlyAllow is set", async () => {
const mockConfig = {
mcpServers: {
"test-server": {
type: "stdio",
command: "node",
args: ["test.js"],
onlyAllow: ["tool1"],
disabledTools: ["tool1"], // onlyAllow should win
},
},
}

vi.mocked(fs.readFile).mockResolvedValue(JSON.stringify(mockConfig))

const mockConnection: ConnectedMcpConnection = {
type: "connected",
server: {
name: "test-server",
type: "stdio",
command: "node",
args: ["test.js"],
source: "global",
} as any,
client: {
request: vi.fn().mockResolvedValue({
tools: [
{ name: "tool1", description: "Tool 1" },
{ name: "tool2", description: "Tool 2" },
],
}),
} as any,
transport: {} as any,
}
mcpHub.connections = [mockConnection]

const tools = await mcpHub["fetchToolsList"]("test-server", "global")

expect(tools.length).toBe(2)
// onlyAllow wins: tool1 is enabled even though it's also in disabledTools
expect(tools[0].enabledForPrompt).toBe(true)
// tool2 is not in onlyAllow so it's disabled
expect(tools[1].enabledForPrompt).toBe(false)
})

it("should fall back to disabledTools behaviour when onlyAllow is absent", async () => {
const mockConfig = {
mcpServers: {
"test-server": {
type: "stdio",
command: "node",
args: ["test.js"],
disabledTools: ["tool2"],
},
},
}

vi.mocked(fs.readFile).mockResolvedValue(JSON.stringify(mockConfig))

const mockConnection: ConnectedMcpConnection = {
type: "connected",
server: {
name: "test-server",
type: "stdio",
command: "node",
args: ["test.js"],
source: "global",
} as any,
client: {
request: vi.fn().mockResolvedValue({
tools: [
{ name: "tool1", description: "Tool 1" },
{ name: "tool2", description: "Tool 2" },
{ name: "tool3", description: "Tool 3" },
],
}),
} as any,
transport: {} as any,
}
mcpHub.connections = [mockConnection]

const tools = await mcpHub["fetchToolsList"]("test-server", "global")

expect(tools.length).toBe(3)
expect(tools[0].enabledForPrompt).toBe(true) // tool1 – not disabled
expect(tools[1].enabledForPrompt).toBe(false) // tool2 – in disabledTools
expect(tools[2].enabledForPrompt).toBe(true) // tool3 – not disabled
})

it("should treat an empty onlyAllow array as if onlyAllow is not set", async () => {
const mockConfig = {
mcpServers: {
"test-server": {
type: "stdio",
command: "node",
args: ["test.js"],
onlyAllow: [], // empty – should NOT disable all tools
disabledTools: ["tool2"],
},
},
}

vi.mocked(fs.readFile).mockResolvedValue(JSON.stringify(mockConfig))

const mockConnection: ConnectedMcpConnection = {
type: "connected",
server: {
name: "test-server",
type: "stdio",
command: "node",
args: ["test.js"],
source: "global",
} as any,
client: {
request: vi.fn().mockResolvedValue({
tools: [
{ name: "tool1", description: "Tool 1" },
{ name: "tool2", description: "Tool 2" },
{ name: "tool3", description: "Tool 3" },
],
}),
} as any,
transport: {} as any,
}
mcpHub.connections = [mockConnection]

const tools = await mcpHub["fetchToolsList"]("test-server", "global")

// Falls back to disabledTools behaviour since onlyAllow is empty
expect(tools.length).toBe(3)
expect(tools[0].enabledForPrompt).toBe(true) // tool1
expect(tools[1].enabledForPrompt).toBe(false) // tool2 – in disabledTools
expect(tools[2].enabledForPrompt).toBe(true) // tool3
})
})

describe("toggleToolEnabledForPrompt", () => {
Expand Down Expand Up @@ -1194,6 +1375,89 @@ describe("McpHub", () => {
expect(writtenConfig.mcpServers["test-server"].disabledTools).toBeDefined()
expect(writtenConfig.mcpServers["test-server"].disabledTools).toContain("new-tool")
})

it("should use disabledTools behavior when onlyAllow is absent or empty", async () => {
// When onlyAllow is not present, toggleToolEnabledForPrompt should modify disabledTools
const mockConfig = {
mcpServers: {
"test-server": {
type: "stdio",
command: "node",
args: ["test.js"],
disabledTools: [],
},
},
}

// Set up mock connection
const mockConnection: ConnectedMcpConnection = {
type: "connected",
server: {
name: "test-server",
config: "test-server-config",
status: "connected",
source: "global",
},
client: {} as any,
transport: {} as any,
}
mcpHub.connections = [mockConnection]

// Mock reading config multiple times
;(fs.readFile as Mock).mockResolvedValue(JSON.stringify(mockConfig))

await mcpHub.toggleToolEnabledForPrompt("test-server", "global", "tool1", false)

const writeCalls = (fs.writeFile as Mock).mock.calls
const callToUse = writeCalls[writeCalls.length - 1]
const writtenConfig = JSON.parse(callToUse[1])

// Without onlyAllow, should modify disabledTools
expect(writtenConfig.mcpServers["test-server"].disabledTools).toContain("tool1")
})

it("should modify onlyAllow list when onlyAllow is active", async () => {
// When onlyAllow is present, toggleToolEnabledForPrompt should modify onlyAllow
const mockConfig = {
mcpServers: {
"test-server": {
type: "stdio",
command: "node",
args: ["test.js"],
onlyAllow: ["toolA", "toolB"],
},
},
}

// Set up mock connection
const mockConnection: ConnectedMcpConnection = {
type: "connected",
server: {
name: "test-server",
config: "test-server-config",
status: "connected",
source: "global",
},
client: {} as any,
transport: {} as any,
}
mcpHub.connections = [mockConnection]

// Mock reading config multiple times
;(fs.readFile as Mock).mockResolvedValue(JSON.stringify(mockConfig))

// Enable toolC (add to onlyAllow)
await mcpHub.toggleToolEnabledForPrompt("test-server", "global", "toolC", true)

const writeCalls = (fs.writeFile as Mock).mock.calls
const callToUse = writeCalls[writeCalls.length - 1]
const writtenConfig = JSON.parse(callToUse[1])

// When onlyAllow is active, should modify onlyAllow not disabledTools
expect(writtenConfig.mcpServers["test-server"].onlyAllow).toContain("toolC")
expect(writtenConfig.mcpServers["test-server"].onlyAllow).toContain("toolA")
expect(writtenConfig.mcpServers["test-server"].onlyAllow).toContain("toolB")
})
})

describe("server disabled state", () => {
Expand Down
Loading