cloud: harden worker→DO trace propagation and explicit sampling

Rhys Sullivan committed 6 months ago

Commit 890f2101 parent 90ed0994 files · +89 −28Browse files

Unified · Split

4 files

apps/cloud/src/env.ts +2 −0
apps/cloud/src/mcp-session.ts +34 −13
apps/cloud/src/mcp.ts +40 −15
apps/cloud/src/server.ts +13 −0

     AXIOM_TOKEN: string;
     AXIOM_DATASET: string;
     AXIOM_TRACES_URL: string;
+    AXIOM_TRACES_SAMPLE_RATIO: string;
   }>;
 type WebEnv = Readonly<Record<string, never>>;
   AXIOM_TOKEN: "",
   AXIOM_DATASET: "executor-cloud",
   AXIOM_TRACES_URL: "https://api.axiom.co/v1/traces",
+  AXIOM_TRACES_SAMPLE_RATIO: "1",
 };
 const SHARED_DEFAULTS: Record<keyof SharedEnv, string> = {

apps/cloud/src/mcp-session.ts+34−13

 // ---------------------------------------------------------------------------
 import { DurableObject, env } from "cloudflare:workers";
+import { createTraceState } from "@opentelemetry/api";
 import { Data, Effect, Layer } from "effect";
 import * as OtelTracer from "@effect/opentelemetry/Tracer";
 import * as Sentry from "@sentry/cloudflare";
   organizationId: string;
 };
+export type IncomingTraceHeaders = {
+  readonly traceparent?: string;
+  readonly tracestate?: string;
+  readonly baggage?: string;
+};
 const HEARTBEAT_MS = 30 * 1000;
 const SESSION_TIMEOUT_MS = 5 * 60 * 1000;
 const TRANSPORT_STATE_KEY = "transport";
     headers: { "content-type": "application/json" },
   });
-// W3C traceparent propagation across the worker→DO boundary. mcp.ts injects
+// W3C propagation across the worker→DO boundary. mcp.ts injects the worker's
-// the worker-side `mcp.request` SpanContext as a `traceparent` header on
+// `traceparent` and forwards incoming `tracestate` / `baggage` headers on
-// forwarded requests (and as a second arg to `init()`). We parse it here and
+// forwarded requests (and as a second arg to `init()`). We parse the context
-// use `OtelTracer.withSpanContext` to stitch the DO's root span under the
+// here and use `OtelTracer.withSpanContext` to stitch the DO's root span
-// worker span so the entire logical request lives in one Axiom trace.
+// under the worker span so the entire logical request lives in one trace.
 const TRACEPARENT_PATTERN = /^([0-9a-f]{2})-([0-9a-f]{32})-([0-9a-f]{16})-([0-9a-f]{2})$/;
 type IncomingSpanContext = {
   readonly traceId: string;
   readonly spanId: string;
   readonly traceFlags: number;
+  readonly traceState?: ReturnType<typeof createTraceState>;
 };
-const parseTraceparent = (value: string | null | undefined): IncomingSpanContext | null => {
+const parseTraceparent = (
+  traceparent: string | null | undefined,
+  tracestate: string | null | undefined,
+): IncomingSpanContext | null => {
+  const value = traceparent;
   if (!value) return null;
   const match = TRACEPARENT_PATTERN.exec(value);
   if (!match) return null;
-  return { traceId: match[2]!, spanId: match[3]!, traceFlags: parseInt(match[4]!, 16) };
+  return {
+    traceId: match[2]!,
+    spanId: match[3]!,
+    traceFlags: parseInt(match[4]!, 16),
+    ...(tracestate ? { traceState: createTraceState(tracestate) } : {}),
+  };
 };
 const withIncomingParent = <A, E, R>(
-  traceparent: string | null | undefined,
+  incoming: IncomingTraceHeaders | null | undefined,
   effect: Effect.Effect<A, E, R>,
 ): Effect.Effect<A, E, R> => {
-  const parsed = parseTraceparent(traceparent);
+  const parsed = parseTraceparent(incoming?.traceparent, incoming?.tracestate);
   return parsed ? OtelTracer.withSpanContext(effect, parsed) : effect;
 };
     });
   }
-  async init(token: McpSessionInit, traceparent?: string): Promise<void> {
+  async init(token: McpSessionInit, incoming?: IncomingTraceHeaders): Promise<void> {
     if (this.initialized) return;
     return Effect.runPromise(
       this.doInitEffect(token).pipe(
         Effect.withSpan("McpSessionDO.init", {
           attributes: { "mcp.auth.organization_id": token.organizationId },
         }),
-        (eff) => withIncomingParent(traceparent, eff),
+        (eff) => withIncomingParent(incoming, eff),
         Effect.provide(DoTelemetryLive),
       ),
     );
     // only (method, session-id presence, response status); rich client
     // fingerprint stays on the edge `mcp.request` span, which shares a
     // trace_id with this one.
-    const traceparent = request.headers.get("traceparent");
+    const incoming = {
+      traceparent: request.headers.get("traceparent") ?? undefined,
+      tracestate: request.headers.get("tracestate") ?? undefined,
+      baggage: request.headers.get("baggage") ?? undefined,
+    } satisfies IncomingTraceHeaders;
     const program = Effect.promise(() => this.dispatchRequest(request)).pipe(
       Effect.tap((response) =>
         Effect.annotateCurrentSpan({
           "mcp.request.session_id_present": !!request.headers.get("mcp-session-id"),
         },
       }),
-      (eff) => withIncomingParent(traceparent, eff),
+      (eff) => withIncomingParent(incoming, eff),
       Effect.provide(DoTelemetryLive),
     );
     return Effect.runPromise(program);

apps/cloud/src/mcp.ts+40−15

 // 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 current
+// to emit a brand-new root trace on every stub call. Ferry the worker span
-// `mcp.request` span's SpanContext across as a W3C `traceparent` string —
+// context across with W3C headers: `traceparent` generated from the active
-// the DO parses it and anchors its own spans under the same trace via
+// Effect span plus passthrough `tracestate` / `baggage` from the inbound
-// `OtelTracer.withSpanContext`.
+// 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 withTraceparentHeader = (request: Request) =>
+const currentPropagationHeaders = (
-  Effect.map(currentTraceparent, (traceparent) => {
+  request: Request,
-    if (!traceparent) return request;
+): Effect.Effect<IncomingPropagationHeaders> =>
-    const headers = new Headers(request.headers);
+  Effect.map(currentTraceparent, (traceparent) => ({
-    headers.set("traceparent", traceparent);
+    traceparent,
-    return new Request(request, { headers });
+    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 });
+};
 /**
  * Forward a request to an existing session DO. Wrapping the DO's `Response`
   Effect.gen(function* () {
     const ns = env.MCP_SESSION;
     const stub = ns.get(ns.idFromString(sessionId));
-    const propagated = yield* withTraceparentHeader(request);
+    const propagation = yield* currentPropagationHeaders(request);
+    const propagated = withPropagationHeaders(request, propagation);
     const raw = yield* Effect.promise(
       () => stub.handleRequest(propagated) as Promise<Response>,
     );
     const ns = env.MCP_SESSION;
     const stub = ns.get(ns.newUniqueId());
-    const traceparent = yield* currentTraceparent;
+    const propagation = yield* currentPropagationHeaders(request);
-    yield* Effect.promise(() => stub.init({ organizationId }, traceparent));
+    yield* Effect.promise(() => stub.init({ organizationId }, propagation));
-    const propagated = yield* withTraceparentHeader(request);
+    const propagated = withPropagationHeaders(request, propagation);
     const raw = yield* Effect.promise(
       () => stub.handleRequest(propagated) as Promise<Response>,
     );

apps/cloud/src/server.ts+13−0

@@ -15,6 +15,12 @@ import { server } from "./env"; · 14 unmodified lines
15	// DOMException "Illegal invocation".	15	// DOMException "Illegal invocation".
16	// ---------------------------------------------------------------------------	16	// ---------------------------------------------------------------------------
17		17
		18	const parseSampleRatio = (value: string): number => {
		19	const n = Number(value);
		20	if (!Number.isFinite(n)) return 1;
		21	return Math.min(1, Math.max(0, n));
		22	};
		23
18	const otelConfig: TraceConfig = {	24	const otelConfig: TraceConfig = {
19	service: { name: "executor-cloud", version: "1.0.0" },	25	service: { name: "executor-cloud", version: "1.0.0" },
20	exporter: {	26	exporter: {
@@ -24,6 +30,13 @@ const otelConfig: TraceConfig = { · 3 unmodified lines
24	"X-Axiom-Dataset": server.AXIOM_DATASET,	30	"X-Axiom-Dataset": server.AXIOM_DATASET,
25	},	31	},
26	},	32	},
		33	sampling: {
		34	headSampler: {
		35	// Keep remote parent decisions and make local sampling policy explicit.
		36	acceptRemote: true,
		37	ratio: parseSampleRatio(server.AXIOM_TRACES_SAMPLE_RATIO),
		38	},
		39	},
27	};	40	};
28		41
29	// otel-cf-workers owns the global TracerProvider. Sentry's OTEL compat shim	42	// otel-cf-workers owns the global TracerProvider. Sentry's OTEL compat shim