History

Clean up MCP invocation boundaries

Rhys Sullivan committed 6 months ago

Commit 573ed371 parent 10d34f81 file · +42 −73Browse files

Unified · Split

packages/plugins/mcp/src/sdk/invoke.ts+42−73
@@ -10,7 +10,7 @@ · 9 unmodified lines
10// 4. Retrying once on connection failure (invalidate + reconnect).10// 4. Retrying once on connection failure (invalidate + reconnect).
11// ---------------------------------------------------------------------------11// ---------------------------------------------------------------------------
12 12
13import { Cause, Effect, Exit, Schema, ScopedCache } from "effect";13import { Cause, Effect, Exit, Option, Predicate, Schema, ScopedCache } from "effect";
14 14
15import { ElicitRequestSchema } from "@modelcontextprotocol/sdk/types.js";15import { ElicitRequestSchema } from "@modelcontextprotocol/sdk/types.js";
16 16
@@ -29,15 +29,13 @@ import type { McpStoredSourceData } from "./types"; · 12 unmodified lines
29// Helpers29// Helpers
30// ---------------------------------------------------------------------------30// ---------------------------------------------------------------------------
31 31
32const asRecord = (value: unknown): Record<string, unknown> =>32const ArgsRecord = Schema.Record(Schema.String, Schema.Unknown);
33 typeof value === "object" && value !== null && !Array.isArray(value)33const decodeArgsRecord = Schema.decodeUnknownOption(ArgsRecord);
34 ? (value as Record<string, unknown>)
35 : {};
36 34
37const connectionCacheKey = (35const argsRecord = (value: unknown): Record<string, unknown> =>
38 sd: McpStoredSourceData,36 Option.getOrElse(decodeArgsRecord(value), () => ({}));
39 invokerScope: string,37
40): string =>38const connectionCacheKey = (sd: McpStoredSourceData, invokerScope: string): string =>
41 sd.transport === "stdio"39 sd.transport === "stdio"
42 ? `stdio:${sd.command}`40 ? `stdio:${sd.command}`
43 : // Remote sources may resolve per-user secrets (OAuth tokens, header41 : // Remote sources may resolve per-user secrets (OAuth tokens, header
@@ -83,42 +81,34 @@ const toElicitationRequest = (params: McpElicitParams): ElicitationRequest => · 39 unmodified lines
83 requestedSchema: params.requestedSchema,81 requestedSchema: params.requestedSchema,
84 });82 });
85 83
86const installElicitationHandler = (84const installElicitationHandler = (client: McpConnection["client"], elicit: Elicit): void => {
87 client: McpConnection["client"],85 client.setRequestHandler(ElicitRequestSchema, async (request: { params: unknown }) => {
88 elicit: Elicit,86 const params = decodeElicitParams(request.params);
89): void => {87 const req = toElicitationRequest(params);
90 client.setRequestHandler(88 // Use runPromiseExit so we can inspect typed failures — `elicit`
91 ElicitRequestSchema,89 // fails with `ElicitationDeclinedError` on decline/cancel, which
92 async (request: { params: unknown }) => {90 // we translate into the equivalent MCP elicit response instead of
93 const params = decodeElicitParams(request.params);91 // surfacing as a JSON-RPC error.
94 const req = toElicitationRequest(params);92 const exit = await Effect.runPromiseExit(elicit(req));
95 // Use runPromiseExit so we can inspect typed failures — `elicit`93 if (Exit.isSuccess(exit)) {
96 // fails with `ElicitationDeclinedError` on decline/cancel, which94 const response = exit.value;
97 // we translate into the equivalent MCP elicit response instead of95 return {
98 // surfacing as a JSON-RPC error.96 action: response.action,
99 const exit = await Effect.runPromiseExit(elicit(req));97 ...(response.action === "accept" && response.content ? { content: response.content } : {}),
100 if (Exit.isSuccess(exit)) {98 };
101 const response = exit.value;99 }
102 return {100 const failure = exit.cause.reasons.find(Cause.isFailReason);
103 action: response.action,101 if (failure) {
104 ...(response.action === "accept" && response.content102 const err = failure.error;
105 ? { content: response.content }103 if (Predicate.isTagged(err, "ElicitationDeclinedError")) {
106 : {}),104 const action =
107 };105 Predicate.hasProperty(err, "action") && err.action === "cancel" ? "cancel" : "decline";
108 }106 return { action };
109 const failure = exit.cause.reasons.find(Cause.isFailReason);
110 if (failure) {
111 const err = failure.error as {
112 readonly _tag?: string;
113 readonly action?: "decline" | "cancel";
114 };
115 if (err._tag === "ElicitationDeclinedError") {
116 return { action: err.action ?? "decline" };
117 }
118 }107 }
119 throw Cause.squash(exit.cause);108 }
120 },109 // oxlint-disable-next-line executor/no-try-catch-or-throw -- boundary: MCP SDK async request handlers signal unexpected failures by rejecting
121 );110 throw Cause.squash(exit.cause);
111 });
122};112};
123 113
124// ---------------------------------------------------------------------------114// ---------------------------------------------------------------------------
@@ -135,12 +125,10 @@ const useConnection = ( · 10 unmodified lines
135 installElicitationHandler(connection.client, elicit);125 installElicitationHandler(connection.client, elicit);
136 return yield* Effect.tryPromise({126 return yield* Effect.tryPromise({
137 try: () => connection.client.callTool({ name: toolName, arguments: args }),127 try: () => connection.client.callTool({ name: toolName, arguments: args }),
138 catch: (cause) =>128 catch: () =>
139 new McpInvocationError({129 new McpInvocationError({
140 toolName,130 toolName,
141 message: `MCP tool call failed for ${toolName}: ${131 message: `MCP tool call failed for ${toolName}`,
142 cause instanceof Error ? cause.message : String(cause)
143 }`,
144 }),132 }),
145 }).pipe(133 }).pipe(
146 Effect.withSpan("plugin.mcp.client.call_tool", {134 Effect.withSpan("plugin.mcp.client.call_tool", {
@@ -163,15 +151,8 @@ export interface InvokeMcpToolInput { · 16 unmodified lines
163 * collapse multiple users onto one shared connection. */151 * collapse multiple users onto one shared connection. */
164 readonly invokerScope: string;152 readonly invokerScope: string;
165 readonly resolveConnector: () => Effect.Effect<McpConnection, McpConnectionError>;153 readonly resolveConnector: () => Effect.Effect<McpConnection, McpConnectionError>;
166 readonly connectionCache: ScopedCache.ScopedCache<154 readonly connectionCache: ScopedCache.ScopedCache<string, McpConnection, McpConnectionError>;
167 string,155 readonly pendingConnectors: Map<string, Effect.Effect<McpConnection, McpConnectionError>>;
168 McpConnection,
169 McpConnectionError
170 >;
171 readonly pendingConnectors: Map<
172 string,
173 Effect.Effect<McpConnection, McpConnectionError>
174 >;
175 readonly elicit: Elicit;156 readonly elicit: Elicit;
176}157}
177 158
@@ -179,12 +160,10 @@ export const invokeMcpTool = ( · 1 unmodified line
179 input: InvokeMcpToolInput,160 input: InvokeMcpToolInput,
180): Effect.Effect<unknown, McpConnectionError | McpInvocationError> => {161): Effect.Effect<unknown, McpConnectionError | McpInvocationError> => {
181 const transport: string =162 const transport: string =
182 input.sourceData.transport === "stdio"163 input.sourceData.transport === "stdio" ? "stdio" : (input.sourceData.remoteTransport ?? "auto");
183 ? "stdio"
184 : (input.sourceData.remoteTransport ?? "auto");
185 return Effect.gen(function* () {164 return Effect.gen(function* () {
186 const cacheKey = connectionCacheKey(input.sourceData, input.invokerScope);165 const cacheKey = connectionCacheKey(input.sourceData, input.invokerScope);
187 const args = asRecord(input.args);166 const args = argsRecord(input.args);
188 167
189 // Register the connector for the cache lookup (side-channel pattern168 // Register the connector for the cache lookup (side-channel pattern
190 // — the ScopedCache lookup closure reads from `pendingConnectors`).169 // — the ScopedCache lookup closure reads from `pendingConnectors`).
@@ -209,12 +188,7 @@ export const invokeMcpTool = ( · 18 unmodified lines
209 }),188 }),
210 );189 );
211 190
212 return yield* useConnection(191 return yield* useConnection(firstConnection, input.toolName, args, input.elicit).pipe(
213 firstConnection,
214 input.toolName,
215 args,
216 input.elicit,
217 ).pipe(
218 // On failure, invalidate the cache and retry once with a fresh192 // On failure, invalidate the cache and retry once with a fresh
219 // connection. Matches the old invoker's retry-once semantics.193 // connection. Matches the old invoker's retry-once semantics.
220 Effect.catch(() =>194 Effect.catch(() =>
@@ -222,12 +196,7 @@ export const invokeMcpTool = ( · 1 unmodified line
222 yield* ScopedCache.invalidate(input.connectionCache, cacheKey);196 yield* ScopedCache.invalidate(input.connectionCache, cacheKey);
223 input.pendingConnectors.set(cacheKey, connector);197 input.pendingConnectors.set(cacheKey, connector);
224 const fresh = yield* ScopedCache.get(input.connectionCache, cacheKey);198 const fresh = yield* ScopedCache.get(input.connectionCache, cacheKey);
225 return yield* useConnection(199 return yield* useConnection(fresh, input.toolName, args, input.elicit);
226 fresh,
227 input.toolName,
228 args,
229 input.elicit,
230 );
231 }).pipe(200 }).pipe(
232 Effect.withSpan("plugin.mcp.invoke.retry", {201 Effect.withSpan("plugin.mcp.invoke.retry", {
233 attributes: {202 attributes: {