feat(core): add runtime MCP controls (#37308)

Co-authored-by: Aiden Cline <rekram1-node@users.noreply.github.com>
Co-authored-by: Aiden Cline <aidenpcline@gmail.com>

packages/core/src/mcp/index.ts

@@ -13,8 +13,10 @@ import { EventV2 } from "../event"
 import { Form } from "../form"
 import { Integration } from "../integration"
 import { IntegrationConnection } from "../integration/connection"
+import { KeyedMutex } from "../effect/keyed-mutex"
 import { Location } from "../location"
 import { waitForAbort } from "../process"
+import { State } from "../state"
 import { MCPClient } from "./client"
 import { MCPOAuth } from "./oauth"
 
@@ -118,6 +120,7 @@ type ServerEntry = {
   prompts?: ReadonlyArray<Prompt>
   // Set when a remote server is registered as an OAuth integration; the credential lives in the global store.
   integrationID?: Integration.ID
+  registration?: State.Registration
 }
 
 // MCP elicitations are Location-scoped, not Session-scoped: the server cannot attribute them to a
@@ -127,6 +130,10 @@ const URL_ELICITATION_FIELD_KEY = "elicitation"
 
 export interface Interface {
   readonly servers: () => Effect.Effect<ServerInfo[]>
+  readonly add: (server: ServerName | string, config: typeof ConfigMCP.Server.Type) => Effect.Effect<void>
+  readonly connect: (server: ServerName | string) => Effect.Effect<void, NotFoundError>
+  readonly disconnect: (server: ServerName | string) => Effect.Effect<void, NotFoundError>
+  readonly remove: (server: ServerName | string) => Effect.Effect<void, NotFoundError>
   readonly tools: () => Effect.Effect<Tool[]>
   readonly callTool: (input: {
     readonly server: ServerName | string
@@ -170,6 +177,9 @@ export const layer = Layer.effect(
     )
     // Later config files win for duplicate server names; per-server timeout overrides globals.
     const runtime = new Map<ServerName, ServerEntry>()
+    // Serializes lifecycle operations per server. Anything taking this lock from a connection
+    // callback must stay forked: lifecycle operations close scopes while holding it, firing onClose.
+    const locks = KeyedMutex.makeUnsafe<ServerName>()
     const urlElicitations = new Map<string, Form.ID>()
     for (const entry of documents) {
       for (const [name, server] of Object.entries(entry.info.mcp?.servers ?? {})) {
@@ -183,14 +193,9 @@ export const layer = Layer.effect(
 
     // Register every remote server as an OAuth integration so credentials live in the global store
     // rather than in committed config. Servers that connect anonymously simply never use the method.
-    const registrations: Array<{
-      readonly name: ServerName
-      readonly remote: typeof ConfigMCP.Remote.Type
-      readonly integrationID: Integration.ID
-      readonly methodID: Integration.MethodID
-    }> = []
-    for (const [name, entry] of runtime) {
-      if (entry.config.type !== "remote" || entry.config.oauth === false) continue
+    const owned = new Set<Integration.ID>()
+    const register = Effect.fnUntraced(function* (name: ServerName, entry: ServerEntry) {
+      if (entry.config.type !== "remote" || entry.config.oauth === false) return
       const remote = entry.config
       // Key identity on name + url, not url alone: two configs for the same url under different names are
       // distinct logical servers that may hold different accounts, so they must not share a credential row.
@@ -200,27 +205,24 @@ export const layer = Layer.effect(
           .update(name + "\u0000" + remote.url)
           .digest("hex")
           .slice(0, 16)
-      entry.integrationID = Integration.ID.make(suffix)
-      registrations.push({
-        name,
-        remote,
-        integrationID: entry.integrationID,
-        methodID: Integration.MethodID.make(suffix),
-      })
-    }
-    if (registrations.length > 0)
-      yield* integration.transform((draft) => {
-        for (const reg of registrations) {
-          draft.update(reg.integrationID, (ref) => {
-            ref.name = reg.name
+      const integrationID = Integration.ID.make(suffix)
+      entry.integrationID = integrationID
+      owned.add(integrationID)
+      const methodID = Integration.MethodID.make(suffix)
+      entry.registration = yield* integration
+        .transform((draft) => {
+          draft.update(integrationID, (ref) => {
+            ref.name = name
           })
           draft.method.update({
-            integrationID: reg.integrationID,
-            method: { id: reg.methodID, type: "oauth", label: reg.name },
-            authorize: () => MCPOAuth.authorize({ name: reg.name, config: reg.remote, methodID: reg.methodID }),
+            integrationID,
+            method: { id: methodID, type: "oauth", label: name },
+            authorize: () => MCPOAuth.authorize({ name, config: remote, methodID }),
           })
-        }
-      })
+        })
+        .pipe(Scope.provide(root))
+    })
+    yield* Effect.forEach(runtime, ([name, entry]) => register(name, entry), { discard: true })
 
     const requireServer = Effect.fnUntraced(function* (server: ServerName | string) {
       const name = ServerName.make(server)
@@ -420,36 +422,44 @@ export const layer = Layer.effect(
         ),
       )
 
-    const watch = (name: ServerName, entry: ServerEntry, connection: MCPClient.Connection) => {
-      connection.onClose(() => {
-        // A reconnect closes the previous scope, but the SDK may fire this onclose after the new
-        // connection is already assigned; ignore the stale close so it can't null out the live client.
-        if (entry.client !== connection) return
-        entry.client = undefined
-        entry.tools = undefined
-        entry.prompts = undefined
-        entry.status = { status: "failed", error: "Connection closed" }
-        fork(events.publish(McpEvent.ToolsChanged, { server: name }).pipe(Effect.ignore))
-        fork(events.publish(McpEvent.ResourcesChanged, { server: name }).pipe(Effect.ignore))
-        fork(events.publish(Command.Event.Updated, {}).pipe(Effect.ignore))
-        fork(events.publish(McpEvent.StatusChanged, { server: name }).pipe(Effect.ignore))
-      })
-      connection.onLog((message) => fork(serverLog(name, message).pipe(Effect.ignore)))
-      connection.onToolsChanged(() => {
+    // Runs a connection callback under the server lock, dropping it if the connection is no longer
+    // the entry's live client, so late SDK callbacks cannot commit obsolete state.
+    const whenLive =
+      (name: ServerName, entry: ServerEntry, connection: MCPClient.Connection) =>
+      <E>(effect: Effect.Effect<void, E>) =>
         fork(
-          refreshTools(name, entry, connection).pipe(
-            Effect.andThen(events.publish(McpEvent.ToolsChanged, { server: name })),
+          Effect.suspend(() => (entry.client === connection ? effect : Effect.void)).pipe(
+            locks.withLock(name),
             Effect.ignore,
           ),
         )
-      })
-      connection.onPromptsChanged(() => {
-        fork(refreshPrompts(name, entry, connection).pipe(Effect.ignore))
-      })
-      connection.onResourcesChanged(() => {
-        if (entry.client !== connection) return
-        fork(events.publish(McpEvent.ResourcesChanged, { server: name }).pipe(Effect.ignore))
-      })
+
+    const watch = (name: ServerName, entry: ServerEntry, connection: MCPClient.Connection) => {
+      const live = whenLive(name, entry, connection)
+      connection.onClose(() =>
+        live(
+          Effect.gen(function* () {
+            entry.client = undefined
+            entry.tools = undefined
+            entry.prompts = undefined
+            entry.status = { status: "failed", error: "Connection closed" }
+            yield* events.publish(McpEvent.ToolsChanged, { server: name }).pipe(Effect.ignore)
+            yield* events.publish(McpEvent.ResourcesChanged, { server: name }).pipe(Effect.ignore)
+            yield* events.publish(Command.Event.Updated, {}).pipe(Effect.ignore)
+            yield* events.publish(McpEvent.StatusChanged, { server: name }).pipe(Effect.ignore)
+          }),
+        ),
+      )
+      connection.onLog((message) => fork(serverLog(name, message).pipe(Effect.ignore)))
+      connection.onToolsChanged(() =>
+        live(
+          refreshTools(name, entry, connection).pipe(
+            Effect.andThen(events.publish(McpEvent.ToolsChanged, { server: name })),
+          ),
+        ),
+      )
+      connection.onPromptsChanged(() => live(refreshPrompts(name, entry, connection)))
+      connection.onResourcesChanged(() => live(events.publish(McpEvent.ResourcesChanged, { server: name })))
     }
 
     const serverLog = (server: ServerName, message: MCPClient.LogMessage) => {
@@ -495,7 +505,7 @@ export const layer = Layer.effect(
           yield* events.publish(McpEvent.ToolsChanged, { server: name }).pipe(Effect.ignore)
           yield* events.publish(McpEvent.ResourcesChanged, { server: name }).pipe(Effect.ignore)
           yield* events.publish(McpEvent.StatusChanged, { server: name }).pipe(Effect.ignore)
-          fork(refreshPrompts(name, entry, result.value.connection).pipe(Effect.ignore))
+          whenLive(name, entry, result.value.connection)(refreshPrompts(name, entry, result.value.connection))
           return
         }
         yield* Scope.close(scope, Exit.void)
@@ -509,6 +519,19 @@ export const layer = Layer.effect(
         yield* events.publish(McpEvent.StatusChanged, { server: name }).pipe(Effect.ignore)
       }).pipe(Effect.ensuring(Deferred.succeed(entry.startup, undefined)))
 
+    const stopServer = Effect.fnUntraced(function* (name: ServerName, entry: ServerEntry) {
+      const scope = entry.scope
+      if (!scope) return
+      entry.scope = undefined
+      entry.client = undefined
+      entry.tools = undefined
+      entry.prompts = undefined
+      yield* Scope.close(scope, Exit.void)
+      yield* events.publish(McpEvent.ToolsChanged, { server: name }).pipe(Effect.ignore)
+      yield* events.publish(McpEvent.ResourcesChanged, { server: name }).pipe(Effect.ignore)
+      yield* events.publish(Command.Event.Updated, {}).pipe(Effect.ignore)
+    })
+
     // Disabled servers settle their startup immediately so queries never block on them.
     for (const [name, entry] of runtime) {
       if (entry.config.disabled) {
@@ -516,27 +539,24 @@ export const layer = Layer.effect(
         Deferred.doneUnsafe(entry.startup, Exit.void)
         continue
       }
-      fork(startServer(name, entry))
+      fork(startServer(name, entry).pipe(locks.withLock(name)))
     }
 
     // Bring a server online (or back to needs_auth) when its integration's credential changes, so an
     // OAuth login takes effect without a restart. Only fires for the integrations we registered.
-    const owned = new Set(registrations.map((reg) => reg.integrationID))
     const reconnect = (integrationID: Integration.ID) =>
       Effect.gen(function* () {
         const match = Array.from(runtime).find(([, entry]) => entry.integrationID === integrationID)
         if (!match) return
-        const [name, entry] = match
-        if (entry.config.disabled) return
-        if (entry.scope) {
-          yield* Scope.close(entry.scope, Exit.void)
-          entry.scope = undefined
-          entry.client = undefined
-          entry.tools = undefined
-          entry.prompts = undefined
-          yield* events.publish(Command.Event.Updated, {}).pipe(Effect.ignore)
-        }
-        yield* startServer(name, entry)
+        const name = match[0]
+        yield* Effect.gen(function* () {
+          // add() or remove() may have replaced or deleted the entry while we waited for the lock.
+          const entry = runtime.get(name)
+          if (!entry || entry.integrationID !== integrationID) return
+          if (entry.status.status === "disabled") return
+          yield* stopServer(name, entry)
+          yield* startServer(name, entry)
+        }).pipe(locks.withLock(name))
       })
     fork(
       events.subscribe(Integration.Event.ConnectionUpdated).pipe(
@@ -546,10 +566,13 @@ export const layer = Layer.effect(
       ),
     )
 
-    const whenAllReady = Effect.forEach(runtime.values(), (entry) => Deferred.await(entry.startup), {
-      concurrency: "unbounded",
-      discard: true,
-    })
+    // Suspend so each await sees current entries; a bare Map iterator is exhausted after one run.
+    const whenAllReady = Effect.suspend(() =>
+      Effect.forEach(Array.from(runtime.values()), (entry) => Deferred.await(entry.startup), {
+        concurrency: "unbounded",
+        discard: true,
+      }),
+    )
     return Service.of({
       servers: Effect.fn("MCP.servers")(function* () {
         const entries = Array.from(runtime).toSorted(([a], [b]) => a.localeCompare(b))
@@ -562,6 +585,66 @@ export const layer = Layer.effect(
           }),
         )
       }),
+      add: Effect.fn("MCP.add")(function* (server, config) {
+        const name = ServerName.make(server)
+        yield* Effect.gen(function* () {
+          const previous = runtime.get(name)
+          if (previous) {
+            yield* stopServer(name, previous)
+            if (previous.integrationID) owned.delete(previous.integrationID)
+            if (previous.registration) yield* previous.registration.dispose
+          }
+          const entry: ServerEntry = {
+            config: { ...config, timeout: { ...timeout, ...config.timeout } },
+            status: { status: "pending" },
+            startup: Deferred.makeUnsafe<void>(),
+          }
+          runtime.set(name, entry)
+          yield* Effect.gen(function* () {
+            yield* register(name, entry)
+            if (config.disabled) {
+              entry.status = { status: "disabled" }
+              yield* events.publish(McpEvent.StatusChanged, { server: name }).pipe(Effect.ignore)
+              return
+            }
+            yield* startServer(name, entry)
+          }).pipe(
+            // Settle startup even when register fails or add is interrupted, so an entry that made it
+            // into runtime can never hang readers awaiting its startup.
+            Effect.ensuring(Effect.sync(() => Deferred.doneUnsafe(entry.startup, Exit.void))),
+          )
+        }).pipe(locks.withLock(name))
+      }),
+      connect: Effect.fn("MCP.connect")(function* (server) {
+        const name = ServerName.make(server)
+        yield* Effect.gen(function* () {
+          const target = yield* requireServer(name)
+          yield* stopServer(name, target.entry)
+          yield* startServer(name, target.entry)
+        }).pipe(locks.withLock(name))
+      }),
+      disconnect: Effect.fn("MCP.disconnect")(function* (server) {
+        const name = ServerName.make(server)
+        yield* Effect.gen(function* () {
+          const target = yield* requireServer(name)
+          yield* stopServer(name, target.entry)
+          target.entry.status = { status: "disabled" }
+          yield* events.publish(McpEvent.StatusChanged, { server: name }).pipe(Effect.ignore)
+        }).pipe(locks.withLock(name))
+      }),
+      remove: Effect.fn("MCP.remove")(function* (server) {
+        const name = ServerName.make(server)
+        yield* Effect.gen(function* () {
+          const target = yield* requireServer(name)
+          yield* stopServer(name, target.entry)
+          if (target.entry.integrationID) owned.delete(target.entry.integrationID)
+          if (target.entry.registration) yield* target.entry.registration.dispose
+          // Credentials are kept: they are keyed by name + url, so re-adding the same server
+          // reuses them without forcing re-auth, matching add()'s replacement semantics.
+          runtime.delete(name)
+          yield* events.publish(McpEvent.StatusChanged, { server: name }).pipe(Effect.ignore)
+        }).pipe(locks.withLock(name))
+      }),
       tools: Effect.fn("MCP.tools")(function* () {
         yield* whenAllReady
         return Array.from(runtime.values())

packages/core/test/fixture/mcp.ts

@@ -9,6 +9,10 @@ export const emptyMcpLayer = Layer.succeed(
   MCP.Service,
   MCP.Service.of({
     servers: () => Effect.succeed([]),
+    add: () => Effect.die("unused mcp.add"),
+    connect: () => Effect.die("unused mcp.connect"),
+    disconnect: () => Effect.die("unused mcp.disconnect"),
+    remove: () => Effect.die("unused mcp.remove"),
     tools: () => Effect.succeed([]),
     callTool: () => Effect.die("unused mcp.callTool"),
     instructions: () => Effect.succeed([]),

packages/core/test/mcp.test.ts

@@ -148,7 +148,10 @@ function resourceServer(
   )
 }
 
-function resourceMcpLayer(url: string, onFormCreated?: (form: Form.Info) => Effect.Effect<void>) {
+function resourceMcpLayer(
+  server: string | typeof ConfigMCP.Server.Type,
+  onFormCreated?: (form: Form.Info) => Effect.Effect<void>,
+) {
   const directory = AbsolutePath.make(import.meta.dir)
   const unusedIntegration = () => Effect.die("unused integration service")
   return MCP.layer.pipe(
@@ -164,7 +167,12 @@ function resourceMcpLayer(url: string, onFormCreated?: (form: Form.Info) => Effe
                   type: "document",
                   info: new Config.Info({
                     mcp: new ConfigMCP.Info({
-                      servers: { resources: new ConfigMCP.Remote({ type: "remote", url, oauth: false }) },
+                      servers: {
+                        resources:
+                          typeof server === "string"
+                            ? new ConfigMCP.Remote({ type: "remote", url: server, oauth: false })
+                            : server,
+                      },
                     }),
                   }),
                 }),
@@ -634,6 +642,120 @@ test("loads and reads MCP resources", async () => {
   )
 })
 
+test("adds, disconnects, and reconnects MCP servers at runtime", async () => {
+  await Effect.runPromise(
+    Effect.scoped(
+      Effect.gen(function* () {
+        yield* Effect.gen(function* () {
+          const service = yield* MCP.Service
+
+          expect((yield* service.servers())[0]?.status).toEqual({ status: "disabled" })
+          expect(yield* service.connect("missing").pipe(Effect.flip)).toBeInstanceOf(MCP.NotFoundError)
+          expect(yield* service.disconnect("missing").pipe(Effect.flip)).toBeInstanceOf(MCP.NotFoundError)
+          yield* service.add(
+            "dynamic",
+            new ConfigMCP.Local({
+              type: "local",
+              command: [process.execPath, path.join(import.meta.dir, "fixture/mcp-output-schema.ts")],
+            }),
+          )
+          expect((yield* service.servers()).find((server) => server.name === "dynamic")?.status).toEqual({
+            status: "connected",
+          })
+
+          yield* service.add(
+            "dynamic",
+            new ConfigMCP.Local({
+              type: "local",
+              command: [process.execPath, path.join(import.meta.dir, "fixture/mcp-output-schema.ts")],
+              disabled: true,
+            }),
+          )
+          expect((yield* service.servers()).find((server) => server.name === "dynamic")?.status).toEqual({
+            status: "disabled",
+          })
+          expect(yield* service.tools()).toEqual([])
+
+          yield* service.connect("dynamic")
+          expect((yield* service.servers()).find((server) => server.name === "dynamic")?.status).toEqual({
+            status: "connected",
+          })
+          yield* service.disconnect("dynamic")
+          expect((yield* service.servers()).find((server) => server.name === "dynamic")?.status).toEqual({
+            status: "disabled",
+          })
+          expect(yield* service.tools()).toEqual([])
+
+          yield* service.connect("dynamic")
+          expect((yield* service.servers()).find((server) => server.name === "dynamic")?.status).toEqual({
+            status: "connected",
+          })
+
+          yield* service.remove("dynamic")
+          expect((yield* service.servers()).some((server) => server.name === "dynamic")).toBe(false)
+          expect(yield* service.tools()).toEqual([])
+          expect(yield* service.remove("dynamic").pipe(Effect.flip)).toBeInstanceOf(MCP.NotFoundError)
+        }).pipe(
+          Effect.provide(
+            resourceMcpLayer(
+              new ConfigMCP.Local({
+                type: "local",
+                command: [process.execPath, path.join(import.meta.dir, "fixture/mcp-output-schema.ts")],
+                disabled: true,
+              }),
+            ),
+          ),
+        )
+      }),
+    ),
+  )
+})
+
+test("serializes concurrent MCP lifecycle operations", async () => {
+  await Effect.runPromise(
+    Effect.scoped(
+      Effect.gen(function* () {
+        yield* Effect.gen(function* () {
+          const service = yield* MCP.Service
+
+          // Whatever order the racing operations land in, the resulting state must be consistent.
+          yield* Effect.all(
+            [
+              service.connect("resources"),
+              service.connect("resources"),
+              service.disconnect("resources"),
+              service.connect("resources"),
+            ],
+            { concurrency: "unbounded", discard: true },
+          )
+          const status = (yield* service.servers()).find((server) => server.name === "resources")?.status
+          const tools = yield* service.tools()
+          expect(status?.status === "connected" || status?.status === "disabled").toBe(true)
+          if (status?.status === "disabled") expect(tools).toEqual([])
+          if (status?.status === "connected") expect(tools.length).toBeGreaterThan(0)
+
+          yield* service.disconnect("resources")
+          expect((yield* service.servers())[0]?.status).toEqual({ status: "disabled" })
+          expect(yield* service.tools()).toEqual([])
+          yield* service.connect("resources")
+          expect((yield* service.servers())[0]?.status).toEqual({ status: "connected" })
+          expect((yield* service.tools()).length).toBeGreaterThan(0)
+        }).pipe(
+          Effect.provide(
+            resourceMcpLayer(
+              new ConfigMCP.Local({
+                type: "local",
+                command: [process.execPath, path.join(import.meta.dir, "fixture/mcp-output-schema.ts")],
+                disabled: true,
+              }),
+            ),
+          ),
+        )
+      }),
+    ),
+  )
+})
+
 it.effect("advertises MCP output schemas to Code Mode", () =>
   Effect.gen(function* () {
     const registry = yield* ToolRegistry.Service