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, Mapping
from dataclasses import dataclass
from typing import Final, Generator, Generic, Iterator, Literal, NewType, Protocol, TypeVar, cast

import pytest
import requests
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 = None
    x_litellm_api_key: str | None = Field(default=None, alias="x-litellm-api-key")


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")


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


# ---------- 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
    # litellm overloads 429 for budget_exceeded too, so keep the body to tell them 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 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

    @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 _headers(headers: BaseModel) -> dict[str, str]:
    dumped: dict[str, object] = headers.model_dump(by_alias=True, exclude_none=True)
    return {key: str(value) for key, value in dumped.items()}


def _params(params: BaseModel | None) -> dict[str, str]:
    if params is None:
        return {}
    dumped: dict[str, object] = params.model_dump(by_alias=True, exclude_none=True)
    return {key: str(value) for key, value in dumped.items()}


TRANSIENT_STATUSES: frozenset[int] = frozenset({529})
RETRY_ATTEMPTS: int = 3
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):
        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()


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)


def get_external[R: BaseModel](
    url: str,
    *,
    response_type: type[R],
    timeout: float = 30.0,
) -> 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."""
    try:
        resp = requests.get(
            url,
            headers={"Accept": "application/json"},
            timeout=timeout,
        )
    except requests.RequestException as exc:
        return NetworkError(message=str(exc))
    return classify(resp, response_type)


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_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)


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:
        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 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=_stream_steps(resp),
    )