From e430ff3d0a3e20ff5b430351c18dbb4397284625 Mon Sep 17 00:00:00 2001 From: Andrew Kent Date: Mon, 12 Oct 2026 11:55:35 -0600 Subject: [PATCH] Fail closed on span customizer errors without dropping records Co-Authored-By: Claude Opus 5.5 --- py/src/braintrust/logger.py | 122 +++--- py/src/braintrust/merge_row_batch.py | 36 +- py/src/braintrust/span_customizer.py | 113 +++-- py/src/braintrust/test_merge_row_batch.py | 72 +++- py/src/braintrust/test_span_customizer.py | 482 ++++++++++++++++++++-- 5 files changed, 701 insertions(+), 124 deletions(-) diff --git a/py/src/braintrust/logger.py b/py/src/braintrust/logger.py index e8243f070..b12079c7c 100644 --- a/py/src/braintrust/logger.py +++ b/py/src/braintrust/logger.py @@ -80,7 +80,7 @@ ) from .git_fields import GitMetadataSettings, RepoInfo from .gitutil import get_past_n_ancestors, get_repo_info -from .merge_row_batch import batch_items, merge_row_batch +from .merge_row_batch import batch_items, merge_row_batch, merge_row_batch_waves from .object import DEFAULT_IS_LEGACY_DATASET, ensure_dataset_record from .parameters import RemoteEvalParameters from .prompt import BRAINTRUST_PARAMS, ImagePart, PromptBlockData, PromptData, PromptMessage, PromptSchema, TextPart @@ -102,7 +102,12 @@ ) from .queue import DEFAULT_QUEUE_SIZE, LogQueue from .serializable_data_class import SerializableDataClass -from .span_customizer import SpanCustomizer, _customize_span_export, _get_span_customizers, _MaskingCustomizer +from .span_customizer import ( + SpanCustomizer, + _customize_span_export, + _get_span_customizers, + _MaskingCustomizer, +) from .span_identifier_v3 import SpanComponentsV3, SpanObjectTypeV3 from .span_identifier_v4 import SpanComponentsV4 from .span_origin import SpanOriginEnvironment, detect_environment, merge_span_origin_context @@ -996,7 +1001,7 @@ def utf8_byte_length(value: str) -> int: return len(value) if value.isascii() else len(value.encode("utf-8")) -_LazyLogRecord = LazyValue[dict[str, Any]] | LazyValue[dict[str, Any] | None] +_LazyLogRecord = LazyValue[dict[str, Any]] class _BackgroundLogger(ABC): @@ -1033,8 +1038,7 @@ def flush(self, batch_size: int | None = None): if not self.logs: return - # Filter dropped records before touching attachments. - logs = [record for item in self.logs if (record := item.get()) is not None] + logs = [item.get() for item in self.logs] # Extract attachments from all logs attachments: list[BaseAttachment] = [] @@ -1046,7 +1050,7 @@ def flush(self, batch_size: int | None = None): def pop(self) -> list[dict[str, Any]]: with self.lock: - logs = [record for item in self.logs if (record := item.get()) is not None] + logs = [item.get() for item in self.logs] self.logs = [] if not logs: @@ -1057,11 +1061,7 @@ def pop(self) -> list[dict[str, Any]]: batch = merge_row_batch(logs) if self._export_customizers: - batch = [ - record - for item in batch - if (record := _customize_span_export(item, self._export_customizers)) is not None - ] + batch = [_customize_span_export(item, self._export_customizers) for item in batch] return batch @@ -1206,40 +1206,28 @@ def flush(self, batch_size: int | None = None): # Drain the queue. wrapped_items = self.queue.drain_all() - all_items, attachments = self._unwrap_lazy_values(wrapped_items) - if len(all_items) == 0: + waves, attachments = self._unwrap_lazy_values(wrapped_items) + if not any(waves): return - # Construct batches of records to flush in parallel. - all_items_with_meta = [stringify_with_overflow_meta(item) for item in all_items] max_request_size_result = self._get_max_request_size() - batches = batch_items( - items=all_items_with_meta, - batch_max_num_items=batch_size, - batch_max_num_bytes=max( - 0, max_request_size_result["max_request_size"] - utf8_byte_length(construct_logs3_data([])) + 1 - ), - get_byte_size=lambda item: item.overflow_meta.byte_size + 1, + batch_max_num_bytes = max( + 0, max_request_size_result["max_request_size"] - utf8_byte_length(construct_logs3_data([])) + 1 ) - - post_promises = [] - try: - post_promises = [ - HTTP_REQUEST_THREAD_POOL.submit(self._submit_logs_request, batch, max_request_size_result) - for batch in batches - ] - except RuntimeError: - # If the thread pool has shut down, e.g. because the process - # is terminating, run the requests the old fashioned way. - for batch in batches: - self._submit_logs_request(batch, max_request_size_result) - - concurrent.futures.wait(post_promises) - # Raise any exceptions from the promises as one group. - post_promise_exceptions = [e for e in (f.exception() for f in post_promises) if e is not None] - if post_promise_exceptions: + # Waves only split same-ID updates around stripped customizer + # failures; send them in order, and each wave's batches in parallel. + post_errors: list[BaseException] = [] + for wave in waves: + batches = batch_items( + items=[stringify_with_overflow_meta(item) for item in wave], + batch_max_num_items=batch_size, + batch_max_num_bytes=batch_max_num_bytes, + get_byte_size=lambda item: item.overflow_meta.byte_size + 1, + ) + post_errors.extend(self._submit_logs_requests(batches, max_request_size_result)) + if post_errors: raise exceptiongroup.BaseExceptionGroup( - f"Encountered the following errors while logging:", post_promise_exceptions + f"Encountered the following errors while logging:", post_errors ) attachment_errors: list[Exception] = [] @@ -1259,27 +1247,48 @@ def flush(self, batch_size: int | None = None): attachment_errors, ) + def _submit_logs_requests( + self, batches: Sequence[Sequence[LogItemWithMeta]], max_request_size_result: dict[str, Any] + ) -> list[BaseException]: + post_promises = [] + errors: list[BaseException] = [] + for index, batch in enumerate(batches): + try: + post_promises.append( + HTTP_REQUEST_THREAD_POOL.submit(self._submit_logs_request, batch, max_request_size_result) + ) + except RuntimeError: + # If the thread pool has shut down, e.g. because the process + # is terminating, run the remaining requests the old fashioned way. + for remaining_batch in batches[index:]: + try: + self._submit_logs_request(remaining_batch, max_request_size_result) + except Exception as e: + errors.append(e) + break + concurrent.futures.wait(post_promises) + return [e for e in (f.exception() for f in post_promises) if e is not None] + errors + def _unwrap_lazy_values( self, wrapped_items: Sequence[_LazyLogRecord] - ) -> tuple[list[dict[str, Any]], list["BaseAttachment"]]: + ) -> tuple[list[list[dict[str, Any]]], list["BaseAttachment"]]: for i in range(self.num_tries): try: - unwrapped_items = [record for item in wrapped_items if (record := item.get()) is not None] - merged_items = merge_row_batch(unwrapped_items) + unwrapped_items = [item.get() for item in wrapped_items] + waves = merge_row_batch_waves(unwrapped_items) # Logger-local hooks run after span hooks and merging. if self._export_customizers: - merged_items = [ - record - for item in merged_items - if (record := _customize_span_export(item, self._export_customizers)) is not None + waves = [ + [_customize_span_export(item, self._export_customizers) for item in wave] for wave in waves ] attachments: list["BaseAttachment"] = [] - for item in merged_items: - _extract_attachments(item, attachments) + for wave in waves: + for item in wave: + _extract_attachments(item, attachments) - return merged_items, attachments + return waves, attachments except Exception as e: errmsg = "Encountered error when constructing records to flush" is_retrying = i + 1 < self.num_tries @@ -1441,8 +1450,8 @@ def _dump_dropped_events(self, wrapped_items): if not (wrapped_items and publish_payloads_dir): return try: - all_items, attachments = self._unwrap_lazy_values(wrapped_items) - items_with_meta = [stringify_with_overflow_meta(item) for item in all_items] + waves, attachments = self._unwrap_lazy_values(wrapped_items) + items_with_meta = [stringify_with_overflow_meta(item) for wave in waves for item in wave] dataStr = construct_logs3_data(items_with_meta) attachment_str = bt_dumps([a.debug_info() for a in attachments]) payload = "{" + f""""data": {dataStr}, "attachments": {attachment_str}""" + "}" @@ -4034,7 +4043,7 @@ def parent_ids(): object_id=parent_object_id.get(), ).object_id_fields() - def compute_record() -> dict[str, Any] | None: + def compute_record() -> dict[str, Any]: record = dict( id=id, **update_event, @@ -4978,11 +4987,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 or dropped. + # local scorers never see content that a customizer redacted or stripped. if not customizers: write_span_cache(serializable_partial_record) - def compute_record() -> dict[str, Any] | None: + def compute_record() -> dict[str, Any]: exporter = _get_exporter() record = dict( **serializable_partial_record, @@ -4997,8 +5006,7 @@ def compute_record() -> dict[str, Any] | None: # logger merges, masks, extracts attachments, or retries delivery. if customizers: customized_record = _customize_span_export(record, customizers) - if customized_record is not None: - write_span_cache(customized_record) + 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 diff --git a/py/src/braintrust/merge_row_batch.py b/py/src/braintrust/merge_row_batch.py index bd3a1ead8..5765d6d49 100644 --- a/py/src/braintrust/merge_row_batch.py +++ b/py/src/braintrust/merge_row_batch.py @@ -2,6 +2,7 @@ from typing import Any, TypeVar from .db_fields import IS_MERGE_FIELD +from .span_customizer import _StrippedSpanExportData T = TypeVar("T") @@ -53,11 +54,22 @@ def _restore_merge_row_skip_fields(row: dict[str, Any], skip_fields: dict[str, A def merge_row_batch(rows: Sequence[dict[str, Any]]) -> list[dict[str, Any]]: + """Like `merge_row_batch_waves`, flattened into one ordered list of rows.""" + return [row for wave in merge_row_batch_waves(rows) for row in wave] + + +def merge_row_batch_waves(rows: Sequence[dict[str, Any]]) -> list[list[dict[str, Any]]]: """Given a batch of rows, merges conflicting rows together to end up with a - set of rows to insert. Returns a set of de-conflicted rows as a flat list. + set of rows to insert. Returns de-conflicted rows split into ordered waves. Note that the returned rows will be the same objects as the input `rows`, meaning they are mutated in place. + SDK-stripped customizer failures are exceptions: they are never mutated and + act as merge barriers for their own row only. Rows for that key before the + barrier, the barrier itself, and rows after it land in successive waves. + A wave never contains two rows with the same key, and without barriers there + is exactly one wave. Exporters must send waves in order to preserve same-ID + update ordering; rows within a wave are independent. There are a few important considerations for the merge procedure: @@ -99,10 +111,24 @@ def merge_row_batch(rows: Sequence[dict[str, Any]]) -> list[dict[str, Any]]: "Logged row is missing an id. This is an internal braintrust error. Please contact us at info@braintrust.dev for help" ) - row_groups: dict[_MergedRowKey, dict[str, Any]] = {} + waves: list[dict[_MergedRowKey, dict[str, Any]]] = [] + # The earliest wave each key may still be written to. + next_wave: dict[_MergedRowKey, int] = {} for row in rows: key = _generate_merged_row_key(row) - existing_row = row_groups.get(key) + wave_index = next_wave.get(key, 0) + if wave_index == len(waves): + waves.append({}) + existing_row = waves[wave_index].get(key) + if isinstance(row, _StrippedSpanExportData): + # Never mutate a memoized failure with neighboring record content. + if existing_row is not None: + wave_index += 1 + if wave_index == len(waves): + waves.append({}) + waves[wave_index][key] = row + next_wave[key] = wave_index + 1 + continue # If there is an existing row and the new row has the IS_MERGE_FIELD == # True property, we merge it with the existing row. Otherwise we can # replace it. @@ -116,9 +142,9 @@ def merge_row_batch(rows: Sequence[dict[str, Any]]) -> list[dict[str, Any]]: if preserve_nomerge: del existing_row[IS_MERGE_FIELD] else: - row_groups[key] = row + waves[wave_index][key] = row - return list(row_groups.values()) + return [list(wave.values()) for wave in waves] def batch_items( diff --git a/py/src/braintrust/span_customizer.py b/py/src/braintrust/span_customizer.py index 4b4ce4c40..d3a26df63 100644 --- a/py/src/braintrust/span_customizer.py +++ b/py/src/braintrust/span_customizer.py @@ -2,7 +2,8 @@ import inspect import logging -from collections.abc import Callable, Sequence +import threading +from collections.abc import Callable, Iterable, Sequence from typing import Any from .db_fields import ( @@ -33,9 +34,15 @@ def on_span_export(self, data: SpanExportData) -> SpanExportData: 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 log a safe error, stop the hook chain, and drop that - outgoing record. Handle recoverable errors within the hook. Only - KeyboardInterrupt and SystemExit propagate. + synchronous returns stop the hook chain and strip the outgoing record. + Only original identity, routing/protocol fields, created, + span_attributes.type, and metrics.start/end survive, plus + context={"customizer_error": True}, replacing all original context. No + attachments from failed records are uploaded. One safe error is logged per + process, without payloads or exception details. Handle recoverable errors + within the hook. Only KeyboardInterrupt and SystemExit propagate, + including while snapshotting fields. Unreadable snapshot fields are + omitted. Dict subclasses are accepted and copied into plain dicts. Records may contain SDK ``Attachment`` objects: integrations convert inline media at capture time. Remove or replace them to prevent their upload; @@ -44,13 +51,18 @@ def on_span_export(self, data: SpanExportData) -> SpanExportData: 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. + fields may be absent and hooks may run before the span ends. Stripping a + record does not retract earlier exports or disable later records. Failed + records retain their original timing without inventing absent timestamps; + they bypass masking and are exported separately from other updates to the + same row. """ return data _span_customizers: tuple[SpanCustomizer, ...] = () +_failure_logged = False +_failure_log_lock = threading.Lock() _PROTECTED_FIELDS = frozenset( ( ID_FIELD, @@ -84,8 +96,11 @@ def on_span_export(self, data: SpanExportData) -> SpanExportData: continue try: data[field] = self._masking_function(data[field]) - except Exception as error: - # Fail closed per field, without leaking exception messages or stacks. + except (KeyboardInterrupt, SystemExit): + raise + except BaseException as error: + # Fail closed per field (including GeneratorExit and CancelledError), + # without leaking exception messages or stacks. message = f"ERROR: Failed to mask field '{field}' - {type(error).__name__}" if field in ("scores", "metrics"): del data[field] @@ -101,7 +116,7 @@ def set_span_customizers(customizers: Sequence[SpanCustomizer] | None) -> None: 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 or dropped result. There is no environment-variable registration. + 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. @@ -125,7 +140,7 @@ def _get_span_customizers() -> tuple[SpanCustomizer, ...]: def _validate_customizer(customizer: Any) -> None: - # Reject misconfiguration eagerly rather than dropping records at export time. + # Reject misconfiguration eagerly rather than stripping 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) @@ -154,24 +169,32 @@ def _restore_protocol_fields(data: SpanExportData, protected: SpanExportData) -> return restored +class _StrippedSpanExportData(dict[str, Any]): + """SDK-only marker: do not merge or mask a memoized failed record.""" + + def _customize_span_export( data: SpanExportData, customizers: Sequence[SpanCustomizer] | None = None -) -> SpanExportData | None: +) -> SpanExportData: if customizers is None: customizers = _span_customizers - if not customizers: + if not customizers or isinstance(data, _StrippedSpanExportData): return data - protected = None - for customizer in customizers: - try: + # Capture identity and timing before even looking up a hook: user-defined + # descriptors and earlier hooks can fail or mutate the record in place. + protected = _snapshot_fields(data, _PROTECTED_FIELDS) + # Non-content fields that keep a stripped initial record recognizable. + preserved = _snapshot_fields(data, ("created",)) + timing = _snapshot_fields(data.get("metrics"), ("start", "end")) + span_type = _snapshot_fields(data.get("span_attributes"), ("type",)) + try: + # Protocol containers may alias span state; don't expose those to hooks. + data = _restore_protocol_fields(data, protected) + for customizer in customizers: 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) result = hook(data) if inspect.isawaitable(result): # Do not execute asynchronous hooks or emit unawaited coroutine warnings. @@ -182,12 +205,52 @@ def _customize_span_export( 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 + return data + + +def _snapshot_fields(source: Any, keys: Iterable[str]) -> dict[str, Any]: + # Best-effort fallback; only process-control exceptions may escape. + snapshot: dict[str, Any] = {} + if not isinstance(source, dict): + return snapshot + for key in keys: + try: + if key not in source: + continue + value = source[key] 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 + continue + try: + snapshot[key] = _copy_protocol_value(value) + except (KeyboardInterrupt, SystemExit): + raise + except BaseException: + # e.g. RecursionError on a self-referential container. + snapshot[key] = value + return snapshot + + +def _log_customizer_failure() -> None: + global _failure_logged + with _failure_log_lock: + if _failure_logged: + return + _failure_logged = True + # Neither exception details nor record data are safe to include here. + logging.getLogger("braintrust").error("Span customizer failed; stripping outgoing span content.") diff --git a/py/src/braintrust/test_merge_row_batch.py b/py/src/braintrust/test_merge_row_batch.py index 745334d73..b3a5c11e7 100644 --- a/py/src/braintrust/test_merge_row_batch.py +++ b/py/src/braintrust/test_merge_row_batch.py @@ -1,7 +1,8 @@ import unittest from braintrust.db_fields import IS_MERGE_FIELD -from braintrust.merge_row_batch import batch_items, merge_row_batch +from braintrust.merge_row_batch import batch_items, merge_row_batch, merge_row_batch_waves +from braintrust.span_customizer import _StrippedSpanExportData class MergeRowBatchTest(unittest.TestCase): @@ -132,6 +133,75 @@ def test_skip_fields(self): ) +def _row(row_id, **fields): + return dict(experiment_id="e0", id=row_id, **fields) + + +def _stripped(row_id, **fields): + return _StrippedSpanExportData(_row(row_id, context={"customizer_error": True}, **fields)) + + +class MergeRowBatchWavesTest(unittest.TestCase): + def test_no_barriers_is_one_wave(self): + rows = [_row("x", a=1), _row("y", a=2), _row("x", b=3, **{IS_MERGE_FIELD: True})] + self.assertEqual(merge_row_batch_waves(rows), [[_row("x", a=1, b=3), _row("y", a=2)]]) + self.assertEqual(merge_row_batch_waves([]), []) + + def test_barrier_only_splits_its_own_row(self): + barrier = _stripped("x", **{IS_MERGE_FIELD: True}) + rows = [ + _row("x", a=1), + _row("y", a=1), + barrier, + _row("y", b=2, **{IS_MERGE_FIELD: True}), + _row("x", c=3, **{IS_MERGE_FIELD: True}), + _row("z", a=1), + _row("x", d=4, **{IS_MERGE_FIELD: True}), + ] + waves = merge_row_batch_waves(rows) + self.assertEqual( + waves, + [ + [_row("x", a=1), _row("y", a=1, b=2), _row("z", a=1)], + [barrier], + [_row("x", c=3, d=4, **{IS_MERGE_FIELD: True})], + ], + ) + self.assertIs(waves[1][0], barrier) + self.assertEqual(barrier, _stripped("x", **{IS_MERGE_FIELD: True})) + + def test_leading_and_consecutive_barriers(self): + # A barrier with no earlier rows for its key needs no extra wave, but + # consecutive barriers for one key must not share a wave. + first = _stripped("x", n=1) + second = _stripped("x", n=2, **{IS_MERGE_FIELD: True}) + rows = [first, _row("y", a=1), second, _row("x", a=1, **{IS_MERGE_FIELD: True}), _stripped("y")] + self.assertEqual( + merge_row_batch_waves(rows), + [ + [first, _row("y", a=1)], + [second, _stripped("y")], + [_row("x", a=1, **{IS_MERGE_FIELD: True})], + ], + ) + + def test_waves_never_repeat_keys_and_flatten_in_order(self): + rows = [ + _row("x", n=0), + _stripped("x"), + _row("y", n=0), + _stripped("y"), + _row("x", n=1, **{IS_MERGE_FIELD: True}), + _stripped("x"), + _row("y", n=1, **{IS_MERGE_FIELD: True}), + ] + waves = merge_row_batch_waves(rows) + for wave in waves: + self.assertEqual(len({row["id"] for row in wave}), len(wave)) + self.assertEqual(merge_row_batch(list(rows)), [row for wave in waves for row in wave]) + self.assertEqual(len(waves), 4) + + class BatchItemsTest(unittest.TestCase): def test_exact_byte_limit(self): self.assertEqual(batch_items(["aa", "bb", "c"], batch_max_num_bytes=4), [["aa", "bb"], ["c"]]) diff --git a/py/src/braintrust/test_span_customizer.py b/py/src/braintrust/test_span_customizer.py index e61fc07db..5c7528673 100644 --- a/py/src/braintrust/test_span_customizer.py +++ b/py/src/braintrust/test_span_customizer.py @@ -5,22 +5,25 @@ from collections import OrderedDict, defaultdict from concurrent.futures import ThreadPoolExecutor, TimeoutError from copy import deepcopy -from threading import Event +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from threading import Event, Thread from types import SimpleNamespace from unittest.mock import MagicMock +import exceptiongroup import pytest from braintrust import Attachment, SpanCustomizer, SpanExportData, auto_instrument, logger, set_span_customizers from braintrust.functions.stream import BraintrustJsonChunk, BraintrustStream from braintrust.integrations.utils import _resolved_attachment_from_bytes -from braintrust.span_customizer import _customize_span_export +from braintrust.span_customizer import _customize_span_export, _StrippedSpanExportData from braintrust.test_helpers import init_test_exp, with_memory_logger # noqa: F401 # type: ignore[reportUnusedImport] from braintrust.trace import LocalTrace from braintrust.util import LazyValue @pytest.fixture(autouse=True) -def reset_customizers(): +def reset_customizers(monkeypatch): + monkeypatch.setattr("braintrust.span_customizer._failure_logged", False) set_span_customizers(None) try: yield @@ -199,7 +202,7 @@ def on_span_export(self, data: SpanExportData) -> SpanExportData: @pytest.mark.parametrize("failure", ["exception", "none", "list", "coroutine", "awaitable"]) -def test_fail_closed_drops_record_stops_hooks_and_logs_safely(failure, caplog): +def test_fail_closed_strips_record_stops_hooks_and_logs_safely(failure, caplog): observed = [] coroutines = [] @@ -236,7 +239,7 @@ def on_span_export(self, data): with caplog.at_level(logging.ERROR, logger="braintrust"): result = _customize_span_export({"id": "private row", "output": "private payload"}) assert observed == [] - assert result is None + assert result == {"id": "private row", "context": {"customizer_error": True}} assert len(caplog.records) == 1 assert caplog.records[0].levelno == logging.ERROR assert caplog.records[0].exc_info is None @@ -244,6 +247,202 @@ def on_span_export(self, data): assert all(inspect.getcoroutinestate(c) == inspect.CORO_CLOSED for c in coroutines) +@pytest.mark.parametrize("timing", [{}, {"start": 12}, {"end": 15}, {"start": 12, "end": 15}]) +def test_failures_preserve_only_original_identity_protocol_and_timing(timing): + protected = { + "id": "row", + "span_id": "span", + "root_span_id": "root", + "span_parents": ["parent"], + "org_id": "org", + "experiment_id": "experiment", + "_is_merge": True, + "_merge_paths": [["metadata"]], + "_parent_id": "parent-row", + "_xact_id": "transaction", + } + record = { + **deepcopy(protected), + "created": "2026-01-01T00:00:00Z", + "metrics": {**timing, "tokens": 100, "time_to_first_token": 1}, + **{ + field: {"private": "payload"} + for field in ( + "name", + "input", + "output", + "expected", + "metadata", + "context", + "attributes", + "events", + "links", + "error", + "tags", + "attachments", + "scores", + ) + }, + "span_attributes": {"type": "llm", "name": "private", "private": "payload"}, + } + + class Mutate(SpanCustomizer): + def on_span_export(self, data): + data["span_parents"].clear() + data["span_attributes"]["type"] = "tool" + data["metrics"].update(start=999, end=1000) + data["context"] = {"private": "earlier hook"} + data["new_field"] = "private" + return data + + class Fail(SpanCustomizer): + def on_span_export(self, data): + data["id"] = "wrong" + data["metrics"].clear() + raise RuntimeError("private") + + result = _customize_span_export(record, [Mutate(), Fail()]) + expected = { + **protected, + "created": "2026-01-01T00:00:00Z", + "span_attributes": {"type": "llm"}, + "context": {"customizer_error": True}, + } + if timing: + expected["metrics"] = timing + assert result == expected + + +def test_hook_lookup_failure_preserves_original_timing(): + class Broken: + @property + def on_span_export(self): + raise GeneratorExit("private") + + assert _customize_span_export({"id": "row", "metrics": {"start": 12}, "input": "private"}, [Broken()]) == { + "id": "row", + "metrics": {"start": 12}, + "context": {"customizer_error": True}, + } + + +class _RaisingMetrics(dict): + def __contains__(self, key): + raise RuntimeError("private") + + +def test_uncopyable_protocol_fields_strip_instead_of_raising(caplog): + parents = ["parent"] + parents.append(parents) + record = {"id": "row", "span_parents": parents, "metrics": _RaisingMetrics(start=1), "input": "private"} + + with caplog.at_level(logging.ERROR, logger="braintrust"): + result = _customize_span_export(record, [_Valid()]) + + # Self-referential containers can't be copied for hooks, so the record is + # stripped, keeping identity as-is and omitting unreadable timing. + assert type(result) is _StrippedSpanExportData + assert result["id"] == "row" + assert result["span_parents"] is parents + assert "metrics" not in result + assert "input" not in result + assert len(caplog.records) == 1 + assert "private" not in caplog.text + + +@pytest.mark.parametrize("operation", ["lookup", "copy"]) +@pytest.mark.parametrize( + "error", [RuntimeError, BaseException, GeneratorExit, asyncio.CancelledError, KeyboardInterrupt, SystemExit] +) +def test_snapshot_failures_preserve_exception_boundaries(operation, error): + class UnreadableMetrics(dict): + def __contains__(self, key): + raise error("private") + + class UncopyableParents(list): + def __iter__(self): + raise error("private") + + record = {"id": "row", "input": "private"} + expected = {"id": "row", "context": {"customizer_error": True}} + if operation == "lookup": + record["metrics"] = UnreadableMetrics(start=1) + else: + parents = UncopyableParents(["parent"]) + record["span_parents"] = parents + expected["span_parents"] = parents + + customizers = [SimpleNamespace(on_span_export=lambda _: None)] + if issubclass(error, (KeyboardInterrupt, SystemExit)): + with pytest.raises(error): + _customize_span_export(record, customizers) + else: + result = _customize_span_export(record, customizers) + assert type(result) is _StrippedSpanExportData + assert result == expected + + +@pytest.mark.parametrize("error", [GeneratorExit, asyncio.CancelledError]) +def test_masking_base_exceptions_fail_closed_per_field(error): + def mask(value): + if value == "secret": + raise error("private") + return value + + result = _customize_span_export( + {"id": "row", "input": "secret", "output": "safe", "scores": {"a": "secret"}}, + [logger._MaskingCustomizer(mask)], + ) + assert type(result) is dict + assert result["input"] == f"ERROR: Failed to mask field 'input' - {error.__name__}" + assert result["output"] == "safe" + assert result["scores"] == {"a": "secret"} + + +@pytest.mark.parametrize("error", [KeyboardInterrupt, SystemExit]) +def test_masking_interpreter_exits_propagate(error): + def mask(_value): + raise error() + + with pytest.raises(error): + _customize_span_export({"id": "row", "input": "secret"}, [logger._MaskingCustomizer(mask)]) + + +def test_failure_logging_is_process_wide_thread_safe_and_private(caplog): + class Broken(SpanCustomizer): + def on_span_export(self, data): + raise ValueError("private exception") + + class Invalid(SpanCustomizer): + def on_span_export(self, data): + return None + + def export(index): + return _customize_span_export({"id": str(index), "input": "private"}, [Broken()]) + + with caplog.at_level(logging.ERROR, logger="braintrust"): + with ThreadPoolExecutor(max_workers=8) as executor: + results = list(executor.map(export, range(32))) + for customizer in (Invalid(), Broken()): + set_span_customizers([customizer]) + assert _customize_span_export({"id": "next"}) == { + "id": "next", + "context": {"customizer_error": True}, + } + set_span_customizers(None) + assert all(result["context"] == {"customizer_error": True} for result in results) + assert len(caplog.records) == 1 + assert caplog.records[0].exc_info is None + assert "private" not in caplog.text + + +def test_stripped_records_keep_marker_with_error_logging_disabled(caplog): + with caplog.at_level(logging.CRITICAL, logger="braintrust"): + result = _customize_span_export({"output": "private"}, [SimpleNamespace(on_span_export=lambda _: None)]) + assert result == {"context": {"customizer_error": True}} + assert caplog.records == [] + + @pytest.mark.parametrize("factory", [OrderedDict, lambda: defaultdict(list)], ids=["ordered-dict", "defaultdict"]) def test_dict_subclass_results_are_accepted_as_plain_dicts(factory): class Replace(SpanCustomizer): @@ -271,8 +470,8 @@ def on_span_export(self, data): @pytest.mark.parametrize("backend", ["memory", "http"]) -@pytest.mark.parametrize("include_healthy", [False, True], ids=["all-dropped", "mixed-batch"]) -def test_dropped_records_skip_attachments_masking_and_upload( +@pytest.mark.parametrize("include_healthy", [False, True], ids=["all-stripped", "mixed-batch"]) +def test_stripped_records_skip_attachments_and_continue_export( monkeypatch, with_memory_logger, test_logger, caplog, backend, include_healthy ): attachment = Attachment(data=b"private", filename="private.txt", content_type="text/plain") @@ -326,19 +525,20 @@ def mask(value): for item in pending: background.queue.put(item) background.flush() - # Reusing a dropped LazyValue must neither invoke nor log its failure again. - rows, attachments = background._unwrap_lazy_values(pending) + # Reusing a stripped LazyValue must neither invoke nor log its failure again. + waves, attachments = background._unwrap_lazy_values(pending) + rows = [row for wave in waves for row in wave] assert attachments == [] - if include_healthy: - sent = json.loads(connection.post.call_args.kwargs["data"])["rows"] - assert [row["id"] for row in sent] == [healthy.id] - else: - connection.post.assert_not_called() - - expected_ids = [healthy.id] if healthy else [] - assert [row["id"] for row in rows] == expected_ids - assert later_hooks == expected_ids - assert invocations == [failed.id, *expected_ids] + sent = [row for call in connection.post.call_args_list for row in json.loads(call.kwargs["data"])["rows"]] + assert [row["id"] for row in sent] == [failed.id, *([healthy.id] if healthy else [])] + + healthy_ids = [healthy.id] if healthy else [] + assert [row["id"] for row in rows] == [failed.id, *healthy_ids] + assert rows[0]["context"] == {"customizer_error": True} + assert "input" not in rows[0] + assert set(rows[0]["metrics"]) == {"start"} + assert later_hooks == healthy_ids + assert invocations == [failed.id, *healthy_ids] upload.assert_not_called() assert attachment not in masked assert bool(masked) == include_healthy @@ -448,11 +648,14 @@ def on_span_export(self, data): assert calls == ["reject", "accepted"] assert accepted == ["accepted"] - assert len(rows) == 1 + assert len(rows) == 2 assert rows[0]["id"] == span.id - assert rows[0]["input"] == "redacted" - assert rows[0]["output"] == "accepted" - assert rows[0]["_is_merge"] is True + assert rows[0]["context"] == {"customizer_error": True} + assert "metrics" not in rows[0] + assert "input" not in rows[0] + assert rows[1]["input"] == "redacted" + assert rows[1]["output"] == "accepted" + assert all(row["_is_merge"] is True for row in rows) assert with_memory_logger.upload_attempts == [] assert "private" not in caplog.text @@ -536,7 +739,9 @@ def mask(value): exported = next(row for row in json.loads(payloads[1])["rows"] if row["id"] == span.id) assert exported["input"] == "redacted" assert exported["output"] == "redacted output" - assert rejected.id not in {row["id"] for row in json.loads(payloads[1])["rows"]} + stripped = next(row for payload in payloads[1:] for row in json.loads(payload)["rows"] if row["id"] == rejected.id) + assert stripped["context"] == {"customizer_error": True} + assert "input" not in stripped upload.assert_not_called() assert len(caplog.records) == 1 assert "private" not in caplog.text @@ -564,8 +769,8 @@ def send(_path, *, data): return background, payloads -@pytest.mark.parametrize("error", [GeneratorExit, asyncio.CancelledError]) -def test_base_exception_hook_failures_drop_only_that_record( +@pytest.mark.parametrize("error", [BaseException, GeneratorExit, asyncio.CancelledError]) +def test_base_exception_hook_failures_strip_only_that_record( monkeypatch, with_memory_logger, test_logger, caplog, error ): invocations = [] @@ -586,10 +791,16 @@ def on_span_export(self, data): with caplog.at_level(logging.ERROR, logger="braintrust"): background.flush() - # The transport retry reuses the memoized drop instead of rerunning the hook. + # The transport retry reuses the memoized stripped record instead of rerunning the hook. + # Different row IDs need no barrier, so both records share one request. assert invocations == [failed.id, healthy.id] assert len(payloads) == 2 - assert [row["id"] for row in json.loads(payloads[1])["rows"]] == [healthy.id] + assert payloads[0] == payloads[1] + rows = json.loads(payloads[1])["rows"] + assert [row["id"] for row in rows] == [failed.id, healthy.id] + assert rows[0]["context"] == {"customizer_error": True} + assert "input" not in rows[0] + assert rows[1]["input"] == "safe" assert len(caplog.records) == 1 assert "private" not in caplog.text @@ -656,8 +867,8 @@ def export(target): if isinstance(target, logger._MemoryBackgroundLogger): target.log(*pending) return target.pop()[0] - rows, _ = target._unwrap_lazy_values(pending) - return rows[0] + waves, _ = target._unwrap_lazy_values(pending) + return waves[0][0] def mask(value): if value is None: @@ -723,7 +934,7 @@ def span_cache_experiment(with_memory_logger): experiment.state.span_cache.dispose() -def test_dropped_cache_record_is_released_and_later_records_are_independent( +def test_stripped_cache_record_is_released_and_later_records_are_independent( with_memory_logger, span_cache_experiment, caplog ): calls = [] @@ -740,13 +951,19 @@ def on_span_export(self, data): span = experiment.start_span(input="private payload") cache = experiment.state.span_cache with caplog.at_level(logging.ERROR, logger="braintrust"): - assert cache.get_by_root_span_id(span.root_span_id) is None - assert cache.get_by_root_span_id(span.root_span_id) is None - assert with_memory_logger.pop() == [] + for _ in range(2): + cached = _cached(experiment, span)[span.span_id] + assert cached.input is None + assert cached.output is None + # Only the SDK span type survives a failure, not its name. + assert cached.span_attributes == {"type": "eval"} + rows = with_memory_logger.pop() + assert rows[0]["context"] == {"customizer_error": True} + assert "input" not in rows[0] assert calls == [span.id] assert len(caplog.records) == 1 assert cache._pending_records == {} - assert span.root_span_id not in cache._root_span_index + assert span.root_span_id in cache._root_span_index assert "private" not in caplog.text span.log(output="safe") @@ -811,7 +1028,7 @@ async def unexpected_flush(): assert exported_manual["metadata"] == {"secret": "redacted"} -@pytest.mark.parametrize("fail", [False, True], ids=["customized", "dropped"]) +@pytest.mark.parametrize("fail", [False, True], ids=["customized", "stripped"]) def test_trace_cache_waits_for_inflight_customization(with_memory_logger, span_cache_experiment, fail): entered = Event() release = Event() @@ -863,3 +1080,196 @@ def test_span_cache_written_eagerly_without_customizers(with_memory_logger, span with experiment.start_span(name="instrumented", internal={"instrumentation": "test-auto"}) as span: span.log(output="output") assert _cached(experiment, span)[span.span_id].output == "output" + + +def test_http_export_smoke_preserves_failure_barriers_and_attachments( + monkeypatch, with_memory_logger, test_logger, caplog +): + requests = [] + invocations = [] + later_hooks = [] + + class Handler(BaseHTTPRequestHandler): + def do_POST(self): + payload = self.rfile.read(int(self.headers["Content-Length"])) + requests.append((self.path, payload)) + # Retry the failed record's first request using the real HTTP client. + self.send_response(503 if len(requests) == 2 else 200) + self.end_headers() + self.wfile.write(b"{}") + + def log_message(self, *_args): + pass + + class Customize(SpanCustomizer): + def on_span_export(self, data): + invocations.append(data["id"]) + if data.get("output") == "private": + data["metrics"]["start"] = 999 + data["metrics"]["end"] = 1000 + data["span_parents"].append("wrong") + raise ValueError("private exception") + return data + + class Observe(SpanCustomizer): + def on_span_export(self, data): + later_hooks.append(data["id"]) + return data + + failed_attachment = Attachment(data=b"private", filename="private.txt", content_type="text/plain") + healthy_attachment = Attachment(data=b"safe", filename="safe.txt", content_type="text/plain") + failed_upload = MagicMock() + healthy_upload = MagicMock(return_value={"upload_status": "done"}) + monkeypatch.setattr(failed_attachment, "upload", failed_upload) + monkeypatch.setattr(healthy_attachment, "upload", healthy_upload) + set_span_customizers([Customize(), Observe()]) + span = test_logger.start_span(input="before") + span.log(output="private", input=failed_attachment, metrics={"start": 12, "end": 15, "tokens": 42}) + span.log(output="after") + healthy = test_logger.start_span(input=healthy_attachment) + pending = list(with_memory_logger.logs) + with_memory_logger.logs.clear() + monkeypatch.setenv("BRAINTRUST_DISABLE_ATEXIT_FLUSH", "1") + monkeypatch.setattr(logger.time, "sleep", lambda _: None) + + with ThreadingHTTPServer(("127.0.0.1", 0), Handler) as server: + thread = Thread(target=server.serve_forever, daemon=True) + thread.start() + connection = logger.HTTPConnection(f"http://127.0.0.1:{server.server_port}") + background = logger._HTTPBackgroundLogger(LazyValue(lambda: connection, use_mutex=False)) + background.sync_flush = True + background.num_tries = 2 + background._max_request_size_result = {"max_request_size": 6_000_000, "can_use_overflow": False} + # Arbitrary masking must not change the SDK's diagnostic or timing. + background.set_masking_function(lambda value: "masked" if not isinstance(value, Attachment) else value) + try: + for item in pending: + background.queue.put(item) + with caplog.at_level(logging.ERROR, logger="braintrust"): + background.flush() + finally: + connection.session.close() + server.shutdown() + thread.join(timeout=5) + + assert len(requests) == 4 + assert all(path == "/logs3" for path, _payload in requests) + assert requests[1] == requests[2] + before, failed, after = [json.loads(requests[index][1])["rows"] for index in (0, 2, 3)] + # The barrier only orders updates to the failed row; other rows go first. + assert [row["id"] for row in before] == [span.id, healthy.id] + assert before[0]["input"] == "masked" + assert failed == [ + { + "id": span.id, + "span_id": span.span_id, + "root_span_id": span.root_span_id, + "span_parents": span.span_parents, + "project_id": "project", + "log_id": "g", + "_is_merge": True, + "metrics": {"start": 12, "end": 15}, + "context": {"customizer_error": True}, + } + ] + assert after == [{**after[0], "output": "masked"}] + assert after[0]["id"] == span.id + # Merging/masking adjacent successes must not poison the memoized failure. + assert pending[1].get() == failed[0] + assert invocations == [span.id, span.id, span.id, healthy.id] + assert later_hooks == [span.id, span.id, healthy.id] + failed_upload.assert_not_called() + healthy_upload.assert_called_once() + assert len(caplog.records) == 1 + assert "private" not in caplog.text + + +def _barrier_pending(): + class Fail(SpanCustomizer): + def on_span_export(self, data): + raise ValueError("private") + + row = {"project_id": "project", "log_id": "g"} + records = [ + {"id": "a", **row, "input": "start"}, + {"id": "b", **row, "input": "other"}, + _customize_span_export({"id": "a", **row, "_is_merge": True, "output": "private"}, [Fail()]), + {"id": "a", **row, "_is_merge": True, "output": "end"}, + ] + return [LazyValue(lambda record=record: record, use_mutex=False) for record in records] + + +def _recording_http_logger(monkeypatch, pending, fails=lambda _data: False): + monkeypatch.setenv("BRAINTRUST_DISABLE_ATEXIT_FLUSH", "1") + monkeypatch.setattr(logger.time, "sleep", lambda _: None) + attempts = [] + + def send(_path, *, data): + attempts.append([(row["id"], row.get("input") or row.get("output")) for row in json.loads(data)["rows"]]) + if fails(data): + raise ConnectionError("permanent failure") + return SimpleNamespace(ok=True) + + connection = MagicMock() + connection.post.side_effect = send + background = logger._HTTPBackgroundLogger(LazyValue(lambda: connection, use_mutex=False)) + background.num_tries = 2 + background.sync_flush = True + background._max_request_size_result = {"max_request_size": 6_000_000, "can_use_overflow": False} + for item in pending: + background.queue.put(item) + return background, attempts + + +class _ShutDownPool: + def submit(self, *_args, **_kwargs): + raise RuntimeError("cannot schedule new futures after shutdown") + + +_BARRIER_WAVES = [[("a", "start"), ("b", "other")], [("a", None)], [("a", "end")]] + + +@pytest.mark.parametrize("pool", ["running", "shut-down"]) +def test_flush_sends_waves_in_order_once_each(monkeypatch, pool): + if pool == "shut-down": + monkeypatch.setattr(logger, "HTTP_REQUEST_THREAD_POOL", _ShutDownPool()) + background, attempts = _recording_http_logger(monkeypatch, _barrier_pending()) + background.flush() + assert attempts == _BARRIER_WAVES + + +@pytest.mark.parametrize("pool", ["running", "shut-down"]) +def test_flush_runtime_error_in_request_does_not_resend_batches(monkeypatch, pool): + if pool == "shut-down": + monkeypatch.setattr(logger, "HTTP_REQUEST_THREAD_POOL", _ShutDownPool()) + background, attempts = _recording_http_logger(monkeypatch, _barrier_pending()) + submit = background._submit_logs_request + calls = 0 + + def submit_or_fail(batch, max_request_size_result): + nonlocal calls + calls += 1 + if calls == 2: + raise RuntimeError("request failed outside its retry loop") + return submit(batch, max_request_size_result) + + monkeypatch.setattr(background, "_submit_logs_request", submit_or_fail) + with pytest.raises(exceptiongroup.BaseExceptionGroup) as raised: + background.flush() + + assert [type(e) for e in raised.value.exceptions] == [RuntimeError] + assert calls == 3 + assert attempts == [_BARRIER_WAVES[0], _BARRIER_WAVES[2]] + + +def test_flush_failed_batch_does_not_skip_later_batches_or_waves(monkeypatch): + background, attempts = _recording_http_logger( + monkeypatch, _barrier_pending(), fails=lambda data: b'"input": "start"' in data + ) + with pytest.raises(exceptiongroup.BaseExceptionGroup) as raised: + background.flush(batch_size=1) + + assert len(raised.value.exceptions) == 1 + # Batches within a wave run concurrently; waves run in order. + assert sorted(attempts[:3]) == [[("a", "start")], [("a", "start")], [("b", "other")]] + assert attempts[3:] == [[("a", None)], [("a", "end")]]