History

refactor(cloud): replace hand-rolled JSON-RPC parsing with Effect Schema

`safeParseJson` + `rpcPayloadAttributes` in mcp.ts were ~50 lines of typeof-guards, `as Record<string, unknown>` casts, and two 2KB JSON blob dumps (`mcp.client.info.json`, `mcp.client.capabilities.json`). Replaces the pile with `Schema.Struct`/`Schema.decodeUnknownOption` plus a per-method switch. Unknown input → `Option.none` → no attrs; known methods emit their narrow attr set. Drops the raw JSON blobs — if a client field matters, add a structured attr key instead of dumping 2KB on every span.

Rhys Sullivan committed 6 months ago

Commit 2d28f021 parent 1f26f471 file · +85 −52Browse files

Unified · Split

apps/cloud/src/mcp.ts+85−52
@@ -17,7 +17,7 @@ · 16 unmodified lines
1717 import { env } from "cloudflare:workers";
1818 import { HttpApp, HttpServerRequest, HttpServerResponse } from "@effect/platform";
1919 import * as Sentry from "@sentry/cloudflare";
20−import { Context, Effect, Layer } from "effect";
20+import { Context, Effect, Layer, Option, Schema } from "effect";
2121 import { createRemoteJWKSet, jwtVerify } from "jose";
2222
2323 import { server } from "./env";
@@ -169,61 +169,94 @@ const dumpHeaders = (request: Request): Record<string, string> => { · 145 unmodified lines
169169 return out;
170170 };
171171
172−type JsonRpcRequestLike = {
173− method?: string;
174− id?: string | number;
175− params?: Record<string, unknown>;
176−};
172+// JSON-RPC shapes — narrow to just the fields we fingerprint. Using Schema
173+// collapses the typeof-guard pile and surfaces "what does an MCP client
174+// actually send us" as declarative types. Unknown/malformed input decodes
175+// to None and contributes no span attrs.
177176
178−const safeParseJson = async (request: Request): Promise<JsonRpcRequestLike | null> => {
179− try {
180− const clone = request.clone();
181− const text = await clone.text();
182− if (!text) return null;
183− const parsed = JSON.parse(text) as JsonRpcRequestLike;
184− if (typeof parsed !== "object" || parsed === null) return null;
185− return parsed;
186− } catch {
187− return null;
188− }
189−};
177+const UnknownRecord = Schema.Record({ key: Schema.String, value: Schema.Unknown });
178+
179+const JsonRpcEnvelope = Schema.Struct({
180+ method: Schema.optional(Schema.String),
181+ id: Schema.optional(Schema.Union(Schema.String, Schema.Number, Schema.Null)),
182+ params: Schema.optional(UnknownRecord),
183+});
184+type JsonRpcEnvelope = typeof JsonRpcEnvelope.Type;
185+
186+const InitializeParams = Schema.Struct({
187+ protocolVersion: Schema.optional(Schema.String),
188+ clientInfo: Schema.optional(
189+ Schema.Struct({
190+ name: Schema.optional(Schema.String),
191+ version: Schema.optional(Schema.String),
192+ title: Schema.optional(Schema.String),
193+ }),
194+ ),
195+ capabilities: Schema.optional(UnknownRecord),
196+});
190197
191−const rpcPayloadAttributes = (payload: JsonRpcRequestLike | null): Record<string, unknown> => {
192− if (!payload) return {};
193− const attrs: Record<string, unknown> = {};
194− if (typeof payload.method === "string") attrs["mcp.rpc.method"] = payload.method;
195− if (payload.id !== undefined) attrs["mcp.rpc.id"] = String(payload.id);
196−
197− const params = (payload.params ?? {}) as Record<string, unknown>;
198−
199− if (payload.method === "initialize") {
200− const clientInfo = (params.clientInfo ?? {}) as Record<string, unknown>;
201− const capabilities = (params.capabilities ?? {}) as Record<string, unknown>;
202− if (typeof clientInfo.name === "string") attrs["mcp.client.name"] = clientInfo.name;
203− if (typeof clientInfo.version === "string") attrs["mcp.client.version"] = clientInfo.version;
204− if (typeof clientInfo.title === "string") attrs["mcp.client.title"] = clientInfo.title;
205− if (typeof params.protocolVersion === "string") {
206− attrs["mcp.client.protocol_version"] = params.protocolVersion;
198+const NamedParams = Schema.Struct({ name: Schema.optional(Schema.String) });
199+const UriParams = Schema.Struct({ uri: Schema.optional(Schema.String) });
200+
201+const decode = <A, I>(schema: Schema.Schema<A, I>, input: unknown): Option.Option<A> =>
202+ Schema.decodeUnknownOption(schema)(input);
203+
204+const readJsonRpcEnvelope = (request: Request): Effect.Effect<Option.Option<JsonRpcEnvelope>> =>
205+ Effect.promise(async () => {
206+ try {
207+ const text = await request.clone().text();
208+ if (!text) return Option.none();
209+ return decode(JsonRpcEnvelope, JSON.parse(text));
210+ } catch {
211+ return Option.none();
207212 }
208− attrs["mcp.client.capability.keys"] = Object.keys(capabilities).sort().join(",");
209− // Bounded JSON dumps for ad-hoc inspection. Length caps guard against a
210− // pathological client bloating a single span.
211− attrs["mcp.client.info.json"] = JSON.stringify(clientInfo).slice(0, 2000);
212− attrs["mcp.client.capabilities.json"] = JSON.stringify(capabilities).slice(0, 2000);
213− } else if (payload.method === "tools/call") {
214− const name = params.name;
215− if (typeof name === "string") attrs["mcp.tool.name"] = name;
216− } else if (payload.method === "resources/read" || payload.method === "resources/subscribe") {
217− const uri = params.uri;
218− if (typeof uri === "string") attrs["mcp.resource.uri"] = uri;
219− } else if (payload.method === "prompts/get") {
220− const name = params.name;
221− if (typeof name === "string") attrs["mcp.prompt.name"] = name;
222− }
213+ });
223214
224− return attrs;
215+const methodAttrs = (envelope: JsonRpcEnvelope): Record<string, unknown> => {
216+ const params = envelope.params ?? {};
217+ switch (envelope.method) {
218+ case "initialize":
219+ return Option.match(decode(InitializeParams, params), {
220+ onNone: () => ({}),
221+ onSome: (init) => ({
222+ ...(init.protocolVersion && { "mcp.client.protocol_version": init.protocolVersion }),
223+ ...(init.clientInfo?.name && { "mcp.client.name": init.clientInfo.name }),
224+ ...(init.clientInfo?.version && { "mcp.client.version": init.clientInfo.version }),
225+ ...(init.clientInfo?.title && { "mcp.client.title": init.clientInfo.title }),
226+ "mcp.client.capability.keys": Object.keys(init.capabilities ?? {}).sort().join(","),
227+ }),
228+ });
229+ case "tools/call":
230+ return Option.match(decode(NamedParams, params), {
231+ onNone: () => ({}),
232+ onSome: ({ name }) => (name ? { "mcp.tool.name": name } : {}),
233+ });
234+ case "resources/read":
235+ case "resources/subscribe":
236+ return Option.match(decode(UriParams, params), {
237+ onNone: () => ({}),
238+ onSome: ({ uri }) => (uri ? { "mcp.resource.uri": uri } : {}),
239+ });
240+ case "prompts/get":
241+ return Option.match(decode(NamedParams, params), {
242+ onNone: () => ({}),
243+ onSome: ({ name }) => (name ? { "mcp.prompt.name": name } : {}),
244+ });
245+ default:
246+ return {};
247+ }
225248 };
226249
250+const rpcAttrs = (envelope: Option.Option<JsonRpcEnvelope>): Record<string, unknown> =>
251+ Option.match(envelope, {
252+ onNone: () => ({}),
253+ onSome: (e) => ({
254+ ...(e.method && { "mcp.rpc.method": e.method }),
255+ ...(e.id !== undefined && e.id !== null && { "mcp.rpc.id": String(e.id) }),
256+ ...methodAttrs(e),
257+ }),
258+ });
259+
227260 const annotateMcpRequest = (
228261 request: Request,
229262 opts: { token: VerifiedToken | null; parseBody: boolean },
@@ -251,11 +284,11 @@ const annotateMcpRequest = ( · 21 unmodified lines
251284 ...dumpHeaders(request),
252285 };
253286
254− const payload = opts.parseBody ? yield* Effect.promise(() => safeParseJson(request)) : null;
287+ const envelope = opts.parseBody ? yield* readJsonRpcEnvelope(request) : Option.none();
255288
256289 yield* Effect.annotateCurrentSpan({
257290 ...baseAttrs,
258− ...rpcPayloadAttributes(payload),
291+ ...rpcAttrs(envelope),
259292 });
260293 });
261294