🤖 refactor: persist workspace MCP overrides locally

ThomasK33 · ThomasK33 · commit 3abb165517be · 2025-12-22T12:10:44.000+01:00
Change-Id: I3850373c6aaa762b3f9973cc18cf1b3d5d22bc7a
Signed-off-by: Thomas Kosiewski &lt;tk@coder.com&gt;
diff --git a/docs/mcp-servers.mdx b/docs/mcp-servers.mdx
@@ -52,6 +52,15 @@ MCP servers have two scopes:
 - **Configuration** is per-project — The `.mux/mcp.jsonc` file lives in your project root and applies to all workspaces created from that project
 - **Runtime instances** are per-workspace — Each workspace runs its own server processes, so state in one workspace doesn't affect another
 
+## Per-workspace overrides
+
+Mux supports per-workspace MCP overrides (enable/disable servers and restrict tool allowlists) without modifying the project-level `.mux/mcp.jsonc`.
+
+These overrides are stored in a workspace-local file: `.mux/mcp.local.jsonc`.
+
+- This file is intended to be **gitignored** (it contains local-only workspace preferences)
+- Older mux versions stored these overrides in `~/.mux/config.json`; mux will migrate them into `.mux/mcp.local.jsonc` on first use
+
 This means you configure servers once per project, but each workspace (branch) gets isolated server instances with independent state.
 
 ## Behavior
diff --git a/src/cli/cli.test.ts b/src/cli/cli.test.ts
@@ -69,6 +69,7 @@ async function createTestServer(authToken?: string): Promise<TestServerHandle> {
     updateService: services.updateService,
     tokenizerService: services.tokenizerService,
     serverService: services.serverService,
+    workspaceMcpOverridesService: services.workspaceMcpOverridesService,
     mcpConfigService: services.mcpConfigService,
     featureFlagService: services.featureFlagService,
     sessionTimingService: services.sessionTimingService,
diff --git a/src/cli/server.test.ts b/src/cli/server.test.ts
@@ -72,6 +72,7 @@ async function createTestServer(): Promise<TestServerHandle> {
     updateService: services.updateService,
     tokenizerService: services.tokenizerService,
     serverService: services.serverService,
+    workspaceMcpOverridesService: services.workspaceMcpOverridesService,
     mcpConfigService: services.mcpConfigService,
     featureFlagService: services.featureFlagService,
     sessionTimingService: services.sessionTimingService,
diff --git a/src/cli/server.ts b/src/cli/server.ts
@@ -89,6 +89,7 @@ const mockWindow: BrowserWindow = {
     tokenizerService: serviceContainer.tokenizerService,
     serverService: serviceContainer.serverService,
     menuEventService: serviceContainer.menuEventService,
+    workspaceMcpOverridesService: serviceContainer.workspaceMcpOverridesService,
     mcpConfigService: serviceContainer.mcpConfigService,
     featureFlagService: serviceContainer.featureFlagService,
     sessionTimingService: serviceContainer.sessionTimingService,
diff --git a/src/common/orpc/schemas/mcp.ts b/src/common/orpc/schemas/mcp.ts
@@ -3,7 +3,7 @@ import { z } from "zod";
 /**
  * Per-workspace MCP overrides.
  *
- * Stored in ~/.mux/config.json under each workspace entry.
+ * Stored per-workspace in <workspace>/.mux/mcp.local.jsonc (workspace-local, intended to be gitignored).
  * Allows workspaces to disable servers or restrict tool allowlists
  * without modifying the project-level .mux/mcp.jsonc.
  */
diff --git a/src/common/orpc/schemas/project.ts b/src/common/orpc/schemas/project.ts
@@ -58,7 +58,8 @@ export const WorkspaceConfigSchema = z.object({
       "Trunk branch used to create/init this agent task workspace (used for restart-safe init on queued tasks).",
   }),
   mcp: WorkspaceMCPOverridesSchema.optional().meta({
-    description: "Per-workspace MCP overrides (disabled servers, tool allowlists)",
+    description:
+      "LEGACY: Per-workspace MCP overrides (migrated to <workspace>/.mux/mcp.local.jsonc)",
   }),
   archivedAt: z.string().optional().meta({
     description:
diff --git a/src/common/types/mcp.ts b/src/common/types/mcp.ts
@@ -53,9 +53,10 @@ export interface CachedMCPTestResult {
 /**
  * Per-workspace MCP overrides.
  *
- * Stored in ~/.mux/config.json under each workspace entry.
- * Allows workspaces to override project-level server enabled/disabled state
- * and restrict tool allowlists.
+ * Stored per-workspace in <workspace>/.mux/mcp.local.jsonc (workspace-local and intended to be gitignored).
+ *
+ * Legacy note: older mux versions stored these overrides in ~/.mux/config.json under each workspace entry.
+ * Newer versions migrate those values into the workspace-local file on first read/write.
  */
 export interface WorkspaceMCPOverrides {
   /**
diff --git a/src/desktop/main.ts b/src/desktop/main.ts
@@ -336,6 +336,7 @@ async function loadServices(): Promise<void> {
     serverService: services.serverService,
     featureFlagService: services.featureFlagService,
     sessionTimingService: services.sessionTimingService,
+    workspaceMcpOverridesService: services.workspaceMcpOverridesService,
     mcpConfigService: services.mcpConfigService,
     mcpServerManager: services.mcpServerManager,
     menuEventService: services.menuEventService,
diff --git a/src/node/config.test.ts b/src/node/config.test.ts
@@ -148,119 +148,4 @@ describe("Config", () => {
       expect(workspace.createdAt).toBe("2025-01-01T00:00:00.000Z");
     });
   });
-
-  describe("workspace MCP overrides", () => {
-    it("should return undefined for non-existent workspace", () => {
-      const result = config.getWorkspaceMCPOverrides("non-existent-id");
-      expect(result).toBeUndefined();
-    });
-
-    it("should return undefined for workspace without MCP overrides", async () => {
-      const projectPath = "/fake/project";
-      const workspacePath = path.join(config.srcDir, "project", "branch");
-
-      fs.mkdirSync(workspacePath, { recursive: true });
-
-      await config.editConfig((cfg) => {
-        cfg.projects.set(projectPath, {
-          workspaces: [{ path: workspacePath, id: "test-ws-id", name: "branch" }],
-        });
-        return cfg;
-      });
-
-      const result = config.getWorkspaceMCPOverrides("test-ws-id");
-      expect(result).toBeUndefined();
-    });
-
-    it("should set and get MCP overrides for a workspace", async () => {
-      const projectPath = "/fake/project";
-      const workspacePath = path.join(config.srcDir, "project", "branch");
-
-      fs.mkdirSync(workspacePath, { recursive: true });
-
-      await config.editConfig((cfg) => {
-        cfg.projects.set(projectPath, {
-          workspaces: [{ path: workspacePath, id: "test-ws-id", name: "branch" }],
-        });
-        return cfg;
-      });
-
-      // Set overrides
-      await config.setWorkspaceMCPOverrides("test-ws-id", {
-        disabledServers: ["server-a", "server-b"],
-        toolAllowlist: { "server-c": ["tool1", "tool2"] },
-      });
-
-      // Get overrides
-      const result = config.getWorkspaceMCPOverrides("test-ws-id");
-      expect(result).toBeDefined();
-      expect(result!.disabledServers).toEqual(["server-a", "server-b"]);
-      expect(result!.toolAllowlist).toEqual({ "server-c": ["tool1", "tool2"] });
-    });
-
-    it("should remove MCP overrides when set to empty", async () => {
-      const projectPath = "/fake/project";
-      const workspacePath = path.join(config.srcDir, "project", "branch");
-
-      fs.mkdirSync(workspacePath, { recursive: true });
-
-      await config.editConfig((cfg) => {
-        cfg.projects.set(projectPath, {
-          workspaces: [
-            {
-              path: workspacePath,
-              id: "test-ws-id",
-              name: "branch",
-              mcp: { disabledServers: ["server-a"] },
-            },
-          ],
-        });
-        return cfg;
-      });
-
-      // Clear overrides
-      await config.setWorkspaceMCPOverrides("test-ws-id", {});
-
-      // Verify overrides are removed
-      const result = config.getWorkspaceMCPOverrides("test-ws-id");
-      expect(result).toBeUndefined();
-
-      // Verify workspace still exists
-      const configData = config.loadConfigOrDefault();
-      const projectConfig = configData.projects.get(projectPath);
-      expect(projectConfig!.workspaces[0].id).toBe("test-ws-id");
-      expect(projectConfig!.workspaces[0].mcp).toBeUndefined();
-    });
-
-    it("should deduplicate disabledServers", async () => {
-      const projectPath = "/fake/project";
-      const workspacePath = path.join(config.srcDir, "project", "branch");
-
-      fs.mkdirSync(workspacePath, { recursive: true });
-
-      await config.editConfig((cfg) => {
-        cfg.projects.set(projectPath, {
-          workspaces: [{ path: workspacePath, id: "test-ws-id", name: "branch" }],
-        });
-        return cfg;
-      });
-
-      // Set with duplicates
-      await config.setWorkspaceMCPOverrides("test-ws-id", {
-        disabledServers: ["server-a", "server-b", "server-a"],
-      });
-
-      // Verify duplicates are removed
-      const result = config.getWorkspaceMCPOverrides("test-ws-id");
-      expect(result!.disabledServers).toHaveLength(2);
-      expect(result!.disabledServers).toContain("server-a");
-      expect(result!.disabledServers).toContain("server-b");
-    });
-
-    it("should throw error when setting overrides for non-existent workspace", async () => {
-      await expect(
-        config.setWorkspaceMCPOverrides("non-existent-id", { disabledServers: ["server-a"] })
-      ).rejects.toThrow("Workspace non-existent-id not found in config");
-    });
-  });
 });
diff --git a/src/node/config.ts b/src/node/config.ts
@@ -589,65 +589,6 @@ export class Config {
     });
   }
 
-  /**
-   * Get MCP overrides for a workspace.
-   * Returns undefined if workspace not found or no overrides set.
-   */
-  getWorkspaceMCPOverrides(workspaceId: string): Workspace["mcp"] | undefined {
-    const config = this.loadConfigOrDefault();
-    for (const [_projectPath, projectConfig] of config.projects) {
-      const workspace = projectConfig.workspaces.find((w) => w.id === workspaceId);
-      if (workspace) {
-        return workspace.mcp;
-      }
-    }
-    return undefined;
-  }
-
-  /**
-   * Set MCP overrides for a workspace.
-   * @throws Error if workspace not found
-   */
-  async setWorkspaceMCPOverrides(workspaceId: string, overrides: Workspace["mcp"]): Promise<void> {
-    await this.editConfig((config) => {
-      for (const [_projectPath, projectConfig] of config.projects) {
-        const workspace = projectConfig.workspaces.find((w) => w.id === workspaceId);
-        if (workspace) {
-          // Normalize: remove empty arrays to keep config clean
-          const normalized = overrides
-            ? {
-                disabledServers:
-                  overrides.disabledServers && overrides.disabledServers.length > 0
-                    ? [...new Set(overrides.disabledServers)] // De-duplicate
-                    : undefined,
-                enabledServers:
-                  overrides.enabledServers && overrides.enabledServers.length > 0
-                    ? [...new Set(overrides.enabledServers)] // De-duplicate
-                    : undefined,
-                toolAllowlist:
-                  overrides.toolAllowlist && Object.keys(overrides.toolAllowlist).length > 0
-                    ? overrides.toolAllowlist
-                    : undefined,
-              }
-            : undefined;
-
-          // Remove mcp field entirely if no overrides
-          if (
-            !normalized?.disabledServers &&
-            !normalized?.enabledServers &&
-            !normalized?.toolAllowlist
-          ) {
-            delete workspace.mcp;
-          } else {
-            workspace.mcp = normalized;
-          }
-          return config;
-        }
-      }
-      throw new Error(`Workspace ${workspaceId} not found in config`);
-    });
-  }
-
   /**
    * Load providers configuration from JSONC file
    * Supports comments in JSONC format
diff --git a/src/node/orpc/context.ts b/src/node/orpc/context.ts
@@ -14,6 +14,7 @@ import type { MenuEventService } from "@/node/services/menuEventService";
 import type { VoiceService } from "@/node/services/voiceService";
 import type { MCPConfigService } from "@/node/services/mcpConfigService";
 import type { ExperimentsService } from "@/node/services/experimentsService";
+import type { WorkspaceMcpOverridesService } from "@/node/services/workspaceMcpOverridesService";
 import type { MCPServerManager } from "@/node/services/mcpServerManager";
 import type { TelemetryService } from "@/node/services/telemetryService";
 import type { FeatureFlagService } from "@/node/services/featureFlagService";
@@ -37,6 +38,7 @@ export interface ORPCContext {
   menuEventService: MenuEventService;
   voiceService: VoiceService;
   mcpConfigService: MCPConfigService;
+  workspaceMcpOverridesService: WorkspaceMcpOverridesService;
   mcpServerManager: MCPServerManager;
   featureFlagService: FeatureFlagService;
   sessionTimingService: SessionTimingService;
diff --git a/src/node/orpc/router.ts b/src/node/orpc/router.ts
@@ -1111,17 +1111,25 @@ export const router = (authToken?: string) => {
         get: t
           .input(schemas.workspace.mcp.get.input)
           .output(schemas.workspace.mcp.get.output)
-          .handler(({ context, input }) => {
-            const overrides = context.config.getWorkspaceMCPOverrides(input.workspaceId);
-            // Return empty object if no overrides (matches schema default)
-            return overrides ?? {};
+          .handler(async ({ context, input }) => {
+            try {
+              return await context.workspaceMcpOverridesService.getOverridesForWorkspace(
+                input.workspaceId
+              );
+            } catch {
+              // Defensive: overrides must never brick workspace UI.
+              return {};
+            }
           }),
         set: t
           .input(schemas.workspace.mcp.set.input)
           .output(schemas.workspace.mcp.set.output)
           .handler(async ({ context, input }) => {
             try {
-              await context.config.setWorkspaceMCPOverrides(input.workspaceId, input.overrides);
+              await context.workspaceMcpOverridesService.setOverridesForWorkspace(
+                input.workspaceId,
+                input.overrides
+              );
               return { success: true, data: undefined };
             } catch (error) {
               const message = error instanceof Error ? error.message : String(error);
diff --git a/src/node/services/aiService.ts b/src/node/services/aiService.ts
@@ -47,7 +47,9 @@ import { buildSystemMessage, readToolInstructions } from "./systemMessage";
 import { getTokenizerForModel } from "@/node/utils/main/tokenizer";
 import type { TelemetryService } from "@/node/services/telemetryService";
 import { getRuntimeTypeForTelemetry, roundToBase2 } from "@/common/telemetry/utils";
+import type { WorkspaceMCPOverrides } from "@/common/types/mcp";
 import type { MCPServerManager, MCPWorkspaceStats } from "@/node/services/mcpServerManager";
+import { WorkspaceMcpOverridesService } from "./workspaceMcpOverridesService";
 import type { TaskService } from "@/node/services/taskService";
 import { buildProviderOptions } from "@/common/utils/ai/providerOptions";
 import type { ThinkingLevel } from "@/common/types/thinking";
@@ -380,6 +382,7 @@ export class AIService extends EventEmitter {
   private readonly historyService: HistoryService;
   private readonly partialService: PartialService;
   private readonly config: Config;
+  private readonly workspaceMcpOverridesService: WorkspaceMcpOverridesService;
   private mcpServerManager?: MCPServerManager;
   private telemetryService?: TelemetryService;
   private readonly initStateManager: InitStateManager;
@@ -394,12 +397,15 @@ export class AIService extends EventEmitter {
     partialService: PartialService,
     initStateManager: InitStateManager,
     backgroundProcessManager?: BackgroundProcessManager,
-    sessionUsageService?: SessionUsageService
+    sessionUsageService?: SessionUsageService,
+    workspaceMcpOverridesService?: WorkspaceMcpOverridesService
   ) {
     super();
     // Increase max listeners to accommodate multiple concurrent workspace listeners
     // Each workspace subscribes to stream events, and we expect >10 concurrent workspaces
     this.setMaxListeners(50);
+    this.workspaceMcpOverridesService =
+      workspaceMcpOverridesService ?? new WorkspaceMcpOverridesService(config);
     this.config = config;
     this.historyService = historyService;
     this.partialService = partialService;
@@ -1128,7 +1134,18 @@ export class AIService extends EventEmitter {
         : runtime.getWorkspacePath(metadata.projectPath, metadata.name);
 
       // Fetch workspace MCP overrides (for filtering servers and tools)
-      const mcpOverrides = this.config.getWorkspaceMCPOverrides(workspaceId);
+      // NOTE: Stored in <workspace>/.mux/mcp.local.jsonc (not ~/.mux/config.json).
+      let mcpOverrides: WorkspaceMCPOverrides | undefined;
+      try {
+        mcpOverrides =
+          await this.workspaceMcpOverridesService.getOverridesForWorkspace(workspaceId);
+      } catch (error) {
+        log.warn("[MCP] Failed to load workspace MCP overrides; continuing without overrides", {
+          workspaceId,
+          error,
+        });
+        mcpOverrides = undefined;
+      }
 
       // Fetch MCP server config for system prompt (before building message)
       // Pass overrides to filter out disabled servers
diff --git a/src/node/services/serviceContainer.ts b/src/node/services/serviceContainer.ts
diff --git a/src/node/services/workspaceMcpOverridesService.test.ts b/src/node/services/workspaceMcpOverridesService.test.ts
diff --git a/src/node/services/workspaceMcpOverridesService.ts b/src/node/services/workspaceMcpOverridesService.ts
diff --git a/tests/ipc/setup.ts b/tests/ipc/setup.ts

Original file line number	Diff line number	Diff line change
`@@ -3,7 +3,7 @@ import { z } from "zod";`
`3`	`3`	`/**`
`4`	`4`	`* Per-workspace MCP overrides.`
`5`	`5`	`*`
`6`		`- * Stored in ~/.mux/config.json under each workspace entry.`
	`6`	`+ * Stored per-workspace in <workspace>/.mux/mcp.local.jsonc (workspace-local, intended to be gitignored).`
`7`	`7`	`* Allows workspaces to disable servers or restrict tool allowlists`
`8`	`8`	`* without modifying the project-level .mux/mcp.jsonc.`
`9`	`9`	`*/`