Merge pull request #4099 from sjawhar/fix/mcp-reload-survival

fix(skill-mcp): allow MCP manager to accept connections after disconnectAll

src/features/skill-mcp-manager/cleanup.ts

@@ -127,6 +127,10 @@ export async function disconnectSession(state: SkillMcpManagerState, sessionID:
 
 export async function disconnectAll(state: SkillMcpManagerState): Promise<void> {
   state.shutdownGeneration++
+  // Temporarily block new connections during cleanup. Reset at the end so that
+  // sessions surviving a plugin reload can reconnect. (Plugin reload calls
+  // disconnectAll via plugin-dispose, but existing sessions' tool closures still
+  // reference this manager instance and must be able to create new connections.)
   state.disposed = true
   stopCleanupTimer(state)
   unregisterProcessCleanup(state)
@@ -141,6 +145,8 @@ export async function disconnectAll(state: SkillMcpManagerState): Promise<void>
   for (const managed of clients) {
     await closeManagedClient(managed)
   }
+
+  state.disposed = false
 }
 
 export async function forceReconnect(state: SkillMcpManagerState, clientKey: string): Promise<boolean> {

src/features/skill-mcp-manager/connection-race.test.ts

@@ -202,20 +202,22 @@ describe("getOrCreateClient disconnectAll race", () => {
     expect(state.clients.has(clientKey)).toBe(false)
   })
 
-  it("#given state after disconnectAll() completed #when getOrCreateClient() is called #then it throws shut down error and registers nothing", async () => {
+  it("#given state after disconnectAll() completed #when getOrCreateClient() is called #then it creates a new client and clears the temporary disposed flag", async () => {
     const state = createState()
     const info = createClientInfo("session-b")
     const clientKey = createClientKey(info)
 
     await disconnectAll(state)
 
-    await expect(getOrCreateClient({ state, clientKey, info, config: stdioConfig })).rejects.toThrow(/has been shut down/)
-    expect(state.clients.size).toBe(0)
+    const client = await getOrCreateClient({ state, clientKey, info, config: stdioConfig })
+
+    expect(client).toBeDefined()
+    expect(state.clients.get(clientKey)?.client).toBe(client)
     expect(state.pendingConnections.size).toBe(0)
     expect(state.inFlightConnections.size).toBe(0)
-    expect(state.disposed).toBe(true)
-    expect(createdClients).toHaveLength(0)
-    expect(createdTransports).toHaveLength(0)
+    expect(state.disposed).toBe(false)
+    expect(createdClients).toHaveLength(1)
+    expect(createdTransports).toHaveLength(1)
   })
 })
 

src/features/skill-mcp-manager/plugin-reload-mcp-survival.test.ts

@@ -0,0 +1,155 @@
+import { afterEach, beforeEach, describe, expect, it, mock } from "bun:test"
+import type { ClaudeCodeMcpServer } from "../claude-code-mcp-loader/types"
+import type { McpClient, McpTransport, SkillMcpClientInfo, SkillMcpManagerState } from "./types"
+import { disconnectAll } from "./cleanup"
+import { getOrCreateClient } from "./connection"
+import { setStdioClientDependenciesForTesting } from "./stdio-client"
+
+/**
+ * This test proves a real-world bug: when the OMO plugin reloads (triggered by a
+ * new OpenCode session opening in a different directory), `disconnectAll()` is called
+ * on the SkillMcpManager via `plugin-dispose`. This sets `state.disposed = true`
+ * permanently, so every subsequent `skill_mcp()` call from sessions that were
+ * active BEFORE the reload fails with "has been shut down, cannot create new
+ * connections."
+ *
+ * The plugin function (`OhMyOpenCodePlugin` in index.ts) is called once per session
+ * directory. Each call disposes the previous plugin via a module-level singleton
+ * `activePluginDispose`. Since the SkillMcpManager is recreated each time, the NEW
+ * session gets a fresh manager — but OLD sessions still hold closures over the
+ * disposed manager.
+ *
+ * The desired behavior: cleaning up all connections (e.g. during a plugin reload)
+ * should NOT permanently prevent the manager from accepting new connections.
+ * Sessions that survive a reload should be able to reconnect.
+ */
+
+const trackedStates: SkillMcpManagerState[] = []
+
+function createMockClient(): McpClient {
+  return {
+    close: mock(async () => {}),
+    connect: mock(async () => {}),
+  } as unknown as McpClient
+}
+
+function createMockTransport(): McpTransport {
+  return {
+    close: mock(async () => {}),
+  } as unknown as McpTransport
+}
+
+function createState(): SkillMcpManagerState {
+  const state: SkillMcpManagerState = {
+    clients: new Map(),
+    pendingConnections: new Map(),
+    disconnectedSessions: new Map(),
+    authProviders: new Map(),
+    cleanupRegistered: false,
+    cleanupInterval: null,
+    cleanupHandlers: [],
+    idleTimeoutMs: 5 * 60 * 1000,
+    shutdownGeneration: 0,
+    inFlightConnections: new Map(),
+    disposed: false,
+  }
+
+  trackedStates.push(state)
+  return state
+}
+
+function createClientInfo(sessionID: string): SkillMcpClientInfo {
+  return {
+    serverName: "whatsapp",
+    skillName: "whatsapp-skill",
+    sessionID,
+  }
+}
+
+function createClientKey(info: SkillMcpClientInfo): string {
+  return `${info.sessionID}:${info.skillName}:${info.serverName}`
+}
+
+const stdioConfig: ClaudeCodeMcpServer = {
+  command: "mock-whatsapp-mcp",
+}
+
+beforeEach(() => {
+  setStdioClientDependenciesForTesting({
+    createClient: () => createMockClient(),
+    createTransport: () => createMockTransport(),
+  })
+})
+
+afterEach(async () => {
+  setStdioClientDependenciesForTesting()
+  for (const state of trackedStates) {
+    state.disposed = false
+    for (const managed of state.clients.values()) {
+      try {
+        await managed.client.close()
+      } catch {}
+      try {
+        await managed.transport.close()
+      } catch {}
+    }
+    state.clients.clear()
+    state.pendingConnections.clear()
+  }
+  trackedStates.length = 0
+})
+
+describe("MCP manager survival across plugin reload", () => {
+  it("#given session A has an active MCP connection #when plugin reloads (disconnectAll) #then session A can still create new connections", async () => {
+    // given: session A established an MCP connection (e.g. WhatsApp)
+    const state = createState()
+    const sessionAInfo = createClientInfo("ses_session_a")
+    const clientKey = createClientKey(sessionAInfo)
+
+    const initialClient = await getOrCreateClient({
+      state,
+      clientKey,
+      info: sessionAInfo,
+      config: stdioConfig,
+    })
+    expect(initialClient).toBeDefined()
+    expect(state.clients.has(clientKey)).toBe(true)
+
+    // when: plugin reloads because a new session opened in a different directory.
+    // In production, index.ts line 37 calls `await activePluginDispose?.()` which
+    // calls `skillMcpManager.disconnectAll()` via plugin-dispose.ts line 33.
+    // This sets state.disposed = true permanently.
+    await disconnectAll(state)
+
+    // then: session A should be able to reconnect.
+    // The old session's tools still reference this manager instance — there is no
+    // mechanism for OpenCode to replace tool closures in existing sessions after a
+    // plugin reload. So the manager must accept new connections.
+    const reconnectedClient = await getOrCreateClient({
+      state,
+      clientKey,
+      info: sessionAInfo,
+      config: stdioConfig,
+    })
+    expect(reconnectedClient).toBeDefined()
+  })
+
+  it("#given no prior connections #when disconnectAll was called (plugin reload) #then new connections should still be possible", async () => {
+    // given: a manager that was part of a previous plugin load cycle
+    const state = createState()
+    await disconnectAll(state)
+
+    // when: a session that survived the reload tries to use MCP
+    const info = createClientInfo("ses_surviving_session")
+    const clientKey = createClientKey(info)
+
+    // then: it should succeed, not throw "has been shut down"
+    const client = await getOrCreateClient({
+      state,
+      clientKey,
+      info,
+      config: stdioConfig,
+    })
+    expect(client).toBeDefined()
+  })
+})