feat(api): add policy-aware HTTP transport

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",
+]

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