support span customizers for otel mode

Author: realark

py/src/braintrust/auto.py

@@ -103,12 +103,10 @@ def auto_instrument(
         livekit_agents: Enable LiveKit Agents instrumentation (default: True)
         pipecat: Enable Pipecat AI instrumentation (default: True)
         typesafe: Enable TypeSafe instrumentation (default: True)
-        span_customizers: Ordered synchronous export customizers for all native SDK
-            spans, including manual spans; dataset rows and feedback are excluded.
+        span_customizers: Ordered synchronous export customizers for native SDK and
+            Braintrust OTel spans, including manual spans; dataset rows and feedback are excluded.
             Copies and replaces the global list when provided; None leaves existing
             configuration unchanged. Pass [] to disable.
-            With BRAINTRUST_OTEL_COMPAT enabled, non-empty customizers log an error
-            and leave the existing configuration unchanged; instrumentation continues.
 
     Returns:
         Dict mapping integration name to whether it was successfully instrumented.

py/src/braintrust/otel/__init__.py

@@ -8,6 +8,7 @@
 from urllib.parse import urljoin
 
 from braintrust.env import BraintrustEnv
+from braintrust.span_customizer import _get_span_customizers
 from braintrust.span_origin import SpanOriginEnvironment, detect_environment, merge_span_origin_context
 
 
@@ -303,6 +304,21 @@ def initialize(self) -> None:
 
     def export(self, spans):
         self._ensure_api_key()
+        customizers = _get_span_customizers()
+        if customizers:
+            from .span_customizer import customize_otel_span
+
+            # Processor exports are already enriched after AI filtering. Direct
+            # exporter users receive the same provenance before customization.
+            spans = [
+                customize_otel_span(
+                    span
+                    if isinstance(span, _SpanWithAttributes)
+                    else _with_span_origin_attributes(span, detect_environment()),
+                    customizers,
+                )
+                for span in spans
+            ]
         return super().export(spans)
 
     def force_flush(self, timeout_millis=30000):

py/src/braintrust/otel/span_customizer.py

@@ -0,0 +1,210 @@
+"""Detached OTel export snapshots for the shared span customizer hook."""
+
+from collections.abc import Mapping
+from copy import deepcopy
+from types import SimpleNamespace
+
+from braintrust.span_customizer import _run_span_export_customizers
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.util.instrumentation import InstrumentationScope
+from opentelemetry.trace import SpanContext, SpanKind, StatusCode, TraceFlags
+
+
+try:
+    from opentelemetry.exporter.otlp.proto.common.trace_encoder import encode_spans
+except ImportError:
+    # OpenTelemetry 1.16 predates the shared HTTP/gRPC encoder package.
+    from opentelemetry.exporter.otlp.proto.http.trace_exporter.encoder import _ProtobufEncoder
+
+    encode_spans = _ProtobufEncoder.encode
+
+
+class _ExportSpan(SimpleNamespace):
+    def get_span_context(self):
+        return self.context
+
+
+def _copy_value(value):
+    if isinstance(value, Mapping):
+        return {key: _copy_value(item) for key, item in value.items()}
+    if isinstance(value, (tuple, list)):
+        return [_copy_value(item) for item in value]
+    if value is None or type(value) in (str, bool, int, float, bytes):
+        return value
+    raise TypeError("Invalid OTel snapshot value")
+
+
+def _context_data(context):
+    if context is None:
+        return None
+    return {
+        "trace_id": context.trace_id,
+        "span_id": context.span_id,
+        "is_remote": context.is_remote,
+        "trace_flags": int(context.trace_flags),
+        "trace_state": dict(context.trace_state.items()),
+    }
+
+
+def _snapshot(span):
+    scope = span.instrumentation_scope
+    return {
+        "name": span.name,
+        "context": _context_data(span.context),
+        "parent": _context_data(span.parent),
+        "start_time": span.start_time,
+        "end_time": span.end_time,
+        "kind": span.kind.value,
+        "attributes": _copy_value(span.attributes or {}),
+        "events": [
+            {
+                "name": event.name,
+                "timestamp": event.timestamp,
+                "attributes": _copy_value(event.attributes or {}),
+                "dropped_attributes": getattr(event, "dropped_attributes", getattr(event.attributes, "dropped", 0)),
+            }
+            for event in span.events
+        ],
+        "links": [
+            {
+                "context": _context_data(link.context),
+                "attributes": _copy_value(link.attributes or {}),
+                "dropped_attributes": getattr(link, "dropped_attributes", getattr(link.attributes, "dropped", 0)),
+            }
+            for link in span.links
+        ],
+        "status": {"status_code": span.status.status_code.value, "description": span.status.description},
+        "resource": {"attributes": _copy_value(span.resource.attributes), "schema_url": span.resource.schema_url},
+        "instrumentation_scope": {
+            "name": scope.name,
+            "version": scope.version,
+            "schema_url": scope.schema_url,
+            "attributes": _copy_value(getattr(scope, "attributes", None) or {}),
+        }
+        if scope
+        else None,
+        "dropped_attributes": span.dropped_attributes,
+        "dropped_events": span.dropped_events,
+        "dropped_links": span.dropped_links,
+    }
+
+
+def _attributes(data):
+    if not isinstance(data, dict):
+        raise TypeError("OTel attributes must be a dict")
+    for key, value in data.items():
+        if type(key) is not str:
+            raise TypeError("OTel attribute keys must be strings")
+        key.encode("utf-8")
+        values = value if isinstance(value, (list, tuple)) else (value,)
+        for item in values:
+            if type(item) not in (str, bool, int, float):
+                raise TypeError("Invalid OTel attribute value")
+            if type(item) is int and not -(2**63) <= item < 2**63:
+                raise ValueError("OTel integer attribute out of range")
+            if type(item) is str:
+                item.encode("utf-8")
+        if values and any(type(item) is not type(values[0]) for item in values):
+            raise TypeError("OTel attribute arrays must be homogeneous")
+    return data
+
+
+def _context(data):
+    if data is None:
+        return None
+    state = data.get("trace_state", {})
+    if not isinstance(state, dict) or any(type(k) is not str or type(v) is not str for k, v in state.items()):
+        raise TypeError("Invalid OTel trace state")
+    return SpanContext(
+        trace_id=data["trace_id"],
+        span_id=data["span_id"],
+        is_remote=data.get("is_remote", False),
+        trace_flags=TraceFlags(data.get("trace_flags", 0)),
+        trace_state=state,
+    )
+
+
+def _scope(data):
+    if data is None:
+        return None
+    attributes = _attributes(data.get("attributes", {}))
+    kwargs = {}
+    if hasattr(InstrumentationScope, "attributes"):
+        kwargs["attributes"] = attributes
+    elif attributes:
+        raise ValueError("Scope attributes require a newer OpenTelemetry SDK")
+    return InstrumentationScope(
+        name=data.get("name", ""), version=data.get("version"), schema_url=data.get("schema_url", ""), **kwargs
+    )
+
+
+def _export_span(data):
+    resource = data.get("resource", {})
+    scope = data.get("instrumentation_scope")
+    status = data.get("status", {})
+    return _ExportSpan(
+        name=data.get("name", ""),
+        context=_context(data["context"]),
+        parent=_context(data["parent"]),
+        start_time=data["start_time"],
+        end_time=data["end_time"],
+        kind=SpanKind(data.get("kind", 0)),
+        attributes=_attributes(data.get("attributes", {})),
+        events=[
+            SimpleNamespace(
+                name=event["name"],
+                timestamp=event["timestamp"],
+                attributes=_attributes(event.get("attributes", {})),
+                dropped_attributes=event.get("dropped_attributes", 0),
+            )
+            for event in data.get("events", [])
+        ],
+        links=[
+            SimpleNamespace(
+                context=_context(link["context"]),
+                attributes=_attributes(link.get("attributes", {})),
+                dropped_attributes=link.get("dropped_attributes", 0),
+            )
+            for link in data.get("links", [])
+        ],
+        status=SimpleNamespace(
+            status_code=StatusCode(status.get("status_code", 0)), description=status.get("description")
+        ),
+        resource=Resource(_attributes(resource.get("attributes", {})), schema_url=resource.get("schema_url", "")),
+        instrumentation_scope=_scope(scope),
+        dropped_attributes=data.get("dropped_attributes", 0),
+        dropped_events=data.get("dropped_events", 0),
+        dropped_links=data.get("dropped_links", 0),
+    )
+
+
+def customize_otel_span(span, customizers):
+    """Customize exactly once at the Braintrust exporter, never other consumers."""
+    data = _snapshot(span)
+    protected = {key: deepcopy(data[key]) for key in ("context", "parent", "start_time", "end_time")}
+    routing = {key: data["attributes"][key] for key in ("braintrust.parent",) if key in data["attributes"]}
+
+    def restore(result):
+        result = _copy_value(result)
+        result.update(deepcopy(protected))
+        attributes = result.setdefault("attributes", {})
+        if not isinstance(attributes, dict):
+            raise TypeError("OTel attributes must be a dict")
+        attributes.pop("braintrust.parent", None)
+        attributes.update(routing)
+        # Validate every hook result before the next hook runs. The upstream encoder
+        # silently skips invalid attributes (and logs their values), so validate those
+        # first; protobuf then checks all remaining wire types and numeric bounds.
+        candidate = _export_span(result)
+        encode_spans([candidate]).SerializeToString()
+        return result
+
+    def stripped_record():
+        stripped = deepcopy(protected)
+        for field in ("context", "parent"):
+            if stripped[field] is not None:
+                stripped[field]["trace_state"] = {}
+        stripped["attributes"] = {**routing, "braintrust.context_json": '{"customizer_error":true}'}
+        return stripped
+
+    return _export_span(_run_span_export_customizers(data, customizers, restore, stripped_record))

py/src/braintrust/span_customizer.py

@@ -1,4 +1,4 @@
-"""Synchronous transformations of native export records."""
+"""Synchronous transformations of native and OpenTelemetry export records."""
 
 import inspect
 import logging
@@ -16,22 +16,31 @@
     PARENT_ID_FIELD,
     TRANSACTION_ID_FIELD,
 )
-from .env import BraintrustEnv
 
 
 __all__ = ["SpanCustomizer", "SpanExportData", "set_span_customizers"]
 
 SpanExportData = dict[str, Any]
-"""A native, possibly incremental span export record, not a live span."""
+"""A native incremental record or OTel export snapshot, never a live span."""
 
 
 class SpanCustomizer:
-    """Extensible hooks for all native SDK spans. Omitted hooks are no-ops."""
+    """Extensible hooks for all SDK spans. Omitted hooks are no-ops."""
 
     def on_span_export(self, data: SpanExportData) -> SpanExportData:
         """Return the original or a replacement export record, synchronously.
 
-        Hooks run in registration order after lazy values resolve, before merging,
+        Native records expose input/output/metadata and may be incremental. OTel
+        exports expose detached dictionaries with name, attributes, events, links,
+        status, resource, instrumentation_scope, context/parent, kind, start_time,
+        end_time, and dropped-item counts (see the SDK README for exact shapes).
+        OTel hooks run after filtering and origin enrichment; original context,
+        parent, timing, and braintrust.parent routing are restored after each hook.
+        OTel failures retain only identity/timing/routing and the attribute
+        braintrust.context_json='{"customizer_error":true}', with empty resource
+        and scope metadata. Both pipelines share one failure diagnostic budget.
+
+        Native hooks run in registration order after lazy values resolve, before merging,
         attachment processing, masking, and serialization. Identity, routing, and
         merge protocol fields are restored after each hook. Exceptions and invalid
         synchronous returns stop the hook chain and strip the outgoing record.
@@ -114,22 +123,16 @@ def set_span_customizers(customizers: Sequence[SpanCustomizer] | None) -> None:
     """Replace the process-wide ordered customizer list with a snapshot.
 
     Configure before logging spans; pass None or an empty sequence to disable.
-    The sequence is copied, but customizer instances are not. Each span record
-    uses the configuration active when it is logged, and retries reuse its
-    transformed or stripped data. There is no environment-variable registration.
+    The sequence is copied, but customizer instances are not. Native records use
+    the configuration active when logged; OTel uses the export-time configuration.
+    Transport retries reuse transformed or stripped data. There is no
+    environment-variable registration.
 
     Raises TypeError for classes, objects without a callable ``on_span_export``,
     and async hooks. Invalid synchronous returns are rejected at export time.
-    Non-empty customizers with BRAINTRUST_OTEL_COMPAT enabled log one error per
-    registration attempt and leave the existing configuration unchanged.
     """
     global _span_customizers
     snapshot = tuple(customizers) if customizers is not None else ()
-    if snapshot and BraintrustEnv.OTEL_COMPAT.get(False):
-        logging.getLogger("braintrust").error(
-            "Span customizers are not yet supported with OTel compat mode (BRAINTRUST_OTEL_COMPAT)"
-        )
-        return
     for customizer in snapshot:
         _validate_customizer(customizer)
     _span_customizers = snapshot
@@ -188,10 +191,36 @@ def _customize_span_export(
     preserved = _snapshot_fields(data, ("created",))
     timing = _snapshot_fields(data.get("metrics"), ("start", "end"))
     span_type = _snapshot_fields(data.get("span_attributes"), ("type",))
+
+    def stripped_record():
+        stripped = _StrippedSpanExportData(protected)
+        stripped.update(preserved)
+        if timing:
+            stripped["metrics"] = timing
+        if span_type:
+            stripped["span_attributes"] = span_type
+        stripped["context"] = {"customizer_error": True}
+        return stripped
+
     try:
         # Protocol containers may alias span state; don't expose those to hooks.
         data = _restore_protocol_fields(data, protected)
-        for customizer in customizers:
+    except (KeyboardInterrupt, SystemExit):
+        raise
+    except BaseException:
+        stripped = stripped_record()
+        _log_customizer_failure()
+        return stripped
+
+    return _run_span_export_customizers(
+        data, customizers, lambda result: _restore_protocol_fields(result, protected), stripped_record
+    )
+
+
+def _run_span_export_customizers(data, customizers, restore, stripped_record):
+    """Shared synchronous hook/failure policy; adapters own their export shape."""
+    for customizer in customizers:
+        try:
             hook = getattr(customizer, "on_span_export", None)
             if hook is None:
                 continue
@@ -204,21 +233,14 @@ def _customize_span_export(
             if not isinstance(result, dict):
                 raise TypeError("Span customizers must return a dict")
             # Always a fresh plain dict, even for dict subclasses.
-            data = _restore_protocol_fields(result, protected)
-    except (KeyboardInterrupt, SystemExit):
-        raise
-    except BaseException:
-        # Memoize this stripped result in the caller's LazyValue, including
-        # GeneratorExit and CancelledError, without rejecting other records.
-        stripped = _StrippedSpanExportData(protected)
-        stripped.update(preserved)
-        if timing:
-            stripped["metrics"] = timing
-        if span_type:
-            stripped["span_attributes"] = span_type
-        stripped["context"] = {"customizer_error": True}
-        _log_customizer_failure()
-        return stripped
+            data = restore(dict(result))
+        except (KeyboardInterrupt, SystemExit):
+            raise
+        except BaseException:
+            # No exception details are safe, including BaseException subclasses.
+            stripped = stripped_record()
+            _log_customizer_failure()
+            return stripped
     return data