History

feat(mcp): effect-native connection pool with ScopedCache, connection reuse test

- Replace hand-rolled Map pool with Effect ScopedCache (TTL, acquireRelease) - Connections cached per source (remote:endpoint / stdio:command) - Auto-reconnect on stale connection (invalidate + retry) - Plugin close tears down cache scope - useMcpConnection follows Effect 'use' pattern with withSpan tracing - Add e2e test verifying session reuse across multiple tool calls - Remove McpClientLike/ElicitableClient — use Client directly - Clean up elicitation bridge with Schema parsing - Clean up manifest.ts (Schema decode, no typeof soup) - Deduplicate in-memory stores (makeInMemoryScopedKv in SDK) - Separate KV namespaces per concern (no prefix hacks) - Discriminated union for ConnectorInput (derives from schema types) - Trim SDK index to only externally-consumed exports - Clean up oauth.ts (shared discovery state, callAuth helper)

Rhys Sullivan committed 7 months ago

Commit aa467b91 parent 528f2c43 files · +309 −144Browse files

Unified · Split

3 files
  1. packages/plugins/mcp/src/sdk/elicitation.test.ts +63 −4
  2. packages/plugins/mcp/src/sdk/invoke.ts +199 −131
  3. packages/plugins/mcp/src/sdk/plugin.ts +47 −9
packages/plugins/mcp/src/sdk/elicitation.test.ts+63−4
@@ -74,10 +74,17 @@ function createTestMcpServer() { · 73 unmodified lines
7474 return server;
7575 }
7676
77+type TestServer = {
78+ readonly url: string;
79+ readonly httpServer: http.Server;
80+ /** Number of MCP sessions created (each connect = 1 session) */
81+ readonly sessionCount: () => number;
82+};
83+
7784 const serveMcpServer = Effect.acquireRelease(
78− Effect.async<{ url: string; httpServer: http.Server }, Error>((resume) => {
85+ Effect.async<TestServer, Error>((resume) => {
7986 const transports = new Map<string, StreamableHTTPServerTransport>();
80− const servers: McpServer[] = [];
87+ let sessions = 0;
8188
8289 const httpServer = http.createServer(async (req, res) => {
8390 const sessionId = req.headers["mcp-session-id"] as string | undefined;
@@ -95,7 +102,7 @@ const serveMcpServer = Effect.acquireRelease( · 11 unmodified lines
95102
96103 // New session — create a fresh McpServer per connection
97104 const mcpServer = createTestMcpServer();
98− servers.push(mcpServer);
105+ sessions++;
99106
100107 const transport = new StreamableHTTPServerTransport({
101108 sessionIdGenerator: () => crypto.randomUUID(),
@@ -111,7 +118,11 @@ const serveMcpServer = Effect.acquireRelease( · 9 unmodified lines
111118 httpServer.listen(0, () => {
112119 const addr = httpServer.address();
113120 const port = typeof addr === "object" && addr ? addr.port : 0;
114− resume(Effect.succeed({ url: `http://127.0.0.1:${port}`, httpServer }));
121+ resume(Effect.succeed({
122+ url: `http://127.0.0.1:${port}`,
123+ httpServer,
124+ sessionCount: () => sessions,
125+ }));
115126 });
116127 }),
117128 ({ httpServer }) =>
@@ -297,4 +308,52 @@ describe("MCP elicitation (end-to-end)", () => { · 179 unmodified lines
297308 });
298309 }),
299310 );
311+
312+ it.scoped(
313+ "connection is reused across multiple tool calls to the same source",
314+ () =>
315+ Effect.gen(function* () {
316+ const server = yield* serveMcpServer;
317+ const executor = yield* makeTestExecutor(server.url);
318+ const tools = yield* executor.tools.list();
319+ const simpleEcho = tools.find((t) => t.name === "simple_echo")!;
320+ const gatedEcho = tools.find((t) => t.name === "gated_echo")!;
321+
322+ // addSource created 1 session during discovery
323+ const sessionsAfterAdd = server.sessionCount();
324+
325+ // First tool call — may create a new session (discovery used a
326+ // different connection that was closed)
327+ yield* executor.tools.invoke(
328+ simpleEcho.id,
329+ { value: "call-1" },
330+ { onElicitation: "accept-all" },
331+ );
332+ const sessionsAfterFirst = server.sessionCount();
333+
334+ // Second call to a different tool on the same source — should reuse
335+ yield* executor.tools.invoke(
336+ simpleEcho.id,
337+ { value: "call-2" },
338+ { onElicitation: "accept-all" },
339+ );
340+ expect(server.sessionCount()).toBe(sessionsAfterFirst);
341+
342+ // Third call to yet another tool on the same source — still reused
343+ yield* executor.tools.invoke(
344+ gatedEcho.id,
345+ { value: "call-3" },
346+ {
347+ onElicitation: () =>
348+ Effect.succeed(
349+ new ElicitationResponse({
350+ action: "accept",
351+ content: { approved: true },
352+ }),
353+ ),
354+ },
355+ );
356+ expect(server.sessionCount()).toBe(sessionsAfterFirst);
357+ }),
358+ );
300359 });
packages/plugins/mcp/src/sdk/invoke.ts+199−131
@@ -2,7 +2,7 @@ · 1 unmodified line
22 // MCP tool invoker — bridges the binding store + MCP client into SDK invoker
33 // ---------------------------------------------------------------------------
44
5−import { Effect, Schema } from "effect";
5+import { Effect, Schema, type ScopedCache } from "effect";
66
77 import type { OAuthClientProvider } from "@modelcontextprotocol/sdk/client/auth.js";
88 import type { OAuthTokens } from "@modelcontextprotocol/sdk/shared/auth.js";
@@ -48,6 +48,10 @@ type Secrets = { · 39 unmodified lines
4848 ) => Effect.Effect<string, unknown>;
4949 };
5050
51+// ---------------------------------------------------------------------------
52+// OAuth provider factory
53+// ---------------------------------------------------------------------------
54+
5155 const makeOAuthProvider = (
5256 accessToken: string,
5357 tokenType: string,
@@ -57,8 +61,8 @@ const makeOAuthProvider = ( · 3 unmodified lines
5761 get clientMetadata() {
5862 return {
5963 redirect_uris: ["http://localhost/oauth/callback"],
60− grant_types: ["authorization_code", "refresh_token"],
61− response_types: ["code"],
64+ grant_types: ["authorization_code", "refresh_token"] as string[],
65+ response_types: ["code"] as string[],
6266 token_endpoint_auth_method: "none" as const,
6367 client_name: "Executor",
6468 };
@@ -81,24 +85,23 @@ const makeOAuthProvider = ( · 16 unmodified lines
8185 });
8286
8387 // ---------------------------------------------------------------------------
84−// Elicitation bridge — MCP elicit request ↔ SDK ElicitationHandler
88+// Elicitation bridge
8589 // ---------------------------------------------------------------------------
8690
87−const McpFormElicitParams = Schema.Struct({
88− mode: Schema.optional(Schema.Literal("form")),
89− message: Schema.String,
90− requestedSchema: Schema.Record({ key: Schema.String, value: Schema.Unknown }),
91−});
92−
93−const McpUrlElicitParams = Schema.Struct({
94− mode: Schema.Literal("url"),
95− message: Schema.String,
96− url: Schema.String,
97− elicitationId: Schema.optional(Schema.String),
98− id: Schema.optional(Schema.String),
99−});
100−
101−const McpElicitParams = Schema.Union(McpUrlElicitParams, McpFormElicitParams);
91+const McpElicitParams = Schema.Union(
92+ Schema.Struct({
93+ mode: Schema.Literal("url"),
94+ message: Schema.String,
95+ url: Schema.String,
96+ elicitationId: Schema.optional(Schema.String),
97+ id: Schema.optional(Schema.String),
98+ }),
99+ Schema.Struct({
100+ mode: Schema.optional(Schema.Literal("form")),
101+ message: Schema.String,
102+ requestedSchema: Schema.Record({ key: Schema.String, value: Schema.Unknown }),
103+ }),
104+);
102105 type McpElicitParams = typeof McpElicitParams.Type;
103106
104107 const decodeElicitParams = Schema.decodeUnknownSync(McpElicitParams);
@@ -124,18 +127,18 @@ const installElicitationHandler = ( · 19 unmodified lines
124127 client.setRequestHandler(
125128 ElicitRequestSchema,
126129 async (request: { params: unknown }) => {
127− const params = decodeElicitParams(request.params);
128− const response = await Effect.runPromise(
129− handler({ toolId, args, request: toElicitationRequest(params) }),
130− );
131−
132− return {
133− action: response.action,
134− ...(response.action === "accept" && response.content
135− ? { content: response.content }
136− : {}),
137− };
138− });
130+ const params = decodeElicitParams(request.params);
131+ const response = await Effect.runPromise(
132+ handler({ toolId, args, request: toElicitationRequest(params) }),
133+ );
134+ return {
135+ action: response.action,
136+ ...(response.action === "accept" && response.content
137+ ? { content: response.content }
138+ : {}),
139+ };
140+ },
141+ );
139142 };
140143
141144 // ---------------------------------------------------------------------------
@@ -221,24 +224,14 @@ const resolveConnectorInput = ( · 79 unmodified lines
221224 });
222225 };
223226
224−/** Execute a single MCP tool call */
225−const callMcpTool = (
226− connection: McpConnection,
227− toolName: string,
228− args: Record<string, unknown>,
229−): Effect.Effect<unknown, ToolInvocationError> =>
230− Effect.tryPromise({
231− try: () =>
232− connection.client.callTool({ name: toolName, arguments: args }),
233− catch: (cause) =>
234− new ToolInvocationError({
235− toolId: "" as ToolId,
236− message: `MCP tool call failed for ${toolName}: ${
237− cause instanceof Error ? cause.message : String(cause)
238− }`,
239− cause,
240− }),
241− });
227+// ---------------------------------------------------------------------------
228+// Connection cache key
229+// ---------------------------------------------------------------------------
230+
231+const connectionCacheKey = (sourceData: McpStoredSourceData): string =>
232+ sourceData.transport === "stdio"
233+ ? `stdio:${sourceData.command}`
234+ : `remote:${sourceData.endpoint}`;
242235
243236 // ---------------------------------------------------------------------------
244237 // Resolve elicitation handler from options
@@ -251,6 +244,34 @@ const resolveElicitationHandler = ( · 6 unmodified lines
251244 ? () => Effect.succeed(new ElicitationResponse({ action: "accept" }))
252245 : options.onElicitation;
253246
247+// ---------------------------------------------------------------------------
248+// Use pattern — wrap MCP Client as an Effect service
249+// ---------------------------------------------------------------------------
250+
251+const useMcpConnection = (
252+ connection: McpConnection,
253+ toolId: ToolId,
254+ toolName: string,
255+ args: Record<string, unknown>,
256+ handler: ElicitationHandler,
257+): Effect.Effect<unknown, ToolInvocationError> =>
258+ Effect.gen(function* () {
259+ installElicitationHandler(connection.client, toolId, args, handler);
260+
261+ return yield* Effect.tryPromise({
262+ try: () =>
263+ connection.client.callTool({ name: toolName, arguments: args }),
264+ catch: (cause) =>
265+ new ToolInvocationError({
266+ toolId,
267+ message: `MCP tool call failed for ${toolName}: ${
268+ cause instanceof Error ? cause.message : String(cause)
269+ }`,
270+ cause,
271+ }),
272+ });
273+ }).pipe(Effect.withSpan(`mcp.callTool.${toolName}`));
274+
254275 // ---------------------------------------------------------------------------
255276 // Public API
256277 // ---------------------------------------------------------------------------
@@ -259,93 +280,140 @@ export const makeMcpInvoker = (opts: { · 2 unmodified lines
259280 readonly bindingStore: McpBindingStore;
260281 readonly secrets: Secrets;
261282 readonly scopeId: ScopeId;
262−}): ToolInvoker => ({
263− resolveAnnotations: () =>
264− Effect.succeed(
265− new ToolAnnotations({
266− requiresApproval: false,
267− }),
268− ),
269−
270− invoke: (toolId: ToolId, args: unknown, options: InvokeOptions) =>
271− Effect.gen(function* () {
272− const entry = yield* opts.bindingStore.get(toolId);
273− if (!entry) {
274− return yield* new ToolInvocationError({
275− toolId,
276− message: `No MCP binding found for tool "${toolId}"`,
277− cause: undefined,
278− });
279− }
283+ readonly connectionCache: ScopedCache.ScopedCache<string, McpConnection, McpConnectionError>;
284+ /** Shared map between cache lookup and invoker — set connector before cache.get */
285+ readonly pendingConnectors: Map<string, Effect.Effect<McpConnection, McpConnectionError>>;
286+}): ToolInvoker & { readonly closeConnections: () => Effect.Effect<void> } => {
287+ const { connectionCache, pendingConnectors } = opts;
288+
289+ return {
290+ resolveAnnotations: () =>
291+ Effect.succeed(
292+ new ToolAnnotations({ requiresApproval: false }),
293+ ),
294+
295+ invoke: (toolId: ToolId, args: unknown, options: InvokeOptions) =>
296+ Effect.gen(function* () {
297+ const entry = yield* opts.bindingStore.get(toolId);
298+ if (!entry) {
299+ return yield* new ToolInvocationError({
300+ toolId,
301+ message: `No MCP binding found for tool "${toolId}"`,
302+ cause: undefined,
303+ });
304+ }
280305
281− const { binding, sourceData } = entry;
306+ const { binding, sourceData } = entry;
307+ const cacheKey = connectionCacheKey(sourceData);
282308
283− const connector = Effect.gen(function* () {
284− const ci = yield* resolveConnectorInput(sourceData, opts.secrets, opts.scopeId);
285− return yield* createMcpConnector(ci);
286− }).pipe(
287− Effect.mapError((err) =>
288− new McpConnectionError({
289− transport: "auto",
290− message: err instanceof Error ? err.message : String(err),
291− }),
292− ),
293− );
309+ // Build the connector and register it for the cache lookup
310+ const connector = resolveConnectorInput(
311+ sourceData,
312+ opts.secrets,
313+ opts.scopeId,
314+ ).pipe(
315+ Effect.flatMap((ci) => createMcpConnector(ci)),
316+ Effect.mapError(
317+ (err) =>
318+ new McpConnectionError({
319+ transport: "auto",
320+ message: err instanceof Error ? err.message : String(err),
321+ }),
322+ ),
323+ );
324+ pendingConnectors.set(cacheKey, connector);
325+
326+ const connection = yield* connectionCache.get(cacheKey).pipe(
327+ Effect.mapError(
328+ (err) =>
329+ new ToolInvocationError({
330+ toolId,
331+ message: `Failed connecting to MCP server: ${
332+ err instanceof Error ? err.message : String(err)
333+ }`,
334+ cause: err,
335+ }),
336+ ),
337+ );
294338
295− const connection = yield* connector.pipe(
296− Effect.mapError(
297− (err) =>
339+ const elicitationHandler = resolveElicitationHandler(options);
340+
341+ return yield* useMcpConnection(
342+ connection,
343+ toolId,
344+ binding.toolName,
345+ asRecord(args),
346+ elicitationHandler,
347+ ).pipe(
348+ // On failure, invalidate the cached connection and retry once
349+ Effect.catchAll((err) =>
350+ Effect.gen(function* () {
351+ yield* connectionCache.invalidate(cacheKey);
352+ pendingConnectors.set(cacheKey, connector);
353+
354+ const freshConnection = yield* connectionCache
355+ .get(cacheKey)
356+ .pipe(
357+ Effect.mapError(
358+ (retryErr) =>
359+ new ToolInvocationError({
360+ toolId,
361+ message: `Failed reconnecting: ${
362+ retryErr instanceof Error
363+ ? retryErr.message
364+ : String(retryErr)
365+ }`,
366+ cause: retryErr,
367+ }),
368+ ),
369+ );
370+
371+ return yield* useMcpConnection(
372+ freshConnection,
373+ toolId,
374+ binding.toolName,
375+ asRecord(args),
376+ elicitationHandler,
377+ );
378+ }),
379+ ),
380+ );
381+ }).pipe(
382+ Effect.scoped,
383+ Effect.map((callResult) => {
384+ const resultRecord = asRecord(callResult);
385+ const isError = resultRecord.isError === true;
386+ return new ToolInvocationResult({
387+ data: isError ? null : (callResult ?? null),
388+ error: isError ? callResult : null,
389+ });
390+ }),
391+ Effect.catchAll((err) => {
392+ if (
393+ typeof err === "object" &&
394+ err !== null &&
395+ "_tag" in err &&
396+ (err as { _tag: string })._tag === "ToolInvocationError"
397+ ) {
398+ return Effect.fail(err as ToolInvocationError);
399+ }
400+ return Effect.fail(
298401 new ToolInvocationError({
299402 toolId,
300− message: `Failed connecting to MCP server: ${err.message}`,
403+ message: `MCP invocation failed: ${
404+ err instanceof Error ? err.message : String(err)
405+ }`,
301406 cause: err,
302407 }),
303− ),
304− );
305−
306− // Install elicitation handler before calling the tool
307− const elicitationHandler = resolveElicitationHandler(options);
308− installElicitationHandler(
309− connection.client,
310− toolId,
311− args,
312− elicitationHandler,
313− );
314−
315− const result = yield* callMcpTool(
316− connection,
317− binding.toolName,
318− asRecord(args),
319− );
320−
321− yield* Effect.promise(() => connection.close().catch(() => {}));
322−
323− const resultRecord = asRecord(result);
324− const isError = resultRecord.isError === true;
408+ );
409+ }),
410+ ),
325411
326− return new ToolInvocationResult({
327− data: isError ? null : (result ?? null),
328− error: isError ? result : null,
329− });
330− }).pipe(
331− Effect.catchAll((err) => {
332− if (
333− typeof err === "object" &&
334− err !== null &&
335− "_tag" in err &&
336− (err as { _tag: string })._tag === "ToolInvocationError"
337− ) {
338− return Effect.fail(err as ToolInvocationError);
339− }
340− return Effect.fail(
341− new ToolInvocationError({
342− toolId,
343− message: `MCP invocation failed: ${
344− err instanceof Error ? err.message : String(err)
345− }`,
346− cause: err,
347− }),
348− );
349− }),
350− ),
351−});
412+ closeConnections: () =>
413+ Effect.sync(() => {
414+ pendingConnectors.clear();
415+ }).pipe(
416+ Effect.flatMap(() => connectionCache.invalidateAll),
417+ ),
418+ };
419+};
packages/plugins/mcp/src/sdk/plugin.ts+47−9
@@ -1,4 +1,4 @@
1−import { Effect } from "effect";
1+import { Effect, Exit, ScopedCache, Duration, Scope } from "effect";
22
33 import type { OAuthClientProvider } from "@modelcontextprotocol/sdk/client/auth.js";
44 import type { OAuthTokens } from "@modelcontextprotocol/sdk/shared/auth.js";
@@ -23,8 +23,10 @@ import { · 18 unmodified lines
2323 } from "./binding-store";
2424 import {
2525 createMcpConnector,
26+ type McpConnection,
2627 type ConnectorInput,
2728 } from "./connection";
29+import { McpConnectionError } from "./errors";
2830 import {
2931 startMcpOAuthAuthorization,
3032 exchangeMcpOAuthCode,
@@ -235,14 +237,48 @@ export const mcpPlugin = (options?: { · 204 unmodified lines
235237 key: "mcp",
236238 init: (ctx: PluginContext) =>
237239 Effect.gen(function* () {
238− yield* ctx.tools.registerInvoker(
239− "mcp",
240− makeMcpInvoker({
241− bindingStore,
242− secrets: ctx.secrets,
243− scopeId: ctx.scope.id,
244− }),
245− );
240+ // Create a long-lived scope for the connection cache
241+ const cacheScope = yield* Scope.make();
242+
243+ // Side-channel for deferred connector lookup (populated before
244+ // each cache.get call so the lookup knows how to connect)
245+ const pendingConnectors = new Map<
246+ string,
247+ Effect.Effect<McpConnection, McpConnectionError>
248+ >();
249+
250+ // ScopedCache: keyed by source identity, acquireRelease manages
251+ // connection lifecycle, TTL evicts idle connections.
252+ const connectionCache = yield* ScopedCache.make({
253+ lookup: (key: string) =>
254+ Effect.acquireRelease(
255+ Effect.suspend(() => {
256+ const connector = pendingConnectors.get(key);
257+ if (!connector) {
258+ return Effect.fail(
259+ new McpConnectionError({
260+ transport: "auto",
261+ message: `No pending connector for key: ${key}`,
262+ }),
263+ );
264+ }
265+ return connector;
266+ }),
267+ (connection) =>
268+ Effect.promise(() => connection.close().catch(() => {})),
269+ ),
270+ capacity: 64,
271+ timeToLive: Duration.minutes(5),
272+ }).pipe(Scope.extend(cacheScope));
273+
274+ const invoker = makeMcpInvoker({
275+ bindingStore,
276+ secrets: ctx.secrets,
277+ scopeId: ctx.scope.id,
278+ connectionCache,
279+ pendingConnectors,
280+ });
281+ yield* ctx.tools.registerInvoker("mcp", invoker);
246282
247283 // Restore source metadata
248284 const savedMetas = yield* bindingStore.listSourceMeta();
@@ -700,6 +736,8 @@ export const mcpPlugin = (options?: { · 451 unmodified lines
700736
701737 close: () =>
702738 Effect.gen(function* () {
739+ yield* invoker.closeConnections();
740+ yield* Scope.close(cacheScope, Exit.void);
703741 for (const sourceId of addedSources.keys()) {
704742 yield* ctx.tools.unregisterBySource(sourceId);
705743 }