From 39c575aafdabec451a710aeab6e40174bb089252 Mon Sep 17 00:00:00 2001 From: Andrew Kent Date: Mon, 21 Sep 2026 02:13:13 -0600 Subject: [PATCH] feat: add span customizers (SDK-316) --- README.md | 38 ++- lib/braintrust.rb | 5 +- lib/braintrust/config.rb | 11 +- lib/braintrust/span_customizer.rb | 18 + lib/braintrust/state.rb | 6 +- lib/braintrust/trace.rb | 8 +- lib/braintrust/trace/span_customizers.rb | 73 ++++ lib/braintrust/trace/span_export_data.rb | 35 ++ lib/braintrust/trace/span_exporter.rb | 48 ++- lib/braintrust/trace/span_origin.rb | 8 +- test/braintrust/trace/span_customizer_test.rb | 320 ++++++++++++++++++ test/braintrust/trace/span_exporter_test.rb | 8 +- test/support/in_memory_exporter.rb | 4 +- 13 files changed, 557 insertions(+), 25 deletions(-) create mode 100644 lib/braintrust/span_customizer.rb create mode 100644 lib/braintrust/trace/span_customizers.rb create mode 100644 lib/braintrust/trace/span_export_data.rb create mode 100644 test/braintrust/trace/span_customizer_test.rb diff --git a/README.md b/README.md index e335c149..06f66515 100644 --- a/README.md +++ b/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 a view of its predecessor's result. The argument is a `Braintrust::Trace::SpanExportData` facade over completed OpenTelemetry span data, 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. + +Mutate and return the facade directly, or return a replacement `OpenTelemetry::SDK::Trace::SpanData`; replacement is not an implicit merge. The facade exposes span field readers and writers, except for identity writers, and `to_span_data` provides the underlying OpenTelemetry representation when constructing a replacement. The exporter provides a writable attributes hash, so assigning or deleting attribute entries does not change another exporter's hash. There is no deep copy: nested strings, arrays, events, links, and resources may still be shared with application code or other exporters. Prefer assigning replacement attribute values over mutating existing values in place. 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. The facade exposes `trace_id`, `span_id`, and `parent_span_id` as read-only, frozen strings and does not expose indexed assignment (`[]=`). Preserve these IDs in replacements too; the exporter checks them 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. Mutations made before a failure are not rolled back. + +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 same span data through `export` again invokes the hooks again on its current state, including prior mutations, 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: @@ -567,7 +603,7 @@ The dev server requires the `rack` gem and a Rack-compatible web server. | [Passenger](https://www.phusionpassenger.com/) | 6.x | | | [WEBrick](https://github.com/ruby/webrick) | Not supported | Does not support server-sent events. | -See examples: [server/eval.ru](./examples/server/eval.ru), +See examples: [server/eval.ru](./examples/server/eval.ru), ## Documentation diff --git a/lib/braintrust.rb b/lib/braintrust.rb index 496fc226..fa0ac05d 100644 --- a/lib/braintrust.rb +++ b/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,6 +43,7 @@ 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, nil] Custom span filter functions + # @param span_customizers [Array, 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 @@ -49,7 +51,7 @@ class Error < StandardError; end # - 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 ) diff --git a/lib/braintrust/config.rb b/lib/braintrust/config.rb index 1aa2a1a0..b28168af 100644 --- a/lib/braintrust/config.rb +++ b/lib/braintrust/config.rb @@ -8,10 +8,10 @@ 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 @@ -19,6 +19,7 @@ def initialize(api_key: nil, org_name: nil, default_project: nil, app_url: nil, @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, nil] Custom span filter functions + # @param span_customizers [Array, 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 diff --git a/lib/braintrust/span_customizer.rb b/lib/braintrust/span_customizer.rb new file mode 100644 index 00000000..0ee73d0a --- /dev/null +++ b/lib/braintrust/span_customizer.rb @@ -0,0 +1,18 @@ +# 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 completed span data through a mutable view with read-only IDs. + # Return this view or a replacement SpanData, never nil. Nested objects are + # shared; prefer replacing attribute values over mutating them in place. + # Exceptions fail the entire batch but do not roll back mutations. + # + # @param span [Braintrust::Trace::SpanExportData] + # @return [Braintrust::Trace::SpanExportData, OpenTelemetry::SDK::Trace::SpanData] + def on_span_export(span) + span + end + end +end diff --git a/lib/braintrust/state.rb b/lib/braintrust/state.rb index f95e5fd9..423ce2ab 100644 --- a/lib/braintrust/state.rb +++ b/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, nil] Custom span filter functions + # @param span_customizers [Array, 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, diff --git a/lib/braintrust/trace.rb b/lib/braintrust/trace.rb index 2c6563fe..b0529906 100644 --- a/lib/braintrust/trace.rb +++ b/lib/braintrust/trace.rb @@ -88,10 +88,16 @@ def self.enable(tracer_provider, state: nil, exporter: nil, config: nil) # Get config from state if available config ||= state.respond_to?(:config) ? state.config : nil + # Customizers only run in SpanExporter; refuse to silently skip them (e.g. redaction). + if exporter && config&.span_customizers&.any? + raise ArgumentError, "span_customizers are not supported with a custom exporter" + end + # 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 diff --git a/lib/braintrust/trace/span_customizers.rb b/lib/braintrust/trace/span_customizers.rb new file mode 100644 index 00000000..aaef9f39 --- /dev/null +++ b/lib/braintrust/trace/span_customizers.rb @@ -0,0 +1,73 @@ +# frozen_string_literal: true + +require_relative "span_export_data" + +module Braintrust + module Trace + # Ordered export-time transformations, independent of origin and transport. + class SpanCustomizers + def initialize(customizers = nil) + @customizers = (customizers || []).dup.freeze + end + + def empty? + @customizers.empty? + end + + # Transform the whole batch before the exporter groups or serializes it. + # Errors propagate to the exporter so it can fail the batch without sending. + def customize(span_data) + return span_data if empty? + + # Hooks may call instrumented clients. Suppress tracing so those spans + # cannot re-enter the exporter and customize themselves forever. + OpenTelemetry::Common::Utilities.untraced do + span_data.map { |span| customize_span(span) } + end + end + + private + + def customize_span(span) + # Snapshot IDs so in-place mutation through to_span_data cannot slip past validation. + trace_id = span.trace_id.dup.freeze + span_id = span.span_id.dup.freeze + parent_span_id = span.parent_span_id.dup.freeze + # Preserve drops caused by SDK limits; entries a customizer deletes are not drops. + dropped_attributes = dropped(span.total_recorded_attributes, span.attributes) + dropped_events = dropped(span.total_recorded_events, span.events) + dropped_links = dropped(span.total_recorded_links, span.links) + view = nil + @customizers.each do |customizer| + next unless customizer.respond_to?(:on_span_export) + view ||= writable_view(span) + + result = customizer.on_span_export(view) + span = result.is_a?(SpanExportData) ? result.to_span_data : result + unless span.is_a?(OpenTelemetry::SDK::Trace::SpanData) + raise TypeError, "SpanCustomizer#on_span_export must return SpanExportData or 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 + view = result.is_a?(SpanExportData) ? result : nil + end + span.total_recorded_attributes = span.attributes&.size.to_i + dropped_attributes + span.total_recorded_events = span.events&.size.to_i + dropped_events + span.total_recorded_links = span.links&.size.to_i + dropped_links + span + end + + # OTel and replacement spans may carry frozen attributes. Only the hash + # needs to be writable; values and other nested objects remain shared. + def writable_view(span) + span.attributes = span.attributes&.dup || {} + SpanExportData.new(span) + end + + def dropped(total, entries) + [total.to_i - entries&.size.to_i, 0].max + end + end + end +end diff --git a/lib/braintrust/trace/span_export_data.rb b/lib/braintrust/trace/span_export_data.rb new file mode 100644 index 00000000..3af86860 --- /dev/null +++ b/lib/braintrust/trace/span_export_data.rb @@ -0,0 +1,35 @@ +# frozen_string_literal: true + +require "forwardable" +require "opentelemetry/sdk" + +module Braintrust + module Trace + # Mutable view of completed span data with read-only identity fields. + # Attribute values and other nested objects are shared, not deep-copied. + class SpanExportData + extend Forwardable + + ID_FIELDS = [:trace_id, :span_id, :parent_span_id].freeze + MUTABLE_FIELDS = (OpenTelemetry::SDK::Trace::SpanData.members - ID_FIELDS).freeze + private_constant :ID_FIELDS, :MUTABLE_FIELDS + + attr_reader(*ID_FIELDS) + def_delegators :@span_data, *MUTABLE_FIELDS, *MUTABLE_FIELDS.map { |field| :"#{field}=" } + def_delegators :@span_data, :hex_trace_id, :hex_span_id, :hex_parent_span_id, :instrumentation_library + + def initialize(span_data) + @span_data = span_data + @trace_id = span_data.trace_id.dup.freeze + @span_id = span_data.span_id.dup.freeze + @parent_span_id = span_data.parent_span_id.dup.freeze + end + + # Access the underlying OpenTelemetry representation, e.g. to build a replacement. + # Replacement IDs are validated by the exporter after each hook. + def to_span_data + @span_data + end + end + end +end diff --git a/lib/braintrust/trace/span_exporter.rb b/lib/braintrust/trace/span_exporter.rb index 32ab2705..e91a1ec8 100644 --- a/lib/braintrust/trace/span_exporter.rb +++ b/lib/braintrust/trace/span_exporter.rb @@ -3,41 +3,77 @@ require "opentelemetry/exporter/otlp" require_relative "../state" require_relative "span_origin" +require_relative "span_customizers" +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 through a mutable view with read-only IDs # - 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 = SpanCustomizers.new(span_customizers) + super(endpoint: endpoint, headers: {"Authorization" => "Bearer #{api_key}"}) end def export(span_data, timeout: nil) + return FAILURE if @shutdown && !@span_customizers.empty? + + groups = prepare_groups(span_data) + return FAILURE unless groups + failed = false - span_data.group_by { |sd| sd.attributes&.[](PARENT_ATTR_KEY) }.each do |parent_value, spans| + groups.each do |parent_value, payload| @headers[PARENT_HEADER] = parent_value if parent_value - failed = true unless super(spans, timeout: timeout) == SUCCESS + result = if @span_customizers.empty? + super(payload, timeout: timeout) + else + send_bytes(payload, timeout: timeout) + end + failed = true unless result == SUCCESS ensure @headers.delete(PARENT_HEADER) end failed ? FAILURE : SUCCESS end + + private + + def prepare_groups(span_data) + # Compose preparation explicitly so origin, hook, and serialization + # failures stay inside the same fail-closed boundary. + span_data = SpanOrigin.enrich_batch(span_data) + span_data = @span_customizers.customize(span_data) + groups = span_data.group_by { |span| span.attributes&.[](PARENT_ATTR_KEY) } + return groups if @span_customizers.empty? + + # Validate serialization for the entire customized batch before sending + # any destination group. Transport retries reuse these encoded bytes. + groups.transform_values do |spans| + encode(spans) || raise(TypeError, "Customized spans must be OTLP serializable") + end + # Hook errors must not kill the BatchSpanProcessor worker thread. + rescue StandardError, ScriptError, SystemStackError => e + raise if @span_customizers.empty? + + Log.error("Failed to prepare customized spans for export: #{e.class}") + nil + end end end end diff --git a/lib/braintrust/trace/span_origin.rb b/lib/braintrust/trace/span_origin.rb index 177ab2fa..fcb479c9 100644 --- a/lib/braintrust/trace/span_origin.rb +++ b/lib/braintrust/trace/span_origin.rb @@ -24,11 +24,15 @@ module SpanOrigin # @param span_data [Array] # @return [Integer] export result from the wrapped exporter def export(span_data, timeout: nil) + super(SpanOrigin.enrich_batch(span_data), timeout: timeout) + end + + # Shared transform for exporters that compose preparation explicitly. + def self.enrich_batch(span_data) # Environment is process-global and stable; read it once per batch # rather than once per span. It is cheap (ENV reads only). environment = Internal::Env.detect_environment - enriched = span_data.map { |sd| SpanOrigin.enrich(sd, environment: environment) } - super(enriched, timeout: timeout) + span_data.map { |sd| enrich(sd, environment: environment) } end # Enrich a single SpanData with span origin provenance. diff --git a/test/braintrust/trace/span_customizer_test.rb b/test/braintrust/trace/span_customizer_test.rb new file mode 100644 index 00000000..264248db --- /dev/null +++ b/test/braintrust/trace/span_customizer_test.rb @@ -0,0 +1,320 @@ +# frozen_string_literal: true + +require "test_helper" + +class Braintrust::Trace::SpanCustomizerTest < Minitest::Test + SUCCESS = OpenTelemetry::SDK::Trace::Export::SUCCESS + FAILURE = OpenTelemetry::SDK::Trace::Export::FAILURE + ENDPOINT = "https://api.ruby-sdk-fixture.com/otel/v1/traces" + + def setup + @requests = [] + @providers = [] + @exporters = [] + stub_request(:post, ENDPOINT).to_return do |request| + body = (request.headers["Content-Encoding"] == "gzip") ? Zlib.gunzip(request.body) : request.body + decoded = Opentelemetry::Proto::Collector::Trace::V1::ExportTraceServiceRequest.decode(body) + @requests << {headers: request.headers, resource_spans: decoded.resource_spans.size, spans: decoded.resource_spans.flat_map { |resource| resource.scope_spans.flat_map { |scope| scope.spans.to_a } }} + {status: 200, body: ""} + end + end + + def teardown + @providers.each(&:shutdown) + @exporters.each(&:shutdown) + end + + def test_order_replacement_routing_and_snapshot_registration_through_init + first = customizer do |span| + origin = JSON.parse(span.attributes.fetch("braintrust.context_json")) + replacement = span.to_span_data.dup + replacement.name = "#{origin.fetch("span_origin").fetch("name")}:redacted" + replacement.attributes = {"braintrust.parent" => "project_name:redacted", "braintrust.input_json" => '"[redacted]"'} + replacement + end + second = customizer do |span| + span.name += ":second" + span.attributes["customized"] = true + span + end + customizers = [Object.new, Braintrust::SpanCustomizer.new, first, second] + provider = make_provider + state = Braintrust.init( + api_key: "test-api-key", default_project: "original", + blocking_login: true, set_global: false, auto_instrument: false, + tracer_provider: provider, span_customizers: customizers + ) + customizers.clear + assert_raises(FrozenError) { state.config.span_customizers.clear } + provider.tracer("manual").start_span("private", attributes: {"braintrust.input_json" => '"secret"'}).finish + provider.force_flush + + assert_equal 1, @requests.size + request = @requests.fetch(0) + assert_equal "project_name:redacted", request[:headers]["X-Bt-Parent"] + assert_equal "Bearer test-api-key", request[:headers]["Authorization"] + span = request[:spans].fetch(0) + assert_equal "braintrust.sdk.ruby:redacted:second", span.name + assert_equal '"[redacted]"', attributes(span).fetch("braintrust.input_json").string_value + assert attributes(span).fetch("customized").bool_value + refute attributes(span).key?("braintrust.context_json"), "replacement must not implicitly merge removed fields" + end + + def test_attribute_replacement_does_not_mutate_application_or_other_exporters + provider = make_provider + memory = OpenTelemetry::SDK::Trace::Export::InMemorySpanExporter.new + provider.add_span_processor(OpenTelemetry::SDK::Trace::Export::SimpleSpanProcessor.new(memory)) + input = +"private" + source_span = provider.tracer("app").start_span("original", attributes: {"input" => [input]}) + source_span.add_event("event", attributes: {"value" => +"private"}) + source_span.finish + source = memory.finished_spans.fetch(0) + hook = customizer do |span| + span.name = "redacted" + span.attributes["input"] = ["redacted"] + span.attributes["added"] = true + span + end + + assert_equal SUCCESS, make_exporter([hook]).export([source_span.to_span_data]) + exported = @requests.fetch(0)[:spans].fetch(0) + assert_equal "redacted", exported.name + assert_equal "redacted", attributes(exported).fetch("input").array_value.values.fetch(0).string_value + assert attributes(exported).fetch("added").bool_value + assert_equal "private", input + assert_equal "original", source.name + assert_equal ["private"], source.attributes.fetch("input") + assert_equal "private", source.events.fetch(0).attributes.fetch("value") + refute source.attributes.key?("braintrust.context_json") + assert_equal source.trace_id, exported.trace_id + assert_equal source.span_id, exported.span_id + end + + def test_exception_after_mutation_sends_none_of_batch + spans = two_destinations + hook = customizer do |span| + if span.name == "second" + span.attributes.fetch("input").replace("changed") + raise "redaction failed" + end + span + end + + assert_equal FAILURE, make_exporter([hook]).export(spans) + assert_empty @requests + assert_equal "changed", spans.last.attributes.fetch("input") + end + + def test_nil_and_non_span_returns_fail_whole_batch + [nil, {}].each do |invalid| + hook = customizer { |span| (span.name == "second") ? invalid : span } + assert_equal FAILURE, make_exporter([hook]).export(two_destinations) + assert_empty @requests + end + end + + def test_unserializable_replacement_fails_before_first_destination_is_sent + hook = customizer do |span| + span.start_timestamp = "invalid" if span.name == "second" + span + end + + assert_equal FAILURE, make_exporter([hook]).export(two_destinations) + assert_empty @requests + end + + def test_identity_is_read_only_while_attributes_are_writable + source = make_span("original") + test = self + hook = customizer do |span| + [:trace_id, :span_id, :parent_span_id].each do |field| + test.assert_raises(NoMethodError) { span.public_send("#{field}=", "changed") } + test.assert_raises(FrozenError) { span.public_send(field).replace("changed") } + end + test.assert_raises(NoMethodError) { span[:trace_id] = "changed" } + span.attributes["input"] = "redacted" + span + end + + assert_equal SUCCESS, make_exporter([hook]).export([source]) + exported = @requests.fetch(0)[:spans].fetch(0) + assert_equal source.trace_id, exported.trace_id + assert_equal source.span_id, exported.span_id + assert_equal "redacted", attributes(exported).fetch("input").string_value + end + + def test_each_hook_must_preserve_all_ids_even_if_later_hook_would_restore_them + [:trace_id, :span_id, :parent_span_id].each do |field| + spans = two_destinations + original = spans.last.public_send(field).dup + later_called = false + mutate = customizer do |span| + if span.name == "second" + replacement = span.to_span_data.dup + replacement.public_send("#{field}=", "x" * original.bytesize) + replacement + else + span + end + end + restore = customizer do |span| + if span.name == "second" + later_called = true + span.public_send("#{field}=", original) + end + span + end + + assert_equal FAILURE, make_exporter([mutate, restore]).export(spans) + assert_empty @requests + refute later_called + assert_equal original, spans.last.public_send(field) + end + end + + def test_parent_identity_survives_successful_replacement + provider = make_provider + tracer = provider.tracer("app") + parent = tracer.start_span("parent") + context = OpenTelemetry::Trace.context_with_span(parent) + child = tracer.start_span("child", with_parent: context) + child.finish + parent.finish + source = child.to_span_data + hook = customizer do |span| + replacement = span.to_span_data.dup + replacement.name = "replacement" + replacement + end + + assert_equal SUCCESS, make_exporter([hook]).export([source]) + exported = @requests.fetch(0)[:spans].fetch(0) + assert_equal source.trace_id, exported.trace_id + assert_equal source.span_id, exported.span_id + assert_equal parent.context.span_id, exported.parent_span_id + end + + def test_resubmission_reruns_customizers_on_current_span_data + calls = 0 + hook = customizer do |span| + calls += 1 + span.name += ":customized" + span + end + customizers = [hook] + exporter = make_exporter(customizers) + customizers.clear + source = make_span("original") + + 2.times { assert_equal SUCCESS, exporter.export([source]) } + assert_equal 2, calls + assert_equal ["original:customized", "original:customized:customized"], @requests.map { |request| request[:spans].fetch(0).name } + assert_equal "original:customized:customized", source.name + end + + def test_batch_from_one_resource_shares_one_resource_spans_entry + tracer = make_provider.tracer("app") + spans = 3.times.map do |i| + tracer.start_span("span#{i}", attributes: {"braintrust.parent" => "project_name:original"}).tap(&:finish).to_span_data + end + + assert_equal SUCCESS, make_exporter([customizer { |span| span }]).export(spans) + assert_equal 1, @requests.fetch(0)[:resource_spans] + assert_equal 3, @requests.fetch(0)[:spans].size + end + + def test_non_standard_errors_from_hooks_fail_closed + [NotImplementedError, SystemStackError].each do |error| + hook = customizer { |_span| raise error } + assert_equal FAILURE, make_exporter([hook]).export(two_destinations) + assert_empty @requests + end + end + + def test_spans_created_by_hooks_are_not_traced + memory = OpenTelemetry::SDK::Trace::Export::InMemorySpanExporter.new + hook_provider = make_provider + hook_provider.add_span_processor(OpenTelemetry::SDK::Trace::Export::SimpleSpanProcessor.new(memory)) + hook = customizer do |span| + hook_provider.tracer("pii-detector").in_span("detect") {} + span + end + + assert_equal SUCCESS, make_exporter([hook]).export([make_span("original")]) + assert_empty memory.finished_spans + end + + def test_later_hook_can_write_attributes_of_replacement_with_frozen_hash + replace = customizer do |span| + replacement = span.to_span_data.dup + replacement.attributes = span.attributes.dup.freeze + replacement + end + write = customizer do |span| + span.attributes["customized"] = true + span + end + + assert_equal SUCCESS, make_exporter([replace, write]).export([make_span("original")]) + assert attributes(@requests.fetch(0)[:spans].fetch(0)).fetch("customized").bool_value + end + + def test_deleted_entries_are_not_reported_as_dropped + tracer = make_provider(span_limits: OpenTelemetry::SDK::Trace::SpanLimits.new(event_count_limit: 2)).tracer("app") + span = tracer.start_span("original", attributes: {"braintrust.parent" => "project_name:original", "secret" => "x"}) + 3.times { |i| span.add_event("event#{i}") } + span.finish + hook = customizer do |view| + view.attributes.delete("secret") + view.events = view.events.first(1) + view + end + + assert_equal SUCCESS, make_exporter([hook]).export([span.to_span_data]) + exported = @requests.fetch(0)[:spans].fetch(0) + assert_equal 0, exported.dropped_attributes_count + assert_equal 1, exported.dropped_events_count, "limit drops are preserved, deletions are not counted" + end + + def test_init_rejects_customizers_with_exporter_override + error = assert_raises(ArgumentError) do + Braintrust.init( + api_key: "test-api-key", default_project: "original", + blocking_login: true, set_global: false, auto_instrument: false, + tracer_provider: make_provider, span_customizers: [customizer { |span| span }], + exporter: OpenTelemetry::SDK::Trace::Export::InMemorySpanExporter.new + ) + end + assert_match(/span_customizers/, error.message) + end + + private + + def customizer(&block) + Object.new.tap { |object| object.define_singleton_method(:on_span_export, &block) } + end + + def make_provider(**options) + OpenTelemetry::SDK::Trace::TracerProvider.new(**options).tap { |provider| @providers << provider } + end + + def make_exporter(customizers) + Braintrust::Trace::SpanExporter.new(endpoint: ENDPOINT, api_key: "test-key", span_customizers: customizers).tap do |exporter| + @exporters << exporter + end + end + + def make_span(name, parent: "project_name:original") + span = make_provider.tracer("app").start_span(name, attributes: {"braintrust.parent" => parent, "input" => +"private"}) + span.finish + span.to_span_data + end + + def two_destinations + [make_span("first", parent: "project_name:first"), make_span("second", parent: "project_name:second")] + end + + def attributes(span) + span.attributes.to_h { |attribute| [attribute.key, attribute.value] } + end +end diff --git a/test/braintrust/trace/span_exporter_test.rb b/test/braintrust/trace/span_exporter_test.rb index c1f69eee..ca5fbc70 100644 --- a/test/braintrust/trace/span_exporter_test.rb +++ b/test/braintrust/trace/span_exporter_test.rb @@ -31,9 +31,7 @@ class RecordingExporter < Braintrust::Trace::SpanExporter def initialize(api_key: "test-key") @calls = [] - # Initialize headers directly — skip super to avoid HTTP setup - @headers = {"Authorization" => "Bearer #{api_key}"} - @shutdown = false + super(endpoint: "https://api.example.test/otel/v1/traces", api_key: api_key) end private @@ -102,11 +100,9 @@ def test_handles_nil_parent end def test_requires_api_key - error = assert_raises(Braintrust::State::MissingAPIKeyError) do + assert_raises(Braintrust::State::MissingAPIKeyError) do Braintrust::Trace::SpanExporter.new(endpoint: "https://api.example.test/otel/v1/traces", api_key: nil) end - - assert_match(/api_key is required/, error.message) end def test_mixed_nil_and_non_nil_parents diff --git a/test/support/in_memory_exporter.rb b/test/support/in_memory_exporter.rb index bd6617d7..ea6bbd1c 100644 --- a/test/support/in_memory_exporter.rb +++ b/test/support/in_memory_exporter.rb @@ -8,8 +8,8 @@ module Support # behaviors as the production SpanExporter - currently span origin decoration # (SpanOrigin), prepended below. # - # Both this and SpanExporter prepend the *same* SpanOrigin module, so the - # behavior under test cannot drift between the production and test exporters. + # Both this prepend and the production SpanExporter call SpanOrigin.enrich_batch, + # so origin enrichment is shared between the production and test exporters. # Tests can therefore assert on origin-decorated SpanData without any network # calls or a real OTLP exporter. class InMemoryExporter < OpenTelemetry::SDK::Trace::Export::InMemorySpanExporter