diff --git a/src/agents/pi-bundle-mcp-runtime.ts b/src/agents/pi-bundle-mcp-runtime.ts index 84537f817cc..f525ff4da6d 100644 --- a/src/agents/pi-bundle-mcp-runtime.ts +++ b/src/agents/pi-bundle-mcp-runtime.ts @@ -79,8 +79,8 @@ async function disposeSession(session: BundleMcpSession) { if (session.transportType === "streamable-http") { await (session.transport as StreamableHTTPClientTransport).terminateSession().catch(() => {}); } - await session.client.close().catch(() => {}); await session.transport.close().catch(() => {}); + await session.client.close().catch(() => {}); } function createCatalogFingerprint(servers: Record): string { @@ -454,6 +454,23 @@ export async function retireSessionMcpRuntime(params: { } } +export async function retireSessionMcpRuntimeForSessionKey(params: { + sessionKey?: string | null; + reason: string; + onError?: (error: unknown, sessionId: string, reason: string) => void; +}): Promise { + const sessionKey = normalizeOptionalString(params.sessionKey); + if (!sessionKey) { + return false; + } + const sessionId = getSessionMcpRuntimeManager().resolveSessionId(sessionKey); + return await retireSessionMcpRuntime({ + sessionId, + reason: params.reason, + onError: params.onError, + }); +} + export async function disposeAllSessionMcpRuntimes(): Promise { await getSessionMcpRuntimeManager().disposeAll(); } diff --git a/src/agents/pi-bundle-mcp-tools.ts b/src/agents/pi-bundle-mcp-tools.ts index 1f8330959a5..62a7fed9d1e 100644 --- a/src/agents/pi-bundle-mcp-tools.ts +++ b/src/agents/pi-bundle-mcp-tools.ts @@ -14,6 +14,7 @@ export { getOrCreateSessionMcpRuntime, getSessionMcpRuntimeManager, retireSessionMcpRuntime, + retireSessionMcpRuntimeForSessionKey, } from "./pi-bundle-mcp-runtime.js"; export { createBundleMcpToolRuntime, diff --git a/src/agents/pi-embedded-runner/run.ts b/src/agents/pi-embedded-runner/run.ts index 11e313e94e5..997cf7c5f4f 100644 --- a/src/agents/pi-embedded-runner/run.ts +++ b/src/agents/pi-embedded-runner/run.ts @@ -50,7 +50,10 @@ import { } from "../model-auth.js"; import { normalizeProviderId } from "../model-selection.js"; import { ensureOpenClawModelsJson } from "../models-config.js"; -import { retireSessionMcpRuntime } from "../pi-bundle-mcp-tools.js"; +import { + retireSessionMcpRuntime, + retireSessionMcpRuntimeForSessionKey, +} from "../pi-bundle-mcp-tools.js"; import { classifyFailoverReason, extractObservedOverflowTokenCount, @@ -2131,15 +2134,23 @@ export async function runEmbeddedPiAgent( await contextEngine.dispose?.(); stopRuntimeAuthRefreshTimer(); if (params.cleanupBundleMcpOnRunEnd === true) { - await retireSessionMcpRuntime({ - sessionId: params.sessionId, + const onError = (error: unknown, sessionId: string) => { + log.warn( + `bundle-mcp cleanup failed after run for ${sessionId}: ${formatErrorMessage(error)}`, + ); + }; + const retiredBySessionKey = await retireSessionMcpRuntimeForSessionKey({ + sessionKey: params.sessionKey, reason: "embedded-run-end", - onError: (error, sessionId) => { - log.warn( - `bundle-mcp cleanup failed after run for ${sessionId}: ${formatErrorMessage(error)}`, - ); - }, + onError, }); + if (!retiredBySessionKey) { + await retireSessionMcpRuntime({ + sessionId: params.sessionId, + reason: "embedded-run-end", + onError, + }); + } } } }); diff --git a/src/agents/subagent-registry-lifecycle.ts b/src/agents/subagent-registry-lifecycle.ts index 1ea1e536648..83777211798 100644 --- a/src/agents/subagent-registry-lifecycle.ts +++ b/src/agents/subagent-registry-lifecycle.ts @@ -9,6 +9,7 @@ import { setDetachedTaskDeliveryStatusByRunId, } from "../tasks/detached-task-runtime.js"; import { normalizeDeliveryContext } from "../utils/delivery-context.shared.js"; +import { retireSessionMcpRuntimeForSessionKey } from "./pi-bundle-mcp-tools.js"; import { type SubagentRunOutcome, withSubagentOutcomeTiming } from "./subagent-announce-output.js"; import { SUBAGENT_ENDED_REASON_COMPLETE, @@ -383,6 +384,20 @@ export function createSubagentRegistryLifecycleController(params: { cleanup: "delete" | "keep"; completedAt: number; }) => { + if (cleanupParams.entry.spawnMode !== "session") { + void retireSessionMcpRuntimeForSessionKey({ + sessionKey: cleanupParams.entry.childSessionKey, + reason: "subagent-run-cleanup", + onError: (error, sessionId) => { + params.warn("failed to retire subagent bundle MCP runtime", { + error: buildSafeLifecycleErrorMeta(error), + sessionId, + runId: maskRunId(cleanupParams.runId), + childSessionKey: maskSessionKey(cleanupParams.entry.childSessionKey), + }); + }, + }); + } if (cleanupParams.cleanup === "delete") { params.clearPendingLifecycleError(cleanupParams.runId); void params.notifyContextEngineSubagentEnded({ @@ -405,6 +420,28 @@ export function createSubagentRegistryLifecycleController(params: { retryDeferredCompletedAnnounces(cleanupParams.runId); }; + const retireRunModeBundleMcpRuntime = (cleanupParams: { + runId: string; + entry: SubagentRunRecord; + reason: string; + }) => { + if (cleanupParams.entry.spawnMode === "session") { + return; + } + void retireSessionMcpRuntimeForSessionKey({ + sessionKey: cleanupParams.entry.childSessionKey, + reason: cleanupParams.reason, + onError: (error, sessionId) => { + params.warn("failed to retire subagent bundle MCP runtime", { + error: buildSafeLifecycleErrorMeta(error), + sessionId, + runId: maskRunId(cleanupParams.runId), + childSessionKey: maskSessionKey(cleanupParams.entry.childSessionKey), + }); + }, + }); + }; + const finalizeSubagentCleanup = async ( runId: string, cleanup: "delete" | "keep", @@ -689,6 +726,12 @@ export function createSubagentRegistryLifecycleController(params: { onWarn: (msg) => params.warn(msg, { runId: entry.runId }), }); + retireRunModeBundleMcpRuntime({ + runId: completeParams.runId, + entry, + reason: "subagent-run-complete", + }); + startSubagentAnnounceCleanupFlow(completeParams.runId, entry); }; diff --git a/src/agents/subagent-spawn.ts b/src/agents/subagent-spawn.ts index aedd8254a98..0614397fae2 100644 --- a/src/agents/subagent-spawn.ts +++ b/src/agents/subagent-spawn.ts @@ -750,6 +750,7 @@ export async function spawnSubagentDirect( idempotencyKey: childIdem, deliver: deliverInitialChildRunDirectly, lane: AGENT_LANE_SUBAGENT, + cleanupBundleMcpOnRunEnd: spawnMode !== "session", extraSystemPrompt: childSystemPrompt, thinking: thinkingOverride, timeout: runTimeoutSeconds, diff --git a/src/agents/tools/agent-step.ts b/src/agents/tools/agent-step.ts index 06ba1e4261e..4d15ed84492 100644 --- a/src/agents/tools/agent-step.ts +++ b/src/agents/tools/agent-step.ts @@ -2,6 +2,7 @@ import crypto from "node:crypto"; import { callGateway } from "../../gateway/call.js"; import { INTERNAL_MESSAGE_CHANNEL } from "../../utils/message-channel.js"; import { resolveNestedAgentLaneForSession } from "../lanes.js"; +import { retireSessionMcpRuntimeForSessionKey } from "../pi-bundle-mcp-tools.js"; import { waitForAgentRunAndReadUpdatedAssistantReply } from "../run-wait.js"; export { readLatestAssistantReply } from "../run-wait.js"; @@ -55,6 +56,12 @@ export async function runAgentStep(params: { sessionKey: params.sessionKey, timeoutMs: Math.min(params.timeoutMs, 60_000), }); + if (result.status === "ok" || result.status === "error") { + await retireSessionMcpRuntimeForSessionKey({ + sessionKey: params.sessionKey, + reason: "nested-agent-step-complete", + }); + } if (result.status !== "ok") { return undefined; } diff --git a/src/gateway/protocol/schema/agent.ts b/src/gateway/protocol/schema/agent.ts index 22fe3816602..7970e1e7a6a 100644 --- a/src/gateway/protocol/schema/agent.ts +++ b/src/gateway/protocol/schema/agent.ts @@ -150,6 +150,7 @@ export const AgentParamsSchema = Type.Object( timeout: Type.Optional(Type.Integer({ minimum: 0 })), bestEffortDeliver: Type.Optional(Type.Boolean()), lane: Type.Optional(Type.String()), + cleanupBundleMcpOnRunEnd: Type.Optional(Type.Boolean()), extraSystemPrompt: Type.Optional(Type.String()), bootstrapContextMode: Type.Optional( Type.Union([Type.Literal("full"), Type.Literal("lightweight")]), diff --git a/src/gateway/server-methods/agent.ts b/src/gateway/server-methods/agent.ts index a8dad53a799..04c709580f3 100644 --- a/src/gateway/server-methods/agent.ts +++ b/src/gateway/server-methods/agent.ts @@ -351,6 +351,7 @@ export const agentHandlers: GatewayRequestHandlers = { idempotencyKey: string; timeout?: number; bestEffortDeliver?: boolean; + cleanupBundleMcpOnRunEnd?: boolean; label?: string; inputProvenance?: InputProvenance; }; @@ -946,6 +947,7 @@ export const agentHandlers: GatewayRequestHandlers = { messageChannel: originMessageChannel, runId, lane: request.lane, + cleanupBundleMcpOnRunEnd: request.cleanupBundleMcpOnRunEnd === true, extraSystemPrompt: request.extraSystemPrompt, bootstrapContextMode: request.bootstrapContextMode, bootstrapContextRunKind: request.bootstrapContextRunKind,