From a595131be054e001c5e73cd8ed372ee9e743f236 Mon Sep 17 00:00:00 2001 From: Starfolk Date: Fri, 14 Aug 2026 20:49:04 +0000 Subject: [PATCH 1/2] feat: add client-side trace sampling --- .changeset/giran-sdk-191-trace-sampling.md | 10 ++ docs/trace-sampling.md | 32 ++++ e2e/scenarios/trace-sampling/scenario.test.ts | 31 ++++ e2e/scenarios/trace-sampling/scenario.ts | 44 +++++ integrations/otel-js/src/context.ts | 19 +- integrations/otel-js/src/otel.ts | 5 +- .../core/channel-tracing-utils.ts | 7 + .../core/channel-tracing.test.ts | 37 ++++ .../instrumentation/core/channel-tracing.ts | 56 ++++-- js/src/instrumentation/core/plugin.ts | 10 ++ js/src/logger.test.ts | 78 ++++++++ js/src/logger.ts | 167 ++++++++++++++++-- js/src/sampling.test.ts | 51 ++++++ js/src/sampling.ts | 80 +++++++++ js/util/span_identifier_v3.test.ts | 16 ++ js/util/span_identifier_v3.ts | 5 + js/util/span_identifier_v4.test.ts | 13 ++ js/util/span_identifier_v4.ts | 6 + 18 files changed, 632 insertions(+), 35 deletions(-) create mode 100644 .changeset/giran-sdk-191-trace-sampling.md create mode 100644 docs/trace-sampling.md create mode 100644 e2e/scenarios/trace-sampling/scenario.test.ts create mode 100644 e2e/scenarios/trace-sampling/scenario.ts create mode 100644 js/src/sampling.test.ts create mode 100644 js/src/sampling.ts create mode 100644 js/util/span_identifier_v3.test.ts diff --git a/.changeset/giran-sdk-191-trace-sampling.md b/.changeset/giran-sdk-191-trace-sampling.md new file mode 100644 index 000000000..bc277df69 --- /dev/null +++ b/.changeset/giran-sdk-191-trace-sampling.md @@ -0,0 +1,10 @@ +--- +"braintrust": minor +"@braintrust/otel": patch +--- + +Add deterministic client-side head sampling for native JavaScript traces with +Logger and root-span `sampleRate` controls. Rejected traces preserve context +without emitting rows, and native/OpenTelemetry propagation now retains raw +trace flags. `Span` adds `isRecording()`; TypeScript consumers that structurally +implement `Span` should add that method. diff --git a/docs/trace-sampling.md b/docs/trace-sampling.md new file mode 100644 index 000000000..8ff3880f2 --- /dev/null +++ b/docs/trace-sampling.md @@ -0,0 +1,32 @@ +# Client-side trace sampling + +Native Braintrust project-log traces can be sampled at their root: + +```ts +const logger = initLogger({ projectName: "production", sampleRate: 0.2 }); +``` + +`sampleRate` is a fraction from `0` through `1` and defaults to `1`. The SDK +uses a deterministic decision derived from the root trace ID, so all native +children make the same decision. A rejected trace still carries trace context +and stable IDs, but `span.isRecording()` is `false` and no rows are queued. + +Use a root-only override for exceptional operations: + +```ts +logger.startSpan({ name: "priority-request", sampleRate: 1 }); +``` + +An existing local, exported, W3C, or OpenTelemetry parent always takes +precedence over either rate. The sampled bit is forwarded in `traceparent` and +native span exports, including unsampled (`00`) continuations. New local +Experiment roots remain recorded; an Experiment that continues an unsampled +parent remains non-recording to preserve distributed trace coherence. + +When OpenTelemetry owns span creation, configure its provider sampler instead. +The `@braintrust/otel` integration preserves the provider's decision and does +not apply `Logger.sampleRate` a second time. + +Sampling is probabilistic head sampling, not a maximum traces-per-second or +bytes-per-minute limiter. Older SDK processes that do not understand the +optional native export-token flags may record an unsampled continuation. diff --git a/e2e/scenarios/trace-sampling/scenario.test.ts b/e2e/scenarios/trace-sampling/scenario.test.ts new file mode 100644 index 000000000..b30dd6768 --- /dev/null +++ b/e2e/scenarios/trace-sampling/scenario.test.ts @@ -0,0 +1,31 @@ +import { expect, test } from "vitest"; +import { + prepareScenarioDir, + resolveScenarioDir, + withScenarioHarness, +} from "../../helpers/scenario-harness"; +import { findLatestSpan } from "../../helpers/trace-selectors"; + +const scenarioDir = await prepareScenarioDir({ + scenarioDir: resolveScenarioDir(import.meta.url), +}); + +test("trace sampling drops whole roots before ingestion and permits a root override", async () => { + await withScenarioHarness(async ({ runScenarioDir, testRunEvents }) => { + await runScenarioDir({ scenarioDir }); + + const events = testRunEvents(); + const keptRoot = findLatestSpan(events, "trace-sampling-kept-root"); + const keptChild = findLatestSpan(events, "trace-sampling-kept-child"); + + expect(keptRoot).toBeDefined(); + expect(keptChild?.span.parentIds).toEqual([keptRoot?.span.id ?? ""]); + expect( + events.some( + (event) => + event.span.name === "trace-sampling-dropped-root" || + event.span.name === "trace-sampling-dropped-child", + ), + ).toBe(false); + }); +}); diff --git a/e2e/scenarios/trace-sampling/scenario.ts b/e2e/scenarios/trace-sampling/scenario.ts new file mode 100644 index 000000000..931707e60 --- /dev/null +++ b/e2e/scenarios/trace-sampling/scenario.ts @@ -0,0 +1,44 @@ +import { initLogger } from "braintrust"; +import { + getTestRunId, + runMain, + scopedName, +} from "../../helpers/scenario-runtime"; + +async function main() { + const testRunId = getTestRunId(); + const droppedLogger = initLogger({ + projectName: scopedName("e2e-trace-sampling-dropped", testRunId), + sampleRate: 0, + }); + await droppedLogger.traced( + (root) => { + root.startSpan({ name: "trace-sampling-dropped-child" }).end(); + }, + { name: "trace-sampling-dropped-root" }, + ); + await droppedLogger.flush(); + + const keptLogger = initLogger({ + projectName: scopedName("e2e-trace-sampling-kept", testRunId), + sampleRate: 0, + }); + await keptLogger.traced( + (root) => { + root + .startSpan({ + name: "trace-sampling-kept-child", + event: { metadata: { testRunId } }, + }) + .end(); + }, + { + name: "trace-sampling-kept-root", + sampleRate: 1, + event: { metadata: { testRunId } }, + }, + ); + await keptLogger.flush(); +} + +runMain(main); diff --git a/integrations/otel-js/src/context.ts b/integrations/otel-js/src/context.ts index 6db697c67..afaff430f 100644 --- a/integrations/otel-js/src/context.ts +++ b/integrations/otel-js/src/context.ts @@ -60,11 +60,24 @@ function isValidSpanContext(spanContext: unknown): boolean { * retrieval by getCurrentSpan(). */ function buildBtOtelContext(span: Span): unknown { - const btSpan = span as { spanId: string; rootSpanId: string }; + const btSpan = span as { + spanId: string; + rootSpanId: string; + _getTraceFlags?: () => string; + isRecording?: () => boolean; + }; + const traceFlags = + typeof btSpan._getTraceFlags === "function" + ? parseInt(btSpan._getTraceFlags(), 16) + : typeof btSpan.isRecording === "function" + ? btSpan.isRecording() + ? 1 + : 0 + : 1; const spanContext = { traceId: btSpan.rootSpanId, spanId: btSpan.spanId, - traceFlags: 1, // sampled + traceFlags, }; const wrappedSpan = otelTrace.wrapSpanContext(spanContext); const currentContext = otelContext.active(); @@ -150,6 +163,7 @@ export class OtelContextManager extends ContextManager { return { rootSpanId: typedBtSpan.rootSpanId, spanParents: [typedBtSpan.spanId], + traceFlags: spanContext.traceFlags.toString(16).padStart(2, "0"), }; } @@ -159,6 +173,7 @@ export class OtelContextManager extends ContextManager { return { rootSpanId: otelTraceId, spanParents: [otelSpanId], + traceFlags: spanContext.traceFlags.toString(16).padStart(2, "0"), }; } diff --git a/integrations/otel-js/src/otel.ts b/integrations/otel-js/src/otel.ts index dc7e23aca..29fa7b619 100644 --- a/integrations/otel-js/src/otel.ts +++ b/integrations/otel-js/src/otel.ts @@ -6,7 +6,6 @@ import { Context, diag, trace, - TraceFlags, propagation, Span, } from "@opentelemetry/api"; @@ -714,7 +713,7 @@ export function contextFromSpanExport(exportStr: string): unknown { traceId: traceIdHex, spanId: spanIdHex, isRemote: true, - traceFlags: TraceFlags?.SAMPLED ?? 1, // SAMPLED flag + traceFlags: parseInt(components.data.trace_flags ?? "01", 16), }; // Create NonRecordingSpan using wrapSpanContext and set in context @@ -1177,11 +1176,13 @@ export function parentFromHeaders( row_id: string; span_id: string; root_span_id: string; + trace_flags?: string; } = { object_type: objectType, row_id: "otel", // Dummy row_id to enable span_id/root_span_id fields span_id: spanIdHex, root_span_id: traceIdHex, + trace_flags: spanContext.traceFlags.toString(16).padStart(2, "0"), }; // Add either object_id or compute_object_metadata_args, not both diff --git a/js/src/instrumentation/core/channel-tracing-utils.ts b/js/src/instrumentation/core/channel-tracing-utils.ts index edf7ea7c6..df17bfa7a 100644 --- a/js/src/instrumentation/core/channel-tracing-utils.ts +++ b/js/src/instrumentation/core/channel-tracing-utils.ts @@ -14,6 +14,13 @@ export type ChannelConfig = { type: string; }; +/** Preserve tracing for custom/older Span implementations that lack the query. */ +export function isSpanRecording(span: Span): boolean { + return ( + (span as Span & { isRecording?: () => boolean }).isRecording?.() ?? true + ); +} + function hasChannelSpanInfo( value: unknown, ): value is SpanInfoCarrier & { span_info: ChannelSpanInfo } { diff --git a/js/src/instrumentation/core/channel-tracing.test.ts b/js/src/instrumentation/core/channel-tracing.test.ts index 55c08f05b..d33c00309 100644 --- a/js/src/instrumentation/core/channel-tracing.test.ts +++ b/js/src/instrumentation/core/channel-tracing.test.ts @@ -203,6 +203,43 @@ describe("traceAsyncChannel current span binding", () => { expect(spans).toHaveLength(0); }); + it("preserves execution context while skipping extraction for a non-recording span", async () => { + initLogger({ + projectName: "channel-tracing-sampled-out", + projectId: "test-project-id", + sampleRate: 0, + }); + const extractInput = vi.fn(() => ({ input: "input", metadata: undefined })); + const extractOutput = vi.fn((result) => result); + const extractMetrics = vi.fn(() => ({})); + const unsubscribe = traceAsyncChannel(testChannels.asyncCall, { + name: "channel-tracing-sampled-out", + type: "function", + extractInput, + extractOutput, + extractMetrics, + }); + + try { + await testChannels.asyncCall.tracePromise( + async () => { + expect(currentSpan().isRecording()).toBe(false); + await Promise.resolve(); + expect(currentSpan().isRecording()).toBe(false); + return { ok: true as const }; + }, + { arguments: [{}] } as any, + ); + } finally { + unsubscribe(); + } + + expect(extractInput).not.toHaveBeenCalled(); + expect(extractOutput).not.toHaveBeenCalled(); + expect(extractMetrics).not.toHaveBeenCalled(); + expect(await backgroundLogger.drain()).toHaveLength(0); + }); + it("uses debug logging when shouldTrace throws", async () => { const consoleErrorSpy = vi .spyOn(console, "error") diff --git a/js/src/instrumentation/core/channel-tracing.ts b/js/src/instrumentation/core/channel-tracing.ts index a8ecc819c..0590f92b9 100644 --- a/js/src/instrumentation/core/channel-tracing.ts +++ b/js/src/instrumentation/core/channel-tracing.ts @@ -26,6 +26,7 @@ import type { import { isAsyncIterable, patchStreamIfNeeded } from "./stream-patcher"; import { buildStartSpanArgs, + isSpanRecording, mergeInputMetadata, type ChannelConfig, } from "./channel-tracing-utils"; @@ -218,18 +219,20 @@ function startSpanForEvent< } const startTime = getCurrentUnixTimestamp(); - try { - const { input, metadata } = config.extractInput( - event.arguments, - event as StartOf, - span, - ); - span.log({ - input, - metadata: mergeInputMetadata(metadata, spanInfoMetadata), - }); - } catch (error) { - debugLogger.error(`Error extracting input for ${channelName}:`, error); + if (isSpanRecording(span)) { + try { + const { input, metadata } = config.extractInput( + event.arguments, + event as StartOf, + span, + ); + span.log({ + input, + metadata: mergeInputMetadata(metadata, spanInfoMetadata), + }); + } catch (error) { + debugLogger.error(`Error extracting input for ${channelName}:`, error); + } } return { span, startTime }; @@ -479,6 +482,12 @@ export function traceAsyncChannel( const asyncEndEvent = event as AsyncEndOf; const { span, startTime } = spanData; + if (!isSpanRecording(span)) { + span.end(); + states.delete(event as object); + return; + } + try { const output = config.extractOutput( asyncEndEvent.result, @@ -561,6 +570,29 @@ export function traceStreamingChannel( const asyncEndEvent = event as AsyncEndOf; const { span, startTime } = spanData; + if (!isSpanRecording(span)) { + if (isAsyncIterable(asyncEndEvent.result)) { + patchStreamIfNeeded(asyncEndEvent.result, { + onComplete: () => { + span.end(); + states.delete(event as object); + }, + onError: () => { + span.end(); + states.delete(event as object); + }, + onCancel: () => { + span.end(); + states.delete(event as object); + }, + }); + } else { + span.end(); + states.delete(event as object); + } + return; + } + if (isAsyncIterable(asyncEndEvent.result)) { let firstChunkTime: number | undefined; const handleStreamError = (error: Error) => { diff --git a/js/src/instrumentation/core/plugin.ts b/js/src/instrumentation/core/plugin.ts index ae09a43f1..991927125 100644 --- a/js/src/instrumentation/core/plugin.ts +++ b/js/src/instrumentation/core/plugin.ts @@ -7,6 +7,7 @@ import type { Span } from "../../logger"; import { getCurrentUnixTimestamp } from "../../util"; import { buildStartSpanArgs, + isSpanRecording, mergeInputMetadata, } from "./channel-tracing-utils"; @@ -122,6 +123,15 @@ export abstract class BasePlugin { const { span, startTime } = spanData; + if (!isSpanRecording(span)) { + span.end(); + spans.delete(event); + return; + } + + if (!isSpanRecording(span)) { + return; + } try { const output = config.extractOutput(event.result, event); const metrics = config.extractMetrics(event.result, startTime, event); diff --git a/js/src/logger.test.ts b/js/src/logger.test.ts index 86a3a1a5c..e2dca0d41 100644 --- a/js/src/logger.test.ts +++ b/js/src/logger.test.ts @@ -426,6 +426,84 @@ test("verify MemoryBackgroundLogger intercepts logs", async () => { _exportsForTestingOnly.clearTestBackgroundLogger(); // can go back to normal }); +describe("native trace sampling", () => { + beforeEach(() => { + _exportsForTestingOnly.simulateLoginForTests(); + }); + + afterEach(() => { + _exportsForTestingOnly.clearTestBackgroundLogger(); + }); + + test("creates a real, non-recording context carrier at rate zero", async () => { + const memoryLogger = _exportsForTestingOnly.useTestBackgroundLogger(); + const logger = initLogger({ + projectName: "sampling", + projectId: "sampling-project", + sampleRate: 0, + }); + + const span = logger.startSpan({ name: "dropped" }); + const child = span.startSpan({ name: "also-dropped", sampleRate: 1 }); + span.log({ input: { ignored: true } }); + child.end(); + span.end(); + + expect(span.isRecording()).toBe(false); + expect(child.isRecording()).toBe(false); + expect(span.id).not.toBe(""); + expect(span.inject().traceparent).toMatch(/-00$/); + await memoryLogger.flush(); + expect(await memoryLogger.drain()).toHaveLength(0); + }); + + test("allows a root override without changing the logger policy", async () => { + const memoryLogger = _exportsForTestingOnly.useTestBackgroundLogger(); + const logger = initLogger({ + projectName: "sampling", + projectId: "sampling-project", + sampleRate: 0, + }); + + const recorded = logger.startSpan({ name: "kept", sampleRate: 1 }); + const dropped = logger.startSpan({ name: "dropped" }); + recorded.end(); + dropped.end(); + + expect(logger.sampleRate).toBe(0); + expect(recorded.isRecording()).toBe(true); + expect(dropped.isRecording()).toBe(false); + await memoryLogger.flush(); + expect(await memoryLogger.drain()).toHaveLength(1); + }); + + test("inherits the parent decision over a local logger rate", () => { + const logger = initLogger({ + projectName: "sampling", + projectId: "sampling-project", + sampleRate: 0, + }); + + logger.traced((parent) => { + expect(parent.isRecording()).toBe(false); + const child = logger.startSpan({ name: "child", sampleRate: 1 }); + expect(child.isRecording()).toBe(false); + }); + }); + + test.each([null, "0.5", NaN, Infinity, -0.1, 1.1])( + "rejects invalid logger sample rate %j", + (sampleRate) => { + expect(() => + initLogger({ + projectName: "sampling", + sampleRate: sampleRate as number, + }), + ).toThrow(RangeError); + }, + ); +}); + test("init validation", () => { expect(() => init({})).toThrow( "Must specify at least one of project or projectId", diff --git a/js/src/logger.ts b/js/src/logger.ts index 9733c3030..996fa0059 100644 --- a/js/src/logger.ts +++ b/js/src/logger.ts @@ -28,6 +28,12 @@ import { parseBaggage, parseTraceparent, } from "./propagation"; +import { + isTraceFlagsSampled, + normalizeTraceFlags, + shouldSampleTraceId, + validateSampleRate, +} from "./sampling"; import { _urljoin, AnyDatasetRecord, @@ -200,6 +206,8 @@ const InlineAttachmentReferenceSchema = z.object({ export interface ContextParentSpanIds { rootSpanId: string; spanParents: string[]; + /** Raw W3C trace-flags from the selected parent, when available. */ + traceFlags?: string; } export class LoginInvalidOrgError extends Error { @@ -276,6 +284,10 @@ const RESUME_SPAN_WITHOUT_INITIAL_WRITE = Symbol( "braintrust.resume-span-without-initial-write", ); const INTERNAL_SPAN_CONTEXT = Symbol("braintrust.internal-span-context"); +const LOGGER_SAMPLE_RATE = Symbol("braintrust.logger-sample-rate"); + +const loggerSampleRates = new WeakMap(); +const spanDefaultSampleRates = new WeakMap(); type InitialSpanWriteAsMergeArg = { readonly [INITIAL_SPAN_WRITE_AS_MERGE]?: true; @@ -286,6 +298,28 @@ type InternalSpanContextArg = { readonly [INTERNAL_SPAN_CONTEXT]?: Record; }; +type InternalLoggerSampleRateArg = { + readonly [LOGGER_SAMPLE_RATE]?: number; +}; + +type ResolvedPropagatedState = Readonly< + Omit & { traceFlags: string } +>; + +function validateStartSpanSampleRate(args?: StartSpanArgs): void { + if (args?.sampleRate !== undefined) { + validateSampleRate(args.sampleRate); + } +} + +function getLoggerSampleRate(logger: object): number { + return loggerSampleRates.get(logger) ?? 1; +} + +function getSpanDefaultSampleRate(span: object): number { + return spanDefaultSampleRates.get(span) ?? 1; +} + export type StartSpanArgs = { name?: string; type?: SpanType; @@ -302,6 +336,11 @@ export type StartSpanArgs = { propagatedEvent?: StartSpanEventArgs; spanId?: string; parentSpanIds?: ParentSpanIds | MultiParentSpanIds; + /** + * Sampling rate for a newly-created project-log trace. Existing parent + * decisions always win, so this never resamples a child. + */ + sampleRate?: number; }; export type EndSpanArgs = { @@ -321,6 +360,8 @@ export interface Exportable { * We suggest using one of the various `traced` methods, instead of creating Spans directly. See {@link Span.traced} for full details. */ export interface Span extends Exportable { + /** Whether this span records trace data, rather than only carrying context. */ + isRecording(): boolean; /** * Row ID of the span. */ @@ -547,9 +588,13 @@ class BraintrustContextManager extends ContextManager { return undefined; } + const traceCarrier = currentSpan as Span & { + _getTraceFlags?: () => string; + }; return { rootSpanId: currentSpan.rootSpanId, spanParents: [currentSpan.spanId], + traceFlags: traceCarrier._getTraceFlags?.(), }; } @@ -611,12 +656,17 @@ export class NoopSpan implements Span { public log(_: ExperimentLogPartialArgs) {} + public isRecording(): boolean { + return false; + } + public logFeedback(_event: Omit) {} public traced( callback: (span: Span) => R, - _1?: StartSpanArgs & SetCurrentArg, + args?: StartSpanArgs & SetCurrentArg, ): R { + validateStartSpanSampleRate(args); return callback(this); } @@ -624,7 +674,8 @@ export class NoopSpan implements Span { return undefined; } - public startSpan(_1?: StartSpanArgs) { + public startSpan(args?: StartSpanArgs) { + validateStartSpanSampleRate(args); return this; } @@ -663,8 +714,9 @@ export class NoopSpan implements Span { public startSpanWithParents( _spanId: string, _spanParents: string[], - _args?: StartSpanArgs, + args?: StartSpanArgs, ): Span { + validateStartSpanSampleRate(args); return this; } @@ -2240,6 +2292,7 @@ export function _internalResumeSpan({ propagatedEvent: (components.data.propagated_event ?? undefined) as | StartSpanEventArgs | undefined, + propagatedState: { traceFlags: components.data.trace_flags }, [RESUME_SPAN_WITHOUT_INITIAL_WRITE]: true, }); } @@ -2490,10 +2543,13 @@ function startSpanParentArgs(args: { | StartSpanEventArgs | undefined); const propagatedState = args.propagatedState ?? parentPropagatedState; - if (propagatedState) { + if (propagatedState || parentComponents.data.trace_flags) { const { braintrustParent: _ignoredBraintrustParent, ...w3cState } = - propagatedState; - argPropagatedState = w3cState; + propagatedState ?? {}; + argPropagatedState = { + ...w3cState, + traceFlags: parentComponents.data.trace_flags ?? w3cState.traceFlags, + }; } } else { argParentObjectId = args.parentObjectId; @@ -2560,6 +2616,11 @@ export class Logger implements Exportable { return this.state; } + /** The immutable default sampling rate for new project-log roots. */ + public get sampleRate(): number { + return getLoggerSampleRate(this); + } + private parentObjectType() { return SpanObjectTypeV3.PROJECT_LOGS; } @@ -2651,6 +2712,7 @@ export class Logger implements Exportable { * See {@link traced} for full details. */ public startSpan(args?: StartSpanArgs): Span { + validateStartSpanSampleRate(args); this.calledStartSpan = true; return this.startSpanImpl(args); } @@ -2661,6 +2723,7 @@ export class Logger implements Exportable { // Sometimes `args` gets passed directly into this function, and it contains an undefined value for `state`. // To ensure that we always use this logger's state, we override the `state` argument no matter what. state: this.state, + [LOGGER_SAMPLE_RATE]: this.sampleRate, ...startSpanParentArgs({ state: this.state, parent: args?.parent, @@ -4687,6 +4750,8 @@ export type InitLoggerOptions = FullLoginOptions & { setCurrent?: boolean; state?: BraintrustState; orgProjectMetadata?: OrgProjectMetadata; + /** Fraction of new project-log traces to record. Defaults to 1. */ + sampleRate?: number; } & AsyncFlushArg; /** @@ -4708,6 +4773,11 @@ export type InitLoggerOptions = FullLoginOptions & { export function initLogger( options: Readonly> = {}, ) { + const configuredSampleRate = options?.sampleRate; + const sampleRate = + configuredSampleRate === undefined + ? 1 + : validateSampleRate(configuredSampleRate); const { projectName, projectId, @@ -4764,6 +4834,7 @@ export function initLogger( computeMetadataArgs, linkArgs, }); + loggerSampleRates.set(ret, sampleRate); if (options.setCurrent ?? true) { state.currentLogger = ret as Logger; } @@ -5497,7 +5568,7 @@ function getSpanParentObjectAndPropagatedState( }; } - const experiment = currentExperiment(); + const experiment = currentExperiment({ state }); if (experiment) { return { parentObject: experiment, propagatedState: undefined }; } @@ -6400,6 +6471,7 @@ function startSpanAndIsLogger( OptionalStateArg & InternalSpanContextArg, ): { span: Span; isSyncFlushLogger: boolean } { + validateStartSpanSampleRate(args); const state = args?.state ?? _globalState; // Resolve the parent object and any forwarded W3C state in one pass, so we @@ -6446,7 +6518,11 @@ function startSpanAndIsLogger( ((parentObject.data.propagated_event ?? undefined) as | StartSpanEventArgs | undefined), - propagatedState, + propagatedState: { + ...(propagatedState ?? {}), + traceFlags: + parentObject.data.trace_flags ?? propagatedState?.traceFlags, + }, }); return { span, @@ -7282,6 +7358,7 @@ export class Experiment * See {@link traced} for full details. */ public startSpan(args?: StartSpanArgs): Span { + validateStartSpanSampleRate(args); this.calledStartSpan = true; return this.startSpanImpl(args); } @@ -7575,6 +7652,8 @@ interface ResolvedSpanIds { spanId: string; rootSpanId: string; spanParents: string[] | undefined; + isRoot: boolean; + traceFlags: string | undefined; } /** @@ -7608,6 +7687,8 @@ function _resolveSpanIds( "parentSpanIds" in parentSpanIds ? parentSpanIds.parentSpanIds : [parentSpanIds.spanId], + isRoot: false, + traceFlags: undefined, }; } @@ -7619,6 +7700,8 @@ function _resolveSpanIds( spanId: resolvedSpanId, rootSpanId: parentInfo.rootSpanId, spanParents: parentInfo.spanParents, + isRoot: false, + traceFlags: parentInfo.traceFlags, }; } } @@ -7640,6 +7723,8 @@ function _resolveSpanIds( spanId: resolvedSpanId, rootSpanId: resolvedRootSpanId, spanParents: undefined, + isRoot: true, + traceFlags: undefined, }; } @@ -7665,12 +7750,10 @@ export class SpanImpl implements Span { private _rootSpanId: string; private _spanParents: string[] | undefined; - // Inbound W3C trace-context state (tracestate + raw traceparent flags) to - // forward on outbound propagation. Captured at the span that received it (via - // extractTraceContextFromHeaders) and inherited by all subspans, so that any - // inject() within the trace re-emits the upstream state unchanged, per the W3C - // Trace Context spec. Not interpreted. - private _propagatedState: PropagatedState | undefined; + // Canonical W3C propagation state for this trace. Its flags byte is both + // forwarded unchanged and the sole source of the native recording decision. + private readonly _propagatedState: ResolvedPropagatedState; + private readonly _recording: boolean; public kind = "span" as const; @@ -7686,10 +7769,11 @@ export class SpanImpl implements Span { propagatedState?: PropagatedState | undefined; } & Omit & InitialSpanWriteAsMergeArg & - InternalSpanContextArg, + InternalSpanContextArg & + InternalLoggerSampleRateArg, ) { + validateStartSpanSampleRate(args); this._state = args.state; - this._propagatedState = args.propagatedState; const instrumentationName = getSpanInstrumentationName(args) ?? INSTRUMENTATION_NAMES.BRAINTRUST_JS_LOGGER; @@ -7761,10 +7845,32 @@ export class SpanImpl implements Span { this._rootSpanId = resolvedIds.rootSpanId; this._spanParents = resolvedIds.spanParents; + const inheritedTraceFlags = + args.propagatedState?.traceFlags ?? resolvedIds.traceFlags; + const traceFlags = resolvedIds.isRoot + ? this.parentObjectType === SpanObjectTypeV3.PROJECT_LOGS + ? shouldSampleTraceId( + this._rootSpanId, + args.sampleRate ?? args[LOGGER_SAMPLE_RATE] ?? 1, + ) + ? "01" + : "00" + : "01" + : normalizeTraceFlags(inheritedTraceFlags, { warnOnInvalid: true }); + this._propagatedState = { + ...(args.propagatedState ?? {}), + traceFlags, + }; + this._recording = isTraceFlagsSampled(traceFlags); + spanDefaultSampleRates.set( + this, + args[LOGGER_SAMPLE_RATE] ?? getSpanDefaultSampleRate(this), + ); + // Deterministic spans can be initialized concurrently by separate // workflow executions, so their first write must not replace later merges. this.isMerge = args[INITIAL_SPAN_WRITE_AS_MERGE] === true; - if (!args[RESUME_SPAN_WITHOUT_INITIAL_WRITE]) { + if (this._recording && !args[RESUME_SPAN_WITHOUT_INITIAL_WRITE]) { this.logInternal({ event, internalData }); } this.isMerge = true; @@ -7784,6 +7890,15 @@ export class SpanImpl implements Span { return this._id; } + public isRecording(): boolean { + return this._recording; + } + + /** @internal Used by the optional OTel bridge and native context manager. */ + public _getTraceFlags(): string { + return this._propagatedState.traceFlags; + } + public get spanId(): string { return this._spanId; } @@ -7797,7 +7912,8 @@ export class SpanImpl implements Span { } public setAttributes(args: Omit): void { - this.logInternal({ internalData: { span_attributes: args } }); + const { sampleRate: _ignoredSampleRate, ...spanAttributes } = args; + this.logInternal({ internalData: { span_attributes: spanAttributes } }); } public setSpanParents(parents: string[]): void { @@ -7817,6 +7933,7 @@ export class SpanImpl implements Span { // set of fields which we want to log in just one of the span rows. internalData?: Partial; }): void { + if (!this._recording) return; const [serializableInternalData, lazyInternalData] = splitLoggingData({ event, internalData, @@ -7879,6 +7996,7 @@ export class SpanImpl implements Span { } public logFeedback(event: Omit): void { + if (!this._recording) return; logFeedbackImpl(this._state, this.parentObjectType, this.parentObjectId, { ...event, id: this.id, @@ -7914,6 +8032,7 @@ export class SpanImpl implements Span { return new SpanImpl({ state: this._state, ...args, + [LOGGER_SAMPLE_RATE]: getSpanDefaultSampleRate(this), ...startSpanParentArgs({ state: this._state, parent: args?.parent, @@ -7939,6 +8058,7 @@ export class SpanImpl implements Span { return new SpanImpl({ state: this._state, ...args, + [LOGGER_SAMPLE_RATE]: getSpanDefaultSampleRate(this), ...startSpanParentArgs({ state: this._state, parent: args?.parent, @@ -7968,7 +8088,9 @@ export class SpanImpl implements Span { public async export(): Promise { // Disable span cache since remote function spans won't be in the local cache - this._state.spanCache.disable(); + if (this._recording) { + this._state.spanCache.disable(); + } return new (getSpanComponentsClass())({ object_type: this.parentObjectType, @@ -7981,6 +8103,7 @@ export class SpanImpl implements Span { span_id: this._spanId, root_span_id: this._rootSpanId, propagated_event: this.propagatedEvent, + trace_flags: this._propagatedState.traceFlags, }).toStr(); } @@ -8033,12 +8156,18 @@ export class SpanImpl implements Span { } public async permalink(): Promise { + if (!this._recording) { + return NOOP_SPAN_PERMALINK; + } return await permalink(await this.export(), { state: this._state, }); } public link(): string { + if (!this._recording) { + return NOOP_SPAN_PERMALINK; + } if (!this.id) { return NOOP_SPAN_PERMALINK; } diff --git a/js/src/sampling.test.ts b/js/src/sampling.test.ts new file mode 100644 index 000000000..a5b2219ec --- /dev/null +++ b/js/src/sampling.test.ts @@ -0,0 +1,51 @@ +import { describe, expect, test } from "vitest"; + +import { + isTraceFlagsSampled, + normalizeTraceFlags, + shouldSampleTraceId, + validateSampleRate, +} from "./sampling"; + +describe("validateSampleRate", () => { + test.each([0, 0.2, 1])("accepts %s", (sampleRate) => { + expect(validateSampleRate(sampleRate)).toBe(sampleRate); + }); + + test.each([-1, 1.1, NaN, Infinity, -Infinity, null, "0.5"])( + "rejects %j", + (sampleRate) => { + expect(() => validateSampleRate(sampleRate)).toThrow(RangeError); + }, + ); +}); + +describe("shouldSampleTraceId", () => { + test("handles exact endpoints", () => { + expect(shouldSampleTraceId("0".repeat(32), 0)).toBe(false); + expect(shouldSampleTraceId("0".repeat(32), 1)).toBe(true); + }); + + test.each([ + [0.5, "0000000000000000007fffffffffffff", false], + [0.5, "00000000000000000080000000000000", true], + [0.25, "000000000000000000bfffffffffffff", false], + [0.25, "000000000000000000c0000000000000", true], + ])("uses exact boundary at rate %s", (rate, traceId, expected) => { + expect(shouldSampleTraceId(traceId, rate)).toBe(expected); + }); + + test("normalizes legacy UUID trace ids", () => { + expect( + shouldSampleTraceId("00000000-0000-0000-0080-000000000000", 0.5), + ).toBe(true); + }); +}); + +describe("trace flags", () => { + test("preserves raw bits while inspecting sampled bit", () => { + expect(normalizeTraceFlags("FF")).toBe("ff"); + expect(isTraceFlagsSampled("02")).toBe(false); + expect(isTraceFlagsSampled("03")).toBe(true); + }); +}); diff --git a/js/src/sampling.ts b/js/src/sampling.ts new file mode 100644 index 000000000..32996c305 --- /dev/null +++ b/js/src/sampling.ts @@ -0,0 +1,80 @@ +import { debugLogger } from "./debug-logger"; + +const RANDOMNESS_BITS = 56n; +const RANDOMNESS_SCALE = 1n << RANDOMNESS_BITS; +const RANDOMNESS_SCALE_NUMBER = 2 ** 56; +const TRACE_ID_RE = /^[0-9a-f]{32}$/; +const TRACE_FLAGS_RE = /^[0-9a-f]{2}$/i; + +let warnedInvalidTraceId = false; +let warnedInvalidTraceFlags = false; + +/** Validate a public root sampling rate. */ +export function validateSampleRate(sampleRate: unknown): number { + if ( + typeof sampleRate !== "number" || + !Number.isFinite(sampleRate) || + sampleRate < 0 || + sampleRate > 1 + ) { + throw new RangeError("sampleRate must be a finite number between 0 and 1"); + } + return sampleRate; +} + +/** + * Return whether a new trace should be recorded. The calculation is stable for + * a trace id, allowing distributed participants and retries to make the same + * decision without mutable random state. + */ +export function shouldSampleTraceId( + traceId: string, + sampleRate: number, +): boolean { + if (sampleRate <= 0) return false; + if (sampleRate >= 1) return true; + + const normalized = traceId.replaceAll("-", "").toLowerCase(); + if (!TRACE_ID_RE.test(normalized)) { + if (!warnedInvalidTraceId) { + warnedInvalidTraceId = true; + debugLogger.warn( + "Unable to deterministically sample an invalid trace id; recording it for compatibility.", + ); + } + return true; + } + + const randomness = BigInt(`0x${normalized.slice(-14)}`); + const keepCount = BigInt(Math.floor(sampleRate * RANDOMNESS_SCALE_NUMBER)); + const threshold = RANDOMNESS_SCALE - keepCount; + return randomness >= threshold; +} + +/** Normalize an inbound flags byte, defaulting old/custom context to sampled. */ +export function normalizeTraceFlags( + traceFlags: string | undefined, + options?: { warnOnInvalid?: boolean }, +): string { + if (traceFlags === undefined) return "01"; + if (TRACE_FLAGS_RE.test(traceFlags)) return traceFlags.toLowerCase(); + + if (options?.warnOnInvalid && !warnedInvalidTraceFlags) { + warnedInvalidTraceFlags = true; + debugLogger.warn( + "Received invalid in-process trace flags; treating the trace as sampled for compatibility.", + ); + } + return "01"; +} + +/** The W3C sampled decision is the least-significant bit of the flags byte. */ +export function isTraceFlagsSampled(traceFlags: string): boolean { + return (parseInt(traceFlags, 16) & 0x01) !== 0; +} + +/** Test-only reset for warn-once state. */ +export function _resetSamplingWarningsForTests(): void { + warnedInvalidTraceId = false; + warnedInvalidTraceFlags = false; +} diff --git a/js/util/span_identifier_v3.test.ts b/js/util/span_identifier_v3.test.ts new file mode 100644 index 000000000..48ca7aa1e --- /dev/null +++ b/js/util/span_identifier_v3.test.ts @@ -0,0 +1,16 @@ +import { expect, test } from "vitest"; + +import { SpanComponentsV3, SpanObjectTypeV3 } from "./span_identifier_v3"; + +test("V3 span identifiers preserve optional trace flags", () => { + const serialized = new SpanComponentsV3({ + object_type: SpanObjectTypeV3.PROJECT_LOGS, + object_id: "project-123", + row_id: "row-456", + span_id: "span-789", + root_span_id: "root-012", + trace_flags: "00", + }).toStr(); + + expect(SpanComponentsV3.fromStr(serialized).data.trace_flags).toBe("00"); +}); diff --git a/js/util/span_identifier_v3.ts b/js/util/span_identifier_v3.ts index f057f798b..170d967a3 100644 --- a/js/util/span_identifier_v3.ts +++ b/js/util/span_identifier_v3.ts @@ -88,6 +88,10 @@ export const spanComponentsV3Schema = z // `propagated_event`. This will required zod-ifying the contents of // sdk/js/util/object.ts. propagated_event: z.record(z.unknown()).nullish(), + trace_flags: z + .string() + .regex(/^[0-9a-fA-F]{2}$/) + .optional(), }) .and( z.union([ @@ -128,6 +132,7 @@ export class SpanComponentsV3 { compute_object_metadata_args: this.data.compute_object_metadata_args || undefined, propagated_event: this.data.propagated_event || undefined, + trace_flags: this.data.trace_flags || undefined, }; const allBuffers: Array = []; allBuffers.push( diff --git a/js/util/span_identifier_v4.test.ts b/js/util/span_identifier_v4.test.ts index 6084607d8..ed968424e 100644 --- a/js/util/span_identifier_v4.test.ts +++ b/js/util/span_identifier_v4.test.ts @@ -82,6 +82,19 @@ describe("SpanComponentsV4", () => { originalData.propagated_event, ); }); + + test("preserves an optional raw trace-flags byte", () => { + const serialized = new SpanComponentsV4({ + object_type: SpanObjectTypeV3.PROJECT_LOGS, + object_id: "project-123", + row_id: "row-456", + span_id: "fedcba0987654321", + root_span_id: "0123456789abcdef0123456789abcdef", + trace_flags: "02", + }).toStr(); + + expect(SpanComponentsV4.fromStr(serialized).data.trace_flags).toBe("02"); + }); }); describe("Hex string compression", () => { diff --git a/js/util/span_identifier_v4.ts b/js/util/span_identifier_v4.ts index e92e3eddf..d93b1d940 100644 --- a/js/util/span_identifier_v4.ts +++ b/js/util/span_identifier_v4.ts @@ -92,6 +92,10 @@ export const spanComponentsV4Schema = z .object({ object_type: spanObjectTypeV3EnumSchema, propagated_event: z.record(z.unknown()).nullish(), + trace_flags: z + .string() + .regex(/^[0-9a-fA-F]{2}$/) + .optional(), }) .and( z.union([ @@ -134,6 +138,7 @@ export class SpanComponentsV4 { compute_object_metadata_args: this.data.compute_object_metadata_args || undefined, propagated_event: this.data.propagated_event || undefined, + trace_flags: this.data.trace_flags || undefined, }; // Filter out undefined values @@ -215,6 +220,7 @@ export class SpanComponentsV4 { jsonObj["span_id"] = v3Components.data.span_id; jsonObj["root_span_id"] = v3Components.data.root_span_id; jsonObj["propagated_event"] = v3Components.data.propagated_event; + jsonObj["trace_flags"] = v3Components.data.trace_flags; } else { // V4 binary format jsonObj["object_type"] = rawBytes[1]; From 67e0e069c6f4acadf6c4b46ec3aefc807c564c7d Mon Sep 17 00:00:00 2001 From: Starfolk Date: Fri, 14 Aug 2026 21:42:11 +0000 Subject: [PATCH 2/2] fix: remove unused sampling helper --- js/src/sampling.ts | 6 ------ 1 file changed, 6 deletions(-) diff --git a/js/src/sampling.ts b/js/src/sampling.ts index 32996c305..a9f811b4e 100644 --- a/js/src/sampling.ts +++ b/js/src/sampling.ts @@ -72,9 +72,3 @@ export function normalizeTraceFlags( export function isTraceFlagsSampled(traceFlags: string): boolean { return (parseInt(traceFlags, 16) & 0x01) !== 0; } - -/** Test-only reset for warn-once state. */ -export function _resetSamplingWarningsForTests(): void { - warnedInvalidTraceId = false; - warnedInvalidTraceFlags = false; -}