fix: address review-work round 2 findings

- MCP teardown race: add shutdownGeneration counter to prevent
  in-flight connections from resurrecting after disconnectAll
- MCP multi-key disconnect race: replace disconnectedSessions Set
  with generation-based Map to track per-session disconnect events
- MCP clients: check shutdownGeneration in stdio/http client
  creators before inserting into state.clients
- BackgroundManager: call clearTaskHistoryWhenParentTasksGone after
  timer-based task removal in scheduleTaskRemoval and notifyParentSession
- BackgroundManager: clean completedTaskSummaries when parent has
  no remaining tasks
- Plugin dispose: remove duplicate tmuxSessionManager.cleanup call
  since BackgroundManager.shutdown already handles it via onShutdown
This commit is contained in:
YeonGyu-Kim
2026-03-11 21:16:37 +09:00
parent 03eaa429ce
commit ff536e992a
23 changed files with 253 additions and 143 deletions
+2
View File
@@ -1149,6 +1149,7 @@ export class BackgroundManager {
if (!parentSessionID) return
if (this.getTasksByParentSession(parentSessionID).length > 0) return
this.taskHistory.clearSession(parentSessionID)
this.completedTaskSummaries.delete(parentSessionID)
}
private scheduleTaskRemoval(taskId: string): void {
@@ -1170,6 +1171,7 @@ export class BackgroundManager {
SessionCategoryRegistry.remove(task.sessionID)
}
log("[background-agent] Removed completed task from memory:", taskId)
this.clearTaskHistoryWhenParentTasksGone(task?.parentSessionID)
}
}, TASK_CLEANUP_DELAY_MS)
+2 -1
View File
@@ -95,7 +95,7 @@ export async function disconnectSession(state: SkillMcpManagerState, sessionID:
}
}
if (hasPendingForSession) {
state.disconnectedSessions.add(sessionID)
state.disconnectedSessions.set(sessionID, (state.disconnectedSessions.get(sessionID) ?? 0) + 1)
}
const keysToRemove: string[] = []
@@ -125,6 +125,7 @@ export async function disconnectSession(state: SkillMcpManagerState, sessionID:
}
export async function disconnectAll(state: SkillMcpManagerState): Promise<void> {
state.shutdownGeneration++
stopCleanupTimer(state)
unregisterProcessCleanup(state)
@@ -73,12 +73,13 @@ function createState(): SkillMcpManagerState {
const state: SkillMcpManagerState = {
clients: new Map(),
pendingConnections: new Map(),
disconnectedSessions: new Set(),
disconnectedSessions: new Map(),
authProviders: new Map(),
cleanupRegistered: false,
cleanupInterval: null,
cleanupHandlers: [],
idleTimeoutMs: 5 * 60 * 1000,
shutdownGeneration: 0,
}
trackedStates.push(state)
@@ -145,11 +146,11 @@ describe("getOrCreateClient disconnect race", () => {
const state = createState()
const info = createClientInfo("session-a")
const clientKey = createClientKey(info)
state.disconnectedSessions.add(info.sessionID)
state.disconnectedSessions.set(info.sessionID, 1)
const client = await getOrCreateClient({ state, clientKey, info, config: stdioConfig })
expect(state.disconnectedSessions.has(info.sessionID)).toBe(false)
expect(state.disconnectedSessions.has(info.sessionID)).toBe(true)
expect(state.clients.get(clientKey)?.client).toBe(client)
expect(createdClients[0]?.close).not.toHaveBeenCalled()
})
@@ -163,3 +164,51 @@ describe("getOrCreateClient disconnect race", () => {
expect(state.clients.size).toBe(0)
})
})
describe("getOrCreateClient disconnectAll race", () => {
it("#given pending connection #when disconnectAll() is called before connection completes #then client is not added to state.clients", async () => {
const state = createState()
const info = createClientInfo("session-a")
const clientKey = createClientKey(info)
const pendingConnect = createDeferred<void>()
pendingConnects.push(pendingConnect)
const clientPromise = getOrCreateClient({ state, clientKey, info, config: stdioConfig })
expect(state.pendingConnections.has(clientKey)).toBe(true)
await disconnectAll(state)
pendingConnect.resolve(undefined)
await expect(clientPromise).rejects.toThrow(/connection completed after shutdown/)
expect(state.clients.has(clientKey)).toBe(false)
})
})
describe("getOrCreateClient multi-key disconnect race", () => {
it("#given 2 pending connections for session A #when disconnectSession(A) before both complete #then both old connections are rejected", async () => {
const state = createState()
const infoKey1 = createClientInfo("session-a")
const infoKey2 = { ...createClientInfo("session-a"), serverName: "server-2" }
const clientKey1 = createClientKey(infoKey1)
const clientKey2 = `${infoKey2.sessionID}:${infoKey2.skillName}:${infoKey2.serverName}`
const pendingConnect1 = createDeferred<void>()
const pendingConnect2 = createDeferred<void>()
pendingConnects.push(pendingConnect1)
pendingConnects.push(pendingConnect2)
const promise1 = getOrCreateClient({ state, clientKey: clientKey1, info: infoKey1, config: stdioConfig })
const promise2 = getOrCreateClient({ state, clientKey: clientKey2, info: infoKey2, config: stdioConfig })
expect(state.pendingConnections.size).toBe(2)
await disconnectSession(state, "session-a")
pendingConnect1.resolve(undefined)
await expect(promise1).rejects.toThrow(/disconnected during MCP connection setup/)
pendingConnect2.resolve(undefined)
await expect(promise2).rejects.toThrow(/disconnected during MCP connection setup/)
expect(state.clients.has(clientKey1)).toBe(false)
expect(state.clients.has(clientKey2)).toBe(false)
})
})
+10 -2
View File
@@ -14,7 +14,6 @@ export async function getOrCreateClient(params: {
config: ClaudeCodeMcpServer
}): Promise<Client> {
const { state, clientKey, info, config } = params
state.disconnectedSessions.delete(info.sessionID)
const existing = state.clients.get(clientKey)
if (existing) {
@@ -31,6 +30,9 @@ export async function getOrCreateClient(params: {
const expandedConfig = expandEnvVarsInObject(config)
let currentConnectionPromise!: Promise<Client>
currentConnectionPromise = (async () => {
const disconnectGenAtStart = state.disconnectedSessions.get(info.sessionID) ?? 0
const shutdownGenAtStart = state.shutdownGeneration
const client = await createClient({ state, clientKey, info, config: expandedConfig })
const isStale = state.pendingConnections.has(clientKey) && state.pendingConnections.get(clientKey) !== currentConnectionPromise
@@ -39,7 +41,13 @@ export async function getOrCreateClient(params: {
throw new Error(`Connection for "${info.sessionID}" was superseded by a newer connection attempt.`)
}
if (state.disconnectedSessions.has(info.sessionID)) {
if (state.shutdownGeneration !== shutdownGenAtStart) {
try { await client.close() } catch {}
throw new Error(`Shutdown occurred during MCP connection for "${info.sessionID}"`)
}
const currentDisconnectGen = state.disconnectedSessions.get(info.sessionID) ?? 0
if (currentDisconnectGen > disconnectGenAtStart) {
await forceReconnect(state, clientKey)
throw new Error(`Session "${info.sessionID}" disconnected during MCP connection setup.`)
}
@@ -18,12 +18,13 @@ function createState(): SkillMcpManagerState {
const state: SkillMcpManagerState = {
clients: new Map(),
pendingConnections: new Map(),
disconnectedSessions: new Set(),
disconnectedSessions: new Map(),
authProviders: new Map(),
cleanupRegistered: false,
cleanupInterval: null,
cleanupHandlers: [],
idleTimeoutMs: 5 * 60 * 1000,
shutdownGeneration: 0,
}
trackedStates.push(state)
@@ -24,6 +24,7 @@ function redactUrl(urlStr: string): string {
export async function createHttpClient(params: SkillMcpClientConnectionParams): Promise<Client> {
const { state, clientKey, info, config } = params
const shutdownGenAtStart = state.shutdownGeneration
if (!config.url) {
throw new Error(`MCP server "${info.serverName}" is configured for HTTP but missing 'url' field.`)
@@ -72,6 +73,12 @@ export async function createHttpClient(params: SkillMcpClientConnectionParams):
)
}
if (state.shutdownGeneration !== shutdownGenAtStart) {
try { await client.close() } catch {}
try { await transport.close() } catch {}
throw new Error(`MCP server "${info.serverName}" connection completed after shutdown`)
}
const managedClient = {
client,
transport,
+2 -1
View File
@@ -10,12 +10,13 @@ export class SkillMcpManager {
private readonly state: SkillMcpManagerState = {
clients: new Map(),
pendingConnections: new Map(),
disconnectedSessions: new Set(),
disconnectedSessions: new Map(),
authProviders: new Map(),
cleanupRegistered: false,
cleanupInterval: null,
cleanupHandlers: [],
idleTimeoutMs: 5 * 60 * 1000,
shutdownGeneration: 0,
}
private getClientKey(info: SkillMcpClientInfo): string {
@@ -14,6 +14,7 @@ function getStdioCommand(config: ClaudeCodeMcpServer, serverName: string): strin
export async function createStdioClient(params: SkillMcpClientConnectionParams): Promise<Client> {
const { state, clientKey, info, config } = params
const shutdownGenAtStart = state.shutdownGeneration
const command = getStdioCommand(config, info.serverName)
const args = config.args ?? []
@@ -55,6 +56,12 @@ export async function createStdioClient(params: SkillMcpClientConnectionParams):
)
}
if (state.shutdownGeneration !== shutdownGenAtStart) {
try { await client.close() } catch {}
try { await transport.close() } catch {}
throw new Error(`MCP server "${info.serverName}" connection completed after shutdown`)
}
const managedClient = {
client,
transport,
+2 -1
View File
@@ -51,12 +51,13 @@ export interface ProcessCleanupHandler {
export interface SkillMcpManagerState {
clients: Map<string, ManagedClient>
pendingConnections: Map<string, Promise<Client>>
disconnectedSessions: Set<string>
disconnectedSessions: Map<string, number>
authProviders: Map<string, McpOAuthProvider>
cleanupRegistered: boolean
cleanupInterval: ReturnType<typeof setInterval> | null
cleanupHandlers: ProcessCleanupHandler[]
idleTimeoutMs: number
shutdownGeneration: number
}
export interface SkillMcpClientConnectionParams {