span hooks

braintrust-java-agent/smoke-test/otel-agent/src/main/java/dev/braintrust/smoketest/otelagent/OtelAgentSmokeTest.java

@@ -113,6 +113,11 @@ private static void assertOtelCollector(MockOtlpCollector collector) throws Exce
         assertTrue(
                 span.endTimeUnixNano > span.startTimeUnixNano,
                 "span endTimeUnixNano <= startTimeUnixNano");
+        if (checkIsRunningWithBraintrustJavaagent()) {
+            assertNotNull(
+                    span.stringAttr("braintrust.context_json"),
+                    "OTel collector span missing Braintrust provenance");
+        }
 
         System.out.println("[otel-agent-smoke-test] OTel collector assertions: OK");
     }
@@ -138,6 +143,9 @@ private static void assertBraintrustBackend(MockOtlpCollector backend) throws Ex
 
         var span = backend.findSpanByName(SPAN_NAME);
         assertNotNull(span, "Braintrust backend missing span with name '" + SPAN_NAME + "'");
+        assertNotNull(
+                span.stringAttr("braintrust.context_json"),
+                "Braintrust backend span missing provenance");
 
         System.out.println("[otel-agent-smoke-test] Braintrust backend assertions: OK");
     }

braintrust-sdk/build.gradle

@@ -145,6 +145,7 @@ dependencies {
 
   testImplementation "org.slf4j:slf4j-simple:${slf4jVersion}"
   testImplementation "io.opentelemetry:opentelemetry-sdk-testing:${otelVersion}"
+  testImplementation 'io.opentelemetry.proto:opentelemetry-proto:1.11.0-alpha'
   testImplementation "org.junit.jupiter:junit-jupiter:${junitVersion}"
   testImplementation "org.junit.jupiter:junit-jupiter-params:${junitVersion}"
   testImplementation 'org.wiremock:wiremock:3.13.1'

braintrust-sdk/src/main/java/dev/braintrust/config/BraintrustConfig.java

@@ -2,13 +2,17 @@
 
 import dev.braintrust.Braintrust;
 import dev.braintrust.api.BraintrustOpenApiClient;
+import dev.braintrust.trace.SpanCustomizer;
 import java.net.URI;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.time.Duration;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Optional;
 import javax.annotation.Nullable;
 import javax.net.ssl.SSLContext;
@@ -64,6 +68,9 @@ public final class BraintrustConfig extends BaseConfig {
     /** Custom X509 trust manager for OTLP exporter. Builder-only field, not backed by envars. */
     private final X509TrustManager x509TrustManager;
 
+    /** Immutable export customizers in registration order. Builder-only, not backed by envars. */
+    private final List<SpanCustomizer> spanCustomizers;
+
     /**
      * Sizes the thread pools the SDK creates for batch work (e.g. eval execution), and so how much
      * of that work runs at once. Used only when the caller has not supplied an executor of their
@@ -123,14 +130,15 @@ public static BraintrustConfig of(String... envOverrides) {
         for (int i = 0; i < envOverrides.length - 1; i = i + 2) {
             overridesMap.put(envOverrides[i], envOverrides[i + 1]);
         }
-        return new BraintrustConfig(overridesMap, null, null);
+        return new BraintrustConfig(overridesMap, null, null, List.of());
     }
 
     @SneakyThrows
     private BraintrustConfig(
             Map<String, String> envOverrides,
             SSLContext sslContext,
-            X509TrustManager x509TrustManager) {
+            X509TrustManager x509TrustManager,
+            List<SpanCustomizer> spanCustomizers) {
         super(envOverrides);
         if (defaultProjectId.isEmpty() && defaultProjectName.isEmpty()) {
             // should never happen
@@ -145,6 +153,7 @@ private BraintrustConfig(
                             + otelMaxQueueSize);
         }
 
+        this.spanCustomizers = List.copyOf(spanCustomizers);
         this.sslContext = sslContext != null ? sslContext : SSLContext.getDefault();
         if (x509TrustManager != null) {
             this.x509TrustManager = x509TrustManager;
@@ -191,6 +200,7 @@ public static class Builder {
         private final Map<String, String> envOverrides = new HashMap<>();
         private SSLContext sslContext;
         private X509TrustManager x509TrustManager;
+        private final List<SpanCustomizer> spanCustomizers = new ArrayList<>();
 
         public Builder apiKey(String value) {
             envOverrides.put("BRAINTRUST_API_KEY", value);
@@ -325,8 +335,20 @@ public Builder otelExportIntervalMillis(int exportIntervalMillis) {
             return this;
         }
 
+        /**
+         * Adds a span customizer and run its hooks over all spans which pass through braintrust
+         * export
+         *
+         * @throws NullPointerException if {@code customizer} is null
+         */
+        public Builder addSpanCustomizer(SpanCustomizer customizer) {
+            spanCustomizers.add(Objects.requireNonNull(customizer, "customizer"));
+            return this;
+        }
+
         public BraintrustConfig build() {
-            return new BraintrustConfig(envOverrides, sslContext, x509TrustManager);
+            return new BraintrustConfig(
+                    envOverrides, sslContext, x509TrustManager, spanCustomizers);
         }
     }
 

braintrust-sdk/src/main/java/dev/braintrust/trace/BraintrustSampler.java

@@ -4,11 +4,14 @@
 import java.util.List;
 
 /**
- * A filter that decides whether a finalized span should be exported. Samplers are evaluated in
- * {@link BraintrustSpanProcessor#onEnd} after all attributes have been set on the span.
+ * A filter that decides whether a span should be enriched and exported. Samplers are evaluated once
+ * in {@link BraintrustSpanProcessor#onEnding}, before enrichment and delegated ending hooks.
  */
 interface BraintrustSampler {
-    /** Returns {@code true} if the span should be exported, {@code false} to discard it. */
+    /**
+     * Returns {@code true} to enrich and export the span, {@code false} to discard it. The span is
+     * still ending; implementations must only read it, not downcast it to mutate it.
+     */
     boolean sample(ReadableSpan span);
 
     /** Keeps only spans that have at least one attribute with a known AI instrumentation prefix. */

braintrust-sdk/src/main/java/dev/braintrust/trace/BraintrustSpanExporter.java

@@ -33,8 +33,18 @@ public CompletableResultCode export(Collection<SpanData> spans) {
             return CompletableResultCode.ofSuccess();
         }
 
-        // Group spans by their parent (project or experiment)
-        var spansByParent = spans.stream().collect(Collectors.groupingBy(this::getParentFromSpan));
+        // Finish customization before sending any group, so failed redaction cannot leak a batch.
+        Map<String, List<SpanData>> spansByParent;
+        try {
+            var exportSpans = spans.stream();
+            if (!config.spanCustomizers().isEmpty()) {
+                exportSpans = exportSpans.map(this::customizeSpan);
+            }
+            spansByParent = exportSpans.collect(Collectors.groupingBy(this::getParentFromSpan));
+        } catch (Exception e) {
+            log.error("Failed to customize spans for export", e);
+            return CompletableResultCode.ofFailure();
+        }
 
         // Export each group with the appropriate x-bt-parent header
         var results =
@@ -49,6 +59,30 @@ public CompletableResultCode export(Collection<SpanData> spans) {
         return combined;
     }
 
+    private SpanData customizeSpan(SpanData span) {
+        var traceId = span.getTraceId();
+        var spanId = span.getSpanId();
+        var parentSpanId = span.getParentSpanId();
+        var current = span;
+        for (var customizer : config.spanCustomizers()) {
+            current = customizer.onSpanExport(current);
+            if (current == null) {
+                throw new IllegalStateException("SpanCustomizer.onSpanExport must not return null");
+            }
+            if (!traceId.equals(current.getTraceId())
+                    || !spanId.equals(current.getSpanId())
+                    || !parentSpanId.equals(current.getParentSpanId())
+                    || !traceId.equals(current.getSpanContext().getTraceId())
+                    || !spanId.equals(current.getSpanContext().getSpanId())
+                    || !parentSpanId.equals(current.getParentSpanContext().getSpanId())) {
+                throw new IllegalStateException(
+                        "SpanCustomizer.onSpanExport must not change trace ID, span ID, or parent"
+                                + " span ID");
+            }
+        }
+        return current;
+    }
+
     private String getParentFromSpan(SpanData span) {
         var parent = span.getAttributes().get(BraintrustSpanProcessor.PARENT);
         if (parent != null) {