`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.
Unified · Split
| @@ -17,7 +17,7 @@ · 16 unmodified lines | |||
| 17 | import { env } from "cloudflare:workers"; | 17 | import { env } from "cloudflare:workers"; |
| 18 | import { HttpApp, HttpServerRequest, HttpServerResponse } from "@effect/platform"; | 18 | import { HttpApp, HttpServerRequest, HttpServerResponse } from "@effect/platform"; |
| 19 | import * as Sentry from "@sentry/cloudflare"; | 19 | import * as Sentry from "@sentry/cloudflare"; |
| 20 | import { Context, Effect, Layer } from "effect"; | 20 | import { Context, Effect, Layer, Option, Schema } from "effect"; |
| 21 | import { createRemoteJWKSet, jwtVerify } from "jose"; | 21 | import { createRemoteJWKSet, jwtVerify } from "jose"; |
| 22 | 22 | ||
| 23 | import { server } from "./env"; | 23 | import { server } from "./env"; |
| @@ -169,61 +169,94 @@ const dumpHeaders = (request: Request): Record<string, string> => { · 145 unmodified lines | |||
| 169 | return out; | 169 | return out; |
| 170 | }; | 170 | }; |
| 171 | 171 | ||
| 172 | type JsonRpcRequestLike = { | 172 | // JSON-RPC shapes — narrow to just the fields we fingerprint. Using Schema |
| 173 | method?: string; | 173 | // collapses the typeof-guard pile and surfaces "what does an MCP client |
| 174 | id?: string | number; | 174 | // actually send us" as declarative types. Unknown/malformed input decodes |
| 175 | params?: Record<string, unknown>; | 175 | // to None and contributes no span attrs. |
| 176 | }; | ||
| 177 | 176 | ||
| 178 | const safeParseJson = async (request: Request): Promise<JsonRpcRequestLike | null> => { | 177 | const UnknownRecord = Schema.Record({ key: Schema.String, value: Schema.Unknown }); |
| 179 | try { | 178 | |
| 180 | const clone = request.clone(); | 179 | const JsonRpcEnvelope = Schema.Struct({ |
| 181 | const text = await clone.text(); | 180 | method: Schema.optional(Schema.String), |
| 182 | if (!text) return null; | 181 | id: Schema.optional(Schema.Union(Schema.String, Schema.Number, Schema.Null)), |
| 183 | const parsed = JSON.parse(text) as JsonRpcRequestLike; | 182 | params: Schema.optional(UnknownRecord), |
| 184 | if (typeof parsed !== "object" || parsed === null) return null; | 183 | }); |
| 185 | return parsed; | 184 | type JsonRpcEnvelope = typeof JsonRpcEnvelope.Type; |
| 186 | } catch { | 185 | |
| 187 | return null; | 186 | const InitializeParams = Schema.Struct({ |
| 188 | } | 187 | protocolVersion: Schema.optional(Schema.String), |
| 189 | }; | 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 | }); | ||
| 190 | 197 | ||
| 191 | const rpcPayloadAttributes = (payload: JsonRpcRequestLike | null): Record<string, unknown> => { | 198 | const NamedParams = Schema.Struct({ name: Schema.optional(Schema.String) }); |
| 192 | if (!payload) return {}; | 199 | const UriParams = Schema.Struct({ uri: Schema.optional(Schema.String) }); |
| 193 | const attrs: Record<string, unknown> = {}; | 200 | |
| 194 | if (typeof payload.method === "string") attrs["mcp.rpc.method"] = payload.method; | 201 | const decode = <A, I>(schema: Schema.Schema<A, I>, input: unknown): Option.Option<A> => |
| 195 | if (payload.id !== undefined) attrs["mcp.rpc.id"] = String(payload.id); | 202 | Schema.decodeUnknownOption(schema)(input); |
| 196 | 203 | ||
| 197 | const params = (payload.params ?? {}) as Record<string, unknown>; | 204 | const readJsonRpcEnvelope = (request: Request): Effect.Effect<Option.Option<JsonRpcEnvelope>> => |
| 198 | 205 | Effect.promise(async () => { | |
| 199 | if (payload.method === "initialize") { | 206 | try { |
| 200 | const clientInfo = (params.clientInfo ?? {}) as Record<string, unknown>; | 207 | const text = await request.clone().text(); |
| 201 | const capabilities = (params.capabilities ?? {}) as Record<string, unknown>; | 208 | if (!text) return Option.none(); |
| 202 | if (typeof clientInfo.name === "string") attrs["mcp.client.name"] = clientInfo.name; | 209 | return decode(JsonRpcEnvelope, JSON.parse(text)); |
| 203 | if (typeof clientInfo.version === "string") attrs["mcp.client.version"] = clientInfo.version; | 210 | } catch { |
| 204 | if (typeof clientInfo.title === "string") attrs["mcp.client.title"] = clientInfo.title; | 211 | return Option.none(); |
| 205 | if (typeof params.protocolVersion === "string") { | ||
| 206 | attrs["mcp.client.protocol_version"] = params.protocolVersion; | ||
| 207 | } | 212 | } |
| 208 | attrs["mcp.client.capability.keys"] = Object.keys(capabilities).sort().join(","); | 213 | }); |
| 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 | } | ||
| 223 | 214 | ||
| 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 | } | ||
| 225 | }; | 248 | }; |
| 226 | 249 | ||
| 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 | |||
| 227 | const annotateMcpRequest = ( | 260 | const annotateMcpRequest = ( |
| 228 | request: Request, | 261 | request: Request, |
| 229 | opts: { token: VerifiedToken | null; parseBody: boolean }, | 262 | opts: { token: VerifiedToken | null; parseBody: boolean }, |
| @@ -251,11 +284,11 @@ const annotateMcpRequest = ( · 21 unmodified lines | |||
| 251 | ...dumpHeaders(request), | 284 | ...dumpHeaders(request), |
| 252 | }; | 285 | }; |
| 253 | 286 | ||
| 254 | const payload = opts.parseBody ? yield* Effect.promise(() => safeParseJson(request)) : null; | 287 | const envelope = opts.parseBody ? yield* readJsonRpcEnvelope(request) : Option.none(); |
| 255 | 288 | ||
| 256 | yield* Effect.annotateCurrentSpan({ | 289 | yield* Effect.annotateCurrentSpan({ |
| 257 | ...baseAttrs, | 290 | ...baseAttrs, |
| 258 | ...rpcPayloadAttributes(payload), | 291 | ...rpcAttrs(envelope), |
| 259 | }); | 292 | }); |
| 260 | }); | 293 | }); |
| 261 | 294 | ||