fix(skill-mcp): allow MCP manager to accept connections after disconnectAll
Also fix CI test isolation: run each test file in its own bun process to prevent
mock.module() leaks across files that were introduced by upstream commit 6acca09b.
This commit is contained in:
@@ -127,6 +127,10 @@ export async function disconnectSession(state: SkillMcpManagerState, sessionID:
|
|||||||
|
|
||||||
export async function disconnectAll(state: SkillMcpManagerState): Promise<void> {
|
export async function disconnectAll(state: SkillMcpManagerState): Promise<void> {
|
||||||
state.shutdownGeneration++
|
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
|
state.disposed = true
|
||||||
stopCleanupTimer(state)
|
stopCleanupTimer(state)
|
||||||
unregisterProcessCleanup(state)
|
unregisterProcessCleanup(state)
|
||||||
@@ -141,6 +145,8 @@ export async function disconnectAll(state: SkillMcpManagerState): Promise<void>
|
|||||||
for (const managed of clients) {
|
for (const managed of clients) {
|
||||||
await closeManagedClient(managed)
|
await closeManagedClient(managed)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
state.disposed = false
|
||||||
}
|
}
|
||||||
|
|
||||||
export async function forceReconnect(state: SkillMcpManagerState, clientKey: string): Promise<boolean> {
|
export async function forceReconnect(state: SkillMcpManagerState, clientKey: string): Promise<boolean> {
|
||||||
|
|||||||
@@ -202,20 +202,22 @@ describe("getOrCreateClient disconnectAll race", () => {
|
|||||||
expect(state.clients.has(clientKey)).toBe(false)
|
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 state = createState()
|
||||||
const info = createClientInfo("session-b")
|
const info = createClientInfo("session-b")
|
||||||
const clientKey = createClientKey(info)
|
const clientKey = createClientKey(info)
|
||||||
|
|
||||||
await disconnectAll(state)
|
await disconnectAll(state)
|
||||||
|
|
||||||
await expect(getOrCreateClient({ state, clientKey, info, config: stdioConfig })).rejects.toThrow(/has been shut down/)
|
const client = await getOrCreateClient({ state, clientKey, info, config: stdioConfig })
|
||||||
expect(state.clients.size).toBe(0)
|
|
||||||
|
expect(client).toBeDefined()
|
||||||
|
expect(state.clients.get(clientKey)?.client).toBe(client)
|
||||||
expect(state.pendingConnections.size).toBe(0)
|
expect(state.pendingConnections.size).toBe(0)
|
||||||
expect(state.inFlightConnections.size).toBe(0)
|
expect(state.inFlightConnections.size).toBe(0)
|
||||||
expect(state.disposed).toBe(true)
|
expect(state.disposed).toBe(false)
|
||||||
expect(createdClients).toHaveLength(0)
|
expect(createdClients).toHaveLength(1)
|
||||||
expect(createdTransports).toHaveLength(0)
|
expect(createdTransports).toHaveLength(1)
|
||||||
})
|
})
|
||||||
})
|
})
|
||||||
|
|
||||||
|
|||||||
@@ -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()
|
||||||
|
})
|
||||||
|
})
|
||||||
Reference in New Issue
Block a user