feat: add span customizers (SDK-316)

py/README.md

@@ -41,6 +41,70 @@ Then run:
 BRAINTRUST_API_KEY=<YOUR_API_KEY> braintrust eval tutorial_eval.py
 ```
 
+## Customizing instrumentation exports
+
+Use synchronous span customizers to redact captured values or add metadata without
+changing provider responses:
+
+```python
+import braintrust
+
+
+class Redact(braintrust.SpanCustomizer):
+    def on_span_export(self, data: braintrust.SpanExportData) -> braintrust.SpanExportData:
+        for field in ("input", "output"):
+            if field in data:
+                data[field] = "[redacted]"
+        data.pop("error", None)
+        return data
+
+
+braintrust.auto_instrument(span_customizers=[Redact()])
+braintrust.init_logger(project="my-project")
+```
+
+For explicit wrappers or integration setup, call
+`braintrust.set_span_customizers([Redact()])` before instrumented calls. Both APIs
+replace the process-wide ordered list with a snapshot of the supplied sequence;
+the customizer objects themselves are not copied. `set_span_customizers(None)` or
+an empty sequence disables customization. Omitting `span_customizers` (or passing
+`None`) to `auto_instrument()` leaves existing configuration unchanged. There is no
+environment-variable registration.
+
+`SpanCustomizer` is an extensible base class with a default no-op
+`on_span_export(data)` method. Hooks run synchronously in registration order, each
+receiving its predecessor's result. Return a plain `dict` containing the data to
+export: either the input record or a replacement. Replacement is **not** a merge;
+retain the content fields you want to export. Returning `None` cannot drop a span.
+
+Hooks apply only to **native instrumentation-created span records**, not manually
+created spans (even children of instrumented spans), dataset rows, feedback, or
+spans exported through a separate OpenTelemetry exporter. They run after lazy
+values resolve and before attachment processing, merging, masking, and transport
+serialization. Existing attachment objects retain their identity; hooks can
+remove them before upload without copying their contents.
+
+A span can produce several incremental records, including records flushed before
+it ends. Input, output, metrics, and errors need not be present together. Removing
+a field does not retract an earlier export, so redact every record containing
+sensitive data. Each record uses the configuration at its first export resolution.
+Retries reuse the transformed record without calling hooks again.
+
+Identity, parentage, routing, and merge protocol fields are restored after every
+hook, including nested protocol arrays. Protected fields are `id`, `span_id`,
+`root_span_id`, `span_parents`, `org_id`, `project_id`, `experiment_id`, `dataset_id`,
+`prompt_session_id`, `log_id`, `function_data`, `_is_merge`, `_merge_paths`,
+`_parent_id`, `_object_delete`, `_array_delete`, and `_xact_id`. Hooks cannot add an
+absent protected field or change the destination.
+
+**Failures are fail-open.** Exceptions, `None`, awaitables, and other invalid return
+types are ignored; later hooks and export continue with the current record.
+In-place content mutations made before failure are retained, not rolled back.
+An asynchronous hook is not awaited. A failing redaction hook can therefore leak
+unredacted data: never throw to block export. Returned content must remain valid
+for SDK serialization. Keep hooks fast, avoid blocking I/O, and make them safe for
+background export threads rather than assuming the caller's thread or context.
+
 ## Optional Extras
 
 Install extras as needed for specific workflows:

py/src/braintrust/__init__.py

@@ -85,5 +85,8 @@ def is_equal(expected, output):
 from .sandbox import RegisterSandboxResult as RegisterSandboxResult
 from .sandbox import SandboxConfig as SandboxConfig
 from .sandbox import register_sandbox as register_sandbox
+from .span_customizer import SpanCustomizer as SpanCustomizer
+from .span_customizer import SpanExportData as SpanExportData
+from .span_customizer import set_span_customizers as set_span_customizers
 from .util import BT_IS_ASYNC_ATTRIBUTE as BT_IS_ASYNC_ATTRIBUTE
 from .util import MarkAsyncWrapper as MarkAsyncWrapper

py/src/braintrust/auto.py

@@ -5,6 +5,7 @@
 """
 
 import logging
+from collections.abc import Sequence
 from contextlib import contextmanager
 
 from braintrust.integrations import (
@@ -40,6 +41,7 @@
     TypeSafeIntegration,
 )
 from braintrust.integrations.base import BaseIntegration
+from braintrust.span_customizer import SpanCustomizer, set_span_customizers
 
 
 __all__ = ["auto_instrument"]
@@ -90,6 +92,7 @@ def auto_instrument(
     livekit_agents: bool = True,
     pipecat: bool = True,
     typesafe: bool = True,
+    span_customizers: Sequence[SpanCustomizer] | None = None,
 ) -> dict[str, bool]:
     """
     Auto-instrument supported AI/ML libraries for Braintrust tracing.
@@ -131,6 +134,9 @@ 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.
 
     Returns:
         Dict mapping integration name to whether it was successfully instrumented.
@@ -176,6 +182,9 @@ def auto_instrument(
         client.models.generate_content(model="gemini-2.0-flash", contents="Hello!")
         ```
     """
+    if span_customizers is not None:
+        set_span_customizers(span_customizers)
+
     results: dict[str, bool] = {}
 
     if openai:

py/src/braintrust/integrations/auto_test_scripts/test_auto_google_discoveryengine.py

@@ -18,7 +18,11 @@
     print("SUCCESS")
     sys.exit(0)
 
-options = {name: False for name in inspect.signature(auto_instrument).parameters}
+options = {
+    name: False
+    for name, parameter in inspect.signature(auto_instrument).parameters.items()
+    if isinstance(parameter.default, bool)
+}
 assert auto_instrument(**options) == {}
 RankServiceClient = None
 if sys.argv[1] == "before":

py/src/braintrust/integrations/pipecat/test_pipecat.py

@@ -8,7 +8,7 @@
 from pathlib import Path
 
 import pytest
-from braintrust import logger
+from braintrust import SpanCustomizer, logger, set_span_customizers
 from braintrust.integrations.pipecat import (
     BraintrustPipecatObserver,
     PipecatIntegration,
@@ -58,6 +58,45 @@ def _single_span(logs, name):
     return matches[0]
 
 
+@pytest.mark.asyncio
+async def test_span_customizer_redacts_incremental_tts_input(memory_logger):
+    TTSStartedFrame = _import("pipecat.frames.frames.TTSStartedFrame")
+    TTSTextFrame = _import("pipecat.frames.frames.TTSTextFrame")
+    TTSStoppedFrame = _import("pipecat.frames.frames.TTSStoppedFrame")
+
+    class Redact(SpanCustomizer):
+        def on_span_export(self, data):
+            if "input" in data:
+                data["input"] = "[redacted]"
+            return data
+
+    observer = BraintrustPipecatObserver()
+    set_span_customizers([Redact()])
+    try:
+        await observer.on_pipeline_started()
+        await observer._handle_frame(TTSStartedFrame(context_id="ctx"))
+        first_frame = TTSTextFrame("private first", aggregated_by="sentence", context_id="ctx")
+        await observer._handle_frame(first_frame)
+        initial = memory_logger.pop()
+        pipeline = _single_span(initial, "pipecat_pipeline")
+        tts = _single_span(initial, "tts_response")
+        assert tts["input"] == "[redacted]"
+        assert tts["span_parents"] == [pipeline["span_id"]]
+        assert first_frame.text == "private first"
+
+        second_frame = TTSTextFrame("private second", aggregated_by="sentence", context_id="ctx")
+        await observer._handle_frame(second_frame)
+        await observer._handle_frame(TTSStoppedFrame(context_id="ctx"))
+        await observer.cleanup()
+        updates = memory_logger.pop()
+        updated_tts = next(row for row in updates if row["id"] == tts["id"])
+        assert updated_tts["input"] == "[redacted]"
+        assert updated_tts["span_id"] == tts["span_id"]
+        assert second_frame.text == "private second"
+    finally:
+        set_span_customizers(None)
+
+
 def _pipeline_worker_kwargs(**overrides):
     PipelineWorker = _import("pipecat.pipeline.worker.PipelineWorker")
     signature = inspect.signature(PipelineWorker)

py/src/braintrust/integrations/pipecat/tracing.py

@@ -9,7 +9,18 @@
     _pcm_to_wav,
     _resolve_audio_attachment_options,
 )
-from braintrust.logger import NOOP_SPAN, Attachment, SpanTypeAttribute, current_span, start_span
+from braintrust.logger import NOOP_SPAN, Attachment, SpanTypeAttribute, current_span
+from braintrust.logger import start_span as _bt_start_span
+
+
+_INSTRUMENTATION = "pipecat-auto"
+
+
+def start_span(*args, **kwargs):
+    internal = dict(kwargs.get("internal") or {})
+    internal.setdefault("instrumentation", _INSTRUMENTATION)
+    kwargs["internal"] = internal
+    return _bt_start_span(*args, **kwargs)
 
 
 try:

py/src/braintrust/logger.py

@@ -99,6 +99,7 @@
 )
 from .queue import DEFAULT_QUEUE_SIZE, LogQueue
 from .serializable_data_class import SerializableDataClass
+from .span_customizer import _customize_span_export
 from .span_identifier_v3 import SpanComponentsV3, SpanObjectTypeV3
 from .span_identifier_v4 import SpanComponentsV4
 from .span_origin import SpanOriginEnvironment, detect_environment, merge_span_origin_context
@@ -4933,14 +4934,20 @@ def log_internal(self, event: dict[str, Any] | None = None, internal_data: dict[
 
         def compute_record() -> dict[str, Any]:
             exporter = _get_exporter()
-            return dict(
+            record = dict(
                 **serializable_partial_record,
                 **{k: v.get() for k, v in lazy_partial_record.items()},
                 **exporter(
                     object_type=self.parent_object_type,
                     object_id=self.parent_object_id.get(),
                 ).object_id_fields(),
             )
+            # Resolve and customize inside the cached LazyValue: every incremental
+            # instrumentation record is transformed once, before the background
+            # logger merges, masks, extracts attachments, or retries delivery.
+            if self._instrumentation != "braintrust-python-logger":
+                return _customize_span_export(record)
+            return record
 
         self.state.global_bg_logger().log(LazyValue(compute_record, use_mutex=False))
 

py/src/braintrust/span_customizer.py

@@ -0,0 +1,118 @@
+"""Synchronous transformations of native instrumentation export records."""
+
+import inspect
+from collections.abc import Sequence
+from typing import Any
+
+from .db_fields import (
+    ARRAY_DELETE_FIELD,
+    ID_FIELD,
+    IS_MERGE_FIELD,
+    MERGE_PATHS_FIELD,
+    OBJECT_DELETE_FIELD,
+    OBJECT_ID_KEYS,
+    PARENT_ID_FIELD,
+    TRANSACTION_ID_FIELD,
+)
+
+
+__all__ = ["SpanCustomizer", "SpanExportData", "set_span_customizers"]
+
+SpanExportData = dict[str, Any]
+"""A native, possibly incremental span export record, not a live span."""
+
+
+class SpanCustomizer:
+    """Extensible hooks for instrumentation-created 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.
+        """
+        return data
+
+
+_span_customizers: tuple[SpanCustomizer, ...] = ()
+_PROTECTED_FIELDS = frozenset(
+    (
+        ID_FIELD,
+        "span_id",
+        "root_span_id",
+        "span_parents",
+        "org_id",
+        *OBJECT_ID_KEYS,
+        IS_MERGE_FIELD,
+        MERGE_PATHS_FIELD,
+        PARENT_ID_FIELD,
+        OBJECT_DELETE_FIELD,
+        ARRAY_DELETE_FIELD,
+        TRANSACTION_ID_FIELD,
+    )
+)
+
+
+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
+    uses the configuration at its first export resolution, and retries reuse its
+    transformed data. There is no environment-variable registration.
+    """
+    global _span_customizers
+    _span_customizers = tuple(customizers) if customizers is not None else ()
+
+
+def _copy_protocol_value(value: Any) -> Any:
+    # Copy only protocol containers. Payloads (especially Attachment objects) must
+    # retain their existing serialization behavior and must not be deep-copied.
+    if isinstance(value, list):
+        return [_copy_protocol_value(item) for item in value]
+    if isinstance(value, dict):
+        return {key: _copy_protocol_value(item) for key, item in value.items()}
+    return value
+
+
+def _restore_protocol_fields(data: SpanExportData, protected: SpanExportData) -> SpanExportData:
+    restored = {key: value for key, value in data.items() if key not in _PROTECTED_FIELDS}
+    # Never expose the snapshot itself: a later hook may mutate nested arrays.
+    restored.update({key: _copy_protocol_value(value) for key, value in protected.items()})
+    return restored
+
+
+def _customize_span_export(data: SpanExportData) -> SpanExportData:
+    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:
+                continue
+            if protected is None:
+                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)
+    return data