diff --git a/src/features/skill-mcp-manager/connection-race.test.ts b/src/features/skill-mcp-manager/connection-race.test.ts index 784a3e189..10e3c6836 100644 --- a/src/features/skill-mcp-manager/connection-race.test.ts +++ b/src/features/skill-mcp-manager/connection-race.test.ts @@ -248,4 +248,44 @@ describe("getOrCreateClient multi-key disconnect race", () => { expect(state.clients.has(clientKey)).toBe(false) 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() + const supersedingConnection = createDeferred>>() + 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 ? 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() + }) }) diff --git a/src/features/skill-mcp-manager/connection.ts b/src/features/skill-mcp-manager/connection.ts index 3be07d853..890444bce 100644 --- a/src/features/skill-mcp-manager/connection.ts +++ b/src/features/skill-mcp-manager/connection.ts @@ -7,6 +7,13 @@ import { createHttpClient } from "./http-client" import { createStdioClient } from "./stdio-client" 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: { state: SkillMcpManagerState clientKey: string @@ -42,12 +49,13 @@ export async function getOrCreateClient(params: { const isStale = state.pendingConnections.has(clientKey) && state.pendingConnections.get(clientKey) !== currentConnectionPromise if (isStale) { - state.clients.delete(clientKey) + removeClientIfCurrent(state, clientKey, client) try { await client.close() } catch {} throw new Error(`Connection for "${info.sessionID}" was superseded by a newer connection attempt.`) } if (state.shutdownGeneration !== shutdownGenAtStart) { + removeClientIfCurrent(state, clientKey, client) try { await client.close() } catch {} throw new Error(`Shutdown occurred during MCP connection for "${info.sessionID}"`) } diff --git a/src/features/skill-mcp-manager/disconnect-cleanup.test.ts b/src/features/skill-mcp-manager/disconnect-cleanup.test.ts index ab584459b..7c6e2bfb8 100644 --- a/src/features/skill-mcp-manager/disconnect-cleanup.test.ts +++ b/src/features/skill-mcp-manager/disconnect-cleanup.test.ts @@ -14,6 +14,8 @@ afterEach(() => { trackedStates.length = 0 }) +const expectedCleanupHandlerCount = process.platform === "win32" ? 3 : 2 + function createState(): SkillMcpManagerState { const state: SkillMcpManagerState = { clients: new Map(), @@ -82,7 +84,7 @@ describe("disconnectSession cleanup registration", () => { // then expect(state.clients.has("session-2:skill-2:server-2")).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("SIGTERM")).toBe(signalTermCountBeforeRegister + 1) }) @@ -124,7 +126,7 @@ describe("disconnectSession cleanup registration", () => { expect(state.clients.size).toBe(0) expect(state.pendingConnections.size).toBe(1) 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("SIGTERM")).toBe(signalTermCountBeforeRegister + 1) })