diff --git a/doc/benchmarks/2026-10-01-mcp-discovery.json b/doc/benchmarks/2026-10-01-mcp-discovery.json new file mode 100644 index 00000000000..82ccd4901c7 --- /dev/null +++ b/doc/benchmarks/2026-10-01-mcp-discovery.json @@ -0,0 +1,476 @@ +{ + "date": "2026-10-01", + "node": "26.4.0", + "baseline": "f2e0f196308629ef05c7f65782243e1713b208f0", + "samplingIntervalMs": 2, + "series": { + "before": [ + { + "tools": 50, + "parallel": 1, + "iteration": 1, + "visible": 55, + "queries": 722, + "elapsedMs": 146, + "peakHeapIncreaseMiB": 82, + "peakRssMiB": 727, + "retainedHeapIncreaseMiB": 2, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + }, + { + "tools": 50, + "parallel": 1, + "iteration": 2, + "visible": 55, + "queries": 722, + "elapsedMs": 151, + "peakHeapIncreaseMiB": 82, + "peakRssMiB": 728, + "retainedHeapIncreaseMiB": 0, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + }, + { + "tools": 300, + "parallel": 1, + "iteration": 1, + "visible": 305, + "queries": 3889, + "elapsedMs": 855, + "peakHeapIncreaseMiB": 424, + "peakRssMiB": 1155, + "retainedHeapIncreaseMiB": 9, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + }, + { + "tools": 300, + "parallel": 1, + "iteration": 2, + "visible": 305, + "queries": 3889, + "elapsedMs": 875, + "peakHeapIncreaseMiB": 409, + "peakRssMiB": 1159, + "retainedHeapIncreaseMiB": -8, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + }, + { + "tools": 900, + "parallel": 1, + "iteration": 1, + "visible": 905, + "queries": 11489, + "elapsedMs": 3631, + "peakHeapIncreaseMiB": 1698, + "peakRssMiB": 2495, + "retainedHeapIncreaseMiB": 22, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + }, + { + "tools": 900, + "parallel": 1, + "iteration": 2, + "visible": 905, + "queries": 11489, + "elapsedMs": 3540, + "peakHeapIncreaseMiB": 1680, + "peakRssMiB": 2508, + "retainedHeapIncreaseMiB": 0, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + } + ], + "after": [ + { + "tools": 50, + "parallel": 1, + "iteration": 1, + "visible": 55, + "queries": 36, + "elapsedMs": 55, + "peakHeapIncreaseMiB": 20, + "peakRssMiB": 563, + "retainedHeapIncreaseMiB": 0, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + }, + { + "tools": 50, + "parallel": 1, + "iteration": 2, + "visible": 55, + "queries": 36, + "elapsedMs": 24, + "peakHeapIncreaseMiB": 19, + "peakRssMiB": 563, + "retainedHeapIncreaseMiB": 0, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + }, + { + "tools": 50, + "parallel": 1, + "iteration": 3, + "visible": 55, + "queries": 36, + "elapsedMs": 24, + "peakHeapIncreaseMiB": 19, + "peakRssMiB": 563, + "retainedHeapIncreaseMiB": 0, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + }, + { + "tools": 50, + "parallel": 16, + "iteration": 1, + "visible": 55, + "queries": 576, + "elapsedMs": 264, + "peakHeapIncreaseMiB": 50, + "peakRssMiB": 564, + "retainedHeapIncreaseMiB": 7, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + }, + { + "tools": 50, + "parallel": 16, + "iteration": 2, + "visible": 55, + "queries": 576, + "elapsedMs": 242, + "peakHeapIncreaseMiB": 45, + "peakRssMiB": 567, + "retainedHeapIncreaseMiB": 1, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + }, + { + "tools": 50, + "parallel": 16, + "iteration": 3, + "visible": 55, + "queries": 576, + "elapsedMs": 268, + "peakHeapIncreaseMiB": 49, + "peakRssMiB": 569, + "retainedHeapIncreaseMiB": 0, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + }, + { + "tools": 300, + "parallel": 1, + "iteration": 1, + "visible": 305, + "queries": 36, + "elapsedMs": 76, + "peakHeapIncreaseMiB": 29, + "peakRssMiB": 570, + "retainedHeapIncreaseMiB": -3, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + }, + { + "tools": 300, + "parallel": 1, + "iteration": 2, + "visible": 305, + "queries": 36, + "elapsedMs": 77, + "peakHeapIncreaseMiB": 30, + "peakRssMiB": 570, + "retainedHeapIncreaseMiB": 2, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + }, + { + "tools": 300, + "parallel": 1, + "iteration": 3, + "visible": 305, + "queries": 36, + "elapsedMs": 76, + "peakHeapIncreaseMiB": 29, + "peakRssMiB": 570, + "retainedHeapIncreaseMiB": 2, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + }, + { + "tools": 300, + "parallel": 16, + "iteration": 1, + "visible": 305, + "queries": 576, + "elapsedMs": 1051, + "peakHeapIncreaseMiB": 108, + "peakRssMiB": 660, + "retainedHeapIncreaseMiB": 23, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + }, + { + "tools": 300, + "parallel": 16, + "iteration": 2, + "visible": 305, + "queries": 576, + "elapsedMs": 1038, + "peakHeapIncreaseMiB": 85, + "peakRssMiB": 661, + "retainedHeapIncreaseMiB": 0, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + }, + { + "tools": 300, + "parallel": 16, + "iteration": 3, + "visible": 305, + "queries": 576, + "elapsedMs": 1048, + "peakHeapIncreaseMiB": 118, + "peakRssMiB": 706, + "retainedHeapIncreaseMiB": 5, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + }, + { + "tools": 900, + "parallel": 1, + "iteration": 1, + "visible": 905, + "queries": 36, + "elapsedMs": 236, + "peakHeapIncreaseMiB": 37, + "peakRssMiB": 707, + "retainedHeapIncreaseMiB": -16, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + }, + { + "tools": 900, + "parallel": 1, + "iteration": 2, + "visible": 905, + "queries": 36, + "elapsedMs": 236, + "peakHeapIncreaseMiB": 37, + "peakRssMiB": 706, + "retainedHeapIncreaseMiB": 4, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + }, + { + "tools": 900, + "parallel": 1, + "iteration": 3, + "visible": 905, + "queries": 36, + "elapsedMs": 238, + "peakHeapIncreaseMiB": 37, + "peakRssMiB": 706, + "retainedHeapIncreaseMiB": 4, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + }, + { + "tools": 900, + "parallel": 16, + "iteration": 1, + "visible": 905, + "queries": 576, + "elapsedMs": 3636, + "peakHeapIncreaseMiB": 169, + "peakRssMiB": 727, + "retainedHeapIncreaseMiB": 54, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + }, + { + "tools": 900, + "parallel": 16, + "iteration": 2, + "visible": 905, + "queries": 576, + "elapsedMs": 3627, + "peakHeapIncreaseMiB": 184, + "peakRssMiB": 767, + "retainedHeapIncreaseMiB": -1, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + }, + { + "tools": 900, + "parallel": 16, + "iteration": 3, + "visible": 905, + "queries": 576, + "elapsedMs": 3645, + "peakHeapIncreaseMiB": 187, + "peakRssMiB": 846, + "retainedHeapIncreaseMiB": 4, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + } + ], + "fourConnections": [ + { + "tools": 900, + "connections": 4, + "parallel": 16, + "iteration": 1, + "visible": 905, + "queries": 624, + "elapsedMs": 3721, + "peakHeapIncreaseMiB": 146, + "peakRssMiB": 686, + "retainedHeapIncreaseMiB": 83, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + }, + { + "tools": 900, + "connections": 4, + "parallel": 16, + "iteration": 2, + "visible": 905, + "queries": 624, + "elapsedMs": 3720, + "peakHeapIncreaseMiB": 174, + "peakRssMiB": 807, + "retainedHeapIncreaseMiB": -4, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + }, + { + "tools": 900, + "connections": 4, + "parallel": 16, + "iteration": 3, + "visible": 905, + "queries": 624, + "elapsedMs": 3688, + "peakHeapIncreaseMiB": 171, + "peakRssMiB": 889, + "retainedHeapIncreaseMiB": 1, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + } + ] + } +} diff --git a/doc/mcp-discovery-performance.md b/doc/mcp-discovery-performance.md new file mode 100644 index 00000000000..bd36964b943 --- /dev/null +++ b/doc/mcp-discovery-performance.md @@ -0,0 +1,69 @@ +# MCP discovery resource bounds + +Tool discovery shares policy inputs for one listing and projects context IDs in SQL. +It does not load task descriptions or run results to decide access. The cache ends +with the listing. Tool execution reads current policy and current rate counters. + +The control-plane process admits two whole listings at a time and queues at most +32 more. Excess requests receive HTTP 503 with `tool_discovery_busy`. Disconnected +listing requests leave the queue or stop scheduling new decisions. Already-started +reads drain before the admission slot is released. This applies only to read-only +discovery; a disconnected tool invocation is not automatically cancelled. + +Discovery audit rows retain the visible tool count and a SHA-256 digest of sorted +names. They do not retain the full name list. Named gateway tokens expire normally; +startup and scheduler sweeps delete at most 500 expired tokens per pass using the +expiry index. Tokens with no expiry and unexpired tokens remain available. The +separate access and activity audit records remain available. + +The stateless gateway and runtime-tools MCP endpoints return HTTP 405 with +`Allow: POST` for GET instead of returning JSON as if it were an SSE stream. +The runtime-tools endpoint still validates its caller before that response. +See the [MCP transport contract](https://modelcontextprotocol.io/specification/2025-03-26/basic/transports). + +## Reproduce the measurements + +Build workspace dependencies, then run from `server/`: + +```sh +node --expose-gc --max-old-space-size=4096 --import tsx scripts/benchmark-tool-gateway-listing.ts --tools 50,300,900 --parallel 1,16 --repeat 3 +node --expose-gc --max-old-space-size=4096 --import tsx scripts/benchmark-tool-gateway-listing.ts --tools 900 --connections 4 --parallel 16 --repeat 3 +``` + +The harness uses a throwaway embedded PostgreSQL database. It makes no model or +provider calls. Each fixture has a 400 KB run snapshot, 140 KB result, approximately +5 KB input schemas, 15 KB connection configurations, one selector per catalog tool, +and 373 policies. Use `--implementation /absolute/path/to/tool-gateway.ts` to replay +an earlier implementation, with its corresponding policy-service import. + +## Verification on 2026-10-01 + +Node 26.4.0 on macOS; forced GC before each measurement; heap/RSS sampled every +2 ms. Baseline: `f2e0f196308629ef05c7f65782243e1713b208f0`. +[Raw measurements](benchmarks/2026-10-01-mcp-discovery.json) retain every sample. + +| Scenario | Queries | Peak extra heap | Elapsed time | +| --- | ---: | ---: | ---: | +| Baseline, 900 tools, one listing | 11,489 | 1,680–1,698 MiB | 3.54–3.63 s | +| Fixed, 900 tools, one listing | 36 | 37–37 MiB | 0.24–0.24 s | +| Fixed, 900 tools, 16 listings | 576 total | 169–187 MiB | 3.63–3.65 s | +| Fixed, 900 tools, four connections, 16 listings | 624 total | 146–174 MiB | 3.69–3.72 s | + +Before the listing optimization, the 50-versus-500-tool regression test failed +with 722 versus 6,422 statements. The connection-row duplication test also failed. +The follow-up regressions failed for full run-row discovery reads, abandoned +requests, full name-list audit records, and expired named tokens. They pass with +the fixes. HTTP coverage exercises initialize, GET/SSE rejection, 16 concurrent +500-tool listings, an allowed remote call, policy revocation after discovery, and +the denied subsequent call. The remote provider response is deterministic. + +The eight runnable MCP browser stories pass against a throwaway production UI +and server, including governed execution and approval journeys. + +These are fixture measurements, not a replay of the incident database. RSS +includes module/runtime overhead and can stay high after V8 frees objects. Query +counts depend on connection/profile sets; uncached rate-limit checks add queries. +Catalog payload memory and policy-matching CPU still grow with the catalog. +The measurements do not establish a safe production heap limit or the share of +active installations affected. Verify restart bursts and sustained traffic on the +actual deployment before reducing its temporary heap allowance. diff --git a/packages/db/src/migrations/0294_chilly_marvel_apes.sql b/packages/db/src/migrations/0294_chilly_marvel_apes.sql new file mode 100644 index 00000000000..97866fafb39 --- /dev/null +++ b/packages/db/src/migrations/0294_chilly_marvel_apes.sql @@ -0,0 +1 @@ +CREATE INDEX IF NOT EXISTS "tool_mcp_gateway_tokens_expiry_idx" ON "tool_mcp_gateway_tokens" USING btree ("expires_at","id"); diff --git a/packages/db/src/migrations/meta/0289_snapshot.json b/packages/db/src/migrations/meta/0294_snapshot.json similarity index 98% rename from packages/db/src/migrations/meta/0289_snapshot.json rename to packages/db/src/migrations/meta/0294_snapshot.json index 1b3fb2249ad..c5115c19742 100644 --- a/packages/db/src/migrations/meta/0289_snapshot.json +++ b/packages/db/src/migrations/meta/0294_snapshot.json @@ -1,6 +1,6 @@ { - "id": "3960ae6b-5bf1-4fd7-91db-6a18fc7eaa19", - "prevId": "770b6d59-6346-4c92-8e49-0b71cf4819ab", + "id": "66a0b3c9-91bb-46b8-a868-a6facedbc14b", + "prevId": "d0c40f4e-c2c5-4efd-bc15-c6ae5632047d", "version": "7", "dialect": "postgresql", "tables": { @@ -3975,6 +3975,599 @@ "checkConstraints": {}, "isRLSEnabled": false }, + "public.browser_use_browsers": { + "name": "browser_use_browsers", + "schema": "", + "columns": { + "id": { + "name": "id", + "type": "uuid", + "primaryKey": true, + "notNull": true, + "default": "gen_random_uuid()" + }, + "company_id": { + "name": "company_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "session_id": { + "name": "session_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "provider_browser_id": { + "name": "provider_browser_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "status": { + "name": "status", + "type": "text", + "primaryKey": false, + "notNull": true, + "default": "'active'" + }, + "expires_at": { + "name": "expires_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": { + "browser_use_browsers_provider_uq": { + "name": "browser_use_browsers_provider_uq", + "columns": [ + { + "expression": "provider_browser_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": true, + "concurrently": false, + "method": "btree", + "with": {} + }, + "browser_use_browsers_session_idx": { + "name": "browser_use_browsers_session_idx", + "columns": [ + { + "expression": "session_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": { + "browser_use_browsers_company_id_companies_id_fk": { + "name": "browser_use_browsers_company_id_companies_id_fk", + "tableFrom": "browser_use_browsers", + "tableTo": "companies", + "columnsFrom": [ + "company_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "no action", + "onUpdate": "no action" + }, + "browser_use_browsers_session_id_browser_use_sessions_id_fk": { + "name": "browser_use_browsers_session_id_browser_use_sessions_id_fk", + "tableFrom": "browser_use_browsers", + "tableTo": "browser_use_sessions", + "columnsFrom": [ + "session_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "no action", + "onUpdate": "no action" + } + }, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "public.browser_use_runs": { + "name": "browser_use_runs", + "schema": "", + "columns": { + "id": { + "name": "id", + "type": "uuid", + "primaryKey": true, + "notNull": true, + "default": "gen_random_uuid()" + }, + "company_id": { + "name": "company_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "session_id": { + "name": "session_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "heartbeat_run_id": { + "name": "heartbeat_run_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "invocation_id": { + "name": "invocation_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "provider_run_id": { + "name": "provider_run_id", + "type": "uuid", + "primaryKey": false, + "notNull": false + }, + "recovery_cursor": { + "name": "recovery_cursor", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "status": { + "name": "status", + "type": "text", + "primaryKey": false, + "notNull": true, + "default": "'creating'" + }, + "detached_until": { + "name": "detached_until", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "event_cursor": { + "name": "event_cursor", + "type": "integer", + "primaryKey": false, + "notNull": true, + "default": 0 + }, + "events_drained": { + "name": "events_drained", + "type": "integer", + "primaryKey": false, + "notNull": true, + "default": 0 + }, + "accounted_cents": { + "name": "accounted_cents", + "type": "integer", + "primaryKey": false, + "notNull": true, + "default": 0 + }, + "progress": { + "name": "progress", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "result": { + "name": "result", + "type": "jsonb", + "primaryKey": false, + "notNull": false + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": { + "browser_use_runs_invocation_uq": { + "name": "browser_use_runs_invocation_uq", + "columns": [ + { + "expression": "invocation_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": true, + "concurrently": false, + "method": "btree", + "with": {} + }, + "browser_use_runs_provider_uq": { + "name": "browser_use_runs_provider_uq", + "columns": [ + { + "expression": "provider_run_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": true, + "concurrently": false, + "method": "btree", + "with": {} + }, + "browser_use_runs_session_idx": { + "name": "browser_use_runs_session_idx", + "columns": [ + { + "expression": "session_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": { + "browser_use_runs_company_id_companies_id_fk": { + "name": "browser_use_runs_company_id_companies_id_fk", + "tableFrom": "browser_use_runs", + "tableTo": "companies", + "columnsFrom": [ + "company_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "no action", + "onUpdate": "no action" + }, + "browser_use_runs_session_id_browser_use_sessions_id_fk": { + "name": "browser_use_runs_session_id_browser_use_sessions_id_fk", + "tableFrom": "browser_use_runs", + "tableTo": "browser_use_sessions", + "columnsFrom": [ + "session_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "no action", + "onUpdate": "no action" + }, + "browser_use_runs_heartbeat_run_id_heartbeat_runs_id_fk": { + "name": "browser_use_runs_heartbeat_run_id_heartbeat_runs_id_fk", + "tableFrom": "browser_use_runs", + "tableTo": "heartbeat_runs", + "columnsFrom": [ + "heartbeat_run_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "no action", + "onUpdate": "no action" + } + }, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "public.browser_use_sessions": { + "name": "browser_use_sessions", + "schema": "", + "columns": { + "id": { + "name": "id", + "type": "uuid", + "primaryKey": true, + "notNull": true, + "default": "gen_random_uuid()" + }, + "company_id": { + "name": "company_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "issue_id": { + "name": "issue_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "agent_id": { + "name": "agent_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "connection_id": { + "name": "connection_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "grant_id": { + "name": "grant_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "provider_session_id": { + "name": "provider_session_id", + "type": "uuid", + "primaryKey": false, + "notNull": false + }, + "status": { + "name": "status", + "type": "text", + "primaryKey": false, + "notNull": true, + "default": "'starting'" + }, + "idle_deadline": { + "name": "idle_deadline", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "stop_requested": { + "name": "stop_requested", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "error": { + "name": "error", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "lease_token": { + "name": "lease_token", + "type": "uuid", + "primaryKey": false, + "notNull": false + }, + "lease_until": { + "name": "lease_until", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "next_poll_at": { + "name": "next_poll_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "updated_at": { + "name": "updated_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": { + "browser_use_sessions_task_idx": { + "name": "browser_use_sessions_task_idx", + "columns": [ + { + "expression": "company_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "issue_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + }, + "browser_use_sessions_provider_uq": { + "name": "browser_use_sessions_provider_uq", + "columns": [ + { + "expression": "provider_session_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": true, + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": { + "browser_use_sessions_company_id_companies_id_fk": { + "name": "browser_use_sessions_company_id_companies_id_fk", + "tableFrom": "browser_use_sessions", + "tableTo": "companies", + "columnsFrom": [ + "company_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "no action", + "onUpdate": "no action" + }, + "browser_use_sessions_issue_id_issues_id_fk": { + "name": "browser_use_sessions_issue_id_issues_id_fk", + "tableFrom": "browser_use_sessions", + "tableTo": "issues", + "columnsFrom": [ + "issue_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "no action", + "onUpdate": "no action" + }, + "browser_use_sessions_agent_id_agents_id_fk": { + "name": "browser_use_sessions_agent_id_agents_id_fk", + "tableFrom": "browser_use_sessions", + "tableTo": "agents", + "columnsFrom": [ + "agent_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "no action", + "onUpdate": "no action" + }, + "browser_use_sessions_connection_id_tool_connections_id_fk": { + "name": "browser_use_sessions_connection_id_tool_connections_id_fk", + "tableFrom": "browser_use_sessions", + "tableTo": "tool_connections", + "columnsFrom": [ + "connection_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "no action", + "onUpdate": "no action" + }, + "browser_use_sessions_grant_id_connection_grants_id_fk": { + "name": "browser_use_sessions_grant_id_connection_grants_id_fk", + "tableFrom": "browser_use_sessions", + "tableTo": "connection_grants", + "columnsFrom": [ + "grant_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "no action", + "onUpdate": "no action" + } + }, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "public.browser_use_settings": { + "name": "browser_use_settings", + "schema": "", + "columns": { + "grant_id": { + "name": "grant_id", + "type": "uuid", + "primaryKey": true, + "notNull": true + }, + "company_id": { + "name": "company_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "allowed_profile_ids": { + "name": "allowed_profile_ids", + "type": "jsonb", + "primaryKey": false, + "notNull": true, + "default": "'[]'::jsonb" + }, + "max_cost_usd": { + "name": "max_cost_usd", + "type": "text", + "primaryKey": false, + "notNull": false + } + }, + "indexes": {}, + "foreignKeys": { + "browser_use_settings_grant_id_connection_grants_id_fk": { + "name": "browser_use_settings_grant_id_connection_grants_id_fk", + "tableFrom": "browser_use_settings", + "tableTo": "connection_grants", + "columnsFrom": [ + "grant_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "no action", + "onUpdate": "no action" + }, + "browser_use_settings_company_id_companies_id_fk": { + "name": "browser_use_settings_company_id_companies_id_fk", + "tableFrom": "browser_use_settings", + "tableTo": "companies", + "columnsFrom": [ + "company_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "no action", + "onUpdate": "no action" + } + }, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, "public.budget_incidents": { "name": "budget_incidents", "schema": "", @@ -12144,6 +12737,373 @@ "checkConstraints": {}, "isRLSEnabled": false }, + "public.company_skill_source_entries": { + "name": "company_skill_source_entries", + "schema": "", + "columns": { + "id": { + "name": "id", + "type": "uuid", + "primaryKey": true, + "notNull": true, + "default": "gen_random_uuid()" + }, + "company_id": { + "name": "company_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "source_id": { + "name": "source_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "path": { + "name": "path", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "name": { + "name": "name", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "description": { + "name": "description", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "inspection": { + "name": "inspection", + "type": "jsonb", + "primaryKey": false, + "notNull": false + }, + "skill_id": { + "name": "skill_id", + "type": "uuid", + "primaryKey": false, + "notNull": false + }, + "selection": { + "name": "selection", + "type": "text", + "primaryKey": false, + "notNull": true, + "default": "'new'" + }, + "present": { + "name": "present", + "type": "boolean", + "primaryKey": false, + "notNull": true, + "default": true + }, + "error": { + "name": "error", + "type": "text", + "primaryKey": false, + "notNull": false + } + }, + "indexes": { + "company_skill_source_entries_source_path_idx": { + "name": "company_skill_source_entries_source_path_idx", + "columns": [ + { + "expression": "source_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "path", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": true, + "concurrently": false, + "method": "btree", + "with": {} + }, + "company_skill_source_entries_company_skill_idx": { + "name": "company_skill_source_entries_company_skill_idx", + "columns": [ + { + "expression": "company_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "skill_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": { + "company_skill_source_entries_company_id_companies_id_fk": { + "name": "company_skill_source_entries_company_id_companies_id_fk", + "tableFrom": "company_skill_source_entries", + "tableTo": "companies", + "columnsFrom": [ + "company_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "cascade", + "onUpdate": "no action" + }, + "company_skill_source_entries_source_id_company_skill_sources_id_fk": { + "name": "company_skill_source_entries_source_id_company_skill_sources_id_fk", + "tableFrom": "company_skill_source_entries", + "tableTo": "company_skill_sources", + "columnsFrom": [ + "source_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "cascade", + "onUpdate": "no action" + }, + "company_skill_source_entries_skill_id_company_skills_id_fk": { + "name": "company_skill_source_entries_skill_id_company_skills_id_fk", + "tableFrom": "company_skill_source_entries", + "tableTo": "company_skills", + "columnsFrom": [ + "skill_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "set null", + "onUpdate": "no action" + } + }, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "public.company_skill_sources": { + "name": "company_skill_sources", + "schema": "", + "columns": { + "id": { + "name": "id", + "type": "uuid", + "primaryKey": true, + "notNull": true, + "default": "gen_random_uuid()" + }, + "company_id": { + "name": "company_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "repository_id": { + "name": "repository_id", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "repository_url": { + "name": "repository_url", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "full_name": { + "name": "full_name", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "tracking_ref": { + "name": "tracking_ref", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "connection_id": { + "name": "connection_id", + "type": "uuid", + "primaryKey": false, + "notNull": false + }, + "excluded_folders": { + "name": "excluded_folders", + "type": "jsonb", + "primaryKey": false, + "notNull": true, + "default": "'[]'::jsonb" + }, + "enabled": { + "name": "enabled", + "type": "boolean", + "primaryKey": false, + "notNull": true, + "default": true + }, + "revision": { + "name": "revision", + "type": "integer", + "primaryKey": false, + "notNull": true, + "default": 0 + }, + "last_attempt_at": { + "name": "last_attempt_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "last_success_at": { + "name": "last_success_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "last_scan_commit": { + "name": "last_scan_commit", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "last_error": { + "name": "last_error", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "lease_token": { + "name": "lease_token", + "type": "uuid", + "primaryKey": false, + "notNull": false + }, + "lease_expires_at": { + "name": "lease_expires_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": { + "company_skill_sources_repository_ref_idx": { + "name": "company_skill_sources_repository_ref_idx", + "columns": [ + { + "expression": "company_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "repository_url", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "tracking_ref", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": true, + "concurrently": false, + "method": "btree", + "with": {} + }, + "company_skill_sources_identity_ref_idx": { + "name": "company_skill_sources_identity_ref_idx", + "columns": [ + { + "expression": "company_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "repository_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "tracking_ref", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": true, + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": { + "company_skill_sources_company_id_companies_id_fk": { + "name": "company_skill_sources_company_id_companies_id_fk", + "tableFrom": "company_skill_sources", + "tableTo": "companies", + "columnsFrom": [ + "company_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "cascade", + "onUpdate": "no action" + }, + "company_skill_sources_connection_id_tool_connections_id_fk": { + "name": "company_skill_sources_connection_id_tool_connections_id_fk", + "tableFrom": "company_skill_sources", + "tableTo": "tool_connections", + "columnsFrom": [ + "connection_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "set null", + "onUpdate": "no action" + } + }, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, "public.company_skill_comments": { "name": "company_skill_comments", "schema": "", @@ -30505,6 +31465,13 @@ "primaryKey": false, "notNull": true }, + "title_needs_generation": { + "name": "title_needs_generation", + "type": "boolean", + "primaryKey": false, + "notNull": true, + "default": false + }, "description": { "name": "description", "type": "text", @@ -45242,6 +46209,27 @@ "method": "btree", "with": {} }, + "tool_mcp_gateway_tokens_expiry_idx": { + "name": "tool_mcp_gateway_tokens_expiry_idx", + "columns": [ + { + "expression": "expires_at", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + }, "tool_mcp_gateway_tokens_gateway_idx": { "name": "tool_mcp_gateway_tokens_gateway_idx", "columns": [ diff --git a/packages/db/src/migrations/meta/_journal.json b/packages/db/src/migrations/meta/_journal.json index 71acfdde0c6..15430d35565 100644 --- a/packages/db/src/migrations/meta/_journal.json +++ b/packages/db/src/migrations/meta/_journal.json @@ -2045,6 +2045,13 @@ "when": 1790796608171, "tag": "0293_broad_rattler", "breakpoints": true + }, + { + "idx": 294, + "version": "7", + "when": 1790891286527, + "tag": "0294_chilly_marvel_apes", + "breakpoints": true } ] } \ No newline at end of file diff --git a/packages/db/src/schema/tool_access.ts b/packages/db/src/schema/tool_access.ts index 54227d932ee..19b93987e03 100644 --- a/packages/db/src/schema/tool_access.ts +++ b/packages/db/src/schema/tool_access.ts @@ -563,6 +563,7 @@ export const toolMcpGatewayTokens = pgTable( }, (table) => [ uniqueIndex("tool_mcp_gateway_tokens_token_hash_uq").on(table.tokenHash), + index("tool_mcp_gateway_tokens_expiry_idx").on(table.expiresAt, table.id), index("tool_mcp_gateway_tokens_gateway_idx").on(table.companyId, table.gatewayId), index("tool_mcp_gateway_tokens_subject_idx").on(table.companyId, table.subjectType, table.subjectId), index("tool_mcp_gateway_tokens_company_expires_idx").on(table.companyId, table.expiresAt), diff --git a/server/scripts/benchmark-tool-gateway-listing.ts b/server/scripts/benchmark-tool-gateway-listing.ts new file mode 100644 index 00000000000..5d460a6cf5a --- /dev/null +++ b/server/scripts/benchmark-tool-gateway-listing.ts @@ -0,0 +1,95 @@ +/** + * Reproducible discovery load probe. Uses a throwaway embedded PostgreSQL DB; + * no provider calls, production data, or model inference. + * node --expose-gc --import tsx scripts/benchmark-tool-gateway-listing.ts + * --tools 50,300,900 --parallel 1,16 --repeat 3 + * --implementation accepts a local tool-gateway module for before/after replay. + */ +import { randomUUID } from "node:crypto"; +import { pathToFileURL } from "node:url"; +import { eq } from "drizzle-orm"; +import { + createDb, heartbeatRuns, toolCatalogEntries, toolConnections, toolPolicies, + toolProfileEntries, startEmbeddedPostgresTestDatabase, +} from "@paperclipai/db"; +import { createListingFixture, recordingDb } from "../src/__tests__/helpers/tool-gateway-listing-fixture.js"; +import type { createToolGatewayService } from "../src/services/tool-gateway.js"; + +function option(name: string, fallback: string) { + const index = process.argv.indexOf(name); + return index < 0 ? fallback : process.argv[index + 1] ?? fallback; +} +function counts(value: string) { + const result = value.split(",").map(Number); + if (result.some((n) => !Number.isInteger(n) || n < 1)) throw new Error("Expected positive integer counts"); + return result; +} +const toolCounts = counts(option("--tools", "50,300,900")); +const parallels = counts(option("--parallel", "1,16")); +const repeats = counts(option("--repeat", "3"))[0]!; +const connections = counts(option("--connections", "1"))[0]!; +const implementation = option("--implementation", ""); +const moduleUrl = implementation ? pathToFileURL(implementation) : new URL("../src/services/tool-gateway.js", import.meta.url); +const factory = (await import(moduleUrl.href)).createToolGatewayService as typeof createToolGatewayService; +const temp = await startEmbeddedPostgresTestDatabase("paperclip-listing-benchmark-"); +const db = createDb(temp.connectionString); +try { + for (const tools of toolCounts) { + const fixture = await createListingFixture(db, tools, { connectionCount: connections }); + await db.update(heartbeatRuns).set({ + contextSnapshot: { issueId: fixture.issue.id, projectId: fixture.project.id, taskMarkdown: "x".repeat(400_000) }, + resultJson: { summary: "x".repeat(140_000) }, + }).where(eq(heartbeatRuns.id, fixture.run.id)); + await db.update(toolConnections).set({ config: { url: "https://8.8.8.8/mcp", notes: "x".repeat(15_000) } }) + .where(eq(toolConnections.companyId, fixture.company.id)); + await db.update(toolCatalogEntries).set({ inputSchema: { + type: "object", properties: { query: { type: "string", description: "x".repeat(5_000) } }, + } }).where(eq(toolCatalogEntries.companyId, fixture.company.id)); + await db.insert(toolProfileEntries).values(fixture.entries.map((entry) => ({ + companyId: fixture.company.id, profileId: fixture.namedGateway.profileId, + selectorType: "catalog_entry" as const, effect: "include" as const, catalogEntryId: entry.id, + }))); + await db.insert(toolPolicies).values(Array.from({ length: 371 }, (_, i) => ({ + companyId: fixture.company.id, name: `Unmatched fixture policy ${i}`, policyType: "block" as const, + selectors: { catalogEntryId: randomUUID() }, priority: 100 + i, + }))); + const input = { gatewayId: fixture.namedGateway.id, bearerToken: fixture.token.token }; + await factory(db).listToolsForNamedGateway(input); + for (const parallel of parallels) { + for (let iteration = 1; iteration <= repeats; iteration += 1) { + global.gc?.(); + const before = process.memoryUsage(); + let peakHeap = before.heapUsed; + let peakRss = before.rss; + const sample = () => { + const memory = process.memoryUsage(); + peakHeap = Math.max(peakHeap, memory.heapUsed); + peakRss = Math.max(peakRss, memory.rss); + }; + const timer = setInterval(sample, 2); + const recorder = recordingDb(db); + const started = performance.now(); + let visible = 0; + try { + const listings = await Promise.all(Array.from({ length: parallel }, () => + factory(recorder.db).listToolsForNamedGateway(input))); + visible = listings[0]!.length; + sample(); + } finally { + clearInterval(timer); + } + const elapsedMs = performance.now() - started; + global.gc?.(); + const after = process.memoryUsage(); + console.log(JSON.stringify({ + tools, connections, parallel, iteration, visible, queries: recorder.statements.length, + elapsedMs: Math.round(elapsedMs), peakHeapIncreaseMiB: Math.round((peakHeap - before.heapUsed) / 1_048_576), + peakRssMiB: Math.round(peakRss / 1_048_576), retainedHeapIncreaseMiB: Math.round((after.heapUsed - before.heapUsed) / 1_048_576), + forcedGc: Boolean(global.gc), fixture: { snapshotBytes: 400_000, schemaBytes: 5_000, policyCount: 373 }, + })); + } + } + } +} finally { + await temp.cleanup(); +} diff --git a/server/src/__tests__/helpers/tool-gateway-listing-fixture.ts b/server/src/__tests__/helpers/tool-gateway-listing-fixture.ts new file mode 100644 index 00000000000..643e72334a8 --- /dev/null +++ b/server/src/__tests__/helpers/tool-gateway-listing-fixture.ts @@ -0,0 +1,190 @@ +import { inArray } from "drizzle-orm"; +import { randomUUID } from "node:crypto"; +import { drizzle } from "drizzle-orm/postgres-js"; +import { + agents, companies, createDb, heartbeatRuns, issues, principalPermissionGrants, + projects, toolApplications, toolCatalogEntries, toolConnections, toolPolicies, + toolProfileBindings, toolProfileEntries, toolProfiles, +} from "@paperclipai/db"; +import { createToolGatewayService } from "../../services/tool-gateway.js"; + +type Db = ReturnType; + +// Real run snapshots carry the task text several times over; a policy check +// must read only the context ids out of it. +const LARGE_TASK_TEXT = "Task description. ".repeat(12_000); + +/** A second handle on the same connection pool that records every statement. */ +export function recordingDb(db: Db) { + const statements: string[] = []; + const statementParams: unknown[][] = []; + const recorded = drizzle(db.$client, { + schema: db._.fullSchema, + logger: { + logQuery: (query: string, params: unknown[]) => { + statements.push(query); + statementParams.push(params); + }, + }, + }) as unknown as Db; + return { db: recorded, statements, statementParams }; +} + +export async function createListingFixture( + db: Db, + toolCount: number, + options: { broadBindings?: boolean; connectionCount?: number } = {}, +) { + const company = await db.insert(companies).values({ + name: `Listing ${randomUUID()}`, + issuePrefix: `LM${randomUUID().slice(0, 6).toUpperCase()}`, + }).returning().then((rows) => rows[0]!); + const agent = await db.insert(agents).values({ + companyId: company.id, + name: `Listing Agent ${randomUUID()}`, + role: "engineer", + adapterType: "process", + adapterConfig: {}, + runtimeConfig: {}, + permissions: {}, + }).returning().then((rows) => rows[0]!); + const project = await db.insert(projects).values({ + companyId: company.id, + name: `Listing Project ${randomUUID()}`, + }).returning().then((rows) => rows[0]!); + const issue = await db.insert(issues).values({ + companyId: company.id, + projectId: project.id, + title: "Listing work", + status: "in_progress", + assigneeAgentId: agent.id, + }).returning().then((rows) => rows[0]!); + const run = await db.insert(heartbeatRuns).values({ + companyId: company.id, + agentId: agent.id, + invocationSource: "assignment", + status: "running", + contextSnapshot: { issueId: issue.id, projectId: project.id, taskMarkdown: LARGE_TASK_TEXT }, + resultJson: { summary: LARGE_TASK_TEXT }, + }).returning().then((rows) => rows[0]!); + const application = await db.insert(toolApplications).values({ + companyId: company.id, + applicationKey: `listing-${randomUUID().slice(0, 8)}`, + name: `Listing MCP ${randomUUID()}`, + type: "mcp_http", + status: "active", + }).returning().then((rows) => rows[0]!); + const connection = await db.insert(toolConnections).values({ + companyId: company.id, + applicationId: application.id, + name: "Listing connection", + uid: `test/${randomUUID()}`, + transport: "mcp_remote", + status: "active", + enabled: true, + healthStatus: "ok", + credentialPolicy: "shared", + config: { url: "https://8.8.8.8/mcp", notes: "connection config ".repeat(500) }, + }).returning().then((rows) => rows[0]!); + const entries = await db.insert(toolCatalogEntries).values( + Array.from({ length: toolCount }, (_, index) => ({ + companyId: company.id, + applicationId: application.id, + connectionId: connection.id, + entryKind: "tool" as const, + name: `tool_${String(index).padStart(4, "0")}`, + toolName: `tool_${String(index).padStart(4, "0")}`, + description: `Fixture tool ${index}`, + inputSchema: { + type: "object", + properties: { query: { type: "string", description: "Query text. ".repeat(40) } }, + }, + riskLevel: index % 3 === 0 ? "write" as const : "read" as const, + isReadOnly: index % 3 !== 0, + isWrite: index % 3 === 0, + status: "active" as const, + versionHash: randomUUID(), + schemaHash: randomUUID(), + })), + ).returning(); + + for (let i = 1; i < (options.connectionCount ?? 1); i += 1) { + const extra = await db.insert(toolConnections).values({ + companyId: company.id, applicationId: application.id, name: `Listing connection ${i}`, + uid: `test/${randomUUID()}`, transport: "mcp_remote", status: "active", enabled: true, + healthStatus: "ok", config: { url: "https://8.8.8.8/mcp" }, + }).returning().then((rows) => rows[0]!); + const ids = entries.filter((_, index) => index % (options.connectionCount ?? 1) === i).map((entry) => entry.id); + if (ids.length > 0) await db.update(toolCatalogEntries).set({ connectionId: extra.id }).where(inArray(toolCatalogEntries.id, ids)); + } + + // The gateway profile allows every tool except one excluded entry. The + // company and agent bindings lose to the narrower gateway binding. + const gatewayProfile = await db.insert(toolProfiles).values({ + companyId: company.id, + profileKey: `gateway-${randomUUID()}`, + name: `Gateway profile ${randomUUID()}`, + defaultAction: "allow", + }).returning().then((rows) => rows[0]!); + await db.insert(toolProfileEntries).values({ + companyId: company.id, + profileId: gatewayProfile.id, + selectorType: "catalog_entry", + catalogEntryId: entries[1]!.id, + effect: "exclude", + }); + const broadProfile = await db.insert(toolProfiles).values({ + companyId: company.id, + profileKey: `broad-${randomUUID()}`, + name: `Broad profile ${randomUUID()}`, + defaultAction: "deny", + }).returning().then((rows) => rows[0]!); + if (options.broadBindings !== false) { + await db.insert(toolProfileBindings).values([ + { companyId: company.id, profileId: broadProfile.id, targetType: "company", targetId: company.id }, + { companyId: company.id, profileId: broadProfile.id, targetType: "agent", targetId: agent.id }, + ]); + } + await db.insert(toolPolicies).values([ + { + companyId: company.id, + name: `Block one tool ${randomUUID()}`, + policyType: "block", + priority: 10, + selectors: { catalogEntryId: entries[2]!.id }, + }, + { + companyId: company.id, + name: `Review writes ${randomUUID()}`, + policyType: "require_approval", + priority: 20, + selectors: { riskLevel: "write" }, + }, + ]); + // A grant for another connection: the listing reads it, but it allows nothing here. + await db.insert(principalPermissionGrants).values({ + companyId: company.id, + principalType: "agent", + principalId: agent.id, + permissionKey: "tools:use", + scope: { connectionId: randomUUID() }, + }); + + const setupGateway = createToolGatewayService(db); + const namedGateway = await setupGateway.createNamedGateway({ + companyId: company.id, + body: { name: `Listing gateway ${randomUUID().slice(0, 8)}`, profileId: gatewayProfile.id }, + }); + const token = await setupGateway.createNamedGatewayToken({ + companyId: company.id, + gatewayId: namedGateway.id, + body: { + name: "Run token", + subjectType: "heartbeat_run", + subjectId: run.id, + clientLabel: "codex", + ownerNote: "", + }, + }); + return { company, agent, project, issue, run, application, connection, entries, namedGateway, token }; +} diff --git a/server/src/__tests__/server-startup-feedback-export.test.ts b/server/src/__tests__/server-startup-feedback-export.test.ts index d5e7bf7853f..f6f8e20dbb3 100644 --- a/server/src/__tests__/server-startup-feedback-export.test.ts +++ b/server/src/__tests__/server-startup-feedback-export.test.ts @@ -37,7 +37,10 @@ const { const completionSweepMock = vi.fn(async () => undefined); const createAppMock = vi.fn(async () => Object.assign((_: unknown, __: unknown) => {}, { locals: { - toolGateway: { sweepActionReviews: vi.fn(async () => ({ scanned: 0 })) }, + toolGateway: { + sweepActionReviews: vi.fn(async () => ({ scanned: 0 })), + cleanupExpiredSessions: vi.fn(async () => ({ deletedCount: 0 })), + }, toolActionDeliveries: { sweepPending: vi.fn(async () => ({ scanned: 0, delivered: 0 })) }, }, }) as never); diff --git a/server/src/__tests__/tool-gateway-discovery-http.test.ts b/server/src/__tests__/tool-gateway-discovery-http.test.ts new file mode 100644 index 00000000000..14429cdb940 --- /dev/null +++ b/server/src/__tests__/tool-gateway-discovery-http.test.ts @@ -0,0 +1,66 @@ +import { randomUUID } from "node:crypto"; +import express from "express"; +import request from "supertest"; +import { afterAll, beforeAll, describe, expect, it } from "vitest"; +import { connectionGrants, createDb, toolPolicies, startEmbeddedPostgresTestDatabase, getEmbeddedPostgresTestSupport } from "@paperclipai/db"; +import { createToolGatewayService } from "../services/tool-gateway.js"; +import { mcpGatewayProtocolRoutes } from "../routes/tool-gateway.js"; +import { createListingFixture } from "./helpers/tool-gateway-listing-fixture.js"; + +const support = await getEmbeddedPostgresTestSupport(); +const suite = support.supported ? describe : describe.skip; +suite("MCP discovery over HTTP", () => { + let db!: ReturnType; + let temp!: Awaited>; + beforeAll(async () => { + temp = await startEmbeddedPostgresTestDatabase("paperclip-discovery-http-"); + db = createDb(temp.connectionString); + }); + afterAll(async () => { await temp?.cleanup(); }); + + it("discovers a large catalog concurrently, calls a tool, and rechecks changed policy", async () => { + const fixture = await createListingFixture(db, 500); + await db.insert(connectionGrants).values({ + companyId: fixture.company.id, connectionId: fixture.connection.id, + kind: "organization", status: "active", isDefault: true, + }); + const calls: string[] = []; + const gateway = createToolGatewayService(db, { + remoteHttpRequest: async (_url, init) => { + const body = JSON.parse(String(init.body)); + if (body.method === "tools/call") calls.push(body.params.name); + return new Response(JSON.stringify({ jsonrpc: "2.0", id: body.id, + result: body.method === "initialize" + ? { protocolVersion: "2025-03-26", capabilities: {}, serverInfo: { name: "fixture", version: "1" } } + : { content: [{ type: "text", text: "fixture result" }] }, + }), { headers: { "content-type": "application/json" } }); + }, + }); + const app = express().use(express.json()).use(mcpGatewayProtocolRoutes(gateway)); + const url = `/mcp/gateways/${fixture.namedGateway.gatewayPublicId}`; + const post = (method: string, params?: unknown) => request(app).post(url) + .set("Authorization", `Bearer ${fixture.token.token}`) + .send({ jsonrpc: "2.0", id: 1, method, params }); + await post("initialize").expect(200); + await request(app).get(url).set("Accept", "text/event-stream").expect("Allow", "POST").expect(405); + const listings = await Promise.all(Array.from({ length: 16 }, () => post("tools/list").expect(200))); + const tools = listings[0]!.body.result.tools as Array<{ name: string; description: string }>; + expect(listings.every((response) => JSON.stringify(response.body.result.tools) === JSON.stringify(tools))).toBe(true); + const descriptors = await gateway.listToolsForNamedGateway({ gatewayId: fixture.namedGateway.id, bearerToken: fixture.token.token }); + expect(descriptors.filter((tool) => tool.connectionId === fixture.connection.id)).toHaveLength(498); + const readName = descriptors.find((tool) => tool.catalogEntryId === fixture.entries[4]!.id)!.name; + const read = tools.find((tool) => tool.name === readName)!; + expect(read).toBeDefined(); + const called = await post("tools/call", { name: read.name, arguments: {} }); + expect(called.status, JSON.stringify(called.body)).toBe(200); + expect(calls).toEqual(["tool_0004"]); + await db.insert(toolPolicies).values({ + companyId: fixture.company.id, name: `Revoke ${randomUUID()}`, policyType: "block", priority: 0, + selectors: { catalogEntryId: fixture.entries[4]!.id }, + }); + await post("tools/call", { name: read.name, arguments: {} }).expect(403); + expect(calls).toHaveLength(1); + const after = await post("tools/list").expect(200); + expect(after.body.result.tools.some((tool: { name: string }) => tool.name === read.name)).toBe(false); + }); +}); diff --git a/server/src/__tests__/tool-gateway-listing-memory.test.ts b/server/src/__tests__/tool-gateway-listing-memory.test.ts index 309eef2e1ab..bc98242e1d9 100644 --- a/server/src/__tests__/tool-gateway-listing-memory.test.ts +++ b/server/src/__tests__/tool-gateway-listing-memory.test.ts @@ -1,7 +1,7 @@ import { randomUUID } from "node:crypto"; -import { drizzle } from "drizzle-orm/postgres-js"; import { afterAll, beforeAll, describe, expect, it } from "vitest"; import { + activityLog, agents, companies, createDb, @@ -16,7 +16,9 @@ import { toolProfileBindings, toolProfileEntries, toolProfiles, + toolMcpGatewayTokens, } from "@paperclipai/db"; +import { eq } from "drizzle-orm"; import type { ToolAccessDecisionInput } from "@paperclipai/shared"; import { createToolAccessDecisionCache, @@ -28,180 +30,13 @@ import { startEmbeddedPostgresTestDatabase, } from "./helpers/embedded-postgres.js"; +import { createListingFixture, recordingDb } from "./helpers/tool-gateway-listing-fixture.js"; + const embeddedPostgresSupport = await getEmbeddedPostgresTestSupport(); const describeEmbeddedPostgres = embeddedPostgresSupport.supported ? describe : describe.skip; type Db = ReturnType; -// Real run snapshots carry the task text several times over; a policy check -// must read only the context ids out of it. -const LARGE_TASK_TEXT = "Task description. ".repeat(12_000); - -/** A second handle on the same connection pool that records every statement. */ -function recordingDb(db: Db) { - const statements: string[] = []; - const statementParams: unknown[][] = []; - const recorded = drizzle(db.$client, { - schema: db._.fullSchema, - logger: { - logQuery: (query: string, params: unknown[]) => { - statements.push(query); - statementParams.push(params); - }, - }, - }) as unknown as Db; - return { db: recorded, statements, statementParams }; -} - -async function createListingFixture( - db: Db, - toolCount: number, - options: { broadBindings?: boolean } = {}, -) { - const company = await db.insert(companies).values({ - name: `Listing ${randomUUID()}`, - issuePrefix: `LM${randomUUID().slice(0, 6).toUpperCase()}`, - }).returning().then((rows) => rows[0]!); - const agent = await db.insert(agents).values({ - companyId: company.id, - name: `Listing Agent ${randomUUID()}`, - role: "engineer", - adapterType: "process", - adapterConfig: {}, - runtimeConfig: {}, - permissions: {}, - }).returning().then((rows) => rows[0]!); - const project = await db.insert(projects).values({ - companyId: company.id, - name: `Listing Project ${randomUUID()}`, - }).returning().then((rows) => rows[0]!); - const issue = await db.insert(issues).values({ - companyId: company.id, - projectId: project.id, - title: "Listing work", - status: "in_progress", - assigneeAgentId: agent.id, - }).returning().then((rows) => rows[0]!); - const run = await db.insert(heartbeatRuns).values({ - companyId: company.id, - agentId: agent.id, - invocationSource: "assignment", - status: "running", - contextSnapshot: { issueId: issue.id, projectId: project.id, taskMarkdown: LARGE_TASK_TEXT }, - resultJson: { summary: LARGE_TASK_TEXT }, - }).returning().then((rows) => rows[0]!); - const application = await db.insert(toolApplications).values({ - companyId: company.id, - applicationKey: `listing-${randomUUID().slice(0, 8)}`, - name: `Listing MCP ${randomUUID()}`, - type: "mcp_http", - status: "active", - }).returning().then((rows) => rows[0]!); - const connection = await db.insert(toolConnections).values({ - companyId: company.id, - applicationId: application.id, - name: "Listing connection", - uid: `test/${randomUUID()}`, - transport: "mcp_remote", - status: "active", - enabled: true, - healthStatus: "ok", - credentialPolicy: "shared", - config: { url: "https://8.8.8.8/mcp", notes: "connection config ".repeat(500) }, - }).returning().then((rows) => rows[0]!); - const entries = await db.insert(toolCatalogEntries).values( - Array.from({ length: toolCount }, (_, index) => ({ - companyId: company.id, - applicationId: application.id, - connectionId: connection.id, - entryKind: "tool" as const, - name: `tool_${String(index).padStart(4, "0")}`, - toolName: `tool_${String(index).padStart(4, "0")}`, - description: `Fixture tool ${index}`, - inputSchema: { - type: "object", - properties: { query: { type: "string", description: "Query text. ".repeat(40) } }, - }, - riskLevel: index % 3 === 0 ? "write" as const : "read" as const, - isReadOnly: index % 3 !== 0, - isWrite: index % 3 === 0, - status: "active" as const, - versionHash: randomUUID(), - schemaHash: randomUUID(), - })), - ).returning(); - - // The gateway profile allows every tool except one excluded entry. The - // company and agent bindings lose to the narrower gateway binding. - const gatewayProfile = await db.insert(toolProfiles).values({ - companyId: company.id, - profileKey: `gateway-${randomUUID()}`, - name: `Gateway profile ${randomUUID()}`, - defaultAction: "allow", - }).returning().then((rows) => rows[0]!); - await db.insert(toolProfileEntries).values({ - companyId: company.id, - profileId: gatewayProfile.id, - selectorType: "catalog_entry", - catalogEntryId: entries[1]!.id, - effect: "exclude", - }); - const broadProfile = await db.insert(toolProfiles).values({ - companyId: company.id, - profileKey: `broad-${randomUUID()}`, - name: `Broad profile ${randomUUID()}`, - defaultAction: "deny", - }).returning().then((rows) => rows[0]!); - if (options.broadBindings !== false) { - await db.insert(toolProfileBindings).values([ - { companyId: company.id, profileId: broadProfile.id, targetType: "company", targetId: company.id }, - { companyId: company.id, profileId: broadProfile.id, targetType: "agent", targetId: agent.id }, - ]); - } - await db.insert(toolPolicies).values([ - { - companyId: company.id, - name: `Block one tool ${randomUUID()}`, - policyType: "block", - priority: 10, - selectors: { catalogEntryId: entries[2]!.id }, - }, - { - companyId: company.id, - name: `Review writes ${randomUUID()}`, - policyType: "require_approval", - priority: 20, - selectors: { riskLevel: "write" }, - }, - ]); - // A grant for another connection: the listing reads it, but it allows nothing here. - await db.insert(principalPermissionGrants).values({ - companyId: company.id, - principalType: "agent", - principalId: agent.id, - permissionKey: "tools:use", - scope: { connectionId: randomUUID() }, - }); - - const setupGateway = createToolGatewayService(db); - const namedGateway = await setupGateway.createNamedGateway({ - companyId: company.id, - body: { name: `Listing gateway ${randomUUID().slice(0, 8)}`, profileId: gatewayProfile.id }, - }); - const token = await setupGateway.createNamedGatewayToken({ - companyId: company.id, - gatewayId: namedGateway.id, - body: { - name: "Run token", - subjectType: "heartbeat_run", - subjectId: run.id, - clientLabel: "codex", - ownerNote: "", - }, - }); - return { company, agent, project, issue, run, application, connection, entries, namedGateway, token }; -} - describeEmbeddedPostgres("tool gateway listing memory", () => { let db!: Db; let tempDb: Awaited> | null = null; @@ -227,6 +62,63 @@ describeEmbeddedPostgres("tool gateway listing memory", () => { return { tools, statements: recorder.statements, statementParams: recorder.statementParams }; } + it("does not load task text when authenticating a named gateway listing", async () => { + const fixture = await createListingFixture(db, 3); + const { statements } = await measureNamedGatewayListing(fixture); + const runReads = statements.filter((query) => query.includes('from "heartbeat_runs"')); + expect(runReads.length).toBeGreaterThan(0); + expect(runReads.filter((query) => /"context_snapshot"\s*(?:,|from\b)/.test(query))).toEqual([]); + }); + + it("stops an abandoned listing before catalog reads", async () => { + const fixture = await createListingFixture(db, 3); + const recorder = recordingDb(db); + const controller = new AbortController(); + controller.abort(); + await expect(createToolGatewayService(recorder.db).listToolsForNamedGateway({ + gatewayId: fixture.namedGateway.id, bearerToken: fixture.token.token, + signal: controller.signal, + })).rejects.toMatchObject({ name: "AbortError" }); + expect(recorder.statements.some((query) => query.includes('from "tool_catalog_entries"'))).toBe(false); + }); + + it("records a fixed-size discovery summary instead of all tool names", async () => { + const fixture = await createListingFixture(db, 50); + const { tools } = await measureNamedGatewayListing(fixture); + const events = await db.select().from(activityLog).where(eq(activityLog.runId, fixture.run.id)); + const event = events.find((row) => row.action === "tool_gateway.discovery"); + expect(event?.details).toMatchObject({ + visibleToolCount: tools.length, visibleToolsHash: expect.stringMatching(/^[a-f0-9]{64}$/), + }); + expect(event?.details).not.toHaveProperty("visibleTools"); + }); + + it("keeps query growth bounded with several connections", async () => { + const small = await createListingFixture(db, 50, { connectionCount: 4 }); + const large = await createListingFixture(db, 500, { connectionCount: 4 }); + const smallListing = await measureNamedGatewayListing(small); + const largeListing = await measureNamedGatewayListing(large); + expect(largeListing.statements.length).toBe(smallListing.statements.length); + expect(largeListing.statements.length).toBeLessThan(100); + expect(largeListing.tools.filter((tool) => tool.catalogEntryId)).toHaveLength(498); + }); + + it("cleans expired named tokens in bounded batches and retains valid tokens", async () => { + const fixture = await createListingFixture(db, 3); + const service = createToolGatewayService(db); + const expired = await service.createNamedGatewayToken({ + companyId: fixture.company.id, gatewayId: fixture.namedGateway.id, + body: { name: "Expired test token", clientLabel: "test", ownerNote: "test" }, + }); + const now = new Date(); + await db.update(toolMcpGatewayTokens).set({ expiresAt: new Date(now.getTime() - 1) }) + .where(eq(toolMcpGatewayTokens.id, expired.id)); + await service.cleanupExpiredSessions({ now }); + expect(await db.select().from(toolMcpGatewayTokens).where(eq(toolMcpGatewayTokens.id, expired.id))).toEqual([]); + expect(await db.select().from(toolMcpGatewayTokens).where(eq(toolMcpGatewayTokens.id, fixture.token.id))).toHaveLength(1); + await expect(service.cleanupExpiredSessions({ now })).resolves.toMatchObject({ deletedCount: 0 }); + }); + it("keeps the query count of a named gateway listing constant from 50 to 500 catalog tools", async () => { const small = await createListingFixture(db, 50); const large = await createListingFixture(db, 500); diff --git a/server/src/index.ts b/server/src/index.ts index 4b917f08a28..92c57e4b95b 100644 --- a/server/src/index.ts +++ b/server/src/index.ts @@ -1293,6 +1293,7 @@ async function startServerWithDatabaseTeardown( await chatCompletionDeliveries.sweepPending().catch((err) => logger.error({ err }, "startup chat completion delivery recovery failed")); await connectionDeliveries.sweepPending(); await app.locals.toolGateway.sweepActionReviews().catch((err: unknown) => logger.error({ err }, "startup tool review recovery failed")); + await app.locals.toolGateway.cleanupExpiredSessions().catch((err: unknown) => logger.error({ err }, "startup gateway token cleanup failed")); await app.locals.toolActionDeliveries.sweepPending().catch((err: unknown) => logger.error({ err }, "startup tool review delivery sweep failed")); await questionResponseDeliveries.sweepPending().then((result) => { if (result.scanned > 0) { @@ -1757,6 +1758,7 @@ async function startServerWithDatabaseTeardown( trackHeartbeatSchedulerWork(chatCompletionDeliveries.sweepPending().catch((err) => logger.error({ err }, "chat completion delivery failed"))); trackHeartbeatSchedulerWork(connectionDeliveries.sweepPending().catch((err) => logger.error({ err }, "connection continuation delivery failed"))); trackHeartbeatSchedulerWork(app.locals.toolGateway.sweepActionReviews().catch((err: unknown) => logger.error({ err }, "tool review recovery failed"))); + trackHeartbeatSchedulerWork(app.locals.toolGateway.cleanupExpiredSessions().catch((err: unknown) => logger.error({ err }, "gateway token cleanup failed"))); trackHeartbeatSchedulerWork(app.locals.toolActionDeliveries.sweepPending().catch((err: unknown) => logger.error({ err }, "tool review delivery sweep failed"))); trackHeartbeatSchedulerWork(questionResponseDeliveries.sweepPending() .then((result) => { diff --git a/server/src/routes/connection-intents.ts b/server/src/routes/connection-intents.ts index ddfaac14703..9d53c87ec96 100644 --- a/server/src/routes/connection-intents.ts +++ b/server/src/routes/connection-intents.ts @@ -59,7 +59,7 @@ export function runtimeConnectionIntentRoutes(db: Db) { router.get("/mcp/runtime-tools", async (req, res) => { await service.validate(runtimeClaims(req)); - res.json({ name: "paperclip-runtime-tools", protocolVersion: "2025-03-26" }); + res.set("Allow", "POST").status(405).end(); }); router.post("/mcp/runtime-tools", async (req, res) => { diff --git a/server/src/routes/openapi.ts b/server/src/routes/openapi.ts index 829064f4468..4895b3e2d24 100644 --- a/server/src/routes/openapi.ts +++ b/server/src/routes/openapi.ts @@ -10255,7 +10255,8 @@ registerCurrentRoute({ method: "get", path: "/mcp/runtime-tools", tags: ["connection-intents"], - summary: "Inspect the heartbeat-bound runtime tools MCP endpoint", + summary: "Reject SSE discovery because the runtime tools endpoint supports POST only", + responses: { 405: { description: "SSE stream is not supported" }, 401: r.unauthorized, 403: r.forbidden }, }); registerCurrentRoute({ diff --git a/server/src/routes/tool-gateway.ts b/server/src/routes/tool-gateway.ts index fc733863491..55818267ac0 100644 --- a/server/src/routes/tool-gateway.ts +++ b/server/src/routes/tool-gateway.ts @@ -55,6 +55,21 @@ function callerHeaders(req: { headers: Record(req: Request, res: Response, work: (signal: AbortSignal) => Promise) { + const controller = new AbortController(); + const aborted = () => controller.abort(); + const closed = () => { if (!res.writableEnded) controller.abort(); }; + req.once("aborted", aborted); + res.once("close", closed); + if (req.aborted || res.destroyed) controller.abort(); + try { + return await work(controller.signal); + } finally { + req.removeListener("aborted", aborted); + res.removeListener("close", closed); + } +} + async function handleMcpGatewayProtocol( req: Request, res: Response, @@ -96,11 +111,12 @@ async function handleMcpGatewayProtocol( return; } if (body.method === "tools/list") { - const tools = await toolGateway.listToolsForNamedGateway({ + const tools = await discoveryRequest(req, res, (signal) => toolGateway.listToolsForNamedGateway({ ...locator, bearerToken: token, callerHeaders: headers, - }); + signal, + })); res.json({ jsonrpc: "2.0", id, @@ -202,6 +218,7 @@ async function handleMcpGatewayProtocol( } res.status(404).json({ jsonrpc: "2.0", id, error: { code: -32601, message: "Method not found" } }); } catch (err) { + if (res.destroyed) return; if (err instanceof ToolGatewayHttpError) { const id = (req.body as { id?: unknown } | undefined)?.id ?? null; // Provider tool failures are MCP tool results, not successful calls or @@ -227,12 +244,8 @@ async function handleMcpGatewayProtocol( export function mcpGatewayProtocolRoutes(toolGateway: ToolGatewayService) { const router = Router(); - router.get("/mcp/gateways/:gatewayPublicId", async (req, res) => { - res.json({ - transport: "streamable_http", - endpoint: `/mcp/gateways/${req.params.gatewayPublicId}`, - authentication: "bearer", - }); + router.get("/mcp/gateways/:gatewayPublicId", async (_req, res) => { + res.set("Allow", "POST").status(405).end(); }); router.post("/mcp/gateways/:gatewayPublicId", async (req, res) => { await handleMcpGatewayProtocol(req, res, toolGateway, { gatewayPublicId: req.params.gatewayPublicId }); @@ -312,6 +325,7 @@ function outcomeCondition(outcome: string) { } function sendGatewayError(res: import("express").Response, err: unknown) { + if (res.destroyed) return; if (err instanceof ToolGatewayHttpError) { res.status(err.status).json({ error: err.message, @@ -450,12 +464,8 @@ export function toolGatewayRoutes(db: Db, toolGateway: ToolGatewayService) { } }); - router.get("/tool-gateway/gateways/:gatewayId/mcp", async (req, res) => { - res.json({ - transport: "streamable_http", - endpoint: `/api/tool-gateway/gateways/${req.params.gatewayId}/mcp`, - authentication: "bearer", - }); + router.get("/tool-gateway/gateways/:gatewayId/mcp", async (_req, res) => { + res.set("Allow", "POST").status(405).end(); }); router.post("/tool-gateway/gateways/:gatewayId/mcp", async (req, res) => { @@ -551,7 +561,7 @@ export function toolGatewayRoutes(db: Db, toolGateway: ToolGatewayService) { res.status(401).json({ error: "Tool gateway session token is required" }); return; } - const tools = await toolGateway.listToolsForSession(token); + const tools = await discoveryRequest(req, res, (signal) => toolGateway.listToolsForSession(token, { signal })); res.json(tools); } catch (err) { sendGatewayError(res, err); diff --git a/server/src/services/chat-github-tools.ts b/server/src/services/chat-github-tools.ts index 572a78a0bfd..f8bcc159ff7 100644 --- a/server/src/services/chat-github-tools.ts +++ b/server/src/services/chat-github-tools.ts @@ -21,6 +21,7 @@ import type { ToolGatewayDescriptor, ToolGatewaySession, } from "./tool-gateway.js"; +import { runContextSnapshotString } from "./tool-access-policy.js"; const objectSchema = ( properties: Record, @@ -304,7 +305,7 @@ export async function githubBotConnectionIdsForRun( const [run] = await db .select({ issueId: heartbeatRuns.nativeIssueId, - context: heartbeatRuns.contextSnapshot, + snapshotIssueId: runContextSnapshotString("issueId"), }) .from(heartbeatRuns) .where( @@ -314,7 +315,7 @@ export async function githubBotConnectionIdsForRun( eq(heartbeatRuns.id, runId), ), ); - const issueId = run?.issueId ?? run?.context?.issueId; + const issueId = run?.issueId ?? run?.snapshotIssueId; if (typeof issueId !== "string") return new Set(); const rows = await db .select({ @@ -356,7 +357,7 @@ export async function githubGuestBotConnectionForSession( ): Promise { if (!session.runId || !session.agentId || !session.issueId) return null; const [run] = await db - .select() + .select({ wakeCommentId: runContextSnapshotString("wakeCommentId") }) .from(heartbeatRuns) .where( and( @@ -365,7 +366,7 @@ export async function githubGuestBotConnectionForSession( eq(heartbeatRuns.id, session.runId), ), ); - const commentId = run?.contextSnapshot?.wakeCommentId; + const commentId = run?.wakeCommentId; if (typeof commentId !== "string") return null; const [source] = await db .select({ diff --git a/server/src/services/tool-access-policy.ts b/server/src/services/tool-access-policy.ts index c5c4edf69ea..827203f82db 100644 --- a/server/src/services/tool-access-policy.ts +++ b/server/src/services/tool-access-policy.ts @@ -125,7 +125,7 @@ function nonBlankString(value: string | null | undefined): string | null { * JSON object, which the run indexes on `context_snapshot ->> 'issueId'` also * expect. */ -function runContextSnapshotString(key: "issueId" | "projectId" | "routineId") { +export function runContextSnapshotString(key: "issueId" | "projectId" | "routineId" | "wakeCommentId") { const field = sql.raw(`'${key}'`); return sql`case when jsonb_typeof(${heartbeatRuns.contextSnapshot} -> ${field}) = 'string' then ${heartbeatRuns.contextSnapshot} ->> ${field} end`; diff --git a/server/src/services/tool-discovery-scheduler.test.ts b/server/src/services/tool-discovery-scheduler.test.ts new file mode 100644 index 00000000000..b50f37079c5 --- /dev/null +++ b/server/src/services/tool-discovery-scheduler.test.ts @@ -0,0 +1,48 @@ +import { describe, expect, it } from "vitest"; +import { createToolDiscoveryScheduler } from "./tool-discovery-scheduler.js"; + +function deferred() { + let resolve!: () => void; + const promise = new Promise((done) => { resolve = done; }); + return { promise, resolve }; +} + +describe("tool discovery admission", () => { + it("bounds whole listings and rejects excess queued requests", async () => { + const scheduler = createToolDiscoveryScheduler(2, 2); + const gate = deferred(); + let active = 0; + let peak = 0; + const work = async () => { + active += 1; + peak = Math.max(peak, active); + await gate.promise; + active -= 1; + }; + const accepted = Array.from({ length: 4 }, () => scheduler.run(work)); + await expect(scheduler.run(work)).rejects.toMatchObject({ name: "ToolDiscoveryBusyError" }); + expect(active).toBe(2); + gate.resolve(); + await Promise.all(accepted); + expect(peak).toBe(2); + await scheduler.run(async () => {}); + }); + + it("removes abandoned queued requests and releases slots after failures", async () => { + const scheduler = createToolDiscoveryScheduler(1, 1); + const gate = deferred(); + const first = scheduler.run(() => gate.promise); + const controller = new AbortController(); + let abandonedStarted = false; + const abandoned = scheduler.run(async () => { abandonedStarted = true; }, controller.signal); + controller.abort(); + await expect(abandoned).rejects.toMatchObject({ name: "AbortError" }); + const next = scheduler.run(async () => { throw new Error("Read failed"); }); + const failure = expect(next).rejects.toThrow("Read failed"); + gate.resolve(); + await first; + await failure; + expect(abandonedStarted).toBe(false); + await scheduler.run(async () => {}); + }); +}); diff --git a/server/src/services/tool-discovery-scheduler.ts b/server/src/services/tool-discovery-scheduler.ts new file mode 100644 index 00000000000..b680729bdd4 --- /dev/null +++ b/server/src/services/tool-discovery-scheduler.ts @@ -0,0 +1,54 @@ +export class ToolDiscoveryBusyError extends Error { + constructor() { + super("Tool discovery is busy. Retry shortly."); + this.name = "ToolDiscoveryBusyError"; + } +} + +/** Limits whole listings, including catalog reads, across gateway service instances. */ +export function createToolDiscoveryScheduler(concurrency = 2, queueLimit = 32) { + if (!Number.isInteger(concurrency) || concurrency < 1 || !Number.isInteger(queueLimit) || queueLimit < 0) { + throw new Error("Invalid tool discovery scheduler limits"); + } + let active = 0; + const queue: Array<{ start: () => void }> = []; + function drain() { + while (active < concurrency && queue.length > 0) { + active += 1; + queue.shift()!.start(); + } + } + function run(work: () => Promise, signal?: AbortSignal): Promise { + if (signal?.aborted) return Promise.reject(signal.reason); + if (active >= concurrency && queue.length >= queueLimit) { + return Promise.reject(new ToolDiscoveryBusyError()); + } + return new Promise((resolve, reject) => { + const job = { + start() { + signal?.removeEventListener("abort", abort); + // Do not free a slot on abort until already-started reads settle. + Promise.resolve().then(() => { + signal?.throwIfAborted(); + return work(); + }).then(resolve, reject).finally(() => { + active -= 1; + drain(); + }); + }, + }; + function abort() { + const index = queue.indexOf(job); + if (index < 0) return; + queue.splice(index, 1); + reject(signal!.reason); + } + queue.push(job); + signal?.addEventListener("abort", abort, { once: true }); + drain(); + }); + } + return { run }; +} + +export const toolDiscoveryScheduler = createToolDiscoveryScheduler(); diff --git a/server/src/services/tool-gateway.ts b/server/src/services/tool-gateway.ts index b6534ca7e9d..f3e7c30093a 100644 --- a/server/src/services/tool-gateway.ts +++ b/server/src/services/tool-gateway.ts @@ -129,9 +129,11 @@ import { } from "./remote-url-credentials.js"; import { createToolAccessDecisionCache, + runContextSnapshotString, toolAccessPolicyService, type ToolAccessDecisionCache, } from "./tool-access-policy.js"; +import { toolDiscoveryScheduler, ToolDiscoveryBusyError } from "./tool-discovery-scheduler.js"; import { commitToolActionReview } from "./tool-action-review.js"; import { issueThreadInteractionService } from "./issue-thread-interactions.js"; import { @@ -470,6 +472,7 @@ async function mapWithConcurrency( items: readonly T[], concurrency: number, mapper: (item: T) => Promise, + signal?: AbortSignal, ): Promise { const results = new Array(items.length); let nextIndex = 0; @@ -480,6 +483,7 @@ async function mapWithConcurrency( const index = nextIndex; nextIndex += 1; try { + signal?.throwIfAborted(); results[index] = await mapper(items[index]!); } catch (error) { // The batch fails with this error, so the other workers stop. @@ -489,9 +493,12 @@ async function mapWithConcurrency( } } - await Promise.all( + const workers = await Promise.allSettled( Array.from({ length: Math.min(concurrency, items.length) }, () => worker()), ); + const failure = workers.find((result) => result.status === "rejected"); + if (failure?.status === "rejected") throw failure.reason; + signal?.throwIfAborted(); return results; } @@ -1599,7 +1606,8 @@ export function createToolGatewayService( companyId: heartbeatRuns.companyId, agentId: heartbeatRuns.agentId, status: heartbeatRuns.status, - contextSnapshot: heartbeatRuns.contextSnapshot, + issueId: runContextSnapshotString("issueId"), + projectId: runContextSnapshotString("projectId"), }) .from(heartbeatRuns) .where(eq(heartbeatRuns.id, input.runId)) @@ -1623,9 +1631,8 @@ export function createToolGatewayService( throw new ToolGatewayHttpError(403, "Run is not active", "run_inactive"); } - const snapshot = asRecord(run.contextSnapshot); - const snapshotIssueId = stringValue(snapshot?.issueId); - const snapshotProjectId = stringValue(snapshot?.projectId); + const snapshotIssueId = stringValue(run.issueId); + const snapshotProjectId = stringValue(run.projectId); if ( (input.issueId && snapshotIssueId && input.issueId !== snapshotIssueId) || (input.projectId && @@ -2722,11 +2729,12 @@ export function createToolGatewayService( tools: readonly T[], inputForTool: (tool: T) => ToolAccessDecisionInput, cache: ToolAccessDecisionCache = createToolAccessDecisionCache(), + signal?: AbortSignal, ): Promise> { return mapWithConcurrency(tools, LISTING_DECISION_CONCURRENCY, async (tool) => ({ tool, decision: await policyService.decide(inputForTool(tool), { cache }), - })); + }), signal); } function policyErrorStatus(decision: ToolAccessDecision) { @@ -2958,7 +2966,23 @@ export function createToolGatewayService( async function listToolsForContext( session: ToolGatewaySession, + signal?: AbortSignal, ): Promise { + try { + return await toolDiscoveryScheduler.run(() => buildToolsForContext(session, signal), signal); + } catch (error) { + if (error instanceof ToolDiscoveryBusyError) { + throw new ToolGatewayHttpError(503, error.message, "tool_discovery_busy"); + } + throw error; + } + } + + async function buildToolsForContext( + session: ToolGatewaySession, + signal?: AbortSignal, + ): Promise { + signal?.throwIfAborted(); if (session.agentId) { await assertAgentInCompany(session.companyId, session.agentId); } @@ -2983,6 +3007,7 @@ export function createToolGatewayService( tools, (tool) => policyInputForTool({ session, tool }), decisionCache, + signal, ); const visibleTools = decisions .filter( @@ -3007,6 +3032,7 @@ export function createToolGatewayService( onDemandTargets, (tool) => policyInputForTool({ session, tool }), decisionCache, + signal, ); if ( targetDecisions.some( @@ -8782,7 +8808,9 @@ export function createToolGatewayService( gatewayPublicId?: string | null; bearerToken: string; callerHeaders?: Record; + signal?: AbortSignal; }): Promise { + input.signal?.throwIfAborted(); const session = await namedGatewaySessionFromBearer({ gatewayId: input.gatewayId ?? null, gatewayPublicId: input.gatewayPublicId ?? null, @@ -8791,7 +8819,7 @@ export function createToolGatewayService( callerHeaders: input.callerHeaders, }); await assertGatewayTokenAction(session, "tools/list"); - const tools = await listToolsForContext(session); + const tools = await listToolsForContext(session, input.signal); await writeAudit({ session, companyId: session.companyId, @@ -8803,7 +8831,7 @@ export function createToolGatewayService( decision: "allow", reasonCode: "named_gateway_discovery_filtered", visibleToolCount: tools.length, - visibleTools: tools.map((tool) => tool.name), + visibleToolsHash: createHash("sha256").update(JSON.stringify(tools.map((tool) => tool.name).sort())).digest("hex"), }, }); return tools; @@ -8882,9 +8910,11 @@ export function createToolGatewayService( async listToolsForSession( sessionToken: string, + options: { signal?: AbortSignal } = {}, ): Promise { + options.signal?.throwIfAborted(); const session = await getActiveSession(sessionToken); - const tools = await listToolsForContext(session); + const tools = await listToolsForContext(session, options.signal); await writeAudit({ session, companyId: session.companyId, @@ -8896,7 +8926,7 @@ export function createToolGatewayService( decision: "allow", reasonCode: "discovery_filtered", visibleToolCount: tools.length, - visibleTools: tools.map((tool) => tool.name), + visibleToolsHash: createHash("sha256").update(JSON.stringify(tools.map((tool) => tool.name).sort())).digest("hex"), }, }); return tools; @@ -11112,6 +11142,16 @@ export function createToolGatewayService( async cleanupExpiredSessions(input: { now?: Date } = {}) { const now = input.now ?? new Date(); + const expiredTokens = await db.select({ id: toolMcpGatewayTokens.id }) + .from(toolMcpGatewayTokens) + .where(lte(toolMcpGatewayTokens.expiresAt, now)) + .orderBy(asc(toolMcpGatewayTokens.expiresAt), asc(toolMcpGatewayTokens.id)) + .limit(500); + if (expiredTokens.length > 0) { + await db.delete(toolMcpGatewayTokens) + .where(and(inArray(toolMcpGatewayTokens.id, expiredTokens.map((token) => token.id)), lte(toolMcpGatewayTokens.expiresAt, now))) + .returning({ id: toolMcpGatewayTokens.id }); + } const rows = await db .delete(toolGatewaySessions) .where(lte(toolGatewaySessions.expiresAt, now)) diff --git a/tests/e2e/mcp-user-stories.spec.ts b/tests/e2e/mcp-user-stories.spec.ts index 93acd041bb4..2585bc1a147 100644 --- a/tests/e2e/mcp-user-stories.spec.ts +++ b/tests/e2e/mcp-user-stories.spec.ts @@ -195,6 +195,10 @@ async function startMockMcp(): Promise { } async function screenshot(page: Page, storyId: string, step: string) { + if (new URL(page.url()).pathname.endsWith("/review")) { + await expect(page.getByRole("heading", { name: "Review", exact: true })).toBeVisible({ timeout: 30_000 }); + await expect(page.getByText("Loading tools…", { exact: true })).toBeHidden(); + } await page.screenshot({ path: `${SCREENSHOT_DIR}/${storyId.toLowerCase()}-${step}.png`, fullPage: true }); } @@ -318,9 +322,10 @@ test.describe.serial("MCP prod Phase 5a user-story harness", () => { expect(pending.actionRequestId).toBeTruthy(); await page.goto(`/${seed.prefix}/apps/${connectionId}/review`); + await expect(page.getByRole("heading", { name: "Waiting for your OK" })).toBeVisible({ timeout: 30_000 }); await screenshot(page, "US-2", "01-review-pending"); - await approveActionRequest(request, seed.companyId, pending.actionRequestId!); + await page.getByRole("button", { name: "Allow once", exact: true }).click(); await pollTestCall(request, connectionId, pending.actionRequestId!, "done"); expect(mock.captures.filter((capture) => capture.method === "tools/call" && capture.toolName === "sheets:update_cell")).toHaveLength(1); await expectAuditEvent(request, seed.companyId, { connectionId, agentId: scout.id, search: "sheets:update_cell" });