mirror of
https://github.com/paperclipai/paperclip.git
synced 2026-10-22 16:51:28 +02:00
fix(tool-gateway): bound discovery admission and retain regression evidence
Add whole-listing admission, abort propagation, scalar context reads, compact discovery audits, indexed expiry sweeps, and GET 405 for the stateless MCP transport. Record red-green tests, an HTTP concurrency and policy-revocation journey, browser approvals, and before/after measurements. Co-Authored-By: Paperclip <noreply@paperclip.ing>
This commit is contained in:
DottaandPaperclip committed
1 parent
054fdf0f0c
commit
c41169095b
21 files changed
+2155
-206
No files matched your search
@@ -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
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
@@ -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.
|
||||
@@ -0,0 +1 @@
|
||||
CREATE INDEX IF NOT EXISTS "tool_mcp_gateway_tokens_expiry_idx" ON "tool_mcp_gateway_tokens" USING btree ("expires_at","id");
|
||||
+990
-2
File diff suppressed because it is too large.
Load diff
+7
@@ -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
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -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),
|
||||
|
||||
@@ -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();
|
||||
}
|
||||
@@ -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<typeof createDb>;
|
||||
|
||||
// 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 };
|
||||
}
|
||||
@@ -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);
|
||||
|
||||
@@ -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<typeof createDb>;
|
||||
let temp!: Awaited<ReturnType<typeof startEmbeddedPostgresTestDatabase>>;
|
||||
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);
|
||||
});
|
||||
});
|
||||
@@ -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<typeof createDb>;
|
||||
|
||||
// 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<ReturnType<typeof startEmbeddedPostgresTestDatabase>> | 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);
|
||||
|
||||
@@ -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) => {
|
||||
|
||||
@@ -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) => {
|
||||
|
||||
@@ -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({
|
||||
|
||||
@@ -55,6 +55,21 @@ function callerHeaders(req: { headers: Record<string, string | string[] | undefi
|
||||
return headers;
|
||||
}
|
||||
|
||||
async function discoveryRequest<T>(req: Request, res: Response, work: (signal: AbortSignal) => Promise<T>) {
|
||||
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);
|
||||
|
||||
@@ -21,6 +21,7 @@ import type {
|
||||
ToolGatewayDescriptor,
|
||||
ToolGatewaySession,
|
||||
} from "./tool-gateway.js";
|
||||
import { runContextSnapshotString } from "./tool-access-policy.js";
|
||||
|
||||
const objectSchema = (
|
||||
properties: Record<string, unknown>,
|
||||
@@ -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<string | null> {
|
||||
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({
|
||||
|
||||
@@ -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<string | null>`case when jsonb_typeof(${heartbeatRuns.contextSnapshot} -> ${field}) = 'string'
|
||||
then ${heartbeatRuns.contextSnapshot} ->> ${field} end`;
|
||||
|
||||
@@ -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<void>((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 () => {});
|
||||
});
|
||||
});
|
||||
@@ -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<T>(work: () => Promise<T>, signal?: AbortSignal): Promise<T> {
|
||||
if (signal?.aborted) return Promise.reject(signal.reason);
|
||||
if (active >= concurrency && queue.length >= queueLimit) {
|
||||
return Promise.reject(new ToolDiscoveryBusyError());
|
||||
}
|
||||
return new Promise<T>((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();
|
||||
@@ -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<T, R>(
|
||||
items: readonly T[],
|
||||
concurrency: number,
|
||||
mapper: (item: T) => Promise<R>,
|
||||
signal?: AbortSignal,
|
||||
): Promise<R[]> {
|
||||
const results = new Array<R>(items.length);
|
||||
let nextIndex = 0;
|
||||
@@ -480,6 +483,7 @@ async function mapWithConcurrency<T, R>(
|
||||
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<T, R>(
|
||||
}
|
||||
}
|
||||
|
||||
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<Array<{ tool: T; decision: ToolAccessDecision }>> {
|
||||
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<ToolGatewayDescriptor[]> {
|
||||
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<ToolGatewayDescriptor[]> {
|
||||
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<string, string | string[] | undefined>;
|
||||
signal?: AbortSignal;
|
||||
}): Promise<ToolGatewayDescriptor[]> {
|
||||
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<ToolGatewayDescriptor[]> {
|
||||
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))
|
||||
|
||||
@@ -195,6 +195,10 @@ async function startMockMcp(): Promise<MockMcpServer> {
|
||||
}
|
||||
|
||||
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" });
|
||||
|
||||
Reference in new issue
Block a user