@bastani/atomic 0.9.4 → 0.9.5-alpha.10

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,12 +1,15 @@
-import type { McpExtensionState } from "./state.ts";
-import type { ToolMetadata } from "./types.ts";
-import { parseUiPromptHandoff } from "./types.ts";
-import { getFailureAgeSeconds } from "./init.ts";
-import { findToolByName, formatSchema } from "./tool-metadata.ts";
-import { truncateAtWord } from "./utils.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) {
@@ -70,6 +73,7 @@ export function executeUiMessages(state: McpExtensionState): ProxyToolResult {
   }
 
   const count = sessions.length;
+  assertMcpStateLease(assertActive);
   state.completedUiSessions = [];
 
   return {
@@ -84,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)) {
@@ -140,26 +145,86 @@ export function executeStatus(state: McpExtensionState): ProxyToolResult {
   };
 }
 
-export function executeDescribe(state: McpExtensionState, toolName: string): ProxyToolResult {
-  let serverName: string | undefined;
-  let toolMeta: ToolMetadata | undefined;
+function findToolMetadata(
+  state: McpExtensionState,
+  toolName: string,
+  server?: string,
+): { serverName: string; toolMeta: ToolMetadata } | null {
+  for (const [serverName, metadata] of state.toolMetadata.entries()) {
+    if (server && serverName !== server) continue;
+    const toolMeta = findToolByName(metadata, toolName);
+    if (toolMeta) return { serverName, toolMeta };
+  }
+  return null;
+}
+
+function normalizeToolAlias(value: string): string {
+  return value.replace(/-/g, "_");
+}
+
+function prefixHydrationCandidates(state: McpExtensionState, toolName: string): string[] {
+  const prefixMode = state.config.settings?.toolPrefix ?? "server";
+  if (prefixMode === "none") return [];
+  const normalizedToolName = normalizeToolAlias(toolName);
+  return Object.keys(state.config.mcpServers)
+    .map((serverName) => ({ serverName, prefix: normalizeToolAlias(getServerPrefix(serverName, prefixMode)) }))
+    .filter(({ prefix }) => prefix.length > 0 && normalizedToolName.startsWith(`${prefix}_`))
+    .sort((a, b) => b.prefix.length - a.prefix.length)
+    .map(({ serverName }) => serverName);
+}
+
+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 }, signal, assertActive);
+    return;
+  }
 
-  for (const [server, metadata] of state.toolMetadata.entries()) {
-    const found = findToolByName(metadata, toolName);
-    if (found) {
-      serverName = server;
-      toolMeta = found;
-      break;
+  const candidates = prefixHydrationCandidates(state, toolName);
+  if (candidates.length > 0) {
+    for (const candidate of candidates) {
+      await hydrateMissingMetadata(state, { server: candidate }, signal, assertActive);
+      if (findToolMetadata(state, toolName, candidate)) return;
     }
+    return;
+  }
+
+  if (findToolMetadata(state, toolName)) return;
+  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, signal, assertActive);
+    assertMcpStateLease(assertActive);
+    found = findToolMetadata(state, toolName, server);
   }
 
-  if (!serverName || !toolMeta) {
+  if (!found) {
     return {
       content: [{ type: "text" as const, text: `Tool "${toolName}" not found. Use mcp({ search: "..." }) to search.` }],
       details: { mode: "describe", error: "tool_not_found", requestedTool: toolName },
     };
   }
 
+  const { serverName, toolMeta } = found;
+
   let text = `${toolMeta.name}\n`;
   text += `Server: ${serverName}\n`;
   if (toolMeta.resourceUri) {
@@ -181,13 +246,17 @@ export function executeDescribe(state: McpExtensionState, toolName: string): Pro
   };
 }
 
-export function executeSearch(
+export async function executeSearch(
   state: McpExtensionState,
   query: string,
   regex?: boolean,
   server?: string,
   includeSchemas?: boolean,
-): ProxyToolResult {
+  signal?: AbortSignal,
+  assertActive?: AssertMcpStateLease,
+): Promise<ProxyToolResult> {
+  signal?.throwIfAborted();
+  assertMcpStateLease(assertActive);
   const showSchemas = includeSchemas !== false;
 
   const matches: Array<{ server: string; tool: ToolMetadata }> = [];
@@ -213,6 +282,8 @@ export function executeSearch(
       details: { mode: "search", error: "invalid_pattern", query },
     };
   }
+  await hydrateMissingMetadata(state, { server }, signal, assertActive);
+  assertMcpStateLease(assertActive);
 
   for (const [serverName, metadata] of state.toolMetadata.entries()) {
     if (server && serverName !== server) continue;
@@ -270,7 +341,14 @@ export function executeSearch(
   };
 }
 
-export function executeList(state: McpExtensionState, server: string): 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.` }],
@@ -278,6 +356,11 @@ export function executeList(state: McpExtensionState, server: string): ProxyTool
     };
   }
 
+  if (!state.toolMetadata.has(server)) {
+    await hydrateMissingMetadata(state, { server }, signal, assertActive);
+    assertMcpStateLease(assertActive);
+  }
+
   const metadata = state.toolMetadata.get(server);
   const toolNames = metadata?.map(m => m.name) ?? [];
   const connection = state.manager.getConnection(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;
+  }
+}