diff --git a/packages/gatekeeper-mcp-portal/README.md b/packages/gatekeeper-mcp-portal/README.md index effb6aa00..e4e643c2f 100644 --- a/packages/gatekeeper-mcp-portal/README.md +++ b/packages/gatekeeper-mcp-portal/README.md @@ -28,7 +28,7 @@ every system the organization has connected in one click. pinning different tools of one upstream server share both the name and the endpoint, so the scope is the only thing that distinguishes them. -The session API — a typed method per tool, plus `callTool`, `getActionResult`, and `listTools` — is +The session API — a typed method per described tool, plus `callTool`, `getActionResult`, and `listTools` — is the same as [`gatekeeper-mcp`](../gatekeeper-mcp/README.md#what-it-provides). Scoping to one server also shrinks what the agent reads: a 57-tool portal generates 57 tool diff --git a/packages/gatekeeper-mcp-portal/src/portal.ts b/packages/gatekeeper-mcp-portal/src/portal.ts index 5291f8d21..efae83602 100644 --- a/packages/gatekeeper-mcp-portal/src/portal.ts +++ b/packages/gatekeeper-mcp-portal/src/portal.ts @@ -26,6 +26,7 @@ import { import type { ToolCatalog } from "@gadgets/mcp-shared/client"; import { classifyTool, + MAX_TOOLS_PER_SERVER, type ServerTrust, } from "@gadgets/mcp-shared/tools"; import { bindingNameFragment, hostOf } from "@gadgets/mcp-shared/util"; @@ -397,7 +398,7 @@ class McpServerConfiguratorUI extends RpcTarget implements McpServerConfigurator // form treats both as having nothing to grant rather than telling them apart. async listServerOptions(): Promise { const { tools, truncated } = await this.#tools(); - if (!looksLikePortal(tools, truncated)) return []; + if (!looksLikePortal(tools, { truncated, cap: MAX_TOOLS_PER_SERVER })) return []; const servers = reconcilePortalServers( await this.#fetchPortalServers(), tools, truncated); @@ -437,7 +438,7 @@ class McpServerConfiguratorUI extends RpcTarget implements McpServerConfigurator const { tools, truncated } = await this.#tools(); requireCompleteCatalogForToolSelection(truncated); const scope: ToolScope = serverId ? { serverId } : {}; - const isPortal = looksLikePortal(tools, truncated); + const isPortal = looksLikePortal(tools, { truncated, cap: MAX_TOOLS_PER_SERVER }); return tools .filter(tool => scopeAllows(scope, tool.name, isPortal)) @@ -517,9 +518,9 @@ export class McpGatekeeperImpl ? `${scope.tools.length} named MCP tool${scope.tools.length === 1 ? "" : "s"} on ` + `${label} \u2014 ${counts}. Other tools are refused.` : scope.serverId - ? `All ${tools.length} MCP tool${plural} of the ` + - `${this.ctx.props.scopeServerName ?? scope.serverId} server on ` + - `${this.ctx.props.serverName} \u2014 ${counts}. Other servers on it are refused.` + ? `All tools of the ${this.ctx.props.scopeServerName ?? scope.serverId} server on ` + + `${this.ctx.props.serverName}; ${tools.length} definition${plural} shown here ` + + `(${counts}). Other servers on it are refused.` : `All ${tools.length} MCP tool${plural} on ${label} \u2014 ${counts}.`; return { diff --git a/packages/gatekeeper-mcp/README.md b/packages/gatekeeper-mcp/README.md index 794231cb6..93317ca88 100644 --- a/packages/gatekeeper-mcp/README.md +++ b/packages/gatekeeper-mcp/README.md @@ -52,12 +52,18 @@ await env.MCP_LINEAR.callTool("search_issues", { query: "state:open" }); Each generated method is a one-line delegate to `callTool`, so the scope check, approval queue, and observation record stay in one place. Tools whose names the RPC layer cannot deliver (`then`, `map`, -`dup`), names that are not identifiers (`2fa`), and both sides of a case collision (`list_issues` -and `listIssues`) get no method and remain callable through `callTool`. - -The agent discovers tools statically: `describeGatekeeper()` sends it the binding name and the full -generated `.d.ts`, where each method carries the tool's own description as JSDoc and states whether -calling it needs approval. `listTools()` exists for runtime enumeration but is rarely needed. +`dup`), names that collide with session methods (`listTools`, `callTool`, `getActionResult`), names +that are not identifiers (`2fa`), and both +sides of a case collision (`list_issues` and `listIssues`) get no method and remain callable through +`callTool`. + +The agent normally discovers tools statically: `describeGatekeeper()` sends it the binding name and +the bounded generated `.d.ts`, where each method carries the tool's own description as JSDoc and +states whether calling it needs approval. A server can publish more definitions than that bounded +catalog holds: `listTools({ search })` returns up to 20 compact summaries from a bounded scan of the +grant, limited to 5,000 tools / 4 MiB; `listTools({ name })` fetches one exact granted definition and +schema under the same bound, and `callTool()` invokes it by wire name. +`listTools()` remains the list of definitions currently described in the generated surface. See `src/types.d.ts` in `@gadgets/mcp-shared` for the base session API. diff --git a/packages/gatekeeper-mcp/src/mcp.ts b/packages/gatekeeper-mcp/src/mcp.ts index 01e736b41..3c8a26a8b 100644 --- a/packages/gatekeeper-mcp/src/mcp.ts +++ b/packages/gatekeeper-mcp/src/mcp.ts @@ -2,7 +2,7 @@ // // Every call is either an observation or an approval-gated action, `readOnlyHint` decides which, // writes are queued rather than performed inline, and per-tool TypeScript is generated from the -// server's schemas so Code Mode works. +// server's schemas so Gadget code gets typed methods. // // The endpoint is whatever a user typed, so annotations never earn auto-approval here and a Gadget // bound to it is owner-only. See `sharing-policy.ts` and the README. @@ -26,6 +26,7 @@ import { import type { ToolCatalog } from "@gadgets/mcp-shared/client"; import { classifyTool, + MAX_TOOLS_PER_SERVER, type ServerTrust, } from "@gadgets/mcp-shared/tools"; import { bindingNameFragment, hostOf } from "@gadgets/mcp-shared/util"; @@ -34,7 +35,7 @@ import { generateSessionTypes, sessionTypeName } from "@gadgets/mcp-shared/schem import { McpAccountBase, type ConnectedServer, type ConnectOutcome } from "@gadgets/mcp-shared/account"; import { generateNonce } from "@gadgets/mcp-shared/connect-nonce"; -import { fetchTools, type ConnectionAccount } from "@gadgets/mcp-shared/connection"; +import { fetchTools, withClient, type ConnectionAccount } from "@gadgets/mcp-shared/connection"; import { McpSessionBase } from "@gadgets/mcp-shared/session"; import { McpFacetBase } from "@gadgets/mcp-shared/facet"; import { looksLikePortal } from "@gadgets/mcp-shared/portal"; @@ -282,8 +283,21 @@ export class GatekeeperUserImpl `do. Connect this endpoint through the MCP Server Portals connector instead.`); } if (scope.tools !== undefined) { + const selected = new Set(scope.tools); validateToolScopeAgainstCatalog( - scope, await fetchTools(this.env, this.#account(), server.endpoint)); + scope, + selected.size === 0 + ? { tools: [], truncated: false } + : await withClient( + this.env, + this.#account(), + server.endpoint, + client => client.listMatchingToolIndex( + selected.size, + tool => selected.has(tool.name), + ), + ), + ); } const props: McpGatekeeperImplProps = { @@ -354,7 +368,10 @@ class McpServerConfiguratorUI extends RpcTarget implements McpServerConfigurator async listToolOptions(): Promise { const { tools, truncated } = await this.#tools(); requireCompleteCatalogForToolSelection(truncated); - const isPortal = looksLikePortal(tools, truncated); + // `fetchTools` lists with the ordinary catalog cap, so that is the cap reaching it would be + // evidence of. Unlike the portal connector, this form refuses a truncated catalog outright + // rather than surveying past it, so `truncated` is already known to be false here. + const isPortal = looksLikePortal(tools, { truncated, cap: MAX_TOOLS_PER_SERVER }); return tools .filter(tool => scopeAllows({}, tool.name, isPortal)) @@ -437,7 +454,8 @@ export class McpGatekeeperImpl const snippet = scope.tools ? `${scope.tools.length} named MCP tool${scope.tools.length === 1 ? "" : "s"} on ` + `${serverName} \u2014 ${counts}. Other tools are refused.` - : `All ${tools.length} MCP tool${plural} on ${serverName} \u2014 ${counts}.`; + : `All tools on ${serverName}; ${tools.length} tool definition${plural} shown here ` + + `(${counts}).`; return { url: this.resourceUrl, diff --git a/packages/mcp-shared/README.md b/packages/mcp-shared/README.md index 7bd031b77..aa89efc83 100644 --- a/packages/mcp-shared/README.md +++ b/packages/mcp-shared/README.md @@ -21,15 +21,16 @@ does so through a named hook (`staticToken`, `mintAccount`), not a private copy. | `client` | Bounded Streamable HTTP transport (`initialize`, `tools/list`, `tools/call`) using official MCP wire types | | `oauth` | Small adapter around the official MCP client's OAuth errors and token revocation gap | | `tools` | The trust boundary: read/action classification, auto-approval eligibility, approval prompts, catalog fingerprinting | -| `schema-to-ts` | JSON Schema to TypeScript, one typed method per tool plus `callTool` overloads | +| `schema-to-ts` | JSON Schema to TypeScript, strict `callTool` overloads plus progressive discovery | | `session-methods` | Installs those methods at runtime, so the generated types are not a fiction | +| `tool-search` | The one query matcher every catalog search uses, so a query cannot mean two things | | `portal` | Gateway detection, tool-name to upstream-server mapping, server listing | | `scope` | The resource-URL scope grammar, and the check every call passes through | | `endpoint` | Validation and host blocklist for a user-supplied endpoint | | `fetch` | Every outbound request; redirects are followed by hand and each hop re-checked, including SDK OAuth fetches | | `account` | Durable Object base and persisted SDK OAuth state: connect, refresh, revocation | | `facet` | Common session, catalog, action, and sharing behavior for connector-owned Durable Object facets | -| `catalog` | One binding's tool list: fetched, cached, scoped to the grant, classified | +| `catalog` | One binding's tool list: fetched, cached, scoped to the grant, classified; plus the cache for tools hydrated past that list | | `connection` | `withClient` — transport sessions, retries, credential-expiry reporting | | `action-store` | Staged to applied/rejected/failed, with a bound on what is retained and a claim so one approval is never sent twice | | `session` | The Gadget-facing capability, and the one path every tool call takes | @@ -109,12 +110,16 @@ Fixed rather than configurable. | Limit | Value | Where | Why | | --- | --- | --- | --- | -| Tools per server | 200 | `tools.ts` | A grant a person can review, and a `.d.ts` an agent can read | +| Described or individually granted tools per server | 200 | `tools.ts` | Bounds the picker and generated `.d.ts`; a server-wide grant can discover additional tools later | | Catalog size | 96 KiB UTF-8 | `client.ts` | Leaves room below Durable Object's 128 KiB per-value limit for the cache wrapper and serialization overhead | +| Filtered discovery scan | 5,000 tools / 4 MiB | `client.ts` | Bounds work spent skipping unrelated tools while searching a large endpoint | +| Search query / results | 200 chars / 20 tools | `tool-search.ts` | Bounds agent-supplied matching work and the summaries returned to it | +| Hydrated definitions per facet | 200 tools / 1 MiB | `catalog.ts` | Bounds definitions fetched individually beyond the described catalog | | Tool description | 4 KB | `client.ts` | As above, per tool, before it reaches storage | | Tool input schema | 20 KB | `client.ts` | Dropped rather than clipped; half a schema is not a schema | -| `tools/list` pages | 50 | `client.ts` | A cursor that never ends would loop until the Worker is killed | +| `tools/list` pages | 50 | `client.ts` | Stops a cursor that never ends; exhaustion truncates catalogs and fails exact/search discovery as a scan limit | | Response body | 1 MiB | `fetch.ts` | Every response is buffered whole before it can be parsed, and a `tools/call` result is otherwise unbounded | +| Bounded outbound operation | 30 seconds | `fetch.ts` | OAuth and discovery callers opt into one deadline covering redirects, pagination, body streaming, and session retry | | Retained result | 128 KB | `action-store.ts` | Held until the Gadget collects it; oversized ones are replaced by a note | | Retained actions | 100 | `action-store.ts` | Records are for collecting a result, not an audit log | | Actions awaiting a decision | 50 | `action-store.ts` | These cannot be pruned, so uncapped they are an unbounded write | diff --git a/packages/mcp-shared/__tests__/client-pagination.test.ts b/packages/mcp-shared/__tests__/client-pagination.test.ts index 8f40ccbe8..f541e760f 100644 --- a/packages/mcp-shared/__tests__/client-pagination.test.ts +++ b/packages/mcp-shared/__tests__/client-pagination.test.ts @@ -1,6 +1,10 @@ import { afterEach, describe, expect, it, vi } from "vitest"; -import { callMayHaveTakenEffect, McpClient } from "../src/client.js"; +import { + callMayHaveTakenEffect, + MAX_TOOL_NAME_CHARS, + McpClient, +} from "../src/client.js"; import { describeCall } from "../src/tools.js"; // Answers every `tools/list` from `pages`, in order, repeating the last one forever. @@ -8,19 +12,39 @@ import { describeCall } from "../src/tools.js"; // Counts requests so a test can prove the client stopped asking rather than merely stopped returning. function stubPages(pages: { tools?: unknown[]; nextCursor?: string }[]): () => number { let calls = 0; - vi.stubGlobal("fetch", async () => { + vi.stubGlobal("fetch", async (_input: unknown, init?: RequestInit) => { const page = pages[Math.min(calls, pages.length - 1)]; calls++; + const id = JSON.parse(String(init?.body)).id; return new Response( - JSON.stringify({ jsonrpc: "2.0", id: calls, result: page }), + JSON.stringify({ jsonrpc: "2.0", id, result: page }), { status: 200, headers: { "Content-Type": "application/json" } }); }); return () => calls; } -afterEach(() => vi.unstubAllGlobals()); +afterEach(() => { + vi.unstubAllGlobals(); + vi.restoreAllMocks(); +}); describe("McpClient.listTools", () => { + it("uses distinct request ids across concurrent client instances", async () => { + const ids: string[] = []; + vi.stubGlobal("fetch", async (_input: unknown, init?: RequestInit) => { + const request = JSON.parse(String(init?.body)); + ids.push(request.id); + return new Response(JSON.stringify({ jsonrpc: "2.0", id: request.id, result: { tools: [] } }), { + status: 200, headers: { "Content-Type": "application/json" }, + }); + }); + const first = new McpClient("https://mcp.example.com/mcp", async () => null); + const second = new McpClient("https://mcp.example.com/mcp", async () => null); + + await Promise.all([first.listTools(10), second.listTools(10)]); + expect(new Set(ids).size).toBe(2); + }); + it("follows cursors until the server stops sending them", async () => { stubPages([ { tools: [{ name: "a" }], nextCursor: "1" }, @@ -32,6 +56,67 @@ describe("McpClient.listTools", () => { expect(truncated).toBe(false); }); + it("treats an empty cursor as an opaque continuation token", async () => { + const cursors: unknown[] = []; + vi.stubGlobal("fetch", async (_input: unknown, init?: RequestInit) => { + const request = JSON.parse(String(init?.body)); + cursors.push(request.params.cursor); + const result = cursors.length === 1 + ? { tools: [{ name: "a" }], nextCursor: "" } + : { tools: [{ name: "b" }] }; + return new Response(JSON.stringify({ jsonrpc: "2.0", id: request.id, result }), { + status: 200, headers: { "Content-Type": "application/json" }, + }); + }); + const client = new McpClient("https://mcp.example.com/mcp", async () => null); + + const { tools, truncated } = await client.listTools(100); + + expect(tools.map(tool => tool.name)).toEqual(["a", "b"]); + expect(truncated).toBe(false); + expect(cursors).toEqual([undefined, ""]); + }); + + it("shares one deadline across every page", async () => { + let now = 0; + let calls = 0; + vi.spyOn(Date, "now").mockImplementation(() => now); + vi.stubGlobal("fetch", async (_input: unknown, init?: RequestInit) => { + calls++; + now += 20; + const id = JSON.parse(String(init?.body)).id; + return new Response(JSON.stringify({ + jsonrpc: "2.0", id, result: { tools: [], nextCursor: "more" }, + }), { status: 200, headers: { "Content-Type": "application/json" } }); + }); + const client = new McpClient( + "https://mcp.example.com/mcp", async () => null, null, { timeoutMs: 30 }); + + await expect(client.listTools(10)).rejects.toThrow(/timed out|timeout/i); + expect(calls).toBe(2); + }); + + it("does not impose a deadline when the caller did not configure one", async () => { + let now = 0; + vi.spyOn(Date, "now").mockImplementation(() => now); + vi.stubGlobal("fetch", async (_input: unknown, init?: RequestInit) => { + now += 60_000; + const request = JSON.parse(String(init?.body)); + return new Response(JSON.stringify({ + jsonrpc: "2.0", + id: request.id, + result: now === 60_000 + ? { tools: [{ name: "a" }], nextCursor: "more" } + : { tools: [{ name: "b" }] }, + }), { status: 200, headers: { "Content-Type": "application/json" } }); + }); + const client = new McpClient("https://mcp.example.com/mcp", async () => null); + + await expect(client.listTools(10)).resolves.toMatchObject({ + tools: [{ name: "a" }, { name: "b" }], + }); + }); + it("stops at maxTools even when the server has more", async () => { const calls = stubPages([{ tools: [{ name: "a" }, { name: "b" }, { name: "c" }], nextCursor: "1" }]); const client = new McpClient("https://mcp.example.com/mcp", async () => null); @@ -42,13 +127,304 @@ describe("McpClient.listTools", () => { expect(calls()).toBe(1); }); - it("gives up on a server that paginates forever without returning tools", async () => { - // `maxTools` cannot stop this on its own: nothing is ever appended, so the tool count never grows - // and the cursor never ends. Without a page cap the loop runs until the Worker is killed. - const calls = stubPages([{ tools: [], nextCursor: "more" }]); + it("applies the tool count limit after filtering", async () => { + stubPages([{ tools: [ + ...Array.from({ length: 250 }, (_, i) => ({ name: `other_${i}` })), + { name: "jira_search" }, + { name: "jira_create" }, + ] }]); + const client = new McpClient("https://mcp.example.com/mcp", async () => null); + const { tools, truncated } = await client.listTools( + 200, + tool => tool.name.startsWith("jira_"), + ); + expect(tools.map(tool => tool.name)).toEqual(["jira_search", "jira_create"]); + expect(truncated).toBe(false); + }); + + it("continues paging until matching tools are found", async () => { + const calls = stubPages([ + { tools: [{ name: "other_a" }], nextCursor: "1" }, + { tools: [{ name: "jira_search" }] }, + ]); + const client = new McpClient("https://mcp.example.com/mcp", async () => null); + const { tools, truncated } = await client.listTools( + 200, + tool => tool.name.startsWith("jira_"), + ); + expect(tools.map(tool => tool.name)).toEqual(["jira_search"]); + expect(truncated).toBe(false); + expect(calls()).toBe(2); + }); + + it("can filter on descriptions before applying catalog budgets", async () => { + stubPages([{ tools: [ + ...Array.from({ length: 250 }, (_, i) => ({ name: `other_${i}` })), + { name: "jira_create_issue", description: "File a ticket in Jira" }, + ] }]); + const client = new McpClient("https://mcp.example.com/mcp", async () => null); + const { tools, truncated } = await client.listTools( + 20, + tool => typeof tool.description === "string" && tool.description.includes("ticket"), + ); + expect(tools.map(tool => tool.name)).toEqual(["jira_create_issue"]); + expect(truncated).toBe(false); + }); + + it("stops paging as soon as an exact tool is found", async () => { + const calls = stubPages([ + { tools: [{ name: "other_a" }], nextCursor: "1" }, + { tools: [{ name: "jira_search" }, { name: "other_b" }], nextCursor: "2" }, + { tools: [{ name: "other_c" }] }, + ]); + const client = new McpClient("https://mcp.example.com/mcp", async () => null); + + await expect(client.findTool("jira_search")).resolves.toMatchObject({ name: "jira_search" }); + expect(calls()).toBe(2); + }); + + it("stops paging when a bounded search has enough matches", async () => { + const calls = stubPages([{ + tools: Array.from({ length: 25 }, (_, i) => ({ name: `jira_tool_${i}` })), + nextCursor: "more", + }]); + const client = new McpClient("https://mcp.example.com/mcp", async () => null); + + const tools = await client.listMatchingToolSummaries( + 20, tool => tool.name.startsWith("jira_")); + expect(tools).toHaveLength(20); + expect(calls()).toBe(1); + }); + + it("returns twenty schema-free search summaries across pages", async () => { + const wireTools = Array.from({ length: 20 }, (_, i) => ({ + name: `jira_tool_${i}`, + title: `Jira tool ${i}`, + description: "x".repeat(4000), + inputSchema: { type: "object", description: "y".repeat(19_000) }, + })); + const calls = stubPages([ + { tools: wireTools.slice(0, 10), nextCursor: "more" }, + { tools: wireTools.slice(10) }, + ]); + const client = new McpClient("https://mcp.example.com/mcp", async () => null); + + const tools = await client.listMatchingToolSummaries(20, () => true); + + expect(tools).toHaveLength(20); + expect(tools[0]).toMatchObject({ name: "jira_tool_0", title: "Jira tool 0" }); + expect(tools.every(tool => tool.inputSchema === undefined)).toBe(true); + expect(calls()).toBe(2); + }); + + it("does not charge filtered-out tools against the byte budget", async () => { + stubPages([{ tools: [ + ...Array.from({ length: 100 }, (_, i) => ({ + name: `other_${i}`, description: "x".repeat(4000), + })), + { name: "jira_search", description: "Search Jira" }, + ] }]); + const client = new McpClient("https://mcp.example.com/mcp", async () => null); + const { tools, truncated } = await client.listTools( + 200, + tool => tool.name.startsWith("jira_"), + ); + expect(tools.map(tool => tool.name)).toEqual(["jira_search"]); + expect(truncated).toBe(false); + }); + + it("can index far more tools than a catalog holds by retaining only names", async () => { + // A catalog of these would exhaust the byte budget long before 400 tools. The index keeps the + stubPages([{ tools: Array.from({ length: 400 }, (_, i) => ({ + name: `server_tool_${i}`, + description: "x".repeat(1000), + inputSchema: { type: "object", description: "y".repeat(1000) }, + annotations: { readOnlyHint: i % 2 === 0 }, + })) }]); + const client = new McpClient("https://mcp.example.com/mcp", async () => null); + const { tools, truncated } = await client.listMatchingToolIndex(500, () => true); + expect(tools).toHaveLength(400); + expect(truncated).toBe(false); + expect(tools[399]).toEqual({ name: "server_tool_399" }); + // The point of the index: no description or schema survives to be stored or shown. + expect(tools[0]).not.toHaveProperty("description"); + expect(tools[0]).not.toHaveProperty("inputSchema"); + }); + + it("can validate selected names after their definitions exhaust the catalog budget", async () => { + const wireTools = Array.from({ length: 8 }, (_, i) => ({ + name: `server_tool_${i}`, + description: "x".repeat(4000), + inputSchema: { type: "object", description: "y".repeat(19_000) }, + })); + stubPages([{ tools: wireTools }]); + const client = new McpClient("https://mcp.example.com/mcp", async () => null); + const selected = new Set(wireTools.map(tool => tool.name)); + const include = (tool: { name: string }) => selected.has(tool.name); + + const catalog = await client.listTools(200, include); + expect(catalog.truncated).toBe(true); + expect(catalog.tools.length).toBeLessThan(wireTools.length); + + const index = await client.listMatchingToolIndex(selected.size, include); + // Stopping once every requested name is found leaves the overall endpoint index incomplete, + // but the returned names are enough to validate this exact grant. + expect(index.truncated).toBe(true); + expect(index.tools.map(tool => tool.name)).toEqual([...selected]); + }); + + it("stops paging once every requested index entry is found", async () => { + const calls = stubPages([{ + tools: [{ name: "jira_search" }, { name: "jira_create" }], + nextCursor: "more", + }]); + const client = new McpClient("https://mcp.example.com/mcp", async () => null); + const requested = new Set(["jira_search", "jira_create"]); + + const index = await client.listMatchingToolIndex( + requested.size, + tool => requested.has(tool.name), + ); + + expect(index.tools.map(tool => tool.name)).toEqual([...requested]); + expect(calls()).toBe(1); + }); + + it("returns a truncated catalog when small pages reach the page limit", async () => { + // One tool per page stays below both the result and scan-tool caps after fifty requests. Without + // a separate page cap, a cursor that never ends would keep the Worker walking indefinitely. + const calls = stubPages([{ tools: [{ name: "other" }], nextCursor: "more" }]); + const client = new McpClient("https://mcp.example.com/mcp", async () => null); + const catalog = await client.listTools(200); + expect(catalog.tools).toHaveLength(50); + expect(catalog.truncated).toBe(true); + expect(calls()).toBe(50); + }); + + it("reports page exhaustion as a scan limit for exact discovery", async () => { + const calls = stubPages([{ tools: [{ name: "other" }], nextCursor: "more" }]); + const client = new McpClient("https://mcp.example.com/mcp", async () => null); + + const error = await client.findTool("missing").catch(caught => caught); + expect(error).toMatchObject({ message: expect.stringMatching(/scan budget/i) }); + expect(callMayHaveTakenEffect(error)).toBe(false); + expect(calls()).toBe(50); + }); + + it("returns a truncated matching index when small pages reach the page limit", async () => { + const calls = stubPages([{ tools: [{ name: "match" }], nextCursor: "more" }]); + const client = new McpClient("https://mcp.example.com/mcp", async () => null); + + const index = await client.listMatchingToolIndex(1000, () => true); + expect(index.tools).toHaveLength(50); + expect(index.truncated).toBe(true); + expect(calls()).toBe(50); + }); + + it("bounds tools scanned while looking for a missing exact name", async () => { + const pages = Array.from({ length: 6 }, (_pageValue, page) => ({ + tools: Array.from( + { length: 1000 }, (_toolValue, i) => ({ name: `other_${page}_${i}` })), + nextCursor: String(page + 1), + })); + const calls = stubPages(pages); + const client = new McpClient("https://mcp.example.com/mcp", async () => null); + + await expect(client.findTool("missing")).rejects.toThrow(/scan budget/i); + expect(calls()).toBe(6); + }); + + it("finds an exact tool before an oversized page crosses the scan limit", async () => { + const calls = stubPages([{ + tools: [ + { name: "target" }, + ...Array.from({ length: 5000 }, (_, i) => ({ name: `other_${i}` })), + ], + }]); + const client = new McpClient("https://mcp.example.com/mcp", async () => null); + + await expect(client.findTool("target")).resolves.toMatchObject({ name: "target" }); + expect(calls()).toBe(1); + }); + + it("bounds bytes scanned from filtered-out tool pages", async () => { + const pages = Array.from({ length: 5 }, (_, page) => ({ + tools: [{ name: `other_${page}`, description: "x".repeat(900_000) }], + nextCursor: String(page + 1), + })); + const calls = stubPages(pages); + const client = new McpClient("https://mcp.example.com/mcp", async () => null); + + await expect(client.findTool("missing")).rejects.toThrow(/scan budget/i); + expect(calls()).toBe(5); + }); + + it("counts ignored JSON-RPC envelope bytes against the scan budget", async () => { + let calls = 0; + vi.stubGlobal("fetch", async (_input: unknown, init?: RequestInit) => { + calls++; + const id = JSON.parse(String(init?.body)).id; + return new Response(JSON.stringify({ + jsonrpc: "2.0", + id, + padding: "x".repeat(900_000), + result: { tools: [], nextCursor: "more" }, + }), { status: 200, headers: { "Content-Type": "application/json" } }); + }); + const client = new McpClient("https://mcp.example.com/mcp", async () => null); + + await expect(client.findTool("missing")).rejects.toThrow(/scan budget/i); + expect(calls).toBe(5); + }); + + it("bounds an unfiltered portal index survey", async () => { + let calls = 0; + vi.stubGlobal("fetch", async (_input: unknown, init?: RequestInit) => { + calls++; + const id = JSON.parse(String(init?.body)).id; + return new Response(JSON.stringify({ + jsonrpc: "2.0", + id, + padding: "x".repeat(900_000), + result: { tools: [], nextCursor: "more" }, + }), { status: 200, headers: { "Content-Type": "application/json" } }); + }); + const client = new McpClient("https://mcp.example.com/mcp", async () => null); + + await expect(client.listMatchingToolIndex(1000, () => true)) + .resolves.toEqual({ tools: [], truncated: true }); + expect(calls).toBe(5); + }); + + it("returns collected catalog matches when later pages exceed the scan budget", async () => { + const pages = [ + { tools: [{ name: "jira_search" }], nextCursor: "1" }, + ...Array.from({ length: 5 }, (_, page) => ({ + tools: [{ name: `other_${page}`, description: "x".repeat(900_000) }], + nextCursor: String(page + 2), + })), + ]; + stubPages(pages); + const client = new McpClient("https://mcp.example.com/mcp", async () => null); + + await expect(client.listTools(200, tool => tool.name.startsWith("jira_"))) + .resolves.toEqual({ tools: [{ name: "jira_search" }], truncated: true }); + }); + + it("fails search rather than presenting a scan-limited prefix as complete", async () => { + const pages = [ + { tools: [{ name: "jira_search" }], nextCursor: "1" }, + ...Array.from({ length: 5 }, (_, page) => ({ + tools: [{ name: `other_${page}`, description: "x".repeat(900_000) }], + nextCursor: String(page + 2), + })), + ]; + stubPages(pages); const client = new McpClient("https://mcp.example.com/mcp", async () => null); - await expect(client.listTools(200)).rejects.toThrow(/kept paginating/); - expect(calls()).toBeLessThanOrEqual(50); + + await expect(client.listMatchingToolSummaries( + 20, tool => tool.name.startsWith("jira_"))) + .rejects.toThrow(/scan budget/i); }); it("ignores nameless entries rather than counting them towards the cap", async () => { @@ -57,6 +433,15 @@ describe("McpClient.listTools", () => { expect((await client.listTools(10)).tools.map(tool => tool.name)).toEqual(["real"]); }); + it("ignores tool names too large to retain or accept back from a Gadget", async () => { + stubPages([{ tools: [ + { name: "x".repeat(MAX_TOOL_NAME_CHARS + 1) }, + { name: "real" }, + ] }]); + const client = new McpClient("https://mcp.example.com/mcp", async () => null); + expect((await client.listTools(10)).tools.map(tool => tool.name)).toEqual(["real"]); + }); + it("drops a description the server did not send as a string", async () => { // A non-string passed the length cap untouched and reached the approval prompt, where // `quoteUntrusted` called `.replace` on it and took down every action call for the connection. @@ -106,17 +491,20 @@ describe("McpClient.listTools", () => { let cancelled = false; let streamController: ReadableStreamDefaultController | undefined; const encoder = new TextEncoder(); - vi.stubGlobal("fetch", async () => new Response(new ReadableStream({ + vi.stubGlobal("fetch", async (_input: unknown, init?: RequestInit) => { + const id = JSON.parse(String(init?.body)).id; + return new Response(new ReadableStream({ start(controller) { streamController = controller; controller.enqueue(encoder.encode( 'data: {"jsonrpc":"2.0","method":"notifications/progress"}\n\n')); controller.enqueue(encoder.encode( - 'data: {"jsonrpc":"2.0","id":1,"result":{"content":[')); + `data: ${JSON.stringify({ jsonrpc: "2.0", id }).slice(0, -1)},"result":{"content":[`)); controller.enqueue(encoder.encode('{"type":"text","text":"done"}]}}\n\r')); }, cancel() { cancelled = true; }, - }), { status: 200, headers: { "Content-Type": "Text/Event-Stream" } })); + }), { status: 200, headers: { "Content-Type": "Text/Event-Stream" } }); + }); const client = new McpClient("https://mcp.example.com/mcp", async () => null); const call = client.callTool("anything", {}); @@ -181,9 +569,12 @@ describe("McpClient.listTools", () => { describe("error text a server wrote", () => { // Answers every request with a JSON-RPC error carrying `message`. function stubError(message: string) { - vi.stubGlobal("fetch", async () => new Response( - JSON.stringify({ jsonrpc: "2.0", id: 1, error: { code: -32000, message } }), - { status: 200, headers: { "Content-Type": "application/json" } })); + vi.stubGlobal("fetch", async (_input: unknown, init?: RequestInit) => { + const id = JSON.parse(String(init?.body)).id; + return new Response( + JSON.stringify({ jsonrpc: "2.0", id, error: { code: -32000, message } }), + { status: 200, headers: { "Content-Type": "application/json" } }); + }); } it("redacts the credential this Worker just sent, if the server echoes it back", async () => { diff --git a/packages/mcp-shared/__tests__/connection.test.ts b/packages/mcp-shared/__tests__/connection.test.ts index 85489481c..de7d029f3 100644 --- a/packages/mcp-shared/__tests__/connection.test.ts +++ b/packages/mcp-shared/__tests__/connection.test.ts @@ -185,3 +185,53 @@ it("reports credentials expired when session recovery is rejected", async () => expect(error).toBeInstanceOf(McpCallNotDispatchedError); expect(expired).toBe(1); }); + +it("gives a retried listing a fresh discovery budget after session recovery", async () => { + let initialPages = 0; + let retryPages = 0; + let recovered = false; + vi.stubGlobal("fetch", async (_input: unknown, init?: RequestInit) => { + const request = JSON.parse(String(init?.body)); + if (request.method === "initialize") { + recovered = true; + return new Response(JSON.stringify({ jsonrpc: "2.0", id: request.id, result: {} }), { + headers: { "Content-Type": "application/json", "Mcp-Session-Id": "new-session" }, + }); + } + if (request.method === "notifications/initialized") { + return new Response(null, { status: 202 }); + } + if (!recovered) { + initialPages++; + if (initialPages === 50) return new Response(null, { status: 404 }); + return new Response(JSON.stringify({ + jsonrpc: "2.0", + id: request.id, + result: { tools: [], nextCursor: String(initialPages) }, + }), { headers: { "Content-Type": "application/json" } }); + } + retryPages++; + return new Response(JSON.stringify({ + jsonrpc: "2.0", + id: request.id, + result: retryPages === 1 + ? { tools: [{ name: "first" }], nextCursor: "more" } + : { tools: [{ name: "second" }] }, + }), { headers: { "Content-Type": "application/json" } }); + }); + const account: ConnectionAccount = { + async getConnection() { + return { authorization: null, sessionId: "expired-session", generation: 1 }; + }, + async assertConnectionCurrent() {}, + async setMcpSessionId() { return true; }, + async noteCredentialsExpired() {}, + }; + + const catalog = await withClient({}, account, "https://mcp.example.com", client => + client.listTools(10)); + expect(catalog.tools.map(tool => tool.name)).toEqual(["first", "second"]); + expect(catalog.truncated).toBe(false); + expect(initialPages).toBe(50); + expect(retryPages).toBe(2); +}); diff --git a/packages/mcp-shared/__tests__/facet.test.ts b/packages/mcp-shared/__tests__/facet.test.ts index 1f157209a..38112314e 100644 --- a/packages/mcp-shared/__tests__/facet.test.ts +++ b/packages/mcp-shared/__tests__/facet.test.ts @@ -1,8 +1,11 @@ -import { expect, it } from "vitest"; +import { expect, it, vi } from "vitest"; import { McpFacetBase } from "../src/facet.js"; import { McpSessionBase } from "../src/session.js"; -import { classifyTool, type ClassifiedTool, type ServerTrust } from "../src/tools.js"; +import { classifyTool, type ServerTrust } from "../src/tools.js"; +import type { McpClient, McpTool } from "../src/client.js"; +import type { ToolScope } from "../src/scope.js"; +import type { ScopedCatalog } from "../src/catalog.js"; import type { ConnectionAccount } from "../src/connection.js"; import type { ResourceDescription } from "@gadgets/workshop-shared/gatekeeper"; @@ -17,11 +20,19 @@ class TestSession extends McpSessionBase {} class TestFacet extends McpFacetBase { - catalog: Promise = Promise.resolve([ - classifyTool({ name: "list_issues", annotations: { readOnlyHint: true } } as never, "byo"), - ]); + catalogResult: Promise = Promise.resolve({ + isPortal: false, + truncated: false, + tools: [ + classifyTool({ name: "list_issues", annotations: { readOnlyHint: true } } as never, "byo"), + ], + }); + catalogReads = 0; + remoteTools: McpTool[] = []; + remoteCalls = 0; + beforeCatalogRead: (() => Promise) | undefined; protected get log() { return log; } protected get trust(): ServerTrust { return "byo"; } @@ -32,12 +43,36 @@ class TestFacet extends McpFacetBase { throw new Error("not used"); } getTypeScriptTypes(): Promise { throw new Error("not used"); } get serverName() { return "Test"; } - override tools() { return this.catalog; } + protected override async catalog() { + this.catalogReads++; + await this.beforeCatalogRead?.(); + return this.catalogResult; + } + override async call( + fn: (client: McpClient) => Promise, + ): Promise { + this.remoteCalls++; + const client = { + findTool: async (name: string) => this.remoteTools.find(tool => tool.name === name), + listTools: async ( + _maxTools: number, + include: (tool: McpTool) => boolean, + ) => ({ tools: this.remoteTools.filter(include), truncated: false }), + listMatchingToolSummaries: async ( + maxTools: number, + include: (tool: McpTool) => boolean, + ) => this.remoteTools.filter(include).slice(0, maxTools), + } as unknown as McpClient; + return fn(client); + } + runDiscoveryTest(operation: () => Promise): Promise { + return this.runDiscovery(operation); + } } -function facet() { +function facet(scope: ToolScope = {}) { const ctx = { - props: { endpoint: "https://example.com/mcp", scope: {} }, + props: { endpoint: "https://example.com/mcp", scope }, storage: { kv: {} }, }; return new TestFacet(ctx as never, {}); @@ -53,7 +88,7 @@ it("builds tool methods and falls back to the plain session when catalog loading const dynamic = await subject.startSession(queue as never); expect("listIssues" in dynamic).toBe(true); - subject.catalog = Promise.reject(new Error("offline")); + subject.catalogResult = Promise.reject(new Error("offline")); const fallback = await subject.startSession(queue as never); expect("listIssues" in fallback).toBe(false); expect(log.warnings).toContain("starting session without per-tool methods"); @@ -63,3 +98,94 @@ it("keeps facets owner-only using the connector's resource label", async () => { await expect(facet().addObserver("observer", {} as never)) .rejects.toThrow(/test server.*only be opened by its owner/s); }); + +it("discovers and resolves tools beyond the initially described catalog", async () => { + const subject = facet(); + subject.catalogResult = Promise.resolve({ tools: [], isPortal: false, truncated: true }); + subject.remoteTools = [ + { name: "search_issues", description: "Search issues", annotations: { readOnlyHint: true } }, + { name: "create_issue", description: "Create an issue" }, + ]; + + await expect(subject.searchTools("search")).resolves.toMatchObject([{ + tool: { name: "search_issues" }, mode: "read", + }]); + await expect(subject.findTool("create_issue")).resolves.toMatchObject({ + tool: { name: "create_issue" }, mode: "action", + }); +}); + +it("answers searches from a complete catalog without rescanning the endpoint", async () => { + const subject = facet(); + subject.remoteTools = [{ name: "search_issues", description: "Remote copy" }]; + + await expect(subject.searchTools("issues")).resolves.toMatchObject([{ + tool: { name: "list_issues" }, mode: "read", + }]); + expect(subject.remoteCalls).toBe(0); +}); + +it("does not hydrate a new portal-native tool past a complete non-portal catalog", async () => { + const subject = facet(); + subject.remoteTools = [{ name: "portal_toggle_servers" }]; + + await expect(subject.findTool("portal_toggle_servers")).resolves.toBeUndefined(); + expect(subject.remoteCalls).toBe(0); +}); + +it("rejects a name outside the grant before loading the catalog or calling the endpoint", async () => { + const subject = facet({ tools: ["allowed"] }); + subject.catalogResult = Promise.resolve({ tools: [], isPortal: false, truncated: false }); + + await expect(subject.findTool("forbidden")).resolves.toBeUndefined(); + expect(subject.catalogReads).toBe(0); + expect(subject.remoteCalls).toBe(0); +}); + +it("does not hydrate portal-native tools from a portal's incomplete catalog", async () => { + const subject = facet(); + subject.catalogResult = Promise.resolve({ tools: [], isPortal: true, truncated: true }); + subject.remoteTools = [{ name: "portal_toggle_servers" }]; + + await expect(subject.findTool("portal_toggle_servers")).resolves.toBeUndefined(); + expect(subject.remoteCalls).toBe(0); +}); + +it("bounds concurrent discovery work across distinct requests", async () => { + const subject = facet(); + let active = 0; + let maxActive = 0; + let release!: () => void; + const gate = new Promise(resolve => { release = resolve; }); + const calls = Array.from({ length: 12 }, () => subject.runDiscoveryTest(async () => { + active++; + maxActive = Math.max(maxActive, active); + await gate; + active--; + })); + + await vi.waitFor(() => expect(active).toBe(4)); + expect(maxActive).toBe(4); + release(); + await Promise.all(calls); + expect(maxActive).toBe(4); +}); + +it("bounds concurrent catalog reads", async () => { + const subject = facet(); + let active = 0; + let release!: () => void; + const gate = new Promise(resolve => { release = resolve; }); + subject.beforeCatalogRead = async () => { + active++; + await gate; + active--; + }; + + const searches = Array.from({ length: 12 }, () => subject.searchTools("issues")); + + await vi.waitFor(() => expect(active).toBeGreaterThan(0)); + expect(active).toBe(4); + release(); + await Promise.all(searches); +}); diff --git a/packages/mcp-shared/__tests__/hydrated-tools.test.ts b/packages/mcp-shared/__tests__/hydrated-tools.test.ts new file mode 100644 index 000000000..77499e6ce --- /dev/null +++ b/packages/mcp-shared/__tests__/hydrated-tools.test.ts @@ -0,0 +1,82 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; + +import { CATALOG_TTL_MS, HydratedTools } from "../src/catalog.js"; +import type { McpTool } from "../src/client.js"; + +const tool = (name: string): McpTool => ({ name }); + +afterEach(() => { + vi.useRealTimers(); +}); + +describe("HydratedTools", () => { + it("fetches a tool once per TTL", async () => { + vi.useFakeTimers(); + const load = vi.fn(async (name: string) => tool(name)); + const cache = new HydratedTools(); + + await expect(cache.resolve("a", load)).resolves.toEqual(tool("a")); + await expect(cache.resolve("a", load)).resolves.toEqual(tool("a")); + expect(load).toHaveBeenCalledTimes(1); + + vi.advanceTimersByTime(CATALOG_TTL_MS + 1); + await expect(cache.resolve("a", load)).resolves.toEqual(tool("a")); + expect(load).toHaveBeenCalledTimes(2); + }); + + it("coalesces concurrent loads of one tool", async () => { + let release!: (value: McpTool) => void; + const load = vi.fn((name: string) => new Promise(resolve => { + release = resolve; + })); + const cache = new HydratedTools(); + + const first = cache.resolve("a", load); + const second = cache.resolve("a", load); + expect(load).toHaveBeenCalledTimes(1); + + release(tool("a")); + await expect(Promise.all([first, second])).resolves.toEqual([tool("a"), tool("a")]); + }); + + it("remembers that a tool does not exist", async () => { + // Without this, a Gadget naming a tool that is not there sends one full paginated listing to the + // endpoint per call -- a cheap way to make this gatekeeper hammer a server on an agent's behalf. + const load = vi.fn(async () => undefined); + const cache = new HydratedTools(); + + await expect(cache.resolve("ghost", load)).resolves.toBeUndefined(); + await expect(cache.resolve("ghost", load)).resolves.toBeUndefined(); + await expect(cache.resolve("ghost", load)).resolves.toBeUndefined(); + expect(load).toHaveBeenCalledTimes(1); + }); + + it("bounds how much a Gadget can make it hold", async () => { + const cache = new HydratedTools(); + const load = vi.fn(async (name: string) => tool(name)); + const total = HydratedTools.MAX_ENTRIES + 10; + for (let i = 0; i < total; i++) await cache.resolve(`t${i}`, load); + expect(load).toHaveBeenCalledTimes(total); + + // The most recent are still cached; the oldest were evicted to make room for them. + load.mockClear(); + await cache.resolve(`t${total - 1}`, load); + expect(load).not.toHaveBeenCalled(); + await cache.resolve("t0", load); + expect(load).toHaveBeenCalledTimes(1); + }); + + it("evicts definitions to keep aggregate retained bytes bounded", async () => { + const cache = new HydratedTools(); + const large = (name: string): McpTool => ({ + name, + description: "x".repeat(4000), + inputSchema: { type: "object", description: "y".repeat(20_000) }, + }); + for (let i = 0; i < 50; i++) await cache.resolve(`large_${i}`, async name => large(name)); + + const reload = vi.fn(async (name: string) => large(name)); + await cache.resolve("large_0", reload); + expect(reload).toHaveBeenCalledTimes(1); + }); +}); diff --git a/packages/mcp-shared/__tests__/portal.test.ts b/packages/mcp-shared/__tests__/portal.test.ts index 6ba477364..a8004d6d5 100644 --- a/packages/mcp-shared/__tests__/portal.test.ts +++ b/packages/mcp-shared/__tests__/portal.test.ts @@ -11,16 +11,22 @@ import { scopeAllows } from "../src/scope.js"; import { MAX_TOOLS_PER_SERVER } from "../src/tools.js"; import type { McpTool } from "../src/client.js"; +// What a caller passes when it fetched with the ordinary catalog cap and the listing completed. +const COMPLETE = { truncated: false, cap: MAX_TOOLS_PER_SERVER }; + describe("portal detection", () => { it("recognizes a portal by its built-in server listing", () => { - expect(looksLikePortal([{ name: "portal_list_servers" }, { name: "gh_list_issues" }])) + expect(looksLikePortal([{ name: "ordinary_tool" }], COMPLETE)).toBe(false); + expect(looksLikePortal( + [{ name: "portal_list_servers" }, { name: "gh_list_issues" }], COMPLETE)) .toBe(true); }); it("does not infer a portal from prefixed-looking names alone", () => { // Plenty of ordinary servers name tools `verb_noun`. Only the portal's own tool is evidence. - expect(looksLikePortal([{ name: "search" }, { name: "create_issue" }])).toBe(false); - expect(looksLikePortal([{ name: "list_issues" }, { name: "create_issue" }])).toBe(false); + expect(looksLikePortal([{ name: "search" }, { name: "create_issue" }], COMPLETE)).toBe(false); + expect(looksLikePortal([{ name: "list_issues" }, { name: "create_issue" }], COMPLETE)) + .toBe(false); }); it("assumes a portal when the catalog was truncated", () => { @@ -29,8 +35,8 @@ describe("portal detection", () => { // from evidence we could not have seen would fail open; the next test shows what it costs. const truncated = Array.from( { length: MAX_TOOLS_PER_SERVER }, (_unused, index) => ({ name: `gh_tool_${index}` })); - expect(looksLikePortal(truncated)).toBe(true); - expect(looksLikePortal(truncated.slice(0, -1))).toBe(false); + expect(looksLikePortal(truncated, COMPLETE)).toBe(true); + expect(looksLikePortal(truncated.slice(0, -1), COMPLETE)).toBe(false); }); it("assumes a portal when the byte budget cut the catalog short", () => { @@ -39,8 +45,8 @@ describe("portal detection", () => { // would then be granted at its bare endpoint, leaving `portal_toggle_servers` callable by a // Gadget that could widen its own reach with it. const short = [{ name: "gh_list_issues" }, { name: "gh_create_issue" }]; - expect(looksLikePortal(short, true)).toBe(true); - expect(looksLikePortal(short, false)).toBe(false); + expect(looksLikePortal(short, { truncated: true, cap: MAX_TOOLS_PER_SERVER })).toBe(true); + expect(looksLikePortal(short, COMPLETE)).toBe(false); }); it("keeps refusing portal-native tools when the listing tool falls outside the cap", () => { @@ -51,7 +57,8 @@ describe("portal detection", () => { ...Array.from( { length: MAX_TOOLS_PER_SERVER - 1 }, (_unused, index) => ({ name: `gh_tool_${index}` })), ]; - expect(scopeAllows({}, "portal_toggle_single_server", looksLikePortal(truncated))).toBe(false); + expect(scopeAllows({}, "portal_toggle_single_server", looksLikePortal(truncated, COMPLETE))) + .toBe(false); }); }); diff --git a/packages/mcp-shared/__tests__/schema-to-ts.test.ts b/packages/mcp-shared/__tests__/schema-to-ts.test.ts index 3e3538e35..fdbecad41 100644 --- a/packages/mcp-shared/__tests__/schema-to-ts.test.ts +++ b/packages/mcp-shared/__tests__/schema-to-ts.test.ts @@ -48,6 +48,23 @@ function expectTypeScriptToCompile(source: string): void { expect(errors).toEqual([]); } +function expectTypeScriptProgramToCompile(source: string): void { + const fileName = "generated.ts"; + const options: ts.CompilerOptions = { noEmit: true, strict: true }; + const host = ts.createCompilerHost(options); + const getSourceFile = host.getSourceFile.bind(host); + host.getSourceFile = (name, languageVersion, onError, shouldCreateNewSourceFile) => + name === fileName + ? ts.createSourceFile(name, source, languageVersion, true, ts.ScriptKind.TS) + : getSourceFile(name, languageVersion, onError, shouldCreateNewSourceFile); + host.fileExists = name => name === fileName || ts.sys.fileExists(name); + host.readFile = name => name === fileName ? source : ts.sys.readFile(name); + const errors = ts.getPreEmitDiagnostics(ts.createProgram([fileName], options, host)) + .filter(diagnostic => diagnostic.file?.fileName === fileName) + .map(diagnostic => ts.flattenDiagnosticMessageText(diagnostic.messageText, "\n")); + expect(errors).toEqual([]); +} + describe("sessionTypeName", () => { const url = "https://acme.example/mcp"; @@ -103,6 +120,34 @@ describe("generateSessionTypes", { timeout: 15_000 }, () => { expect(output).toContain(`export interface ${name} {`); }); + it("keeps known tool overloads strict while accepting dynamic names", () => { + const output = generate([tool({ + name: "search", + inputSchema: { + type: "object", + properties: { q: { type: "string" } }, + required: ["q"], + }, + })], MCP_BASE_TYPES); + const name = sessionTypeName("acme-crm", "https://acme.example/mcp"); + expectTypeScriptProgramToCompile(`${output} +declare const session: ${name}; +// @ts-expect-error known tool requires its arguments +session.callTool("search"); +// @ts-expect-error known tool keeps its generated schema +session.callTool("search", { q: 123 }); +declare const discovered: string; +session.callTool(discovered, { anything: true }); +declare const options: McpToolListOptions; +session.listTools(options); +// @ts-expect-error search and name are mutually exclusive +session.listTools({ search: "issues", name: "search" }); +declare const ambiguousOptions: { search: string; name: string }; +// @ts-expect-error variables containing both selectors are also rejected +session.listTools(ambiguousOptions); +`); + }); + it("prepends the base types verbatim so the file is self-contained", () => { expect(generate([])).toContain("// base"); }); @@ -399,11 +444,23 @@ describe("generateSessionTypes", { timeout: 15_000 }, () => { expect(output).toContain('callTool(name: "ping", args?: Record)'); }); - it("always exposes listTools and getActionResult", () => { + it("always exposes discovery, generic calls, and action results", () => { const output = generate([]); expect(output).toContain("listTools(): Promise;"); + expect(output).toContain( + "listTools(options: { search: string; name?: never }): Promise;"); + expect(output).toContain( + "listTools(options: { name: string; search?: never }): Promise;"); + expect(output).toContain("callTool("); expect(output).toContain("getActionResult(actionId: number): Promise;"); }); + + it("keeps discovery stable when an upstream tool collides with its method name", () => { + const output = generate([tool({ name: "search_tools" })]); + expect(output).toContain("searchTools(): Promise;"); + expect(output).toContain("listTools(options: { search: string; name?: never })"); + expect(output).toContain('callTool(name: "search_tools"'); + }); }); // The generated types promise methods that `session-methods.ts` actually installs. These pin the two diff --git a/packages/mcp-shared/__tests__/scope.test.ts b/packages/mcp-shared/__tests__/scope.test.ts index cf3d266c7..9680ea642 100644 --- a/packages/mcp-shared/__tests__/scope.test.ts +++ b/packages/mcp-shared/__tests__/scope.test.ts @@ -10,6 +10,7 @@ import { scopeAllows, validateToolScopeAgainstCatalog, } from "../src/scope.js"; +import { MAX_TOOLS_PER_SERVER } from "../src/tools.js"; const ENDPOINT = "https://portal.example.com/mcp"; @@ -70,6 +71,15 @@ describe("parseToolScope", () => { const emptied = { serverId: undefined, tools: [] }; expect(parseToolScope(formatToolScope(ENDPOINT, emptied))).toEqual(emptied); }); + + it("rejects oversized named-tool grants before they reach a connector", () => { + const fragment = new URLSearchParams(); + for (let i = 0; i <= MAX_TOOLS_PER_SERVER; i++) fragment.append("tool", `tool_${i}`); + expect(() => parseToolScope(`${ENDPOINT}#${fragment}`)).toThrow(/at most 200/); + expect(() => formatToolScope(ENDPOINT, { + tools: Array.from({ length: MAX_TOOLS_PER_SERVER + 1 }, (_, i) => `tool_${i}`), + })).toThrow(/at most 200/); + }); }); describe("formatToolScope", () => { @@ -173,9 +183,15 @@ describe("validateToolScopeAgainstCatalog", () => { .not.toThrow(); }); - it("refuses to validate named tools from a truncated catalog", () => { + it("accepts named tools that were found before an exact-name listing stopped", () => { expect(() => validateToolScopeAgainstCatalog( { tools: ["gh_list_issues"] }, { ...catalog, truncated: true })) + .not.toThrow(); + }); + + it("refuses an unresolved named tool when its listing was truncated", () => { + expect(() => validateToolScopeAgainstCatalog( + { tools: ["missing"] }, { ...catalog, truncated: true })) .toThrow(/truncated/i); }); diff --git a/packages/mcp-shared/__tests__/session-methods-e2e.test.ts b/packages/mcp-shared/__tests__/session-methods-e2e.test.ts index 1e85987f2..d9e067641 100644 --- a/packages/mcp-shared/__tests__/session-methods-e2e.test.ts +++ b/packages/mcp-shared/__tests__/session-methods-e2e.test.ts @@ -24,7 +24,10 @@ it("every method the .d.ts promises is really installed, and routes to the right trust: "byo", tools, }); const promised = [...dts.matchAll(/^ {2}([a-z]\w*)\(/gm)].map(match => match[1]) - .filter(name => !["listTools", "callTool", "getActionResult"].includes(name)); + .filter(name => ![ + "listTools", "callTool", + "getActionResult", + ].includes(name)); class Base { called: string | null = null; diff --git a/packages/mcp-shared/__tests__/session-methods.test.ts b/packages/mcp-shared/__tests__/session-methods.test.ts index aca385b22..4559f1abb 100644 --- a/packages/mcp-shared/__tests__/session-methods.test.ts +++ b/packages/mcp-shared/__tests__/session-methods.test.ts @@ -67,6 +67,16 @@ describe("toolMethodNames", () => { } }); + it("preserves delegates that predate progressive discovery", () => { + expect([...toolMethodNames([ + tool("search_tools"), tool("describe_tool"), tool("call_discovered_tool"), + ])]).toEqual([ + ["searchTools", "search_tools"], + ["describeTool", "describe_tool"], + ["callDiscoveredTool", "call_discovered_tool"], + ]); + }); + it("drops both sides of a collision rather than shadowing one", () => { const names = toolMethodNames([tool("list_issues"), tool("listIssues"), tool("search")]); expect([...names]).toEqual([["search", "search"]]); diff --git a/packages/mcp-shared/__tests__/session.test.ts b/packages/mcp-shared/__tests__/session.test.ts index 8077622f6..081e83ff4 100644 --- a/packages/mcp-shared/__tests__/session.test.ts +++ b/packages/mcp-shared/__tests__/session.test.ts @@ -1,6 +1,7 @@ import { expect, it } from "vitest"; import { McpSessionBase, type McpSessionHost, type StoredAction } from "../src/session.js"; +import { MAX_TOOL_NAME_CHARS } from "../src/client.js"; import { classifyTool } from "../src/tools.js"; it("reports an execution failure distinctly from a rejected approval", async () => { @@ -40,7 +41,7 @@ it("tells an agent to return a pending action so its approval can appear in chat serverName: "Jira", endpoint: "https://mcp.example.com", scope: { serverId: "jira" }, - tools: async () => [entry], + findTool: async () => entry, stageAction: () => staged, discardStagedAction() {}, actionKindFor: () => ({ tag: "jira:create", label: "Create issue" }), @@ -54,3 +55,205 @@ it("tells an agent to return a pending action so its approval can appear in chat expect(result.message).toContain("return from this executeCode call"); expect(result.message).not.toMatch(/poll/i); }); + +it("searches progressively discovered tools and records the catalog read", async () => { + const found = classifyTool({ + name: "jira_search_issues", + description: "Search Jira issues", + annotations: { readOnlyHint: true }, + }, "byo"); + const observations: unknown[] = []; + const host = { + serverName: "Jira", + endpoint: "https://mcp.example.com", + scope: { serverId: "jira" }, + searchTools: async () => [found], + } as unknown as McpSessionHost; + const queue = { + authorizeObservation: (description: unknown) => { observations.push(description); }, + }; + const session = new McpSessionBase(host, queue as never); + + await expect(session.listTools({ search: "issues" })).resolves.toEqual([{ + name: "jira_search_issues", + description: "Search Jira issues", + mode: "read", + classifiedBy: "server-annotation", + inputSchema: undefined, + title: undefined, + }]); + expect(observations).toHaveLength(1); +}); + +it("calls a tool resolved beyond the initial generated catalog", async () => { + const expanded = classifyTool({ + name: "jira_search_issues", + annotations: { readOnlyHint: true }, + }, "byo"); + const host = { + serverName: "Jira", + endpoint: "https://mcp.example.com", + scope: { serverId: "jira" }, + tools: async () => [], + findTool: async () => expanded, + call: async (fn: (client: never) => Promise) => fn({ + callTool: async () => ({ content: [{ type: "text", text: "PROJ-1" }] }), + } as never), + } as unknown as McpSessionHost; + const queue = { authorizeObservation() {} }; + const session = new McpSessionBase(host, queue as never); + + await expect(session.callTool("jira_search_issues", { query: "open" })).resolves.toMatchObject({ + status: "ok", + text: "PROJ-1", + }); +}); + +it("identifies the tool in a describe observation", async () => { + const observations: { description: string }[] = []; + const host = { + serverName: "Jira", + endpoint: "https://mcp.example.com", + scope: { serverId: "jira" }, + findTool: async () => classifyTool({ name: "jira_search_issues" }, "byo"), + } as unknown as McpSessionHost; + const queue = { + authorizeObservation: (description: { description: string }) => { + observations.push(description); + }, + }; + const session = new McpSessionBase(host, queue as never); + + await session.listTools({ name: "jira_search_issues" }); + expect(observations[0].description).toContain("jira_search_issues"); +}); + +it("names the grant, not the server, when a scoped binding lacks the tool", async () => { + // On a scoped binding the tool may well exist on the endpoint. "No such tool" would send an agent + // hunting for a typo it will not find, so both entry points have to say the same thing. + const host = { + serverName: "Jira", + endpoint: "https://mcp.example.com", + scope: { serverId: "jira" }, + findTool: async () => undefined, + } as unknown as McpSessionHost; + const session = new McpSessionBase(host, { authorizeObservation() {} } as never); + + await expect(session.listTools({ name: "gh_list_issues" })).resolves.toEqual([]); + await expect(session.callTool("gh_list_issues")) + .rejects.toThrow('This binding does not grant a tool named "gh_list_issues".'); +}); + +it("says the server has no such tool when the whole endpoint was granted", async () => { + const host = { + serverName: "Jira", + endpoint: "https://mcp.example.com", + scope: {}, + findTool: async () => undefined, + } as unknown as McpSessionHost; + const session = new McpSessionBase(host, { authorizeObservation() {} } as never); + + await expect(session.listTools({ name: "nope" })).resolves.toEqual([]); +}); + +it("records what was searched, with the agent's text defused", async () => { + // The query is the agent's, and an observation is read by a person: left alone it can close the + // markdown it sits in and carry on in the record's own voice. + const observations: { description: string }[] = []; + const host = { + serverName: "Jira", + endpoint: "https://mcp.example.com", + scope: { serverId: "jira" }, + searchTools: async () => [], + } as unknown as McpSessionHost; + const queue = { + authorizeObservation: (d: { description: string }) => { observations.push(d); }, + }; + const session = new McpSessionBase(host, queue as never); + + await session.listTools({ search: "issues `**Approved**`" }); + expect(observations[0].description).toContain("issues Approved"); + expect(observations[0].description).toContain("returned 0 match(es)"); + expect(observations[0].description).not.toContain("**Approved**"); +}); + +it("returns the same compact summary shape from a complete local catalog", async () => { + const host = { + serverName: "Jira", + endpoint: "https://mcp.example.com", + scope: { serverId: "jira" }, + searchTools: async () => [classifyTool({ + name: "jira_search_issues", + description: "x".repeat(4000), + inputSchema: { type: "object" }, + }, "byo")], + } as unknown as McpSessionHost; + const session = new McpSessionBase(host, { authorizeObservation() {} } as never); + + const [summary] = await session.listTools({ search: "issues" }); + + expect(summary.description).toBe(`${"x".repeat(256)}\u2026`); + expect(summary).not.toHaveProperty("inputSchema"); +}); + +it("refuses an empty or oversized query before calling the endpoint", async () => { + let searches = 0; + const searchTools = async () => { searches++; return []; }; + const host = { + serverName: "Jira", + endpoint: "https://mcp.example.com", + scope: {}, + searchTools, + } as unknown as McpSessionHost; + const session = new McpSessionBase(host, { authorizeObservation() {} } as never); + + await expect(session.listTools({ search: " " })).rejects.toThrow(/non-empty query/); + await expect(session.listTools({ search: " _ - " })).rejects.toThrow(/search terms/); + // Bounded on the trimmed text, which is what is actually searched and recorded. + await expect(session.listTools({ search: `${" ".repeat(50)}${"x".repeat(201)}` })) + .rejects.toThrow(/at most 200 characters/); + await expect(session.listTools({ search: ` ${"x".repeat(200)} ` })).resolves.toEqual([]); + expect(searches).toBe(1); +}); + +it("refuses ambiguous progressive list options", async () => { + const host = { serverName: "Jira", endpoint: "https://mcp.example.com", scope: {} } as unknown as McpSessionHost; + const session = new McpSessionBase(host, { authorizeObservation() {} } as never); + + await expect(session.listTools({ name: "jira_search", search: "jira" } as never)) + .rejects.toThrow(/exactly one/); + await expect(session.listTools({} as never)).rejects.toThrow(/exactly one/); +}); + +it("treats optional selectors set to undefined as absent", async () => { + const found = classifyTool({ name: "jira_search", annotations: { readOnlyHint: true } }, "byo"); + const host = { + serverName: "Jira", + endpoint: "https://mcp.example.com", + scope: {}, + searchTools: async () => [found], + findTool: async () => found, + } as unknown as McpSessionHost; + const session = new McpSessionBase(host, { authorizeObservation() {} } as never); + + await expect(session.listTools({ search: "jira", name: undefined })) + .resolves.toHaveLength(1); + await expect(session.listTools({ name: "jira_search", search: undefined })) + .resolves.toHaveLength(1); +}); + +it("refuses oversized tool names before consulting the host", async () => { + let finds = 0; + const host = { + serverName: "Jira", + endpoint: "https://mcp.example.com", + scope: {}, + findTool: async () => { finds++; return undefined; }, + } as unknown as McpSessionHost; + const session = new McpSessionBase(host, { authorizeObservation() {} } as never); + const oversized = "x".repeat(MAX_TOOL_NAME_CHARS + 1); + + await expect(session.listTools({ name: oversized })).rejects.toThrow(/tool name.*at most/i); + await expect(session.callTool(oversized)).rejects.toThrow(/tool name.*at most/i); + expect(finds).toBe(0); +}); diff --git a/packages/mcp-shared/__tests__/tool-search.test.ts b/packages/mcp-shared/__tests__/tool-search.test.ts new file mode 100644 index 000000000..ac1f6aeb8 --- /dev/null +++ b/packages/mcp-shared/__tests__/tool-search.test.ts @@ -0,0 +1,52 @@ +import { describe, expect, it } from "vitest"; + +import { matchesToolQuery, toolQueryTerms } from "../src/tool-search.js"; + +const match = ( + tool: { name: string; title?: string; description?: string }, + query: string, +) => matchesToolQuery(tool, toolQueryTerms(query)); + +describe("tool query matching", () => { + it("treats word separators alike on both sides of the comparison", () => { + // The obvious query has to find the obvious tool. MCP tool names are snake_case and a portal + // prefixes each with `{server_id}_`, so a matcher that took `_` literally would answer nothing + // for every one of these spellings. + const tool = { name: "github_list_issues" }; + expect(match(tool, "list issues")).toBe(true); + expect(match(tool, "list_issues")).toBe(true); + expect(match(tool, "list-issues")).toBe(true); + expect(match(tool, "listIssues")).toBe(true); + expect(match(tool, " list issues ")).toBe(true); + }); + + it("requires every term, in any order and any case", () => { + const tool = { name: "gh_create_issue", description: "Open a new issue on a repository" }; + expect(match(tool, "issue repository")).toBe(true); + expect(match(tool, "REPOSITORY Issue")).toBe(true); + expect(match(tool, "issue milestone")).toBe(false); + }); + + it("matches on title and description, not only the name", () => { + const tool = { name: "gh_x1", title: "Create issue", description: "Files a bug report" }; + expect(match(tool, "create")).toBe(true); + expect(match(tool, "bug report")).toBe(true); + }); + + it("matches nothing outside the bounded prefix of a field", () => { + // The bound is what keeps matching cheap against text that arrives straight from the endpoint, + // where a description has not been clamped yet and is whatever the server chose to send. + const buried = { name: "gh_x2", description: `${"padding ".repeat(1000)} needle` }; + expect(match(buried, "needle")).toBe(false); + expect(match(buried, "padding")).toBe(true); + }); + + it("does not turn a separator-only query into match-everything", () => { + expect(toolQueryTerms(" _ - ")).toEqual([]); + expect(match({ name: "anything" }, " _ - ")).toBe(false); + }); + + it("deduplicates normalized terms", () => { + expect(toolQueryTerms("issue ISSUE issue_issue")).toEqual(["issue"]); + }); +}); diff --git a/packages/mcp-shared/package.json b/packages/mcp-shared/package.json index ebfb272f8..26ef8d9cd 100644 --- a/packages/mcp-shared/package.json +++ b/packages/mcp-shared/package.json @@ -25,6 +25,7 @@ "./session": "./src/session.ts", "./session-methods": "./src/session-methods.ts", "./sharing-policy": "./src/sharing-policy.ts", + "./tool-search": "./src/tool-search.ts", "./tools": "./src/tools.ts", "./types": "./src/types.d.ts", "./user": "./src/user.ts", diff --git a/packages/mcp-shared/src/base-types.ts b/packages/mcp-shared/src/base-types.ts index 02552cc5c..7bf47cc42 100644 --- a/packages/mcp-shared/src/base-types.ts +++ b/packages/mcp-shared/src/base-types.ts @@ -11,8 +11,8 @@ export const MCP_BASE_TYPES = `// Base types for MCP-server sessions. // // These are prepended to every generated per-server \`.d.ts\` (see \`schema-to-ts.ts\`), so a workspace's -// coding agent always has them in scope. One method per tool, plus \`callTool\` overloads, is -// generated from the server's own tool catalog and appended below this file's contents. +// coding agent always has them in scope. One method per tool and \`callTool\` overloads are generated +// from the server's own tool catalog and appended below this file's contents. /** A block of content returned by an MCP tool. */ export type McpContent = @@ -81,7 +81,15 @@ export type McpToolInfo = { * action. Recorded so an audit can find every call that was trusted on the server's word. */ classifiedBy: "server-annotation" | "default"; - /** JSON Schema for the tool's arguments, exactly as the server published it. */ + /** JSON Schema for the tool's arguments when it fits the connector's definition budget. */ inputSchema?: unknown; }; + +/** Bounded search result. Request the exact name through \`listTools({ name })\` for its schema. */ +export type McpToolSummary = Omit; + +/** Progressive catalog lookup through the existing \`listTools\` session method. */ +export type McpToolListOptions = + | { search: string; name?: never } + | { name: string; search?: never }; `; diff --git a/packages/mcp-shared/src/catalog.ts b/packages/mcp-shared/src/catalog.ts index 52f19a8ad..74be7c1a3 100644 --- a/packages/mcp-shared/src/catalog.ts +++ b/packages/mcp-shared/src/catalog.ts @@ -1,17 +1,23 @@ -// The tool catalog of one binding: fetched, cached, scoped, and classified. Where a grant's -// `ToolScope` becomes the set of tools a Gadget can name at all. +// The described tool catalog of one binding: fetched, cached, scoped, and classified. A broad grant +// can cover definitions omitted by this bounded catalog; `HydratedTools` remembers those when the +// facet discovers them directly from the endpoint. // // The cache is per gatekeeper facet, so per binding rather than per account. That wastes // `tools/list` calls, but a shared cache would be a channel between two otherwise unrelated // bindings, and a scoped grant is meant not to see what a wider one fetched. -import type { McpTool } from "./client.js"; +import { isValidToolName, type McpTool } from "./client.js"; import { fetchTools, type ConnectionAccount, type ConnectionEnv } from "./connection.js"; -import { looksLikePortal } from "./portal.js"; +import { looksLikePortal, toolBelongsToServer } from "./portal.js"; import { scopeAllows, type ToolScope } from "./scope.js"; import type { McpLog } from "./log.js"; -import { catalogRevision, classifyTool, type ClassifiedTool, type ServerTrust } - from "./tools.js"; +import { + catalogRevision, + classifyTool, + MAX_TOOLS_PER_SERVER, + type ClassifiedTool, + type ServerTrust, +} from "./tools.js"; /** How long a fetched tool catalog is reused before the server is asked again. */ export const CATALOG_TTL_MS = 5 * 60 * 1000; @@ -50,23 +56,42 @@ export type CatalogRequest = { * state, so withdrawing the tier takes effect without a reconnect. See `ServerTrust`. */ trust: ServerTrust; + /** Absolute deadline shared with the facet operation that requested this catalog. */ + deadline?: number; +}; + +/** A binding's currently described tools and whether their endpoint is a portal. */ +export type ScopedCatalog = { + /** Scoped and classified definitions retained in the bounded catalog. */ + tools: ClassifiedTool[]; + /** Whether the endpoint is known or conservatively assumed to be a portal. */ + isPortal: boolean; + /** False when this catalog proves an absent tool name does not exist. */ + truncated: boolean; }; /** - * Returns the tools this binding may call, refreshing from the server when the cache is stale. + * Returns the tools this binding may call and its endpoint kind, refreshing stale cached data. * * A changed catalog is adopted rather than pinned, since refusing to see new tools would break * working Gadgets, but the change is logged and a scoped binding cannot widen: a tool list is a set * of names and a server scope is a name prefix. */ -export async function scopedTools(request: CatalogRequest): Promise { +export async function scopedCatalog(request: CatalogRequest): Promise { const cached = request.store.get("catalog"); let tools = cached?.tools; let truncated = cached?.truncated ?? false; if (!cached || Date.now() - cached.fetchedAt > CATALOG_TTL_MS) { try { - const fetched = await fetchTools(request.env, request.account, request.endpoint); + const serverId = request.scope.serverId; + const fetched = await fetchTools( + request.env, + request.account, + request.endpoint, + serverId ? tool => toolBelongsToServer(tool.name, serverId) : undefined, + { deadline: request.deadline }, + ); const revision = await catalogRevision(fetched.tools); if (cached && cached.revision !== revision) { request.log.info("server tool catalog changed", { @@ -100,8 +125,104 @@ export async function scopedTools(request: CatalogRequest): Promise scopeAllows(request.scope, tool.name, isPortal)) - .map(tool => classifyTool(tool, request.trust)); + // A server-scoped grant can only be minted by the portal connector, and its filtered listing no + // longer contains the `portal_list_servers` that `looksLikePortal` looks for. Deciding from the + // scope first is what keeps the `portal_*` exclusion in force: without it, a `#server=portal` + // grant would filter down to exactly the portal's own tools, find no evidence of a portal, and let + // a Gadget call `portal_toggle_servers` to widen its own reach. + const isPortal = request.scope.serverId !== undefined + || looksLikePortal(all, { truncated, cap: MAX_TOOLS_PER_SERVER }); + return { + isPortal, + truncated, + tools: all + .filter(tool => scopeAllows(request.scope, tool.name, isPortal)) + .map(tool => classifyTool(tool, request.trust)), + }; +} + +/** + * Memory-only cache for definitions fetched beyond a binding's bounded catalog. + * + * Broad grants may cover tools past the catalog cut, so definitions and absences are remembered for + * one TTL to avoid repeating a paginated lookup on every call. + */ +export class HydratedTools { + // Insertion-ordered, which is what makes the eviction below the oldest entry. + #entries = new Map(); + #loads = new Map>(); + #totalBytes = 0; + + /** Maximum definitions and remembered absences retained per facet activation. */ + static readonly MAX_ENTRIES = MAX_TOOLS_PER_SERVER; + /** Aggregate retained definition budget per facet activation. */ + static readonly MAX_BYTES = 1024 * 1024; + + #fresh(name: string): { tool: McpTool | null } | undefined { + const cached = this.#entries.get(name); + if (!cached) return undefined; + if (Date.now() - cached.fetchedAt > CATALOG_TTL_MS) { + this.#entries.delete(name); + this.#totalBytes -= cached.bytes; + return undefined; + } + return cached; + } + + #remember(name: string, tool: McpTool | null): void { + const existing = this.#entries.get(name); + if (existing) { + this.#entries.delete(name); + this.#totalBytes -= existing.bytes; + } + const bytes = new TextEncoder().encode(JSON.stringify(tool)).byteLength; + while (this.#entries.size >= HydratedTools.MAX_ENTRIES + || this.#totalBytes + bytes > HydratedTools.MAX_BYTES) { + const oldest = this.#entries.keys().next(); + if (oldest.done) break; + const removed = this.#entries.get(oldest.value); + this.#entries.delete(oldest.value); + this.#totalBytes -= removed?.bytes ?? 0; + } + if (bytes > HydratedTools.MAX_BYTES) return; + this.#entries.set(name, { fetchedAt: Date.now(), tool, bytes }); + this.#totalBytes += bytes; + } + + #load( + name: string, + load: (name: string) => Promise, + ): Promise { + const existing = this.#loads.get(name); + if (existing) return existing; + if (this.#loads.size >= HydratedTools.MAX_ENTRIES) { + throw new Error("Too many MCP tool definitions are already loading."); + } + + const pending = load(name).then(loaded => { + this.#remember(name, loaded ?? null); + return loaded; + }).finally(() => { + if (this.#loads.get(name) === pending) this.#loads.delete(name); + }); + this.#loads.set(name, pending); + return pending; + } + + /** + * Returns the named tool, coalescing concurrent loads and caching a result for one TTL. + * + * `load` is expected to enforce the caller's scope: this cache is keyed by name alone and holds + * whatever it is given, so it can neither widen nor narrow a grant. A failed load is not cached, + * allowing a later call to recover from a transient endpoint failure. + */ + async resolve( + name: string, + load: (name: string) => Promise, + ): Promise { + if (!isValidToolName(name)) return undefined; + const cached = this.#fresh(name); + if (cached) return cached.tool ?? undefined; + return this.#load(name, load); + } } diff --git a/packages/mcp-shared/src/client.ts b/packages/mcp-shared/src/client.ts index 3995972d1..b3bfad3d9 100644 --- a/packages/mcp-shared/src/client.ts +++ b/packages/mcp-shared/src/client.ts @@ -14,8 +14,8 @@ // 401/403/404 classification below are the SSRF and response-size boundary for this connector. import { - FetchNotStartedError, guardedFetch, + FetchNotStartedError, MAX_RESPONSE_BYTES, readTextCapped, type FetchOptions, @@ -45,7 +45,29 @@ export type ToolCatalog = { truncated: boolean; }; -// Only to make a non-terminating cursor terminate; generous next to MAX_TOOLS_PER_SERVER. +/** One tool identity retained to validate an exact name against a large endpoint. */ +export type IndexedTool = { + /** Exact wire name advertised by the endpoint. */ + name: string; +}; + +/** A bounded survey of endpoint tool identities without schemas or descriptions. */ +export type ToolIndex = { + /** Retained tool identities. */ + tools: IndexedTool[]; + /** Whether count, byte, page, or scan limits cut the survey short. */ + truncated: boolean; +}; + +/** Longest tool name retained from a server or accepted from a Gadget. */ +export const MAX_TOOL_NAME_CHARS = 512; + +/** Whether a value is a non-empty tool name small enough to retain or look up. */ +export function isValidToolName(value: unknown): value is string { + return typeof value === "string" && value.length > 0 && value.length <= MAX_TOOL_NAME_CHARS; +} + +// Independent bound for a non-terminating cursor that returns only tiny or empty pages. const MAX_TOOL_PAGES = 50; // Caps on the size of a tool catalog, as opposed to its length. `maxTools` alone bounds nothing: @@ -60,8 +82,16 @@ const MAX_TOOL_PAGES = 50; // Oversized text is truncated rather than the tool dropped, since the name is what a grant is made // of and a callable tool with a clipped description beats one that vanished. const MAX_TOOL_DESCRIPTION_CHARS = 4000; +// Search is a discovery hint; full schemas come from `listTools({ name })`. Keeping each text field small +// lets ordinary twenty-result searches stay well below the catalog byte budget. +const MAX_TOOL_SEARCH_SUMMARY_CHARS = 256; const MAX_TOOL_SCHEMA_CHARS = 20_000; const MAX_CATALOG_BYTES = 96 * 1024; +// Filtered discovery can skip every tool on a page, so its retained-result budget would otherwise +// permit fifty maximum-sized responses. Bound what was inspected independently of what matched. +const MAX_SCANNED_TOOL_BYTES = 4 * 1024 * 1024; +const MAX_SCANNED_TOOLS = 5000; +const encoder = new TextEncoder(); /** Content block returned by a tool call. */ export type McpContentBlock = ContentBlock; @@ -76,6 +106,9 @@ export type McpToolAnnotations = ToolAnnotations; /** One entry of the server's `tools/list` response. */ export type McpWireTool = Tool; +/** Selects which wire tools count toward a bounded catalog. */ +export type McpToolFilter = (tool: McpWireTool) => boolean; + export type McpTool = { name: string; title?: string; @@ -199,7 +232,7 @@ type JsonRpcResponse = { error?: { code: number; message: string; data?: unknown }; }; -function extractResponse(bodyText: string, id: number): JsonRpcResponse { +function extractResponse(bodyText: string, id: number | string): JsonRpcResponse { let parsed: JsonRpcResponse; try { parsed = JSON.parse(bodyText) as JsonRpcResponse; @@ -214,7 +247,10 @@ function extractResponse(bodyText: string, id: number): JsonRpcResponse { // Reads completed SSE events until this request's response arrives. Streamable HTTP servers may // keep the stream open after responding, so waiting for EOF would hang an otherwise-complete call. -async function readSseResponse(response: Response, id: number): Promise { +async function readSseResponse( + response: Response, + id: number | string, +): Promise<{ parsed: JsonRpcResponse; bytes: number }> { if (!response.body) { throw new McpProtocolError("MCP server's event stream contained no response to the request."); } @@ -249,7 +285,7 @@ async function readSseResponse(response: Response, id: number): Promise undefined); - return parsed; + return { parsed, bytes: total }; } } } finally { @@ -281,13 +317,31 @@ function clampText(value: unknown, max: number): string | undefined { // Trims one tool down to what is worth keeping, before it reaches storage or the agent. A schema too // large to render is dropped rather than clipped, so the generated method degrades to // `Record`. -function clampTool(tool: McpWireTool): McpTool { +// Keeps only the hints this gatekeeper understands, so a server cannot attach unbounded text to a +// tool under `annotations` and have it stored, indexed, or rendered. Each retained field is a +// boolean or absent, which is what makes an index entry's size predictable. +function clampAnnotations( + annotations: McpToolAnnotations | undefined, +): McpToolAnnotations | undefined { + return annotations && { + readOnlyHint: typeof annotations.readOnlyHint === "boolean" + ? annotations.readOnlyHint : undefined, + destructiveHint: typeof annotations.destructiveHint === "boolean" + ? annotations.destructiveHint : undefined, + idempotentHint: typeof annotations.idempotentHint === "boolean" + ? annotations.idempotentHint : undefined, + openWorldHint: typeof annotations.openWorldHint === "boolean" + ? annotations.openWorldHint : undefined, + }; +} + +/** Reduces one untrusted wire tool to the bounded fields this gatekeeper understands. */ +export function clampToolDefinition(tool: McpWireTool | McpTool): McpTool { const schema = tool.inputSchema && typeof tool.inputSchema === "object" ? tool.inputSchema as JsonSchema : undefined; const oversized = schema !== undefined && JSON.stringify(schema).length > MAX_TOOL_SCHEMA_CHARS; - const annotations = tool.annotations; return { // Pick known fields rather than spreading an untrusted JSON object. Unknown extensions are not // used anywhere, and retaining one would let it bypass every per-field cap before caching. @@ -295,12 +349,22 @@ function clampTool(tool: McpWireTool): McpTool { title: clampText(tool.title, MAX_TOOL_DESCRIPTION_CHARS), description: clampText(tool.description, MAX_TOOL_DESCRIPTION_CHARS), inputSchema: oversized ? undefined : schema, - annotations: annotations && { - readOnlyHint: annotations.readOnlyHint, - destructiveHint: annotations.destructiveHint, - idempotentHint: annotations.idempotentHint, - openWorldHint: annotations.openWorldHint, - }, + annotations: clampAnnotations(tool.annotations), + }; +} + +// Reduces one tool to an index entry. Bounded by its already-validated name. +function indexTool(tool: McpWireTool): IndexedTool { + return { name: tool.name }; +} + +/** Reduces one tool to the bounded, schema-free form returned by search. */ +export function clampToolSummary(tool: McpWireTool | McpTool): McpTool { + return { + name: tool.name, + title: clampText(tool.title, MAX_TOOL_SEARCH_SUMMARY_CHARS), + description: clampText(tool.description, MAX_TOOL_SEARCH_SUMMARY_CHARS), + annotations: clampAnnotations(tool.annotations), }; } @@ -315,6 +379,7 @@ export class McpClient { #endpoint: string; #getAuthorization: AuthorizationProvider; #fetchOptions: FetchOptions; + #requestPrefix = crypto.randomUUID(); #requestId = 0; /** Transport session id, assigned by the server during `initialize`. Persist and pass it back. */ @@ -329,7 +394,9 @@ export class McpClient { this.#endpoint = endpoint; this.#getAuthorization = getAuthorization; this.sessionId = sessionId ?? null; - this.#fetchOptions = fetchOptions; + this.#fetchOptions = fetchOptions.timeoutMs !== undefined && fetchOptions.deadline === undefined + ? { ...fetchOptions, deadline: Date.now() + fetchOptions.timeoutMs } + : fetchOptions; } // The credential most recently sent, kept only so it can be recognised if it comes back. See @@ -396,8 +463,14 @@ export class McpClient { return response; } - async #call(method: string, params?: unknown): Promise { - const id = ++this.#requestId; + async #callMeasured( + method: string, + params?: unknown, + ): Promise<{ result: T; responseBytes: number }> { + // A transport session is persisted on the account and can be used by several short-lived client + // instances concurrently. Prefixing IDs per instance prevents two active requests from both + // being JSON-RPC id 1 and confusing the server's SSE response routing. + const id = `${this.#requestPrefix}:${++this.#requestId}`; const response = await this.#post({ jsonrpc: "2.0", id, method, params }); if (!response.ok) { @@ -411,8 +484,11 @@ export class McpClient { const contentType = (response.headers.get("Content-Type") ?? "").toLowerCase(); let parsed: JsonRpcResponse; + let responseBytes: number; if (contentType.includes("text/event-stream")) { - parsed = await readSseResponse(response, id); + const measured = await readSseResponse(response, id); + parsed = measured.parsed; + responseBytes = measured.bytes; } else { // Capped: a JSON tool result is server-controlled and unbounded, and has to be buffered whole // before it can be parsed. The catalog limits above bound what is kept, not what arrives. @@ -424,6 +500,7 @@ export class McpClient { `MCP server's response to "${method}" was too large to read: ` + `${err instanceof Error ? err.message : String(err)}`); } + responseBytes = encoder.encode(bodyText).byteLength; parsed = extractResponse(bodyText, id); } @@ -432,7 +509,11 @@ export class McpClient { `MCP server rejected "${method}": ${this.#quoteServerText(parsed.error.message)}`, parsed.error.code, method === "tools/call" ? "unknown" : "declined"); } - return parsed.result as T; + return { result: parsed.result as T, responseBytes }; + } + + async #call(method: string, params?: unknown): Promise { + return (await this.#callMeasured(method, params)).result; } // Prepares text the server wrote for inclusion in an error message. @@ -480,37 +561,97 @@ export class McpClient { } /** - * Lists every tool the server offers, following `nextCursor` pagination to exhaustion. Bounded by - * pages, tool count, and bytes: a server answering with empty pages and a fresh cursor each time - * would otherwise loop until the Worker's limits killed it. + * Lists tools the server offers, following pagination to exhaustion. + * + * `include`, when present, is applied before count and byte budgets so an aggregator's unrelated + * tools cannot crowd the requested server or exact grant names out of the bounded result. */ - async listTools(maxTools: number): Promise { - const tools: McpTool[] = []; + async listTools(maxTools: number, include?: McpToolFilter): Promise { + return this.#list(maxTools, include, clampToolDefinition); + } + + /** Collects at most `maxTools` matching index entries without scanning later pages. */ + async listMatchingToolIndex(maxTools: number, include: McpToolFilter): Promise { + return this.#list(maxTools, include, indexTool, true); + } + + /** Finds one exact tool without reading pages after the match. */ + async findTool(name: string): Promise { + if (!isValidToolName(name)) return undefined; + return (await this.#list( + 1, tool => tool.name === name, clampToolDefinition, true, true)).tools[0]; + } + + /** Collects at most `maxTools` bounded matching summaries without scanning later pages. */ + async listMatchingToolSummaries(maxTools: number, include: McpToolFilter): Promise { + return (await this.#list(maxTools, include, clampToolSummary, true, true)).tools; + } + + // The shared listing loop. `project` decides how much of each tool is retained, and therefore how + // much of the byte budget each one costs; the budget itself is applied to whatever it returns. + async #list( + maxTools: number, + include: McpToolFilter | undefined, + project: (tool: McpWireTool) => T, + stopWhenFull = false, + failOnScanLimit = false, + ): Promise<{ tools: T[]; truncated: boolean }> { + const tools: T[] = []; let budget = MAX_CATALOG_BYTES; + let scannedBytes = 0; + let scannedTools = 0; let cursor: string | undefined; + const scanLimit = (): { tools: T[]; truncated: boolean } => { + if (failOnScanLimit) { + throw new McpProtocolError( + "MCP tool discovery exceeded its scan budget.", undefined, "declined"); + } + return { tools, truncated: true }; + }; for (let page = 0; page < MAX_TOOL_PAGES; page++) { - const body = await this.#call<{ tools?: McpWireTool[]; nextCursor?: string }>( - "tools/list", cursor ? { cursor } : {}); - for (const tool of body.tools ?? []) { + const measured = await this.#callMeasured<{ tools?: McpWireTool[]; nextCursor?: string }>( + "tools/list", cursor === undefined ? {} : { cursor }); + const body = measured.result; + scannedBytes += measured.responseBytes; + if (scannedBytes > MAX_SCANNED_TOOL_BYTES) return scanLimit(); + const pageTools = body.tools ?? []; + const remainingTools = Math.max(0, MAX_SCANNED_TOOLS - scannedTools); + const scanCount = Math.min(pageTools.length, remainingTools); + scannedTools += scanCount; + for (let index = 0; index < scanCount; index++) { + const tool = pageTools[index]; + if (!isValidToolName(tool?.name)) continue; + if (include && !include(tool)) continue; // A cap was reached with tools still arriving, so the catalog is known to be incomplete. // Reported rather than inferred from `tools.length`, since the byte budget can stop the // listing well short of `maxTools` and leaves no trace in the array itself. if (tools.length >= maxTools || budget <= 0) return { tools, truncated: true }; - if (typeof tool?.name !== "string" || tool.name.length === 0) continue; - const trimmed = clampTool(tool); + const trimmed = project(tool); // Include a comma's byte for every array member. Brackets are covered by the 32 KiB storage // headroom. Refuse the whole next tool rather than storing half a schema or an invalid value. - const bytes = new TextEncoder().encode(JSON.stringify(trimmed)).byteLength + 1; + const bytes = encoder.encode(JSON.stringify(trimmed)).byteLength + 1; if (bytes > budget) return { tools, truncated: true }; budget -= bytes; tools.push(trimmed); + if (stopWhenFull && tools.length >= maxTools) { + // Discovery only asks for a bounded prefix and does not need to prove whether another + // match exists. The public discovery methods discard `truncated`; ordinary listings retain + // the precise semantics above and still paginate to completion. + return { tools, truncated: true }; + } + } + // Inspect the bounded prefix before reporting the scan limit: an exact requested tool may be + // the first entry in a page whose remaining entries cross the aggregate tool budget. + if (scanCount < pageTools.length) return scanLimit(); + const nextCursor = body.nextCursor; + if (nextCursor === undefined) return { tools, truncated: false }; + if (typeof nextCursor !== "string") { + throw new McpProtocolError('MCP server returned an invalid "tools/list" cursor.'); } - cursor = body.nextCursor; - if (!cursor) return { tools, truncated: false }; + cursor = nextCursor; } - throw new McpProtocolError( - `MCP server kept paginating "tools/list" past ${MAX_TOOL_PAGES} pages.`); + return scanLimit(); } /** Invokes one tool. A tool-level failure arrives as `isError`, not as a thrown error. */ diff --git a/packages/mcp-shared/src/connection.ts b/packages/mcp-shared/src/connection.ts index 4ba993833..da3c59fd8 100644 --- a/packages/mcp-shared/src/connection.ts +++ b/packages/mcp-shared/src/connection.ts @@ -12,6 +12,7 @@ import { McpCallNotDispatchedError, McpClient, McpSessionExpiredError, + type McpToolFilter, type ToolCatalog, } from "./client.js"; @@ -27,6 +28,8 @@ export type ConnectionEnv = InsecureEnv & { export type WithClientOptions = { /** False for a call that may have taken effect, so a dropped session is not retried. See above. */ retryOnExpiry?: boolean; + /** Absolute deadline shared with time spent waiting for a discovery slot. */ + deadline?: number; }; /** @@ -103,7 +106,10 @@ export async function withClient( await account.assertConnectionCurrent(endpoint, generation); } return authorization; - }, sessionId, fetchOptions(env)); + }, sessionId, { + ...fetchOptions(env), + deadline: options.deadline, + }); let persistedSessionId = sessionId; const persistSession = async (): Promise => { @@ -186,8 +192,12 @@ export async function withClient( * listing short, since callers infer what the endpoint is from what it does and does not offer. */ export async function fetchTools( - env: ConnectionEnv, account: ConnectionAccount, endpoint: string, + env: ConnectionEnv, + account: ConnectionAccount, + endpoint: string, + include?: McpToolFilter, + options?: WithClientOptions, ): Promise { return withClient(env, account, endpoint, - client => client.listTools(MAX_TOOLS_PER_SERVER)); + client => client.listTools(MAX_TOOLS_PER_SERVER, include), options); } diff --git a/packages/mcp-shared/src/facet.ts b/packages/mcp-shared/src/facet.ts index 0dcf124b3..009594dbe 100644 --- a/packages/mcp-shared/src/facet.ts +++ b/packages/mcp-shared/src/facet.ts @@ -11,7 +11,12 @@ import type { } from "@gadgets/workshop-shared/gatekeeper"; import { ActionStore, REVERT_UNSUPPORTED_MESSAGE } from "./action-store.js"; -import { CATALOG_TTL_MS, scopedTools } from "./catalog.js"; +import { + CATALOG_TTL_MS, + HydratedTools, + scopedCatalog, + type ScopedCatalog, +} from "./catalog.js"; import type { McpClient } from "./client.js"; import { withClient, @@ -20,11 +25,18 @@ import { type WithClientOptions, } from "./connection.js"; import type { McpLog } from "./log.js"; -import { formatToolScope, type ToolScope } from "./scope.js"; +import { DEFAULT_REQUEST_TIMEOUT_MS } from "./fetch.js"; +import { formatToolScope, scopeAllows, type ToolScope } from "./scope.js"; +import { matchesToolQuery, toolQueryTerms, MAX_SEARCH_RESULTS } from "./tool-search.js"; import { McpSessionBase, type McpSessionHost, type StoredAction } from "./session.js"; import { installToolMethods } from "./session-methods.js"; import { observerRefusalMessage } from "./sharing-policy.js"; -import { actionKindFor, type ClassifiedTool, type ServerTrust } from "./tools.js"; +import { + actionKindFor, + classifyTool, + type ClassifiedTool, + type ServerTrust, +} from "./tools.js"; type FacetProps = { endpoint: string; @@ -36,16 +48,22 @@ type SessionConstructor = new ( queue: RpcStub, ) => Session; +const MAX_CONCURRENT_DISCOVERIES = 4; +const MAX_QUEUED_DISCOVERIES = 32; + /** Common session, catalog, action, and sharing behavior for connector-owned MCP facets. */ export abstract class McpFacetBase< Env extends ConnectionEnv, Props extends FacetProps, Session extends McpSessionBase, > extends DurableObject implements Gatekeeper, McpSessionHost { - #toolsPromise: Promise | undefined; + #catalogPromise: Promise | undefined; #toolsFetchedAt = 0; #toolsTrust: ServerTrust | undefined; #actionStore: ActionStore | undefined; + #hydrated = new HydratedTools(); + #activeDiscoveries = 0; + #waitingDiscoveries: Array<() => void> = []; #actions(): ActionStore { return this.#actionStore ??= new ActionStore(this.ctx.storage.sql); @@ -93,14 +111,14 @@ export abstract class McpFacetBase< return formatToolScope(this.endpoint, this.scope); } - /** Returns this facet's scoped and classified tool catalog. */ - tools(): Promise { + /** Returns this facet's scoped catalog and endpoint kind. */ + protected catalog(deadline?: number): Promise { const trust = this.trust; - if (!this.#toolsPromise || this.#toolsTrust !== trust + if (!this.#catalogPromise || this.#toolsTrust !== trust || Date.now() - this.#toolsFetchedAt > CATALOG_TTL_MS) { this.#toolsFetchedAt = Date.now(); this.#toolsTrust = trust; - this.#toolsPromise = scopedTools({ + const load = (operationDeadline: number) => scopedCatalog({ store: this.ctx.storage.kv, log: this.log, env: this.env, @@ -108,12 +126,93 @@ export abstract class McpFacetBase< endpoint: this.endpoint, scope: this.scope, trust, - }).catch(err => { - this.#toolsPromise = undefined; + deadline: operationDeadline, + }); + const loading = deadline === undefined ? this.runDiscovery(load) : load(deadline); + this.#catalogPromise = loading.catch(err => { + this.#catalogPromise = undefined; throw err; }); } - return this.#toolsPromise; + return this.#catalogPromise; + } + + /** Returns this facet's scoped and classified tool definitions. */ + async tools(): Promise { + return (await this.catalog()).tools; + } + + /** Runs Gadget-triggered catalog I/O within one facet-wide concurrency bound. */ + protected async runDiscovery(operation: (deadline: number) => Promise): Promise { + const deadline = Date.now() + DEFAULT_REQUEST_TIMEOUT_MS; + if (this.#activeDiscoveries >= MAX_CONCURRENT_DISCOVERIES) { + if (this.#waitingDiscoveries.length >= MAX_QUEUED_DISCOVERIES) { + throw new Error("Too many MCP discovery requests are already in progress."); + } + await new Promise((resolve, reject) => { + const resume = () => { + clearTimeout(timer); + resolve(); + }; + const timer = setTimeout(() => { + const index = this.#waitingDiscoveries.indexOf(resume); + if (index >= 0) this.#waitingDiscoveries.splice(index, 1); + reject(new Error("Timed out waiting to discover MCP tools.")); + }, Math.max(0, deadline - Date.now())); + this.#waitingDiscoveries.push(resume); + }); + } else { + this.#activeDiscoveries++; + } + + try { + return await operation(deadline); + } finally { + const next = this.#waitingDiscoveries.shift(); + if (next) next(); + else this.#activeDiscoveries--; + } + } + + /** Searches the endpoint for granted tools by name, title, and description. */ + async searchTools(query: string): Promise { + const terms = toolQueryTerms(query); + return this.runDiscovery(async deadline => { + const catalog = await this.catalog(deadline); + if (!catalog.truncated) { + return catalog.tools.filter(entry => matchesToolQuery(entry.tool, terms)) + .slice(0, MAX_SEARCH_RESULTS); + } + const { isPortal } = catalog; + const tools = await this.call( + client => client.listMatchingToolSummaries( + MAX_SEARCH_RESULTS, + tool => scopeAllows(this.scope, tool.name, isPortal) && matchesToolQuery(tool, terms), + ), + { deadline }, + ); + return tools.map(tool => classifyTool(tool, this.trust)); + }); + } + + /** Resolves one granted tool, fetching it when the described catalog omitted it. */ + async findTool(name: string): Promise { + // Grant restrictions can be enforced without loading anything. A portal-native exclusion needs + // the endpoint kind below, except for a server scope, which by definition belongs to a portal. + if (!scopeAllows(this.scope, name, this.scope.serverId !== undefined)) return undefined; + + return this.runDiscovery(async deadline => { + const catalog = await this.catalog(deadline); + if (!scopeAllows(this.scope, name, catalog.isPortal)) return undefined; + const described = catalog.tools.find(entry => entry.tool.name === name); + if (described) return described; + if (!catalog.truncated) return undefined; + + const load = (candidate: string) => + this.call(client => client.findTool(candidate), { deadline }); + const tool = await this.#hydrated.resolve(name, load); + return tool && classifyTool(tool, this.trust); + }); } /** Returns action kinds that this facet's current catalog permits auto-approving. */ @@ -176,7 +275,10 @@ export abstract class McpFacetBase< } /** Runs a call against this facet's endpoint and account. */ - call(fn: (client: McpClient) => Promise, options?: WithClientOptions): Promise { + call( + fn: (client: McpClient) => Promise, + options?: WithClientOptions, + ): Promise { return withClient(this.env, this.account(), this.endpoint, fn, options); } diff --git a/packages/mcp-shared/src/portal.ts b/packages/mcp-shared/src/portal.ts index 31a1c45db..29c78bc02 100644 --- a/packages/mcp-shared/src/portal.ts +++ b/packages/mcp-shared/src/portal.ts @@ -8,7 +8,6 @@ // portal. See the MCP Server Portals connector's README. import type { McpTool } from "./client.js"; -import { MAX_TOOLS_PER_SERVER } from "./tools.js"; /** * The portal's built-in server-listing tool. Its presence in `tools/list` is what identifies an @@ -46,16 +45,20 @@ export function isPortalNativeTool(name: string): boolean { /** * True when this tool list came from a portal. * - * A truncated catalog counts as a portal regardless of what is in it: `tools/list` is unordered, so + * A truncated listing counts as a portal regardless of what is in it: `tools/list` is unordered, so * answering "not a portal" because the evidence fell past the cut would fail open on the `portal_*` * exclusion above -- a real portal would be granted at its bare endpoint, and a Gadget holding that - * grant could call `portal_toggle_servers` to widen its own reach. `truncated` covers the byte - * budget as well as the tool count, which stops the listing without leaving a short array behind. + * grant could call `portal_toggle_servers` to widen its own reach. + * + * The explicit bounds form avoids guessing. `truncated` covers the byte budget, while `cap` is the + * tool count the caller requested. Callers use different caps for ordinary catalogs and wide portal + * indexes, so each must supply the bound it actually used. */ export function looksLikePortal( - tools: Pick[], truncated = false, + tools: readonly Pick[], + bounds: { truncated: boolean; cap: number }, ): boolean { - if (truncated || tools.length >= MAX_TOOLS_PER_SERVER) return true; + if (bounds.truncated || tools.length >= bounds.cap) return true; return tools.some(tool => tool.name === PORTAL_LIST_SERVERS_TOOL); } diff --git a/packages/mcp-shared/src/schema-to-ts.ts b/packages/mcp-shared/src/schema-to-ts.ts index 9182e5925..9bad02eea 100644 --- a/packages/mcp-shared/src/schema-to-ts.ts +++ b/packages/mcp-shared/src/schema-to-ts.ts @@ -8,11 +8,12 @@ // `toMethodName` rather than each spelling the rule out: a type advertising a method that does not // exist is worse than no type at all. // -// `callTool` is still generated as an overload set. It is the escape hatch for tools whose names -// cannot become methods, and the stable way to call a tool whose name a server later changes. +// `callTool` generates precise overloads for described tools plus a generic overload for names +// discovered later through `listTools`. import type { JsonSchema } from "./client.js"; import { toMethodName, toolMethodNames } from "./session-methods.js"; +import { MAX_SEARCH_RESULTS } from "./tool-search.js"; import type { ClassifiedTool, ServerTrust } from "./tools.js"; // Depth limit for recursive/self-referential schemas; deeper nodes degrade to `unknown`. @@ -339,8 +340,9 @@ export function generateSessionTypes(args: { lines.push("/**"); lines.push(` * Session for the "${serverName}" MCP server.`); lines.push(" *"); - lines.push(` * ${readTools.length} tool(s) are read-only and return results immediately, recorded`); - lines.push(" * as observations. The remaining " + actionTools.length + " tool(s) are treated as actions:"); + lines.push(` * Of the ${args.tools.length} currently described tool(s), ${readTools.length} are read-only`); + lines.push(" * and return results immediately as observations. The remaining " + actionTools.length); + lines.push(" * described tool(s) are treated as actions:"); lines.push(" * `callTool` queues them for approval and returns `{ status: \"pending\" }`; the result"); lines.push(" * becomes available through `getActionResult` once a human approves."); lines.push(" * When using this session from `executeCode`, return from that executeCode call as soon as"); @@ -356,8 +358,13 @@ export function generateSessionTypes(args: { lines.push(" * publish it as a blueprint so they connect their own account."); lines.push(" */"); lines.push(`export interface ${typeName} {`); - lines.push(" /** Lists the tools this session exposes, including their read/action classification. */"); + lines.push(" /** Lists currently described tools. */"); lines.push(" listTools(): Promise;"); + lines.push(` /** Searches for up to ${MAX_SEARCH_RESULTS} matching tool summaries. */`); + lines.push(" listTools(options: { search: string; name?: never }): Promise;"); + lines.push(" /** Returns zero or one exact granted tool definition by wire name. */"); + lines.push(" listTools(options: { name: string; search?: never }): Promise;"); + lines.push(" listTools(options: McpToolListOptions): Promise;"); lines.push(""); // One named method per tool, which is how a Gadget is expected to call them. @@ -382,8 +389,8 @@ export function generateSessionTypes(args: { lines.push(" /**"); lines.push(" * Calls a tool by its exact name, as the server publishes it."); lines.push(" *"); - lines.push(" * Equivalent to the named methods above, and the only way to reach a tool that has"); - lines.push(" * none. Prefer it when a tool name must survive the server renaming its tools."); + lines.push(" * Equivalent to the named methods above, with static argument checking for tools that"); + lines.push(" * cannot have a named method."); lines.push(" */"); for (const { tool } of args.tools) { switch (argumentStyle(tool.inputSchema)) { @@ -399,6 +406,14 @@ export function generateSessionTypes(args: { } } } + lines.push(" /** Calls a dynamically discovered tool by exact wire name. */"); + const knownToolNames = args.tools.map(({ tool }) => quote(tool.name)).join(" | ") || "never"; + lines.push(" callTool("); + lines.push(" name: Name,"); + lines.push(` ...args: Name extends ${knownToolNames}`); + lines.push(" ? [args: never]"); + lines.push(" : [args?: Record]"); + lines.push(" ): Promise;"); lines.push(""); lines.push(" /**"); diff --git a/packages/mcp-shared/src/scope.ts b/packages/mcp-shared/src/scope.ts index 9f508772f..d581e0c0e 100644 --- a/packages/mcp-shared/src/scope.ts +++ b/packages/mcp-shared/src/scope.ts @@ -15,6 +15,7 @@ import { toolBelongsToServer, type PortalServer, } from "./portal.js"; +import { MAX_TOOLS_PER_SERVER } from "./tools.js"; /** The breadth of one binding's grant over its endpoint. */ export type ToolScope = { @@ -90,12 +91,16 @@ export function parseToolScope(resourceUrl: string | URL): ToolScope { } catch { return {}; } + const rawTools = params.getAll("tool"); + if (rawTools.length > MAX_TOOLS_PER_SERVER) { + throw new Error(`A grant can name at most ${MAX_TOOLS_PER_SERVER} MCP tools.`); + } return { serverId: params.has("server") ? params.get("server")!.trim() : undefined, tools: params.has("tools") ? [] : params.has("tool") - ? params.getAll("tool").map(name => name.trim()).filter(Boolean) + ? rawTools.map(name => name.trim()).filter(Boolean) : undefined, }; } @@ -106,6 +111,9 @@ export function parseToolScope(resourceUrl: string | URL): ToolScope { * it would undo the fail-closed parse above on the round trip. */ export function formatToolScope(endpoint: string, scope: ToolScope): string { + if ((scope.tools?.length ?? 0) > MAX_TOOLS_PER_SERVER) { + throw new Error(`A grant can name at most ${MAX_TOOLS_PER_SERVER} MCP tools.`); + } const params = new URLSearchParams(); if (scope.serverId !== undefined) params.set("server", scope.serverId); for (const tool of scope.tools ?? []) params.append("tool", tool); @@ -144,10 +152,6 @@ export function validateToolScopeAgainstCatalog( catalog: ToolCatalog, reportedServers: PortalServer[] = [], ): PortalServer | undefined { - if (catalog.truncated && scope.tools !== undefined) { - throw new Error("The current tool catalog is truncated, so this grant cannot be validated."); - } - let server: PortalServer | undefined; const serverId = scope.serverId; if (serverId !== undefined) { @@ -166,6 +170,9 @@ export function validateToolScopeAgainstCatalog( throw new Error(`Tool "${name}" does not belong to portal server "${serverId}".`); } if (!names.has(name)) { + if (catalog.truncated) { + throw new Error("The current tool catalog is truncated, so this grant cannot be validated."); + } throw new Error(`Tool "${name}" is absent from the current tool catalog.`); } } diff --git a/packages/mcp-shared/src/session.ts b/packages/mcp-shared/src/session.ts index a279f8e5f..d40fb99ab 100644 --- a/packages/mcp-shared/src/session.ts +++ b/packages/mcp-shared/src/session.ts @@ -8,11 +8,28 @@ import { RpcTarget, type RpcStub } from "cloudflare:workers"; import type { ActionDescription, ActionKind, ApprovalQueue } from "@gadgets/workshop-shared/gatekeeper"; -import type { McpClient } from "./client.js"; +import { + MAX_TOOL_NAME_CHARS, + type McpClient, +} from "./client.js"; import type { WithClientOptions } from "./connection.js"; import { isWholeEndpoint, type ToolScope } from "./scope.js"; -import { describeCall, toCallResult, toolInfo, type ClassifiedTool } from "./tools.js"; -import type { McpCallResult, McpToolInfo } from "./types"; +import { toolQueryTerms, MAX_QUERY_CHARS, MAX_SEARCH_RESULTS } from "./tool-search.js"; +import { + codeSpan, + describeCall, + plainInline, + toCallResult, + toolInfo, + toolSummary, + type ClassifiedTool, +} from "./tools.js"; +import type { + McpCallResult, + McpToolInfo, + McpToolListOptions, + McpToolSummary, +} from "./types"; /** * A queued tool call, awaiting a decision. Persisted by the host in its own storage; the session @@ -40,6 +57,7 @@ export type StoredAction = { retryable?: boolean; /** Populated once applied; delivered to the Gadget as an observation. */ result?: Extract; + /** Terminal failure reason retained for later collection. */ error?: string; }; @@ -53,10 +71,17 @@ export interface McpSessionHost { /** How much of the endpoint this binding may call. Only used to word the "no such tool" error. */ readonly scope: ToolScope; + /** Returns the bounded described catalog. */ tools(): Promise; - + /** Searches bounded summaries across the granted endpoint scope. */ + searchTools(query: string): Promise; + /** Finds one granted tool definition by exact wire name. */ + findTool(name: string): Promise; /** Runs `fn` against an initialized client for this binding's endpoint. */ - call(fn: (client: McpClient) => Promise, options?: WithClientOptions): Promise; + call( + fn: (client: McpClient) => Promise, + options?: WithClientOptions, + ): Promise; /** The approval-kind tag for one tool, namespaced so pre-approvals cannot cross servers. */ actionKindFor(toolName: string): ActionKind; @@ -66,6 +91,15 @@ export interface McpSessionHost { lookupAction(id: number): StoredAction | undefined; } +function requireToolName(method: string, name: unknown): asserts name is string { + if (typeof name !== "string" || name.length === 0) { + throw new Error(`${method}() requires a tool name.`); + } + if (name.length > MAX_TOOL_NAME_CHARS) { + throw new Error(`${method}() tool name must be at most ${MAX_TOOL_NAME_CHARS} characters.`); + } +} + /** * The Gadget-facing session. A named method per tool is installed on a per-grant subclass (see * `session-methods.ts`), each a one-line delegate to `callTool`. Connectors subclass this and apply @@ -85,37 +119,79 @@ export class McpSessionBase extends RpcTarget { (this.#queue as RpcStub & { [Symbol.dispose](): void })[Symbol.dispose](); } - async listTools(): Promise { - const tools = await this.#host.tools(); + async listTools( + options?: McpToolListOptions, + ): Promise { + if (options === undefined) { + const tools = await this.#host.tools(); + await this.#queue.authorizeObservation({ + title: `${this.#host.serverName}: list tools`, + description: + `Read the tool catalog of the MCP server **${this.#host.serverName}** ` + + `(\`${this.#host.endpoint}\`).`, + }); + return tools.map(toolInfo); + } + if (typeof options !== "object" || options === null) { + throw new Error("listTools() options must select search or name."); + } + const hasName = options.name !== undefined; + const hasSearch = options.search !== undefined; + if (hasName === hasSearch) { + throw new Error("listTools() options must select exactly one of search or name."); + } + if (hasName) { + requireToolName("listTools", options.name); + const tool = await this.#host.findTool(options.name); + await this.#queue.authorizeObservation({ + title: `${this.#host.serverName}: find tool`, + description: + `Looked for ${codeSpan(options.name, MAX_TOOL_NAME_CHARS)} on the MCP server ` + + `**${this.#host.serverName}** (\`${this.#host.endpoint}\`).`, + }); + return tool ? [toolInfo(tool)] : []; + } + const trimmed = typeof options.search === "string" ? options.search.trim() : ""; + if (trimmed.length === 0) { + throw new Error("listTools({ search }) requires a non-empty query."); + } + if (trimmed.length > MAX_QUERY_CHARS) { + throw new Error(`listTools({ search }) query must be at most ${MAX_QUERY_CHARS} characters.`); + } + if (toolQueryTerms(trimmed).length === 0) { + throw new Error("listTools({ search }) requires one or more search terms."); + } + const tools = await this.#host.searchTools(trimmed); await this.#queue.authorizeObservation({ - title: `${this.#host.serverName}: list tools`, + title: `${this.#host.serverName}: search tools`, description: - `Read the tool catalog of the MCP server **${this.#host.serverName}** ` + - `(\`${this.#host.endpoint}\`).`, + `Searched the tool catalog of the MCP server **${this.#host.serverName}** ` + + // The query is the agent's text, so it is flattened before being quoted: an observation is + // read by a person, and a query is as able to forge structure as a tool description is. + `(\`${this.#host.endpoint}\`) for \`${plainInline(trimmed)}\` ` + + `and returned ${tools.length} match(es), up to a limit of ${MAX_SEARCH_RESULTS}.`, }); - return tools.map(toolInfo); + return tools.map(toolSummary); + } + + // Worded from the grant's point of view: on a scoped binding the tool may well exist on the server, + // and "no such tool" would send an agent looking for a typo it will not find. + #noSuchToolMessage(name: string): string { + return isWholeEndpoint(this.#host.scope) + ? `The MCP server "${this.#host.serverName}" has no tool named "${name}".` + : `This binding does not grant a tool named "${name}".`; } async callTool(name: string, args?: Record): Promise { - if (typeof name !== "string" || name.length === 0) { - throw new Error("callTool() requires a tool name."); - } + requireToolName("callTool", name); const toolArgs = args ?? {}; if (typeof toolArgs !== "object" || Array.isArray(toolArgs)) { throw new Error("callTool() arguments must be an object."); } const host = this.#host; - const tools = await host.tools(); - const entry = tools.find(candidate => candidate.tool.name === name); - if (!entry) { - // Worded from the grant's point of view: on a scoped binding the tool may exist on the server, - // and "no such tool" would send an agent looking for a typo. - const available = tools.map(candidate => candidate.tool.name).join(", "); - throw new Error(isWholeEndpoint(host.scope) - ? `The MCP server "${host.serverName}" has no tool named "${name}". Available: ${available}` - : `This binding grants only these tools: ${available}.`); - } + const entry = await host.findTool(name); + if (!entry) throw new Error(this.#noSuchToolMessage(name)); const described = describeCall({ serverName: host.serverName, diff --git a/packages/mcp-shared/src/tool-search.ts b/packages/mcp-shared/src/tool-search.ts new file mode 100644 index 000000000..a988dcc39 --- /dev/null +++ b/packages/mcp-shared/src/tool-search.ts @@ -0,0 +1,51 @@ +// Matching a text query against a tool, shared by every catalog search. +// +// One implementation on purpose. Two connectors search the same way over different sources -- one +// filters definitions already in hand, while endpoint discovery applies it as each page is parsed -- +// and a query answered differently depending on the path would be a difference nobody asked for and +// nobody could see. + +import type { McpTool } from "./client.js"; + +/** Longest accepted agent-supplied search query. */ +export const MAX_QUERY_CHARS = 200; + +/** Most tool summaries returned by one search. */ +export const MAX_SEARCH_RESULTS = 20; + +// Longest run of one server-supplied field considered when matching. +// +// The bound is load-bearing on the path that matches tools as they arrive from the endpoint, before +// per-field clamping: a description there is whatever the server sent, and every term is tested +// against it. Applied unconditionally rather than only on that path, so no caller has to know which +// of its inputs were already clamped in order to be safe. +const MAX_SEARCHABLE_FIELD_CHARS = 4000; + +// Normalizes text for matching. Word separators become spaces so that a query typed as +// `list issues`, `list_issues`, or `listIssues` reaches a tool named `github_list_issues`: MCP tool +// names are overwhelmingly snake_case or camelCase, and a portal prefixes each with `{server_id}_`, +// so leaving word boundaries in place would make the obvious query miss the obvious tool. +function normalize(text: string): string { + return text.slice(0, MAX_SEARCHABLE_FIELD_CHARS) + .replace(/([a-z0-9])([A-Z])/g, "$1 $2") + .toLowerCase() + .replace(/[_\-\s]+/g, " "); +} + +/** Splits a query into normalized terms that a matching tool must all contain. */ +export function toolQueryTerms(query: string): string[] { + return [...new Set(normalize(query.slice(0, MAX_QUERY_CHARS)).split(" ").filter(Boolean))]; +} + +/** Returns whether a tool's name, title, and description collectively match every term. */ +export function matchesToolQuery( + tool: Pick, + terms: readonly string[], +): boolean { + if (terms.length === 0) return false; + const text = [tool.name, tool.title, tool.description] + .filter((value): value is string => typeof value === "string") + .map(normalize) + .join(" "); + return terms.every(term => text.includes(term)); +} diff --git a/packages/mcp-shared/src/tools.ts b/packages/mcp-shared/src/tools.ts index fd72ce97f..b89dfbc27 100644 --- a/packages/mcp-shared/src/tools.ts +++ b/packages/mcp-shared/src/tools.ts @@ -2,8 +2,13 @@ // Nothing outside this file reads a tool's `annotations`. import type { ActionKind } from "@gadgets/workshop-shared/gatekeeper"; -import type { McpContentBlock, McpTool, McpToolCallResult } from "./client.js"; -import type { McpCallResult, McpToolInfo } from "./types"; +import { + clampToolSummary, + type McpContentBlock, + type McpTool, + type McpToolCallResult, +} from "./client.js"; +import type { McpCallResult, McpToolInfo, McpToolSummary } from "./types"; import { hexEncode } from "./util.js"; /** @@ -82,10 +87,7 @@ export function classifyTool(tool: McpTool, trust: ServerTrust): ClassifiedTool }; } -/** - * The tool as a Gadget sees it. `classifiedBy` is carried through so an audit can find every call - * that was trusted on the server's word. - */ +/** The tool as a Gadget sees it, retaining the source of its read/action classification. */ export function toolInfo(entry: ClassifiedTool): McpToolInfo { return { name: entry.tool.name, @@ -97,9 +99,21 @@ export function toolInfo(entry: ClassifiedTool): McpToolInfo { }; } +/** The bounded, schema-free form returned by catalog search. */ +export function toolSummary(entry: ClassifiedTool): McpToolSummary { + const tool = clampToolSummary(entry.tool); + return { + name: tool.name, + title: tool.title, + description: tool.description, + mode: entry.mode, + classifiedBy: entry.classifiedBy, + }; +} + /** - * The approval-policy identity of one tool on one binding. `scopeTag` is caller-supplied so that two - * connectors using the same binding id cannot share pre-approvals. + * Returns the approval-policy identity of one tool on one binding. `scopeTag` prevents two + * connectors using the same binding id from sharing pre-approvals. */ export function actionKindFor(scopeTag: string, toolName: string): ActionKind { return { tag: `${encodeURIComponent(scopeTag)}:${encodeURIComponent(toolName)}`, label: toolName }; @@ -179,13 +193,12 @@ function quoteUntrusted(text: string, max: number): string { return clipped.split("\n").map(line => `> ${line}`).join("\n"); } -// Renders server-chosen text inside a Markdown code span. -// -// Tool names and endpoints are placed in backticks so the approver can see them exactly as sent, but -// a name is as server-controlled as a description: one containing a backtick closes the span and -// everything after it becomes prose the server wrote in the prompt's own voice. Backticks are -// dropped and the text is flattened, so what is shown cannot be more than one inline span. -function codeSpan(text: string, max = MAX_INLINE_TEXT): string { +/** + * Renders server-chosen text inside a bounded Markdown code span. + * + * Backticks are dropped and whitespace is flattened so the value cannot escape into prompt prose. + */ +export function codeSpan(text: string, max = MAX_INLINE_TEXT): string { const cleaned = text.replace(/`/g, "").replace(/\s+/g, " ").trim(); const clipped = cleaned.length > max ? `${cleaned.slice(0, max)}\u2026` : cleaned; return `\`${clipped || "(unnamed)"}\``; @@ -194,10 +207,11 @@ function codeSpan(text: string, max = MAX_INLINE_TEXT): string { // Longest server-chosen name or endpoint shown inline in a prompt. const MAX_INLINE_TEXT = 120; -// Renders server-chosen text as inline prose, with the characters that would let it forge structure -// removed. `account.ts` already does this to a server's reported name before storing it; this is the -// same guard at the point of use, for the callers that pass a name from somewhere else. -function plainInline(text: string, max = MAX_INLINE_TEXT): string { +/** + * Renders untrusted text as inline prose, removing characters that could forge Markdown structure. + * Exported for observation records quoting agent-chosen text such as search queries. + */ +export function plainInline(text: string, max = MAX_INLINE_TEXT): string { const cleaned = text.replace(/[`*_[\]()#>|]/g, "").replace(/\s+/g, " ").trim(); const clipped = cleaned.length > max ? `${cleaned.slice(0, max)}\u2026` : cleaned; return clipped || "(unnamed)"; diff --git a/packages/mcp-shared/src/types.d.ts b/packages/mcp-shared/src/types.d.ts index 8085c908a..2a785092d 100644 --- a/packages/mcp-shared/src/types.d.ts +++ b/packages/mcp-shared/src/types.d.ts @@ -1,8 +1,8 @@ // Base types for MCP-server sessions. // // These are prepended to every generated per-server `.d.ts` (see `schema-to-ts.ts`), so a workspace's -// coding agent always has them in scope. One method per tool, plus `callTool` overloads, is -// generated from the server's own tool catalog and appended below this file's contents. +// coding agent always has them in scope. One method per tool and `callTool` overloads are generated +// from the server's own tool catalog and appended below this file's contents. /** A block of content returned by an MCP tool. */ export type McpContent = @@ -71,6 +71,14 @@ export type McpToolInfo = { * action. Recorded so an audit can find every call that was trusted on the server's word. */ classifiedBy: "server-annotation" | "default"; - /** JSON Schema for the tool's arguments, exactly as the server published it. */ + /** JSON Schema for the tool's arguments when it fits the connector's definition budget. */ inputSchema?: unknown; }; + +/** Bounded search result. Request the exact name through `listTools({ name })` for its schema. */ +export type McpToolSummary = Omit; + +/** Progressive catalog lookup through the existing `listTools` session method. */ +export type McpToolListOptions = + | { search: string; name?: never } + | { name: string; search?: never };