feat: add client-side trace sampling

.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.

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.

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);
+  });
+});

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);

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"),
     };
   }
 

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

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 } {

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")

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<TChannel>,
-      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<TChannel>,
+        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<TChannel extends AnyAsyncChannel>(
       const asyncEndEvent = event as AsyncEndOf<TChannel>;
       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<TChannel extends AnyAsyncChannel>(
       const asyncEndEvent = event as AsyncEndOf<TChannel>;
       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) => {

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);