cloud: harden worker→DO trace propagation and explicit sampling

Author: RhysSullivan

apps/cloud/src/env.ts

@@ -36,6 +36,7 @@ type ServerEnv = SharedEnv &
     AXIOM_TOKEN: string;
     AXIOM_DATASET: string;
     AXIOM_TRACES_URL: string;
+    AXIOM_TRACES_SAMPLE_RATIO: string;
   }>;
 
 type WebEnv = Readonly<Record<string, never>>;
@@ -55,6 +56,7 @@ const SERVER_DEFAULTS: Record<keyof ServerEnv, string> = {
   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

@@ -3,6 +3,7 @@
 // ---------------------------------------------------------------------------
 
 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";
@@ -36,6 +37,12 @@ export type McpSessionInit = {
   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";
@@ -59,31 +66,41 @@ const jsonRpcError = (status: number, code: number, message: string) =>
     headers: { "content-type": "application/json" },
   });
 
-// W3C traceparent propagation across the worker→DO boundary. mcp.ts injects
-// the worker-side `mcp.request` SpanContext as a `traceparent` header on
-// forwarded requests (and as a second arg to `init()`). We parse it here and
-// use `OtelTracer.withSpanContext` to stitch the DO's root span under the
-// worker span so the entire logical request lives in one Axiom trace.
+// W3C propagation across the worker→DO boundary. mcp.ts injects the worker's
+// `traceparent` and forwards incoming `tracestate` / `baggage` headers on
+// forwarded requests (and as a second arg to `init()`). We parse the context
+// here and use `OtelTracer.withSpanContext` to stitch the DO's root span
+// 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;
 };
 
@@ -250,14 +267,14 @@ export class McpSessionDO extends DurableObject {
     });
   }
 
-  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),
       ),
     );
@@ -355,7 +372,11 @@ export class McpSessionDO extends DurableObject {
     // 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({
@@ -368,7 +389,7 @@ export class McpSessionDO extends DurableObject {
           "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

@@ -440,23 +440,47 @@ const peekAndAnnotate = (response: Response): Effect.Effect<Response> =>
 
 // 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
-// `mcp.request` span's SpanContext across as a W3C `traceparent` string —
-// the DO parses it and anchors its own spans under the same trace via
-// `OtelTracer.withSpanContext`.
+// 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 withTraceparentHeader = (request: Request) =>
-  Effect.map(currentTraceparent, (traceparent) => {
-    if (!traceparent) return request;
-    const headers = new Headers(request.headers);
-    headers.set("traceparent", traceparent);
-    return new Request(request, { headers });
-  });
+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 });
+};
 
 /**
  * Forward a request to an existing session DO. Wrapping the DO's `Response`
@@ -467,7 +491,8 @@ const forwardToExistingSession = (request: Request, sessionId: string, peek: boo
   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>,
     );
@@ -487,9 +512,9 @@ const dispatchPost = (request: Request, token: VerifiedToken) =>
 
     const ns = env.MCP_SESSION;
     const stub = ns.get(ns.newUniqueId());
-    const traceparent = yield* currentTraceparent;
-    yield* Effect.promise(() => stub.init({ organizationId }, traceparent));
-    const propagated = yield* withTraceparentHeader(request);
+    const propagation = yield* currentPropagationHeaders(request);
+    yield* Effect.promise(() => stub.init({ organizationId }, propagation));
+    const propagated = withPropagationHeaders(request, propagation);
     const raw = yield* Effect.promise(
       () => stub.handleRequest(propagated) as Promise<Response>,
     );

apps/cloud/src/server.ts

@@ -15,6 +15,12 @@ import { server } from "./env";
 // DOMException "Illegal invocation".
 // ---------------------------------------------------------------------------
 
+const parseSampleRatio = (value: string): number => {
+  const n = Number(value);
+  if (!Number.isFinite(n)) return 1;
+  return Math.min(1, Math.max(0, n));
+};
+
 const otelConfig: TraceConfig = {
   service: { name: "executor-cloud", version: "1.0.0" },
   exporter: {
@@ -24,6 +30,13 @@ const otelConfig: TraceConfig = {
       "X-Axiom-Dataset": server.AXIOM_DATASET,
     },
   },
+  sampling: {
+    headSampler: {
+      // Keep remote parent decisions and make local sampling policy explicit.
+      acceptRemote: true,
+      ratio: parseSampleRatio(server.AXIOM_TRACES_SAMPLE_RATIO),
+    },
+  },
 };
 
 // otel-cf-workers owns the global TracerProvider. Sentry's OTEL compat shim