748 lines · 27.8 KB · TypeScript
RawDownload
// ---------------------------------------------------------------------------// Cloud MCP handler — Effect-native HTTP app for /mcp + /.well-known/*// ---------------------------------------------------------------------------//// Built on Effect v4's unstable HTTP `HttpEffect.toWebHandler`. start.ts's// mcpRequestMiddleware calls `mcpFetch` and falls through to `next()` when it// returns `null` (non-MCP path) so TanStack Start keeps routing.//// Streaming passthrough — the MCP session Durable Object returns a `Response`// whose body is a `ReadableStream` (SSE). We wrap that `Response` in// `HttpServerResponse.raw(response)`; the platform's `toWeb` conversion// recognises `body.body instanceof Response` and returns it as-is (only// merging headers we set on the outer response, which is none), so the// underlying `ReadableStream` passes through untouched.// --------------------------------------------------------------------------- import * as Cloudflare from "alchemy/Cloudflare/Workers/Runtime";import { HttpEffect, HttpServerRequest, HttpServerResponse } from "effect/unstable/http";import { Cause, Context, Effect, Layer, Option, Predicate, Result, Schema } from "effect"; import { createCachedRemoteJWKSet } from "./jwks-cache";import { captureCause } from "./observability";import { TelemetryLive } from "./services/telemetry";import {  McpJwtVerificationError,  verifyWorkOSMcpAccessToken,  type VerifiedToken,} from "./mcp-auth";import { authorizeOrganization } from "./auth/authorize-organization";import { UserStoreService } from "./auth/context";import { CoreSharedServices } from "./api/core-shared-services";import { DbService } from "./services/db";import { peekAndAnnotate } from "./mcp/response-peek";import {  authTemporarilyUnavailable,  CORS_ALLOW_ORIGIN,  jsonResponse,  jsonRpcError,  unauthorized,} from "./mcp/responses"; // ---------------------------------------------------------------------------// Constants// --------------------------------------------------------------------------- // Module-scope cache survives across MCP requests within the same worker// isolate. AuthKit's JWKS rotates on the order of hours/days, so a 1h TTL// dominates the upstream cooldown without sacrificing rotation safety —// `createCachedRemoteJWKSet` force-refreshes on key-not-found inside its// resolver. Production telemetry showed ~222 fetches/8h with p99 1.7s on// the previous default-cooldown setup; this collapses that to ~1 per// isolate-hour.const jwksByAuthkitDomain = new Map<string, ReturnType<typeof createCachedRemoteJWKSet>>(); const getJwks = (authkitDomain: string) => {  const existing = jwksByAuthkitDomain.get(authkitDomain);  if (existing) return existing;  const next = createCachedRemoteJWKSet(new URL(`${authkitDomain}/oauth2/jwks`));  jwksByAuthkitDomain.set(authkitDomain, next);  return next;}; const BEARER_PREFIX = "Bearer ";const INTERNAL_ACCOUNT_ID_HEADER = "x-executor-mcp-account-id";const INTERNAL_ORGANIZATION_ID_HEADER = "x-executor-mcp-organization-id"; const CORS_PREFLIGHT_HEADERS = {  ...CORS_ALLOW_ORIGIN,  "access-control-allow-methods": "GET, POST, DELETE, OPTIONS",  "access-control-allow-headers":    "authorization, content-type, mcp-session-id, accept, mcp-protocol-version",  "access-control-expose-headers": "mcp-session-id",} as const; const MCP_PATH = "/mcp";const PROTECTED_RESOURCE_METADATA_PATH = "/.well-known/oauth-protected-resource/mcp"; const getMcpConfig = Effect.gen(function* () {  const env = yield* Cloudflare.WorkerEnvironment.typed<Env>();  const authkitDomain = env.MCP_AUTHKIT_DOMAIN ?? "https://signin.executor.sh";  const resourceOrigin = env.MCP_RESOURCE_ORIGIN ?? "https://executor.sh";  return {    authkitDomain,    resourceOrigin,    workosClientId: env.WORKOS_CLIENT_ID,    protectedResourceMetadataUrl: `${resourceOrigin}${PROTECTED_RESOURCE_METADATA_PATH}`,    resourceUrl: `${resourceOrigin}${MCP_PATH}`,    jwks: getJwks(authkitDomain),  };}); type McpUnauthorizedReason = "missing_bearer" | "invalid_token"; type McpAuthorizedResult = {  readonly _tag: "Authorized";  readonly token: VerifiedToken;}; type McpUnauthorizedResult = {  readonly _tag: "Unauthorized";  readonly reason: McpUnauthorizedReason;  readonly description?: string;}; export type McpAuthResult = McpAuthorizedResult | McpUnauthorizedResult; export const mcpAuthorized = (token: VerifiedToken): McpAuthorizedResult => ({  _tag: "Authorized",  token,}); export const mcpUnauthorized = (  reason: McpUnauthorizedReason,  description?: string,): McpUnauthorizedResult => ({  _tag: "Unauthorized",  reason,  description,}); const corsPreflight = HttpServerResponse.empty({  status: 204,  headers: CORS_PREFLIGHT_HEADERS,}); // ---------------------------------------------------------------------------// Auth// --------------------------------------------------------------------------- export class McpAuth extends Context.Service<  McpAuth,  {    readonly verifyBearer: (      request: Request,    ) => Effect.Effect<McpAuthResult, McpJwtVerificationError, Cloudflare.WorkerEnvironment>;  }>()("@executor-js/cloud/McpAuth") {} export class McpOrganizationAuth extends Context.Service<  McpOrganizationAuth,  {    readonly authorize: (      accountId: string,      organizationId: string,    ) => Effect.Effect<boolean, unknown, Cloudflare.WorkerEnvironment>;  }>()("@executor-js/cloud/McpOrganizationAuth") {} const verifyJwt = (token: string) =>  Effect.gen(function* () {    const config = yield* getMcpConfig;    return yield* verifyWorkOSMcpAccessToken(token, config.jwks, {      issuer: config.authkitDomain,      audience: config.workosClientId,    });  }); const DbLive = DbService.Live;const UserStoreLive = UserStoreService.Live.pipe(Layer.provide(DbLive));const McpOrganizationAuthServices = Layer.mergeAll(DbLive, UserStoreLive, CoreSharedServices); export const McpOrganizationAuthLive = Layer.succeed(McpOrganizationAuth)({  authorize: (accountId, organizationId) =>    authorizeOrganization(accountId, organizationId).pipe(      Effect.map((org) => org !== null),      Effect.provide(McpOrganizationAuthServices),    ),}); export const McpAuthLive = Layer.succeed(McpAuth)({  verifyBearer: Effect.fn("mcp.auth.verify_bearer")(function* (request) {    const authHeader = request.headers.get("authorization");    if (!authHeader?.startsWith(BEARER_PREFIX)) {      yield* Effect.annotateCurrentSpan({ "mcp.auth.outcome": "missing_bearer" });      return mcpUnauthorized("missing_bearer");    }    const verified = yield* verifyJwt(authHeader.slice(BEARER_PREFIX.length)).pipe(      Effect.catchTag("McpJwtVerificationError", (error) => {        if (error.reason === "system") return Effect.fail(error);        return Effect.gen(function* () {          yield* Effect.annotateCurrentSpan({            "mcp.auth.outcome": "invalid",            "mcp.auth.invalid_reason": error.reason,          });          return mcpUnauthorized(            "invalid_token",            error.reason === "expired" ? "The access token expired" : "The access token is invalid",          );        });      }),    );    if (!verified) return mcpUnauthorized("invalid_token", "The access token is invalid");    if (Predicate.isTagged(verified, "Unauthorized")) return verified;    if (!verified.accountId) {      yield* Effect.annotateCurrentSpan({ "mcp.auth.outcome": "missing_subject" });      return mcpUnauthorized("invalid_token", "The access token is invalid");    }    yield* Effect.annotateCurrentSpan({      "mcp.auth.outcome": "verified",      "mcp.auth.has_organization": !!verified.organizationId,    });    return mcpAuthorized(verified);  }),}); // ---------------------------------------------------------------------------// Client fingerprint capture// ---------------------------------------------------------------------------// Annotates the current Effect span with everything we can learn about a// connecting MCP client: the parsed JSON-RPC body, whitelisted request// headers, CF request metadata, and verified-JWT claims. Runs before dispatch// so unauthorized requests still get fingerprinted.// --------------------------------------------------------------------------- type CfRequestMetadata = {  country?: string;  city?: string;  region?: string;  timezone?: string;  asn?: number;  asOrganization?: string;  tlsVersion?: string;  tlsCipher?: string;  httpProtocol?: string;  colo?: string;}; const requestWithCf = (request: Request): Request & { cf?: CfRequestMetadata } =>  request as Request & { cf?: CfRequestMetadata }; const getCfMeta = (request: Request): CfRequestMetadata => requestWithCf(request).cf ?? {}; const HEADERS_TO_DUMP = [  "accept",  "accept-encoding",  "accept-language",  "cache-control",  "content-type",  "mcp-protocol-version",  "origin",  "referer",  "sec-fetch-dest",  "sec-fetch-mode",  "sec-fetch-site",  "user-agent",  "x-client-name",  "x-client-version",  "x-requested-with",] as const; const dumpHeaders = (request: Request): Record<string, string> => {  const out: Record<string, string> = {};  for (const name of HEADERS_TO_DUMP) {    const value = request.headers.get(name);    if (value !== null) out[`mcp.http.header.${name}`] = value;  }  const authHeader = request.headers.get("authorization");  if (authHeader) {    out["mcp.http.header.authorization.scheme"] = authHeader.split(" ", 1)[0] ?? "";    out["mcp.http.header.authorization.length"] = String(authHeader.length);  }  // Record the full header name list too — surfaces anything unexpected  // without us having to enumerate every possibility up front.  out["mcp.http.header.names"] = Array.from(request.headers.keys()).sort().join(",");  return out;}; // JSON-RPC shapes — narrow to just the fields we fingerprint. Using Schema// collapses the typeof-guard pile and surfaces "what does an MCP client// actually send us" as declarative types. Unknown/malformed input decodes// to None and contributes no span attrs. const UnknownRecord = Schema.Record(Schema.String, Schema.Unknown); const JsonRpcEnvelope = Schema.Struct({  method: Schema.optional(Schema.String),  id: Schema.optional(Schema.Union([Schema.String, Schema.Number, Schema.Null])),  params: Schema.optional(UnknownRecord),  // Responses to server-initiated requests arrive as POST bodies too —  // notably elicitation replies (`result.action = "accept" | "decline" | "cancel"`).  result: Schema.optional(UnknownRecord),});type JsonRpcEnvelope = typeof JsonRpcEnvelope.Type; const ElicitationReplyResult = Schema.Struct({  action: Schema.optional(Schema.Literals(["accept", "decline", "cancel"])),}); const InitializeParams = Schema.Struct({  protocolVersion: Schema.optional(Schema.String),  clientInfo: Schema.optional(    Schema.Struct({      name: Schema.optional(Schema.String),      version: Schema.optional(Schema.String),      title: Schema.optional(Schema.String),    }),  ),  capabilities: Schema.optional(UnknownRecord),}); const NamedParams = Schema.Struct({ name: Schema.optional(Schema.String) });const UriParams = Schema.Struct({ uri: Schema.optional(Schema.String) }); const decodeJsonRpcEnvelopeString = Schema.decodeUnknownOption(  Schema.fromJsonString(JsonRpcEnvelope),);const decodeInitializeParams = Schema.decodeUnknownOption(InitializeParams);const decodeNamedParams = Schema.decodeUnknownOption(NamedParams);const decodeUriParams = Schema.decodeUnknownOption(UriParams);const decodeElicitationReplyResult = Schema.decodeUnknownOption(ElicitationReplyResult); const isMcpAuthorized = (value: McpAuthResult): value is McpAuthorizedResult =>  Predicate.isTagged(value, "Authorized");const isMcpUnauthorized = (value: McpAuthResult): value is McpUnauthorizedResult =>  Predicate.isTagged(value, "Unauthorized"); const readJsonRpcEnvelope = (request: Request): Effect.Effect<Option.Option<JsonRpcEnvelope>> =>  Effect.tryPromise({    try: () => request.clone().text(),    catch: () => undefined,  }).pipe(    Effect.map((text) => (text ? decodeJsonRpcEnvelopeString(text) : Option.none())),    Effect.catchCause(() => Effect.succeed(Option.none())),    Effect.withSpan("mcp.request.read_json_rpc"),  ); const methodAttrs = (envelope: JsonRpcEnvelope): Record<string, unknown> => {  const params = envelope.params ?? {};  switch (envelope.method) {    case "initialize":      return Option.match(decodeInitializeParams(params), {        onNone: () => ({}),        onSome: (init) => ({          ...(init.protocolVersion && { "mcp.client.protocol_version": init.protocolVersion }),          ...(init.clientInfo?.name && { "mcp.client.name": init.clientInfo.name }),          ...(init.clientInfo?.version && { "mcp.client.version": init.clientInfo.version }),          ...(init.clientInfo?.title && { "mcp.client.title": init.clientInfo.title }),          "mcp.client.capability.keys": Object.keys(init.capabilities ?? {})            .sort()            .join(","),        }),      });    case "tools/call":      return Option.match(decodeNamedParams(params), {        onNone: () => ({}),        onSome: ({ name }) => (name ? { "mcp.tool.name": name } : {}),      });    case "resources/read":    case "resources/subscribe":      return Option.match(decodeUriParams(params), {        onNone: () => ({}),        onSome: ({ uri }) => (uri ? { "mcp.resource.uri": uri } : {}),      });    case "prompts/get":      return Option.match(decodeNamedParams(params), {        onNone: () => ({}),        onSome: ({ name }) => (name ? { "mcp.prompt.name": name } : {}),      });    default:      return {};  }}; const replyAttrs = (envelope: JsonRpcEnvelope): Record<string, unknown> => {  if (!envelope.result || envelope.method) return {};  return Option.match(decodeElicitationReplyResult(envelope.result), {    onNone: () => ({}),    onSome: ({ action }) => (action ? { "mcp.elicitation.action": action } : {}),  });}; const rpcAttrs = (envelope: Option.Option<JsonRpcEnvelope>): Record<string, unknown> =>  Option.match(envelope, {    onNone: () => ({}),    onSome: (e) => ({      ...(e.method && { "mcp.rpc.method": e.method }),      ...(e.id !== undefined && e.id !== null && { "mcp.rpc.id": String(e.id) }),      ...methodAttrs(e),      ...replyAttrs(e),    }),  }); const annotateMcpRequest = (  request: Request,  opts: { token: VerifiedToken | null; parseBody: boolean },): Effect.Effect<void> =>  Effect.gen(function* () {    const cf = getCfMeta(request);    const baseAttrs: Record<string, unknown> = {      "mcp.request.method": request.method,      "mcp.request.session_id_present": !!request.headers.get("mcp-session-id"),      "mcp.request.session_id": request.headers.get("mcp-session-id") ?? "",      "mcp.auth.has_bearer": (request.headers.get("authorization") ?? "").startsWith(BEARER_PREFIX),      "mcp.auth.verified": !!opts.token,      "mcp.auth.organization_id": opts.token?.organizationId ?? "",      "mcp.auth.account_id": opts.token?.accountId ?? "",      "cf.country": cf.country ?? "",      "cf.city": cf.city ?? "",      "cf.region": cf.region ?? "",      "cf.timezone": cf.timezone ?? "",      "cf.asn": cf.asn ?? 0,      "cf.as_organization": cf.asOrganization ?? "",      "cf.tls_version": cf.tlsVersion ?? "",      "cf.tls_cipher": cf.tlsCipher ?? "",      "cf.http_protocol": cf.httpProtocol ?? "",      "cf.colo": cf.colo ?? "",      ...dumpHeaders(request),    };     const envelope = opts.parseBody ? yield* readJsonRpcEnvelope(request) : Option.none();    const attrs = {      ...baseAttrs,      ...rpcAttrs(envelope),      "mcp.request.parse_body": opts.parseBody,    };     yield* Effect.annotateCurrentSpan(attrs);    yield* Effect.annotateCurrentSpan(attrs).pipe(Effect.withSpan("mcp.request.annotate"));  }); // ---------------------------------------------------------------------------// OAuth metadata endpoints// --------------------------------------------------------------------------- const protectedResourceMetadata = Effect.gen(function* () {  const config = yield* getMcpConfig;  return jsonResponse({    resource: config.resourceUrl,    authorization_servers: [config.authkitDomain],    bearer_methods_supported: ["header"],    scopes_supported: [],  });}); const authorizationServerMetadata = Effect.gen(function* () {  const config = yield* getMcpConfig;  return yield* Effect.tryPromise({    try: async () => {      const res = await fetch(`${config.authkitDomain}/.well-known/oauth-authorization-server`);      if (!res.ok) return jsonResponse({ error: "upstream_error" }, 502);      return jsonResponse(await res.json());    },    catch: () => undefined,  });}).pipe(Effect.catchCause(() => Effect.succeed(jsonResponse({ error: "upstream_error" }, 502)))); // ---------------------------------------------------------------------------// DO dispatch// --------------------------------------------------------------------------- // Worker and DO run in separate isolates with independent WebSdk tracer// providers. Neither one can see the other's OTEL context, so the DO used// to emit a brand-new root trace on every stub call. Ferry the worker span// context across with W3C headers: `traceparent` generated from the active// Effect span plus passthrough `tracestate` / `baggage` from the inbound// request.type IncomingPropagationHeaders = {  readonly traceparent?: string;  readonly tracestate?: string;  readonly baggage?: string;}; const currentTraceparent = Effect.map(Effect.currentSpan, (span) => {  if (!span || !span.traceId || !span.spanId) return undefined;  const flags = span.sampled ? "01" : "00";  return `00-${span.traceId}-${span.spanId}-${flags}`;}).pipe(Effect.orElseSucceed(() => undefined)); const currentPropagationHeaders = (request: Request): Effect.Effect<IncomingPropagationHeaders> =>  Effect.map(currentTraceparent, (traceparent) => ({    traceparent,    tracestate: request.headers.get("tracestate") ?? undefined,    baggage: request.headers.get("baggage") ?? undefined,  })); const withPropagationHeaders = (  request: Request,  propagation: IncomingPropagationHeaders,): Request => {  const headers = new Headers(request.headers);  if (propagation.traceparent) {    headers.set("traceparent", propagation.traceparent);  }  if (propagation.tracestate) {    headers.set("tracestate", propagation.tracestate);  }  if (propagation.baggage) {    headers.set("baggage", propagation.baggage);  }  return new Request(request, { headers });}; const withVerifiedIdentityHeaders = (request: Request, token: VerifiedToken): Request => {  const headers = new Headers(request.headers);  headers.set(INTERNAL_ACCOUNT_ID_HEADER, token.accountId);  headers.set(INTERNAL_ORGANIZATION_ID_HEADER, token.organizationId ?? "");  return new Request(request, { headers });}; const withMcpResponseHeaders = (response: Response): Response => {

Lines 1–500 of 748.