Batch toolkit policy resolution to bound MCP connect time
Unified · Split
| @@ -0,0 +1,103 @@ | ||
| 1 | +// The toolkit-connect performance bug, reproduced with the REAL opencode | |
| 2 | +// binary in a REAL terminal. A user adds a dozen sources to their workspace, | |
| 3 | +// makes a toolkit, and points OpenCode at /mcp/toolkits/<slug>. OpenCode runs | |
| 4 | +// its own OAuth (discovery, DCR, PKCE) and then `mcp list`. The whole session | |
| 5 | +// runs in one recorded PTY: the run's terminal.cast replays exactly what a | |
| 6 | +// user at a shell sees. | |
| 7 | +// | |
| 8 | +// Before the fix, building the toolkit session walks the WHOLE catalog with a | |
| 9 | +// per-tool policy resolution (an N+1 that scales with total tools, not toolkit | |
| 10 | +// size), so `mcp list` blows past OpenCode's connect timeout and the server | |
| 11 | +// shows "failed". After the fix the walk is two batched reads and the same | |
| 12 | +// command shows "connected". The catalog is seeded over the public API exactly | |
| 13 | +// as a user would build it. | |
| 14 | +import { join } from "node:path"; | |
| 15 | + | |
| 16 | +import { expect } from "@effect/vitest"; | |
| 17 | +import { Effect } from "effect"; | |
| 18 | + | |
| 19 | +import { scenario } from "../src/scenario"; | |
| 20 | +import { Api, Cli, OpenCode, RunDir, Target } from "../src/services"; | |
| 21 | +import { catalogApi, seedLargeCatalog } from "../scenarios/support/large-catalog"; | |
| 22 | + | |
| 23 | +const SERVER_NAME = "executor"; | |
| 24 | + | |
| 25 | +scenario( | |
| 26 | + "Toolkits · the real OpenCode binary connects to a toolkit over a large catalog", | |
| 27 | + { timeout: 240_000 }, | |
| 28 | + Effect.scoped( | |
| 29 | + Effect.gen(function* () { | |
| 30 | + const target = yield* Target; | |
| 31 | + const opencode = yield* OpenCode; | |
| 32 | + const runDir = yield* RunDir; | |
| 33 | + const cli = yield* Cli; | |
| 34 | + const { client: makeClient } = yield* Api; | |
| 35 | + | |
| 36 | + const identity = yield* target.newIdentity(); | |
| 37 | + const email = identity.credentials?.email ?? identity.label; | |
| 38 | + const client = yield* makeClient(catalogApi, identity); | |
| 39 | + | |
| 40 | + // Build a production-shaped workspace: one real spec + ten more sources, | |
| 41 | + // a toolkit scoped to one of them. ~3,300 tools across 11 sources — large | |
| 42 | + // enough that the pre-fix per-tool policy N+1 pushes the toolkit connect | |
| 43 | + // past OpenCode's client timeout. | |
| 44 | + const seeded = yield* seedLargeCatalog(client); | |
| 45 | + | |
| 46 | + yield* Effect.gen(function* () { | |
| 47 | + const toolkitUrl = new URL( | |
| 48 | + `/mcp/toolkits/${seeded.toolkitSlug}`, | |
| 49 | + target.baseUrl, | |
| 50 | + ).toString(); | |
| 51 | + const home = opencode.makeHome(SERVER_NAME, toolkitUrl); | |
| 52 | + // First-run database migration happens off camera. | |
| 53 | + yield* Effect.sync(() => opencode.warmUp(home)); | |
| 54 | + | |
| 55 | + yield* cli.session( | |
| 56 | + ["bash", "--norc"], | |
| 57 | + async (term) => { | |
| 58 | + await term.screen.waitForText("$", { timeoutMs: 10_000 }); | |
| 59 | + | |
| 60 | + const outputAfter = (text: string, line: string): string | null => { | |
| 61 | + const echoed = text.lastIndexOf(line); | |
| 62 | + if (echoed === -1) return null; | |
| 63 | + const after = text.slice(echoed + line.length); | |
| 64 | + return after.trimEnd().endsWith("\n$") ? after : null; | |
| 65 | + }; | |
| 66 | + const sh = async (line: string, timeoutMs: number) => { | |
| 67 | + await term.keyboard.type(line); | |
| 68 | + await term.keyboard.press("Enter"); | |
| 69 | + const snapshot = await term.screen.waitUntil( | |
| 70 | + (current) => outputAfter(current.text, line) !== null, | |
| 71 | + { timeoutMs }, | |
| 72 | + ); | |
| 73 | + return outputAfter(snapshot.text, line) ?? ""; | |
| 74 | + }; | |
| 75 | + | |
| 76 | + // OpenCode completes MCP OAuth for real: discovery, DCR, PKCE, its | |
| 77 | + // own scope request, its own token store. | |
| 78 | + const consent = opencode.completeOAuthConsent(home, email, home.openedUrls().length); | |
| 79 | + const auth = await sh(`opencode mcp auth ${SERVER_NAME}`, 90_000); | |
| 80 | + await consent; | |
| 81 | + expect(auth, "opencode mcp auth completes").not.toContain("failed"); | |
| 82 | + | |
| 83 | + // The load-bearing line: listing the toolkit server forces OpenCode | |
| 84 | + // to establish the session, which builds the execute-tool | |
| 85 | + // description over the whole catalog. Pre-fix this is where the | |
| 86 | + // per-tool N+1 times the client out; post-fix it returns promptly. | |
| 87 | + const listed = await sh("opencode mcp list", 120_000); | |
| 88 | + expect( | |
| 89 | + listed, | |
| 90 | + "OpenCode connects to the toolkit even with a large catalog behind it", | |
| 91 | + ).toContain("connected"); | |
| 92 | + }, | |
| 93 | + { | |
| 94 | + cwd: home.projectDir, | |
| 95 | + env: { ...home.env, PS1: "$ ", BASH_SILENCE_DEPRECATION_WARNING: "1" }, | |
| 96 | + record: join(runDir, "terminal.cast"), | |
| 97 | + viewport: { cols: 100, rows: 40 }, | |
| 98 | + }, | |
| 99 | + ); | |
| 100 | + }).pipe(Effect.ensuring(seeded.cleanup)); | |
| 101 | + }), | |
| 102 | + ), | |
| 103 | +); | |
| @@ -0,0 +1,186 @@ | ||
| 1 | +// Cloud: a toolkit MCP endpoint must connect in time that does NOT scale with | |
| 2 | +// the size of the surrounding catalog. | |
| 3 | +// | |
| 4 | +// When a client (OpenCode, Claude Code, mcporter) connects to | |
| 5 | +// /mcp/toolkits/<slug>, the server builds the `execute` tool's description, | |
| 6 | +// which lists the workspace's connections, which (for a toolkit session) runs | |
| 7 | +// the policy engine over the WHOLE catalog to decide tool visibility. A | |
| 8 | +// per-tool policy resolution there is an N+1 that scales with total catalog | |
| 9 | +// size, not toolkit size: a workspace with thousands of tools across a dozen | |
| 10 | +// sources pushes the connect past the MCP client's connect timeout, and the | |
| 11 | +// toolkit appears permanently "failed" even though nothing is broken. | |
| 12 | +// | |
| 13 | +// This scenario seeds a production-shaped catalog (one real OpenAPI spec plus | |
| 14 | +// enough synthetic sources to look like a working workspace, ~3,300 tools over | |
| 15 | +// 11 sources) and asserts the catalog size adds only a small, bounded cost to | |
| 16 | +// connect. The control is a fresh identity with a near-empty catalog: it pays | |
| 17 | +// the same OAuth + MCP handshake, so the DIFFERENCE isolates the catalog-walk | |
| 18 | +// cost the fix removes. Before the fix this delta is tens of seconds (and in | |
| 19 | +// production, a hard timeout); after it, it is sub-second. | |
| 20 | + | |
| 21 | +import { expect } from "@effect/vitest"; | |
| 22 | +import { Effect } from "effect"; | |
| 23 | +import { AuthTemplateSlug, ConnectionName, IntegrationSlug } from "@executor-js/sdk/shared"; | |
| 24 | + | |
| 25 | +import { scenario } from "../src/scenario"; | |
| 26 | +import { Api, Mcp, Target } from "../src/services"; | |
| 27 | +import type { McpSurface } from "../src/surfaces/mcp"; | |
| 28 | +import type { Identity } from "../src/target"; | |
| 29 | +import { | |
| 30 | + catalogApi, | |
| 31 | + seedLargeCatalog, | |
| 32 | + unique, | |
| 33 | + type SeededCatalog, | |
| 34 | +} from "../scenarios/support/large-catalog"; | |
| 35 | + | |
| 36 | +const toolkitUrl = (baseUrl: string, slug: string): string => | |
| 37 | + new URL(`/mcp/toolkits/${slug}`, baseUrl).toString(); | |
| 38 | + | |
| 39 | +// The extra wall-clock a ~3,300-tool catalog is allowed to add to a toolkit | |
| 40 | +// connect over a near-empty one. Post-fix the catalog walk is a couple of | |
| 41 | +// batched reads (~tens of ms); pre-fix it is 2 uncached reads PER TOOL, tens | |
| 42 | +// of seconds. 10s sits an order of magnitude below the regression and well | |
| 43 | +// above OAuth/scheduling jitter, so it is decisive without being flaky. | |
| 44 | +const MAX_CATALOG_CONNECT_OVERHEAD_MS = 10_000; | |
| 45 | + | |
| 46 | +// One-operation OpenAPI spec for the control identity: a real toolkit session, | |
| 47 | +// minimal catalog, so its connect time is "the handshake without the walk". | |
| 48 | +const tinySpec = (baseUrl: string): string => | |
| 49 | + JSON.stringify({ | |
| 50 | + openapi: "3.0.3", | |
| 51 | + info: { title: "Tiny API", version: "1.0.0" }, | |
| 52 | + servers: [{ url: baseUrl }], | |
| 53 | + paths: { | |
| 54 | + "/ping/{id}": { | |
| 55 | + get: { | |
| 56 | + operationId: "getPing", | |
| 57 | + security: [{ apiKey: [] }], | |
| 58 | + parameters: [{ name: "id", in: "path", required: true, schema: { type: "string" } }], | |
| 59 | + responses: { "200": { description: "ok" } }, | |
| 60 | + }, | |
| 61 | + }, | |
| 62 | + }, | |
| 63 | + components: { securitySchemes: { apiKey: { type: "apiKey", in: "header", name: "x-tok" } } }, | |
| 64 | + }); | |
| 65 | + | |
| 66 | +/** Time how long a fresh MCP client takes to connect to a toolkit endpoint and | |
| 67 | + * read its advertised tools (the path that pays the catalog-walk cost). */ | |
| 68 | +const timeToolkitConnect = (mcp: McpSurface, identity: Identity, url: string) => | |
| 69 | + Effect.gen(function* () { | |
| 70 | + const session = mcp.session(identity, { url }); | |
| 71 | + const startedAt = Date.now(); | |
| 72 | + const defs = yield* session.describeTools(); | |
| 73 | + const elapsedMs = Date.now() - startedAt; | |
| 74 | + return { elapsedMs, toolNames: defs.map((d) => d.name) }; | |
| 75 | + }); | |
| 76 | + | |
| 77 | +scenario( | |
| 78 | + "Toolkits · connect time does not scale with catalog size", | |
| 79 | + { timeout: 240_000 }, | |
| 80 | + Effect.scoped( | |
| 81 | + Effect.gen(function* () { | |
| 82 | + const target = yield* Target; | |
| 83 | + const mcp = yield* Mcp; | |
| 84 | + const { client: makeClient } = yield* Api; | |
| 85 | + | |
| 86 | + // --- control: a fresh identity with a near-empty catalog --------------- | |
| 87 | + const controlIdentity = yield* target.newIdentity(); | |
| 88 | + const controlClient = yield* makeClient(catalogApi, controlIdentity); | |
| 89 | + const controlIntegration = unique("tiny"); | |
| 90 | + const controlToolkitName = unique("control-kit"); | |
| 91 | + | |
| 92 | + const control = yield* Effect.gen(function* () { | |
| 93 | + yield* controlClient.openapi.addSpec({ | |
| 94 | + payload: { | |
| 95 | + spec: { kind: "blob", value: tinySpec("https://tiny.example") }, | |
| 96 | + slug: IntegrationSlug.make(controlIntegration), | |
| 97 | + baseUrl: "https://tiny.example", | |
| 98 | + authenticationTemplate: [ | |
| 99 | + { | |
| 100 | + slug: "apiKey", | |
| 101 | + type: "apiKey", | |
| 102 | + headers: { "x-tok": [{ type: "variable", name: "token" }] }, | |
| 103 | + }, | |
| 104 | + ], | |
| 105 | + }, | |
| 106 | + }); | |
| 107 | + yield* controlClient.connections.create({ | |
| 108 | + payload: { | |
| 109 | + owner: "org", | |
| 110 | + name: ConnectionName.make("conn0"), | |
| 111 | + integration: IntegrationSlug.make(controlIntegration), | |
| 112 | + template: AuthTemplateSlug.make("apiKey"), | |
| 113 | + value: "unused-token", | |
| 114 | + }, | |
| 115 | + }); | |
| 116 | + const toolkit = yield* controlClient.toolkits.create({ | |
| 117 | + payload: { owner: "org", name: controlToolkitName }, | |
| 118 | + }); | |
| 119 | + yield* controlClient.toolkits.createConnection({ | |
| 120 | + params: { toolkitId: toolkit.id }, | |
| 121 | + payload: { pattern: `${controlIntegration}.org.conn0.*` }, | |
| 122 | + }); | |
| 123 | + return yield* timeToolkitConnect( | |
| 124 | + mcp, | |
| 125 | + controlIdentity, | |
| 126 | + toolkitUrl(target.baseUrl, toolkit.slug), | |
| 127 | + ); | |
| 128 | + }).pipe( | |
| 129 | + Effect.ensuring( | |
| 130 | + controlClient.openapi | |
| 131 | + .removeSpec({ params: { slug: controlIntegration } }) | |
| 132 | + .pipe(Effect.ignore), | |
| 133 | + ), | |
| 134 | + ); | |
| 135 | + | |
| 136 | + expect(control.toolNames, "control toolkit advertises the execute tool").toContain("execute"); | |
| 137 | + | |
| 138 | + // --- subject: a fresh identity with a large, production-shaped catalog - | |
| 139 | + const bigIdentity = yield* target.newIdentity(); | |
| 140 | + const bigClient = yield* makeClient(catalogApi, bigIdentity); | |
| 141 | + let seededCatalog: SeededCatalog | undefined; | |
| 142 | + | |
| 143 | + yield* Effect.gen(function* () { | |
| 144 | + // 1 real source (Vercel, 322 ops) + 10 synthetic = 11 sources, ~2,200 | |
| 145 | + // tools. Sized so the pre-fix N+1 connect (~27s here) still completes | |
| 146 | + // under the MCP client's connect timeout, so the regression surfaces as | |
| 147 | + // a failed ASSERTION on the overhead rather than a connect crash. | |
| 148 | + const seeded = yield* seedLargeCatalog(bigClient, { | |
| 149 | + includeRealSpec: true, | |
| 150 | + syntheticSources: 10, | |
| 151 | + opsPerSource: 190, | |
| 152 | + }); | |
| 153 | + seededCatalog = seeded; | |
| 154 | + expect( | |
| 155 | + seeded.toolCount, | |
| 156 | + "the seeded catalog is large enough to expose the N+1 (thousands of tools)", | |
| 157 | + ).toBeGreaterThan(2_000); | |
| 158 | + expect( | |
| 159 | + seeded.integrationSlugs.length, | |
| 160 | + "the catalog spans a production-like number of sources", | |
| 161 | + ).toBeGreaterThanOrEqual(11); | |
| 162 | + | |
| 163 | + const big = yield* timeToolkitConnect( | |
| 164 | + mcp, | |
| 165 | + bigIdentity, | |
| 166 | + toolkitUrl(target.baseUrl, seeded.toolkitSlug), | |
| 167 | + ); | |
| 168 | + expect(big.toolNames, "large-catalog toolkit advertises the execute tool").toContain( | |
| 169 | + "execute", | |
| 170 | + ); | |
| 171 | + | |
| 172 | + const overhead = big.elapsedMs - control.elapsedMs; | |
| 173 | + expect( | |
| 174 | + overhead, | |
| 175 | + `catalog of ${seeded.toolCount} tools across ${seeded.integrationSlugs.length} sources ` + | |
| 176 | + `added ${overhead}ms to connect (big ${big.elapsedMs}ms vs control ${control.elapsedMs}ms); ` + | |
| 177 | + `the per-tool policy N+1 would make this scale into tens of seconds`, | |
| 178 | + ).toBeLessThan(MAX_CATALOG_CONNECT_OVERHEAD_MS); | |
| 179 | + }).pipe( | |
| 180 | + Effect.ensuring( | |
| 181 | + Effect.suspend(() => seededCatalog?.cleanup ?? Effect.void).pipe(Effect.ignore), | |
| 182 | + ), | |
| 183 | + ); | |
| 184 | + }), | |
| 185 | + ), | |
| 186 | +); | |
| @@ -0,0 +1,220 @@ | ||
| 1 | +// Seeds a production-shaped catalog for the toolkit-policy performance | |
| 2 | +// scenarios, plus the pure builders behind it. | |
| 3 | +// | |
| 4 | +// Why a large catalog at all: a toolkit MCP endpoint resolves tool visibility | |
| 5 | +// by running the policy engine over the WHOLE catalog on connect (the | |
| 6 | +// getDescription -> connections.list -> toolsList path), not just the toolkit's | |
| 7 | +// own connections. A per-tool resolution there is an N+1 that scales with total | |
| 8 | +// catalog size, so a realistic workspace (a real spec plus enough sources to | |
| 9 | +// look like one) is what surfaces the regression. | |
| 10 | +// | |
| 11 | +// `catalogApi` is exported so scenarios build their client from the SAME | |
| 12 | +// composition `seedLargeCatalog` is typed against — that keeps the seeding | |
| 13 | +// fully typed (no structural-client gymnastics) while staying DRY across the | |
| 14 | +// deterministic guard and the OpenCode recording. | |
| 15 | +import { randomBytes } from "node:crypto"; | |
| 16 | +import { readFileSync } from "node:fs"; | |
| 17 | +import { fileURLToPath } from "node:url"; | |
| 18 | + | |
| 19 | +import { Effect } from "effect"; | |
| 20 | +import type { HttpApi, HttpApiClient } from "effect/unstable/httpapi"; | |
| 21 | +import { composePluginApi } from "@executor-js/api/server"; | |
| 22 | +import { openApiHttpPlugin } from "@executor-js/plugin-openapi/api"; | |
| 23 | +import { toolkitsPlugin } from "@executor-js/plugin-toolkits/server"; | |
| 24 | +import { AuthTemplateSlug, ConnectionName, IntegrationSlug } from "@executor-js/sdk/shared"; | |
| 25 | + | |
| 26 | +export const unique = (prefix: string): string => `${prefix}_${randomBytes(4).toString("hex")}`; | |
| 27 | + | |
| 28 | +const VERCEL_SPEC_PATH = fileURLToPath( | |
| 29 | + new URL("../../../packages/plugins/openapi/fixtures/vercel.json", import.meta.url), | |
| 30 | +); | |
| 31 | + | |
| 32 | +/** The real Vercel OpenAPI fixture (322 operations) as a JSON string. */ | |
| 33 | +export const vercelSpecText = (): string => readFileSync(VERCEL_SPEC_PATH, "utf-8"); | |
| 34 | + | |
| 35 | +/** An OpenAPI 3 doc with `ops` independent GET operations, shaped like a real | |
| 36 | + * REST surface (path params, an apiKey scheme, JSON responses). */ | |
| 37 | +export const syntheticSpec = (title: string, ops: number, baseUrl: string): string => { | |
| 38 | + const paths: Record<string, unknown> = {}; | |
| 39 | + for (let i = 0; i < ops; i++) { | |
| 40 | + paths[`/resource${i}/{id}`] = { | |
| 41 | + get: { | |
| 42 | + operationId: `get_resource_${i}`, | |
| 43 | + summary: `Fetch resource ${i}`, | |
| 44 | + security: [{ apiKey: [] }], | |
| 45 | + parameters: [{ name: "id", in: "path", required: true, schema: { type: "string" } }], | |
| 46 | + responses: { | |
| 47 | + "200": { | |
| 48 | + description: "ok", | |
| 49 | + content: { | |
| 50 | + "application/json": { | |
| 51 | + schema: { type: "object", properties: { id: { type: "string" } } }, | |
| 52 | + }, | |
| 53 | + }, | |
| 54 | + }, | |
| 55 | + }, | |
| 56 | + }, | |
| 57 | + }; | |
| 58 | + } | |
| 59 | + return JSON.stringify({ | |
| 60 | + openapi: "3.0.3", | |
| 61 | + info: { title, version: "1.0.0" }, | |
| 62 | + servers: [{ url: baseUrl }], | |
| 63 | + paths, | |
| 64 | + components: { securitySchemes: { apiKey: { type: "apiKey", in: "header", name: "x-tok" } } }, | |
| 65 | + }); | |
| 66 | +}; | |
| 67 | + | |
| 68 | +export interface SeedSource { | |
| 69 | + readonly slug: string; | |
| 70 | + readonly specText: string; | |
| 71 | + readonly baseUrl: string; | |
| 72 | + readonly description?: string; | |
| 73 | +} | |
| 74 | + | |
| 75 | +export interface SeedPlan { | |
| 76 | + readonly sources: ReadonlyArray<SeedSource>; | |
| 77 | + /** The first source's slug — the one the toolkit is scoped to. */ | |
| 78 | + readonly firstSlug: string; | |
| 79 | + /** Connection pattern the toolkit binds (the first source, org, conn0). */ | |
| 80 | + readonly toolkitConnectionPattern: string; | |
| 81 | +} | |
| 82 | + | |
| 83 | +export interface SeedOptions { | |
| 84 | + /** Synthetic sources on top of the real spec. Default 10. */ | |
| 85 | + readonly syntheticSources?: number; | |
| 86 | + /** Operations per synthetic source. Default 300. */ | |
| 87 | + readonly opsPerSource?: number; | |
| 88 | + /** Include the real Vercel fixture (322 ops) as one source. Default true. */ | |
| 89 | + readonly includeRealSpec?: boolean; | |
| 90 | +} | |
| 91 | + | |
| 92 | +/** | |
| 93 | + * Plan a large, production-shaped catalog. With the defaults (1 real + 10 | |
| 94 | + * synthetic sources, 300 ops each) it is ~3,300 tools across 11 sources, where | |
| 95 | + * the per-tool policy N+1 turns a toolkit connect from sub-second into a 30s+ | |
| 96 | + * client timeout. Each source gets exactly one org connection at `conn0`. | |
| 97 | + */ | |
| 98 | +export const planLargeCatalog = (options: SeedOptions = {}): SeedPlan => { | |
| 99 | + const syntheticSources = options.syntheticSources ?? 10; | |
| 100 | + const opsPerSource = options.opsPerSource ?? 300; | |
| 101 | + const includeRealSpec = options.includeRealSpec ?? true; | |
| 102 | + | |
| 103 | + const sources: SeedSource[] = []; | |
| 104 | + if (includeRealSpec) { | |
| 105 | + sources.push({ | |
| 106 | + slug: unique("vercel"), | |
| 107 | + specText: vercelSpecText(), | |
| 108 | + baseUrl: "https://api.vercel.com", | |
| 109 | + description: "Vercel API", | |
| 110 | + }); | |
| 111 | + } | |
| 112 | + for (let s = 0; s < syntheticSources; s++) { | |
| 113 | + sources.push({ | |
| 114 | + slug: unique("svc"), | |
| 115 | + specText: syntheticSpec(`Service ${s}`, opsPerSource, "https://service.example"), | |
| 116 | + baseUrl: "https://service.example", | |
| 117 | + }); | |
| 118 | + } | |
| 119 | + | |
| 120 | + const firstSlug = sources[0]!.slug; | |
| 121 | + return { | |
| 122 | + sources, | |
| 123 | + firstSlug, | |
| 124 | + toolkitConnectionPattern: `${firstSlug}.org.conn0.*`, | |
| 125 | + }; | |
| 126 | +}; | |
| 127 | + | |
| 128 | +// --------------------------------------------------------------------------- | |
| 129 | +// Typed seeding. Scenarios build their client from `catalogApi` so the client | |
| 130 | +// type here matches theirs exactly. | |
| 131 | +// --------------------------------------------------------------------------- | |
| 132 | + | |
| 133 | +/** The plugin API the seeding (and the scenarios) speak: OpenAPI specs + | |
| 134 | + * toolkits. Build the scenario's client from THIS so types line up. */ | |
| 135 | +export const catalogApi = composePluginApi([openApiHttpPlugin(), toolkitsPlugin()] as const); | |
| 136 | + | |
| 137 | +type GroupsOf<A> = A extends HttpApi.HttpApi<infer _Id, infer Groups> ? Groups : never; | |
| 138 | +export type CatalogClient = HttpApiClient.Client<GroupsOf<typeof catalogApi>>; | |
| 139 | + | |
| 140 | +export interface SeededCatalog { | |
| 141 | + /** The org toolkit slug to point an MCP client at: /mcp/toolkits/<slug>. */ | |
| 142 | + readonly toolkitSlug: string; | |
| 143 | + readonly toolkitId: string; | |
| 144 | + /** Total tools in the identity's catalog after seeding (the N+1 multiplier). */ | |
| 145 | + readonly toolCount: number; | |
| 146 | + /** Integration slugs created, for assertions / debugging. */ | |
| 147 | + readonly integrationSlugs: ReadonlyArray<string>; | |
| 148 | + /** Remove every spec + toolkit this seeder created. */ | |
| 149 | + readonly cleanup: Effect.Effect<void>; | |
| 150 | +} | |
| 151 | + | |
| 152 | +/** | |
| 153 | + * Build the planned catalog under the current identity (one org connection per | |
| 154 | + * source) and a toolkit scoped to the first source, then report the total tool | |
| 155 | + * count and a finalizer. The toolkit's own surface is small; the connect cost | |
| 156 | + * the scenarios guard comes from the policy engine walking the whole catalog. | |
| 157 | + */ | |
| 158 | +export const seedLargeCatalog = ( | |
| 159 | + client: CatalogClient, | |
| 160 | + options: SeedOptions = {}, | |
| 161 | +): Effect.Effect<SeededCatalog, unknown> => | |
| 162 | + Effect.gen(function* () { | |
| 163 | + const plan = planLargeCatalog(options); | |
| 164 | + | |
| 165 | + for (const source of plan.sources) { | |
| 166 | + yield* client.openapi.addSpec({ | |
| 167 | + payload: { | |
| 168 | + spec: { kind: "blob", value: source.specText }, | |
| 169 | + slug: IntegrationSlug.make(source.slug), | |
| 170 | + baseUrl: source.baseUrl, | |
| 171 | + ...(source.description ? { description: source.description } : {}), | |
| 172 | + authenticationTemplate: [ | |
| 173 | + { | |
| 174 | + slug: "apiKey", | |
| 175 | + type: "apiKey", | |
| 176 | + headers: { "x-tok": [{ type: "variable", name: "token" }] }, | |
| 177 | + }, | |
| 178 | + ], | |
| 179 | + }, | |
| 180 | + }); | |
| 181 | + yield* client.connections.create({ | |
| 182 | + payload: { | |
| 183 | + owner: "org", | |
| 184 | + name: ConnectionName.make("conn0"), | |
| 185 | + integration: IntegrationSlug.make(source.slug), | |
| 186 | + template: AuthTemplateSlug.make("apiKey"), | |
| 187 | + value: "unused-token", | |
| 188 | + }, | |
| 189 | + }); | |
| 190 | + } | |
| 191 | + | |
| 192 | + const allTools = yield* client.tools.list({ query: {} }); | |
| 193 | + const toolCount = allTools.length; | |
| 194 | + | |
| 195 | + const toolkit = yield* client.toolkits.create({ | |
| 196 | + payload: { owner: "org", name: unique("perf-kit") }, | |
| 197 | + }); | |
| 198 | + yield* client.toolkits.createConnection({ | |
| 199 | + params: { toolkitId: toolkit.id }, | |
| 200 | + payload: { pattern: plan.toolkitConnectionPattern }, | |
| 201 | + }); | |
| 202 | + | |
| 203 | + const integrationSlugs = plan.sources.map((source) => source.slug); | |
| 204 | + const cleanup = Effect.gen(function* () { | |
| 205 | + yield* client.toolkits.remove({ params: { toolkitId: toolkit.id } }).pipe(Effect.ignore); | |
| 206 | + yield* Effect.forEach( | |
| 207 | + integrationSlugs, | |
| 208 | + (slug) => client.openapi.removeSpec({ params: { slug } }).pipe(Effect.ignore), | |
| 209 | + { discard: true }, | |
| 210 | + ); | |
| 211 | + }).pipe(Effect.ignore); | |
| 212 | + | |
| 213 | + return { | |
| 214 | + toolkitSlug: toolkit.slug, | |
| 215 | + toolkitId: toolkit.id, | |
| 216 | + toolCount, | |
| 217 | + integrationSlugs, | |
| 218 | + cleanup, | |
| 219 | + }; | |
| 220 | + }); | |
| @@ -66,10 +66,15 @@ export const makeOpenCodeHome = ( · 65 unmodified lines | ||
| 66 | 66 | : {}), |
| 67 | 67 | }), |
| 68 | 68 | ); |
| 69 | − // OpenCode launches the OAuth URL via `open`; the shim records it instead. | |
| 70 | − writeFileSync(join(binDir, "open"), `#!/bin/sh\necho "$@" >> ${openedUrlsFile}\nexit 0\n`, { | |
| 71 | − mode: 0o755, | |
| 72 | − }); | |
| 69 | + // OpenCode launches the OAuth URL via the platform browser opener; the shim | |
| 70 | + // records it instead. macOS uses `open`; Linux uses `xdg-open` (and a few | |
| 71 | + // fallback names), so shim them all to the same recorder — otherwise the URL | |
| 72 | + // is never captured on Linux and consent never completes. | |
| 73 | + for (const opener of ["open", "xdg-open", "www-browser", "x-www-browser"]) { | |
| 74 | + writeFileSync(join(binDir, opener), `#!/bin/sh\necho "$@" >> ${openedUrlsFile}\nexit 0\n`, { | |
| 75 | + mode: 0o755, | |
| 76 | + }); | |
| 77 | + } | |
| 73 | 78 | |
| 74 | 79 | return { |
| 75 | 80 | projectDir, |
| @@ -2442,6 +2442,13 @@ export const createExecutor = <const TPlugins extends readonly AnyPlugin[] = rea · 2441 unmodified lines | ||
| 2442 | 2442 | readonly kind: "provider"; |
| 2443 | 2443 | readonly provider: ToolPolicyProvider; |
| 2444 | 2444 | readonly rules: readonly ToolPolicyProviderRule[] | null; |
| 2445 | + } | |
| 2446 | + | { | |
| 2447 | + readonly kind: "prepared"; | |
| 2448 | + readonly resolve: (input: { | |
| 2449 | + readonly toolId: string; | |
| 2450 | + readonly defaultRequiresApproval?: boolean; | |
| 2451 | + }) => EffectivePolicy; | |
| 2445 | 2452 | }; |
| 2446 | 2453 | |
| 2447 | 2454 | const compareProviderPolicyRule = ( |
| @@ -2477,19 +2484,26 @@ export const createExecutor = <const TPlugins extends readonly AnyPlugin[] = rea · 29 unmodified lines | ||
| 2477 | 2484 | |
| 2478 | 2485 | const listActivePolicyRuleSet = (): Effect.Effect<ActivePolicyRuleSet, StorageFailure> => |
| 2479 | 2486 | activeToolPolicyProvider |
| 2480 | − ? activeToolPolicyProvider.resolve | |
| 2481 | − ? Effect.succeed({ | |
| 2482 | − kind: "provider" as const, | |
| 2483 | − provider: activeToolPolicyProvider, | |
| 2484 | − rules: null, | |
| 2485 | − }) | |
| 2486 | − : activeToolPolicyProvider.list().pipe( | |
| 2487 | − Effect.map((rules) => ({ | |
| 2487 | + ? // Batched per-operation resolver: fetch all policy + connection state | |
| 2488 | + // once, then resolve every tool in this operation against that | |
| 2489 | + // snapshot. Avoids the per-tool resolve N+1 on the list surface. | |
| 2490 | + activeToolPolicyProvider.prepare | |
| 2491 | + ? activeToolPolicyProvider | |
| 2492 | + .prepare() | |
| 2493 | + .pipe(Effect.map((resolve) => ({ kind: "prepared" as const, resolve }))) | |
| 2494 | + : activeToolPolicyProvider.resolve | |
| 2495 | + ? Effect.succeed({ | |
| 2488 | 2496 | kind: "provider" as const, |
| 2489 | − provider: activeToolPolicyProvider!, | |
| 2490 | − rules, | |
| 2491 | − })), | |
| 2492 | − ) | |
| 2497 | + provider: activeToolPolicyProvider, | |
| 2498 | + rules: null, | |
| 2499 | + }) | |
| 2500 | + : activeToolPolicyProvider.list().pipe( | |
| 2501 | + Effect.map((rules) => ({ | |
| 2502 | + kind: "provider" as const, | |
| 2503 | + provider: activeToolPolicyProvider!, | |
| 2504 | + rules, | |
| 2505 | + })), | |
| 2506 | + ) | |
| 2493 | 2507 | : core |
| 2494 | 2508 | .findMany("tool_policy", {}) |
| 2495 | 2509 | .pipe(Effect.map((rows) => ({ kind: "global" as const, rows }))); |
| @@ -2499,13 +2513,20 @@ export const createExecutor = <const TPlugins extends readonly AnyPlugin[] = rea · 3 unmodified lines | ||
| 2499 | 2513 | ruleSet: ActivePolicyRuleSet, |
| 2500 | 2514 | defaultRequiresApproval?: boolean, |
| 2501 | 2515 | ): Effect.Effect<EffectivePolicy, StorageFailure> => |
| 2502 | − ruleSet.kind === "provider" | |
| 2503 | − ? ruleSet.provider.resolve | |
| 2504 | − ? ruleSet.provider.resolve({ toolId, defaultRequiresApproval }) | |
| 2505 | − : Effect.succeed(resolveProviderPolicyFromRules(toolId, ruleSet.rules ?? [])) | |
| 2506 | − : Effect.succeed( | |
| 2507 | − resolveEffectivePolicy(toolId, ruleSet.rows, ownerRankForRow, defaultRequiresApproval), | |
| 2508 | − ); | |
| 2516 | + ruleSet.kind === "prepared" | |
| 2517 | + ? Effect.succeed(ruleSet.resolve({ toolId, defaultRequiresApproval })) | |
| 2518 | + : ruleSet.kind === "provider" | |
| 2519 | + ? ruleSet.provider.resolve | |
| 2520 | + ? ruleSet.provider.resolve({ toolId, defaultRequiresApproval }) | |
| 2521 | + : Effect.succeed(resolveProviderPolicyFromRules(toolId, ruleSet.rules ?? [])) | |
| 2522 | + : Effect.succeed( | |
| 2523 | + resolveEffectivePolicy( | |
| 2524 | + toolId, | |
| 2525 | + ruleSet.rows, | |
| 2526 | + ownerRankForRow, | |
| 2527 | + defaultRequiresApproval, | |
| 2528 | + ), | |
| 2529 | + ); | |
| 2509 | 2530 | |
| 2510 | 2531 | // ------------------------------------------------------------------ |
| 2511 | 2532 | // Tools (read surface) |
| @@ -111,6 +111,26 @@ export interface ToolPolicyProvider { · 110 unmodified lines | ||
| 111 | 111 | readonly toolId: string; |
| 112 | 112 | readonly defaultRequiresApproval?: boolean; |
| 113 | 113 | }) => Effect.Effect<EffectivePolicy, StorageFailure>; |
| 114 | + /** | |
| 115 | + * Batched per-operation resolver. When defined, core calls `prepare` once at | |
| 116 | + * the start of an operation (a single tools/list or tools/call), fetching all | |
| 117 | + * the underlying policy + connection state in one pass, and reuses the | |
| 118 | + * returned pure resolver for every tool in that operation. This avoids the | |
| 119 | + * per-tool `resolve` N+1 (2 uncached storage reads per tool) that scales with | |
| 120 | + * the total catalog size on `toolsList`. | |
| 121 | + * | |
| 122 | + * The resolver is intentionally per-operation scoped, not memoized on the | |
| 123 | + * provider: the provider instance is session-scoped (lives across many | |
| 124 | + * requests), so caching on it would serve stale policy state. Each operation | |
| 125 | + * gets a fresh snapshot. | |
| 126 | + */ | |
| 127 | + readonly prepare?: () => Effect.Effect< | |
| 128 | + (input: { | |
| 129 | + readonly toolId: string; | |
| 130 | + readonly defaultRequiresApproval?: boolean; | |
| 131 | + }) => EffectivePolicy, | |
| 132 | + StorageFailure | |
| 133 | + >; | |
| 114 | 134 | } |
| 115 | 135 | |
| 116 | 136 | // --------------------------------------------------------------------------- |
| @@ -509,6 +509,38 @@ const makeToolkitsExtension = (ctx: PluginCtx<ToolkitStorage>) => { · 508 unmodified lines | ||
| 509 | 509 | return resolveToolkitPolicy(toolId, connections, policies, defaultRequiresApproval); |
| 510 | 510 | }); |
| 511 | 511 | |
| 512 | + // Batched form of `resolvePolicyForSlug`: fetch the toolkit, its policies, and | |
| 513 | + // its connections ONCE, then hand back a pure resolver core can run for every | |
| 514 | + // tool in a single tools/list or tools/call. `resolvePolicyForSlug` re-fetches | |
| 515 | + // policies + connections on every tool, which is the per-tool N+1 that scales | |
| 516 | + // with the whole catalog on the list surface. This is byte-for-byte the same | |
| 517 | + // resolution, just hoisted out of the loop. | |
| 518 | + const preparePolicyResolverForSlug = ( | |
| 519 | + slug: string, | |
| 520 | + ): Effect.Effect< | |
| 521 | + (input: { | |
| 522 | + readonly toolId: string; | |
| 523 | + readonly defaultRequiresApproval?: boolean; | |
| 524 | + }) => EffectivePolicy, | |
| 525 | + StorageFailure | |
| 526 | + > => | |
| 527 | + Effect.gen(function* () { | |
| 528 | + const toolkit = yield* getBySlugEntry(slug); | |
| 529 | + if (!toolkit) return () => blockedPolicy(); | |
| 530 | + const isOrg = toolkit.owner === "org"; | |
| 531 | + const policies = yield* listPoliciesForRecord(toolkit.data.id); | |
| 532 | + const connections = yield* listConnectionsForRecord(toolkit.data.id); | |
| 533 | + return (input: { readonly toolId: string; readonly defaultRequiresApproval?: boolean }) => { | |
| 534 | + if (isOrg && isPersonalDynamicToolId(input.toolId)) return blockedPolicy(); | |
| 535 | + return resolveToolkitPolicy( | |
| 536 | + input.toolId, | |
| 537 | + connections, | |
| 538 | + policies, | |
| 539 | + input.defaultRequiresApproval, | |
| 540 | + ); | |
| 541 | + }; | |
| 542 | + }); | |
| 543 | + | |
| 512 | 544 | return { |
| 513 | 545 | list, |
| 514 | 546 | create, |
| @@ -527,6 +559,7 @@ const makeToolkitsExtension = (ctx: PluginCtx<ToolkitStorage>) => { · 12 unmodified lines | ||
| 527 | 559 | removeConnection, |
| 528 | 560 | policyRulesForSlug, |
| 529 | 561 | resolvePolicyForSlug, |
| 562 | + preparePolicyResolverForSlug, | |
| 530 | 563 | }; |
| 531 | 564 | }; |
| 532 | 565 | |
| @@ -638,12 +671,18 @@ const ToolkitsHandlers = HttpApiBuilder.group(ExecutorApiWithToolkits, "toolkits · 105 unmodified lines | ||
| 638 | 671 | ); |
| 639 | 672 | |
| 640 | 673 | const makePolicyProvider = ( |
| 641 | − extension: Pick<ToolkitsExtension, "policyRulesForSlug" | "resolvePolicyForSlug">, | |
| 674 | + extension: Pick< | |
| 675 | + ToolkitsExtension, | |
| 676 | + "policyRulesForSlug" | "resolvePolicyForSlug" | "preparePolicyResolverForSlug" | |
| 677 | + >, | |
| 642 | 678 | slug: string, |
| 643 | 679 | ): ToolPolicyProvider => ({ |
| 644 | 680 | list: () => extension.policyRulesForSlug(slug), |
| 645 | 681 | resolve: ({ toolId, defaultRequiresApproval }) => |
| 646 | 682 | extension.resolvePolicyForSlug(slug, toolId, defaultRequiresApproval), |
| 683 | + // Preferred path: core calls this once per operation, so the toolkit's | |
| 684 | + // policies + connections are fetched once instead of once per tool. | |
| 685 | + prepare: () => extension.preparePolicyResolverForSlug(slug), | |
| 647 | 686 | }); |
| 648 | 687 | |
| 649 | 688 | export const toolkitsPlugin = definePlugin((options: ToolkitsPluginOptions = {}) => { |