Enchancements for #11971

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<void> {
 		// Find the connection with matching name and source
@@ -1899,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

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", () => {