fix(skill-mcp-manager): guard stale client cleanup
Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
This commit is contained in:
@@ -248,4 +248,44 @@ describe("getOrCreateClient multi-key disconnect race", () => {
|
|||||||
expect(state.clients.has(clientKey)).toBe(false)
|
expect(state.clients.has(clientKey)).toBe(false)
|
||||||
expect(createdClients[0]?.close).toHaveBeenCalledTimes(1)
|
expect(createdClients[0]?.close).toHaveBeenCalledTimes(1)
|
||||||
})
|
})
|
||||||
|
|
||||||
|
it("#given a superseded pending connection #when a newer client already replaced the map entry #then the stale cleanup does not delete the newer client", async () => {
|
||||||
|
const state = createState()
|
||||||
|
const info = createClientInfo("session-a")
|
||||||
|
const clientKey = createClientKey(info)
|
||||||
|
const pendingConnect = createDeferred<void>()
|
||||||
|
const supersedingConnection = createDeferred<Awaited<ReturnType<typeof getOrCreateClient>>>()
|
||||||
|
pendingConnects.push(pendingConnect)
|
||||||
|
|
||||||
|
const newerClient = new MockClient(
|
||||||
|
{ name: "newer-client", version: "1.0.0" },
|
||||||
|
{ capabilities: {} },
|
||||||
|
)
|
||||||
|
const newerTransport = new MockStdioClientTransport({ command: "mock-mcp-server" })
|
||||||
|
let replacedEntry = false
|
||||||
|
const originalSet = state.clients.set.bind(state.clients)
|
||||||
|
Reflect.set(state.clients, "set", (key: string, value: SkillMcpManagerState["clients"] extends Map<string, infer TValue> ? TValue : never) => {
|
||||||
|
originalSet(key, value)
|
||||||
|
if (!replacedEntry && key === clientKey) {
|
||||||
|
replacedEntry = true
|
||||||
|
originalSet(key, {
|
||||||
|
client: newerClient as never,
|
||||||
|
transport: newerTransport as never,
|
||||||
|
skillName: info.skillName,
|
||||||
|
lastUsedAt: Date.now(),
|
||||||
|
connectionType: "stdio",
|
||||||
|
})
|
||||||
|
}
|
||||||
|
return state.clients
|
||||||
|
})
|
||||||
|
|
||||||
|
const clientPromise = getOrCreateClient({ state, clientKey, info, config: stdioConfig })
|
||||||
|
state.pendingConnections.set(clientKey, supersedingConnection.promise)
|
||||||
|
|
||||||
|
pendingConnect.resolve(undefined)
|
||||||
|
|
||||||
|
await expect(clientPromise).rejects.toThrow(/superseded by a newer connection attempt/)
|
||||||
|
expect(state.clients.get(clientKey)?.client.close).toBe(newerClient.close)
|
||||||
|
expect(newerClient.close).not.toHaveBeenCalled()
|
||||||
|
})
|
||||||
})
|
})
|
||||||
|
|||||||
@@ -7,6 +7,13 @@ import { createHttpClient } from "./http-client"
|
|||||||
import { createStdioClient } from "./stdio-client"
|
import { createStdioClient } from "./stdio-client"
|
||||||
import type { SkillMcpClientConnectionParams, SkillMcpClientInfo, SkillMcpManagerState } from "./types"
|
import type { SkillMcpClientConnectionParams, SkillMcpClientInfo, SkillMcpManagerState } from "./types"
|
||||||
|
|
||||||
|
function removeClientIfCurrent(state: SkillMcpManagerState, clientKey: string, client: Client): void {
|
||||||
|
const managed = state.clients.get(clientKey)
|
||||||
|
if (managed?.client === client) {
|
||||||
|
state.clients.delete(clientKey)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
export async function getOrCreateClient(params: {
|
export async function getOrCreateClient(params: {
|
||||||
state: SkillMcpManagerState
|
state: SkillMcpManagerState
|
||||||
clientKey: string
|
clientKey: string
|
||||||
@@ -42,12 +49,13 @@ export async function getOrCreateClient(params: {
|
|||||||
|
|
||||||
const isStale = state.pendingConnections.has(clientKey) && state.pendingConnections.get(clientKey) !== currentConnectionPromise
|
const isStale = state.pendingConnections.has(clientKey) && state.pendingConnections.get(clientKey) !== currentConnectionPromise
|
||||||
if (isStale) {
|
if (isStale) {
|
||||||
state.clients.delete(clientKey)
|
removeClientIfCurrent(state, clientKey, client)
|
||||||
try { await client.close() } catch {}
|
try { await client.close() } catch {}
|
||||||
throw new Error(`Connection for "${info.sessionID}" was superseded by a newer connection attempt.`)
|
throw new Error(`Connection for "${info.sessionID}" was superseded by a newer connection attempt.`)
|
||||||
}
|
}
|
||||||
|
|
||||||
if (state.shutdownGeneration !== shutdownGenAtStart) {
|
if (state.shutdownGeneration !== shutdownGenAtStart) {
|
||||||
|
removeClientIfCurrent(state, clientKey, client)
|
||||||
try { await client.close() } catch {}
|
try { await client.close() } catch {}
|
||||||
throw new Error(`Shutdown occurred during MCP connection for "${info.sessionID}"`)
|
throw new Error(`Shutdown occurred during MCP connection for "${info.sessionID}"`)
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -14,6 +14,8 @@ afterEach(() => {
|
|||||||
trackedStates.length = 0
|
trackedStates.length = 0
|
||||||
})
|
})
|
||||||
|
|
||||||
|
const expectedCleanupHandlerCount = process.platform === "win32" ? 3 : 2
|
||||||
|
|
||||||
function createState(): SkillMcpManagerState {
|
function createState(): SkillMcpManagerState {
|
||||||
const state: SkillMcpManagerState = {
|
const state: SkillMcpManagerState = {
|
||||||
clients: new Map(),
|
clients: new Map(),
|
||||||
@@ -82,7 +84,7 @@ describe("disconnectSession cleanup registration", () => {
|
|||||||
// then
|
// then
|
||||||
expect(state.clients.has("session-2:skill-2:server-2")).toBe(true)
|
expect(state.clients.has("session-2:skill-2:server-2")).toBe(true)
|
||||||
expect(state.cleanupRegistered).toBe(true)
|
expect(state.cleanupRegistered).toBe(true)
|
||||||
expect(state.cleanupHandlers).toHaveLength(2)
|
expect(state.cleanupHandlers).toHaveLength(expectedCleanupHandlerCount)
|
||||||
expect(process.listenerCount("SIGINT")).toBe(signalIntCountBeforeRegister + 1)
|
expect(process.listenerCount("SIGINT")).toBe(signalIntCountBeforeRegister + 1)
|
||||||
expect(process.listenerCount("SIGTERM")).toBe(signalTermCountBeforeRegister + 1)
|
expect(process.listenerCount("SIGTERM")).toBe(signalTermCountBeforeRegister + 1)
|
||||||
})
|
})
|
||||||
@@ -124,7 +126,7 @@ describe("disconnectSession cleanup registration", () => {
|
|||||||
expect(state.clients.size).toBe(0)
|
expect(state.clients.size).toBe(0)
|
||||||
expect(state.pendingConnections.size).toBe(1)
|
expect(state.pendingConnections.size).toBe(1)
|
||||||
expect(state.cleanupRegistered).toBe(true)
|
expect(state.cleanupRegistered).toBe(true)
|
||||||
expect(state.cleanupHandlers).toHaveLength(2)
|
expect(state.cleanupHandlers).toHaveLength(expectedCleanupHandlerCount)
|
||||||
expect(process.listenerCount("SIGINT")).toBe(signalIntCountBeforeRegister + 1)
|
expect(process.listenerCount("SIGINT")).toBe(signalIntCountBeforeRegister + 1)
|
||||||
expect(process.listenerCount("SIGTERM")).toBe(signalTermCountBeforeRegister + 1)
|
expect(process.listenerCount("SIGTERM")).toBe(signalTermCountBeforeRegister + 1)
|
||||||
})
|
})
|
||||||
|
|||||||
Reference in New Issue
Block a user