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

@@ -36,6 +36,7 @@ type ServerEnv = SharedEnv & · 35 unmodified lines
36	36	AXIOM_TOKEN: string;
37	37	AXIOM_DATASET: string;
38	38	AXIOM_TRACES_URL: string;
	39	+ AXIOM_TRACES_SAMPLE_RATIO: string;
39	40	}>;
40	41
41	42	type WebEnv = Readonly<Record<string, never>>;
@@ -55,6 +56,7 @@ const SERVER_DEFAULTS: Record<keyof ServerEnv, string> = { · 13 unmodified lines
55	56	AXIOM_TOKEN: "",
56	57	AXIOM_DATASET: "executor-cloud",
57	58	AXIOM_TRACES_URL: "https://api.axiom.co/v1/traces",
	59	+ AXIOM_TRACES_SAMPLE_RATIO: "1",
58	60	};
59	61
60	62	const SHARED_DEFAULTS: Record<keyof SharedEnv, string> = {

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

@@ -3,6 +3,7 @@ · 2 unmodified lines
3	3	// ---------------------------------------------------------------------------
4	4
5	5	import { DurableObject, env } from "cloudflare:workers";
	6	+import { createTraceState } from "@opentelemetry/api";
6	7	import { Data, Effect, Layer } from "effect";
7	8	import * as OtelTracer from "@effect/opentelemetry/Tracer";
8	9	import * as Sentry from "@sentry/cloudflare";
@@ -36,6 +37,12 @@ export type McpSessionInit = { · 27 unmodified lines
36	37	organizationId: string;
37	38	};
38	39
	40	+export type IncomingTraceHeaders = {
	41	+ readonly traceparent?: string;
	42	+ readonly tracestate?: string;
	43	+ readonly baggage?: string;
	44	+};
	45	+
39	46	const HEARTBEAT_MS = 30 * 1000;
40	47	const SESSION_TIMEOUT_MS = 5 * 60 * 1000;
41	48	const TRANSPORT_STATE_KEY = "transport";
@@ -59,31 +66,41 @@ const jsonRpcError = (status: number, code: number, message: string) => · 17 unmodified lines
59	66	headers: { "content-type": "application/json" },
60	67	});
61	68
62		−// W3C traceparent propagation across the worker→DO boundary. mcp.ts injects
63		−// the worker-side `mcp.request` SpanContext as a `traceparent` header on
64		−// forwarded requests (and as a second arg to `init()`). We parse it here and
65		−// use `OtelTracer.withSpanContext` to stitch the DO's root span under the
66		−// worker span so the entire logical request lives in one Axiom trace.
	69	+// W3C propagation across the worker→DO boundary. mcp.ts injects the worker's
	70	+// `traceparent` and forwards incoming `tracestate` / `baggage` headers on
	71	+// forwarded requests (and as a second arg to `init()`). We parse the context
	72	+// here and use `OtelTracer.withSpanContext` to stitch the DO's root span
	73	+// under the worker span so the entire logical request lives in one trace.
67	74	const TRACEPARENT_PATTERN = /^([0-9a-f]{2})-([0-9a-f]{32})-([0-9a-f]{16})-([0-9a-f]{2})$/;
68	75
69	76	type IncomingSpanContext = {
70	77	readonly traceId: string;
71	78	readonly spanId: string;
72	79	readonly traceFlags: number;
	80	+ readonly traceState?: ReturnType<typeof createTraceState>;
73	81	};
74	82
75		−const parseTraceparent = (value: string \| null \| undefined): IncomingSpanContext \| null => {
	83	+const parseTraceparent = (
	84	+ traceparent: string \| null \| undefined,
	85	+ tracestate: string \| null \| undefined,
	86	+): IncomingSpanContext \| null => {
	87	+ const value = traceparent;
76	88	if (!value) return null;
77	89	const match = TRACEPARENT_PATTERN.exec(value);
78	90	if (!match) return null;
79		− return { traceId: match[2]!, spanId: match[3]!, traceFlags: parseInt(match[4]!, 16) };
	91	+ return {
	92	+ traceId: match[2]!,
	93	+ spanId: match[3]!,
	94	+ traceFlags: parseInt(match[4]!, 16),
	95	+ ...(tracestate ? { traceState: createTraceState(tracestate) } : {}),
	96	+ };
80	97	};
81	98
82	99	const withIncomingParent = <A, E, R>(
83		− traceparent: string \| null \| undefined,
	100	+ incoming: IncomingTraceHeaders \| null \| undefined,
84	101	effect: Effect.Effect<A, E, R>,
85	102	): Effect.Effect<A, E, R> => {
86		− const parsed = parseTraceparent(traceparent);
	103	+ const parsed = parseTraceparent(incoming?.traceparent, incoming?.tracestate);
87	104	return parsed ? OtelTracer.withSpanContext(effect, parsed) : effect;
88	105	};
89	106
@@ -250,14 +267,14 @@ export class McpSessionDO extends DurableObject { · 160 unmodified lines
250	267	});
251	268	}
252	269
253		− async init(token: McpSessionInit, traceparent?: string): Promise<void> {
	270	+ async init(token: McpSessionInit, incoming?: IncomingTraceHeaders): Promise<void> {
254	271	if (this.initialized) return;
255	272	return Effect.runPromise(
256	273	this.doInitEffect(token).pipe(
257	274	Effect.withSpan("McpSessionDO.init", {
258	275	attributes: { "mcp.auth.organization_id": token.organizationId },
259	276	}),
260		− (eff) => withIncomingParent(traceparent, eff),
	277	+ (eff) => withIncomingParent(incoming, eff),
261	278	Effect.provide(DoTelemetryLive),
262	279	),
263	280	);
@@ -355,7 +372,11 @@ export class McpSessionDO extends DurableObject { · 91 unmodified lines
355	372	// only (method, session-id presence, response status); rich client
356	373	// fingerprint stays on the edge `mcp.request` span, which shares a
357	374	// trace_id with this one.
358		− const traceparent = request.headers.get("traceparent");
	375	+ const incoming = {
	376	+ traceparent: request.headers.get("traceparent") ?? undefined,
	377	+ tracestate: request.headers.get("tracestate") ?? undefined,
	378	+ baggage: request.headers.get("baggage") ?? undefined,
	379	+ } satisfies IncomingTraceHeaders;
359	380	const program = Effect.promise(() => this.dispatchRequest(request)).pipe(
360	381	Effect.tap((response) =>
361	382	Effect.annotateCurrentSpan({
@@ -368,7 +389,7 @@ export class McpSessionDO extends DurableObject { · 6 unmodified lines
368	389	"mcp.request.session_id_present": !!request.headers.get("mcp-session-id"),
369	390	},
370	391	}),
371		− (eff) => withIncomingParent(traceparent, eff),
	392	+ (eff) => withIncomingParent(incoming, eff),
372	393	Effect.provide(DoTelemetryLive),
373	394	);
374	395	return Effect.runPromise(program);

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

@@ -440,23 +440,47 @@ const peekAndAnnotate = (response: Response): Effect.Effect<Response> => · 439 unmodified lines
440	440
441	441	// Worker and DO run in separate isolates with independent WebSdk tracer
442	442	// providers. Neither one can see the other's OTEL context, so the DO used
443		−// to emit a brand-new root trace on every stub call. Ferry the current
444		−// `mcp.request` span's SpanContext across as a W3C `traceparent` string —
445		−// the DO parses it and anchors its own spans under the same trace via
446		−// `OtelTracer.withSpanContext`.
	443	+// to emit a brand-new root trace on every stub call. Ferry the worker span
	444	+// context across with W3C headers: `traceparent` generated from the active
	445	+// Effect span plus passthrough `tracestate` / `baggage` from the inbound
	446	+// request.
	447	+type IncomingPropagationHeaders = {
	448	+ readonly traceparent?: string;
	449	+ readonly tracestate?: string;
	450	+ readonly baggage?: string;
	451	+};
	452	+
447	453	const currentTraceparent = Effect.map(Effect.currentSpan, (span) => {
448	454	if (!span \|\| !span.traceId \|\| !span.spanId) return undefined;
449	455	const flags = span.sampled ? "01" : "00";
450	456	return `00-${span.traceId}-${span.spanId}-${flags}`;
451	457	}).pipe(Effect.orElseSucceed(() => undefined));
452	458
453		−const withTraceparentHeader = (request: Request) =>
454		− Effect.map(currentTraceparent, (traceparent) => {
455		− if (!traceparent) return request;
456		− const headers = new Headers(request.headers);
457		− headers.set("traceparent", traceparent);
458		− return new Request(request, { headers });
459		− });
	459	+const currentPropagationHeaders = (
	460	+ request: Request,
	461	+): Effect.Effect<IncomingPropagationHeaders> =>
	462	+ Effect.map(currentTraceparent, (traceparent) => ({
	463	+ traceparent,
	464	+ tracestate: request.headers.get("tracestate") ?? undefined,
	465	+ baggage: request.headers.get("baggage") ?? undefined,
	466	+ }));
	467	+
	468	+const withPropagationHeaders = (
	469	+ request: Request,
	470	+ propagation: IncomingPropagationHeaders,
	471	+): Request => {
	472	+ const headers = new Headers(request.headers);
	473	+ if (propagation.traceparent) {
	474	+ headers.set("traceparent", propagation.traceparent);
	475	+ }
	476	+ if (propagation.tracestate) {
	477	+ headers.set("tracestate", propagation.tracestate);
	478	+ }
	479	+ if (propagation.baggage) {
	480	+ headers.set("baggage", propagation.baggage);
	481	+ }
	482	+ return new Request(request, { headers });
	483	+};
460	484
461	485	/**
462	486	* 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 · 4 unmodified lines
467	491	Effect.gen(function* () {
468	492	const ns = env.MCP_SESSION;
469	493	const stub = ns.get(ns.idFromString(sessionId));
470		− const propagated = yield* withTraceparentHeader(request);
	494	+ const propagation = yield* currentPropagationHeaders(request);
	495	+ const propagated = withPropagationHeaders(request, propagation);
471	496	const raw = yield* Effect.promise(
472	497	() => stub.handleRequest(propagated) as Promise<Response>,
473	498	);
@@ -487,9 +512,9 @@ const dispatchPost = (request: Request, token: VerifiedToken) => · 13 unmodified lines
487	512
488	513	const ns = env.MCP_SESSION;
489	514	const stub = ns.get(ns.newUniqueId());
490		− const traceparent = yield* currentTraceparent;
491		− yield* Effect.promise(() => stub.init({ organizationId }, traceparent));
492		− const propagated = yield* withTraceparentHeader(request);
	515	+ const propagation = yield* currentPropagationHeaders(request);
	516	+ yield* Effect.promise(() => stub.init({ organizationId }, propagation));
	517	+ const propagated = withPropagationHeaders(request, propagation);
493	518	const raw = yield* Effect.promise(
494	519	() => stub.handleRequest(propagated) as Promise<Response>,
495	520	);

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

@@ -15,6 +15,12 @@ import { server } from "./env"; · 14 unmodified lines
15	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	24	const otelConfig: TraceConfig = {
19	25	service: { name: "executor-cloud", version: "1.0.0" },
20	26	exporter: {
@@ -24,6 +30,13 @@ const otelConfig: TraceConfig = { · 3 unmodified lines
24	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	42	// otel-cf-workers owns the global TracerProvider. Sentry's OTEL compat shim