@codemation/host 0.6.0 → 0.7.0

package/src/infrastructure/persistence/PrismaWorkflowRunRepository.ts

@@ -22,6 +22,8 @@ import { inject, injectable } from "@codemation/core";
 import type { WorkflowRunRepository } from "../../domain/runs/WorkflowRunRepository";
 import type { Prisma } from "../../../prisma-generated/prisma-postgresql-client/client.js";
 import { PrismaDatabaseClientToken, type PrismaDatabaseClient } from "./PrismaDatabaseClient";
+import type { WorkflowSnapshotRepository } from "./PrismaWorkflowSnapshotRepository";
+import { PrismaWorkflowSnapshotRepository } from "./PrismaWorkflowSnapshotRepository";
 
 type ExecutionInstanceRow = {
   instanceId: string;
@@ -80,7 +82,10 @@ type RunSlotProjectionRow = {
 
 @injectable()
 export class PrismaWorkflowRunRepository implements WorkflowRunRepository, WorkflowExecutionRepository {
-  constructor(@inject(PrismaDatabaseClientToken) private readonly prisma: PrismaDatabaseClient) {}
+  constructor(
+    @inject(PrismaDatabaseClientToken) private readonly prisma: PrismaDatabaseClient,
+    @inject(PrismaWorkflowSnapshotRepository) private readonly snapshotRepo: WorkflowSnapshotRepository,
+  ) {}
 
   async createRun(args: {
     runId: RunId;
@@ -96,6 +101,14 @@ export class PrismaWorkflowRunRepository implements WorkflowRunRepository, Workf
   }): Promise<void> {
     const now = new Date().toISOString();
     const testContext = args.executionOptions?.testContext;
+    const snapshotJson = args.workflowSnapshot ? JSON.stringify(args.workflowSnapshot) : null;
+    const workflowSnapshotId =
+      snapshotJson !== null
+        ? await this.snapshotRepo.findOrCreate({
+            workflowId: args.workflowId,
+            snapshotJson,
+          })
+        : null;
     await this.prisma.run.create({
       data: {
         runId: args.runId,
@@ -108,7 +121,8 @@ export class PrismaWorkflowRunRepository implements WorkflowRunRepository, Workf
         revision: 0,
         outputsByNodeJson: JSON.stringify({}),
         controlJson: args.control ? JSON.stringify(args.control) : null,
-        workflowSnapshotJson: args.workflowSnapshot ? JSON.stringify(args.workflowSnapshot) : null,
+        workflowSnapshotJson: snapshotJson,
+        workflowSnapshotId,
         policySnapshotJson: args.policySnapshot ? JSON.stringify(args.policySnapshot) : null,
         engineCountersJson: args.engineCounters ? JSON.stringify(args.engineCounters) : null,
         mutableStateJson: args.mutableState ? JSON.stringify(args.mutableState) : null,
@@ -290,6 +304,14 @@ export class PrismaWorkflowRunRepository implements WorkflowRunRepository, Workf
     const workItems = this.buildWorkItems(state, now);
     const instances = this.buildExecutionInstances(state);
     const projectionJson = this.buildProjectionSlotStatesJson(state);
+    const snapshotJson = state.workflowSnapshot ? JSON.stringify(state.workflowSnapshot) : null;
+    const workflowSnapshotId =
+      snapshotJson !== null
+        ? await this.snapshotRepo.findOrCreate({
+            workflowId: state.workflowId,
+            snapshotJson,
+          })
+        : null;
 
     await this.prisma.$transaction(async (tx) => {
       await tx.runWorkItem.deleteMany({ where: { runId: state.runId } });
@@ -389,7 +411,8 @@ export class PrismaWorkflowRunRepository implements WorkflowRunRepository, Workf
           parentJson: state.parent ? JSON.stringify(state.parent) : null,
           executionOptionsJson: state.executionOptions ? JSON.stringify(state.executionOptions) : null,
           controlJson: state.control ? JSON.stringify(state.control) : null,
-          workflowSnapshotJson: state.workflowSnapshot ? JSON.stringify(state.workflowSnapshot) : null,
+          workflowSnapshotJson: snapshotJson,
+          workflowSnapshotId,
           policySnapshotJson: state.policySnapshot ? JSON.stringify(state.policySnapshot) : null,
           engineCountersJson: state.engineCounters ? JSON.stringify(state.engineCounters) : null,
           mutableStateJson: state.mutableState ? JSON.stringify(state.mutableState) : null,

package/src/infrastructure/persistence/PrismaWorkflowSnapshotRepository.ts

@@ -0,0 +1,48 @@
+import { createHash } from "node:crypto";
+import { inject, injectable } from "@codemation/core";
+import { PrismaDatabaseClientToken, type PrismaDatabaseClient } from "./PrismaDatabaseClient";
+
+export interface WorkflowSnapshotRepository {
+  /**
+   * Returns the id of an existing snapshot matching (workflowId, snapshotHash), or creates
+   * a new one from the provided snapshotJson. Deduplication is by content hash.
+   */
+  findOrCreate(args: Readonly<{ workflowId: string; snapshotJson: string }>): Promise<string>;
+}
+
+@injectable()
+export class PrismaWorkflowSnapshotRepository implements WorkflowSnapshotRepository {
+  constructor(
+    @inject(PrismaDatabaseClientToken)
+    private readonly prisma: PrismaDatabaseClient,
+  ) {}
+
+  async findOrCreate(args: Readonly<{ workflowId: string; snapshotJson: string }>): Promise<string> {
+    const snapshotHash = createHash("sha256").update(args.snapshotJson, "utf8").digest("hex");
+    const existing = await this.prisma.workflowSnapshot.findUnique({
+      where: { workflowId_snapshotHash: { workflowId: args.workflowId, snapshotHash } },
+      select: { id: true },
+    });
+    if (existing) {
+      return existing.id;
+    }
+    const id = crypto.randomUUID();
+    await this.prisma.workflowSnapshot.upsert({
+      where: { workflowId_snapshotHash: { workflowId: args.workflowId, snapshotHash } },
+      create: {
+        id,
+        workflowId: args.workflowId,
+        snapshotHash,
+        snapshotJson: args.snapshotJson,
+        createdAt: new Date().toISOString(),
+      },
+      update: {},
+    });
+    // Re-fetch so the returned id is the winner under concurrent inserts
+    const row = await this.prisma.workflowSnapshot.findUniqueOrThrow({
+      where: { workflowId_snapshotHash: { workflowId: args.workflowId, snapshotHash } },
+      select: { id: true },
+    });
+    return row.id;
+  }
+}

package/src/mcp/AgentMcpIntegrationImpl.ts

@@ -0,0 +1,344 @@
+import type { ToolSet } from "ai";
+import {
+  AgentBindError,
+  CodemationTelemetryAttributeNames,
+  ConnectionInvocationIdFactory,
+  ConnectionNodeIdFactory,
+  inject,
+  injectable,
+  type AgentMcpIntegration,
+  type AgentMcpToolMap,
+  type ConnectionInvocationAppendArgs,
+  type JsonValue,
+  type McpServerDeclaration,
+  type NeedsReconsentEvent,
+  type NodeActivationId,
+  type NodeIterationId,
+  type ConnectionInvocationId,
+  type TelemetrySpanEventRecord,
+} from "@codemation/core";
+import { ApplicationTokens } from "../applicationTokens";
+import type { LoggerFactory } from "../application/logging/Logger";
+import { McpServerCatalog } from "./McpServerCatalog";
+import { McpConnectionPool } from "./McpConnectionPool";
+import type { CredentialStore } from "../domain/credentials/CredentialServices";
+
+/**
+ * Host-side implementation of AgentMcpIntegration.
+ *
+ * Resolves the credential binding for each declared MCP server via the standard
+ * credential-binding table — the binding lives on the MCP connection node itself
+ * (slot key `"credential"`), matching ChatModel/Tool connection nodes. Opens pool
+ * connections and returns a ToolSet map with execute callbacks wrapped for
+ * telemetry + 403 detection.
+ */
+@injectable()
+export class AgentMcpIntegrationImpl implements AgentMcpIntegration {
+  constructor(
+    @inject(McpServerCatalog) private readonly catalog: McpServerCatalog,
+    @inject(McpConnectionPool) private readonly pool: McpConnectionPool,
+    @inject(ApplicationTokens.CredentialStore) private readonly credentialStore: CredentialStore,
+    @inject(ApplicationTokens.LoggerFactory) private readonly loggers: LoggerFactory,
+  ) {}
+
+  async prepareMcpTools(args: Parameters<AgentMcpIntegration["prepareMcpTools"]>[0]): Promise<AgentMcpToolMap> {
+    const {
+      workflowId,
+      agentNodeId,
+      serverIds,
+      pinnedMcpTools: _pinnedMcpTools,
+      emitSpanEvent,
+      startChildSpan,
+      appendMcpInvocation,
+      parentAgentActivationId,
+      iterationId,
+      itemIndex,
+      parentInvocationId,
+    } = args;
+
+    const result = new Map<string, Readonly<Record<string, unknown>>>();
+    const logger = this.loggers.create("AgentMcpIntegrationImpl");
+
+    for (const serverId of serverIds) {
+      const decl = this.catalog.get(serverId);
+      if (!decl) {
+        throw new AgentBindError(`MCP server "${serverId}" not found in catalog`);
+      }
+
+      const credentialInstanceId = await this.resolveCredentialInstanceId(workflowId, agentNodeId, serverId);
+
+      // Validate scopes before opening the connection.
+      await this.validateScopes(decl, credentialInstanceId);
+
+      // Lazy-open via pool (single-flight, cached after first open).
+      await this.pool.getClient(credentialInstanceId, serverId);
+
+      // Fetch tool list from pool (cached after first fetch).
+      const rawTools = await this.pool.getTools(credentialInstanceId, serverId);
+
+      // Wrap each tool's execute for telemetry and 403 detection.
+      const wrappedTools = this.wrapToolExecutes({
+        tools: rawTools as ToolSet,
+        serverId,
+        credentialInstanceId,
+        agentNodeId,
+        emitSpanEvent,
+        startChildSpan,
+        logger,
+        appendMcpInvocation,
+        parentAgentActivationId,
+        iterationId,
+        itemIndex,
+        parentInvocationId,
+      });
+
+      result.set(serverId, wrappedTools as unknown as Readonly<Record<string, unknown>>);
+    }
+
+    return result;
+  }
+
+  /**
+   * Looks up the credential binding for the MCP connection node and verifies the
+   * referenced credential instance still exists.
+   */
+  private async resolveCredentialInstanceId(workflowId: string, agentNodeId: string, serverId: string): Promise<string> {
+    const mcpNodeId = ConnectionNodeIdFactory.mcpConnectionNodeId(agentNodeId, serverId);
+    const binding = await this.credentialStore.getBinding({ workflowId, nodeId: mcpNodeId, slotKey: "credential" });
+    if (!binding) {
+      throw new AgentBindError(
+        `MCP server "${serverId}" has no credential bound on connection node "${mcpNodeId}". ` +
+          `Bind a credential instance via the canvas credential dropdown before activation.`,
+      );
+    }
+    const instance = await this.credentialStore.getInstance(binding.instanceId);
+    if (!instance) {
+      throw new AgentBindError(
+        `Credential instance "${binding.instanceId}" not found for mcpServer "${serverId}" (connection node "${mcpNodeId}")`,
+      );
+    }
+    return instance.instanceId;
+  }
+
+  /**
+   * Validates that the credential instance's granted scopes cover the server's requiredScopes.
+   * Scopes are read from the OAuth2 material record (populated by the broker push endpoint).
+   */
+  private async validateScopes(decl: McpServerDeclaration, credentialInstanceId: string): Promise<void> {
+    if (!decl.requiredScopes?.length) {
+      return;
+    }
+
+    const material = await this.credentialStore.getOAuth2Material(credentialInstanceId);
+    const grantedScopes = new Set(material?.scopes ?? []);
+    const missing = decl.requiredScopes.filter((s) => !grantedScopes.has(s));
+
+    if (missing.length > 0) {
+      throw new AgentBindError(
+        `Credential instance "${credentialInstanceId}" lacks required scopes for server "${decl.id}": ${missing.join(", ")}. ` +
+          `Reconnect the credential to grant the missing scopes.`,
+      );
+    }
+  }
+
+  /**
+   * Returns a new ToolSet where each tool's execute callback is replaced with a wrapped version
+   * that:
+   * - Opens a child telemetry span tagged with mcp.server_id and mcp.tool_name.
+   * - Calls the original tool's execute (from @ai-sdk/mcp), which internally calls the MCP server.
+   * - On 403 / permission errors: emits a NeedsReconsentEvent span event, closes the span with
+   *   error status, and re-throws a descriptive error. The agent turn continues for other tools.
+   */
+  private wrapToolExecutes(args: {
+    tools: ToolSet;
+    serverId: string;
+    credentialInstanceId: string;
+    agentNodeId: string;
+    emitSpanEvent: (event: TelemetrySpanEventRecord) => void;
+    startChildSpan: (args: { name: string; attributes?: Record<string, string> }) => {
+      end: (args?: { status?: "ok" | "error"; statusMessage?: string }) => void;
+    };
+    logger: ReturnType<LoggerFactory["create"]>;
+    appendMcpInvocation?: (args: ConnectionInvocationAppendArgs) => Promise<void>;
+    parentAgentActivationId?: NodeActivationId;
+    iterationId?: NodeIterationId;
+    itemIndex?: number;
+    parentInvocationId?: ConnectionInvocationId;
+  }): ToolSet {
+    const {
+      tools,
+      serverId,
+      credentialInstanceId,
+      agentNodeId,
+      emitSpanEvent,
+      startChildSpan,
+      logger,
+      appendMcpInvocation,
+      parentAgentActivationId,
+      iterationId,
+      itemIndex,
+      parentInvocationId,
+    } = args;
+    const wrapped: Record<string, ToolSet[string]> = {};
+    const checkPermissionError = (err: unknown): boolean => this.isPermissionError(err);
+    const connectionNodeId = ConnectionNodeIdFactory.mcpConnectionNodeId(agentNodeId, serverId);
+
+    for (const [toolName, toolDef] of Object.entries(tools)) {
+      const originalExecute = (toolDef as { execute?: (input: unknown) => Promise<unknown> }).execute;
+      const wrappedDef = {
+        ...toolDef,
+        execute: async (input: unknown): Promise<unknown> => {
+          const span = startChildSpan({
+            name: "mcp.tool_call",
+            attributes: {
+              [CodemationTelemetryAttributeNames.mcpServerId]: serverId,
+              [CodemationTelemetryAttributeNames.mcpToolName]: toolName,
+            },
+          });
+          const invocationId = ConnectionInvocationIdFactory.create();
+          const startedAtIso = new Date().toISOString();
+          const baseRecord = {
+            invocationId,
+            connectionNodeId,
+            parentAgentNodeId: agentNodeId,
+            parentAgentActivationId: parentAgentActivationId ?? agentNodeId,
+            iterationId,
+            itemIndex,
+            parentInvocationId,
+            subjectName: toolName,
+          };
+          const summarizedInput = this.summarizeForInvocation(input);
+          if (appendMcpInvocation) {
+            await appendMcpInvocation({
+              ...baseRecord,
+              status: "running",
+              managedInput: summarizedInput,
+              queuedAt: startedAtIso,
+              startedAt: startedAtIso,
+              statusLabel: `calling ${toolName}`,
+            });
+          }
+          try {
+            if (!originalExecute) {
+              throw new Error(`MCP tool "${toolName}" on server "${serverId}" has no execute callback`);
+            }
+            const result = await originalExecute(input);
+            span.end({ status: "ok" });
+            if (appendMcpInvocation) {
+              const finishedAtIso = new Date().toISOString();
+              await appendMcpInvocation({
+                ...baseRecord,
+                status: "completed",
+                managedInput: summarizedInput,
+                managedOutput: this.summarizeForInvocation(result),
+                queuedAt: startedAtIso,
+                startedAt: startedAtIso,
+                finishedAt: finishedAtIso,
+              });
+            }
+            return result;
+          } catch (error) {
+            if (checkPermissionError(error)) {
+              const event: NeedsReconsentEvent = {
+                serverId,
+                credentialInstanceId,
+              };
+              const spanEvent: TelemetrySpanEventRecord = {
+                name: "mcp.needs_reconsent",
+                attributes: {
+                  "mcp.server_id": serverId,
+                  "mcp.credential_instance_id": credentialInstanceId,
+                },
+              };
+              emitSpanEvent(spanEvent);
+              span.end({ status: "error", statusMessage: "MCP tool permission error" });
+              logger.warn(
+                `AgentMcpIntegrationImpl: permission error from MCP tool "${toolName}" on server "${serverId}". ` +
+                  `NeedsReconsentEvent emitted for credential instance "${credentialInstanceId}".`,
+                error instanceof Error ? error : undefined,
+              );
+              const wrapped = new Error(
+                `MCP tool "${toolName}" on server "${serverId}" returned a permission error. ` +
+                  `Reconnect the credential "${credentialInstanceId}" via the Connect flow. ` +
+                  `needsReconsent: ${JSON.stringify(event satisfies NeedsReconsentEvent)}`,
+                { cause: error },
+              );
+              if (appendMcpInvocation) {
+                await appendMcpInvocation({
+                  ...baseRecord,
+                  status: "failed",
+                  managedInput: summarizedInput,
+                  error: { message: wrapped.message, name: wrapped.name },
+                  queuedAt: startedAtIso,
+                  startedAt: startedAtIso,
+                  finishedAt: new Date().toISOString(),
+                });
+              }
+              // The event carries the structured data; the agent turn continues for other tools.
+              throw wrapped;
+            }
+            const effectiveMessage = error instanceof Error ? error.message : String(error);
+            span.end({
+              status: "error",
+              statusMessage: effectiveMessage,
+            });
+            if (appendMcpInvocation) {
+              await appendMcpInvocation({
+                ...baseRecord,
+                status: "failed",
+                managedInput: summarizedInput,
+                error: { message: effectiveMessage, name: error instanceof Error ? error.name : undefined },
+                queuedAt: startedAtIso,
+                startedAt: startedAtIso,
+                finishedAt: new Date().toISOString(),
+              });
+            }
+            throw error;
+          }
+        },
+      };
+      wrapped[toolName] = wrappedDef as unknown as ToolSet[string];
+    }
+
+    return wrapped as ToolSet;
+  }
+
+  private summarizeForInvocation(value: unknown): JsonValue | undefined {
+    if (value === undefined) return undefined;
+    try {
+      const serialized = JSON.stringify(value);
+      if (serialized.length > 1024) {
+        return { truncated: true, preview: serialized.slice(0, 1024) };
+      }
+      return JSON.parse(serialized) as JsonValue;
+    } catch {
+      return undefined;
+    }
+  }
+
+  /**
+   * Detects 403 / MCP-level permission / scope-insufficiency errors from callTool.
+   * The exact shape depends on how @ai-sdk/mcp surfaces them — we check for HTTP status
+   * 403, MCP error codes (insufficient_scope / UNAUTHORIZED), and common message patterns.
+   */
+  private isPermissionError(error: unknown): boolean {
+    if (!(error instanceof Error)) {
+      return false;
+    }
+    const msg = error.message.toLowerCase();
+    // HTTP 403 from the MCP transport
+    if (msg.includes("403") || msg.includes("forbidden")) {
+      return true;
+    }
+    // MCP-level error codes
+    if (msg.includes("insufficient_scope") || msg.includes("unauthorized") || msg.includes("unauthenticated")) {
+      return true;
+    }
+    // Check error name or code
+    const candidate = error as Error & { statusCode?: number; code?: string };
+    if (candidate.statusCode === 403 || candidate.code === "EUNAUTHORIZED") {
+      return true;
+    }
+    return false;
+  }
+}

package/src/mcp/McpClientFactory.ts

@@ -0,0 +1,29 @@
+import { injectable } from "@codemation/core";
+import { experimental_createMCPClient } from "@ai-sdk/mcp";
+import type { MCPClient } from "@ai-sdk/mcp";
+
+export type McpClientOpenArgs = Readonly<{
+  url: string;
+  headers: Record<string, string>;
+}>;
+
+export interface McpClientFactory {
+  open(args: McpClientOpenArgs): Promise<MCPClient>;
+}
+
+/**
+ * Default implementation — delegates to @ai-sdk/mcp's experimental_createMCPClient
+ * using the streamable HTTP transport.
+ */
+@injectable()
+export class DefaultMcpClientFactory implements McpClientFactory {
+  async open(args: McpClientOpenArgs): Promise<MCPClient> {
+    return experimental_createMCPClient({
+      transport: {
+        type: "http",
+        url: args.url,
+        headers: args.headers,
+      },
+    });
+  }
+}

package/src/mcp/McpConnectionPool.ts

@@ -0,0 +1,184 @@
+import { inject, injectable } from "@codemation/core";
+import { ApplicationTokens } from "../applicationTokens";
+import type { LoggerFactory } from "../application/logging/Logger";
+import { McpServerCatalog } from "./McpServerCatalog";
+import { DefaultMcpClientFactory } from "./McpClientFactory";
+import type { McpClientFactory } from "./McpClientFactory";
+import { CredentialOAuth2MaterialReader } from "../credentials/CredentialOAuth2MaterialReader";
+import type { MCPClient, McpToolSet } from "./McpConnectionPool.types";
+
+/** Mutable internal pool entry (toolsCache may be filled lazily). */
+type MutablePoolEntry = {
+  client: MCPClient;
+  toolsCache: McpToolSet | null;
+  openedAt: Date;
+};
+
+@injectable()
+export class McpConnectionPool {
+  /** Key: `${credentialInstanceId}:${serverId}` */
+  private readonly pool = new Map<string, MutablePoolEntry>();
+  /**
+   * In-flight open promises — prevents a double-open race when two callers request
+   * the same (credentialInstanceId, serverId) pair concurrently.
+   */
+  private readonly inFlight = new Map<string, Promise<MutablePoolEntry>>();
+
+  constructor(
+    @inject(McpServerCatalog) private readonly catalog: McpServerCatalog,
+    @inject(CredentialOAuth2MaterialReader) private readonly oauth2Material: CredentialOAuth2MaterialReader,
+    @inject(ApplicationTokens.LoggerFactory) private readonly loggers: LoggerFactory,
+    @inject(DefaultMcpClientFactory) private readonly clientFactory: McpClientFactory,
+  ) {}
+
+  /**
+   * Returns a live MCP client for the given credential instance + server.
+   * Opens a new connection lazily; subsequent calls with the same pair return the cached client.
+   * Two concurrent calls for the same pair share a single open operation (single-flight).
+   */
+  async getClient(credentialInstanceId: string, serverId: string): Promise<MCPClient> {
+    const entry = await this.getOrOpenEntry(credentialInstanceId, serverId);
+    return entry.client;
+  }
+
+  /**
+   * Returns the tools/list result for the given credential instance + server,
+   * with toolDescriptionOverrides from the declaration applied.
+   * Fetches and caches once per pool entry; subsequent calls return the cached value.
+   * Used by Story 10's BM25 indexer.
+   */
+  async getTools(credentialInstanceId: string, serverId: string): Promise<McpToolSet> {
+    const entry = await this.getOrOpenEntry(credentialInstanceId, serverId);
+    if (!entry.toolsCache) {
+      const raw = await entry.client.tools();
+      const decl = this.catalog.get(serverId);
+      entry.toolsCache = this.applyOverrides(raw, decl?.toolDescriptionOverrides);
+    }
+    return entry.toolsCache;
+  }
+
+  /**
+   * Closes all pool entries for a credential instance.
+   * Call this when the credential is revoked or disconnected.
+   * Token refresh does NOT require closing — OAuthFlowExecutor
+   * keeps the stored token fresh; the next open will read the current token.
+   *
+   * Resolves after all matched clients have completed close(), so callers can
+   * await this before re-connecting or cleaning up downstream state.
+   *
+   * TODO(story-credential-lifecycle): Wire this method to the credential lifecycle event.
+   * CredentialDisconnectedError (packages/host/src/credentials/refresh/CredentialDisconnectedError.ts)
+   * is thrown on dead refresh tokens but is an error, not a broadcast event — there is no
+   * event bus for credential lifecycle today. When a credential-disconnected event mechanism
+   * is introduced, call closeForCredential(credentialInstanceId) from its handler so that
+   * stale MCP pool entries are cleaned up on credential revocation.
+   */
+  async closeForCredential(credentialInstanceId: string): Promise<void> {
+    const logger = this.loggers.create("McpConnectionPool");
+    const prefix = `${credentialInstanceId}:`;
+    const toClose: Array<[string, MutablePoolEntry]> = [];
+    for (const [key, entry] of this.pool.entries()) {
+      if (key.startsWith(prefix)) {
+        toClose.push([key, entry]);
+        this.pool.delete(key);
+        logger.info(`McpConnectionPool: closed pool entry on credential revocation (key=${key})`);
+      }
+    }
+    await Promise.allSettled(
+      toClose.map(([key, entry]) =>
+        entry.client.close().catch((e: unknown) => {
+          logger.warn(
+            `McpConnectionPool: error closing client on credential revocation (key=${key})`,
+            e instanceof Error ? e : undefined,
+          );
+        }),
+      ),
+    );
+  }
+
+  /**
+   * Closes all pool entries. Called on host shutdown.
+   */
+  async closeAll(): Promise<void> {
+    await Promise.allSettled([...this.pool.values()].map((e) => e.client.close()));
+    this.pool.clear();
+    this.inFlight.clear();
+  }
+
+  private async getOrOpenEntry(credentialInstanceId: string, serverId: string): Promise<MutablePoolEntry> {
+    const key = this.poolKey(credentialInstanceId, serverId);
+    const cached = this.pool.get(key);
+    if (cached) {
+      return cached;
+    }
+    const existing = this.inFlight.get(key);
+    if (existing) {
+      return existing;
+    }
+    const openPromise = this.open(credentialInstanceId, serverId, key).finally(() => {
+      this.inFlight.delete(key);
+    });
+    this.inFlight.set(key, openPromise);
+    return openPromise;
+  }
+
+  private async open(credentialInstanceId: string, serverId: string, key: string): Promise<MutablePoolEntry> {
+    const decl = this.catalog.get(serverId);
+    if (!decl) {
+      throw new Error(`McpConnectionPool: MCP server "${serverId}" not found in catalog`);
+    }
+    // D1: HTTP-only in managed mode. The catalog already blocks stdio at merge time, but we
+    // double-check here as a defensive guard in case a declaration bypasses the catalog (e.g.
+    // a future in-memory test harness or a dynamically injected server that skips catalog
+    // validation). Transport is an attribute of the declaration, not the env var.
+    if (decl.transport !== "http") {
+      throw new Error(
+        `McpConnectionPool: MCP server "${serverId}" uses transport "${decl.transport}" which is not allowed in managed mode. ` +
+          `Only "http" transport is supported. For stdio, set CODEMATION_ALLOW_STDIO_MCP=true in a self-hosted environment.`,
+      );
+    }
+
+    // Read OAuth material directly. The bearer is baked into the client's headers at open
+    // time (per-open, not per-call) — @ai-sdk/mcp v1.0.42 does not support per-request header
+    // injection. LIMITATION: a pool entry uses a stale bearer after the access token expires;
+    // OAuthFlowExecutor refreshes stored material in the background, but the entry must be
+    // closed via closeForCredential and re-opened for the refreshed token to take effect.
+    const accessToken = await this.readAccessToken(credentialInstanceId, serverId);
+    const headers: Record<string, string> = {
+      ...(decl.staticHeaders ?? {}),
+      authorization: `Bearer ${accessToken}`,
+    };
+
+    const client = await this.clientFactory.open({ url: decl.url, headers });
+    const entry: MutablePoolEntry = { client, toolsCache: null, openedAt: new Date() };
+    this.pool.set(key, entry);
+    return entry;
+  }
+
+  private poolKey(credentialInstanceId: string, serverId: string): string {
+    return `${credentialInstanceId}:${serverId}`;
+  }
+
+  private async readAccessToken(credentialInstanceId: string, serverId: string): Promise<string> {
+    const material = await this.oauth2Material.readMaterial(credentialInstanceId);
+    if (!material.accessToken) {
+      throw new Error(
+        `McpConnectionPool: credential instance "${credentialInstanceId}" has no access token — reconnect the credential bound to MCP server "${serverId}"`,
+      );
+    }
+    return material.accessToken;
+  }
+
+  private applyOverrides(tools: McpToolSet, overrides?: Record<string, string>): McpToolSet {
+    if (!overrides) {
+      return tools;
+    }
+    const result = { ...tools };
+    for (const [name, description] of Object.entries(overrides)) {
+      if (result[name]) {
+        result[name] = { ...result[name], description };
+      }
+    }
+    return result;
+  }
+}

package/src/mcp/McpConnectionPool.types.ts

@@ -0,0 +1,12 @@
+import type { MCPClient } from "@ai-sdk/mcp";
+
+export type { MCPClient };
+
+/** The ToolSet shape returned by MCPClient.tools() with 'automatic' schema resolution. */
+export type McpToolSet = Awaited<ReturnType<MCPClient["tools"]>>;
+
+export type McpPoolEntry = Readonly<{
+  client: MCPClient;
+  toolsCache: McpToolSet | null;
+  openedAt: Date;
+}>;