fix(mcp): disconnect owned MCP manager on session dispose
This commit is contained in:
@@ -2576,6 +2576,9 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
|
||||
if (watchdogFiles && watchdogFiles.length > 0) {
|
||||
advisorWatchdogPrompt = watchdogFiles.join("\n\n");
|
||||
}
|
||||
// Owned only when this session created the manager; subagents receive a
|
||||
// parent's manager via `options.mcpManager` and MUST NOT disconnect it.
|
||||
const ownedMcpManager = options.mcpManager ? undefined : mcpManager;
|
||||
session = new AgentSession({
|
||||
advisorWatchdogPrompt,
|
||||
agent,
|
||||
@@ -2621,6 +2624,7 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
|
||||
return out;
|
||||
}
|
||||
: undefined,
|
||||
disconnectOwnedMcpManager: ownedMcpManager ? () => ownedMcpManager.disconnectAll() : undefined,
|
||||
mcpDiscoveryEnabled,
|
||||
initialSelectedMCPToolNames,
|
||||
defaultSelectedMCPToolNames,
|
||||
|
||||
@@ -487,6 +487,13 @@ export interface AgentSessionConfig {
|
||||
advisorReadOnlyTools?: AgentTool[];
|
||||
/** Preloaded watchdog prompt content for the advisor. */
|
||||
advisorWatchdogPrompt?: string;
|
||||
/**
|
||||
* Disconnect this session's OWNED MCP manager on dispose. Provided only when
|
||||
* the session created the manager (top-level sessions); subagents reuse a
|
||||
* parent's manager via `options.mcpManager` and omit this so a child's
|
||||
* teardown never tears down the shared servers.
|
||||
*/
|
||||
disconnectOwnedMcpManager?: () => Promise<void>;
|
||||
}
|
||||
|
||||
/** Options for AgentSession.prompt() */
|
||||
@@ -1154,6 +1161,7 @@ export class AgentSession {
|
||||
| undefined;
|
||||
#getMcpServerInstructions: (() => Map<string, string> | undefined) | undefined;
|
||||
#reloadSshTool: (() => Promise<AgentTool | null>) | undefined;
|
||||
#disconnectOwnedMcpManager: (() => Promise<void>) | undefined;
|
||||
#requestedToolNames: ReadonlySet<string> | undefined;
|
||||
#baseSystemPrompt: string[];
|
||||
/**
|
||||
@@ -1515,6 +1523,7 @@ export class AgentSession {
|
||||
this.#rebuildSystemPrompt = config.rebuildSystemPrompt;
|
||||
this.#getMcpServerInstructions = config.getMcpServerInstructions;
|
||||
this.#reloadSshTool = config.reloadSshTool;
|
||||
this.#disconnectOwnedMcpManager = config.disconnectOwnedMcpManager;
|
||||
this.#baseSystemPrompt = this.agent.state.systemPrompt;
|
||||
this.#promptModelKey = this.#currentPromptModelKey();
|
||||
this.#mcpDiscoveryEnabled = config.mcpDiscoveryEnabled ?? false;
|
||||
@@ -3824,6 +3833,18 @@ export class AgentSession {
|
||||
this.#releasePowerAssertion();
|
||||
await this.sessionManager.close();
|
||||
this.#closeAllProviderSessions("dispose");
|
||||
// Disconnect the MCP manager this session OWNS so its stdio servers are
|
||||
// not orphaned at exit. Best-effort: a failure here must never throw out
|
||||
// of dispose. Only owning (top-level) sessions provide this callback;
|
||||
// subagents reuse a parent's manager and must not tear it down. Idempotent
|
||||
// with the deferred-discovery disconnect in `createAgentSession`.
|
||||
if (this.#disconnectOwnedMcpManager) {
|
||||
try {
|
||||
await this.#disconnectOwnedMcpManager();
|
||||
} catch (error) {
|
||||
logger.warn("Failed to disconnect owned MCP manager during dispose", { error: String(error) });
|
||||
}
|
||||
}
|
||||
// Flush the retain queue BEFORE clearing the session's pointer so
|
||||
// `HindsightRetainQueue.#doFlush` still sees `session.getHindsightSessionState() === state`.
|
||||
// Reversed, the spliced batch survives just long enough to fail the
|
||||
|
||||
@@ -151,4 +151,62 @@ describe("createAgentSession deferred MCP auto discovery", () => {
|
||||
expect(session.getActiveToolNames()).not.toContain("search_tool_bm25");
|
||||
expect(session.isMCPDiscoveryEnabled()).toBe(false);
|
||||
}, 40_000);
|
||||
|
||||
it("disconnects the owned MCP manager when a top-level session disposes", async () => {
|
||||
writeMcpConfig();
|
||||
const { session, mcpManager } = await createAgentSession({
|
||||
...baseOptions(),
|
||||
toolNames: ["read", "edit", "bash"],
|
||||
});
|
||||
expect(mcpManager).toBeDefined();
|
||||
if (!mcpManager) throw new Error("expected owning session to create an MCPManager");
|
||||
try {
|
||||
// Let the deferred connect FINISH first, so a later disconnectAll can
|
||||
// only originate from dispose() itself — not the mid-connect disposal
|
||||
// path (covered by the test above). Genuine integration wait: discovery
|
||||
// connects a real subprocess fire-and-forget with no awaitable signal,
|
||||
// and fake timers cannot drive a child process; poll the live session
|
||||
// with a generous ceiling, exiting the instant discovery flips on.
|
||||
const deadline = Date.now() + 30_000;
|
||||
while (!session.isMCPDiscoveryEnabled() && Date.now() < deadline) {
|
||||
await Bun.sleep(50);
|
||||
}
|
||||
expect(session.isMCPDiscoveryEnabled()).toBe(true);
|
||||
expect(mcpManager.getConnectedServers()).toContain("many");
|
||||
|
||||
const disconnectSpy = spyOn(mcpManager, "disconnectAll");
|
||||
await session.dispose();
|
||||
expect(disconnectSpy).toHaveBeenCalled();
|
||||
} finally {
|
||||
// dispose() already tore it down; this is idempotent belt-and-braces.
|
||||
await mcpManager.disconnectAll();
|
||||
}
|
||||
}, 40_000);
|
||||
|
||||
it("does not disconnect a reused parent MCP manager when a child session disposes", async () => {
|
||||
writeMcpConfig();
|
||||
const parent = await createAgentSession({
|
||||
...baseOptions(),
|
||||
toolNames: ["read", "edit", "bash"],
|
||||
});
|
||||
expect(parent.mcpManager).toBeDefined();
|
||||
if (!parent.mcpManager) throw new Error("expected parent session to create an MCPManager");
|
||||
const parentManager = parent.mcpManager;
|
||||
try {
|
||||
// A subagent-style session reuses the parent's manager via
|
||||
// `mcpManager` and therefore does NOT own it.
|
||||
const child = await createAgentSession({
|
||||
...baseOptions(),
|
||||
hasUI: false,
|
||||
toolNames: ["read"],
|
||||
mcpManager: parentManager,
|
||||
});
|
||||
const disconnectSpy = spyOn(parentManager, "disconnectAll");
|
||||
await child.session.dispose();
|
||||
expect(disconnectSpy).not.toHaveBeenCalled();
|
||||
} finally {
|
||||
await parent.session.dispose();
|
||||
await parentManager.disconnectAll();
|
||||
}
|
||||
}, 40_000);
|
||||
});
|
||||
|
||||
Reference in New Issue
Block a user