litellm/tests/e2e/e2e_http.py

"""The ONLY module permitted to call ``requests.*``.

Enforced by tests/code_coverage_tests/check_e2e_no_raw_requests.py. Every request
body / query / header / response is a pydantic model; outcomes are a tagged union
(``Result[R]``) so callers ``match`` on them instead of catching exceptions.

``forward`` relays one provider-bound request for the provider edge and buffers
the whole body; ``forward_stream`` relays the same request but hands back the
response head plus a lazy iterator over the upstream's own transfer chunks, which
is what lets a recording keep the split points a streamed response arrived on.

Named e2e_http (not http) so it does not shadow the stdlib ``http`` package that
requests itself imports.
"""

from __future__ import annotations

import time
from collections.abc import Callable, Generator, Iterator, Mapping
from contextlib import contextmanager
from contextvars import ContextVar
from dataclasses import dataclass
from typing import Final, Generic, Literal, NewType, Protocol, TypeVar, cast

import pytest
import requests
from e2e_metadata import step
from pydantic import BaseModel, ConfigDict, Field

URL = NewType("URL", str)


class Headers(BaseModel):
    """Base for header models. Subclasses may alias to hyphenated header names
    (e.g. ``x-litellm-api-key``); serialization uses by_alias."""

    model_config = ConfigDict(populate_by_name=True)


class AuthHeaders(Headers):
    # litellm accepts either; set whichever the call needs, leave the other None.
    authorization: str | None = Field(default=None, repr=False)
    x_litellm_api_key: str | None = Field(default=None, alias="x-litellm-api-key", repr=False)


class AnthropicHeaders(AuthHeaders):
    """Auth plus the ``anthropic-version`` header the Anthropic-native
    /v1/messages and /v1/messages/count_tokens routes expect. It is harmless on
    the other providers the proxy routes to, and matches what Claude Code sends
    on its own internal calls."""

    anthropic_version: str = Field(default="2023-06-01", alias="anthropic-version")
    x_litellm_session_id: str | None = Field(default=None, serialization_alias="x-litellm-session-id")


class PartialBody(BaseModel):
    """A body for a partial-update route (absent = keep, null = clear): a field left
    unset is omitted from the wire, and a field set to None is sent as JSON null."""


class NoBody(BaseModel):
    """Empty body/query for routes that take none."""


class FileUploadForm(BaseModel):
    """Multipart form fields for POST /v1/files. The file bytes are passed
    separately; `model` is not here because the proxy reads it from the query
    (?model=) not the form."""

    purpose: str = "batch"
    target_model_names: str | None = None
    custom_llm_provider: str | None = None
    passthrough: bool | None = None


# ---------- Result types ----------

R = TypeVar("R", bound=BaseModel)


class Success(BaseModel, Generic[R]):
    kind: Literal["success"] = "success"
    status_code: int
    data: R


class NetworkError(BaseModel):
    kind: Literal["network"] = "network"
    message: str


class UnauthorizedError(BaseModel):
    kind: Literal["unauthorized"] = "unauthorized"
    # litellm 401s for key auth, model access, and tag routing alike, so keep the body to tell them apart.
    body: str = ""


class RateLimitedError(BaseModel):
    kind: Literal["rate_limited"] = "rate_limited"
    retry_after_seconds: int | None = None
    # keep the body so callers can tell limiter kinds apart.
    body: str = ""


class ValidationError(BaseModel):
    kind: Literal["validation"] = "validation"
    message: str


class UnknownApiError(BaseModel):
    kind: Literal["unknown"] = "unknown"
    status_code: int
    body: str


type Result[R: BaseModel] = (
    Success[R] | NetworkError | UnauthorizedError | RateLimitedError | ValidationError | UnknownApiError
)


class ProbeResult(BaseModel):
    """A route's reachability: status + body, no schema validation. Healthy ==
    route exists (not 404) and the handler did not crash (not 5xx)."""

    status_code: int
    body: str

    @property
    def healthy(self) -> bool:
        return 200 <= self.status_code < 500 and self.status_code != 404


class ExternalWrite(BaseModel):
    """Outcome of a call to a non-proxy API (an identity provider's admin API, a
    secret manager) that answers with a status, on create a Location header naming
    the new resource, and a body kept as text rather than parsed as JSON."""

    status_code: int
    location: str = ""
    body: str = ""

    @property
    def ok(self) -> bool:
        return 200 <= self.status_code < 300


class StreamingResponse(BaseModel):
    """Raw outcome for calls whose body is provider-native or streamed: status, the
    x-litellm-call-id header, the x-litellm-response-cost header (StandardLogging
    response_cost), the content-type (which tells streaming `text/event-stream` from
    non-streaming `application/json`), the response headers (lowercased names, e.g.
    the x-ratelimit-* pacing headers and retry-after on a 429), and the body.
    SpendLogs.request_id is the completion body id, not call_id. Used by passthrough
    and streaming, where one validated JSON model does not fit."""

    status_code: int
    call_id: str | None = None  # x-litellm-call-id header
    response_cost: float | None = None  # x-litellm-response-cost header
    content_type: str | None = None
    headers: dict[str, str] = {}
    body: str
    chunks: int = 0  # streamed events (0 for non-streaming)
    stream_events: list[str] = []
    stream_event_arrivals: list[float] = []
    # First in-stream error event, if any. A streamed call commits its HTTP 200
    # before the upstream completes, so upstream failures (e.g. insufficient
    # quota) arrive as SSE error events inside an otherwise-successful response;
    # the consumed body is elided, so this is the only place they surface.
    stream_error: str | None = None
    stream_done: bool = False
    stream_done_positions: tuple[int, ...] = ()

    @property
    def ok(self) -> bool:
        return 200 <= self.status_code < 300

    @property
    def is_streaming(self) -> bool:
        return "text/event-stream" in (self.content_type or "")


class BinaryStream(BaseModel):
    """Outcome of consuming a binary chunked response (e.g. TTS audio) as a stream.

    Unlike StreamingResponse, which line-splits an SSE text body, this iterates the
    raw bytes with iter_content and reports how many non-empty chunks arrived and
    the total byte count, so a caller can assert customer-observable streaming
    (multiple chunks, real bytes) without decoding the payload."""

    status_code: int
    content_type: str | None = None
    call_id: str | None = None
    transfer_encoding: str | None = None
    content_length: str | None = None
    error_body: str | None = None
    chunk_count: int = 0
    total_bytes: int = 0

    @property
    def ok(self) -> bool:
        return 200 <= self.status_code < 300

    @property
    def chunked(self) -> bool:
        return "chunked" in (self.transfer_encoding or "")


class SseResponse(Protocol):
    @property
    def status_code(self) -> int: ...

    @property
    def headers(self) -> Mapping[str, str]: ...

    @property
    def text(self) -> str: ...

    def iter_lines(self) -> Iterator[bytes]: ...


def _hdr(resp: SseResponse, name: str) -> str | None:
    value = resp.headers.get(name)
    return value if isinstance(value, str) else None


def unwrap[R: BaseModel](result: Result[R]) -> R:
    match result:
        case Success(data=data):
            return data
        case _:
            raise AssertionError(result)


def unwrap_status[R: BaseModel](result: Result[R], expected_status: int) -> R:
    """Like unwrap, but also pins the exact HTTP status the success came back on,
    for routes whose contract is a specific 2xx (e.g. 201 Created on a submission)."""
    match result:
        case Success(status_code=status_code, data=data) if status_code == expected_status:
            return data
        case Success(status_code=status_code):
            raise AssertionError(f"expected HTTP {expected_status}, got {status_code}")
        case _:
            raise AssertionError(result)


def is_ok[R: BaseModel](result: Result[R]) -> bool:
    match result:
        case Success():
            return True
        case _:
            return False


def require_successful_call(result: StreamingResponse) -> None:
    """A call that should have succeeded but didn't is a hard failure, never a skip:
    if the proxy can't make a call it's expected to, the test must fail."""
    if result.ok:
        return
    pytest.fail(f"upstream call failed (status {result.status_code}); body={result.body[:300]}")


def assert_client_error(result: StreamingResponse, context: str) -> None:
    assert 400 <= result.status_code < 500, f"{context}: expected 4xx, got {result.status_code}: {result.body[:300]}"


def assert_auth_denied(result: StreamingResponse, context: str) -> None:
    assert result.status_code in (401, 403), (
        f"{context}: expected 401/403, got {result.status_code}: {result.body[:300]}"
    )


def wire_body(json: BaseModel) -> dict[str, object]:
    if isinstance(json, PartialBody):
        return json.model_dump(by_alias=True, exclude_unset=True)
    return json.model_dump(by_alias=True, exclude_none=True)


def _flat(model: BaseModel) -> dict[str, str]:
    dumped: dict[str, object] = model.model_dump(by_alias=True, exclude_none=True)
    return {key: str(value) for key, value in dumped.items()}


def _headers(headers: BaseModel) -> dict[str, str]:
    return _flat(headers)


def _params(params: BaseModel | None) -> dict[str, str]:
    return _flat(params) if params is not None else {}


TRANSIENT_STATUSES: frozenset[int] = frozenset({529})
RETRY_ATTEMPTS: int = 3
_QUALIFICATION: Final[ContextVar[bool]] = ContextVar("e2e_qualification", default=False)


def retry_attempts(default: int) -> int:
    return 1 if _QUALIFICATION.get() else default


@contextmanager
def without_retries() -> Generator[None]:
    token: Final = _QUALIFICATION.set(True)
    try:
        yield
    finally:
        _QUALIFICATION.reset(token)


RETRY_BACKOFF_SECONDS: float = 0.5


class RetryableResponse(Protocol):
    status_code: int

    def close(self) -> None: ...


def request_with_retry[T: RetryableResponse](
    issue: Callable[[], T], *, sleep: Callable[[float], None] = time.sleep
) -> T:
    """Bounded retry on statuses attributable to the PROVIDER, never the proxy.

    The system under test is the proxy, so the transport may only absorb
    statuses the proxy itself cannot emit; today that is exactly 529, the
    Anthropic overloaded_error passed through verbatim (their own SDK retries
    it too). 500/502/503/504 stay first-class failures: at this layer a 5xx
    from the proxy is indistinguishable from one it relayed, and retrying them
    could mask an intermittently failing proxy. Widen the set only for a
    status litellm provably never originates, with an observed flake in hand.

    Also deliberately NOT retried: 429, because this suite asserts the proxy's
    own rate-limit and budget 429s; network errors and timeouts, because a
    hang should surface as a hang instead of doubling the wall clock. Every
    retry prints, so flakiness stays visible in the run log instead of
    vanishing into green."""
    for attempt in range(1, retry_attempts(RETRY_ATTEMPTS)):
        resp = issue()
        if resp.status_code not in TRANSIENT_STATUSES:
            return resp
        delay = RETRY_BACKOFF_SECONDS * (1 << (attempt - 1))
        print(
            f"e2e-http: transient {resp.status_code}; retry {attempt}/{RETRY_ATTEMPTS - 1} in {delay}s",
            flush=True,
        )
        resp.close()
        sleep(delay)
    return issue()


PROVIDER_RATE_LIMIT_MARKER: Final = "litellm.RateLimitError"
PROVIDER_RATE_LIMIT_ATTEMPTS: Final = 4
PROVIDER_RATE_LIMIT_BACKOFF_SECONDS: Final = 5.0


def tolerate_provider_rate_limit[R: BaseModel](
    issue: Callable[[], Result[R]],
    *,
    attempts: int = PROVIDER_RATE_LIMIT_ATTEMPTS,
    sleep: Callable[[float], None] = time.sleep,
) -> Result[R]:
    """Retry a call up to `attempts` times while the proxy relays the provider's own 429;
    any other outcome, the proxy's own 429 included, comes back at once."""
    for attempt in range(1, attempts):
        match issue():
            case RateLimitedError(body=body, retry_after_seconds=retry_after) if PROVIDER_RATE_LIMIT_MARKER in body:
                delay = retry_after or PROVIDER_RATE_LIMIT_BACKOFF_SECONDS * (1 << (attempt - 1))
                print(
                    f"e2e-http: provider rate limit relayed by the proxy; retry {attempt}/{attempts - 1} in {delay}s",
                    flush=True,
                )
                sleep(delay)
            case result:
                return result
    return issue()


class ProxyErrorDetail(BaseModel):
    message: str
    type: str
    code: str
    param: str | None = None


class _ProxyErrorBody(BaseModel):
    error: ProxyErrorDetail


def proxy_error(body: str) -> ProxyErrorDetail:
    """The proxy's own error envelope (`{"error": {message, type, param, code}}`) parsed off a rejected call."""
    return _ProxyErrorBody.model_validate_json(body).error


def relayed_provider_rate_limit(outcome: RateLimitedError) -> ProxyErrorDetail | None:
    """The provider's own 429 as the proxy relayed it, or None when the 429 is the proxy's own."""
    if PROVIDER_RATE_LIMIT_MARKER not in outcome.body:
        return None
    return _ProxyErrorBody.model_validate_json(outcome.body).error


class ClassifiableResponse(Protocol):
    """What classifying an outcome reads off a response. requests.Response satisfies
    it, and so does a fake, so the classification rules are testable on their own."""

    @property
    def status_code(self) -> int: ...

    @property
    def ok(self) -> bool: ...

    @property
    def text(self) -> str: ...

    @property
    def content(self) -> bytes: ...

    def json(self) -> object: ...


def classify[R: BaseModel](resp: ClassifiableResponse, response_type: type[R]) -> Result[R]:
    if resp.status_code == 401:
        return UnauthorizedError(body=resp.text)
    if resp.status_code == 429:
        return RateLimitedError(body=resp.text)
    if not resp.ok:
        return UnknownApiError(status_code=resp.status_code, body=resp.text)
    try:
        payload: Final[object] = resp.json() if resp.content else {}
        return Success(status_code=resp.status_code, data=response_type.model_validate(payload))
    except Exception as exc:  # noqa: BLE001 - any parse/validation failure is a value
        return ValidationError(message=str(exc))


def post[R: BaseModel](
    url: URL,
    *,
    headers: BaseModel,
    json: BaseModel,
    response_type: type[R],
    timeout: float = 30.0,
) -> Result[R]:
    try:
        resp = request_with_retry(
            lambda: requests.post(
                str(url),
                headers=_headers(headers),
                json=wire_body(json),
                timeout=timeout,
            )
        )
    except requests.RequestException as exc:
        return NetworkError(message=str(exc))
    return classify(resp, response_type)


def get[R: BaseModel](
    url: URL,
    *,
    headers: BaseModel,
    params: BaseModel,
    response_type: type[R],
    timeout: float = 30.0,
) -> Result[R]:
    try:
        resp = request_with_retry(
            lambda: requests.get(
                str(url),
                headers=_headers(headers),
                params=params.model_dump(by_alias=True, exclude_none=True),
                timeout=timeout,
            )
        )
    except requests.RequestException as exc:
        return NetworkError(message=str(exc))
    return classify(resp, response_type)


@step("GET the external URL {url}")
def get_external[R: BaseModel](
    url: str,
    *,
    response_type: type[R],
    headers: BaseModel | None = None,
    timeout: float = 30.0,
    ca_bundle: str | None = None,
) -> Result[R]:
    """GET an absolute URL outside the proxy (e.g. a public /.well-known document).
    Unlike the transport wrappers there is no proxy base url and no proxy auth; the
    response still gets the same tagged-union classification as every other call.
    `ca_bundle` is a PEM path for a server the default trust store does not cover."""
    try:
        resp = requests.get(
            url,
            headers={"Accept": "application/json", **(_headers(headers) if headers is not None else {})},
            timeout=timeout,
            verify=ca_bundle if ca_bundle is not None else True,
        )
    except requests.RequestException as exc:
        return NetworkError(message=str(exc))
    return classify(resp, response_type)


def post_external[R: BaseModel](
    url: str,
    *,
    json: BaseModel,
    response_type: type[R],
    headers: BaseModel | None = None,
    timeout: float = 30.0,
) -> Result[R]:
    """POST JSON to an absolute URL outside the proxy (e.g. a mock server's control
    route), classified the same way as every other call."""
    try:
        resp = requests.post(
            url,
            headers={"Accept": "application/json", **(_headers(headers) if headers is not None else {})},
            json=wire_body(json),
            timeout=timeout,
        )
    except requests.RequestException as exc:
        return NetworkError(message=str(exc))
    return classify(resp, response_type)


def post_form_external[R: BaseModel](
    url: str,
    *,
    form: BaseModel,
    response_type: type[R],
    headers: BaseModel | None = None,
    timeout: float = 30.0,
) -> Result[R]:
    """POST an absolute URL outside the proxy as `application/x-www-form-urlencoded`,
    the encoding OAuth 2 token endpoints take. Like get_external: no proxy base url,
    no proxy auth, and the same tagged-union classification as every other call."""
    try:
        resp = requests.post(
            url,
            data=_flat(form),
            headers=_headers(headers) if headers is not None else None,
            timeout=timeout,
        )
    except requests.RequestException as exc:
        return NetworkError(message=str(exc))
    return classify(resp, response_type)


def post_json_external(
    url: str,
    *,
    headers: BaseModel,
    json: BaseModel,
    timeout: float = 30.0,
) -> ExternalWrite:
    """POST an absolute URL outside the proxy under its own bearer, for an API that
    answers a create with a status and a Location header rather than a JSON body."""
    try:
        resp = requests.post(
            url,
            headers=_headers(headers),
            json=json.model_dump(by_alias=True, exclude_none=True),
            timeout=timeout,
        )
    except requests.RequestException as exc:
        return ExternalWrite(status_code=-1, body=str(exc))
    return ExternalWrite(
        status_code=resp.status_code,
        location=resp.headers.get("Location", ""),
        body=resp.text,
    )


def send_text_external(
    method: Literal["GET", "POST", "PATCH"],
    url: str,
    *,
    headers: BaseModel,
    content: str | None = None,
    timeout: float = 30.0,
) -> ExternalWrite:
    """Send an absolute URL outside the proxy a raw text body (or none) and keep the
    answer as text, for an API that takes and returns neither JSON nor forms: CyberArk
    Conjur takes a secret value or a YAML policy and returns a secret as its raw value."""
    try:
        resp = requests.request(
            method,
            url,
            headers=_headers(headers),
            data=content.encode() if content is not None else None,
            timeout=timeout,
        )
    except requests.RequestException as exc:
        return ExternalWrite(status_code=-1, body=str(exc))
    return ExternalWrite(status_code=resp.status_code, body=resp.text)


def delete_external(url: str, *, headers: BaseModel, timeout: float = 30.0) -> ExternalWrite:
    try:
        resp = requests.delete(url, headers=_headers(headers), timeout=timeout)
    except requests.RequestException as exc:
        return ExternalWrite(status_code=-1, body=str(exc))
    return ExternalWrite(status_code=resp.status_code, body=resp.text)


def delete[R: BaseModel](
    url: URL,
    *,
    headers: BaseModel,
    json: BaseModel,
    response_type: type[R],
    params: BaseModel | None = None,
    timeout: float = 30.0,
) -> Result[R]:
    try:
        resp = request_with_retry(
            lambda: requests.delete(
                str(url),
                headers=_headers(headers),
                json=wire_body(json),
                params=_params(params),
                timeout=timeout,
            )
        )
    except requests.RequestException as exc:
        return NetworkError(message=str(exc))
    return classify(resp, response_type)


def patch[R: BaseModel](
    url: URL,
    *,
    headers: BaseModel,
    json: BaseModel,
    response_type: type[R],
    timeout: float = 30.0,
) -> Result[R]:
    try:
        resp = request_with_retry(
            lambda: requests.patch(
                str(url),
                headers=_headers(headers),
                json=wire_body(json),
                timeout=timeout,
            )
        )
    except requests.RequestException as exc:
        return NetworkError(message=str(exc))
    return classify(resp, response_type)


def put[R: BaseModel](
    url: URL,
    *,
    headers: BaseModel,
    json: BaseModel,
    response_type: type[R],
    timeout: float = 30.0,
) -> Result[R]:
    try:
        resp = request_with_retry(
            lambda: requests.put(
                str(url),
                headers=_headers(headers),
                json=wire_body(json),
                timeout=timeout,
            )
        )
    except requests.RequestException as exc:
        return NetworkError(message=str(exc))
    return classify(resp, response_type)


def probe(url: URL, *, headers: BaseModel, params: BaseModel, timeout: float = 30.0) -> ProbeResult:
    try:
        resp = request_with_retry(
            lambda: requests.get(
                str(url),
                headers=_headers(headers),
                params=params.model_dump(by_alias=True, exclude_none=True),
                timeout=timeout,
            )
        )
    except requests.RequestException as exc:
        return ProbeResult(status_code=-1, body=str(exc))
    return ProbeResult(status_code=resp.status_code, body=resp.text)


def _parse_response_cost(resp: SseResponse) -> float | None:
    raw = _hdr(resp, "x-litellm-response-cost")
    if raw is None or raw == "":
        return None
    try:
        return float(raw)
    except ValueError:
        return None


_SSE_DATA_PREFIX: Final = b"data: "
_SSE_DONE: Final = "[DONE]"


def _is_stream_error_line(line: bytes) -> bool:
    return (
        line.startswith(b"event: error")
        or b'"type":"error"' in line
        or b'"type": "error"' in line
        or line.startswith(b'data: {"error"')
    )


def streaming_outcome(
    resp: SseResponse, stream: bool, *, sent_at: float, clock: Callable[[], float] = time.monotonic
) -> StreamingResponse:
    call_id: Final = _hdr(resp, "x-litellm-call-id")
    response_cost: Final = _parse_response_cost(resp)
    content_type: Final = _hdr(resp, "content-type")
    headers: Final = {name.lower(): value for name, value in resp.headers.items()}
    if not stream or not (200 <= resp.status_code < 300):
        return StreamingResponse(
            status_code=resp.status_code,
            call_id=call_id,
            response_cost=response_cost,
            content_type=content_type,
            headers=headers,
            body=resp.text,
        )
    stamped: Final = tuple((line, clock() - sent_at) for line in resp.iter_lines() if line)
    payloads: Final = tuple(
        (line.removeprefix(_SSE_DATA_PREFIX).decode(errors="replace"), arrived)
        for line, arrived in stamped
        if line.startswith(_SSE_DATA_PREFIX)
    )
    events: Final = tuple((payload, arrived) for payload, arrived in payloads if payload != _SSE_DONE)
    return StreamingResponse(
        status_code=resp.status_code,
        call_id=call_id,
        response_cost=response_cost,
        content_type=content_type,
        headers=headers,
        body="<streamed>",
        chunks=len(stamped),
        stream_events=[payload for payload, _ in events],
        stream_event_arrivals=[arrived for _, arrived in events],
        stream_done=any(payload == _SSE_DONE for payload, _ in payloads),
        stream_done_positions=tuple(index for index, (payload, _) in enumerate(payloads) if payload == _SSE_DONE),
        stream_error=next(
            (line.decode(errors="replace")[:300] for line, _ in stamped if _is_stream_error_line(line)),
            None,
        ),
    )


def send(
    url: URL,
    *,
    headers: BaseModel,
    json: BaseModel,
    params: BaseModel | None = None,
    stream: bool = False,
    timeout: float = 60.0,
) -> StreamingResponse:
    """Raw POST returning the unparsed HTTP outcome: status, full body, and the
    x-litellm-call-id header. For native/passthrough bodies and for calls judged by
    status rather than a typed JSON model (e.g. a budget block is a non-2xx). With
    ``stream=True`` the SSE body is consumed and its events counted instead."""
    sent_at: Final = time.monotonic()
    try:
        resp = request_with_retry(
            lambda: requests.post(
                str(url),
                headers=_headers(headers),
                params=_params(params),
                json=wire_body(json),
                stream=stream,
                timeout=timeout,
            )
        )
    except requests.RequestException as exc:
        return StreamingResponse(status_code=-1, body=str(exc))
    return streaming_outcome(resp, stream, sent_at=sent_at)


class AbandonedRequest(BaseModel):
    """A non-streaming request whose socket the client closed ``after`` seconds in,
    before the proxy had answered."""

    kind: Literal["abandoned"] = "abandoned"
    after: float


def abandon(
    url: URL, *, headers: BaseModel, json: BaseModel, after: float, connect_timeout: float = 10.0
) -> AbandonedRequest | StreamingResponse:
    """POST and close the connection ``after`` seconds if no response head has arrived
    by then; returns the response instead when the proxy answered first."""
    sent_at: Final = time.monotonic()
    session: Final = requests.Session()
    try:
        resp = session.post(
            str(url),
            headers=_headers(headers),
            json=wire_body(json),
            timeout=(connect_timeout, after),
        )
    except requests.exceptions.ReadTimeout:
        return AbandonedRequest(after=after)
    finally:
        session.close()
    return streaming_outcome(resp, False, sent_at=sent_at)


def stream(url: URL, *, headers: BaseModel, json: BaseModel, timeout: float = 60.0) -> StreamingResponse:
    """Streaming (SSE) call: consumes the stream counting events, and captures the
    x-litellm-call-id + content-type headers. Body is elided."""
    return send(url, headers=headers, json=json, stream=True, timeout=timeout)


def upload[R: BaseModel](
    url: URL,
    *,
    headers: BaseModel,
    form: BaseModel,
    filename: str,
    content: bytes,
    file_content_type: str = "application/jsonl",
    file_field: str = "file",
    params: BaseModel | None = None,
    response_type: type[R],
    timeout: float = 60.0,
) -> Result[R]:
    """Multipart POST for file-bearing routes (/v1/files, /v1/audio/transcriptions,
    /v1/images/edits). Form fields come from `form`, the file bytes are sent as the
    `file_field` part with `file_content_type`, and `params` carries any query
    routing (e.g. ?model=). requests sets the multipart Content-Type itself."""
    dumped: dict[str, object] = form.model_dump(by_alias=True, exclude_none=True)
    data = {key: str(value) for key, value in dumped.items()}
    try:
        resp = request_with_retry(
            lambda: requests.post(
                str(url),
                headers=_headers(headers),
                params=_params(params),
                data=data,
                files={file_field: (filename, content, file_content_type)},
                timeout=timeout,
            )
        )
    except requests.RequestException as exc:
        return NetworkError(message=str(exc))
    return classify(resp, response_type)


def stream_binary(
    url: URL,
    *,
    headers: BaseModel,
    json: BaseModel,
    chunk_size: int = 8192,
    timeout: float = 60.0,
) -> BinaryStream:
    """POST that consumes a binary chunked response (e.g. TTS audio) as a stream,
    counting non-empty chunks and total bytes with iter_content. A non-2xx status
    short-circuits with the counts left at zero so the caller can fail loudly."""
    try:
        resp = requests.post(
            str(url),
            headers=_headers(headers),
            json=wire_body(json),
            stream=True,
            timeout=timeout,
        )
    except requests.RequestException as exc:
        return BinaryStream(status_code=-1, error_body=str(exc)[:300])
    with resp:
        content_type = _hdr(resp, "content-type")
        call_id = _hdr(resp, "x-litellm-call-id")
        transfer_encoding = _hdr(resp, "transfer-encoding")
        content_length = _hdr(resp, "content-length")
        if not (200 <= resp.status_code < 300):
            return BinaryStream(
                status_code=resp.status_code,
                content_type=content_type,
                call_id=call_id,
                transfer_encoding=transfer_encoding,
                content_length=content_length,
                error_body=resp.text[:300],
            )
        raw_chunks = cast("Iterator[bytes]", resp.iter_content(chunk_size=chunk_size))
        chunks = tuple(chunk for chunk in raw_chunks if chunk)
        return BinaryStream(
            status_code=resp.status_code,
            content_type=content_type,
            call_id=call_id,
            transfer_encoding=transfer_encoding,
            content_length=content_length,
            chunk_count=len(chunks),
            total_bytes=sum(len(chunk) for chunk in chunks),
        )


def download(url: URL, *, headers: BaseModel, timeout: float = 60.0) -> StreamingResponse:
    """Raw GET for file content (/v1/files/{id}/content): provider-native bytes, no
    schema. Returns the decoded body and the x-litellm-call-id header."""
    try:
        resp = requests.get(str(url), headers=_headers(headers), timeout=timeout)
    except requests.RequestException as exc:
        return StreamingResponse(status_code=-1, body=str(exc))
    return StreamingResponse(
        status_code=resp.status_code,
        call_id=_hdr(resp, "x-litellm-call-id"),
        content_type=_hdr(resp, "content-type"),
        body=resp.text,
    )


class RawResponse(BaseModel):
    """A verbatim upstream HTTP response for the provider edge (provider_edge.py):
    status, lowercased headers, raw bytes. No Result classification because the
    edge relays provider errors to the proxy untouched."""

    status_code: int
    headers: dict[str, str]
    body: bytes


def forward(
    method: str,
    url: str,
    *,
    headers: dict[str, str],
    body: bytes | None,
    timeout: float = 60.0,
) -> RawResponse | NetworkError:
    """Relay one provider-bound request verbatim for the provider edge's record
    mode. No retries, no redirects, no schema: the proxy owns retry policy and
    the recorded bundle must hold exactly what the provider returned."""
    try:
        resp = requests.request(method, url, headers=headers, data=body, timeout=timeout, allow_redirects=False)
    except requests.RequestException as exc:
        return NetworkError(message=str(exc))
    return RawResponse(
        status_code=resp.status_code,
        headers={name.lower(): value for name, value in resp.headers.items()},
        body=resp.content,
    )


@dataclass(frozen=True, slots=True)
class StreamChunk:
    """One transfer chunk of a response body, exactly as the upstream framed it."""

    data: bytes


@dataclass(frozen=True, slots=True)
class StreamTruncation:
    """The body ended without its terminator, i.e. the upstream hung up mid-message.
    Always the last step, and ``reason`` is the transport's own description of it."""

    reason: str


type StreamStep = StreamChunk | StreamTruncation


@dataclass(frozen=True, slots=True)
class StreamHead:
    """An upstream response whose head has arrived and whose body has not been read.

    A dataclass rather than a BaseModel because it owns a live socket: ``steps`` is
    consumed once, in order, and closing it closes the underlying response."""

    status_code: int
    headers: dict[str, str]
    steps: Generator[StreamStep, None, None]


def _stream_steps(resp: requests.Response) -> Generator[StreamStep, None, None]:
    """The body as the upstream framed it, one step per transfer chunk.

    ``chunk_size=None`` is the whole point: urllib3 then returns exactly one piece
    per wire chunk, so the provider's split points survive into the recording. Any
    integer would re-slice the body into fixed-size pieces instead. Empty pieces are
    dropped because a zero-length chunk is the terminator on the wire, and a failure
    part way through becomes a final truncation step rather than an exception, since
    the chunks already delivered are exactly what makes a mid-stream failure
    different from a request that never streamed at all."""
    try:
        yield StreamChunk(b"")
        for piece in cast("Iterator[bytes]", resp.iter_content(chunk_size=None)):
            if piece:
                yield StreamChunk(data=piece)
    except requests.RequestException as exc:
        yield StreamTruncation(reason=str(exc))
    finally:
        resp.close()


def primed_steps(steps: Generator[StreamStep, None, None]) -> Generator[StreamStep, None, None]:
    first: Final = next(steps)
    assert isinstance(first, StreamChunk) and first.data == b""
    return steps


@dataclass(frozen=True, slots=True, repr=False)
class PreparedForward:
    request: requests.PreparedRequest
    url: str
    headers: dict[str, str]


def prepare_forward(
    method: str,
    url: str,
    headers: dict[str, str],
    body: bytes | None,
) -> PreparedForward | NetworkError:
    try:
        with requests.Session() as session:
            request: Final = session.prepare_request(requests.Request(method, url, headers=headers, data=body))
    except requests.RequestException as exc:
        return NetworkError(message=str(exc))
    assert request.url is not None
    return PreparedForward(request, request.url, dict(request.headers))


def forward_prepared_stream(prepared: PreparedForward, timeout: float) -> StreamHead | NetworkError:
    try:
        with requests.Session() as session:
            settings: Final = session.merge_environment_settings(prepared.url, {}, True, None, None)
            resp: Final = session.send(prepared.request, timeout=timeout, allow_redirects=False, **settings)
    except requests.RequestException as exc:
        return NetworkError(message=str(exc))
    return StreamHead(
        resp.status_code,
        {name.lower(): value for name, value in resp.headers.items()},
        primed_steps(_stream_steps(resp)),
    )


def open_stream(url: URL, *, headers: BaseModel, json: BaseModel, timeout: float = 60.0) -> StreamHead | NetworkError:
    """POST a streaming request and return the moment its response head arrives,
    leaving the body unread behind ``StreamHead.steps``. For a test that must keep
    one request in flight while it sends others: the head carries the routing
    headers (x-litellm-model-id), and draining ``steps`` ends the request."""
    return forward_stream(
        "POST",
        str(url),
        headers={**_headers(headers), "Content-Type": "application/json"},
        body=json.model_dump_json(by_alias=True, exclude_none=True).encode(),
        timeout=timeout,
    )


def forward_stream(
    method: str,
    url: str,
    *,
    headers: dict[str, str],
    body: bytes | None,
    timeout: float = 60.0,
) -> StreamHead | NetworkError:
    """Relay one provider-bound request for the provider edge and return as soon as
    the response head arrives, with the body left unread behind ``StreamHead.steps``.

    Same contract as ``forward`` otherwise: no retries, no redirects, no schema. A
    failure before the head arrives is still a ``NetworkError``; one raised while the
    body streams arrives as the last step. With ``stream=True`` the timeout bounds
    each socket read rather than the whole body, which is the right bound for a
    stream and strictly more permissive for a long generation."""
    try:
        resp = requests.request(
            method,
            url,
            headers=headers,
            data=body,
            timeout=timeout,
            allow_redirects=False,
            stream=True,
        )
    except requests.RequestException as exc:
        return NetworkError(message=str(exc))
    return StreamHead(
        status_code=resp.status_code,
        headers={name.lower(): value for name, value in resp.headers.items()},
        steps=primed_steps(_stream_steps(resp)),
    )