diff --git a/.changeset/trace-context-propagation-race.md b/.changeset/trace-context-propagation-race.md new file mode 100644 index 000000000..1689e57d6 --- /dev/null +++ b/.changeset/trace-context-propagation-race.md @@ -0,0 +1,5 @@ +--- +"braintrust": patch +--- + +fix: Fix OTEL trace context propagation race diff --git a/js/src/framework.test.ts b/js/src/framework.test.ts index cb0940092..095a2d28b 100644 --- a/js/src/framework.test.ts +++ b/js/src/framework.test.ts @@ -19,8 +19,10 @@ import { _exportsForTestingOnly, BraintrustState, initLogger, + injectTraceContext, TestBackgroundLogger, } from "./logger"; +import { parseBaggage } from "./propagation"; import { configureNode } from "./node/config"; import type { ProgressReporter } from "./reporters/types"; import { InternalAbortError } from "./util"; @@ -2131,6 +2133,64 @@ test("Eval with parent flushes evaluator state, not global state", async () => { _exportsForTestingOnly.simulateLogoutForTests(); }); +// Regression: the experiment id resolves asynchronously (POST +// /api/experiment/register), but `Span.inject` reads it synchronously to build +// the `braintrust.parent` baggage entry. If Eval starts tasks before the id +// lands, the first task propagates trace identity with no destination and the +// receiving process silently starts a fresh local trace. +test("Eval resolves the experiment id before tasks run so injected context is routable", async () => { + const state = await _exportsForTestingOnly.simulateLoginForTests(); + vi.spyOn(state, "login").mockResolvedValue(state as never); + vi.spyOn(_exportsForTestingOnly.isomorph, "getRepoInfo").mockResolvedValue( + undefined, + ); + vi.spyOn( + _exportsForTestingOnly.isomorph, + "getPastNAncestors", + ).mockResolvedValue([]); + vi.spyOn(state.appConn(), "post_json").mockResolvedValue({ + project: { id: "project-id", name: "test-inject-project" }, + experiment: { + id: "experiment-id", + project_id: "project-id", + name: "test-inject-experiment", + public: false, + }, + }); + _exportsForTestingOnly.useTestBackgroundLogger(); + + const carriers: Record[] = []; + + await Eval( + "test-inject-project", + { + data: [{ input: 1 }, { input: 2 }], + task: (input: number) => { + carriers.push(injectTraceContext()); + return input * 2; + }, + scores: [], + state, + // Keep the run hermetic: score summarization is a separate server round trip. + summarizeScores: false, + }, + { returnResults: false }, + ); + + expect(carriers).toHaveLength(2); + for (const carrier of carriers) { + expect(carrier["traceparent"]).toBeDefined(); + // The first task must carry the parent, not just the later ones. + expect(parseBaggage(carrier["baggage"])["braintrust.parent"]).toBe( + "experiment_id:experiment-id", + ); + } + + _exportsForTestingOnly.clearTestBackgroundLogger(); + _exportsForTestingOnly.simulateLogoutForTests(); + vi.restoreAllMocks(); +}); + test("classifier-only evaluator populates classifications field", async () => { const result = await Eval( "test-classifier-only", diff --git a/js/src/framework.ts b/js/src/framework.ts index 4b4b58386..bb59b02dd 100644 --- a/js/src/framework.ts +++ b/js/src/framework.ts @@ -815,14 +815,13 @@ export async function Eval< { disabled: Boolean(options.parent || options.noSendLogs) }, ); - // Ensure experiment ID is resolved before tasks start for OTEL parent attribute support - // The Experiment constructor starts resolution (fire-and-forget), but we await here to ensure completion - // Only needed when OTEL compat mode is enabled - if ( - experiment && - typeof process !== "undefined" && - globalThis.BRAINTRUST_CONTEXT_MANAGER !== undefined - ) { + // Resolve the experiment ID before any task (and therefore any span) exists. + // Everything that reads the parent synchronously depends on it: OTEL parent + // attributes, and the `braintrust.parent` baggage entry that `Span.inject` / + // `injectTraceContext` emit. Without this, the first span of a run would + // propagate trace identity with no destination, and the receiving process + // would silently start a fresh local trace instead. + if (experiment) { await experiment._waitForId(); } diff --git a/js/src/logger.ts b/js/src/logger.ts index 22388cd4b..5ab5b7083 100644 --- a/js/src/logger.ts +++ b/js/src/logger.ts @@ -8462,11 +8462,24 @@ export class SpanImpl implements Span { ): T | Record { const resolvedCarrier = carrier ?? {}; try { + const braintrustParent = + this._getOtelParent() ?? this._propagatedState?.braintrustParent; + if (!braintrustParent) { + // Symmetric with the receive-side warning in resolveW3cParent: + // surface this in the process that caused it, rather than leaving the + // consumer to debug a trace that arrived with no destination. + debugLogger + .forState(this._state) + .warn( + "Injecting trace context without braintrust.parent because the span's " + + "destination is not available yet. The receiver will start a new " + + "local trace instead of continuing this one.", + ); + } _injectIntoCarrier(resolvedCarrier, { traceId: this._rootSpanId, spanId: this._spanId, - braintrustParent: - this._getOtelParent() ?? this._propagatedState?.braintrustParent, + braintrustParent, propagatedState: this._propagatedState, }); } catch (e) { diff --git a/js/src/propagation.test.ts b/js/src/propagation.test.ts index 17b07dbac..062cb16d2 100644 --- a/js/src/propagation.test.ts +++ b/js/src/propagation.test.ts @@ -7,7 +7,7 @@ * propagation path (no `@opentelemetry/api` dependency). */ -import { afterEach, beforeEach, describe, expect, test } from "vitest"; +import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; import { _exportsForTestingOnly, _injectIntoCarrier, @@ -27,6 +27,10 @@ import { parseBaggage, parseTraceparent, } from "./propagation"; +import { + resetDebugLoggerForTests, + setGlobalDebugLogLevel, +} from "./debug-logger"; import { SpanComponentsV3 } from "../util/span_identifier_v3"; import { SpanComponentsV4 } from "../util/span_identifier_v4"; import { SpanObjectTypeV3 } from "../util/index"; @@ -476,6 +480,7 @@ test("getHeader reads Web Headers-style objects", () => { // --------------------------------------------------------------------------- // const PROJECT_NAME = "propagation-test"; +const EXPERIMENT_ID = "propagation-test-experiment"; describe("inject / extract / round-trip", () => { let memoryLogger: ReturnType< @@ -722,6 +727,47 @@ describe("inject / extract / round-trip", () => { }); }); + // An experiment's id resolves asynchronously (POST /api/experiment/register), + // so a span created before it lands has no `braintrust.parent` to inject. The + // trace identity still propagates, which is how this used to fail silently: + // the receiver saw a traceparent it could not route and started a fresh local + // trace instead. Warn on the send side, where the problem actually is. + test("experiment span with an unresolved id warns and omits the parent", () => { + const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}); + setGlobalDebugLogLevel("warn"); + try { + const experiment = + _exportsForTestingOnly.initTestExperiment(EXPERIMENT_ID); + const span = experiment.startSpan({ name: "task" }); + const carrier = span.inject>({}); + span.end(); + + expect(carrier[TRACEPARENT_HEADER]).toMatch(TRACEPARENT_RE); + expect(BAGGAGE_HEADER in carrier).toBe(false); + expect(warnSpy).toHaveBeenCalledWith( + "[braintrust]", + expect.stringContaining("without braintrust.parent"), + ); + } finally { + resetDebugLoggerForTests(); + warnSpy.mockRestore(); + } + }); + + test("experiment span injects the parent once the id is resolved", async () => { + const experiment = + _exportsForTestingOnly.initTestExperiment(EXPERIMENT_ID); + await experiment._waitForId(); + + const span = experiment.startSpan({ name: "task" }); + const carrier = span.inject>({}); + span.end(); + + expect(parseBaggage(carrier[BAGGAGE_HEADER])).toEqual({ + [BRAINTRUST_PARENT_KEY]: `experiment_id:${EXPERIMENT_ID}`, + }); + }); + test("no braintrust parent injects traceparent without baggage", () => { const carrier: Record = {}; _injectIntoCarrier(carrier, {