@bastani/atomic 0.9.5-alpha.9 → 0.9.5

package/dist/builtin/mcp/proxy-call.ts

@@ -1,14 +1,19 @@
 import type { ToolInfo } from "@bastani/atomic";
-import type { McpExtensionState } from "./state.ts";
-import type { McpContent, ToolMetadata } from "./types.ts";
-import { getServerPrefix } from "./types.ts";
-import { lazyConnect, updateServerMetadata, updateMetadataCache, getFailureAgeSeconds, updateStatusBar } from "./init.ts";
-import { getToolNames, findToolByName, formatSchema } from "./tool-metadata.ts";
-import { transformMcpContent } from "./tool-registrar.ts";
-import { maybeStartUiSession, type UiSessionRuntime } from "./ui-session.ts";
-import { unflattenToolArguments } from "./utils.ts";
-import { attemptAutoAuth, getAuthRequiredMessage } from "./proxy-auth.ts";
-import type { ProxyToolResult } from "./proxy-types.ts";
+import type { McpExtensionState } from "./state.js";
+import type { McpContent, ToolMetadata } from "./types.js";
+import { getServerPrefix } from "./types.js";
+import { lazyConnect, updateServerMetadata, updateMetadataCache, getFailureAgeSeconds, updateStatusBar } from "./init.js";
+import { getToolNames, findToolByName, formatSchema } from "./tool-metadata.js";
+import { transformMcpContent } from "./tool-registrar.js";
+import { maybeStartUiSession, type UiSessionRuntime } from "./ui-session.js";
+import { unflattenToolArguments } from "./utils.js";
+import { attemptAutoAuth, getAuthRequiredMessage, type AutoAuthResult } from "./proxy-auth.js";
+import type { ProxyToolResult } from "./proxy-types.js";
+import { CallToolResultSchema } from "@modelcontextprotocol/sdk/types.js";
+import { waitForCaller, waitForCallerWithLateCleanup } from "./caller-wait.js";
+import { notifyUiCancellation, rethrowHostAbortAfterUiCancellation } from "./apps-cancellation.js";
+import { asCallToolResult } from "./call-tool-result.js";
+import { assertMcpStateLease, McpStateChangedError, type AssertMcpStateLease } from "./state-lease.js";
 
 export async function executeCall(
   state: McpExtensionState,
@@ -16,7 +21,12 @@ export async function executeCall(
   args?: Record<string, unknown>,
   serverOverride?: string,
   getPiTools?: () => ToolInfo[],
+  signal?: AbortSignal,
+  startAutoAuth: (state: McpExtensionState, serverName: string) => Promise<AutoAuthResult> = attemptAutoAuth,
+  assertActive?: AssertMcpStateLease,
 ): Promise<ProxyToolResult> {
+  signal?.throwIfAborted();
+  assertMcpStateLease(assertActive);
   let serverName: string | undefined = serverOverride;
   let toolMeta: ToolMetadata | undefined;
   let autoAuthAttempted = false;
@@ -43,14 +53,18 @@ export async function executeCall(
   }
 
   if (serverName && !toolMeta) {
-    const connected = await lazyConnect(state, serverName);
+    const connected = await waitForCaller(() => lazyConnect(state, serverName!), signal);
+    signal?.throwIfAborted();
+    assertMcpStateLease(assertActive);
     if (connected) {
       toolMeta = findToolByName(state.toolMetadata.get(serverName), toolName);
     } else {
       const needsAuthConnection = state.manager.getConnection(serverName);
       if (needsAuthConnection?.status === "needs-auth") {
         autoAuthAttempted = true;
-        const autoAuth = await attemptAutoAuth(state, serverName);
+        const autoAuth = await waitForCaller(() => startAutoAuth(state, serverName!), signal);
+        signal?.throwIfAborted();
+        assertMcpStateLease(assertActive);
         if (autoAuth.status === "failed") {
           return {
             content: [{ type: "text" as const, text: autoAuth.message }],
@@ -58,9 +72,13 @@ export async function executeCall(
           };
         }
         if (autoAuth.status === "success") {
-          await state.manager.close(serverName);
+          await waitForCaller(() => state.manager.close(serverName!), signal);
+          signal?.throwIfAborted();
+          assertMcpStateLease(assertActive);
           state.failureTracker.delete(serverName);
-          const connectedAfterAuth = await lazyConnect(state, serverName);
+          const connectedAfterAuth = await waitForCaller(() => lazyConnect(state, serverName!), signal);
+          signal?.throwIfAborted();
+          assertMcpStateLease(assertActive);
           if (connectedAfterAuth) {
             toolMeta = findToolByName(state.toolMetadata.get(serverName), toolName);
             if (!toolMeta) {
@@ -106,10 +124,14 @@ export async function executeCall(
       const failedAgo = getFailureAgeSeconds(state, configuredServer);
       if (failedAgo !== null && existingConnection?.status !== "needs-auth") continue;
 
-      let connected = await lazyConnect(state, configuredServer);
+      let connected = await waitForCaller(() => lazyConnect(state, configuredServer), signal);
+      signal?.throwIfAborted();
+      assertMcpStateLease(assertActive);
       if (!connected && state.manager.getConnection(configuredServer)?.status === "needs-auth" && !autoAuthAttempted) {
         autoAuthAttempted = true;
-        const autoAuth = await attemptAutoAuth(state, configuredServer);
+        const autoAuth = await waitForCaller(() => startAutoAuth(state, configuredServer), signal);
+        signal?.throwIfAborted();
+        assertMcpStateLease(assertActive);
         if (autoAuth.status === "failed") {
           return {
             content: [{ type: "text" as const, text: autoAuth.message }],
@@ -117,9 +139,13 @@ export async function executeCall(
           };
         }
         if (autoAuth.status === "success") {
-          await state.manager.close(configuredServer);
+          await waitForCaller(() => state.manager.close(configuredServer), signal);
+          signal?.throwIfAborted();
+          assertMcpStateLease(assertActive);
           state.failureTracker.delete(configuredServer);
-          connected = await lazyConnect(state, configuredServer);
+          connected = await waitForCaller(() => lazyConnect(state, configuredServer), signal);
+          signal?.throwIfAborted();
+          assertMcpStateLease(assertActive);
         }
       }
 
@@ -162,7 +188,9 @@ export async function executeCall(
   if (connection?.status === "needs-auth") {
     if (!autoAuthAttempted) {
       autoAuthAttempted = true;
-      const autoAuth = await attemptAutoAuth(state, serverName);
+      const autoAuth = await waitForCaller(() => startAutoAuth(state, serverName!), signal);
+      signal?.throwIfAborted();
+      assertMcpStateLease(assertActive);
       if (autoAuth.status === "failed") {
         return {
           content: [{ type: "text" as const, text: autoAuth.message }],
@@ -170,7 +198,9 @@ export async function executeCall(
         };
       }
       if (autoAuth.status === "success") {
-        await state.manager.close(serverName);
+        await waitForCaller(() => state.manager.close(serverName!), signal);
+        signal?.throwIfAborted();
+        assertMcpStateLease(assertActive);
         state.failureTracker.delete(serverName);
         connection = state.manager.getConnection(serverName);
       }
@@ -202,14 +232,19 @@ export async function executeCall(
     }
 
     try {
+      assertMcpStateLease(assertActive);
       if (state.ui) {
         state.ui.setStatus("mcp", `MCP: connecting to ${serverName}...`);
       }
-      connection = await state.manager.connect(serverName, definition);
+      connection = await waitForCaller(() => state.manager.connect(serverName, definition), signal);
+      signal?.throwIfAborted();
+      assertMcpStateLease(assertActive);
       if (connection.status === "needs-auth") {
         if (!autoAuthAttempted) {
           autoAuthAttempted = true;
-          const autoAuth = await attemptAutoAuth(state, serverName);
+          const autoAuth = await waitForCaller(() => startAutoAuth(state, serverName!), signal);
+          signal?.throwIfAborted();
+          assertMcpStateLease(assertActive);
           if (autoAuth.status === "failed") {
             return {
               content: [{ type: "text" as const, text: autoAuth.message }],
@@ -217,8 +252,12 @@ export async function executeCall(
             };
           }
           if (autoAuth.status === "success") {
-            await state.manager.close(serverName);
-            connection = await state.manager.connect(serverName, definition);
+            await waitForCaller(() => state.manager.close(serverName!), signal);
+            signal?.throwIfAborted();
+            assertMcpStateLease(assertActive);
+            connection = await waitForCaller(() => state.manager.connect(serverName, definition), signal);
+            signal?.throwIfAborted();
+            assertMcpStateLease(assertActive);
           }
         }
 
@@ -230,6 +269,7 @@ export async function executeCall(
           };
         }
       }
+      assertMcpStateLease(assertActive);
       state.failureTracker.delete(serverName);
       updateServerMetadata(state, serverName);
       updateMetadataCache(state, serverName);
@@ -246,6 +286,8 @@ export async function executeCall(
         };
       }
     } catch (error) {
+      signal?.throwIfAborted();
+      assertMcpStateLease(assertActive);
       state.failureTracker.set(serverName, Date.now());
       updateStatusBar(state);
       const message = error instanceof Error ? error.message : String(error);
@@ -259,11 +301,15 @@ export async function executeCall(
   let uiSession: UiSessionRuntime | null = null;
 
   try {
+    assertMcpStateLease(assertActive);
     state.manager.touch(serverName);
     state.manager.incrementInFlight(serverName);
 
     if (toolMeta.resourceUri) {
-      const result = await connection.client.readResource({ uri: toolMeta.resourceUri });
+      assertMcpStateLease(assertActive);
+      const result = await connection.client.readResource({ uri: toolMeta.resourceUri }, { signal });
+      signal?.throwIfAborted();
+      assertMcpStateLease(assertActive);
       const content = (result.contents ?? []).map(c => ({
         type: "text" as const,
         text: "text" in c ? c.text : ("blob" in c ? `[Binary data: ${(c as { mimeType?: string }).mimeType ?? "unknown"}]` : JSON.stringify(c)),
@@ -275,15 +321,24 @@ export async function executeCall(
     }
 
     uiSession = toolMeta.uiResourceUri
-      ? await maybeStartUiSession(state, {
-          serverName,
-          toolName: toolMeta.originalName,
-          toolArgs: args ?? {},
-          uiResourceUri: toolMeta.uiResourceUri,
-          streamMode: toolMeta.uiStreamMode,
-        })
+      ? await waitForCallerWithLateCleanup(
+          () => maybeStartUiSession(state, {
+            serverName,
+            toolName: toolMeta.originalName,
+            toolArgs: args ?? {},
+            uiResourceUri: toolMeta.uiResourceUri!,
+            streamMode: toolMeta.uiStreamMode,
+          }),
+          signal,
+          (lateSession) => {
+            if (lateSession && !lateSession.reused) lateSession.close("caller cancelled during UI startup");
+          },
+        )
       : null;
+    signal?.throwIfAborted();
+    assertMcpStateLease(assertActive);
 
+    assertMcpStateLease(assertActive);
     const resultPromise = connection.client.callTool({
       name: toolMeta.originalName,
       // Normalize provider-flattened argument keys (e.g. Gemini's `keywords[0]`)
@@ -292,11 +347,14 @@ export async function executeCall(
       // preserved unless the schema proves the head is a container.
       arguments: unflattenToolArguments(args, toolMeta.inputSchema),
       _meta: uiSession?.requestMeta,
-    });
+    }, CallToolResultSchema, { signal });
 
     if (toolMeta.uiResourceUri) {
-      const result = await resultPromise;
-      uiSession?.sendToolResult(result as unknown as import("@modelcontextprotocol/sdk/types.js").CallToolResult);
+      const sdkResult = await resultPromise;
+      signal?.throwIfAborted();
+      assertMcpStateLease(assertActive);
+      const result = asCallToolResult(sdkResult);
+      uiSession?.sendToolResult(result);
       const mcpContent = (result.content ?? []) as McpContent[];
       const content = transformMcpContent(mcpContent);
 
@@ -326,7 +384,10 @@ export async function executeCall(
       };
     }
 
-    const result = await resultPromise;
+    const sdkResult = await resultPromise;
+    signal?.throwIfAborted();
+    assertMcpStateLease(assertActive);
+    const result = asCallToolResult(sdkResult);
 
     const mcpContent = (result.content ?? []) as McpContent[];
     const content = transformMcpContent(mcpContent);
@@ -353,8 +414,11 @@ export async function executeCall(
       details: { mode: "call", mcpResult: result, server: serverName, tool: toolMeta.originalName },
     };
   } catch (error) {
+    await rethrowHostAbortAfterUiCancellation(signal, uiSession);
     const message = error instanceof Error ? error.message : String(error);
-    uiSession?.sendToolCancelled(message);
+    await notifyUiCancellation(uiSession, message);
+    signal?.throwIfAborted();
+    assertMcpStateLease(assertActive);
 
     let errorWithSchema = `Failed to call tool: ${message}`;
     if (toolMeta.inputSchema) {
@@ -366,10 +430,17 @@ export async function executeCall(
       details: { mode: "call", error: "call_failed", message },
     };
   } finally {
-    if (uiSession?.reused) {
-      uiSession.close();
+    let staleState = false;
+    try {
+      assertMcpStateLease(assertActive);
+    } catch (error) {
+      if (error instanceof McpStateChangedError) staleState = true;
+      else throw error;
+    }
+    if (uiSession && (uiSession.reused || staleState)) {
+      uiSession.close(staleState ? "stale MCP session" : undefined);
     }
     state.manager.decrementInFlight(serverName);
-    state.manager.touch(serverName);
+    if (!staleState) state.manager.touch(serverName);
   }
 }

package/dist/builtin/mcp/proxy-connect.ts

@@ -1,11 +1,21 @@
-import type { McpExtensionState } from "./state.ts";
-import { updateMetadataCache, updateStatusBar } from "./init.ts";
-import { buildToolMetadata } from "./tool-metadata.ts";
-import { attemptAutoAuth, getAuthRequiredMessage } from "./proxy-auth.ts";
-import { executeList } from "./proxy-info-modes.ts";
-import type { ProxyToolResult } from "./proxy-types.ts";
+import type { McpExtensionState } from "./state.js";
+import { updateMetadataCache, updateStatusBar } from "./init.js";
+import { buildToolMetadata } from "./tool-metadata.js";
+import { attemptAutoAuth, getAuthRequiredMessage, type AutoAuthResult } from "./proxy-auth.js";
+import { executeList } from "./proxy-info-modes.js";
+import type { ProxyToolResult } from "./proxy-types.js";
+import { waitForCaller } from "./caller-wait.js";
+import { assertMcpStateLease, type AssertMcpStateLease } from "./state-lease.js";
 
-export async function executeConnect(state: McpExtensionState, serverName: string): Promise<ProxyToolResult> {
+export async function executeConnect(
+  state: McpExtensionState,
+  serverName: string,
+  signal?: AbortSignal,
+  startAutoAuth: (state: McpExtensionState, serverName: string) => Promise<AutoAuthResult> = attemptAutoAuth,
+  assertActive?: AssertMcpStateLease,
+): Promise<ProxyToolResult> {
+  signal?.throwIfAborted();
+  assertMcpStateLease(assertActive);
   const definition = state.config.mcpServers[serverName];
   if (!definition) {
     return {
@@ -15,12 +25,15 @@ export async function executeConnect(state: McpExtensionState, serverName: strin
   }
 
   try {
-    if (state.ui) {
-      state.ui.setStatus("mcp", `MCP: connecting to ${serverName}...`);
-    }
-    let connection = await state.manager.connect(serverName, definition);
+    assertMcpStateLease(assertActive);
+    if (state.ui) state.ui.setStatus("mcp", `MCP: connecting to ${serverName}...`);
+    let connection = await waitForCaller(() => state.manager.connect(serverName, definition), signal);
+    signal?.throwIfAborted();
+    assertMcpStateLease(assertActive);
     if (connection.status === "needs-auth") {
-      const autoAuth = await attemptAutoAuth(state, serverName);
+      const autoAuth = await waitForCaller(() => startAutoAuth(state, serverName), signal);
+      signal?.throwIfAborted();
+      assertMcpStateLease(assertActive);
       if (autoAuth.status === "failed") {
         return {
           content: [{ type: "text" as const, text: autoAuth.message }],
@@ -28,8 +41,12 @@ export async function executeConnect(state: McpExtensionState, serverName: strin
         };
       }
       if (autoAuth.status === "success") {
-        await state.manager.close(serverName);
-        connection = await state.manager.connect(serverName, definition);
+        await waitForCaller(() => state.manager.close(serverName), signal);
+        signal?.throwIfAborted();
+        assertMcpStateLease(assertActive);
+        connection = await waitForCaller(() => state.manager.connect(serverName, definition), signal);
+        signal?.throwIfAborted();
+        assertMcpStateLease(assertActive);
       }
       if (connection.status === "needs-auth") {
         const message = getAuthRequiredMessage(state, serverName);
@@ -39,14 +56,17 @@ export async function executeConnect(state: McpExtensionState, serverName: strin
         };
       }
     }
+    assertMcpStateLease(assertActive);
     const prefix = state.config.settings?.toolPrefix ?? "server";
     const { metadata } = buildToolMetadata(connection.tools, connection.resources, definition, serverName, prefix);
     state.toolMetadata.set(serverName, metadata);
     updateMetadataCache(state, serverName);
     state.failureTracker.delete(serverName);
     updateStatusBar(state);
-    return executeList(state, serverName);
+    return executeList(state, serverName, signal, assertActive);
   } catch (error) {
+    signal?.throwIfAborted();
+    assertMcpStateLease(assertActive);
     state.failureTracker.set(serverName, Date.now());
     updateStatusBar(state);
     const message = error instanceof Error ? error.message : String(error);

package/dist/builtin/mcp/proxy-info-modes.ts

@@ -1,13 +1,15 @@
-import type { McpExtensionState } from "./state.ts";
-import type { ToolMetadata } from "./types.ts";
-import { getServerPrefix, parseUiPromptHandoff } from "./types.ts";
-import { getFailureAgeSeconds } from "./init.ts";
-import { findToolByName, formatSchema } from "./tool-metadata.ts";
-import { truncateAtWord } from "./utils.ts";
-import { hydrateMissingMetadata } from "./metadata-hydration.ts";
-import type { ProxyToolResult } from "./proxy-types.ts";
-
-export function executeUiMessages(state: McpExtensionState): ProxyToolResult {
+import type { McpExtensionState } from "./state.js";
+import type { ToolMetadata } from "./types.js";
+import { getServerPrefix, parseUiPromptHandoff } from "./types.js";
+import { getFailureAgeSeconds } from "./init.js";
+import { findToolByName, formatSchema } from "./tool-metadata.js";
+import { truncateAtWord } from "./utils.js";
+import { hydrateMissingMetadata } from "./metadata-hydration.js";
+import type { ProxyToolResult } from "./proxy-types.js";
+import { assertMcpStateLease, type AssertMcpStateLease } from "./state-lease.js";
+
+export function executeUiMessages(state: McpExtensionState, assertActive?: AssertMcpStateLease): ProxyToolResult {
+  assertMcpStateLease(assertActive);
   const sessions = state.completedUiSessions;
 
   if (sessions.length === 0) {
@@ -71,6 +73,7 @@ export function executeUiMessages(state: McpExtensionState): ProxyToolResult {
   }
 
   const count = sessions.length;
+  assertMcpStateLease(assertActive);
   state.completedUiSessions = [];
 
   return {
@@ -85,7 +88,8 @@ export function executeUiMessages(state: McpExtensionState): ProxyToolResult {
   };
 }
 
-export function executeStatus(state: McpExtensionState): ProxyToolResult {
+export function executeStatus(state: McpExtensionState, assertActive?: AssertMcpStateLease): ProxyToolResult {
+  assertMcpStateLease(assertActive);
   const servers: Array<{ name: string; status: string; toolCount: number; failedAgo: number | null }> = [];
 
   for (const name of Object.keys(state.config.mcpServers)) {
@@ -169,33 +173,46 @@ function prefixHydrationCandidates(state: McpExtensionState, toolName: string):
     .map(({ serverName }) => serverName);
 }
 
-async function hydrateDescribeMetadata(state: McpExtensionState, toolName: string, server?: string): Promise<void> {
+async function hydrateDescribeMetadata(
+  state: McpExtensionState,
+  toolName: string,
+  server?: string,
+  signal?: AbortSignal,
+  assertActive?: AssertMcpStateLease,
+): Promise<void> {
+  signal?.throwIfAborted();
+  assertMcpStateLease(assertActive);
   if (server) {
-    await hydrateMissingMetadata(state, { server });
+    await hydrateMissingMetadata(state, { server }, signal, assertActive);
     return;
   }
 
   const candidates = prefixHydrationCandidates(state, toolName);
   if (candidates.length > 0) {
     for (const candidate of candidates) {
-      await hydrateMissingMetadata(state, { server: candidate });
+      await hydrateMissingMetadata(state, { server: candidate }, signal, assertActive);
       if (findToolMetadata(state, toolName, candidate)) return;
     }
     return;
   }
 
   if (findToolMetadata(state, toolName)) return;
-  await hydrateMissingMetadata(state);
+  await hydrateMissingMetadata(state, undefined, signal, assertActive);
 }
 
 export async function executeDescribe(
   state: McpExtensionState,
   toolName: string,
   server?: string,
+  signal?: AbortSignal,
+  assertActive?: AssertMcpStateLease,
 ): Promise<ProxyToolResult> {
+  signal?.throwIfAborted();
+  assertMcpStateLease(assertActive);
   let found = findToolMetadata(state, toolName, server);
   if (!found) {
-    await hydrateDescribeMetadata(state, toolName, server);
+    await hydrateDescribeMetadata(state, toolName, server, signal, assertActive);
+    assertMcpStateLease(assertActive);
     found = findToolMetadata(state, toolName, server);
   }
 
@@ -235,7 +252,11 @@ export async function executeSearch(
   regex?: boolean,
   server?: string,
   includeSchemas?: boolean,
+  signal?: AbortSignal,
+  assertActive?: AssertMcpStateLease,
 ): Promise<ProxyToolResult> {
+  signal?.throwIfAborted();
+  assertMcpStateLease(assertActive);
   const showSchemas = includeSchemas !== false;
 
   const matches: Array<{ server: string; tool: ToolMetadata }> = [];
@@ -261,7 +282,8 @@ export async function executeSearch(
       details: { mode: "search", error: "invalid_pattern", query },
     };
   }
-  await hydrateMissingMetadata(state, { server });
+  await hydrateMissingMetadata(state, { server }, signal, assertActive);
+  assertMcpStateLease(assertActive);
 
   for (const [serverName, metadata] of state.toolMetadata.entries()) {
     if (server && serverName !== server) continue;
@@ -319,7 +341,14 @@ export async function executeSearch(
   };
 }
 
-export async function executeList(state: McpExtensionState, server: string): Promise<ProxyToolResult> {
+export async function executeList(
+  state: McpExtensionState,
+  server: string,
+  signal?: AbortSignal,
+  assertActive?: AssertMcpStateLease,
+): Promise<ProxyToolResult> {
+  signal?.throwIfAborted();
+  assertMcpStateLease(assertActive);
   if (!state.config.mcpServers[server]) {
     return {
       content: [{ type: "text" as const, text: `Server "${server}" not found. Use mcp({}) to see available servers.` }],
@@ -328,7 +357,8 @@ export async function executeList(state: McpExtensionState, server: string): Pro
   }
 
   if (!state.toolMetadata.has(server)) {
-    await hydrateMissingMetadata(state, { server });
+    await hydrateMissingMetadata(state, { server }, signal, assertActive);
+    assertMcpStateLease(assertActive);
   }
 
   const metadata = state.toolMetadata.get(server);

package/dist/builtin/mcp/proxy-modes.ts

@@ -1,3 +1,3 @@
-export { executeUiMessages, executeStatus, executeDescribe, executeSearch, executeList } from "./proxy-info-modes.ts";
-export { executeConnect } from "./proxy-connect.ts";
-export { executeCall } from "./proxy-call.ts";
+export { executeUiMessages, executeStatus, executeDescribe, executeSearch, executeList } from "./proxy-info-modes.js";
+export { executeConnect } from "./proxy-connect.js";
+export { executeCall } from "./proxy-call.js";

package/dist/builtin/mcp/session-cleanup-barrier.ts

@@ -0,0 +1,43 @@
+/**
+ * Retain teardown work across MCP lifecycle generations without allowing a
+ * non-abortable producer to block every later generation forever.
+ */
+type CleanupTask = PromiseLike<object | string | void>;
+
+const DEFAULT_CLEANUP_DEADLINE_MS = 2_000;
+
+function observeWithin(task: CleanupTask, deadlineMs: number, onTimeout: () => void): Promise<void> {
+  const observed = Promise.resolve(task).then(() => undefined, () => undefined);
+  let timer: ReturnType<typeof setTimeout> | undefined;
+  const deadline = new Promise<void>((resolve) => {
+    timer = setTimeout(() => {
+      onTimeout();
+      resolve();
+    }, deadlineMs);
+    timer.unref?.();
+  });
+  return Promise.race([observed, deadline]).finally(() => {
+    if (timer) clearTimeout(timer);
+  });
+}
+
+export class McpSessionCleanupBarrier {
+  private settled: Promise<void> = Promise.resolve();
+
+  constructor(private readonly deadlineMs = DEFAULT_CLEANUP_DEADLINE_MS) {}
+
+  wait(): Promise<void> {
+    return this.settled;
+  }
+
+  retain(tasks: readonly (CleanupTask | null | undefined)[]): Promise<void> {
+    const observed = tasks.filter((task): task is CleanupTask => task !== null && task !== undefined);
+    const prior = this.settled;
+    const next = Promise.all([
+      observeWithin(prior, this.deadlineMs, () => console.error("MCP: prior cleanup exceeded its deadline; continuing with fenced state")),
+      ...observed.map((task) => observeWithin(task, this.deadlineMs, () => console.error("MCP: cleanup task exceeded its deadline; continuing with fenced state"))),
+    ]).then(() => undefined);
+    this.settled = next;
+    return next;
+  }
+}

package/dist/builtin/mcp/state-lease.ts

@@ -0,0 +1,12 @@
+export type AssertMcpStateLease = () => void;
+
+export class McpStateChangedError extends Error {
+  constructor() {
+    super("MCP session changed during execution");
+    this.name = "McpStateChangedError";
+  }
+}
+
+export function assertMcpStateLease(assertActive?: AssertMcpStateLease): void {
+  assertActive?.();
+}

package/dist/builtin/mcp/ui-server.ts

@@ -7,12 +7,12 @@ import type {
   CallToolRequest,
   CallToolResult,
 } from "@modelcontextprotocol/sdk/types.js";
-import type { ConsentManager } from "./consent-manager.ts";
-import { ServerError, wrapError } from "./errors.ts";
-import { buildHostHtmlTemplate, buildCspMetaContent, applyCspMeta } from "./host-html-template.ts";
-import { logger } from "./logger.ts";
-import { parseBody, sendJson, validateTokenBody, validateTokenQuery } from "./ui-server-http.ts";
-import type { McpServerManager } from "./server-manager.ts";
+import type { ConsentManager } from "./consent-manager.js";
+import { ServerError, wrapError } from "./errors.js";
+import { buildHostHtmlTemplate, buildCspMetaContent, applyCspMeta } from "./host-html-template.js";
+import { logger } from "./logger.js";
+import { parseBody, sendJson, validateTokenBody, validateTokenQuery } from "./ui-server-http.js";
+import type { McpServerManager } from "./server-manager.js";
 import {
   extractUiPromptText,
   getVisualizationStreamEnvelope,
@@ -26,7 +26,7 @@ import {
   type UiResourceContent,
   type UiSessionMessages,
   type UiStreamSummary,
-} from "./types.ts";
+} from "./types.js";
 
 const ABANDONED_GRACE_MS = 60_000;
 const WATCHDOG_INTERVAL_MS = 5_000;
@@ -58,7 +58,7 @@ export interface UiServerHandle {
   sendToolInput: (args: Record<string, unknown>) => void;
   sendToolResult: (result: CallToolResult) => void;
   sendResultPatch: (result: CallToolResult) => void;
-  sendToolCancelled: (reason: string) => void;
+  sendToolCancelled: (reason: string) => void | Promise<void>;
   sendHostContext: (context: UiHostContext) => void;
   /** Get accumulated messages from this session */
   getSessionMessages: () => UiSessionMessages;