litellm/tests/e2e/e2e_http.py
mateo-berri ec47bbaaaa feat(e2e): record and replay streamed provider responses chunk-for-chunk
The record/replay harness stored a streamed provider response as one
buffered body, so a replayed stream arrived coalesced and the
/v1/messages streaming test could not be edge-wired. Keep each SSE
transfer chunk in the bundle in the order the provider sent it (a new
streamed response shape at BUNDLE_FORMAT_VERSION 4) so replay reproduces
the provider's split points, the recorded usage chunk keeps its
position, and a mid-stream upstream error replays as the same
mid-stream error rather than a clean body.

Resolves LIT-5742
2026-08-24 12:51:44 -07:00

774 lines
26 KiB
Python

"""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
from dataclasses import dataclass
from typing import 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 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] = []
# 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 "")
def _hdr(resp: requests.Response, 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 _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()
def _classify[R: BaseModel](
resp: requests.Response, 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:
return Success(status_code=resp.status_code, data=response_type.model_validate(resp.json()))
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=json.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[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=json.model_dump(by_alias=True, exclude_none=True),
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=json.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 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=json.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 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: requests.Response) -> 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
def _streaming_outcome(resp: requests.Response, stream: bool) -> StreamingResponse:
call_id = _hdr(resp, "x-litellm-call-id")
response_cost = _parse_response_cost(resp)
content_type = _hdr(resp, "content-type")
headers = {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,
)
lines = cast("Iterator[bytes]", resp.iter_lines())
chunks = 0
stream_error: str | None = None
stream_events: list[str] = []
stream_done = False
for line in lines:
if not line:
continue
chunks += 1
decoded_line = line.decode(errors="replace")
if decoded_line.startswith("data: "):
payload = decoded_line.removeprefix("data: ")
if payload == "[DONE]":
stream_done = True
else:
stream_events.append(payload)
if stream_error is None and (
line.startswith(b"event: error")
or b'"type":"error"' in line
or b'"type": "error"' in line
or line.startswith(b'data: {"error"')
):
stream_error = line.decode(errors="replace")[:300]
return StreamingResponse(
status_code=resp.status_code,
call_id=call_id,
response_cost=response_cost,
content_type=content_type,
headers=headers,
body="<streamed>",
chunks=chunks,
stream_events=stream_events,
stream_done=stream_done,
stream_error=stream_error,
)
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."""
try:
resp = request_with_retry(
lambda: requests.post(
str(url),
headers=_headers(headers),
params=_params(params),
json=json.model_dump(by_alias=True, exclude_none=True),
stream=stream,
timeout=timeout,
)
)
except requests.RequestException as exc:
return StreamingResponse(status_code=-1, body=str(exc))
return _streaming_outcome(resp, stream)
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=json.model_dump(by_alias=True, exclude_none=True),
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),
)