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
1010 // 4. Retrying once on connection failure (invalidate + reconnect).
1111 // ---------------------------------------------------------------------------
1212
13−import { Cause, Effect, Exit, Schema, ScopedCache } from "effect";
13+import { Cause, Effect, Exit, Option, Predicate, Schema, ScopedCache } from "effect";
1414
1515 import { ElicitRequestSchema } from "@modelcontextprotocol/sdk/types.js";
1616
@@ -29,15 +29,13 @@ import type { McpStoredSourceData } from "./types"; · 12 unmodified lines
2929 // Helpers
3030 // ---------------------------------------------------------------------------
3131
32−const asRecord = (value: unknown): Record<string, unknown> =>
33− typeof value === "object" && value !== null && !Array.isArray(value)
34− ? (value as Record<string, unknown>)
35− : {};
32+const ArgsRecord = Schema.Record(Schema.String, Schema.Unknown);
33+const decodeArgsRecord = Schema.decodeUnknownOption(ArgsRecord);
3634
37−const connectionCacheKey = (
38− sd: McpStoredSourceData,
39− invokerScope: string,
40−): string =>
35+const argsRecord = (value: unknown): Record<string, unknown> =>
36+ Option.getOrElse(decodeArgsRecord(value), () => ({}));
37+
38+const connectionCacheKey = (sd: McpStoredSourceData, invokerScope: string): string =>
4139 sd.transport === "stdio"
4240 ? `stdio:${sd.command}`
4341 : // Remote sources may resolve per-user secrets (OAuth tokens, header
@@ -83,42 +81,34 @@ const toElicitationRequest = (params: McpElicitParams): ElicitationRequest => · 39 unmodified lines
8381 requestedSchema: params.requestedSchema,
8482 });
8583
86−const installElicitationHandler = (
87− client: McpConnection["client"],
88− elicit: Elicit,
89−): void => {
90− client.setRequestHandler(
91− ElicitRequestSchema,
92− async (request: { params: unknown }) => {
93− const params = decodeElicitParams(request.params);
94− const req = toElicitationRequest(params);
95− // Use runPromiseExit so we can inspect typed failures — `elicit`
96− // fails with `ElicitationDeclinedError` on decline/cancel, which
97− // we translate into the equivalent MCP elicit response instead of
98− // surfacing as a JSON-RPC error.
99− const exit = await Effect.runPromiseExit(elicit(req));
100− if (Exit.isSuccess(exit)) {
101− const response = exit.value;
102− return {
103− action: response.action,
104− ...(response.action === "accept" && response.content
105− ? { content: response.content }
106− : {}),
107− };
108− }
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− }
84+const installElicitationHandler = (client: McpConnection["client"], elicit: Elicit): void => {
85+ client.setRequestHandler(ElicitRequestSchema, async (request: { params: unknown }) => {
86+ const params = decodeElicitParams(request.params);
87+ const req = toElicitationRequest(params);
88+ // Use runPromiseExit so we can inspect typed failures — `elicit`
89+ // fails with `ElicitationDeclinedError` on decline/cancel, which
90+ // we translate into the equivalent MCP elicit response instead of
91+ // surfacing as a JSON-RPC error.
92+ const exit = await Effect.runPromiseExit(elicit(req));
93+ if (Exit.isSuccess(exit)) {
94+ const response = exit.value;
95+ return {
96+ action: response.action,
97+ ...(response.action === "accept" && response.content ? { content: response.content } : {}),
98+ };
99+ }
100+ const failure = exit.cause.reasons.find(Cause.isFailReason);
101+ if (failure) {
102+ const err = failure.error;
103+ if (Predicate.isTagged(err, "ElicitationDeclinedError")) {
104+ const action =
105+ Predicate.hasProperty(err, "action") && err.action === "cancel" ? "cancel" : "decline";
106+ return { action };
118107 }
119− throw Cause.squash(exit.cause);
120− },
121− );
108+ }
109+ // oxlint-disable-next-line executor/no-try-catch-or-throw -- boundary: MCP SDK async request handlers signal unexpected failures by rejecting
110+ throw Cause.squash(exit.cause);
111+ });
122112 };
123113
124114 // ---------------------------------------------------------------------------
@@ -135,12 +125,10 @@ const useConnection = ( · 10 unmodified lines
135125 installElicitationHandler(connection.client, elicit);
136126 return yield* Effect.tryPromise({
137127 try: () => connection.client.callTool({ name: toolName, arguments: args }),
138− catch: (cause) =>
128+ catch: () =>
139129 new McpInvocationError({
140130 toolName,
141− message: `MCP tool call failed for ${toolName}: ${
142− cause instanceof Error ? cause.message : String(cause)
143− }`,
131+ message: `MCP tool call failed for ${toolName}`,
144132 }),
145133 }).pipe(
146134 Effect.withSpan("plugin.mcp.client.call_tool", {
@@ -163,15 +151,8 @@ export interface InvokeMcpToolInput { · 16 unmodified lines
163151 * collapse multiple users onto one shared connection. */
164152 readonly invokerScope: string;
165153 readonly resolveConnector: () => Effect.Effect<McpConnection, McpConnectionError>;
166− readonly connectionCache: ScopedCache.ScopedCache<
167− string,
168− McpConnection,
169− McpConnectionError
170− >;
171− readonly pendingConnectors: Map<
172− string,
173− Effect.Effect<McpConnection, McpConnectionError>
174− >;
154+ readonly connectionCache: ScopedCache.ScopedCache<string, McpConnection, McpConnectionError>;
155+ readonly pendingConnectors: Map<string, Effect.Effect<McpConnection, McpConnectionError>>;
175156 readonly elicit: Elicit;
176157 }
177158
@@ -179,12 +160,10 @@ export const invokeMcpTool = ( · 1 unmodified line
179160 input: InvokeMcpToolInput,
180161 ): Effect.Effect<unknown, McpConnectionError | McpInvocationError> => {
181162 const transport: string =
182− input.sourceData.transport === "stdio"
183− ? "stdio"
184− : (input.sourceData.remoteTransport ?? "auto");
163+ input.sourceData.transport === "stdio" ? "stdio" : (input.sourceData.remoteTransport ?? "auto");
185164 return Effect.gen(function* () {
186165 const cacheKey = connectionCacheKey(input.sourceData, input.invokerScope);
187− const args = asRecord(input.args);
166+ const args = argsRecord(input.args);
188167
189168 // Register the connector for the cache lookup (side-channel pattern
190169 // — the ScopedCache lookup closure reads from `pendingConnectors`).
@@ -209,12 +188,7 @@ export const invokeMcpTool = ( · 18 unmodified lines
209188 }),
210189 );
211190
212− return yield* useConnection(
213− firstConnection,
214− input.toolName,
215− args,
216− input.elicit,
217− ).pipe(
191+ return yield* useConnection(firstConnection, input.toolName, args, input.elicit).pipe(
218192 // On failure, invalidate the cache and retry once with a fresh
219193 // connection. Matches the old invoker's retry-once semantics.
220194 Effect.catch(() =>
@@ -222,12 +196,7 @@ export const invokeMcpTool = ( · 1 unmodified line
222196 yield* ScopedCache.invalidate(input.connectionCache, cacheKey);
223197 input.pendingConnectors.set(cacheKey, connector);
224198 const fresh = yield* ScopedCache.get(input.connectionCache, cacheKey);
225− return yield* useConnection(
226− fresh,
227− input.toolName,
228− args,
229− input.elicit,
230− );
199+ return yield* useConnection(fresh, input.toolName, args, input.elicit);
231200 }).pipe(
232201 Effect.withSpan("plugin.mcp.invoke.retry", {
233202 attributes: {