fix(mcp): apply timeouts to catalog requests (#31618)

2026-06-10 14:38:08 -05:00
parent e10012527b
commit ffcb568912
2 changed files with 170 additions and 164 deletions
@@ -0,0 +1,138 @@
+import { Client } from "@modelcontextprotocol/sdk/client/index.js"
+import {
+  CallToolResultSchema,
+  ListToolsResultSchema,
+  ToolSchema,
+  type Tool as MCPToolDef,
+} from "@modelcontextprotocol/sdk/types.js"
+import { dynamicTool, jsonSchema, type JSONSchema7, type Tool } from "ai"
+import { Effect } from "effect"
+
+const DEFAULT_TIMEOUT = 30_000
+const MAX_LIST_PAGES = 1_000
+
+const TolerantListToolsResultSchema = ListToolsResultSchema.extend({
+  tools: ToolSchema.omit({ outputSchema: true }).array(),
+})
+
+export async function paginate<T, R extends { nextCursor?: string }>(
+  list: (cursor?: string) => Promise<R>,
+  items: (result: R) => T[],
+) {
+  const result: T[] = []
+  const cursors = new Set<string>()
+  let cursor: string | undefined
+
+  for (let page = 0; page < MAX_LIST_PAGES; page++) {
+    const page = await list(cursor)
+    result.push(...items(page))
+    if (page.nextCursor === undefined) return result
+    if (cursors.has(page.nextCursor)) throw new Error(`MCP list returned duplicate cursor: ${page.nextCursor}`)
+    cursors.add(page.nextCursor)
+    cursor = page.nextCursor
+  }
+
+  throw new Error(`MCP list exceeded ${MAX_LIST_PAGES} pages`)
+}
+
+export function defs(client: Client, timeout?: number) {
+  return listTools(client, timeout ?? DEFAULT_TIMEOUT).pipe(Effect.catch(() => Effect.void))
+}
+
+export function convertTool(mcpTool: MCPToolDef, client: Client, timeout?: number): Tool {
+  const inputSchema: JSONSchema7 = {
+    ...(mcpTool.inputSchema as JSONSchema7),
+    type: "object",
+    properties: (mcpTool.inputSchema.properties ?? {}) as JSONSchema7["properties"],
+    additionalProperties: false,
+  }
+
+  return dynamicTool({
+    description: mcpTool.description ?? "",
+    inputSchema: jsonSchema(inputSchema),
+    execute: (args: unknown, options) =>
+      client.callTool(
+        {
+          name: mcpTool.name,
+          arguments: (args || {}) as Record<string, unknown>,
+        },
+        CallToolResultSchema,
+        {
+          resetTimeoutOnProgress: true,
+          signal: options.abortSignal,
+          timeout,
+        },
+      ),
+  })
+}
+
+export function fetch<T extends { name: string }>(
+  clientName: string,
+  client: Client,
+  list: (client: Client) => Promise<T[]>,
+  label: string,
+) {
+  return Effect.tryPromise({
+    try: () => list(client),
+    catch: (error) => error,
+  }).pipe(
+    Effect.tapError((error) =>
+      Effect.logWarning(`failed to get ${label}`, {
+        clientName,
+        error: error instanceof Error ? error.message : String(error),
+      }),
+    ),
+    Effect.map((items) => {
+      const sanitizedClient = sanitize(clientName)
+      return Object.fromEntries(
+        items.map((item) => [sanitizedClient + ":" + sanitize(item.name), { ...item, client: clientName }]),
+      )
+    }),
+    Effect.orElseSucceed(() => undefined),
+  )
+}
+
+export const sanitize = (value: string) => value.replace(/[^a-zA-Z0-9_-]/g, "_")
+
+export function prompts(client: Client, timeout?: number) {
+  if (!client.getServerCapabilities()?.prompts) return Promise.resolve([])
+  return paginate(
+    (cursor) => client.listPrompts(cursor === undefined ? undefined : { cursor }, { timeout }),
+    (result) => result.prompts,
+  )
+}
+
+export function resources(client: Client, timeout?: number) {
+  if (!client.getServerCapabilities()?.resources) return Promise.resolve([])
+  return paginate(
+    (cursor) => client.listResources(cursor === undefined ? undefined : { cursor }, { timeout }),
+    (result) => result.resources,
+  )
+}
+
+function listTools(client: Client, timeout: number) {
+  return Effect.tryPromise({
+    try: () =>
+      paginate(
+        async (cursor) => {
+          const params = cursor === undefined ? undefined : { cursor }
+          try {
+            return await client.listTools(params, { timeout })
+          } catch (error) {
+            if (!(error instanceof Error) || !isOutputSchemaValidationError(error)) throw error
+            return client.request({ method: "tools/list", params }, TolerantListToolsResultSchema, { timeout })
+          }
+        },
+        (result) => result.tools,
+      ),
+    catch: (error) => (error instanceof Error ? error : new Error(String(error))),
+  })
+}
+
+function isOutputSchemaValidationError(error: Error) {
+  return /can't resolve reference|resolves to more than one schema|outputSchema|schema.*reference|reference.*schema/i.test(
+    error.message,
+  )
+}
+
+export * as McpCatalog from "./catalog"