From 0ebe1d5b1b81fa2fffce4cdf21ed8c1f319b81e1 Mon Sep 17 00:00:00 2001 From: YeonGyu-Kim Date: Thu, 12 Mar 2026 00:28:19 +0900 Subject: [PATCH] fix: address review-work round 6 findings (dispose isolation, event dispatch, disconnectedSessions ref-counting) --- src/features/skill-mcp-manager/cleanup.ts | 1 + .../skill-mcp-manager/connection-race.test.ts | 8 ++- src/features/skill-mcp-manager/connection.ts | 8 +++ .../disconnect-cleanup.test.ts | 1 + src/features/skill-mcp-manager/manager.ts | 1 + src/features/skill-mcp-manager/types.ts | 1 + src/plugin-dispose.test.ts | 56 +++++++++++++++++++ src/plugin-dispose.ts | 20 ++++++- src/plugin/event.ts | 1 + 9 files changed, 91 insertions(+), 6 deletions(-) diff --git a/src/features/skill-mcp-manager/cleanup.ts b/src/features/skill-mcp-manager/cleanup.ts index 500294123..c217826f6 100644 --- a/src/features/skill-mcp-manager/cleanup.ts +++ b/src/features/skill-mcp-manager/cleanup.ts @@ -134,6 +134,7 @@ export async function disconnectAll(state: SkillMcpManagerState): Promise state.clients.clear() state.pendingConnections.clear() state.disconnectedSessions.clear() + state.inFlightConnections.clear() state.authProviders.clear() for (const managed of clients) { diff --git a/src/features/skill-mcp-manager/connection-race.test.ts b/src/features/skill-mcp-manager/connection-race.test.ts index c571e9e30..d56db4033 100644 --- a/src/features/skill-mcp-manager/connection-race.test.ts +++ b/src/features/skill-mcp-manager/connection-race.test.ts @@ -80,6 +80,7 @@ function createState(): SkillMcpManagerState { cleanupHandlers: [], idleTimeoutMs: 5 * 60 * 1000, shutdownGeneration: 0, + inFlightConnections: new Map(), } trackedStates.push(state) @@ -136,13 +137,13 @@ describe("getOrCreateClient disconnect race", () => { await expect(clientPromise).rejects.toThrow(/disconnected during MCP connection setup/) expect(state.clients.has(clientKey)).toBe(false) expect(state.pendingConnections.has(clientKey)).toBe(false) - expect(state.disconnectedSessions.has(info.sessionID)).toBe(true) + expect(state.disconnectedSessions.has(info.sessionID)).toBe(false) expect(createdClients).toHaveLength(1) expect(createdClients[0]?.close).toHaveBeenCalledTimes(1) expect(createdTransports[0]?.close).toHaveBeenCalledTimes(1) }) - it("#given session A in disconnectedSessions #when new connection is requested for session A #then connection proceeds normally and disconnectedSessions entry is retained for pending race protection", async () => { + it("#given session A in disconnectedSessions #when new connection completes with no remaining pending #then disconnectedSessions entry is cleaned up", async () => { const state = createState() const info = createClientInfo("session-a") const clientKey = createClientKey(info) @@ -150,7 +151,7 @@ describe("getOrCreateClient disconnect race", () => { const client = await getOrCreateClient({ state, clientKey, info, config: stdioConfig }) - expect(state.disconnectedSessions.has(info.sessionID)).toBe(true) + expect(state.disconnectedSessions.has(info.sessionID)).toBe(false) expect(state.clients.get(clientKey)?.client).toBe(client) expect(createdClients[0]?.close).not.toHaveBeenCalled() }) @@ -210,5 +211,6 @@ describe("getOrCreateClient multi-key disconnect race", () => { expect(state.clients.has(clientKey1)).toBe(false) expect(state.clients.has(clientKey2)).toBe(false) + expect(state.disconnectedSessions.has("session-a")).toBe(false) }) }) diff --git a/src/features/skill-mcp-manager/connection.ts b/src/features/skill-mcp-manager/connection.ts index 52b9d24da..c4af1b705 100644 --- a/src/features/skill-mcp-manager/connection.ts +++ b/src/features/skill-mcp-manager/connection.ts @@ -29,6 +29,7 @@ export async function getOrCreateClient(params: { const expandedConfig = expandEnvVarsInObject(config) let currentConnectionPromise!: Promise + state.inFlightConnections.set(info.sessionID, (state.inFlightConnections.get(info.sessionID) ?? 0) + 1) currentConnectionPromise = (async () => { const disconnectGenAtStart = state.disconnectedSessions.get(info.sessionID) ?? 0 const shutdownGenAtStart = state.shutdownGeneration @@ -64,6 +65,13 @@ export async function getOrCreateClient(params: { if (state.pendingConnections.get(clientKey) === currentConnectionPromise) { state.pendingConnections.delete(clientKey) } + const remaining = (state.inFlightConnections.get(info.sessionID) ?? 1) - 1 + if (remaining <= 0) { + state.inFlightConnections.delete(info.sessionID) + state.disconnectedSessions.delete(info.sessionID) + } else { + state.inFlightConnections.set(info.sessionID, remaining) + } } } diff --git a/src/features/skill-mcp-manager/disconnect-cleanup.test.ts b/src/features/skill-mcp-manager/disconnect-cleanup.test.ts index 352efe59b..bd293940b 100644 --- a/src/features/skill-mcp-manager/disconnect-cleanup.test.ts +++ b/src/features/skill-mcp-manager/disconnect-cleanup.test.ts @@ -25,6 +25,7 @@ function createState(): SkillMcpManagerState { cleanupHandlers: [], idleTimeoutMs: 5 * 60 * 1000, shutdownGeneration: 0, + inFlightConnections: new Map(), } trackedStates.push(state) diff --git a/src/features/skill-mcp-manager/manager.ts b/src/features/skill-mcp-manager/manager.ts index 1aa981482..52f141553 100644 --- a/src/features/skill-mcp-manager/manager.ts +++ b/src/features/skill-mcp-manager/manager.ts @@ -17,6 +17,7 @@ export class SkillMcpManager { cleanupHandlers: [], idleTimeoutMs: 5 * 60 * 1000, shutdownGeneration: 0, + inFlightConnections: new Map(), } private getClientKey(info: SkillMcpClientInfo): string { diff --git a/src/features/skill-mcp-manager/types.ts b/src/features/skill-mcp-manager/types.ts index 5ee15a76d..17c867799 100644 --- a/src/features/skill-mcp-manager/types.ts +++ b/src/features/skill-mcp-manager/types.ts @@ -58,6 +58,7 @@ export interface SkillMcpManagerState { cleanupHandlers: ProcessCleanupHandler[] idleTimeoutMs: number shutdownGeneration: number + inFlightConnections: Map } export interface SkillMcpClientConnectionParams { diff --git a/src/plugin-dispose.test.ts b/src/plugin-dispose.test.ts index 57758cc17..e95184b4f 100644 --- a/src/plugin-dispose.test.ts +++ b/src/plugin-dispose.test.ts @@ -116,4 +116,60 @@ describe("createPluginDispose", () => { expect(disconnectAllSpy).toHaveBeenCalledTimes(1) expect(disposeHooksSpy).toHaveBeenCalledTimes(1) }) + + test("#given backgroundManager.shutdown() throws #when dispose() is called #then skillMcpManager.disconnectAll() and disposeHooks() are still called", async () => { + // given + const backgroundManager = { + shutdown: async (): Promise => { + throw new Error("shutdown failed") + }, + } + const skillMcpManager = { + disconnectAll: async (): Promise => {}, + } + const disposeHooksCalls: number[] = [] + const disconnectAllSpy = spyOn(skillMcpManager, "disconnectAll") + const dispose = createPluginDispose({ + backgroundManager, + skillMcpManager, + disposeHooks: (): void => { + disposeHooksCalls.push(1) + }, + }) + + // when + await dispose() + + // then + expect(disconnectAllSpy).toHaveBeenCalledTimes(1) + expect(disposeHooksCalls).toHaveLength(1) + }) + + test("#given skillMcpManager.disconnectAll() throws #when dispose() is called #then disposeHooks() is still called", async () => { + // given + const backgroundManager = { + shutdown: async (): Promise => {}, + } + const skillMcpManager = { + disconnectAll: async (): Promise => { + throw new Error("disconnectAll failed") + }, + } + const disposeHooksCalls: number[] = [] + const shutdownSpy = spyOn(backgroundManager, "shutdown") + const dispose = createPluginDispose({ + backgroundManager, + skillMcpManager, + disposeHooks: (): void => { + disposeHooksCalls.push(1) + }, + }) + + // when + await dispose() + + // then + expect(shutdownSpy).toHaveBeenCalledTimes(1) + expect(disposeHooksCalls).toHaveLength(1) + }) }) diff --git a/src/plugin-dispose.ts b/src/plugin-dispose.ts index 6d0d65b75..d7a2f2640 100644 --- a/src/plugin-dispose.ts +++ b/src/plugin-dispose.ts @@ -1,3 +1,5 @@ +import { log } from "./shared" + export type PluginDispose = () => Promise export function createPluginDispose(args: { @@ -19,9 +21,21 @@ export function createPluginDispose(args: { } disposePromise = (async (): Promise => { - await backgroundManager.shutdown() - await skillMcpManager.disconnectAll() - disposeHooks() + try { + await backgroundManager.shutdown() + } catch (error) { + log("[plugin-dispose] backgroundManager.shutdown() error:", error) + } + try { + await skillMcpManager.disconnectAll() + } catch (error) { + log("[plugin-dispose] skillMcpManager.disconnectAll() error:", error) + } + try { + disposeHooks() + } catch (error) { + log("[plugin-dispose] disposeHooks() error:", error) + } })() await disposePromise diff --git a/src/plugin/event.ts b/src/plugin/event.ts index 1cd23f375..6b4789b28 100644 --- a/src/plugin/event.ts +++ b/src/plugin/event.ts @@ -190,6 +190,7 @@ export function createEventHandler(args: { await Promise.resolve(hooks.compactionTodoPreserver?.event?.(input)); await Promise.resolve(hooks.writeExistingFileGuard?.event?.(input)); await Promise.resolve(hooks.atlasHook?.handler?.(input)); + await Promise.resolve(hooks.autoSlashCommand?.event?.(input)); }; const recentSyntheticIdles = new Map();