From 456a473b282074e7f72ff5fa0198fef4b893e7c7 Mon Sep 17 00:00:00 2001 From: dmjdarshan Date: Sat, 21 Mar 2026 11:28:33 +0530 Subject: [PATCH] Enchancements for #11971 --- src/services/mcp/McpHub.ts | 40 +++++++++-- src/services/mcp/__tests__/McpHub.spec.ts | 83 +++++++++++++++++++++++ 2 files changed, 116 insertions(+), 7 deletions(-) diff --git a/src/services/mcp/McpHub.ts b/src/services/mcp/McpHub.ts index 32fc84bb50..26a54005e9 100644 --- a/src/services/mcp/McpHub.ts +++ b/src/services/mcp/McpHub.ts @@ -1785,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 { // Find the connection with matching name and source @@ -1899,10 +1899,36 @@ export class McpHub { isEnabled: boolean, ): Promise { 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 diff --git a/src/services/mcp/__tests__/McpHub.spec.ts b/src/services/mcp/__tests__/McpHub.spec.ts index e56c80bff8..25a36c5a56 100644 --- a/src/services/mcp/__tests__/McpHub.spec.ts +++ b/src/services/mcp/__tests__/McpHub.spec.ts @@ -1375,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", () => {