diff --git a/py/src/braintrust/auto.py b/py/src/braintrust/auto.py index 04fc5e303..c29a41408 100644 --- a/py/src/braintrust/auto.py +++ b/py/src/braintrust/auto.py @@ -103,12 +103,12 @@ 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 spans + and OTel spans sent through OtelExporter or BraintrustSpanProcessor, + including manual spans; generic OTLP exporters, 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. diff --git a/py/src/braintrust/otel/__init__.py b/py/src/braintrust/otel/__init__.py index 2ef361161..92221209c 100644 --- a/py/src/braintrust/otel/__init__.py +++ b/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 @@ -132,6 +133,10 @@ class AISpanProcessor: Only filtered spans and root spans will be forwarded to the inner processor. This dramatically reduces telemetry volume while preserving important observability. + This filter does not run Braintrust span customizers. Wrap OtelExporter + instead of a generic OTLPSpanExporter to customize Braintrust exports, + or use BraintrustSpanProcessor for the complete pipeline. + Example: > processor = AISpanProcessor(BatchSpanProcessor(OTLPSpanExporter())) > provider = TracerProvider() @@ -210,6 +215,9 @@ class OtelExporter(OTLPSpanExporter): """ A subclass of OTLPSpanExporter configured for Braintrust. + Runs the registered Braintrust span customizers before OTLP serialization. + Generic OTLPSpanExporter instances do not run these hooks. + For most use cases, consider using the Processor class instead, which provides a more convenient all-in-one interface. @@ -303,6 +311,11 @@ 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 + + spans = [customize_otel_span(span, customizers) for span in spans] return super().export(spans) def force_flush(self, timeout_millis=30000): diff --git a/py/src/braintrust/otel/span_customizer.py b/py/src/braintrust/otel/span_customizer.py new file mode 100644 index 000000000..20ff7294f --- /dev/null +++ b/py/src/braintrust/otel/span_customizer.py @@ -0,0 +1,360 @@ +"""Detached OTel export snapshots for the shared span customizer hook.""" + +from collections.abc import Mapping +from copy import deepcopy +from dataclasses import dataclass, field +from types import SimpleNamespace + +from braintrust.span_customizer import _log_customizer_failure, _run_span_export_customizers +from opentelemetry.sdk.util.instrumentation import InstrumentationScope +from opentelemetry.trace import SpanContext, SpanKind, StatusCode, TraceFlags, TraceState +from opentelemetry.trace.span import _is_valid_pair + + +try: + from opentelemetry.exporter.otlp.proto.common._internal import _encode_key_value +except ImportError: + # OpenTelemetry 1.16 predates the shared HTTP/gRPC encoder package. + from opentelemetry.exporter.otlp.proto.http.trace_exporter.encoder import _encode_key_value + + +def _encodable(key, value): + try: + _encode_key_value(key, value) + except Exception: + return False + return True + + +# Older encoders can't send some values their SDK keeps (e.g. None inside arrays) +# and skip those attributes on export. Drop them before hooks run so only values a +# hook introduces fail closed. Newer encoders accept everything the SDK keeps. +_DROP_UNENCODABLE = not _encodable("probe", [None]) + + +def _snapshot_attributes(attributes): + attributes = _copy_value(attributes or {}) + if _DROP_UNENCODABLE: + return {key: value for key, value in attributes.items() if _encodable(key, value)} + return attributes + + +class _ExportSpan(SimpleNamespace): + def get_span_context(self): + return self.context + + +def _metadata_key(value): + """Hash nested metadata without conflating OTLP scalar types.""" + if isinstance(value, dict): + return (dict, frozenset((key, _metadata_key(item)) for key, item in value.items())) + if isinstance(value, list): + return (list, tuple(_metadata_key(item) for item in value)) + if isinstance(value, float): + return (float, value.hex()) + return (type(value), value) + + +# OTLP groups resources and scopes by these keys. Keep detached attributes +# losslessly, without SDK constructors that can sanitize or log their contents. +@dataclass(frozen=True) +class _ExportResource: + attributes: dict = field(compare=False) + schema_url: str | None + _attributes_key: tuple = field(init=False, repr=False) + + def __post_init__(self): + object.__setattr__(self, "_attributes_key", _metadata_key(self.attributes)) + + +@dataclass(frozen=True) +class _ExportScope: + name: str + version: str | None + schema_url: str | None + attributes: dict = field(compare=False) + _attributes_key: tuple = field(init=False, repr=False) + + def __post_init__(self): + object.__setattr__(self, "_attributes_key", _metadata_key(self.attributes)) + + +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 + # SDK attributes may retain primitive subclasses (for example IntEnum). + # Use the base implementations to detach them without invoking overrides. + if isinstance(value, str): + return str.__str__(value) + if isinstance(value, int): + return int.__int__(value) + if isinstance(value, float): + return float.__float__(value) + if isinstance(value, bytes): + # bytes.__bytes__ is unavailable on Python 3.10. + return bytes.__getitem__(value, slice(None)) + 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": _snapshot_attributes(span.attributes), + "events": [ + { + "name": event.name, + "timestamp": event.timestamp, + "attributes": _snapshot_attributes(event.attributes), + "dropped_attributes": getattr(event, "dropped_attributes", getattr(event.attributes, "dropped", 0)), + } + for event in span.events + ], + "links": [ + { + "context": _context_data(link.context), + "attributes": _snapshot_attributes(link.attributes), + "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": _resource_data(span.resource), + "instrumentation_scope": _scope_data(scope), + "dropped_attributes": span.dropped_attributes, + "dropped_events": span.dropped_events, + "dropped_links": span.dropped_links, + } + + +def _resource_data(resource): + return {"attributes": _snapshot_attributes(resource.attributes), "schema_url": resource.schema_url} + + +def _scope_data(scope): + if not scope: + return None + return { + "name": scope.name, + "version": scope.version, + "schema_url": scope.schema_url, + "attributes": _snapshot_attributes(getattr(scope, "attributes", None)), + } + + +def _string(value, optional=False): + if value is None and optional: + return value + if type(value) is not str: + raise TypeError("Expected an OTel string") + value.encode("utf-8") + return value + + +def _integer(value, bits, signed=False): + if type(value) is not int: + raise TypeError("Expected an OTel integer") + low, high = (-(2 ** (bits - 1)), 2 ** (bits - 1)) if signed else (0, 2**bits) + if not low <= value < high: + raise ValueError("OTel integer out of range") + return value + + +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") + # The installed encoder defines which values are valid. Its attribute + # encoding logs and skips failures, so encode directly to fail closed. + _encode_key_value(key, value) + return data + + +def _check_context(data): + if not isinstance(data, dict): + raise TypeError("OTel links require a span context") + state = data.get("trace_state", {}) + if not isinstance(state, dict) or len(state) > 32: + raise TypeError("Invalid OTel trace state") + for key, value in state.items(): + # TraceState logs invalid pairs, so reject them before construction. + if type(key) is not str or type(value) is not str or not _is_valid_pair(key, value): + raise ValueError("Invalid OTel trace state") + _integer(data["trace_id"], 128) + _integer(data["span_id"], 64) + _integer(data.get("trace_flags", 0), 8) + if type(data.get("is_remote", False)) is not bool: + raise TypeError("Invalid OTel span context") + + +def _check_scope(data): + if data is None: + return + _string(data.get("name", "")) + _string(data.get("version"), optional=True) + _string(data.get("schema_url", ""), optional=True) + if _attributes(data.get("attributes", {})) and not hasattr(InstrumentationScope, "attributes"): + raise ValueError("Scope attributes require a newer OpenTelemetry SDK") + + +def _check_export_data(data): + """Validate a hook result against what the OTLP encoder accepts, without encoding the span.""" + _string(data.get("name", "")) + SpanKind(data.get("kind", 0)) + _attributes(data.get("attributes", {})) + for event in data.get("events", []): + _string(event["name"]) + _integer(event["timestamp"], 64) + _attributes(event.get("attributes", {})) + _integer(event.get("dropped_attributes", 0), 32) + for link in data.get("links", []): + _check_context(link["context"]) + _attributes(link.get("attributes", {})) + _integer(link.get("dropped_attributes", 0), 32) + status = data.get("status", {}) + StatusCode(status.get("status_code", 0)) + _string(status.get("description"), optional=True) + resource = data.get("resource", {}) + _attributes(resource.get("attributes", {})) + _string(resource.get("schema_url", ""), optional=True) + _check_scope(data.get("instrumentation_scope")) + for key in ("dropped_attributes", "dropped_events", "dropped_links"): + _integer(data.get(key, 0), 32) + + +def _context(data): + if data is None: + return None + 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=TraceState(list(data.get("trace_state", {}).items())), + ) + + +def _scope(data): + if data is None: + return None + return _ExportScope( + name=data.get("name", ""), + version=data.get("version"), + schema_url=data.get("schema_url", ""), + attributes=data.get("attributes", {}), + ) + + +def _export_span(data): + 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=data.get("attributes", {}), + events=[ + SimpleNamespace( + name=event["name"], + timestamp=event["timestamp"], + attributes=event.get("attributes", {}), + dropped_attributes=event.get("dropped_attributes", 0), + ) + for event in data.get("events", []) + ], + links=[ + SimpleNamespace( + context=_context(link["context"]), + 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=_ExportResource( + attributes=data.get("resource", {}).get("attributes", {}), + schema_url=data.get("resource", {}).get("schema_url", ""), + ), + instrumentation_scope=_scope(data.get("instrumentation_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.""" + protected = { + "context": _context_data(span.context), + "parent": _context_data(span.parent), + "start_time": span.start_time, + "end_time": span.end_time, + } + routing = {key: value for key, value in (span.attributes or {}).items() if key == "braintrust.parent"} + + 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 + + def stripped_span(): + _log_customizer_failure() + return _export_span(stripped_record()) + + 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. + _check_export_data(result) + return result + + try: + data = _snapshot(span) + except (KeyboardInterrupt, SystemExit): + raise + except BaseException: + # Unexportable source data fails closed for this span, not the batch. + return stripped_span() + + data = _run_span_export_customizers(data, customizers, restore, stripped_record) + try: + return _export_span(data) + except (KeyboardInterrupt, SystemExit): + raise + except BaseException: + return stripped_span() diff --git a/py/src/braintrust/span_customizer.py b/py/src/braintrust/span_customizer.py index d3a26df63..e77108660 100644 --- a/py/src/braintrust/span_customizer.py +++ b/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,36 @@ 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 "OTel snapshot format" in the SDK + README). Kind and status codes are integer enum values, not enum objects. + Primitive subclasses in OTel snapshots are copied as plain built-ins. + OTel hooks run only in OtelExporter, including the exporter used by + BraintrustSpanProcessor; generic OTLPSpanExporter instances do not run them. + BraintrustSpanProcessor filtering and origin enrichment happen first. + 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 +128,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 +196,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 +238,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(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 diff --git a/py/src/braintrust/test_otel.py b/py/src/braintrust/test_otel.py index c79c0813a..ce4fbac83 100644 --- a/py/src/braintrust/test_otel.py +++ b/py/src/braintrust/test_otel.py @@ -1256,3 +1256,610 @@ def test_otel_inject_extracted_by_native(): with logger.start_span(name="native", parent=parent) as span: assert span.root_span_id == format(trace_id, "032x") assert span.span_parents == [format(span_id, "016x")] + + +@pytest.fixture +def customized_otlp_receiver(monkeypatch): + if not OTEL_INSTALLED: + pytest.skip("OpenTelemetry SDK not installed") + from braintrust import set_span_customizers + from google.protobuf import symbol_database + from opentelemetry.proto.collector.trace.v1 import trace_service_pb2 + + request_type = symbol_database.Default().GetSymbol( + f"{trace_service_pb2.DESCRIPTOR.package}.ExportTraceServiceRequest" + ) + + monkeypatch.setattr("braintrust.span_customizer._failure_logged", False) + requests = [] + + def receive(method, path, body, headers): + request = request_type() + request.ParseFromString(body) + requests.append(request) + return 200, {}, b"" + + set_span_customizers(None) + try: + with scripted_server(receive) as (url, _): + yield url, requests + finally: + set_span_customizers(None) + + +def _wire_spans(requests): + return [ + (resource, scope, span) + for request in requests + for resource in request.resource_spans + for scope in resource.scope_spans + for span in scope.spans + ] + + +def _sensitive_otel_spans(): + from opentelemetry.sdk.resources import Resource + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.sdk.trace.export import SimpleSpanProcessor + from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + from opentelemetry.trace import Link, Status, StatusCode + + memory = InMemorySpanExporter() + provider = TracerProvider(resource=Resource({"private.resource": "secret"}), shutdown_on_exit=False) + provider.add_span_processor(SimpleSpanProcessor(memory)) + tracer = provider.get_tracer("private.scope", "private.version", schema_url="private.schema") + with tracer.start_as_current_span("private.parent") as parent: + with tracer.start_as_current_span( + "private.child", links=[Link(parent.get_span_context(), {"private.link": "secret"})] + ) as child: + child.set_attributes( + { + "braintrust.parent": "experiment_id:destination", + "braintrust.input_json": '"private.input"', + "private.array": ["first", "second"], + } + ) + child.add_event("private.event", {"private.event": "secret"}, timestamp=123) + child.set_status(Status(StatusCode.ERROR, "private.status")) + spans = memory.get_finished_spans() + provider.shutdown() + return spans + + +@pytest.mark.parametrize("replacement", [False, True]) +def test_otel_customizers_real_http_snapshot_replacement_and_identity(customized_otlp_receiver, replacement): + from braintrust import SpanCustomizer, set_span_customizers + from braintrust.otel import OtelExporter + + url, requests = customized_otlp_receiver + spans = _sensitive_otel_spans() + seen = [] + + class Redact(SpanCustomizer): + def on_span_export(self, data): + assert type(data) is dict + # Direct exporters don't add provenance, with or without customizers. + assert "braintrust.context_json" not in data["attributes"] + seen.append(data["name"]) + data["context"]["span_id"] = 1 + data["context"]["trace_id"] = 2 + if data["parent"]: + data["parent"]["span_id"] = 3 + data["start_time"] = 4 + data["end_time"] = 5 + data["attributes"]["braintrust.parent"] = "project_name:wrong" + if "private.array" in data["attributes"]: + data["attributes"]["private.array"][0] = "safe" + data["events"][0]["attributes"].clear() + data["links"][0]["attributes"].clear() + data["resource"]["attributes"].clear() + data["instrumentation_scope"]["name"] = "safe.scope" + if replacement: + return {"name": "redacted", "attributes": {"safe": True}} + data["name"] = "redacted" + data["attributes"] = {"safe": True} + data["events"] = [] + data["links"] = [] + data.pop("status") + data.pop("resource") + data.pop("instrumentation_scope") + return data + + class Observe(SpanCustomizer): + def on_span_export(self, data): + original = spans[len(observed)] + assert data["context"]["span_id"] == original.context.span_id + assert data["start_time"] == original.start_time + observed.append(data) + return data + + observed = [] + set_span_customizers([Redact(), Observe()]) + exporter = OtelExporter(url=url, api_key="test") + try: + exporter.export(spans) + finally: + exporter.shutdown() + assert seen == ["private.child", "private.parent"] + assert len(observed) == 2 + for (resource, scope, wire), original in zip(_wire_spans(requests), spans): + assert wire.name == "redacted" + assert wire.span_id == original.context.span_id.to_bytes(8, "big") + assert wire.trace_id == original.context.trace_id.to_bytes(16, "big") + assert wire.parent_span_id == (original.parent.span_id.to_bytes(8, "big") if original.parent else b"") + assert (wire.start_time_unix_nano, wire.end_time_unix_nano) == (original.start_time, original.end_time) + assert not wire.events and not wire.links and not wire.status.message + assert not resource.resource.attributes and not scope.scope.name + attrs = {attr.key: attr.value for attr in wire.attributes} + assert attrs["safe"].bool_value + assert "braintrust.context_json" not in attrs + if original.parent: + assert attrs["braintrust.parent"].string_value == "experiment_id:destination" + child = spans[0] + assert child.name == "private.child" + assert child.attributes["private.array"] == ("first", "second") + assert child.events[0].attributes["private.event"] == "secret" + assert child.links[0].attributes["private.link"] == "secret" + assert child.resource.attributes["private.resource"] == "secret" + assert child.instrumentation_scope.name == "private.scope" + + +@pytest.mark.parametrize( + "failure", + [ + "exception", + "base_exception", + "none", + "coroutine", + "attributes", + "attribute_value", + "integer_overflow", + "invalid_utf8_key", + "event", + "link", + "null_link_context", + "status", + "resource", + "scope", + "link_trace_state", + "name", + "cycle", + ], +) +def test_otel_customizer_fail_closed_real_http_mixed_batch(customized_otlp_receiver, caplog, failure): + import logging + + from braintrust import SpanCustomizer, set_span_customizers + from braintrust.otel import OtelExporter + from braintrust.span_customizer import _customize_span_export + + url, requests = customized_otlp_receiver + spans = _sensitive_otel_spans() + calls = [] + + async def asynchronous(): + return {} + + class Break(SpanCustomizer): + def on_span_export(self, data): + if data.get("name") == "private.parent": + return {"name": "healthy"} + data["name"] = "private.partial" + if failure == "exception": + raise ValueError("private.exception") + if failure == "base_exception": + raise GeneratorExit("private.exception") + if failure == "none": + return None + if failure == "coroutine": + return asynchronous() + if failure == "cycle": + data["attributes"] = data + else: + key, value = { + "attributes": ("attributes", []), + "attribute_value": ("attributes", {"private": object()}), + "integer_overflow": ("attributes", {"private": 2**64}), + "invalid_utf8_key": ("attributes", {"private\ud800": "secret"}), + "event": ("events", [{"name": "private", "timestamp": -1}]), + "link": ("links", [{"context": {"trace_id": -1, "span_id": 1}}]), + "null_link_context": ("links", [{"context": None}]), + "status": ("status", {"status_code": 999, "description": "private"}), + "resource": ("resource", {"attributes": {"private": object()}}), + "scope": ("instrumentation_scope", {"attributes": {"private": object()}}), + "link_trace_state": ( + "links", + [{"context": {"trace_id": 1, "span_id": 1, "trace_state": {"private": "secret,"}}}], + ), + "name": ("name", {}), + }[failure] + data[key] = value + return data + + class Observe(SpanCustomizer): + def on_span_export(self, data): + calls.append(data["name"]) + return data + + set_span_customizers([Break(), Observe()]) + exporter = OtelExporter(url=url, api_key="test") + with caplog.at_level(logging.ERROR): + try: + exporter.export(spans) + exporter.export(spans) + # Native and OTel share the same once-only diagnostic budget. + _customize_span_export({"id": "row"}, [Break()]) + finally: + exporter.shutdown() + assert calls == ["healthy", "healthy"] + assert len(caplog.records) == 1 + assert "private" not in caplog.text + assert len(requests) == 2 + for request in requests: + assert len(request.resource_spans) == 1 + encoded = request.SerializeToString() + assert b"private" not in encoded and b"secret" not in encoded + records = _wire_spans([request]) + assert len(records) == 2 + for resource, scope, wire in records: + assert not resource.resource.attributes and not resource.schema_url + assert not scope.scope.name and not scope.scope.version and not scope.schema_url + assert not wire.events and not wire.links and not wire.status.message + if wire.name == "healthy": + continue + assert not wire.name and not wire.trace_state + assert wire.span_id == spans[0].context.span_id.to_bytes(8, "big") + assert wire.trace_id == spans[0].context.trace_id.to_bytes(16, "big") + assert wire.parent_span_id == spans[0].parent.span_id.to_bytes(8, "big") + assert (wire.start_time_unix_nano, wire.end_time_unix_nano) == (spans[0].start_time, spans[0].end_time) + assert {a.key: a.value.string_value for a in wire.attributes} == { + "braintrust.parent": "experiment_id:destination", + "braintrust.context_json": '{"customizer_error":true}', + } + + +@pytest.mark.parametrize("compat", [False, True]) +def test_otel_processor_customizes_after_filter_once_without_mutating_other_exporters( + customized_otlp_receiver, monkeypatch, compat +): + from braintrust import SpanCustomizer, set_span_customizers + from braintrust.otel import BraintrustSpanProcessor + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.sdk.trace.export import SimpleSpanProcessor + from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + + monkeypatch.setenv("BRAINTRUST_OTEL_COMPAT", str(compat).lower()) + url, requests = customized_otlp_receiver + seen = [] + + class Redact(SpanCustomizer): + def on_span_export(self, data): + seen.append(data["name"]) + assert "span_origin" in json.loads(data["attributes"]["braintrust.context_json"]) + return {"name": "redacted"} + + set_span_customizers([Redact()]) + provider = TracerProvider(shutdown_on_exit=False) + provider.add_span_processor( + BraintrustSpanProcessor( + api_key="test", + api_url=url, + SpanProcessor=SimpleSpanProcessor, + filter_ai_spans=True, + ) + ) + other = InMemorySpanExporter() + provider.add_span_processor(SimpleSpanProcessor(other)) + tracer = provider.get_tracer(__name__) + with tracer.start_as_current_span("gen_ai.root"): + tracer.start_span("gen_ai.child").end() + tracer.start_span("filtered").end() + provider.shutdown() + assert seen == ["gen_ai.child", "gen_ai.root"] + assert [wire.name for _, _, wire in _wire_spans(requests)] == ["redacted", "redacted"] + assert [span.name for span in other.get_finished_spans()] == ["gen_ai.child", "filtered", "gen_ai.root"] + + +def test_otel_customizer_preserves_successful_snapshot_content(customized_otlp_receiver): + from braintrust import SpanCustomizer, set_span_customizers + from braintrust.otel import OtelExporter + + url, requests = customized_otlp_receiver + spans = _sensitive_otel_spans() + + class Customize(SpanCustomizer): + def on_span_export(self, data): + if data["events"]: + data["events"][0]["name"] = "safe.event" + data["events"][0]["attributes"] = {"safe.event": True} + data["links"][0]["attributes"] = {"safe.link": True} + data["status"]["description"] = "safe.status" + data["resource"]["attributes"] = {"safe.resource": "value"} + data["instrumentation_scope"]["name"] = "safe.scope" + return data + + set_span_customizers([Customize()]) + exporter = OtelExporter(url=url, api_key="test") + try: + exporter.export(spans) + finally: + exporter.shutdown() + resource, scope, child = _wire_spans(requests)[0] + assert child.name == "private.child" + assert child.events[0].name == "safe.event" and child.events[0].time_unix_nano == 123 + assert child.events[0].attributes[0].value.bool_value + assert child.links[0].attributes[0].value.bool_value + assert child.status.message == "safe.status" + assert resource.resource.attributes[0].key == "safe.resource" + assert scope.scope.name == "safe.scope" + + +@pytest.mark.parametrize("experiment", [False, True]) +def test_customizers_cover_native_and_otel_compat_roots_and_children( + customized_otlp_receiver, monkeypatch, experiment +): + from braintrust import SpanCustomizer, set_span_customizers + from braintrust.logger import _internal_with_memory_background_logger + from braintrust.otel import BraintrustSpanProcessor + from braintrust.test_helpers import init_test_exp, init_test_logger + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.sdk.trace.export import SimpleSpanProcessor + + monkeypatch.setenv("BRAINTRUST_OTEL_COMPAT", "true") + url, requests = customized_otlp_receiver + shapes = set() + + class Redact(SpanCustomizer): + def on_span_export(self, data): + if "attributes" in data: + shapes.add("otel") + data["attributes"].pop("braintrust.input_json", None) + else: + shapes.add("native") + data.pop("input", None) + return data + + set_span_customizers([Redact()]) + provider = TracerProvider(shutdown_on_exit=False) + provider.add_span_processor( + BraintrustSpanProcessor(api_key="test", api_url=url, SpanProcessor=SimpleSpanProcessor) + ) + tracer = provider.get_tracer(__name__) + try: + with _internal_with_memory_background_logger() as memory: + owner = init_test_exp("customized", "project") if experiment else init_test_logger("customized") + with owner.start_span(name="native.root", input="private") as root: + with root.start_span(name="native.child", input="private"): + with tracer.start_as_current_span("otel.root", attributes={"braintrust.input_json": '"private"'}): + tracer.start_span("otel.child", attributes={"braintrust.input_json": '"private"'}).end() + rows = memory.pop() + finally: + provider.shutdown() + assert shapes == {"native", "otel"} + assert len(rows) == 2 and all("input" not in row for row in rows) + wires = _wire_spans(requests) + assert len(wires) == 2 + assert all("braintrust.input_json" not in {a.key for a in wire.attributes} for _, _, wire in wires) + + +def test_otel_noop_customizer_exports_same_request_as_no_customizers(customized_otlp_receiver): + from http import HTTPStatus + + from braintrust import SpanCustomizer, set_span_customizers + from braintrust.otel import OtelExporter + from opentelemetry import trace as otel_trace + from opentelemetry.sdk.resources import Resource + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.sdk.trace.export import SimpleSpanProcessor + from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + from opentelemetry.trace import NonRecordingSpan, SpanContext, TraceFlags, TraceState + + url, requests = customized_otlp_receiver + + class Label(str): + pass + + class Measure(float): + pass + + class Binary(bytes): + pass + + # Newer SDKs keep None, bytes, mappings, and mixed arrays. Older SDKs keep None + # inside arrays, which their encoders can't send and skip without customizers. + attributes = { + "none": None, + "bytes": b"value", + "mapping": {"nested": {"value": 1}}, + "mixed": [1, "two"], + "none_in_array": ["one", None], + "string": "value", + "http.status_code": HTTPStatus.OK, + "label": Label("value"), + "measure": Measure(1.5), + "binary_subclass": Binary(b"value"), + "status_array": [HTTPStatus.OK, HTTPStatus.NOT_FOUND], + } + + memory = InMemorySpanExporter() + provider = TracerProvider(resource=Resource({"service.name": "svc", "mapping": {"a": 1}}), shutdown_on_exit=False) + provider.add_span_processor(SimpleSpanProcessor(memory)) + tracer = provider.get_tracer("scope", "1.0") + remote = SpanContext( + trace_id=1, span_id=2, is_remote=True, trace_flags=TraceFlags(1), trace_state=TraceState([("k", "v")]) + ) + with tracer.start_as_current_span("root", context=otel_trace.set_span_in_context(NonRecordingSpan(remote))): + with tracer.start_as_current_span("child") as child: + for key, value in attributes.items(): + child.set_attribute(key, value) + child.add_event("event", {"mapping": {"a": 1}} if "mapping" in attributes else {"a": 1}) + spans = memory.get_finished_spans() + provider.shutdown() + + class Noop(SpanCustomizer): + def on_span_export(self, data): + assert data["context"]["trace_state"] == {"k": "v"} + return data + + exporter = OtelExporter(url=url, api_key="test") + try: + exporter.export(spans) + set_span_customizers([Noop()]) + exporter.export(spans) + finally: + exporter.shutdown() + assert "none_in_array" in spans[0].attributes + baseline, customized = requests + assert len(customized.resource_spans) == 1 + assert customized == baseline + + +@pytest.mark.parametrize("field", ["resource", "instrumentation_scope"]) +def test_otel_customizer_preserves_encoder_valid_metadata(customized_otlp_receiver, caplog, field): + from braintrust import SpanCustomizer, set_span_customizers + from braintrust.otel import OtelExporter + from opentelemetry.sdk.trace.export import SpanExportResult + from opentelemetry.sdk.util.instrumentation import InstrumentationScope + + scope_attributes_supported = hasattr(InstrumentationScope, "attributes") + if field == "instrumentation_scope" and not scope_attributes_supported: + pytest.skip("This OTel encoder does not export scope attributes") + + url, requests = customized_otlp_receiver + spans = _sensitive_otel_spans() + attributes = {"mixed": [1, "two"]} + if scope_attributes_supported: + attributes["nested"] = {"value": "safe"} + + class Customize(SpanCustomizer): + def on_span_export(self, data): + # Dictionary insertion order must not split equivalent OTLP groups. + data[field]["attributes"] = dict(reversed(attributes.items())) if data["parent"] else attributes + return data + + set_span_customizers([Customize()]) + exporter = OtelExporter(url=url, api_key="test") + try: + assert exporter.export(spans) == SpanExportResult.SUCCESS + finally: + exporter.shutdown() + + assert not caplog.records + assert len(requests) == 1 + assert len(requests[0].resource_spans) == 1 + assert len(requests[0].resource_spans[0].scope_spans) == 1 + records = _wire_spans(requests) + assert len(records) == 2 + for resource, scope, wire in records: + assert wire.name in {"private.child", "private.parent"} + metadata = resource.resource if field == "resource" else scope.scope + attrs = {attribute.key: attribute.value for attribute in metadata.attributes} + assert set(attrs) == set(attributes) + mixed = attrs["mixed"].array_value.values + assert len(mixed) == 2 + assert mixed[0].int_value == 1 + assert mixed[1].string_value == "two" + if scope_attributes_supported: + nested = attrs["nested"].kvlist_value.values + assert len(nested) == 1 + assert nested[0].key == "value" + assert nested[0].value.string_value == "safe" + assert spans[0].resource.attributes["private.resource"] == "secret" + + +@pytest.mark.parametrize("field", ["resource", "instrumentation_scope"]) +def test_otel_customizer_groups_equal_metadata_without_merging_distinct_values(customized_otlp_receiver, field): + from braintrust import SpanCustomizer, set_span_customizers + from braintrust.otel import OtelExporter + from opentelemetry.sdk.resources import Resource + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.sdk.trace.export import SimpleSpanProcessor, SpanExportResult + from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + from opentelemetry.sdk.util.instrumentation import InstrumentationScope + + url, requests = customized_otlp_receiver + variants = { + "same-a": {"schema_url": "safe"}, + "same-b": {"schema_url": "safe"}, + "schema": {"schema_url": "different"}, + } + if field == "resource" or hasattr(InstrumentationScope, "attributes"): + variants.update( + { + "bool": {"schema_url": "safe", "attributes": {"value": True}}, + "int": {"schema_url": "safe", "attributes": {"value": 1}}, + "float": {"schema_url": "safe", "attributes": {"value": 1.0}}, + "array": {"schema_url": "safe", "attributes": {"value": [1]}}, + } + ) + if field == "instrumentation_scope": + variants.update({"name": {"name": "different"}, "version": {"version": "2"}}) + + memory = InMemorySpanExporter() + provider = TracerProvider(resource=Resource({}), shutdown_on_exit=False) + provider.add_span_processor(SimpleSpanProcessor(memory)) + tracer = provider.get_tracer("source", "1") + for name in variants: + tracer.start_span(name).end() + provider.shutdown() + + class Customize(SpanCustomizer): + def on_span_export(self, data): + data[field] = variants[data["name"]] + return data + + set_span_customizers([Customize()]) + exporter = OtelExporter(url=url, api_key="test") + try: + assert exporter.export(memory.get_finished_spans()) == SpanExportResult.SUCCESS + finally: + exporter.shutdown() + + assert len(requests) == 1 + if field == "resource": + groups = [ + frozenset(span.name for scope in resource.scope_spans for span in scope.spans) + for resource in requests[0].resource_spans + ] + else: + assert len(requests[0].resource_spans) == 1 + groups = [frozenset(span.name for span in scope.spans) for scope in requests[0].resource_spans[0].scope_spans] + assert len(groups) == len(variants) - 1 + assert set(groups) == {frozenset({"same-a", "same-b"})} | { + frozenset({name}) for name in variants if name not in {"same-a", "same-b"} + } + + +def test_otel_customizer_unexportable_span_strips_only_that_span(customized_otlp_receiver, caplog): + import logging + + from braintrust import SpanCustomizer, set_span_customizers + from braintrust.otel import OtelExporter, _SpanWithAttributes + from opentelemetry.sdk.trace.export import SpanExportResult + + url, requests = customized_otlp_receiver + healthy, original = _sensitive_otel_spans() + broken = _SpanWithAttributes(original, {**original.attributes, "private": object()}) + calls = [] + + class Observe(SpanCustomizer): + def on_span_export(self, data): + calls.append(data["name"]) + return data + + set_span_customizers([Observe()]) + exporter = OtelExporter(url=url, api_key="test") + with caplog.at_level(logging.ERROR): + try: + assert exporter.export([healthy, broken]) == SpanExportResult.SUCCESS + finally: + exporter.shutdown() + assert calls == ["private.child"] + assert len(caplog.records) == 1 and "private" not in caplog.text + wires = {wire.span_id: wire for _, _, wire in _wire_spans(requests)} + assert wires[healthy.context.span_id.to_bytes(8, "big")].name == "private.child" + stripped = wires[original.context.span_id.to_bytes(8, "big")] + assert not stripped.name + assert {a.key: a.value.string_value for a in stripped.attributes} == { + "braintrust.context_json": '{"customizer_error":true}' + } diff --git a/py/src/braintrust/test_span_customizer.py b/py/src/braintrust/test_span_customizer.py index 5c7528673..8a922fc29 100644 --- a/py/src/braintrust/test_span_customizer.py +++ b/py/src/braintrust/test_span_customizer.py @@ -95,37 +95,15 @@ def test_registration_rejects_misconfigured_customizers(invalid): assert _customize_span_export({"id": "row", "output": "secret"}) == {"id": "row", "output": "valid"} -def test_registration_rejected_in_otel_compat_mode(monkeypatch, caplog): - # Keep auto-instrumentation independent of optional provider libraries. +def test_registration_supported_in_otel_compat_mode(monkeypatch, caplog): monkeypatch.setattr("braintrust.auto._instrument_integration", lambda _module, _class: False) - valid = _Valid() - set_span_customizers([valid]) monkeypatch.setenv("BRAINTRUST_OTEL_COMPAT", "true") - with caplog.at_level(logging.ERROR, logger="braintrust"): - # A different hook makes replacement of the previous registration observable. - set_span_customizers([SpanCustomizer()]) - assert len(caplog.records) == 1 - assert caplog.records[0].levelno == logging.ERROR - assert _customize_span_export({"output": "secret"}) == {"output": "valid"} - - auto_instrument(span_customizers=[SpanCustomizer()]) - assert len(caplog.records) == 2 - assert caplog.records[1].levelno == logging.ERROR - # Record processing must neither replace the previous hooks nor repeat the diagnostic. - for output in ("first", "second"): - assert _customize_span_export({"output": output}) == {"output": "valid"} - assert len(caplog.records) == 2 - - # Both clearing forms remain silent, including through auto_instrument. - auto_instrument(span_customizers=[]) - assert _customize_span_export({"output": "secret"}) == {"output": "secret"} - monkeypatch.setenv("BRAINTRUST_OTEL_COMPAT", "false") - set_span_customizers([valid]) - assert _customize_span_export({"output": "secret"}) == {"output": "valid"} - monkeypatch.setenv("BRAINTRUST_OTEL_COMPAT", "true") - set_span_customizers(None) - assert _customize_span_export({"output": "secret"}) == {"output": "secret"} - assert len(caplog.records) == 2 + set_span_customizers([SpanCustomizer()]) + auto_instrument(span_customizers=[_Valid()]) + assert _customize_span_export({"output": "secret"}) == {"output": "valid"} + auto_instrument(span_customizers=[]) + assert _customize_span_export({"output": "secret"}) == {"output": "secret"} + assert not caplog.records def test_duck_typed_customizers_are_accepted():