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
5The MCP plugin **already pools connections within a session DO**. Production
6telemetry attributing 17% of `plugin.mcp.connection.acquire` calls to a
7fresh `plugin.mcp.connection.handshake` represents *cold-start sessions*,
8not redundant in-session handshakes. No structural change is required;
9this PR ships a strict regression test that pins the existing contract
10plus a `plugin.mcp.cache_hit` span attribute so future telemetry can
11distinguish 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
18const 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
32Properties:
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"`
52that asserts the underlying MCP server sees no new HTTP sessions across
53three sequential `executor.tools.invoke` calls — passes against current
54code.
55
56This PR adds two stricter regression cases under
57`packages/plugins/mcp/src/sdk/connection-pool.test.ts`:
58
591. Five sequential invokes of the same tool — one handshake total.
602. Different tools on the same source — still one handshake, different
61 tool ids hit the same cache key.
62
63Both pass without any source change. They serve as a deterministic
64contract pin so a future refactor that breaks pooling (e.g., scoping
65the cache to the per-invoke scope, dropping `Scope.extend(cacheScope)`,
66mutating the cache key per call) trips immediately in CI.
67
68## Re-reading the production trace
69
70Trace `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
77Span nesting plus `resolve_connector ≈ acquire ≈ handshake` durations is
78the cache-miss path: `resolveConnector` only runs (and emits its span)
79when the cache lookup actually invokes the lookup closure, which only
80happens on a miss.
81
828h aggregate:
83- `plugin.mcp.connection.handshake` count 4
84- `plugin.mcp.connection.acquire` count 24
85- `mcp.plugin.resolve_connector` count 4
86
87Acquire count is 6× handshake count. The cache *is* hitting on 20 of 24
88calls. The 4 misses are cold-start sessions (each MCP session DO's first
89tool call, plus any session where the connection went stale and was
90invalidated by the retry path). Six tool calls per session is consistent
91with normal MCP usage.
92
93## Why we can't drive misses below the cold-start floor
94
95Every miss observed in prod is structurally unavoidable inside a single
96DO:
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
107To go below this floor we'd need either:
108
1091. **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.
1132. **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
1191. `packages/plugins/mcp/src/sdk/connection-pool.test.ts` — two
120 strict regression tests pinning the per-session pooling contract.
1212. `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
126No 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
24import { describe, expect, it } from "@effect/vitest";
25import { Effect } from "effect";
26import * as http from "node:http";
27import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
28import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js";
29import { z } from "zod";
30
31import { createExecutor, makeTestConfig } from "@executor/sdk";
32
33import { mcpPlugin } from "./plugin";
34
35// ---------------------------------------------------------------------------
36// Test MCP server — counts session connects (each = one cold handshake)
37// ---------------------------------------------------------------------------
38
39function 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
64type 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
71const 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
126const 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
146describe("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
194 const connector = input.resolveConnector();194 const connector = input.resolveConnector();
195 input.pendingConnectors.set(cacheKey, connector);195 input.pendingConnectors.set(cacheKey, connector);
196 196
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
197 const firstConnection = yield* input.connectionCache.get(cacheKey).pipe(204 const firstConnection = yield* input.connectionCache.get(cacheKey).pipe(
198 Effect.withSpan("plugin.mcp.connection.acquire", {205 Effect.withSpan("plugin.mcp.connection.acquire", {
199 attributes: {206 attributes: {
200 "plugin.mcp.transport": transport,207 "plugin.mcp.transport": transport,
201 "plugin.mcp.cache_key": cacheKey,208 "plugin.mcp.cache_key": cacheKey,
202 "plugin.mcp.attempt": 1,209 "plugin.mcp.attempt": 1,
210 "plugin.mcp.cache_hit": cacheHit,
203 },211 },
204 }),212 }),
205 );213 );