feat: add span customizers (SDK-316)

README.md

@@ -20,6 +20,7 @@ This is the official Ruby SDK for [Braintrust](https://www.braintrust.dev), for
   - [Supported providers](#supported-providers)
   - [Manually applying instrumentation](#manually-applying-instrumentation)
   - [Creating custom spans](#creating-custom-spans)
+  - [Span customizers](#span-customizers)
   - [Attachments](#attachments)
   - [Viewing traces](#viewing-traces)
 - [Evals](#evals)
@@ -120,6 +121,7 @@ Braintrust.init
 | `filter_ai_spans` | `ENV['BRAINTRUST_OTEL_FILTER_AI_SPANS']` | Only export AI-related spans                                                |
 | `org_name`        | `ENV['BRAINTRUST_ORG_NAME']`             | Organization name                                                           |
 | `set_global`      | `true`                                   | Set as global state. Set to `false` for isolated instances                  |
+| `span_customizers` | `[]`                                   | Ordered objects with optional synchronous export hooks (see [Span customizers](#span-customizers)) |
 
 **Example with options:**
 
@@ -192,6 +194,40 @@ tracer.in_span("process-request") do |span|
 end
 ```
 
+### Span customizers
+
+Register customizers before creating spans to redact or transform outgoing telemetry:
+
+```ruby
+require "braintrust"
+
+class RedactContent < Braintrust::SpanCustomizer
+  def on_span_export(span)
+    %w[braintrust.input_json braintrust.output_json].each do |key|
+      span.attributes[key] = JSON.generate("[redacted]") if span.attributes.key?(key)
+    end
+    span
+  end
+end
+
+Braintrust.init(
+  default_project: "my-project",
+  span_customizers: [RedactContent.new]
+)
+```
+
+`Braintrust::SpanCustomizer` is an extensible base class with a no-op `on_span_export`. Any object may be registered; an omitted hook is also a no-op. `Braintrust::Config.new` / `.from_env`, `Braintrust::State.from_env`, and a directly constructed `Braintrust::Trace::SpanExporter` also accept `span_customizers:`. Registration is programmatic, not environment-based. Configuration and exporters retain frozen copies of the ordered list, not copies of the customizer objects.
+
+Hooks run synchronously in registration order, each receiving its predecessor's result. The argument is a completed `OpenTelemetry::SDK::Trace::SpanData` snapshot, after Braintrust origin metadata is added but before destination grouping or OTLP serialization. Hooks apply to **all completed spans reaching the Braintrust exporter**, including manual, evaluation, and instrumented spans that pass any configured span filters. An unrelated exporter supplied through `exporter:` is not wrapped with these hooks.
+
+You may mutate and return the snapshot or return a replacement `SpanData`; replacement is not an implicit merge. Nested span data is copied so customization does not modify application/provider return values, live spans, or another exporter's data. OTel resource objects retain their normal API; replace resources with `OpenTelemetry::SDK::Resources::Resource.create(...)` when changing resource attributes. Added attributes, events, and links have their recorded counts adjusted to prevent negative OTLP dropped counts.
+
+Always return valid, serializable span data, never `nil` to drop a span. Preserve `trace_id`, `span_id`, and `parent_span_id`; the exporter checks these after every hook. You may change `braintrust.parent` to route the result to another project or experiment. Existing destination/header behavior otherwise remains unchanged.
+
+Customization is **fail-closed for the whole batch**: an exception, invalid return, changed protected ID, or serialization failure returns an export failure and sends none of that batch. The SDK logs the failure without falling back to unredacted originals. Hook failures do not affect application return values.
+
+Keep hooks fast and avoid blocking I/O; they may run on a background export thread. OTLP transport retries reuse the already customized bytes. Submitting the batch through `export` again invokes the hooks again on fresh snapshots, so do not assume one invocation per logical span. Without customizers, the existing export path is unchanged.
+
 ### Attachments
 
 Log binary data (images, PDFs, audio) in your traces:

lib/braintrust.rb

@@ -2,6 +2,7 @@
 
 require_relative "braintrust/version"
 require_relative "braintrust/config"
+require_relative "braintrust/span_customizer"
 require_relative "braintrust/state"
 require_relative "braintrust/trace"
 require_relative "braintrust/api"
@@ -42,14 +43,15 @@ class Error < StandardError; end
   # @param tracer_provider [TracerProvider, nil] Optional tracer provider to use instead of creating one
   # @param filter_ai_spans [Boolean, nil] Enable AI span filtering (overrides BRAINTRUST_OTEL_FILTER_AI_SPANS env var)
   # @param span_filter_funcs [Array<Proc>, nil] Custom span filter functions
+  # @param span_customizers [Array<SpanCustomizer>, nil] Ordered synchronous export customizers
   # @param exporter [Exporter, nil] Optional exporter override (for testing)
   # @param auto_instrument [Boolean, Hash, nil] Auto-instrumentation config:
   #   - nil (default): use BRAINTRUST_AUTO_INSTRUMENT env var, default true if not set
   #   - true: explicitly enable
   #   - false: explicitly disable
   #   - Hash with :only or :except keys for filtering
   # @return [State] the created state
-  def self.init(api_key: nil, org_name: nil, default_project: nil, app_url: nil, api_url: nil, set_global: true, blocking_login: false, enable_tracing: true, tracer_provider: nil, filter_ai_spans: nil, span_filter_funcs: nil, exporter: nil, auto_instrument: nil)
+  def self.init(api_key: nil, org_name: nil, default_project: nil, app_url: nil, api_url: nil, set_global: true, blocking_login: false, enable_tracing: true, tracer_provider: nil, filter_ai_spans: nil, span_filter_funcs: nil, span_customizers: nil, exporter: nil, auto_instrument: nil)
     state = State.from_env(
       api_key: api_key,
       org_name: org_name,
@@ -61,6 +63,7 @@ def self.init(api_key: nil, org_name: nil, default_project: nil, app_url: nil, a
       tracer_provider: tracer_provider,
       filter_ai_spans: filter_ai_spans,
       span_filter_funcs: span_filter_funcs,
+      span_customizers: span_customizers,
       exporter: exporter
     )
 

lib/braintrust/config.rb

@@ -8,17 +8,18 @@ module Braintrust
   # and allows overriding with explicit options
   class Config
     attr_reader :api_key, :org_name, :default_project, :app_url, :api_url,
-      :filter_ai_spans, :span_filter_funcs
+      :filter_ai_spans, :span_filter_funcs, :span_customizers
 
     def initialize(api_key: nil, org_name: nil, default_project: nil, app_url: nil, api_url: nil,
-      filter_ai_spans: nil, span_filter_funcs: nil)
+      filter_ai_spans: nil, span_filter_funcs: nil, span_customizers: nil)
       @api_key = api_key
       @org_name = org_name
       @default_project = default_project
       @app_url = app_url
       @api_url = api_url
       @filter_ai_spans = filter_ai_spans
       @span_filter_funcs = span_filter_funcs || []
+      @span_customizers = (span_customizers || []).dup.freeze
     end
 
     # Create a Config from environment variables, with option overrides
@@ -30,9 +31,10 @@ def initialize(api_key: nil, org_name: nil, default_project: nil, app_url: nil,
     # @param api_url [String, nil] API URL (overrides BRAINTRUST_API_URL env var)
     # @param filter_ai_spans [Boolean, nil] Enable AI span filtering (overrides BRAINTRUST_OTEL_FILTER_AI_SPANS env var)
     # @param span_filter_funcs [Array<Proc>, nil] Custom span filter functions
+    # @param span_customizers [Array<SpanCustomizer>, nil] Ordered export customizers (copied and frozen)
     # @return [Config] the created config
     def self.from_env(api_key: nil, org_name: nil, default_project: nil, app_url: nil, api_url: nil,
-      filter_ai_spans: nil, span_filter_funcs: nil)
+      filter_ai_spans: nil, span_filter_funcs: nil, span_customizers: nil)
       # Parse filter_ai_spans from ENV if not explicitly provided
       env_filter_ai_spans = ENV["BRAINTRUST_OTEL_FILTER_AI_SPANS"]
       filter_ai_spans_value = if filter_ai_spans.nil?
@@ -48,7 +50,8 @@ def self.from_env(api_key: nil, org_name: nil, default_project: nil, app_url: ni
         app_url: app_url || ENV["BRAINTRUST_APP_URL"] || "https://www.braintrust.dev",
         api_url: api_url || ENV["BRAINTRUST_API_URL"] || "https://api.braintrust.dev",
         filter_ai_spans: filter_ai_spans_value,
-        span_filter_funcs: span_filter_funcs
+        span_filter_funcs: span_filter_funcs,
+        span_customizers: span_customizers
       )
     end
   end

lib/braintrust/span_customizer.rb

@@ -0,0 +1,17 @@
+# frozen_string_literal: true
+
+module Braintrust
+  # Optional synchronous hooks that transform outgoing telemetry, not live spans.
+  # Subclass this class or supply any object implementing the desired hooks.
+  class SpanCustomizer
+    # Transform a completed, isolated OpenTelemetry SpanData snapshot.
+    # Return this snapshot or a replacement SpanData, never nil. Preserve its
+    # trace_id, span_id and parent_span_id. Exceptions fail the entire batch.
+    #
+    # @param span [OpenTelemetry::SDK::Trace::SpanData]
+    # @return [OpenTelemetry::SDK::Trace::SpanData]
+    def on_span_export(span)
+      span
+    end
+  end
+end

lib/braintrust/state.rb

@@ -24,9 +24,10 @@ class MissingAPIKeyError < ArgumentError; end
     # @param tracer_provider [TracerProvider, nil] Optional tracer provider to use
     # @param filter_ai_spans [Boolean, nil] Enable AI span filtering
     # @param span_filter_funcs [Array<Proc>, nil] Custom span filter functions
+    # @param span_customizers [Array<SpanCustomizer>, nil] Ordered synchronous export customizers
     # @param exporter [Exporter, nil] Optional exporter override (for testing)
     # @return [State] the created state
-    def self.from_env(api_key: nil, org_name: nil, default_project: nil, app_url: nil, api_url: nil, blocking_login: false, enable_tracing: true, tracer_provider: nil, filter_ai_spans: nil, span_filter_funcs: nil, exporter: nil)
+    def self.from_env(api_key: nil, org_name: nil, default_project: nil, app_url: nil, api_url: nil, blocking_login: false, enable_tracing: true, tracer_provider: nil, filter_ai_spans: nil, span_filter_funcs: nil, span_customizers: nil, exporter: nil)
       require_relative "config"
       config = Config.from_env(
         api_key: api_key,
@@ -35,7 +36,8 @@ def self.from_env(api_key: nil, org_name: nil, default_project: nil, app_url: ni
         app_url: app_url,
         api_url: api_url,
         filter_ai_spans: filter_ai_spans,
-        span_filter_funcs: span_filter_funcs
+        span_filter_funcs: span_filter_funcs,
+        span_customizers: span_customizers
       )
       new(
         api_key: config.api_key,

lib/braintrust/trace.rb

@@ -91,7 +91,8 @@ def self.enable(tracer_provider, state: nil, exporter: nil, config: nil)
       # Create OTLP HTTP exporter unless override provided
       exporter ||= SpanExporter.new(
         endpoint: "#{state.api_url}/otel/v1/traces",
-        api_key: state.api_key
+        api_key: state.api_key,
+        span_customizers: config&.span_customizers
       )
 
       # Use SimpleSpanProcessor for InMemorySpanExporter (testing), BatchSpanProcessor for production

lib/braintrust/trace/span_exporter.rb

@@ -3,41 +3,107 @@
 require "opentelemetry/exporter/otlp"
 require_relative "../state"
 require_relative "span_origin"
+require_relative "../logger"
 
 module Braintrust
   module Trace
     # Custom OTLP exporter for the Braintrust backend. On export it:
-    # - stamps span origin provenance onto each SpanData (via the prepended SpanOrigin behavior)
+    # - stamps span origin provenance onto each SpanData
+    # - runs optional customizers on isolated export snapshots
     # - groups spans by braintrust.parent and sets the x-bt-parent header per group,
     #   so the backend routes them to the correct experiment/project
     #
     # Thread safety: BatchSpanProcessor serializes export() calls via its
     # @export_mutex, so @headers mutation here is safe.
     class SpanExporter < OpenTelemetry::Exporter::OTLP::Exporter
-      prepend SpanOrigin
-
       PARENT_ATTR_KEY = SpanProcessor::PARENT_ATTR_KEY
       PARENT_HEADER = "x-bt-parent"
 
       SUCCESS = OpenTelemetry::SDK::Trace::Export::SUCCESS
       FAILURE = OpenTelemetry::SDK::Trace::Export::FAILURE
 
-      def initialize(endpoint:, api_key:)
+      def initialize(endpoint:, api_key:, span_customizers: nil)
         raise State::MissingAPIKeyError, "api_key is required" if api_key.nil? || api_key.empty?
 
+        @span_customizers = (span_customizers || []).dup.freeze
+
         super(endpoint: endpoint, headers: {"Authorization" => "Bearer #{api_key}"})
       end
 
       def export(span_data, timeout: nil)
+        customizers = @span_customizers
+        customize = customizers && !customizers.empty?
+        environment = Internal::Env.detect_environment
+        if customize
+          return FAILURE if @shutdown
+
+          begin
+            span_data = span_data.map do |span|
+              # SpanData includes nested mutable strings, attributes, events,
+              # links and resources. A shallow dup would leak hook mutations to
+              # application data or another exporter. Only SDK-owned data is
+              # marshaled here; no externally supplied serialized bytes are read.
+              # Resources cache an Enumerator, which cannot be marshaled.
+              snapshot = span.dup
+              snapshot.resource = span.resource.attribute_enumerator.to_h
+              snapshot = Marshal.load(Marshal.dump(snapshot))
+              snapshot.resource = OpenTelemetry::SDK::Resources::Resource.create(snapshot.resource)
+              snapshot = SpanOrigin.enrich(snapshot, environment: environment)
+              snapshot.attributes = snapshot.attributes.dup
+              customize_span(snapshot, customizers)
+            end
+            groups = span_data.group_by { |sd| sd.attributes&.[](PARENT_ATTR_KEY) }
+            # Validate serialization for the entire batch before any group can
+            # leave the process. Reuse these bytes rather than encoding twice.
+            encoded_groups = groups.transform_values { |spans| encode(spans) }
+            raise TypeError, "Customized spans must be OTLP serializable" if encoded_groups.any? { |_, bytes| bytes.nil? }
+          rescue => e
+            Log.error("Failed to customize spans for export: #{e.class}")
+            return FAILURE
+          end
+        else
+          span_data = span_data.map { |span| SpanOrigin.enrich(span, environment: environment) }
+          groups = span_data.group_by { |sd| sd.attributes&.[](PARENT_ATTR_KEY) }
+        end
+
         failed = false
-        span_data.group_by { |sd| sd.attributes&.[](PARENT_ATTR_KEY) }.each do |parent_value, spans|
+        groups.each do |parent_value, spans|
           @headers[PARENT_HEADER] = parent_value if parent_value
-          failed = true unless super(spans, timeout: timeout) == SUCCESS
+          result = if customize
+            send_bytes(encoded_groups.fetch(parent_value), timeout: timeout)
+          else
+            super(spans, timeout: timeout)
+          end
+          failed = true unless result == SUCCESS
         ensure
           @headers.delete(PARENT_HEADER)
         end
         failed ? FAILURE : SUCCESS
       end
+
+      private
+
+      def customize_span(span, customizers)
+        trace_id = span.trace_id.dup
+        span_id = span.span_id.dup
+        parent_span_id = span.parent_span_id.dup
+        customizers.each do |customizer|
+          next unless customizer.respond_to?(:on_span_export)
+
+          span = customizer.on_span_export(span)
+          unless span.is_a?(OpenTelemetry::SDK::Trace::SpanData)
+            raise TypeError, "SpanCustomizer#on_span_export must return SpanData"
+          end
+          unless span.trace_id == trace_id && span.span_id == span_id && span.parent_span_id == parent_span_id
+            raise ArgumentError, "SpanCustomizer#on_span_export must preserve trace, span and parent IDs"
+          end
+        end
+        # Added entries cannot have negative dropped counts in OTLP.
+        span.total_recorded_attributes = [span.total_recorded_attributes, span.attributes&.size.to_i].max
+        span.total_recorded_events = [span.total_recorded_events, span.events&.size.to_i].max
+        span.total_recorded_links = [span.total_recorded_links, span.links&.size.to_i].max
+        span
+      end
     end
   end
 end