@opengeni/api-router 0.7.3 → 0.11.1

package/src/mcp/toolspace.ts

@@ -1,6 +1,6 @@
 import { Client } from "@modelcontextprotocol/sdk/client/index.js";
 import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js";
-import type { FetchLike, Transport } from "@modelcontextprotocol/sdk/shared/transport.js";
+import type { Transport } from "@modelcontextprotocol/sdk/shared/transport.js";
 import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
 import { environmentsEncryptionKeyBytes, type McpServerConfig } from "@opengeni/config";
 import {
@@ -17,15 +17,35 @@ import {
 import {
   buildConnectionTokenResolver,
   buildHostConnectionTokenResolver,
+  clearPendingSessionToolspaceCall,
+  getActiveSessionTurnForExecution,
   getSessionRootId,
-  getSessionTurn,
   listSessionMcpServerMetadata,
   listSessionMcpServersForRun,
+  registerPendingSessionToolCall,
   requireSession,
-  reserveToolspaceCallForTurn,
+  reserveToolspaceCallForAttempt,
   type ResolveConnectionCredentialResult,
 } from "@opengeni/db";
-import { appendAndPublishEvents } from "@opengeni/events";
+import { appendAndPublishEvents, appendAndPublishTurnEventsFenced } from "@opengeni/events";
+import { undiciFetch, type FetchLike } from "@opengeni/network";
+import {
+  MCP_MAX_AGGREGATE_TOOL_LIST_BYTES,
+  MCP_MAX_AGGREGATE_TOOL_LIST_ENTRIES,
+  MCP_MAX_CONCURRENT_SERVER_OPERATIONS,
+  MCP_MAX_TOOL_RESULT_BYTES,
+  McpAggregateToolListBudget,
+  McpPayloadTooLargeError,
+  assertMcpPayloadWithinBytes,
+  assertMcpServerSelectionWithinBounds,
+  assertMcpToolListWithinBounds,
+  boundedParallelMap,
+  cancelMcpResponseBody,
+  guardedMcpFetch,
+  mcpSerializedSizeBytes,
+} from "@opengeni/runtime/mcp-network";
+import { Buffer } from "node:buffer";
+import { createHash } from "node:crypto";
 
 export type ToolspaceCallResult = CallToolResult;
 
@@ -73,19 +93,109 @@ const FIRST_PARTY_PROXY_IDS = new Set(["files", "docs"]);
 // to every upstream on every call.
 const TOOLSPACE_TOOL_LIST_TTL_MS = 30_000;
 const TOOLSPACE_TOOL_LIST_CACHE_MAX_ENTRIES = 2_000;
-const toolListCache = new Map<string, { expiresAt: number; entries: ToolListingEntry[] }>();
+const TOOLSPACE_TOOL_LIST_CACHE_MAX_BYTES = 64 * 1024 * 1024;
 
-type ToolListingEntry = {
+export type ToolListingEntry = {
   serverId: string;
   tool: McpTool;
   requireApproval: McpServerConfig["requireApproval"];
 };
 
+type ToolListCacheValue = {
+  expiresAt: number;
+  entries: ToolListingEntry[];
+  sizeBytes: number;
+};
+
+/** Deterministic LRU bounded by both key count and serialized retained bytes. */
+export class ToolspaceToolListCache {
+  private readonly values = new Map<string, ToolListCacheValue>();
+  private retainedBytes = 0;
+
+  constructor(
+    private readonly maxEntries = TOOLSPACE_TOOL_LIST_CACHE_MAX_ENTRIES,
+    private readonly maxBytes = TOOLSPACE_TOOL_LIST_CACHE_MAX_BYTES,
+    private readonly ttlMs = TOOLSPACE_TOOL_LIST_TTL_MS,
+  ) {
+    if (maxEntries < 1 || maxBytes < 1 || ttlMs < 1) {
+      throw new Error("toolspace cache limits must be positive");
+    }
+  }
+
+  read(key: string, now = Date.now()): ToolListingEntry[] | null {
+    const hit = this.values.get(key);
+    if (!hit) return null;
+    if (hit.expiresAt <= now) {
+      this.delete(key);
+      return null;
+    }
+    this.values.delete(key);
+    this.values.set(key, hit);
+    return hit.entries;
+  }
+
+  write(key: string, entries: ToolListingEntry[], now = Date.now()): boolean {
+    const sizeBytes = Buffer.byteLength(key) + mcpSerializedSizeBytes(entries);
+    if (sizeBytes > this.maxBytes) return false;
+    // A rejected replacement is a no-op: preserve the current safe value
+    // until the candidate has passed its own size check.
+    this.delete(key);
+    for (const [existingKey, value] of this.values) {
+      if (value.expiresAt <= now) this.delete(existingKey);
+    }
+    while (this.values.size >= this.maxEntries || this.retainedBytes + sizeBytes > this.maxBytes) {
+      const oldestKey = this.values.keys().next().value as string | undefined;
+      if (oldestKey === undefined) break;
+      this.delete(oldestKey);
+    }
+    this.values.set(key, { expiresAt: now + this.ttlMs, entries, sizeBytes });
+    this.retainedBytes += sizeBytes;
+    return true;
+  }
+
+  clear(): void {
+    this.values.clear();
+    this.retainedBytes = 0;
+  }
+
+  snapshot(): { entries: number; bytes: number; keys: string[] } {
+    return {
+      entries: this.values.size,
+      bytes: this.retainedBytes,
+      keys: [...this.values.keys()],
+    };
+  }
+
+  private delete(key: string): void {
+    const existing = this.values.get(key);
+    if (!existing) return;
+    this.values.delete(key);
+    this.retainedBytes -= existing.sizeBytes;
+  }
+}
+
+const toolListCache = new ToolspaceToolListCache();
+
+type ToolspaceAuthority = {
+  sessionId: string;
+};
+
+type ToolspaceAttemptAuthority = ToolspaceAuthority & {
+  turnId: string;
+  attemptId: string;
+  executionGeneration: number;
+};
+
+function toolspaceAuthorityForGrant(grant: AccessGrant): ToolspaceAuthority | null {
+  const sessionId = grant.metadata?.sessionId;
+  return typeof sessionId === "string" ? { sessionId } : null;
+}
+
 export function isToolspaceGrant(settings: ApiRouteDeps["settings"], grant: AccessGrant): boolean {
   return (
     settings.toolspaceEnabled &&
     hasPermission(grant.permissions, "toolspace:call") &&
-    typeof grant.metadata?.sessionId === "string"
+    toolspaceAuthorityForGrant(grant) !== null
   );
 }
 
@@ -97,7 +207,24 @@ export async function prepareToolspaceMcpSurface(input: {
   if (!isToolspaceGrant(deps.settings, grant)) {
     return null;
   }
-  const sessionId = grant.metadata!.sessionId as string;
+  const authority = toolspaceAuthorityForGrant(grant);
+  if (!authority) {
+    return null;
+  }
+  const { sessionId } = authority;
+  const activeTurn = await getActiveSessionTurnForExecution(deps.db, grant.workspaceId, sessionId);
+  // Recovering/waiting-capacity attempts retain ownership pointers, but they
+  // are not currently executing model code. Do not even enumerate upstream
+  // tools until the turn has returned to the running state.
+  if (!activeTurn?.activeAttemptId || activeTurn.status !== "running") {
+    return emptyToolspaceSurface(sessionId, grant.subjectId);
+  }
+  const attemptAuthority: ToolspaceAttemptAuthority = {
+    sessionId,
+    turnId: activeTurn.id,
+    attemptId: activeTurn.activeAttemptId,
+    executionGeneration: activeTurn.executionGeneration,
+  };
   const session = await requireSession(deps.db, grant.workspaceId, sessionId);
   let rootSessionId = sessionId;
   if (deps.connectionCredentials?.mcpCredentials) {
@@ -114,17 +241,24 @@ export async function prepareToolspaceMcpSurface(input: {
   // Proxyable ids: everything selected except the first-party OpenGeni tool
   // server and the first-party MCP proxies, both of which would re-enter /mcp.
   const proxyableIds = [...selectedIds].filter((id) => toolspaceCanProxyServerId(id));
+  assertMcpServerSelectionWithinBounds(proxyableIds);
   if (proxyableIds.length === 0) {
     return emptyToolspaceSurface(sessionId, grant.subjectId);
   }
-
   // The registry (decrypted session servers + capability/pack expansion) is a
   // handful of DB reads with no upstream dials. Build it at most once per
   // request, and only when we actually need it (a cache-miss listing or a real
   // tools/call), so a cache-hit request does no registry work.
-  let registryPromise: Promise<Map<string, McpServerConfig>> | null = null;
-  const getRegistry = () =>
-    (registryPromise ??= buildToolspaceRegistry(deps, grant.workspaceId, sessionId));
+  const registryPromises = new Map<string, Promise<Map<string, McpServerConfig>>>();
+  const getRegistry = (attemptId: string) => {
+    const existing = registryPromises.get(attemptId);
+    if (existing) {
+      return existing;
+    }
+    const created = buildToolspaceRegistry(deps, grant.workspaceId, sessionId, attemptId);
+    registryPromises.set(attemptId, created);
+    return created;
+  };
 
   const listing = await resolveToolListing({
     deps,
@@ -132,11 +266,18 @@ export async function prepareToolspaceMcpSurface(input: {
     sessionId,
     rootSessionId,
     proxyableIds,
-    activeTurnId: session.activeTurnId ?? null,
-    getRegistry,
+    activeTurn,
+    getRegistry: () => getRegistry(attemptAuthority.attemptId),
   });
   const tools = listing.map((entry) =>
-    toolspaceToolFor({ deps, grant, sessionId, rootSessionId, entry, getRegistry }),
+    toolspaceToolFor({
+      deps,
+      grant,
+      authority: attemptAuthority,
+      rootSessionId,
+      entry,
+      getRegistry,
+    }),
   );
 
   return {
@@ -157,6 +298,7 @@ async function buildToolspaceRegistry(
   deps: ApiRouteDeps,
   workspaceId: string,
   sessionId: string,
+  attemptId: string,
 ): Promise<Map<string, McpServerConfig>> {
   const runtimeSettings = await settingsWithEnabledCapabilityMcpServers(
     deps.db,
@@ -167,6 +309,7 @@ async function buildToolspaceRegistry(
     deps,
     workspaceId,
     sessionId,
+    attemptId,
     runtimeSettings,
   );
   return new Map(withSessionServers.mcpServers.map((server) => [server.id, server]));
@@ -183,17 +326,10 @@ async function resolveToolListing(input: {
   sessionId: string;
   rootSessionId: string;
   proxyableIds: string[];
-  activeTurnId: string | null;
+  activeTurn: SessionTurn;
   getRegistry: () => Promise<Map<string, McpServerConfig>>;
 }): Promise<ToolListingEntry[]> {
-  const { deps, grant, sessionId, rootSessionId, proxyableIds, activeTurnId, getRegistry } = input;
-  if (!activeTurnId) {
-    return [];
-  }
-  const activeTurn = await getSessionTurn(deps.db, grant.workspaceId, activeTurnId);
-  if (!activeTurn || activeTurn.sessionId !== sessionId) {
-    return [];
-  }
+  const { deps, grant, sessionId, rootSessionId, proxyableIds, activeTurn, getRegistry } = input;
   // Host credentials can be initiator-specific. A prior turn's tool list must
   // never be reused under a different frozen authority in the same session.
   const cacheKey = await toolListCacheKey(
@@ -208,11 +344,21 @@ async function resolveToolListing(input: {
     return cached;
   }
   const registry = await getRegistry();
+  const aggregateBudget = new McpAggregateToolListBudget(
+    "aggregate Toolspace tool list",
+    MCP_MAX_AGGREGATE_TOOL_LIST_ENTRIES,
+    MCP_MAX_AGGREGATE_TOOL_LIST_BYTES,
+  );
   const entries: ToolListingEntry[] = [];
-  for (const serverId of proxyableIds) {
+  // Each mapper commits its source contribution synchronously immediately
+  // after the upstream result is bounded. No per-provider result arrays are
+  // retained by boundedParallelMap, and a failed aggregate replacement never
+  // reaches the cache or the exposed MCP surface.
+  await boundedParallelMap(proxyableIds, MCP_MAX_CONCURRENT_SERVER_OPERATIONS, async (serverId) => {
     const config = registry.get(serverId);
     if (!config || !toolspaceCanProxyServer(config)) {
-      continue;
+      aggregateBudget.replace(serverId, []);
+      return;
     }
     const connection = await connectToolspaceServer({
       deps,
@@ -223,22 +369,37 @@ async function resolveToolListing(input: {
       turn: activeTurn,
     }).catch(() => null);
     if (!connection) {
-      continue;
+      aggregateBudget.replace(serverId, []);
+      return;
     }
     try {
       const listed = await connection.client
         .listTools(undefined, toolspaceRequestOptions(config))
         .catch(() => ({ tools: [] }));
-      for (const tool of listed.tools as McpTool[]) {
-        if (!tool?.name || !allowedByConfig(config, tool.name)) {
-          continue;
-        }
-        entries.push({ serverId, tool, requireApproval: config.requireApproval });
+      let boundedTools: readonly McpTool[];
+      try {
+        boundedTools = assertMcpToolListWithinBounds(listed.tools as McpTool[]) as McpTool[];
+      } catch (error) {
+        deps.observability?.warn("toolspace upstream tool list exceeded safety limit", {
+          serverId,
+          errorClass: error instanceof Error ? error.name : typeof error,
+        });
+        aggregateBudget.replace(serverId, []);
+        return;
       }
+      const sourceEntries = boundedTools
+        .filter((tool) => Boolean(tool?.name) && allowedByConfig(config, tool.name))
+        .map((tool) => ({
+          serverId,
+          tool,
+          requireApproval: config.requireApproval,
+        }));
+      aggregateBudget.replace(serverId, sourceEntries);
+      entries.push(...sourceEntries);
     } finally {
       await connection.close();
     }
-  }
+  });
   writeToolListCache(cacheKey, entries);
   return entries;
 }
@@ -260,42 +421,25 @@ async function toolListCacheKey(
   const authority = JSON.stringify({
     turnId: turn.id,
     executionGeneration: turn.executionGeneration,
+    attemptId: turn.activeAttemptId,
     initiator: turn.initiator,
   });
   return `${workspaceId}:${sessionId}:${signature}:${authority}`;
 }
 
 function readToolListCache(key: string): ToolListingEntry[] | null {
-  const hit = toolListCache.get(key);
-  if (!hit) {
-    return null;
-  }
-  if (hit.expiresAt <= Date.now()) {
-    toolListCache.delete(key);
-    return null;
-  }
-  return hit.entries;
+  return toolListCache.read(key);
 }
 
 function writeToolListCache(key: string, entries: ToolListingEntry[]): void {
-  if (toolListCache.size >= TOOLSPACE_TOOL_LIST_CACHE_MAX_ENTRIES) {
-    const now = Date.now();
-    for (const [existingKey, value] of toolListCache) {
-      if (value.expiresAt <= now) {
-        toolListCache.delete(existingKey);
-      }
-    }
-    if (toolListCache.size >= TOOLSPACE_TOOL_LIST_CACHE_MAX_ENTRIES) {
-      toolListCache.clear();
-    }
-  }
-  toolListCache.set(key, { expiresAt: Date.now() + TOOLSPACE_TOOL_LIST_TTL_MS, entries });
+  toolListCache.write(key, entries);
 }
 
 async function settingsWithSessionMcpServersForToolspace(
   deps: ApiRouteDeps,
   workspaceId: string,
   sessionId: string,
+  attemptId: string,
   settings: ApiRouteDeps["settings"],
 ): Promise<ApiRouteDeps["settings"]> {
   const encryptionKey = environmentsEncryptionKeyBytes(settings);
@@ -314,6 +458,7 @@ async function settingsWithSessionMcpServersForToolspace(
     deps.db,
     workspaceId,
     sessionId,
+    attemptId,
     encryptionKey ?? null,
   );
   if (servers.length === 0) {
@@ -349,9 +494,10 @@ async function connectToolspaceServer(input: {
   rootSessionId: string;
   turn: SessionTurn;
 }): Promise<ConnectedToolspaceServer> {
+  const guardedFetch = guardedMcpFetch(input.deps.settings, undiciFetch);
   const baseFetch: FetchLike = input.config.connectionRef
-    ? connectionBrokerFetch(globalThis.fetch, input)
-    : globalThis.fetch;
+    ? connectionBrokerFetch(guardedFetch, input)
+    : guardedFetch;
   const client = new Client(
     { name: `opengeni-toolspace-${input.config.id}`, version: "1.0.0" },
     { capabilities: {} },
@@ -362,7 +508,12 @@ async function connectToolspaceServer(input: {
       headers: toolspaceServerHeaders(input.config),
     },
   });
-  await client.connect(transport as unknown as Transport, toolspaceRequestOptions(input.config));
+  try {
+    await client.connect(transport as unknown as Transport, toolspaceRequestOptions(input.config));
+  } catch (error) {
+    await client.close().catch(() => undefined);
+    throw error;
+  }
   return {
     config: input.config,
     client,
@@ -375,12 +526,13 @@ async function connectToolspaceServer(input: {
 function toolspaceToolFor(input: {
   deps: ApiRouteDeps;
   grant: AccessGrant;
-  sessionId: string;
+  authority: ToolspaceAttemptAuthority;
   rootSessionId: string;
   entry: ToolListingEntry;
-  getRegistry: () => Promise<Map<string, McpServerConfig>>;
+  getRegistry: (attemptId: string) => Promise<Map<string, McpServerConfig>>;
 }): ToolspaceRegisteredTool {
-  const { deps, grant, sessionId, rootSessionId, entry, getRegistry } = input;
+  const { deps, grant, authority, rootSessionId, entry, getRegistry } = input;
+  const { sessionId } = authority;
   const { serverId, tool } = entry;
   const name = prefixedMcpToolName(serverId, tool.name);
   const approvalRequired = mcpToolRequiresApproval(entry.requireApproval, tool.name);
@@ -392,10 +544,7 @@ function toolspaceToolFor(input: {
     ...(description ? { description } : {}),
     ...(tool.inputSchema ? { inputSchema: tool.inputSchema } : {}),
     call: async (args) => {
-      if (approvalRequired) {
-        return mcpError(APPROVAL_REQUIRED_MESSAGE);
-      }
-      const reservation = await reserveActiveTurnCall(deps, grant.workspaceId, sessionId);
+      const reservation = await reserveExactAttemptCall(deps, grant, authority);
       if (reservation.status === "no_active_turn") {
         return mcpError(TOOLSPACE_NO_ACTIVE_TURN_MESSAGE);
       }
@@ -406,9 +555,9 @@ function toolspaceToolFor(input: {
       }
       const turnId = reservation.turn.id;
       // Dial only the ONE server this tool belongs to, from the freshly-built
-      // registry, and re-check policy against that live config (the listing may
-      // have been served from a slightly stale cache entry).
-      const registry = await getRegistry();
+      // exact-attempt registry. Listing policy is descriptive only; a stale MCP
+      // surface can never decide authorization for a successor attempt.
+      const registry = await getRegistry(authority.attemptId);
       const config = registry.get(serverId);
       if (!config || !toolspaceCanProxyServer(config) || !allowedByConfig(config, tool.name)) {
         return mcpError(`upstream tool failed: ${name}`);
@@ -429,39 +578,91 @@ function toolspaceToolFor(input: {
       }
       try {
         const callId = crypto.randomUUID();
-        await appendAndPublishEvents(deps.db, deps.bus, grant.workspaceId, sessionId, [
-          {
-            type: "agent.toolCall.created",
-            turnId,
-            producerId: grant.subjectId,
-            payload: {
-              id: callId,
-              name,
-              arguments: args,
-              origin: "toolspace",
-              subjectId: grant.subjectId,
-              raw: {
-                type: "toolspace_call",
-                serverId,
-                toolName: tool.name,
+        const receipt = {
+          accountId: grant.accountId,
+          workspaceId: grant.workspaceId,
+          sessionId,
+          turnId,
+          executionGeneration: authority.executionGeneration,
+          attemptId: authority.attemptId,
+          callId,
+        };
+        const registered = await registerPendingSessionToolCall(deps.db, {
+          ...receipt,
+          callType: "toolspace_call",
+          callItem: {
+            type: "toolspace_call",
+            id: callId,
+            name,
+            arguments: toolspaceAuditSummary(args),
+            serverId,
+            toolName: tool.name,
+          },
+        });
+        if (!registered.accepted || !registered.registered) {
+          return mcpError(TOOLSPACE_NO_ACTIVE_TURN_MESSAGE);
+        }
+        const created = await appendAndPublishTurnEventsFenced(
+          deps.db,
+          deps.bus,
+          grant.workspaceId,
+          sessionId,
+          turnId,
+          authority.executionGeneration,
+          authority.attemptId,
+          [
+            {
+              type: "agent.toolCall.created",
+              turnId,
+              turnGeneration: authority.executionGeneration,
+              turnAttemptId: authority.attemptId,
+              producerId: grant.subjectId,
+              payload: {
+                id: callId,
+                name,
+                arguments: args,
+                origin: "toolspace",
+                subjectId: grant.subjectId,
+                raw: {
+                  type: "toolspace_call",
+                  serverId,
+                  toolName: tool.name,
+                },
               },
             },
-          },
-        ]);
+          ],
+        );
+        if (!created.accepted) {
+          return mcpError(TOOLSPACE_NO_ACTIVE_TURN_MESSAGE);
+        }
         const output = await callRemoteTool(deps, connection, tool.name, args);
-        await appendAndPublishEvents(deps.db, deps.bus, grant.workspaceId, sessionId, [
-          {
-            type: "agent.toolCall.output",
-            turnId,
-            producerId: grant.subjectId,
-            payload: {
-              id: callId,
-              output,
-              origin: "toolspace",
-              subjectId: grant.subjectId,
+        const completed = await appendAndPublishTurnEventsFenced(
+          deps.db,
+          deps.bus,
+          grant.workspaceId,
+          sessionId,
+          turnId,
+          authority.executionGeneration,
+          authority.attemptId,
+          [
+            {
+              type: "agent.toolCall.output",
+              turnId,
+              turnGeneration: authority.executionGeneration,
+              turnAttemptId: authority.attemptId,
+              producerId: grant.subjectId,
+              payload: {
+                id: callId,
+                output: toolspaceAuditSummary(output),
+                origin: "toolspace",
+                subjectId: grant.subjectId,
+              },
             },
-          },
-        ]);
+          ],
+        );
+        if (completed.accepted) {
+          await clearPendingSessionToolspaceCall(deps.db, receipt);
+        }
         return output;
       } finally {
         await connection.close();
@@ -477,7 +678,7 @@ async function callRemoteTool(
   args: Record<string, unknown>,
 ): Promise<ToolspaceCallResult> {
   try {
-    return (await server.client.callTool(
+    const output = (await server.client.callTool(
       {
         name: toolName,
         arguments: args,
@@ -485,7 +686,17 @@ async function callRemoteTool(
       undefined,
       toolspaceRequestOptions(server.config),
     )) as ToolspaceCallResult;
+    assertMcpPayloadWithinBytes(output, MCP_MAX_TOOL_RESULT_BYTES, "MCP tool result");
+    return output;
   } catch (error) {
+    if (error instanceof McpPayloadTooLargeError) {
+      deps.observability?.warn("toolspace upstream tool result exceeded safety limit", {
+        serverId: server.config.id,
+        toolName,
+        errorClass: error.name,
+      });
+      return mcpError("upstream tool result exceeded the safety limit");
+    }
     if (isToolspaceAuthNeededError(error)) {
       return mcpError(TOOLSPACE_AUTH_NEEDED_MESSAGE);
     }
@@ -501,34 +712,49 @@ async function callRemoteTool(
   }
 }
 
+function toolspaceAuditSummary(value: unknown): {
+  redacted: true;
+  sizeBytes: number;
+  sha256: string;
+} {
+  let serialized: string;
+  try {
+    serialized = JSON.stringify(value) ?? "null";
+  } catch {
+    serialized = "[unserializable]";
+  }
+  return {
+    redacted: true,
+    sizeBytes: Buffer.byteLength(serialized),
+    sha256: createHash("sha256").update(serialized).digest("hex"),
+  };
+}
+
 type ToolspaceReservation =
   | { status: "ok"; turn: SessionTurn }
   | { status: "no_active_turn" }
   | { status: "budget_exhausted" };
 
-async function reserveActiveTurnCall(
+async function reserveExactAttemptCall(
   deps: ApiRouteDeps,
-  workspaceId: string,
-  sessionId: string,
+  grant: AccessGrant,
+  authority: ToolspaceAttemptAuthority,
 ): Promise<ToolspaceReservation> {
-  const session = await requireSession(deps.db, workspaceId, sessionId);
-  if (!session.activeTurnId) {
-    return { status: "no_active_turn" };
-  }
-  const reservation = await reserveToolspaceCallForTurn(
-    deps.db,
-    workspaceId,
-    sessionId,
-    session.activeTurnId,
-    deps.settings.toolspaceMaxCallsPerTurn,
-  );
+  const reservation = await reserveToolspaceCallForAttempt(deps.db, {
+    accountId: grant.accountId,
+    workspaceId: grant.workspaceId,
+    sessionId: authority.sessionId,
+    turnId: authority.turnId,
+    executionGeneration: authority.executionGeneration,
+    attemptId: authority.attemptId,
+    limit: deps.settings.toolspaceMaxCallsPerTurn,
+  });
   if (!reservation.reserved) {
-    return { status: "budget_exhausted" };
+    return reservation.reason === "budget_exhausted"
+      ? { status: "budget_exhausted" }
+      : { status: "no_active_turn" };
   }
-  const turn = await getSessionTurn(deps.db, workspaceId, session.activeTurnId);
-  return turn && turn.sessionId === sessionId
-    ? { status: "ok", turn }
-    : { status: "no_active_turn" };
+  return { status: "ok", turn: reservation.turn };
 }
 
 function selectedMcpServerIds(tools: ToolRef[], sessionServerIds: string[]): Set<string> {
@@ -599,7 +825,11 @@ type McpRequestInfo = {
   toolName?: string;
 };
 
-function connectionBrokerFetch(
+function mcpRequestDestinationUrl(input: string | URL | Request): string {
+  return new URL(input instanceof Request ? input.url : input.toString()).toString();
+}
+
+export function connectionBrokerFetch(
   baseFetch: FetchLike,
   input: {
     deps: ApiRouteDeps;
@@ -630,10 +860,12 @@ function connectionBrokerFetch(
     : buildConnectionTokenResolver(input.deps.db, input.deps.settings);
   return async (requestInput, init) => {
     const request = await mcpRequestInfo(requestInput, init);
+    const destinationUrl = mcpRequestDestinationUrl(requestInput);
     const first = await resolveCredential({
       workspaceId: input.grant.workspaceId,
       serverId: input.config.id,
       connectionRef,
+      destinationUrl,
       forceRefresh: false,
       ...(request.toolName ? { toolName: request.toolName } : {}),
       subjectId: input.grant.subjectId,
@@ -646,10 +878,12 @@ function connectionBrokerFetch(
       withConnectionHeaders(requestInput, init, first.headers),
     );
     if (response.status === 401) {
+      await cancelMcpResponseBody(response);
       const refreshed = await resolveCredential({
         workspaceId: input.grant.workspaceId,
         serverId: input.config.id,
         connectionRef,
+        destinationUrl,
         forceRefresh: true,
         ...(request.toolName ? { toolName: request.toolName } : {}),
         subjectId: input.grant.subjectId,
@@ -657,12 +891,30 @@ function connectionBrokerFetch(
       if (refreshed.status === "auth_needed") {
         return await authNeededFetchResponse(input, request, refreshed);
       }
-      return await baseFetch(
+      const retry = await baseFetch(
         fetchInputForAttempt(requestInput),
         withConnectionHeaders(requestInput, init, refreshed.headers),
       );
+      if (retry.status === 401) {
+        await cancelMcpResponseBody(retry);
+        return await authNeededFetchResponse(
+          input,
+          request,
+          authNeededFromStatus(input.config, refreshed, "expired"),
+        );
+      }
+      if (retry.status === 403) {
+        await cancelMcpResponseBody(retry);
+        return await authNeededFetchResponse(
+          input,
+          request,
+          authNeededFromStatus(input.config, refreshed, "insufficient_scope"),
+        );
+      }
+      return retry;
     }
     if (response.status === 403) {
+      await cancelMcpResponseBody(response);
       return await authNeededFetchResponse(
         input,
         request,
@@ -748,8 +1000,19 @@ async function authNeededFetchResponse(
   return new Response("Authentication required for MCP server connection", { status: 401 });
 }
 
-async function mcpRequestInfo(_input: string | URL, init?: RequestInit): Promise<McpRequestInfo> {
-  const body = typeof init?.body === "string" ? init.body : "";
+async function mcpRequestInfo(
+  input: string | URL | Request,
+  init?: RequestInit,
+): Promise<McpRequestInfo> {
+  const body =
+    typeof init?.body === "string"
+      ? init.body
+      : input instanceof Request && (init?.method ?? input.method).toUpperCase() === "POST"
+        ? await input
+            .clone()
+            .text()
+            .catch(() => "")
+        : "";
   if (!body) {
     return {};
   }
@@ -779,19 +1042,21 @@ async function mcpRequestInfo(_input: string | URL, init?: RequestInit): Promise
 }
 
 function withConnectionHeaders(
-  _input: string | URL,
+  input: string | URL | Request,
   init: RequestInit | undefined,
   authHeaders: Record<string, string>,
 ): RequestInit {
-  const headers = new Headers(init?.headers);
+  const headers = new Headers(
+    init?.headers ?? (input instanceof Request ? input.headers : undefined),
+  );
   for (const [name, value] of Object.entries(authHeaders)) {
     headers.set(name, value);
   }
   return { ...init, headers };
 }
 
-function fetchInputForAttempt(input: string | URL): string | URL {
-  return input;
+function fetchInputForAttempt(input: string | URL | Request): string | URL | Request {
+  return input instanceof Request ? input.clone() : input;
 }
 
 function isToolspaceAuthNeededError(error: unknown): boolean {