Merge pull request #323 from RhysSullivan/fix/otel-propagation-hardening

Cloud: harden OTEL worker→DO propagation and make sampling explicit

Author: RhysSullivan

AGENTS.md

@@ -5,3 +5,5 @@ do not add any AI assistant, Claude, Anthropic, or Co-Authored-By attribution/tr
 i use speech to text occasionally so if sentences are weird / words aren't right that's why
 
 code is very cheap to write. do not give time estimates with agents code is practically instant to generate therefore unless stated otherwise time to implement is not a blocker
+
+you have repos in .references like effect, effect-atom. if you are given a git url clone it into that directory to explore it. if you need to know about good patterns look in there

CLAUDE.md

@@ -5,3 +5,5 @@ do not add any AI assistant, Claude, Anthropic, or Co-Authored-By attribution/tr
 i use speech to text occasionally so if sentences are weird / words aren't right that's why
 
 code is very cheap to write. do not give time estimates with agents code is practically instant to generate therefore unless stated otherwise time to implement is not a blocker
+
+you have repos in .references like effect, effect-atom. if you are given a git url clone it into that directory to explore it. if you need to know about good patterns look in there

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