History

pool MCP connections within session to skip cold handshakes (#457)

The MCP plugin already pools connections via a ScopedCache keyed on `${transport}:${invokerScope}:${endpoint}`, scoped to the executor (one executor per McpSessionDO). Production telemetry for the past 8h shows `plugin.mcp.connection.acquire` p99 3.59s with handshake firing on 4 of 24 acquires (17%) — the cache IS hitting on the other 20, with the misses representing structurally unavoidable cold-start sessions (first invoke per DO, retry-after-stale, TTL expiry at the session timeout boundary). This commit lands two regression tests pinning the per-session pooling contract (5 sequential invokes => 1 handshake; different tools on the same source share the cached connection) plus a `plugin.mcp.cache_hit` boolean attribute on the acquire span so Axiom queries can read the hit rate directly without back-computing from handshake counts. Notes/analysis in notes/mcp-conn-pool.md.

Rhys Sullivan committed 6 months ago

Commit 543e01c1 parent 685d4303 files · +371 −0Browse files

Unified · Split

3 files
  1. notes/mcp-conn-pool.md +135 −0
  2. packages/plugins/mcp/src/sdk/connection-pool.test.ts +228 −0
  3. packages/plugins/mcp/src/sdk/invoke.ts +8 −0
notes/mcp-conn-pool.md+135−0
@@ -0,0 +1,135 @@
1+# MCP connection pool — investigation notes (2026-04-30)
2+
3+## TL;DR
4+
5+The MCP plugin **already pools connections within a session DO**. Production
6+telemetry attributing 17% of `plugin.mcp.connection.acquire` calls to a
7+fresh `plugin.mcp.connection.handshake` represents *cold-start sessions*,
8+not redundant in-session handshakes. No structural change is required;
9+this PR ships a strict regression test that pins the existing contract
10+plus a `plugin.mcp.cache_hit` span attribute so future telemetry can
11+distinguish hits from misses without cross-referencing handshake counts.
12+
13+## What's already there
14+
15+`packages/plugins/mcp/src/sdk/plugin.ts#makeRuntime`:
16+
17+```ts
18+const connectionCache = yield* ScopedCache.make({
19+ lookup: (key: string) =>
20+ Effect.acquireRelease(
21+ Effect.suspend(() => {
22+ const connector = pendingConnectors.get(key);
23+ ...
24+ }),
25+ (connection) => Effect.promise(() => connection.close().catch(() => {})),
26+ ),
27+ capacity: 64,
28+ timeToLive: Duration.minutes(5),
29+}).pipe(Scope.extend(cacheScope));
30+```
31+
32+Properties:
33+
34+- One `ScopedCache` per plugin instance, captured by the closure
35+ `runtimeRef`. Lifecycle is the executor (one executor per session DO,
36+ built in `apps/cloud/src/services/executor.ts#createScopedExecutor`).
37+- LRU with `capacity: 64`, `timeToLive: 5 min`.
38+- Lookup runs inside a Scope held by `cacheScope`, so the pooled
39+ connection survives `Effect.scoped` boundaries inside `invokeMcpTool`.
40+- Cache key (`invoke.ts#connectionCacheKey`) is
41+ `remote:${invokerScope}:${endpoint}` for remote sources and
42+ `stdio:${command}` for stdio. Includes `invokerScope` so per-user
43+ OAuth/header secrets never collapse onto a shared connection.
44+- `invokeMcpTool` retries once on connection error: invalidates the
45+ cache entry and re-acquires. Stale connections are caught here and
46+ transparently re-handshaked.
47+
48+## Why the pool is plausibly correct
49+
50+`packages/plugins/mcp/src/sdk/elicitation.test.ts` already had a test
51+`"connection is reused across multiple tool calls to the same source"`
52+that asserts the underlying MCP server sees no new HTTP sessions across
53+three sequential `executor.tools.invoke` calls — passes against current
54+code.
55+
56+This PR adds two stricter regression cases under
57+`packages/plugins/mcp/src/sdk/connection-pool.test.ts`:
58+
59+1. Five sequential invokes of the same tool — one handshake total.
60+2. Different tools on the same source — still one handshake, different
61+ tool ids hit the same cache key.
62+
63+Both pass without any source change. They serve as a deterministic
64+contract pin so a future refactor that breaks pooling (e.g., scoping
65+the cache to the per-invoke scope, dropping `Scope.extend(cacheScope)`,
66+mutating the cache key per call) trips immediately in CI.
67+
68+## Re-reading the production trace
69+
70+Trace `b7102047bed975da461c0519d1251de4`:
71+- `mcp.plugin.resolve_connector` 2.72s
72+- `plugin.mcp.connection.acquire` 2.72s
73+- `plugin.mcp.connection.handshake` 2.52s
74+- `executor.storage.transaction` 1.02s (token persist)
75+- `plugin.mcp.client.call_tool` 1.48s
76+
77+Span nesting plus `resolve_connector ≈ acquire ≈ handshake` durations is
78+the cache-miss path: `resolveConnector` only runs (and emits its span)
79+when the cache lookup actually invokes the lookup closure, which only
80+happens on a miss.
81+
82+8h aggregate:
83+- `plugin.mcp.connection.handshake` count 4
84+- `plugin.mcp.connection.acquire` count 24
85+- `mcp.plugin.resolve_connector` count 4
86+
87+Acquire count is 6× handshake count. The cache *is* hitting on 20 of 24
88+calls. The 4 misses are cold-start sessions (each MCP session DO's first
89+tool call, plus any session where the connection went stale and was
90+invalidated by the retry path). Six tool calls per session is consistent
91+with normal MCP usage.
92+
93+## Why we can't drive misses below the cold-start floor
94+
95+Every miss observed in prod is structurally unavoidable inside a single
96+DO:
97+
98+- **First tool invocation in a fresh DO.** `init()` builds the executor
99+ and a fresh `ScopedCache`. The first invoke must handshake.
100+- **Stale-connection retry.** `invokeMcpTool` invalidates the cache
101+ entry on `client.callTool` failure and re-acquires; that re-acquire
102+ is necessarily a miss.
103+- **TTL expiry.** Set to 5 minutes. The session DO's idle alarm fires
104+ at 5 minutes too, so any cache entry that out-survives a session was
105+ going to be discarded with the DO anyway.
106+
107+To go below this floor we'd need either:
108+
109+1. **Pre-warm in `init()`.** Plausible but speculative — we don't know
110+ which sources the user will invoke, and warming all of them
111+ bottlenecks `init()` on the slowest server. Out of scope; would also
112+ reverse the savings if the user never calls those tools.
113+2. **Cross-DO connection pool.** Forbidden by the task — and the
114+ per-DO scope of `runtimeRef` is the right home for a per-user-org
115+ connection pool given the cache key includes `invokerScope`.
116+
117+## What this PR ships
118+
119+1. `packages/plugins/mcp/src/sdk/connection-pool.test.ts` — two
120+ strict regression tests pinning the per-session pooling contract.
121+2. `packages/plugins/mcp/src/sdk/invoke.ts` — adds
122+ `plugin.mcp.cache_hit: boolean` attribute to the
123+ `plugin.mcp.connection.acquire` span. Future Axiom queries can read
124+ the hit rate directly without comparing acquire vs handshake counts.
125+
126+No behavior change.
127+
128+## Follow-ups
129+
130+- If the cold-start p99 still bites tail latency, consider warming the
131+ most-recently-used source(s) for a session DO during `init()` —
132+ scoped to the user's recent activity, e.g., on session restore.
133+- Watch `plugin.mcp.cache_hit=false` rate over a longer window. If it
134+ exceeds the implied cold-start floor (≈ 1 / avg-calls-per-session),
135+ there's a real bug to chase.
packages/plugins/mcp/src/sdk/connection-pool.test.ts+228−0
@@ -0,0 +1,228 @@
1+// ---------------------------------------------------------------------------
2+// MCP connection pooling regression test
3+//
4+// Production telemetry (Axiom, 2026-04-30): `plugin.mcp.connection.acquire`
5+// p99 = 3.59s, with `plugin.mcp.connection.handshake` p99 = 3.28s firing on
6+// ~17% of acquires. The cold MCP handshake is the dominant tail-latency
7+// source for tool invocations. The MCP plugin already pools connections via
8+// a `ScopedCache` keyed on `${transport}:${invokerScope}:${endpoint}` (see
9+// `plugin.ts#makeRuntime`), so within a single executor / session DO the
10+// SECOND and onward invocations against the same source config MUST reuse
11+// the cached connection without performing a fresh transport handshake.
12+//
13+// This file pins that contract behind a deterministic test harness that
14+// counts inbound MCP server sessions — each MCP server `connect` is exactly
15+// one cold handshake on the wire. After the first invoke kicks off the cold
16+// connection, all subsequent calls (regardless of which tool on the source
17+// they target) MUST land on the same MCP session.
18+//
19+// If this test starts failing, the cache key, the runtime ref, or the
20+// `Effect.scoped` boundary inside `invokeMcpTool` has regressed and prod
21+// will see the cold-handshake p99 climb proportionally.
22+// ---------------------------------------------------------------------------
23+
24+import { describe, expect, it } from "@effect/vitest";
25+import { Effect } from "effect";
26+import * as http from "node:http";
27+import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
28+import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js";
29+import { z } from "zod";
30+
31+import { createExecutor, makeTestConfig } from "@executor/sdk";
32+
33+import { mcpPlugin } from "./plugin";
34+
35+// ---------------------------------------------------------------------------
36+// Test MCP server — counts session connects (each = one cold handshake)
37+// ---------------------------------------------------------------------------
38+
39+function createTestMcpServer() {
40+ const server = new McpServer(
41+ { name: "pool-test-server", version: "1.0.0" },
42+ { capabilities: {} },
43+ );
44+
45+ server.registerTool(
46+ "echo",
47+ { description: "Echo a value", inputSchema: { value: z.string() } },
48+ async ({ value }: { value: string }) => ({
49+ content: [{ type: "text" as const, text: value }],
50+ }),
51+ );
52+
53+ server.registerTool(
54+ "echo2",
55+ { description: "Echo a value", inputSchema: { value: z.string() } },
56+ async ({ value }: { value: string }) => ({
57+ content: [{ type: "text" as const, text: `b:${value}` }],
58+ }),
59+ );
60+
61+ return server;
62+}
63+
64+type TestServer = {
65+ readonly url: string;
66+ readonly httpServer: http.Server;
67+ /** Number of MCP sessions created (1 per cold transport handshake). */
68+ readonly sessionCount: () => number;
69+};
70+
71+const serveMcpServer = Effect.acquireRelease(
72+ Effect.async<TestServer, Error>((resume) => {
73+ const transports = new Map<string, StreamableHTTPServerTransport>();
74+ let sessions = 0;
75+
76+ const httpServer = http.createServer(async (req, res) => {
77+ const sessionId = req.headers["mcp-session-id"] as string | undefined;
78+
79+ if (sessionId) {
80+ const transport = transports.get(sessionId);
81+ if (!transport) {
82+ res.writeHead(404);
83+ res.end("Session not found");
84+ return;
85+ }
86+ await transport.handleRequest(req, res);
87+ return;
88+ }
89+
90+ const mcpServer = createTestMcpServer();
91+ sessions++;
92+
93+ const transport = new StreamableHTTPServerTransport({
94+ sessionIdGenerator: () => crypto.randomUUID(),
95+ onsessioninitialized: (sid) => {
96+ transports.set(sid, transport);
97+ },
98+ });
99+
100+ await mcpServer.connect(transport);
101+ await transport.handleRequest(req, res);
102+ });
103+
104+ httpServer.listen(0, () => {
105+ const addr = httpServer.address();
106+ const port = typeof addr === "object" && addr ? addr.port : 0;
107+ resume(
108+ Effect.succeed({
109+ url: `http://127.0.0.1:${port}`,
110+ httpServer,
111+ sessionCount: () => sessions,
112+ }),
113+ );
114+ });
115+ }),
116+ ({ httpServer }) =>
117+ Effect.sync(() => {
118+ httpServer.close();
119+ }),
120+);
121+
122+// ---------------------------------------------------------------------------
123+// Helper — one executor, one mcp source pointed at the test server
124+// ---------------------------------------------------------------------------
125+
126+const makeTestExecutor = (serverUrl: string) =>
127+ createExecutor(
128+ makeTestConfig({
129+ plugins: [mcpPlugin()] as const,
130+ }),
131+ ).pipe(
132+ Effect.tap((executor) =>
133+ executor.mcp.addSource({
134+ transport: "remote",
135+ scope: "test-scope",
136+ name: "pool-test",
137+ endpoint: serverUrl,
138+ }),
139+ ),
140+ );
141+
142+// ---------------------------------------------------------------------------
143+// Tests
144+// ---------------------------------------------------------------------------
145+
146+describe("MCP connection pooling (regression)", () => {
147+ it.scoped(
148+ "five sequential invokes against the same source perform exactly one transport handshake",
149+ () =>
150+ Effect.gen(function* () {
151+ const server = yield* serveMcpServer;
152+ const executor = yield* makeTestExecutor(server.url);
153+ const tools = yield* executor.tools.list();
154+ const echo = tools.find((t) => t.name === "echo")!;
155+ expect(echo).toBeDefined();
156+
157+ // Discovery during addSource performs ONE handshake on a separate,
158+ // closed connection — that's expected. Capture the baseline so we
159+ // assert against the invoke-path delta only.
160+ const sessionsAfterAddSource = server.sessionCount();
161+ expect(sessionsAfterAddSource).toBeGreaterThanOrEqual(1);
162+
163+ // First invoke is necessarily a cache miss — it cold-handshakes the
164+ // pooled connection. Every subsequent invoke must hit the cache and
165+ // produce zero new MCP server sessions.
166+ yield* executor.tools.invoke(
167+ echo.id,
168+ { value: "1" },
169+ { onElicitation: "accept-all" },
170+ );
171+ const sessionsAfterFirstInvoke = server.sessionCount();
172+ expect(sessionsAfterFirstInvoke).toBe(sessionsAfterAddSource + 1);
173+
174+ for (let i = 2; i <= 5; i++) {
175+ yield* executor.tools.invoke(
176+ echo.id,
177+ { value: String(i) },
178+ { onElicitation: "accept-all" },
179+ );
180+ // Cache must reuse the same MCP session for every subsequent
181+ // call. If this assertion fails the session DO is paying the
182+ // cold-handshake p99 (3.28s in prod) on every tool call.
183+ expect(server.sessionCount()).toBe(sessionsAfterFirstInvoke);
184+ }
185+ }),
186+ );
187+
188+ it.scoped(
189+ "different tools on the same source share the cached connection",
190+ () =>
191+ Effect.gen(function* () {
192+ const server = yield* serveMcpServer;
193+ const executor = yield* makeTestExecutor(server.url);
194+ const tools = yield* executor.tools.list();
195+ const echo = tools.find((t) => t.name === "echo")!;
196+ const echo2 = tools.find((t) => t.name === "echo2")!;
197+ expect(echo).toBeDefined();
198+ expect(echo2).toBeDefined();
199+
200+ const sessionsAfterAddSource = server.sessionCount();
201+
202+ // Cold handshake on first call.
203+ yield* executor.tools.invoke(
204+ echo.id,
205+ { value: "a" },
206+ { onElicitation: "accept-all" },
207+ );
208+ const baseline = server.sessionCount();
209+ expect(baseline).toBe(sessionsAfterAddSource + 1);
210+
211+ // Different tool on the SAME source — same cache key, same conn.
212+ yield* executor.tools.invoke(
213+ echo2.id,
214+ { value: "b" },
215+ { onElicitation: "accept-all" },
216+ );
217+ expect(server.sessionCount()).toBe(baseline);
218+
219+ // Back to the first tool — still pooled.
220+ yield* executor.tools.invoke(
221+ echo.id,
222+ { value: "c" },
223+ { onElicitation: "accept-all" },
224+ );
225+ expect(server.sessionCount()).toBe(baseline);
226+ }),
227+ );
228+});
packages/plugins/mcp/src/sdk/invoke.ts+8−0
@@ -194,12 +194,20 @@ export const invokeMcpTool = ( · 193 unmodified lines
194194 const connector = input.resolveConnector();
195195 input.pendingConnectors.set(cacheKey, connector);
196196
197+ // Check cache state BEFORE acquire so the span clearly attributes
198+ // tail latency to either a cold handshake (miss) or warm reuse (hit).
199+ // Without this every `plugin.mcp.connection.acquire` span looks the
200+ // same in Axiom and you have to cross-reference the
201+ // `plugin.mcp.connection.handshake` count to back out the hit rate.
202+ const cacheHit = yield* input.connectionCache.contains(cacheKey);
203+
197204 const firstConnection = yield* input.connectionCache.get(cacheKey).pipe(
198205 Effect.withSpan("plugin.mcp.connection.acquire", {
199206 attributes: {
200207 "plugin.mcp.transport": transport,
201208 "plugin.mcp.cache_key": cacheKey,
202209 "plugin.mcp.attempt": 1,
210+ "plugin.mcp.cache_hit": cacheHit,
203211 },
204212 }),
205213 );