span customizer enhancements

Author: realark

.agents/skills/sdk-integrations/SKILL.md

@@ -190,7 +190,7 @@ def start_span(*args, **kwargs):
     return _bt_start_span(*args, **kwargs)
 ```
 
-Module-level `start_span(...)` calls flow through automatically. Calls that go through a `Logger` / `Experiment` / parent-`Span` instance (`logger.start_span(...)`, `parent.start_span(...)`) need `internal={"instrumentation": _INSTRUMENTATION}` explicitly. Tests assert `span["context"]["span_origin"]["instrumentation"]["name"] == "<provider>-auto"`; `SpanImpl._instrumentation` exposes the resolved value.
+Module-level `start_span(...)` calls flow through automatically. Calls that go through a `Logger` / `Experiment` / parent-`Span` instance (`logger.start_span(...)`, `parent.start_span(...)`) need `internal={"instrumentation": _INSTRUMENTATION}` explicitly. Tests assert the exported provenance: `span["context"]["span_origin"]["instrumentation"]["name"] == "<provider>-auto"`.
 
 ## Metrics
 

py/src/braintrust/auto.py

@@ -103,9 +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 instrumentation
-            spans. Copies and replaces the global list when provided; None leaves
-            existing configuration unchanged. Pass [] to disable.
+        span_customizers: Ordered synchronous export customizers for all native SDK
+            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.
 
     Returns:
         Dict mapping integration name to whether it was successfully instrumented.

py/src/braintrust/logger.py

@@ -999,9 +999,12 @@ def utf8_byte_length(value: str) -> int:
     return len(value.encode("utf-8"))
 
 
+_LazyLogRecord = LazyValue[dict[str, Any]] | LazyValue[dict[str, Any] | None]
+
+
 class _BackgroundLogger(ABC):
     @abstractmethod
-    def log(self, *args: LazyValue[dict[str, Any]]) -> None:
+    def log(self, *args: _LazyLogRecord) -> None:
         pass
 
     @abstractmethod
@@ -1019,7 +1022,7 @@ def __init__(self):
     def enforce_queue_size_limit(self, enforce: bool) -> None:
         pass
 
-    def log(self, *args: LazyValue[dict[str, Any]]) -> None:
+    def log(self, *args: _LazyLogRecord) -> None:
         with self.lock:
             self.logs.extend(args)
 
@@ -1033,8 +1036,8 @@ def flush(self, batch_size: int | None = None):
             if not self.logs:
                 return
 
-            # Unwrap lazy values and extract attachments
-            logs = [l.get() for l in self.logs]
+            # Filter dropped records before touching attachments.
+            logs = [record for item in self.logs if (record := item.get()) is not None]
 
             # Extract attachments from all logs
             attachments: list[BaseAttachment] = []
@@ -1046,7 +1049,7 @@ def flush(self, batch_size: int | None = None):
 
     def pop(self):
         with self.lock:
-            logs = [l.get() for l in self.logs]  # unwrap the LazyValues
+            logs = [record for item in self.logs if (record := item.get()) is not None]
             self.logs = []
 
             if not logs:
@@ -1057,7 +1060,11 @@ def pop(self):
             batch = merge_row_batch(logs)
 
             if self._export_customizers:
-                batch = [_customize_span_export(item, self._export_customizers) for item in batch]
+                batch = [
+                    record
+                    for item in batch
+                    if (record := _customize_span_export(item, self._export_customizers)) is not None
+                ]
 
             return batch
 
@@ -1098,7 +1105,7 @@ def __init__(self, api_conn: LazyValue[HTTPConnection]):
         self.started = False
 
         self.logger = logging.getLogger("braintrust")
-        self.queue: "LogQueue[LazyValue[dict[str, Any]]]" = LogQueue(maxsize=self.queue_maxsize)
+        self.queue: "LogQueue[_LazyLogRecord]" = LogQueue(maxsize=self.queue_maxsize)
 
         # Counter for tracking overflow uploads (useful for testing)
         self._overflow_upload_count = 0
@@ -1113,7 +1120,7 @@ def enforce_queue_size_limit(self, enforce: bool) -> None:
         """
         self.queue.enforce_queue_size_limit(enforce)
 
-    def log(self, *args: LazyValue[dict[str, Any]]) -> None:
+    def log(self, *args: _LazyLogRecord) -> None:
         self._start()
         dropped_items = []
         for event in args:
@@ -1254,16 +1261,20 @@ def flush(self, batch_size: int | None = None):
                 )
 
     def _unwrap_lazy_values(
-        self, wrapped_items: Sequence[LazyValue[dict[str, Any]]]
+        self, wrapped_items: Sequence[_LazyLogRecord]
     ) -> tuple[list[dict[str, Any]], list["BaseAttachment"]]:
         for i in range(self.num_tries):
             try:
-                unwrapped_items = [item.get() for item in wrapped_items]
+                unwrapped_items = [record for item in wrapped_items if (record := item.get()) is not None]
                 merged_items = merge_row_batch(unwrapped_items)
 
-                # Logger-local hooks run after instrumentation hooks and merging.
+                # Logger-local hooks run after span hooks and merging.
                 if self._export_customizers:
-                    merged_items = [_customize_span_export(item, self._export_customizers) for item in merged_items]
+                    merged_items = [
+                        record
+                        for item in merged_items
+                        if (record := _customize_span_export(item, self._export_customizers)) is not None
+                    ]
 
                 attachments: list["BaseAttachment"] = []
                 for item in merged_items:
@@ -2489,8 +2500,8 @@ def set_masking_function(masking_function: Callable[[Any], Any] | None) -> None:
     """
     Set a global masking function that will be applied to all logged data before sending to Braintrust.
     The masking function will be applied after records are merged but before they are sent to the backend.
-    Internally, masking is a logger-local export customizer that runs after instrumentation
-    customizers and also covers manually logged records.
+    Internally, masking is a logger-local export customizer that runs after span
+    customizers and also covers non-span records such as datasets and feedback.
 
     :param masking_function: A function that takes a JSON-serializable object and returns a masked version.
                            Set to None to disable masking.
@@ -3983,6 +3994,7 @@ def _update_span_impl(
     )
 
     update_event = bt_safe_deep_copy(update_event)
+    customizers = _get_span_customizers()
 
     def parent_ids():
         exporter = _get_exporter()
@@ -3991,15 +4003,16 @@ def parent_ids():
             object_id=parent_object_id.get(),
         ).object_id_fields()
 
-    def compute_record():
-        return dict(
+    def compute_record() -> dict[str, Any] | None:
+        record = dict(
             id=id,
             **update_event,
             **parent_ids(),
             **{
                 IS_MERGE_FIELD: True,
             },
         )
+        return _customize_span_export(record, customizers) if customizers else record
 
     _state.global_bg_logger().log(LazyValue(compute_record, use_mutex=False))
 
@@ -4815,6 +4828,7 @@ def __init__(
             created=datetime.datetime.now(datetime.timezone.utc).isoformat(),
         )
         internal = internal or {}
+        # Integration consumers use this provenance independently of export hooks.
         self._instrumentation = internal.get("instrumentation") or "braintrust-python-logger"
         internal_data["context"] = merge_span_origin_context(
             caller_location or {},
@@ -4903,7 +4917,7 @@ def log_internal(self, event: dict[str, Any] | None = None, internal_data: dict[
 
         # Snapshot at log time so the span cache and the export agree on whether
         # (and how) this record is customized.
-        customizers = _get_span_customizers() if self._instrumentation != "braintrust-python-logger" else ()
+        customizers = _get_span_customizers()
         pending_cache_key = (
             object()
             if customizers
@@ -4933,11 +4947,11 @@ def write_span_cache(record: dict[str, Any]) -> None:
             self.state.span_cache.queue_write(self.root_span_id, self.span_id, cached_span)
 
         # Customized records are cached after export customization instead, so
-        # local scorers never see content that a customizer redacted.
+        # local scorers never see content that a customizer redacted or dropped.
         if not customizers:
             write_span_cache(serializable_partial_record)
 
-        def compute_record() -> dict[str, Any]:
+        def compute_record() -> dict[str, Any] | None:
             exporter = _get_exporter()
             record = dict(
                 **serializable_partial_record,
@@ -4948,13 +4962,15 @@ def compute_record() -> dict[str, Any]:
                 ).object_id_fields(),
             )
             # Resolve and customize inside the cached LazyValue: every incremental
-            # instrumentation record is transformed once, before the background
+            # span record is transformed once, before the background
             # logger merges, masks, extracts attachments, or retries delivery.
             if customizers:
-                record = _customize_span_export(record, customizers)
-                write_span_cache(record)
+                customized_record = _customize_span_export(record, customizers)
+                if customized_record is not None:
+                    write_span_cache(customized_record)
                 if pending_cache_key is not None:
                     self.state.span_cache._forget_pending_record(self.root_span_id, pending_cache_key)
+                return customized_record
             return record
 
         # Cache readers and the publisher share resolution, including the cache

py/src/braintrust/span_cache.py

@@ -118,8 +118,8 @@ def __init__(self, disabled: bool = False):
         self._root_span_index: set[str] = set()
         # Buffer for pending writes
         self._write_buffer: list[DiskSpanRecord] = []
-        # Unresolved customized exports must be cached before a trace read returns.
-        self._pending_records: dict[str, dict[object, LazyValue[dict[str, Any]]]] = {}
+        # Unresolved exports must be cached or dropped before a trace read returns.
+        self._pending_records: dict[str, dict[object, LazyValue[dict[str, Any] | None]]] = {}
         self._pending_lock = threading.Lock()
 
     def disable(self) -> None:
@@ -203,7 +203,7 @@ def cleanup_all_caches():
 
             atexit.register(cleanup_all_caches)
 
-    def _track_pending_record(self, root_span_id: str, key: object, record: LazyValue[dict[str, Any]]) -> None:
+    def _track_pending_record(self, root_span_id: str, key: object, record: LazyValue[dict[str, Any] | None]) -> None:
         with self._pending_lock:
             if not self.disabled:
                 self._pending_records.setdefault(root_span_id, {})[key] = record

py/src/braintrust/span_customizer.py

@@ -1,6 +1,7 @@
 """Synchronous transformations of native export records."""
 
 import inspect
+import logging
 from collections.abc import Callable, Sequence
 from typing import Any
 
@@ -14,6 +15,7 @@
     PARENT_ID_FIELD,
     TRANSACTION_ID_FIELD,
 )
+from .env import BraintrustEnv
 
 
 __all__ = ["SpanCustomizer", "SpanExportData", "set_span_customizers"]
@@ -23,15 +25,27 @@
 
 
 class SpanCustomizer:
-    """Extensible hooks for instrumentation-created spans. Omitted hooks are no-ops."""
+    """Extensible hooks for all native 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,
         attachment processing, masking, and serialization. Identity, routing, and
         merge protocol fields are restored after each hook. Exceptions and invalid
-        returns fail open: in-place edits survive and later hooks still run.
+        synchronous returns log a safe error, stop the hook chain, and drop that
+        outgoing record. Handle recoverable errors within the hook. Only
+        KeyboardInterrupt and SystemExit propagate.
+
+        Records may contain SDK ``Attachment`` objects: integrations convert inline
+        media at capture time. Remove or replace them to prevent their upload;
+        hooks that serialize data must handle them.
+
+        Manual roots and children, instrumented spans, logger spans, and experiment
+        spans are included. Dataset rows and feedback records are not. Each hook
+        receives a possibly incremental record, not a completed logical span;
+        fields may be absent and hooks may run before the span ends. Dropping a
+        record does not retract earlier exports or disable later records.
         """
         return data
 
@@ -84,16 +98,21 @@ def on_span_export(self, data: SpanExportData) -> SpanExportData:
 def set_span_customizers(customizers: Sequence[SpanCustomizer] | None) -> None:
     """Replace the process-wide ordered customizer list with a snapshot.
 
-    Configure before making instrumented calls; pass None or an empty sequence to
-    disable. The sequence is copied, but customizer instances are not. Each record
+    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 data. There is no environment-variable registration.
+    transformed data or dropped result. There is no environment-variable registration.
 
     Raises TypeError for classes, objects without a callable ``on_span_export``,
-    and async hooks, since those would otherwise silently export unredacted data.
+    and async hooks. Invalid synchronous returns are rejected at export time.
+    Raises NotImplementedError for non-empty customizers when BRAINTRUST_OTEL_COMPAT is enabled.
     """
     global _span_customizers
     snapshot = tuple(customizers) if customizers is not None else ()
+    if snapshot and BraintrustEnv.OTEL_COMPAT.get(False):
+        raise NotImplementedError(
+            "Span customizers are not yet supported with OTel compat mode (BRAINTRUST_OTEL_COMPAT)"
+        )
     for customizer in snapshot:
         _validate_customizer(customizer)
     _span_customizers = snapshot
@@ -104,8 +123,7 @@ def _get_span_customizers() -> tuple[SpanCustomizer, ...]:
 
 
 def _validate_customizer(customizer: Any) -> None:
-    # Reject misconfiguration eagerly: at export time these would silently fail
-    # open and ship unredacted data.
+    # Reject misconfiguration eagerly rather than dropping records at export time.
     if isinstance(customizer, type):
         raise TypeError(f"Span customizers must be instances, not classes; did you mean {customizer.__name__}()?")
     hook = getattr(customizer, "on_span_export", None)
@@ -136,15 +154,14 @@ def _restore_protocol_fields(data: SpanExportData, protected: SpanExportData) ->
 
 def _customize_span_export(
     data: SpanExportData, customizers: Sequence[SpanCustomizer] | None = None
-) -> SpanExportData:
+) -> SpanExportData | None:
     if customizers is None:
         customizers = _span_customizers
     if not customizers:
         return data
 
     protected = None
     for customizer in customizers:
-        candidate = data
         try:
             hook = getattr(customizer, "on_span_export", None)
             if hook is None:
@@ -153,19 +170,22 @@ def _customize_span_export(
                 protected = {key: _copy_protocol_value(data[key]) for key in _PROTECTED_FIELDS if key in data}
                 # Protocol containers may alias span state; don't expose those to hooks.
                 data = _restore_protocol_fields(data, protected)
-                candidate = data
             result = hook(data)
             if inspect.isawaitable(result):
                 # Do not execute asynchronous hooks or emit unawaited coroutine warnings.
                 if inspect.iscoroutine(result):
                     result.close()
-            elif type(result) is dict:
-                candidate = result
-        except Exception:
-            # Deliberately fail open, retaining in-place changes. Never log the
-            # exception, since its message may itself contain sensitive payloads.
-            pass
-
-        if protected is not None:
-            data = _restore_protocol_fields(candidate, protected)
+                raise TypeError("Span customizers must return a synchronous dict")
+            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:
+            # Treat GeneratorExit, CancelledError, etc. as hook failures so the drop
+            # is memoized and cannot escape flush with the rest of the batch.
+            # Neither exception details nor record data are safe to include here.
+            logging.getLogger("braintrust").error("Span customizer failed; dropping outgoing span record.")
+            return None
     return data