Skip to content

Commit 7d4f4c4

Browse files
Qardclaude
Qard
and
claude
authored
fix(otel): propagate TracingChannel span context in OTEL compat mode (#1724)
## Problem When `setupOtelCompat()` is active, `OtelContextManager` did not expose `[BRAINTRUST_CURRENT_SPAN_STORE]`, so `bindCurrentSpanStoreToStart()` silently skipped context binding. Auto-instrumented spans (Anthropic, OpenAI, Google GenAI) had no parent context in OTEL compat mode. ## Solution Make `ContextManager` a thin abstraction over whichever ALS is in use: - **`wrapSpanForStore(span)`** added to `ContextManager` base class (default: identity — returns span directly) - **`OtelContextManager`** overrides with two-mode behavior: - **OTEL ALS available** (`AsyncLocalStorageContextManager`): exposes OTEL's internal `_asyncLocalStorage`, `wrapSpanForStore` returns an OTEL `Context` wrapping the span - **OTEL ALS unavailable** (`AsyncHooksContextManager`): falls back to its own `IsoAsyncLocalStorage<Span>`, `wrapSpanForStore` returns the span directly — same behavior as `BraintrustContextManager` - **`bindCurrentSpanStoreToStart`** (and `google-genai-plugin.ts`) call `contextManager.wrapSpanForStore(span)` in the bindStore transform - `getCurrentSpan()` checks both `otelContext.active().getValue(BT_SPAN_KEY)` and the fallback ALS, so spans propagated by either path are visible ## Additional cleanup - Extract `buildBtOtelContext()` helper (eliminates 40-line duplication in `runInContext`) - Unify `BT_SPAN_KEY` / `BT_PARENT_KEY` as module-level constants - Add `CurrentSpanStore = IsoAsyncLocalStorage<unknown>` type alias - Export `CurrentSpanStore` from `braintrust` package ## Tests 8 new tests in `otel-compat.test.ts`: - OTEL ALS exposed via `BRAINTRUST_CURRENT_SPAN_STORE` ✓ - `wrapSpanForStore` returns Context → `currentSpan()` works ✓ - `store.run()` propagation ✓ - Nested runs maintain span chain ✓ - OTEL active span has matching IDs ✓ - `AsyncHooksContextManager` fallback: ALS always exposed ✓ - `AsyncHooksContextManager` fallback: `wrapSpanForStore` returns span directly ✓ - `AsyncHooksContextManager` fallback: `getCurrentSpan()` reads from fallback ALS ✓ All tests pass on both otel-v1 (OTel 1.x) and otel-v2 (OTel 2.x). 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
1 parent 40f98bb commit 7d4f4c4

6 files changed

Lines changed: 363 additions & 76 deletions

File tree

‎integrations/otel-js/src/context.ts‎

Lines changed: 94 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,20 @@
11
import {
22
ContextManager,
3+
BRAINTRUST_CURRENT_SPAN_STORE,
4+
_internalIso as iso,
35
type ContextParentSpanIds,
6+
type CurrentSpanStore,
47
type Span,
58
} from "braintrust";
69

710
import { trace as otelTrace, context as otelContext } from "@opentelemetry/api";
811
import { getOtelParentFromSpan } from "./otel";
912

13+
// eslint-disable-next-line @typescript-eslint/no-explicit-any
14+
const BT_SPAN_KEY = "braintrust_span" as any;
15+
// eslint-disable-next-line @typescript-eslint/no-explicit-any
16+
const BT_PARENT_KEY = "braintrust.parent" as any;
17+
1018
function isOtelSpan(span: unknown): span is {
1119
spanContext: () => { spanId: string; traceId: string };
1220
} {
@@ -46,7 +54,76 @@ function isValidSpanContext(spanContext: unknown): boolean {
4654
);
4755
}
4856

57+
/**
58+
* Builds an OTEL Context containing a NonRecordingSpan wrapper around the
59+
* Braintrust span's IDs, plus the BT span stored under BT_SPAN_KEY for
60+
* retrieval by getCurrentSpan().
61+
*/
62+
function buildBtOtelContext(span: Span): unknown {
63+
const btSpan = span as { spanId: string; rootSpanId: string };
64+
const spanContext = {
65+
traceId: btSpan.rootSpanId,
66+
spanId: btSpan.spanId,
67+
traceFlags: 1, // sampled
68+
};
69+
const wrappedSpan = otelTrace.wrapSpanContext(spanContext);
70+
const currentContext = otelContext.active();
71+
let newContext = otelTrace.setSpan(currentContext, wrappedSpan);
72+
newContext = newContext.setValue(BT_SPAN_KEY, span);
73+
74+
if (isBraintrustSpan(span)) {
75+
// eslint-disable-next-line @typescript-eslint/consistent-type-assertions
76+
const parentValue = getOtelParentFromSpan(span as never);
77+
if (parentValue) {
78+
newContext = newContext.setValue(BT_PARENT_KEY, parentValue);
79+
}
80+
}
81+
82+
return newContext;
83+
}
84+
4985
export class OtelContextManager extends ContextManager {
86+
/** Fallback ALS used when the OTEL context manager doesn't expose _asyncLocalStorage. */
87+
private _ownAls: CurrentSpanStore | undefined;
88+
89+
private _getOtelAls(): CurrentSpanStore | undefined {
90+
// eslint-disable-next-line @typescript-eslint/no-explicit-any
91+
return (otelContext as any)._getContextManager?.()._asyncLocalStorage;
92+
}
93+
94+
constructor() {
95+
super();
96+
// Expose whichever ALS is in use via BRAINTRUST_CURRENT_SPAN_STORE so that
97+
// TracingChannel's bindStore can propagate span context. We prefer OTEL's own
98+
// ALS (AsyncLocalStorageContextManager._asyncLocalStorage) so that spans
99+
// stored by runStores are visible to OTEL's context APIs. If the active OTEL
100+
// context manager doesn't expose an ALS (e.g. AsyncHooksContextManager), we
101+
// fall back to our own IsoAsyncLocalStorage<Span> and behave like the default
102+
// BraintrustContextManager for TracingChannel binding.
103+
//
104+
// A lazy getter is required because the global OTEL context manager may not be
105+
// registered until after this instance is constructed.
106+
const self = this;
107+
Object.defineProperty(this, BRAINTRUST_CURRENT_SPAN_STORE, {
108+
get(): CurrentSpanStore {
109+
const otelAls = self._getOtelAls();
110+
if (otelAls) return otelAls;
111+
if (!self._ownAls) self._ownAls = iso.newAsyncLocalStorage<unknown>();
112+
return self._ownAls;
113+
},
114+
configurable: true,
115+
enumerable: false,
116+
});
117+
}
118+
119+
wrapSpanForStore(span: Span): unknown {
120+
// When using OTEL's ALS the stored value must be an OTEL Context, not a raw
121+
// Span, so that OTEL's own context propagation sees a valid Context object.
122+
// When using our own fallback ALS we store the Span directly (default mode).
123+
if (this._getOtelAls()) return buildBtOtelContext(span);
124+
return span;
125+
}
126+
50127
getParentSpanIds(): ContextParentSpanIds | undefined {
51128
const currentSpan = otelTrace.getActiveSpan();
52129
if (!currentSpan || !isOtelSpan(currentSpan)) {
@@ -59,8 +136,7 @@ export class OtelContextManager extends ContextManager {
59136
}
60137

61138
// Check if this is a wrapped BT span
62-
// eslint-disable-next-line @typescript-eslint/consistent-type-assertions, @typescript-eslint/no-explicit-any
63-
const btSpan = otelContext?.active().getValue?.("braintrust_span" as any);
139+
const btSpan = otelContext?.active().getValue?.(BT_SPAN_KEY);
64140
if (
65141
btSpan &&
66142
currentSpan.constructor.name === "NonRecordingSpan" &&
@@ -88,46 +164,12 @@ export class OtelContextManager extends ContextManager {
88164

89165
runInContext<R>(span: Span, callback: () => R): R {
90166
try {
91-
if (
92-
typeof span === "object" &&
93-
span !== null &&
94-
"spanId" in span &&
95-
"rootSpanId" in span
96-
) {
97-
// eslint-disable-next-line @typescript-eslint/consistent-type-assertions
98-
const btSpan = span as { spanId: string; rootSpanId: string };
99-
100-
// Create a span context for the NonRecordingSpan
101-
const spanContext = {
102-
traceId: btSpan.rootSpanId,
103-
spanId: btSpan.spanId,
104-
traceFlags: 1, // sampled
105-
};
106-
107-
// Wrap the span context
108-
const wrappedContext = otelTrace.wrapSpanContext(spanContext);
109-
110-
// Get current context and add both the wrapped span and the BT span
111-
const currentContext = otelContext.active();
112-
let newContext = otelTrace.setSpan(currentContext, wrappedContext);
113-
// eslint-disable-next-line @typescript-eslint/consistent-type-assertions, @typescript-eslint/no-explicit-any
114-
newContext = newContext.setValue("braintrust_span" as any, span);
115-
116-
// Get parent value and store it in context (matching Python's behavior)
117-
if (isBraintrustSpan(span)) {
167+
if (isBraintrustSpan(span)) {
168+
return otelContext.with(
118169
// eslint-disable-next-line @typescript-eslint/consistent-type-assertions
119-
const parentValue = getOtelParentFromSpan(span as never);
120-
if (parentValue) {
121-
newContext = newContext.setValue(
122-
// eslint-disable-next-line @typescript-eslint/consistent-type-assertions, @typescript-eslint/no-explicit-any
123-
"braintrust.parent" as any,
124-
parentValue,
125-
);
126-
}
127-
}
128-
129-
// Run the callback in the new context
130-
return otelContext.with(newContext, callback);
170+
buildBtOtelContext(span) as Parameters<typeof otelContext.with>[0],
171+
callback,
172+
);
131173
}
132174
} catch (error) {
133175
console.warn("Failed to run in OTEL context:", error);
@@ -137,9 +179,9 @@ export class OtelContextManager extends ContextManager {
137179
}
138180

139181
getCurrentSpan(): Span | undefined {
140-
// eslint-disable-next-line @typescript-eslint/consistent-type-assertions, @typescript-eslint/no-explicit-any
141-
const btSpan = otelContext.active().getValue?.("braintrust_span" as any);
142-
182+
// Check OTEL context first — this covers both runInContext and the OTEL-ALS
183+
// TracingChannel path where runStores stores a Context under BT_SPAN_KEY.
184+
const btSpan = otelContext.active().getValue?.(BT_SPAN_KEY);
143185
if (
144186
btSpan &&
145187
typeof btSpan === "object" &&
@@ -150,6 +192,14 @@ export class OtelContextManager extends ContextManager {
150192
// eslint-disable-next-line @typescript-eslint/consistent-type-assertions
151193
return btSpan as Span;
152194
}
195+
196+
// If we're using the fallback ALS (non-ALS OTEL context manager), spans are
197+
// stored directly in our own ALS by TracingChannel's runStores.
198+
if (this._ownAls) {
199+
const stored = this._ownAls.getStore();
200+
if (isBraintrustSpan(stored)) return stored;
201+
}
202+
153203
return undefined;
154204
}
155205
}

0 commit comments

Comments
 (0)