From 81ab3cb1f0c934439670954f6394fbbe15c50e3d Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Tue, 4 Aug 2026 15:48:08 -0400 Subject: [PATCH] feat(api): add policy-aware HTTP transport Add operation-level retry policies, bounded status and transport retries, and structured API errors while leaving legacy HTTPConnection call sites unchanged. Include real-socket coverage for retry safety, Retry-After handling, deadlines, custom transports, and session lifecycle. Next, add BraintrustClient bootstrap, endpoint routing, and the resource service facade on top of this transport. --- py/src/braintrust/api/__init__.py | 22 +- py/src/braintrust/api/_transport.py | 281 +++++++++++++++- py/src/braintrust/api/errors.py | 135 ++++++++ py/src/braintrust/api/policies.py | 55 ++++ py/src/braintrust/api/test_transport.py | 412 ++++++++++++++++++++++++ 5 files changed, 901 insertions(+), 4 deletions(-) create mode 100644 py/src/braintrust/api/errors.py create mode 100644 py/src/braintrust/api/policies.py create mode 100644 py/src/braintrust/api/test_transport.py diff --git a/py/src/braintrust/api/__init__.py b/py/src/braintrust/api/__init__.py index 57c4cb00..fef1d172 100644 --- a/py/src/braintrust/api/__init__.py +++ b/py/src/braintrust/api/__init__.py @@ -1,3 +1,23 @@ """Braintrust API client package.""" -__all__: list[str] = [] +from .errors import ( + BraintrustAPIError, + BraintrustHTTPError, + BraintrustResponseError, + BraintrustRetryExhaustedError, + BraintrustTransportError, + BraintrustTransportRetryExhaustedError, +) +from .policies import RetryMode, RetryPolicy + + +__all__ = [ + "BraintrustAPIError", + "BraintrustHTTPError", + "BraintrustResponseError", + "BraintrustRetryExhaustedError", + "BraintrustTransportError", + "BraintrustTransportRetryExhaustedError", + "RetryMode", + "RetryPolicy", +] diff --git a/py/src/braintrust/api/_transport.py b/py/src/braintrust/api/_transport.py index fb326742..6dd96c0f 100644 --- a/py/src/braintrust/api/_transport.py +++ b/py/src/braintrust/api/_transport.py @@ -1,9 +1,12 @@ -"""Legacy HTTP transport primitives used by the Braintrust SDK.""" +"""Legacy and policy-aware HTTP transport primitives for the Braintrust SDK.""" +import datetime +import logging import sys import time -from collections.abc import Mapping -from typing import Any +from collections.abc import Callable, Mapping +from email.utils import parsedate_to_datetime +from typing import Any, NoReturn import requests import urllib3 @@ -12,6 +15,17 @@ from ..env import BraintrustEnv from ..util import _urljoin, response_raise_for_status +from .errors import ( + BraintrustHTTPError, + BraintrustResponseError, + BraintrustRetryExhaustedError, + BraintrustTransportError, + BraintrustTransportRetryExhaustedError, +) +from .policies import RetryMode, RetryPolicy + + +logger = logging.getLogger(__name__) class RetryRequestExceptionsAdapter(HTTPAdapter): @@ -158,3 +172,264 @@ def patch_json(self, object_type: str, args: Mapping[str, Any] | None = None) -> resp = self.patch(f"/{object_type.lstrip('/')}", json=args) response_raise_for_status(resp) return resp.json() + + +class Transport: + """Policy-aware HTTP request engine for new API services. + + Existing ``HTTPConnection`` call sites deliberately do not use this class yet, + so adding endpoint policies does not change their behavior during migration. + Injected sessions and adapters own retries by default; callers may explicitly + enable the SDK loop when they know the injected transport does not retry. + """ + + def __init__( + self, + *, + session: requests.Session | None = None, + adapter: HTTPAdapter | None = None, + enable_sdk_retries: bool | None = None, + sleep: Callable[[float], None] = time.sleep, + monotonic: Callable[[], float] = time.monotonic, + wall_clock: Callable[[], float] = time.time, + ): + custom_transport = session is not None or adapter is not None + self._owns_session = session is None + self.session = session if session is not None else requests.Session() + self._sdk_retries_enabled = not custom_transport if enable_sdk_retries is None else enable_sdk_retries + if adapter is not None: + self.session.mount("http://", adapter) + self.session.mount("https://", adapter) + self._sleep = sleep + self._monotonic = monotonic + self._wall_clock = wall_clock + + def close(self) -> None: + if self._owns_session: + self.session.close() + + def __enter__(self) -> "Transport": + return self + + def __exit__(self, *_: Any) -> None: + self.close() + + def request( + self, + method: str, + url: str, + *, + params: Mapping[str, Any] | None = None, + json: Any = None, + data: Any = None, + headers: Mapping[str, str] | None = None, + retry_mode: RetryMode = RetryMode.NONE, + retry_policy: RetryPolicy | None = None, + stream: bool = False, + **kwargs: Any, + ) -> requests.Response: + method = method.upper() + policy = retry_policy or RetryPolicy.for_mode(retry_mode) + + replay_safe = retry_mode in (RetryMode.SAFE_READ, RetryMode.IDEMPOTENT_WRITE) + body_replayable = _request_body_is_replayable(data, kwargs.get("files")) + sdk_retries_enabled = replay_safe and body_replayable and self._sdk_retries_enabled and policy.max_attempts > 1 + if replay_safe and not body_replayable: + logger.debug("Disabling SDK retries for %s %s because its request body is not replayable", method, url) + max_attempts = policy.max_attempts if sdk_retries_enabled else 1 + started_at = self._monotonic() + + for attempt in range(1, max_attempts + 1): + remaining = self._remaining_budget(policy, started_at) + if remaining is not None and remaining <= 0: + # This is reachable only when an injected clock advances between + # attempts; ordinary retry delays are checked before sleeping. + error = BraintrustTransportRetryExhaustedError( + method=method, url=url, attempts=attempt - 1, retryable=True + ) + raise error + attempt_timeout = min(policy.timeout, remaining) if remaining is not None else policy.timeout + + try: + response = self.session.request( + method, + url, + params=params, + json=json, + data=data, + headers=headers, + timeout=attempt_timeout, + stream=stream, + **kwargs, + ) + except requests.exceptions.RequestException as exc: + if not _is_retryable_request_exception(exc): + error = BraintrustTransportError(method=method, url=url, attempts=attempt, retryable=False) + raise error from exc + if attempt >= max_attempts: + error_type = ( + BraintrustTransportRetryExhaustedError if sdk_retries_enabled else BraintrustTransportError + ) + error = error_type(method=method, url=url, attempts=attempt, retryable=sdk_retries_enabled) + raise error from exc + + delay = _retry_delay(policy, attempt, None) + if not self._can_wait(policy, started_at, delay): + error = BraintrustTransportRetryExhaustedError( + method=method, url=url, attempts=attempt, retryable=True + ) + raise error from exc + logger.debug( + "Retrying %s %s after transport error (attempt %d/%d) in %.3fs", + method, + url, + attempt, + max_attempts, + delay, + ) + self._sleep(delay) + continue + + if response.status_code < 400: + setattr(response, "_braintrust_attempts", attempt) + return response + + retry_after = _parse_retry_after(response.headers.get("Retry-After"), self._wall_clock()) + transient_status = response.status_code in policy.retryable_statuses + retryable = replay_safe and body_replayable and transient_status + if not (sdk_retries_enabled and transient_status): + _raise_http_error( + response, + method=method, + url=url, + attempts=attempt, + retryable=retryable, + retry_after=retry_after, + exhausted=False, + ) + + if attempt >= max_attempts: + _raise_http_error( + response, + method=method, + url=url, + attempts=attempt, + retryable=True, + retry_after=retry_after, + exhausted=True, + ) + + delay = _retry_delay(policy, attempt, retry_after) + if not self._can_wait(policy, started_at, delay): + _raise_http_error( + response, + method=method, + url=url, + attempts=attempt, + retryable=True, + retry_after=retry_after, + exhausted=True, + close_response=True, + ) + + logger.debug( + "Retrying %s %s after HTTP %d (attempt %d/%d) in %.3fs; Retry-After=%r parsed=%r", + method, + url, + response.status_code, + attempt, + max_attempts, + delay, + response.headers.get("Retry-After"), + retry_after, + ) + response.close() + self._sleep(delay) + + raise AssertionError("retry loop exited unexpectedly") + + def request_json(self, method: str, url: str, **kwargs: Any) -> Any: + response = self.request(method, url, **kwargs) + try: + return response.json() + except ValueError as exc: + error = BraintrustResponseError( + method=method.upper(), + url=response.url or url, + status_code=response.status_code, + response_body=response.text, + response_headers=response.headers, + attempts=getattr(response, "_braintrust_attempts", 1), + ) + raise error from exc + + def _remaining_budget(self, policy: RetryPolicy, started_at: float) -> float | None: + if policy.max_elapsed_time is None: + return None + return policy.max_elapsed_time - (self._monotonic() - started_at) + + def _can_wait(self, policy: RetryPolicy, started_at: float, delay: float) -> bool: + remaining = self._remaining_budget(policy, started_at) + return remaining is None or delay < remaining + + +def _retry_delay(policy: RetryPolicy, attempt: int, retry_after: float | None) -> float: + if retry_after is not None: + return retry_after + return min(policy.max_backoff, policy.backoff_factor * (2 ** (attempt - 1))) + + +def _request_body_is_replayable(data: Any, files: Any) -> bool: + return files is None and (data is None or isinstance(data, (bytes, str))) + + +def _is_retryable_request_exception(exc: requests.exceptions.RequestException) -> bool: + return isinstance(exc, (requests.exceptions.ConnectionError, requests.exceptions.Timeout)) and not isinstance( + exc, requests.exceptions.SSLError + ) + + +def _parse_retry_after(value: str | None, wall_time: float) -> float | None: + if value is None: + return None + value = value.strip() + if value.isdigit(): + return float(value) + try: + retry_at = parsedate_to_datetime(value) + if retry_at.tzinfo is None: + retry_at = retry_at.replace(tzinfo=datetime.timezone.utc) + return max(0.0, retry_at.timestamp() - wall_time) + except (TypeError, ValueError, OverflowError, OSError): + return None + + +def _raise_http_error( + response: requests.Response, + *, + method: str, + url: str, + attempts: int, + retryable: bool, + retry_after: float | None, + exhausted: bool, + close_response: bool = False, +) -> NoReturn: + error_type = BraintrustRetryExhaustedError if exhausted else BraintrustHTTPError + error = error_type( + method=method, + url=response.url or url, + status_code=response.status_code, + response_body=response.text, + response_headers=response.headers, + attempts=attempts, + retryable=retryable, + retry_after=retry_after, + ) + if close_response: + response.close() + try: + response.raise_for_status() + except requests.exceptions.HTTPError as cause: + raise error from cause + raise error diff --git a/py/src/braintrust/api/errors.py b/py/src/braintrust/api/errors.py new file mode 100644 index 00000000..d4fca223 --- /dev/null +++ b/py/src/braintrust/api/errors.py @@ -0,0 +1,135 @@ +"""Structured errors raised by the Braintrust API client.""" + +from collections.abc import Mapping +from types import MappingProxyType + +from ..util import AugmentedHTTPError + + +_RESPONSE_HEADER_ALLOWLIST = frozenset( + { + "content-type", + "retry-after", + "x-bt-internal-trace-id", + "x-vercel-id", + "request-id", + "x-request-id", + } +) +_REQUEST_ID_HEADERS = ("x-bt-internal-trace-id", "x-vercel-id", "x-request-id", "request-id") + + +class BraintrustAPIError(Exception): + """Base class for errors raised by the Braintrust API client.""" + + +class BraintrustHTTPError(BraintrustAPIError, AugmentedHTTPError): + """An HTTP error with status, request, and retry context.""" + + method: str + url: str + status_code: int + response_body: str + request_id: str | None + request_id_header: str | None + response_headers: Mapping[str, str] + attempts: int + retryable: bool + retry_after: float | None + retry_after_header: str | None + + def __init__( + self, + *, + method: str, + url: str, + status_code: int, + response_body: str, + response_headers: Mapping[str, str], + attempts: int, + retryable: bool, + retry_after: float | None = None, + ): + headers = { + normalized_name: value + for name, value in response_headers.items() + if (normalized_name := name.lower()) in _RESPONSE_HEADER_ALLOWLIST + } + request_id_header = next((name for name in _REQUEST_ID_HEADERS if headers.get(name)), None) + request_id = headers.get(request_id_header) if request_id_header else None + self.method = method + self.url = url + self.status_code = status_code + self.response_body = response_body + self.request_id = request_id + self.request_id_header = request_id_header + self.response_headers = MappingProxyType(headers) + self.attempts = attempts + self.retryable = retryable + self.retry_after = retry_after + self.retry_after_header = headers.get("retry-after") + message = f"{method} {url} failed with HTTP {status_code} after {attempts} attempt(s)" + if request_id: + message += f" (request ID: {request_id})" + if self.response_body: + message += f": {self.response_body}" + super().__init__(message) + + +class BraintrustRetryExhaustedError(BraintrustHTTPError): + """A retryable HTTP response exhausted its operation policy.""" + + +class BraintrustTransportError(BraintrustAPIError): + """A request failed before a usable HTTP response was received.""" + + method: str + url: str + attempts: int + retryable: bool + + def __init__(self, *, method: str, url: str, attempts: int, retryable: bool): + self.method = method + self.url = url + self.attempts = attempts + self.retryable = retryable + super().__init__(f"{method} {url} failed after {attempts} attempt(s) without an HTTP response") + + +class BraintrustTransportRetryExhaustedError(BraintrustTransportError): + """Transport exceptions exhausted the operation's retry policy.""" + + +class BraintrustResponseError(BraintrustAPIError): + """A successful HTTP response could not be decoded.""" + + method: str + url: str + status_code: int + response_body: str + response_headers: Mapping[str, str] + attempts: int + + def __init__( + self, + *, + method: str, + url: str, + status_code: int, + response_body: str, + response_headers: Mapping[str, str], + attempts: int, + ): + self.method = method + self.url = url + self.status_code = status_code + self.response_body = response_body + self.response_headers = MappingProxyType( + { + normalized_name: value + for name, value in response_headers.items() + if (normalized_name := name.lower()) in _RESPONSE_HEADER_ALLOWLIST + } + ) + self.attempts = attempts + super().__init__(f"Could not decode the response from {method} {url} as JSON") diff --git a/py/src/braintrust/api/policies.py b/py/src/braintrust/api/policies.py new file mode 100644 index 00000000..8bc9b57e --- /dev/null +++ b/py/src/braintrust/api/policies.py @@ -0,0 +1,55 @@ +"""Retry policies for Braintrust API operations.""" + +import enum +from dataclasses import dataclass + + +DEFAULT_RETRYABLE_STATUSES = frozenset({408, 429, 500, 502, 503, 504}) +DEFAULT_MAX_ATTEMPTS = 4 +DEFAULT_MAX_ELAPSED_TIME = 60.0 +DEFAULT_BACKOFF_FACTOR = 0.5 +DEFAULT_MAX_BACKOFF = 10.0 + + +class RetryMode(enum.Enum): + """The replay safety classification for an API operation.""" + + NONE = "none" + SAFE_READ = "safe_read" + IDEMPOTENT_WRITE = "idempotent_write" + LOG_INGESTION = "log_ingestion" + + +@dataclass(frozen=True) +class RetryPolicy: + """Bounded retry settings for one API operation.""" + + max_attempts: int = DEFAULT_MAX_ATTEMPTS + max_elapsed_time: float | None = DEFAULT_MAX_ELAPSED_TIME + timeout: float = 20.0 + retryable_statuses: frozenset[int] = DEFAULT_RETRYABLE_STATUSES + backoff_factor: float = DEFAULT_BACKOFF_FACTOR + max_backoff: float = DEFAULT_MAX_BACKOFF + + def __post_init__(self) -> None: + if self.max_attempts < 1: + raise ValueError("max_attempts must be at least 1") + if self.max_elapsed_time is not None and self.max_elapsed_time <= 0: + raise ValueError("max_elapsed_time must be positive") + if self.timeout <= 0: + raise ValueError("timeout must be positive") + if self.max_attempts > 1: + if self.max_elapsed_time is None: + raise ValueError("a retrying policy must have a maximum elapsed time") + if self.max_elapsed_time <= self.timeout: + raise ValueError("max_elapsed_time must be greater than its timeout") + if self.backoff_factor < 0 or self.max_backoff < 0: + raise ValueError("backoff settings cannot be negative") + + @classmethod + def for_mode(cls, mode: RetryMode) -> "RetryPolicy": + if mode in (RetryMode.NONE, RetryMode.LOG_INGESTION): + return cls(max_attempts=1, max_elapsed_time=None, timeout=60.0) + if mode in (RetryMode.SAFE_READ, RetryMode.IDEMPOTENT_WRITE): + return cls() + raise ValueError(f"Unknown retry mode: {mode!r}") diff --git a/py/src/braintrust/api/test_transport.py b/py/src/braintrust/api/test_transport.py new file mode 100644 index 00000000..96e99791 --- /dev/null +++ b/py/src/braintrust/api/test_transport.py @@ -0,0 +1,412 @@ +import contextlib +import datetime +import http.server +import io +import socketserver +import threading +import time +from email.utils import format_datetime + +import pytest +import requests +from braintrust.api import ( + BraintrustHTTPError, + BraintrustResponseError, + BraintrustRetryExhaustedError, + BraintrustTransportError, + BraintrustTransportRetryExhaustedError, + RetryMode, + RetryPolicy, +) +from braintrust.api._transport import Transport +from braintrust.util import AugmentedHTTPError +from requests.adapters import HTTPAdapter +from urllib3.util.retry import Retry + + +class FakeClock: + def __init__(self): + self.monotonic_time = 0.0 + self.wall_time = 1_800_000_000.0 + self.sleeps = [] + + def monotonic(self): + return self.monotonic_time + + def time(self): + return self.wall_time + + def sleep(self, delay): + self.sleeps.append(delay) + self.monotonic_time += delay + self.wall_time += delay + + +@contextlib.contextmanager +def scripted_server(script): + class ScriptedHandler(http.server.BaseHTTPRequestHandler): + request_count = 0 + requests = [] + + def log_message(self, format, *args): + pass + + def do_GET(self): + self._handle() + + def do_POST(self): + self._handle() + + def _handle(self): + request_number = type(self).request_count + type(self).request_count += 1 + content_length = int(self.headers.get("Content-Length", "0")) + body = self.rfile.read(content_length) if content_length else b"" + type(self).requests.append((self.command, self.path, body)) + action = script[min(request_number, len(script) - 1)] + + if action == "close": + self.connection.close() + return + + if action[0] == "sleep": + _, delay, status, headers, response_body = action + time.sleep(delay) + else: + status, headers, response_body = action + + self.send_response(status) + for name, value in headers.items(): + self.send_header(name, value) + self.send_header("Content-Length", str(len(response_body))) + self.end_headers() + try: + self.wfile.write(response_body) + except BrokenPipeError: + pass + + server = socketserver.ThreadingTCPServer(("127.0.0.1", 0), ScriptedHandler) + server.daemon_threads = True + thread = threading.Thread(target=server.serve_forever, kwargs={"poll_interval": 0.01}, daemon=True) + thread.start() + + try: + yield f"http://127.0.0.1:{server.server_address[1]}", ScriptedHandler + finally: + server.shutdown() + server.server_close() + + +def make_transport(clock=None, adapter=None): + clock = clock or FakeClock() + return Transport( + adapter=adapter, + sleep=clock.sleep, + monotonic=clock.monotonic, + wall_clock=clock.time, + ) + + +class TrackingAdapter(HTTPAdapter): + def __init__(self): + super().__init__() + self.close_count = 0 + + def close(self): + self.close_count += 1 + super().close() + + +class TrackingSession(requests.Session): + def __init__(self): + super().__init__() + self.close_count = 0 + + def close(self): + self.close_count += 1 + super().close() + + +def test_transport_closes_owned_session(): + adapter = TrackingAdapter() + + with Transport(adapter=adapter) as transport: + assert transport.session is not None + + assert adapter.close_count > 0 + + +def test_transport_does_not_close_injected_session(): + session = TrackingSession() + + with Transport(session=session) as transport: + assert transport.session is session + + assert session.close_count == 0 + session.close() + + +def test_retry_policy_defaults_and_validation(): + safe_read = RetryPolicy.for_mode(RetryMode.SAFE_READ) + assert safe_read.max_attempts == 4 + assert safe_read.max_elapsed_time == 60 + assert safe_read.timeout == 20 + assert safe_read.retryable_statuses == frozenset({408, 429, 500, 502, 503, 504}) + + for mode in (RetryMode.NONE, RetryMode.LOG_INGESTION): + policy = RetryPolicy.for_mode(mode) + assert policy.max_attempts == 1 + assert policy.timeout == 60 + + with pytest.raises(ValueError, match="greater than its timeout"): + RetryPolicy(max_attempts=2, max_elapsed_time=20, timeout=20) + + +def test_safe_logical_post_retries_transient_status(): + script = [ + (503, {}, b'{"error":"unavailable"}'), + (200, {"Content-Type": "application/json"}, b'{"ok":true}'), + ] + clock = FakeClock() + with scripted_server(script) as (url, handler): + result = make_transport(clock).request_json( + "POST", f"{url}/btql", json={"query": "select 1"}, retry_mode=RetryMode.SAFE_READ + ) + + assert result == {"ok": True} + assert handler.request_count == 2 + assert [request[0] for request in handler.requests] == ["POST", "POST"] + assert clock.sleeps == [0.5] + + +def test_none_and_log_ingestion_never_retry(): + for mode in (RetryMode.NONE, RetryMode.LOG_INGESTION): + with scripted_server([(429, {"Retry-After": "0"}, b"limited"), (200, {}, b"ok")]) as (url, handler): + with pytest.raises(BraintrustHTTPError) as exc_info: + make_transport().request("POST", url, retry_mode=mode) + + assert not isinstance(exc_info.value, BraintrustRetryExhaustedError) + assert exc_info.value.status_code == 429 + assert exc_info.value.attempts == 1 + assert handler.request_count == 1 + + +@pytest.mark.parametrize("status", [400, 401, 403, 501, 505, 507, 508, 511]) +def test_non_retryable_statuses_are_attempted_once(status): + with scripted_server([(status, {}, b"failed"), (200, {}, b"ok")]) as (url, handler): + with pytest.raises(BraintrustHTTPError) as exc_info: + make_transport().request("GET", url, retry_mode=RetryMode.SAFE_READ) + + assert not isinstance(exc_info.value, BraintrustRetryExhaustedError) + assert handler.request_count == 1 + + +def test_retry_after_delay_seconds_is_used_directly(): + clock = FakeClock() + with scripted_server([(429, {"Retry-After": "2"}, b"limited"), (200, {}, b"ok")]) as (url, handler): + response = make_transport(clock).request("GET", url, retry_mode=RetryMode.SAFE_READ) + + assert response.status_code == 200 + assert handler.request_count == 2 + assert clock.sleeps == [2] + + +def test_retry_after_http_date_uses_injected_wall_clock(): + clock = FakeClock() + retry_at = datetime.datetime.fromtimestamp(clock.wall_time + 5, tz=datetime.timezone.utc) + with scripted_server( + [(503, {"Retry-After": format_datetime(retry_at, usegmt=True)}, b"later"), (200, {}, b"ok")] + ) as (url, _): + make_transport(clock).request("GET", url, retry_mode=RetryMode.SAFE_READ) + + assert clock.sleeps == [5] + + +@pytest.mark.parametrize("retry_after", [None, "not-a-delay"]) +def test_absent_or_malformed_retry_after_uses_backoff(retry_after): + headers = {} if retry_after is None else {"Retry-After": retry_after} + clock = FakeClock() + with scripted_server([(429, headers, b"limited"), (200, {}, b"ok")]) as (url, _): + make_transport(clock).request("GET", url, retry_mode=RetryMode.SAFE_READ) + + assert clock.sleeps == [0.5] + + +def test_retry_after_larger_than_budget_raises_without_sleeping(): + clock = FakeClock() + policy = RetryPolicy(max_attempts=2, max_elapsed_time=2, timeout=1) + with scripted_server([(429, {"Retry-After": "5"}, b"limited")]) as (url, handler): + with pytest.raises(BraintrustRetryExhaustedError) as exc_info: + make_transport(clock).request("GET", url, retry_mode=RetryMode.SAFE_READ, retry_policy=policy) + + error = exc_info.value + assert error.status_code == 429 + assert error.attempts == 1 + assert error.retry_after == 5 + assert error.retry_after_header == "5" + assert handler.request_count == 1 + assert clock.sleeps == [] + + +def test_retry_exhaustion_preserves_structured_http_context_and_compatibility(): + policy = RetryPolicy(max_attempts=2, max_elapsed_time=2, timeout=1) + response_headers = { + "x-bt-internal-trace-id": "trace-id", + "x-vercel-id": "deployment-id", + "Set-Cookie": "secret=cookie", + "Content-Type": "application/json", + } + body = b'{"error":"failed","api_key":"super-secret"}' + with scripted_server([(500, response_headers, body)]) as (url, handler): + with pytest.raises(BraintrustRetryExhaustedError) as exc_info: + make_transport().request("GET", f"{url}/test", retry_mode=RetryMode.SAFE_READ, retry_policy=policy) + + error = exc_info.value + assert isinstance(error, AugmentedHTTPError) + assert isinstance(error.__cause__, requests.exceptions.HTTPError) + assert error.method == "GET" + assert error.url == f"{url}/test" + assert error.status_code == 500 + assert error.attempts == 2 + assert error.retryable is True + assert error.request_id == "trace-id" + assert error.request_id_header == "x-bt-internal-trace-id" + assert error.response_headers["content-type"] == "application/json" + assert "set-cookie" not in error.response_headers + assert error.response_body == body.decode() + assert str(error).endswith(body.decode()) + assert handler.request_count == 2 + + +def test_transport_exception_retries_and_final_error_chains_cause(): + policy = RetryPolicy(max_attempts=2, max_elapsed_time=2, timeout=1) + with scripted_server(["close"]) as (url, handler): + with pytest.raises(BraintrustTransportRetryExhaustedError) as exc_info: + make_transport().request("GET", url, retry_mode=RetryMode.SAFE_READ, retry_policy=policy) + + error = exc_info.value + assert isinstance(error, BraintrustTransportError) + assert error.attempts == 2 + assert error.retryable is True + assert isinstance(error.__cause__, requests.exceptions.RequestException) + assert handler.request_count == 2 + + +def test_none_wraps_transport_exception_without_retrying(): + with scripted_server(["close"]) as (url, handler): + with pytest.raises(BraintrustTransportError) as exc_info: + make_transport().request("POST", url, retry_mode=RetryMode.NONE) + + assert not isinstance(exc_info.value, BraintrustTransportRetryExhaustedError) + assert exc_info.value.attempts == 1 + assert handler.request_count == 1 + + +def test_deterministic_request_error_is_not_retried(): + clock = FakeClock() + with pytest.raises(BraintrustTransportError) as exc_info: + make_transport(clock).request("GET", "ftp://example.com", retry_mode=RetryMode.SAFE_READ) + + assert not isinstance(exc_info.value, BraintrustTransportRetryExhaustedError) + assert exc_info.value.attempts == 1 + assert exc_info.value.retryable is False + assert isinstance(exc_info.value.__cause__, requests.exceptions.InvalidSchema) + assert clock.sleeps == [] + + +def test_safe_read_timeout_leaves_room_for_retry(): + policy = RetryPolicy(max_attempts=2, max_elapsed_time=1, timeout=0.1, backoff_factor=0) + script = [ + ("sleep", 0.3, 200, {}, b"late"), + (200, {}, b"ok"), + ] + with scripted_server(script) as (url, handler): + response = Transport().request("GET", url, retry_mode=RetryMode.SAFE_READ, retry_policy=policy) + + assert response.status_code == 200 + assert handler.request_count == 2 + + +def test_invalid_json_is_not_retried_and_preserves_response_context(): + with scripted_server([(200, {"Content-Type": "application/json"}, b"not json")]) as (url, handler): + with pytest.raises(BraintrustResponseError) as exc_info: + make_transport().request_json("GET", url, retry_mode=RetryMode.SAFE_READ) + + assert exc_info.value.status_code == 200 + assert exc_info.value.response_body == "not json" + assert exc_info.value.attempts == 1 + assert handler.request_count == 1 + assert exc_info.value.__cause__ is not None + + +def test_file_like_request_body_disables_sdk_retries(): + body = io.BytesIO(b"payload") + with scripted_server([(500, {}, b"failed"), (200, {}, b"ok")]) as (url, handler): + with pytest.raises(BraintrustHTTPError) as exc_info: + make_transport().request("POST", url, data=body, retry_mode=RetryMode.SAFE_READ) + + assert not isinstance(exc_info.value, BraintrustRetryExhaustedError) + assert exc_info.value.retryable is False + assert handler.request_count == 1 + assert handler.requests[0][2] == b"payload" + + +def test_file_upload_disables_sdk_retries(): + files = {"file": ("payload.txt", io.BytesIO(b"payload"))} + with scripted_server([(500, {}, b"failed"), (200, {}, b"ok")]) as (url, handler): + with pytest.raises(BraintrustHTTPError) as exc_info: + make_transport().request("POST", url, files=files, retry_mode=RetryMode.IDEMPOTENT_WRITE) + + assert not isinstance(exc_info.value, BraintrustRetryExhaustedError) + assert handler.request_count == 1 + + +def retrying_adapter(): + return HTTPAdapter( + max_retries=Retry( + total=1, + status=1, + backoff_factor=0, + status_forcelist={500}, + allowed_methods=None, + raise_on_status=False, + ) + ) + + +def test_custom_adapter_disables_sdk_retry_loop(): + with scripted_server([(500, {}, b"failed")]) as (url, handler): + with pytest.raises(BraintrustHTTPError) as exc_info: + make_transport(adapter=retrying_adapter()).request("GET", url, retry_mode=RetryMode.SAFE_READ) + + assert not isinstance(exc_info.value, BraintrustRetryExhaustedError) + assert exc_info.value.attempts == 1 + assert handler.request_count == 2 # The adapter's two attempts, with no outer-loop multiplication. + + +def test_injected_session_disables_sdk_retry_loop(): + session = requests.Session() + session.mount("http://", retrying_adapter()) + with scripted_server([(500, {}, b"failed")]) as (url, handler): + with pytest.raises(BraintrustHTTPError) as exc_info: + Transport(session=session).request("GET", url, retry_mode=RetryMode.SAFE_READ) + + assert not isinstance(exc_info.value, BraintrustRetryExhaustedError) + assert exc_info.value.attempts == 1 + assert handler.request_count == 2 + + +def test_non_retrying_custom_adapter_can_delegate_retries_to_sdk(): + clock = FakeClock() + transport = Transport( + adapter=HTTPAdapter(), + enable_sdk_retries=True, + sleep=clock.sleep, + monotonic=clock.monotonic, + wall_clock=clock.time, + ) + with scripted_server([(500, {}, b"failed"), (200, {}, b"ok")]) as (url, handler): + response = transport.request("GET", url, retry_mode=RetryMode.SAFE_READ) + + assert response.status_code == 200 + assert handler.request_count == 2