From c618b8f7f7dd520e199022111fdb94c1ebee9b8c Mon Sep 17 00:00:00 2001 From: Rhys Sullivan <39114868+RhysSullivan@users.noreply.github.com> Date: Sun, 19 Jul 2026 23:15:06 -0700 Subject: [PATCH] Reuse downstream MCP sessions across tool calls --- .changeset/mcp-session-reuse.md | 5 + e2e/scenarios/mcp-session-state.test.ts | 157 ++++++++++++++++++ .../mcp/src/sdk/connection-pool.test.ts | 156 +++++++++++++++++ .../plugins/mcp/src/sdk/connection-pool.ts | 136 +++++++++++++++ .../plugins/mcp/src/sdk/elicitation.test.ts | 131 ++++++--------- packages/plugins/mcp/src/sdk/errors.ts | 4 + packages/plugins/mcp/src/sdk/invoke.ts | 42 +++-- .../mcp/src/sdk/multi-placement-auth.test.ts | 8 +- packages/plugins/mcp/src/sdk/plugin.ts | 40 ++++- packages/plugins/mcp/src/testing/server.ts | 33 +++- 10 files changed, 615 insertions(+), 97 deletions(-) create mode 100644 .changeset/mcp-session-reuse.md create mode 100644 e2e/scenarios/mcp-session-state.test.ts create mode 100644 packages/plugins/mcp/src/sdk/connection-pool.test.ts create mode 100644 packages/plugins/mcp/src/sdk/connection-pool.ts diff --git a/.changeset/mcp-session-reuse.md b/.changeset/mcp-session-reuse.md new file mode 100644 index 000000000..11b7d32f9 --- /dev/null +++ b/.changeset/mcp-session-reuse.md @@ -0,0 +1,5 @@ +--- +"@executor-js/plugin-mcp": patch +--- + +Reuse downstream MCP sessions across tool calls. Remote MCP invocations now lease connections from a per-plugin-instance pool (one idle session per resolved credential identity, exclusive per invoke, 5-minute idle TTL) instead of dialing a fresh connection per call, so servers that key state by `Mcp-Session-Id` (workspace selection and similar) see consecutive calls in the same session. A reused session rejected with HTTP 404 is redialed once transparently; stdio transports and endpoint probing remain per-call. diff --git a/e2e/scenarios/mcp-session-state.test.ts b/e2e/scenarios/mcp-session-state.test.ts new file mode 100644 index 000000000..21dfa0f08 --- /dev/null +++ b/e2e/scenarios/mcp-session-state.test.ts @@ -0,0 +1,157 @@ +// Cross-target: downstream MCP session continuity — the "an agent can use a +// stateful MCP server" promise. Servers like Render's MCP key state by +// Mcp-Session-Id (select_workspace, then every later call reads the +// selection). If executor dials a fresh downstream connection per tool call, +// every call lands in a brand-new session: select_workspace succeeds, and the +// very next call reports no workspace selected. This scenario drives that +// journey against a session-stateful MCP fixture and asserts state set by one +// tool call is visible to the next call in the same execution. +import { randomBytes } from "node:crypto"; + +import { expect } from "@effect/vitest"; +import { Effect, Predicate } from "effect"; +import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; +import { composePluginApi } from "@executor-js/api/server"; +import { mcpHttpPlugin } from "@executor-js/plugin-mcp/api"; +import { serveMcpServer } from "@executor-js/plugin-mcp/testing"; +import { AuthTemplateSlug, ConnectionName, IntegrationSlug } from "@executor-js/sdk/shared"; + +import { scenario } from "../src/scenario"; +import { Api, Target } from "../src/services"; + +const api = composePluginApi([mcpHttpPlugin()] as const); + +const freshSlug = (prefix: string): string => `${prefix}_${randomBytes(4).toString("hex")}`; + +// A session-stateful MCP server, shaped like Render's workspace selection: +// `select_workspace` stores the choice in the session (each MCP session gets +// its own server instance from the factory, exactly like state keyed by +// Mcp-Session-Id), and `get_workspace` answers from that same session state. +const makeSessionStatefulMcpServer = () => () => { + const server = new McpServer( + { name: "session-stateful-test-server", version: "1.0.0" }, + { capabilities: {} }, + ); + // Per-session: lives in this server instance, created per MCP session. + let selected: string | null = null; + + server.registerTool( + "select_workspace", + { + description: "Selects the workspace for this MCP session", + inputSchema: {}, + }, + async () => { + selected = "ws-e2e"; + return { content: [{ type: "text" as const, text: "selected:ws-e2e" }] }; + }, + ); + + server.registerTool( + "get_workspace", + { + description: "Returns the workspace selected earlier in this MCP session", + inputSchema: {}, + }, + async () => ({ + content: [ + { + type: "text" as const, + text: selected === null ? "no workspace selected" : `workspace:${selected}`, + }, + ], + }), + ); + + return server; +}; + +// One sandbox execution, two dependent tool calls — the smallest agent journey +// that relies on downstream session state. +const selectThenReadCode = (slug: string) => ` +const selected = await tools.${slug}.org.main.select_workspace({}); +const read = await tools.${slug}.org.main.get_workspace({}); +return JSON.stringify({ selected, read }); +`; + +scenario( + "MCP · session state set by one tool call is visible to the next call", + {}, + Effect.scoped( + Effect.gen(function* () { + const target = yield* Target; + const { client: makeApiClient } = yield* Api; + const identity = yield* target.newIdentity(); + const client = yield* makeApiClient(api, identity); + const slug = freshSlug("mcp_state"); + + const server = yield* serveMcpServer(makeSessionStatefulMcpServer()); + + yield* client.mcp.addServer({ + payload: { + transport: "remote", + name: "Session-stateful MCP", + endpoint: server.url, + slug, + remoteTransport: "streamable-http", + }, + }); + + yield* Effect.gen(function* () { + yield* client.connections.create({ + payload: { + owner: "org", + name: ConnectionName.make("main"), + integration: IntegrationSlug.make(slug), + template: AuthTemplateSlug.make("none"), + value: "", + }, + }); + + // Only observe the execution's own traffic, not discovery's. + yield* server.clearRequests; + + const executed = yield* client.executions.execute({ + payload: { code: selectThenReadCode(slug), autoApprove: true }, + }); + expect(executed.status, "the two-call execution completed").toBe("completed"); + + // THE promise: the selection made by the first call is what the second + // call reads. A per-call connection lands the second call in a fresh + // downstream session, which answers "no workspace selected". + expect(executed.text, "the first call selected the workspace").toContain("selected:ws-e2e"); + expect(executed.text, "the second call sees the same session's selection").toContain( + "workspace:ws-e2e", + ); + + // Wire-level corroboration from the fixture's request ledger: every + // session-bound request of this execution carried ONE Mcp-Session-Id. + const sessionIds = new Set( + (yield* server.requests) + .map((request) => request.sessionId) + .filter(Predicate.isNotUndefined), + ); + expect([...sessionIds].length, "both tool calls rode a single downstream MCP session").toBe( + 1, + ); + }).pipe( + Effect.ensuring( + Effect.gen(function* () { + yield* client.connections + .remove({ + params: { + owner: "org", + integration: IntegrationSlug.make(slug), + name: ConnectionName.make("main"), + }, + }) + .pipe(Effect.ignore); + yield* client.mcp + .removeServer({ params: { slug: IntegrationSlug.make(slug) } }) + .pipe(Effect.ignore); + }), + ), + ); + }), + ), +); diff --git a/packages/plugins/mcp/src/sdk/connection-pool.test.ts b/packages/plugins/mcp/src/sdk/connection-pool.test.ts new file mode 100644 index 000000000..64bb0a4eb --- /dev/null +++ b/packages/plugins/mcp/src/sdk/connection-pool.test.ts @@ -0,0 +1,156 @@ +import { describe, expect, it } from "@effect/vitest"; +import { Effect } from "effect"; +import { ElicitationResponse, type Elicit } from "@executor-js/sdk"; + +import { createMcpConnector } from "./connection"; +import { createMcpConnectionPool } from "./connection-pool"; +import { invokeMcpTool } from "./invoke"; +import { makeEchoMcpServer, serveMcpServer } from "../testing"; + +const acceptAll: Elicit = () => + Effect.succeed(ElicitationResponse.make({ action: "accept", content: { approved: true } })); + +const invoke = (input: { + readonly endpoint: string; + readonly pool: ReturnType; + readonly toolName: string; + readonly args: Record; + readonly elicit?: Elicit; +}) => + invokeMcpTool({ + toolId: input.toolName, + toolName: input.toolName, + args: input.args, + transport: "streamable-http", + connector: createMcpConnector({ + transport: "remote", + endpoint: input.endpoint, + remoteTransport: "streamable-http", + }), + connectionPool: input.pool, + connectionPoolKey: input.endpoint, + elicit: input.elicit ?? acceptAll, + }); + +describe("MCP connection pool", () => { + it.effect("reuses one session for sequential invokes", () => + Effect.scoped( + Effect.gen(function* () { + const server = yield* serveMcpServer(() => makeEchoMcpServer()); + const pool = createMcpConnectionPool(); + + const first = yield* invoke({ + endpoint: server.endpoint, + pool, + toolName: "echo", + args: { value: "first" }, + }); + const second = yield* invoke({ + endpoint: server.endpoint, + pool, + toolName: "echo", + args: { value: "second" }, + }); + + expect(first).toMatchObject({ content: [{ type: "text", text: "first" }] }); + expect(second).toMatchObject({ content: [{ type: "text", text: "second" }] }); + expect(server.sessionCount()).toBe(1); + yield* pool.close(); + }), + ), + ); + + it.effect("leases separate sessions to concurrent invokes", () => + Effect.scoped( + Effect.gen(function* () { + const server = yield* serveMcpServer(() => makeEchoMcpServer()); + const pool = createMcpConnectionPool(); + + const results = yield* Effect.all( + [ + invoke({ + endpoint: server.endpoint, + pool, + toolName: "echo", + args: { value: "left" }, + }), + invoke({ + endpoint: server.endpoint, + pool, + toolName: "echo", + args: { value: "right" }, + }), + ], + { concurrency: "unbounded" }, + ); + + expect(results).toEqual([ + expect.objectContaining({ content: [{ type: "text", text: "left" }] }), + expect.objectContaining({ content: [{ type: "text", text: "right" }] }), + ]); + expect(server.sessionCount()).toBe(2); + yield* pool.close(); + }), + ), + ); + + it.effect("redials once when a reused session has expired", () => + Effect.scoped( + Effect.gen(function* () { + const server = yield* serveMcpServer(() => makeEchoMcpServer()); + const pool = createMcpConnectionPool(); + + yield* invoke({ + endpoint: server.endpoint, + pool, + toolName: "echo", + args: { value: "before" }, + }); + yield* server.forgetSessions; + const after = yield* invoke({ + endpoint: server.endpoint, + pool, + toolName: "echo", + args: { value: "after" }, + }); + + expect(after).toMatchObject({ content: [{ type: "text", text: "after" }] }); + expect(server.sessionCount()).toBe(2); + yield* pool.close(); + }), + ), + ); + + it.effect("drops a connection after a transport-level invocation failure", () => + Effect.scoped( + Effect.gen(function* () { + const server = yield* serveMcpServer(() => makeEchoMcpServer()); + const pool = createMcpConnectionPool(); + + yield* invoke({ + endpoint: server.endpoint, + pool, + toolName: "echo", + args: { value: "before" }, + }); + yield* server.rejectNextSessionRequest(500); + yield* invoke({ + endpoint: server.endpoint, + pool, + toolName: "echo", + args: { value: "fails" }, + }).pipe(Effect.flip); + const after = yield* invoke({ + endpoint: server.endpoint, + pool, + toolName: "echo", + args: { value: "after" }, + }); + + expect(after).toMatchObject({ content: [{ type: "text", text: "after" }] }); + expect(server.sessionCount()).toBe(2); + yield* pool.close(); + }), + ), + ); +}); diff --git a/packages/plugins/mcp/src/sdk/connection-pool.ts b/packages/plugins/mcp/src/sdk/connection-pool.ts new file mode 100644 index 000000000..63f636551 --- /dev/null +++ b/packages/plugins/mcp/src/sdk/connection-pool.ts @@ -0,0 +1,136 @@ +import { Cause, Effect, Exit, Predicate } from "effect"; + +import type { McpConnection, McpConnector } from "./connection"; +import type { McpInvocationError } from "./errors"; + +const IDLE_TTL_MS = 5 * 60 * 1_000; + +type IdleConnection = { + readonly connection: McpConnection; + readonly idleSince: number; +}; + +type ConnectionLease = { + readonly connection: McpConnection; + readonly reused: boolean; +}; + +const closeQuietly = (connection: McpConnection): Effect.Effect => + Effect.tryPromise(() => connection.close()).pipe(Effect.ignore); + +const isMcpInvocationError = (error: unknown): error is McpInvocationError => + Predicate.isTagged(error, "McpInvocationError"); + +const isDeadConnectionFailure = (error: unknown): boolean => { + if (Predicate.isTagged(error, "McpConnectionError")) return true; + if (!isMcpInvocationError(error)) return false; + return ( + error.transportFailure === true || + error.status === 400 || + error.status === 404 || + error.status === 408 + ); +}; + +const shouldDropConnection = (exit: Exit.Exit): boolean => { + if (Exit.isSuccess(exit)) return false; + const failures = exit.cause.reasons.filter(Cause.isFailReason); + if (failures.length === 0) return true; + return failures.some((failure) => isDeadConnectionFailure(failure.error)); +}; + +/** A per-plugin-instance pool that gives each invocation an exclusive MCP + * connection lease while retaining at most one idle session per identity. */ +export interface McpConnectionPool { + readonly withConnection: ( + key: string, + connector: McpConnector, + use: (connection: McpConnection) => Effect.Effect, + ) => Effect.Effect, R>; + /** Closes and removes every currently idle connection. Leased connections + * remain owned by their in-flight invocation and follow normal release. */ + readonly close: () => Effect.Effect; +} + +/** Creates an MCP connection pool with lazy five-minute idle eviction and one + * automatic fresh-dial retry for a reused session rejected with HTTP 404. */ +export const createMcpConnectionPool = (): McpConnectionPool => { + const idle = new Map(); + + const acquire = (key: string, connector: McpConnector, forceFresh: boolean) => + Effect.gen(function* () { + if (forceFresh) { + const connection = yield* connector; + return { connection, reused: false } satisfies ConnectionLease; + } + + const taken = yield* Effect.sync(() => { + const entry = idle.get(key); + if (!entry) return { connection: undefined, expired: undefined }; + idle.delete(key); + if (Date.now() - entry.idleSince >= IDLE_TTL_MS) { + return { connection: undefined, expired: entry.connection }; + } + return { connection: entry.connection, expired: undefined }; + }); + + if (taken.expired) yield* closeQuietly(taken.expired); + if (taken.connection) { + return { connection: taken.connection, reused: true } satisfies ConnectionLease; + } + + const connection = yield* connector; + return { connection, reused: false } satisfies ConnectionLease; + }); + + const release = (key: string, lease: ConnectionLease, exit: Exit.Exit) => + Effect.gen(function* () { + if (shouldDropConnection(exit)) { + yield* closeQuietly(lease.connection); + return; + } + + const displaced = yield* Effect.sync(() => { + if (idle.has(key)) return lease.connection; + idle.set(key, { connection: lease.connection, idleSince: Date.now() }); + return undefined; + }); + if (displaced) yield* closeQuietly(displaced); + }); + + const withConnection: McpConnectionPool["withConnection"] = (key, connector, use) => { + let reused = false; + const run = (forceFresh: boolean) => + Effect.acquireUseRelease( + acquire(key, connector, forceFresh), + (lease) => { + if (!forceFresh) reused = lease.reused; + return use(lease.connection); + }, + (lease, exit) => release(key, lease, exit), + ); + + return run(false).pipe( + Effect.catch((error) => + reused && isMcpInvocationError(error) && error.status === 404 + ? run(true) + : Effect.fail(error), + ), + ); + }; + + const close = () => + Effect.gen(function* () { + const connections = yield* Effect.sync(() => { + const connections = [...idle.values()].map((entry) => entry.connection); + idle.clear(); + return connections; + }); + yield* Effect.forEach(connections, closeQuietly, { + concurrency: "unbounded", + discard: true, + }); + }); + + return { withConnection, close }; +}; diff --git a/packages/plugins/mcp/src/sdk/elicitation.test.ts b/packages/plugins/mcp/src/sdk/elicitation.test.ts index 9c0a6e50a..4ca64f32a 100644 --- a/packages/plugins/mcp/src/sdk/elicitation.test.ts +++ b/packages/plugins/mcp/src/sdk/elicitation.test.ts @@ -1,5 +1,5 @@ import { describe, expect, it } from "@effect/vitest"; -import { Effect, Schema } from "effect"; +import { Effect, Predicate, Schema, Semaphore } from "effect"; import { CfWorkerJsonSchemaValidator } from "@modelcontextprotocol/sdk/validation/cfworker"; import type { JsonSchemaType } from "@modelcontextprotocol/sdk/validation/types"; @@ -26,7 +26,10 @@ import { makeElicitationMcpServer, serveMcpServer } from "../testing"; const isFormElicitation = Schema.is(FormElicitation); -const serveElicitationTestServer = serveMcpServer(makeElicitationMcpServer); +const elicitationServerSemaphore = Semaphore.makeUnsafe(1); +const serveElicitationTestServer = Effect.acquireRelease(elicitationServerSemaphore.take(1), () => + elicitationServerSemaphore.release(1).pipe(Effect.asVoid), +).pipe(Effect.flatMap(() => serveMcpServer(makeElicitationMcpServer))); const schemaValidator = new CfWorkerJsonSchemaValidator({ shortcircuit: false }); @@ -56,8 +59,9 @@ const TEMPLATE = AuthTemplateSlug.make("none"); // --------------------------------------------------------------------------- const makeTestExecutor = (serverUrl: string) => - createExecutor( - makeTestConfig({ plugins: [memoryCredentialsPlugin(), mcpPlugin()] as const }), + Effect.acquireRelease( + createExecutor(makeTestConfig({ plugins: [memoryCredentialsPlugin(), mcpPlugin()] as const })), + (executor) => executor.close().pipe(Effect.ignore), ).pipe( Effect.tap((executor) => Effect.gen(function* () { @@ -85,7 +89,7 @@ const findTool = (tools: readonly Tool[], name: string): Tool => // --------------------------------------------------------------------------- describe("MCP elicitation (end-to-end)", () => { - it.effect("form elicitation accepted → tool returns approved result", () => + it.effect("reuses the session while installing each invoke's elicitation handler", () => Effect.gen(function* () { const server = yield* serveElicitationTestServer; const executor = yield* makeTestExecutor(server.url); @@ -93,12 +97,19 @@ describe("MCP elicitation (end-to-end)", () => { const tools = yield* executor.tools.list(); const gatedEcho = findTool(tools, "gated_echo"); - const elicitationMessages: string[] = []; + const acceptedMessages: string[] = []; + const declinedMessages: string[] = []; + let capturedAddress: string | undefined; + let capturedArgs: unknown; + let capturedRequest: unknown; const options: InvokeOptions = { onElicitation: (ctx) => { + capturedAddress = String(ctx.address); + capturedArgs = ctx.args; + capturedRequest = ctx.request; if (isFormElicitation(ctx.request)) { - elicitationMessages.push(ctx.request.message); + acceptedMessages.push(ctx.request.message); } return Effect.succeed( ElicitationResponse.make({ @@ -109,36 +120,50 @@ describe("MCP elicitation (end-to-end)", () => { }, }; - const result = yield* executor.execute(gatedEcho.address, { value: "hello" }, options); - - expect(result).toMatchObject({ - ok: true, - data: { content: [{ type: "text", text: "approved:hello" }] }, - }); - expect(elicitationMessages.length).toBeGreaterThanOrEqual(1); - expect(elicitationMessages.some((m) => m.includes('Approve echo for "hello"?'))).toBe(true); - }), - ); - - it.effect("form elicitation declined → tool returns denied result", () => - Effect.gen(function* () { - const server = yield* serveElicitationTestServer; - const executor = yield* makeTestExecutor(server.url); - const tools = yield* executor.tools.list(); - const gatedEcho = findTool(tools, "gated_echo"); - - const result = yield* executor.execute( + yield* server.clearRequests; + const accepted = yield* executor.execute(gatedEcho.address, { value: "hello" }, options); + const declined = yield* executor.execute( gatedEcho.address, { value: "nope" }, { - onElicitation: () => Effect.succeed(ElicitationResponse.make({ action: "decline" })), + onElicitation: (ctx) => { + if (isFormElicitation(ctx.request)) { + declinedMessages.push(ctx.request.message); + } + return Effect.succeed(ElicitationResponse.make({ action: "decline" })); + }, }, ); + yield* executor.close(); - expect(result).toMatchObject({ + expect(accepted).toMatchObject({ + ok: true, + data: { content: [{ type: "text", text: "approved:hello" }] }, + }); + expect(declined).toMatchObject({ ok: true, data: { content: [{ type: "text", text: "denied:nope" }] }, }); + expect(acceptedMessages).toEqual(['Approve echo for "hello"?']); + expect(declinedMessages).toEqual(['Approve echo for "nope"?']); + expect(capturedAddress).toBe(String(gatedEcho.address)); + expect(capturedArgs).toEqual({ value: "hello" }); + expect(isFormElicitation(capturedRequest)).toBe(true); + const form = capturedRequest as FormElicitation; + expect(form.message).toContain('Approve echo for "hello"?'); + expect(form.requestedSchema).toEqual({ + type: "object", + properties: { + approved: { type: "boolean", title: "Approve" }, + }, + required: ["approved"], + }); + const sessionIds = new Set( + (yield* server.requests) + .map((request) => request.sessionId) + .filter(Predicate.isNotUndefined), + ); + expect(sessionIds.size).toBe(1); }), ); @@ -316,54 +341,4 @@ describe("MCP elicitation (end-to-end)", () => { expect(integration?.description).toBe("Gmail"); }), ); - - it.effect("handler receives correct address, args, and FormElicitation schema", () => - Effect.gen(function* () { - const server = yield* serveElicitationTestServer; - const executor = yield* makeTestExecutor(server.url); - const tools = yield* executor.tools.list(); - const gatedEcho = findTool(tools, "gated_echo"); - - let capturedAddress: string | undefined; - let capturedArgs: unknown; - let capturedRequest: unknown; - - yield* executor.execute( - gatedEcho.address, - { value: "ctx-test" }, - { - onElicitation: (ctx) => { - capturedAddress = String(ctx.address); - capturedArgs = ctx.args; - capturedRequest = ctx.request; - return Effect.succeed( - ElicitationResponse.make({ - action: "accept", - content: { approved: true }, - }), - ); - }, - }, - ); - - expect(capturedAddress).toBe(String(gatedEcho.address)); - expect(capturedArgs).toEqual({ value: "ctx-test" }); - expect(isFormElicitation(capturedRequest)).toBe(true); - - const form = capturedRequest as FormElicitation; - expect(form.message).toContain('Approve echo for "ctx-test"?'); - expect(form.requestedSchema).toEqual({ - type: "object", - properties: { - approved: { type: "boolean", title: "Approve" }, - }, - required: ["approved"], - }); - }), - ); - - // removed: "connection is reused across multiple tool calls" — v2 does not - // cache MCP client connections across invocations (Hyperdrive request-scoped - // rule; each invoke dials and closes its own connection). Session reuse is no - // longer a property of the plugin, so the assertion no longer applies. }); diff --git a/packages/plugins/mcp/src/sdk/errors.ts b/packages/plugins/mcp/src/sdk/errors.ts index 7ea487d14..e3bba2faf 100644 --- a/packages/plugins/mcp/src/sdk/errors.ts +++ b/packages/plugins/mcp/src/sdk/errors.ts @@ -38,6 +38,10 @@ export class McpInvocationError extends Data.TaggedError("McpInvocationError")<{ readonly toolName: string; readonly message: string; readonly status?: number; + /** The SDK failure came from the HTTP/SSE transport rather than a JSON-RPC + * protocol response. Used only to decide whether a leased session is safe + * to return to the idle pool. */ + readonly transportFailure?: boolean; /** The server rejected the call as an unknown tool (protocol error), which * means the persisted catalog has drifted from the server's live tool set. */ readonly unknownTool?: boolean; diff --git a/packages/plugins/mcp/src/sdk/invoke.ts b/packages/plugins/mcp/src/sdk/invoke.ts index e14998fab..3e5aa7695 100644 --- a/packages/plugins/mcp/src/sdk/invoke.ts +++ b/packages/plugins/mcp/src/sdk/invoke.ts @@ -2,9 +2,12 @@ // MCP tool invocation — shared helper called from plugin.invokeTool. // // Responsible for: -// 1. Dialing a fresh MCP client connection for the call (no DB-connection -// caching — request-scoped per the Hyperdrive rule; each invoke acquires -// and releases its own connection). +// 1. Leasing remote connections from a per-plugin-instance pool when one is +// provided. MCP streamable-http is a series of HTTP requests keyed by the +// application-level `Mcp-Session-Id`, not a raw pooled socket, and servers +// legitimately retain state in that session. The pool keeps one idle +// connection per resolved identity and leases it exclusively per invoke; +// stdio and callers without a pool remain strictly per-call. // 2. Installing a per-invocation `ElicitRequestSchema` handler that bridges // MCP's elicit capability into the host's elicit function threaded via // `InvokeToolInput.elicit`. @@ -30,6 +33,7 @@ import { import { McpConnectionError, McpInvocationError, McpOAuthReauthorizationRequired } from "./errors"; import type { McpConnection, McpConnector } from "./connection"; +import type { McpConnectionPool } from "./connection-pool"; import { httpStatusFromCause, insufficientScopeFromCause } from "./http-status"; // --------------------------------------------------------------------------- @@ -181,13 +185,17 @@ const useConnection = ( message: `MCP tool call failed for ${toolName}`, status: 403, insufficientScope: true, + transportFailure: true, }); } const status = httpStatusFromCause(cause); + // oxlint-disable-next-line executor/no-instanceof-tagged-error -- boundary: MCP SDK protocol failures are its McpError subclass; transport failures use other error shapes + const protocolFailure = cause instanceof McpError; return new McpInvocationError({ toolName, message: `MCP tool call failed for ${toolName}`, ...(status === undefined ? {} : { status }), + ...(!protocolFailure ? { transportFailure: true } : {}), ...(isUnknownToolCause(cause, toolName) ? { unknownTool: true } : {}), ...(status === 403 && insufficientScopeFromCause(cause) ? { insufficientScope: true } @@ -211,8 +219,12 @@ export interface InvokeMcpToolInput { readonly toolName: string; readonly args: unknown; readonly transport: string; - /** Dials a fresh connection. The connection is closed after the call. */ + /** Dials a fresh connection when no reusable remote session is available. */ readonly connector: McpConnector; + /** Optional per-plugin-instance pool and resolved remote identity key. Both + * must be present to enable reuse; otherwise the connection is per-call. */ + readonly connectionPool?: McpConnectionPool; + readonly connectionPoolKey?: string; readonly elicit: Elicit; /** Fired when the server sends `notifications/tools/list_changed` during * the call window. Synchronous and non-throwing by contract; the caller @@ -228,6 +240,20 @@ export const invokeMcpTool = ( > => Effect.gen(function* () { const args = argsRecord(input.args); + const use = (connection: McpConnection) => + useConnection(connection, input.toolName, args, input.elicit, input.onToolListChanged); + + if (input.connectionPool && input.connectionPoolKey) { + return yield* input.connectionPool.withConnection( + input.connectionPoolKey, + input.connector.pipe( + Effect.withSpan("plugin.mcp.connection.acquire", { + attributes: { "plugin.mcp.transport": input.transport }, + }), + ), + use, + ); + } const connection = yield* Effect.acquireRelease( input.connector.pipe( @@ -248,13 +274,7 @@ export const invokeMcpTool = ( ), ); - return yield* useConnection( - connection, - input.toolName, - args, - input.elicit, - input.onToolListChanged, - ); + return yield* use(connection); }).pipe( Effect.scoped, Effect.withSpan("plugin.mcp.invoke", { diff --git a/packages/plugins/mcp/src/sdk/multi-placement-auth.test.ts b/packages/plugins/mcp/src/sdk/multi-placement-auth.test.ts index 4729eb39e..590c05273 100644 --- a/packages/plugins/mcp/src/sdk/multi-placement-auth.test.ts +++ b/packages/plugins/mcp/src/sdk/multi-placement-auth.test.ts @@ -169,7 +169,9 @@ describe("MCP multi-placement auth", () => { { marker: "h" }, { onElicitation: "accept-all" }, ); - const headerRequests = (yield* server.requests).slice(beforeHeader); + const headerRequests = (yield* server.requests) + .slice(beforeHeader) + .filter((request) => request.method === "POST"); expect(headerRequests.every((r) => r.authorization === "Bearer header-secret")).toBe(true); expect(headerRequests.every((r) => !r.url.includes("auth_token="))).toBe(true); @@ -179,7 +181,9 @@ describe("MCP multi-placement auth", () => { { marker: "q" }, { onElicitation: "accept-all" }, ); - const queryRequests = (yield* server.requests).slice(beforeQuery); + const queryRequests = (yield* server.requests) + .slice(beforeQuery) + .filter((request) => request.method === "POST"); expect(queryRequests.every((r) => r.url.includes("auth_token=query-secret"))).toBe(true); expect(queryRequests.every((r) => r.authorization === undefined)).toBe(true); }), diff --git a/packages/plugins/mcp/src/sdk/plugin.ts b/packages/plugins/mcp/src/sdk/plugin.ts index 4459f2bc1..623c262a9 100644 --- a/packages/plugins/mcp/src/sdk/plugin.ts +++ b/packages/plugins/mcp/src/sdk/plugin.ts @@ -41,6 +41,7 @@ import { } from "@executor-js/sdk/http-auth"; import { createMcpConnector, type ConnectorInput, type McpConnector } from "./connection"; +import { createMcpConnectionPool } from "./connection-pool"; import { discoverTools } from "./discover"; import { McpConnectionError, @@ -615,6 +616,33 @@ const buildConnectorInput = ( }); }; +// Record fields are key-sorted before stringifying: the key only guards +// same-identity reuse (a mismatch merely costs a fresh dial), but insertion +// order must not make two equal identities look distinct. +const sortedRecord = ( + record: Record | undefined, +): Record => + Object.fromEntries( + Object.entries(record ?? {}) + .filter((entry): entry is [string, string | null] => entry[1] !== undefined) + .sort(([a], [b]) => (a < b ? -1 : a > b ? 1 : 0)), + ); + +const connectionPoolKey = ( + input: Extract, + template: string, + values: Record, +): string => + JSON.stringify({ + endpoint: input.endpoint, + transport: input.transport, + remoteTransport: input.remoteTransport, + headers: sortedRecord(input.headers), + queryParams: sortedRecord(input.queryParams), + template, + values: sortedRecord(values), + }); + // --------------------------------------------------------------------------- // Declared auth methods — project the stored MCP config into the catalog's // plugin-agnostic `AuthMethodDescriptor[]`, one per declared method. Pure and @@ -699,6 +727,7 @@ export interface McpPluginOptions { export const mcpPlugin = definePlugin((options?: McpPluginOptions) => { const allowStdio = options?.dangerouslyAllowStdioMCP ?? false; + const connectionPool = createMcpConnectionPool(); const presetEntries = ( allowStdio @@ -731,6 +760,7 @@ export const mcpPlugin = definePlugin((options?: McpPluginOptions) => { // factory reads `allowStdio` and gates the stdio tab + presets. clientConfig: { allowStdio }, storage: () => ({}), + close: () => connectionPool.close(), extension: (ctx: PluginCtx) => { const httpClientLayer = options?.httpClientLayer ?? ctx.httpClientLayer; @@ -1266,13 +1296,18 @@ export const mcpPlugin = definePlugin((options?: McpPluginOptions) => { } } - const connector: McpConnector = yield* buildConnectorInput( + const connectorInput = yield* buildConnectorInput( parsed, credential.values, String(credential.template), allowStdio, options?.httpClientLayer ?? ctx.httpClientLayer, - ).pipe(Effect.map((ci) => createMcpConnector(ci))); + ); + const connector: McpConnector = createMcpConnector(connectorInput); + const poolKey = + connectorInput.transport === "remote" + ? connectionPoolKey(connectorInput, String(credential.template), credential.values) + : undefined; const connectionRef = { owner: credential.owner, @@ -1293,6 +1328,7 @@ export const mcpPlugin = definePlugin((options?: McpPluginOptions) => { args, transport, connector, + ...(poolKey === undefined ? {} : { connectionPool, connectionPoolKey: poolKey }), elicit, onToolListChanged: () => { toolListChanged = true; diff --git a/packages/plugins/mcp/src/testing/server.ts b/packages/plugins/mcp/src/testing/server.ts index 4e864666a..0b6edf0f0 100644 --- a/packages/plugins/mcp/src/testing/server.ts +++ b/packages/plugins/mcp/src/testing/server.ts @@ -12,6 +12,10 @@ export type McpTestServer = { readonly sessionCount: () => number; readonly requests: Effect.Effect; readonly clearRequests: Effect.Effect; + /** Drops all server-side session registrations without notifying clients. */ + readonly forgetSessions: Effect.Effect; + /** Rejects the next request carrying an MCP session id with this status. */ + readonly rejectNextSessionRequest: (status: number) => Effect.Effect; }; export type McpTestRequest = { @@ -64,9 +68,11 @@ export const serveMcpServer = (factory: () => McpServer, options: McpTestServerO Effect.acquireRelease( Effect.gen(function* () { const transports = new Map(); + const allTransports = new Set(); const requests = yield* Ref.make([]); const path = options.path ?? "/"; let sessions = 0; + let nextSessionRequestStatus: number | undefined; const handleMcpRequest = ( request: http.IncomingMessage, @@ -130,6 +136,13 @@ export const serveMcpServer = (factory: () => McpServer, options: McpTestServerO } } + if (sessionId && request.method === "POST" && nextSessionRequestStatus !== undefined) { + const status = nextSessionRequestStatus; + nextSessionRequestStatus = undefined; + writeText(response, status, `Forced HTTP ${status}`); + return; + } + const existingTransport = sessionId ? transports.get(sessionId) : undefined; if (sessionId && !existingTransport) { writeText(response, 404, "Session not found"); @@ -150,6 +163,7 @@ export const serveMcpServer = (factory: () => McpServer, options: McpTestServerO transports.set(sid, transport); }, }); + allTransports.add(transport); sessions += 1; const mcpServer = factory(); @@ -202,21 +216,32 @@ export const serveMcpServer = (factory: () => McpServer, options: McpTestServerO sessionCount: () => sessions, requests: Ref.get(requests), clearRequests: Ref.set(requests, []), + forgetSessions: Effect.sync(() => transports.clear()), + rejectNextSessionRequest: (status: number) => + Effect.sync(() => { + nextSessionRequestStatus = status; + }), close: Effect.gen(function* () { - for (const transport of transports.values()) { + for (const transport of allTransports) { yield* Effect.tryPromise({ try: () => transport.close(), catch: (cause) => new McpTestServerError({ cause }), }).pipe(Effect.ignore); } - yield* Effect.sync(() => { - nodeServer.close(); + yield* Effect.callback((resume) => { + nodeServer.close(() => resume(Effect.void)); nodeServer.closeAllConnections?.(); }); + // closeAllConnections destroys sockets out from under in-flight SDK + // reads; give those rejections a beat to settle before the scope + // ends so they surface inside the test, not as unhandled rejections. + // A raw timer, not Effect.sleep: this runs inside it.effect tests + // whose TestClock never advances wall-clock sleeps. + yield* Effect.promise(() => new Promise((resolve) => setTimeout(resolve, 25))); }), }; }), - (server) => server.close, + (server) => server.close.pipe(Effect.ignore), ).pipe(Effect.map(({ close: _close, ...server }) => server)); export const serveMcpServerWithOAuth = (