diff --git a/tests/e2e/logging/callback_config.py b/tests/e2e/logging/callback_config.py new file mode 100644 index 00000000000..dc033a14b20 --- /dev/null +++ b/tests/e2e/logging/callback_config.py @@ -0,0 +1,224 @@ +"""Turn a logging integration on for the duration of a test, then put the proxy +back the way it was found. + +A litellm logging integration is a process-wide callback, not a per-request +option, so a delivery test can only assert against a proxy that has the +integration registered for the event it cares about. Rather than depend on how +the proxy under test happened to be launched, a test declares what it needs: +``callback_enabled`` reads the registered callbacks back from +/get/config/callbacks, registers the missing one through /config/update, waits +until the proxy reports it, and unregisters exactly what it registered on the +way out. A proxy that already ships the integration is left untouched, so this +is a no-op wherever the destination is already wired into the deployment. + +Success and failure are separate registrations in litellm and are removed by +separate routes: /config/update unions into success_callback (callbacks are +additive there, so a shorter list cannot remove one) and /config/callback/delete +is the only way back out, while failure_callback is written wholesale. Every +write here is therefore a read-modify-write of the list as it stands at that +moment, adding or dropping just this caller's entry. Restoring a snapshot taken +at setup would unregister whatever a concurrently running test had registered in +the meantime, which on a shared proxy is a real way to break someone else's run. + +A read-modify-write is still not atomic, and nothing available here can make it +so: litellm exposes a whole-list write and a server-side read-remove-write, with +no per-entry update and no conditional write, so there is no compare-and-set to +build ownership on, and a reference count would be one more read-then-write over +the same shared state. What is available is detection. Every write re-reads the +list afterwards and fails, naming the entries, if anything that was registered +before it and belongs to somebody else has gone, which turns a silent change to +a shared proxy's logging configuration into a diagnosable failure. Entries that +appear only after a write are somebody else's later registration, not damage, +and are left alone. + +All of this assumes one sequential pytest session per proxy, which is what the +rest of the harness assumes anyway; the session-scoped proxy fixture and the +destructive spend-log truncate in the root conftest are both single-session by +construction. Run against a shared proxy under xdist the detection still fires, +but it reports the collision rather than preventing it. + +/get/config/callbacks also returns each integration's resolved credentials. The +model here keeps only the name and the event type, so no secret is ever parsed +into a test. +""" + +from __future__ import annotations + +import time +from collections.abc import Iterator +from typing import Literal, assert_never + +import pytest +from pydantic import BaseModel, ConfigDict + +from e2e_config import POLL_INTERVAL, POLL_TIMEOUT +from e2e_http import NoBody, unwrap +from logging_client import LoggingClient + +type CallbackEvent = Literal["success", "failure"] + + +def _matching_types(event: CallbackEvent) -> frozenset[str]: + """The /get/config/callbacks ``type`` values that register for ``event``: + litellm_settings.callbacks registers for both and is reported as + success_and_failure.""" + match event: + case "success": + return frozenset({"success", "success_and_failure"}) + case "failure": + return frozenset({"failure", "success_and_failure"}) + case _: + assert_never(event) + + +class _ConfiguredCallback(BaseModel): + """One /get/config/callbacks entry. The route also returns the + integration's resolved credentials under ``variables``; leaving them + unmodelled keeps them out of the test process.""" + + model_config = ConfigDict(extra="ignore") + + name: str + type: str + + +class _ConfiguredCallbacks(BaseModel): + model_config = ConfigDict(extra="ignore") + + callbacks: list[_ConfiguredCallback] = [] + + +class _CallbackLists(BaseModel): + """The litellm_settings slice /config/update needs. Only the list for the + event being changed is sent; the other stays None and is dropped.""" + + success_callback: list[str] | None = None + failure_callback: list[str] | None = None + + +class _ConfigUpdateBody(BaseModel): + litellm_settings: _CallbackLists + + +class _ConfigUpdateAck(BaseModel): + model_config = ConfigDict(extra="ignore") + + message: str + + +class _CallbackDeleteBody(BaseModel): + callback_name: str + + +class _CallbackDeleteAck(BaseModel): + model_config = ConfigDict(extra="ignore") + + removed_callback: str + + +def registered_callbacks(client: LoggingClient, event: CallbackEvent) -> tuple[str, ...]: + """The integrations the proxy reports as registered for ``event``.""" + reported = unwrap( + client.proxy.transport.get( + "/get/config/callbacks", + headers=client.proxy.transport.master, + params=NoBody(), + response_type=_ConfiguredCallbacks, + ) + ) + wanted = _matching_types(event) + return tuple(callback.name for callback in reported.callbacks if callback.type in wanted) + + +def _write_callbacks(client: LoggingClient, event: CallbackEvent, names: tuple[str, ...]) -> None: + settings = ( + _CallbackLists(success_callback=list(names)) + if event == "success" + else _CallbackLists(failure_callback=list(names)) + ) + ack = unwrap( + client.proxy.transport.post( + "/config/update", + headers=client.proxy.transport.master, + json=_ConfigUpdateBody(litellm_settings=settings), + response_type=_ConfigUpdateAck, + ) + ) + assert "success" in ack.message.lower(), f"POST /config/update must acknowledge the write; got {ack.message!r}" + + +#: Why a vanished entry is worth stopping for, appended to both loss reports. +_LOST_ENTRIES = ( + "the callback list has no atomic per-entry update, so a registration made between this " + "read-modify-write's read and its write is overwritten rather than merged; those entries are " + "gone from the proxy's configuration and anything relying on them is now logged differently" +) + + +def _register(client: LoggingClient, name: str, event: CallbackEvent) -> None: + before = registered_callbacks(client, event) + _write_callbacks(client, event, before + (name,)) + deadline = time.monotonic() + POLL_TIMEOUT + while True: + settled = frozenset(registered_callbacks(client, event)) + if name in settled: + lost = frozenset(before) - settled + if lost: + pytest.fail( + f"registering {name!r} for {event} calls dropped {sorted(lost)}: {_LOST_ENTRIES}" + ) + return + if time.monotonic() >= deadline: + pytest.fail(f"the proxy never reported the {name!r} callback registered for {event} calls") + time.sleep(POLL_INTERVAL) + + +def _unregister(client: LoggingClient, name: str, event: CallbackEvent) -> frozenset[str]: + """Drop only this test's own entry, and report anything else that vanished. + + The list is re-read here and the name filtered out of what is registered + *now*, never replaced with the snapshot taken at setup: the proxy is shared, + so writing back a stale whole list would silently unregister a callback + something else added in between. The re-read afterwards catches the case that + remains, where something registered inside the window between this read and + this write and the write overwrote it. Returns those lost names rather than + failing, so the caller can finish unregistering every event before it + reports; leaving half the events registered would be a worse outcome than + the loss being reported one moment later.""" + before = registered_callbacks(client, event) + if event == "failure": + _write_callbacks(client, event, tuple(n for n in before if n != name)) + else: + _ = client.proxy.transport.post( + "/config/callback/delete", + headers=client.proxy.transport.master, + json=_CallbackDeleteBody(callback_name=name), + response_type=_CallbackDeleteAck, + ) + return (frozenset(before) - {name}) - frozenset(registered_callbacks(client, event)) + + +def callback_enabled(client: LoggingClient, name: str, *, events: tuple[CallbackEvent, ...]) -> Iterator[None]: + """Guarantee ``name`` is registered for every event in ``events`` while the + generator is suspended, and no longer registered for the ones this call + added once it resumes. Every write is a read-modify-write of the live list + followed by a re-read, so concurrent registrations by other tests survive, + and one lost to the gap between the read and the write is reported instead + of disappearing. Registration happens inside the try, so a failure part way + through a multi-event registration still unregisters what it managed to add. + Drive it from a fixture with ``yield from``.""" + added: tuple[CallbackEvent, ...] = tuple( + event for event in events if name not in registered_callbacks(client, event) + ) + try: + for event in added: + _register(client, name, event) + yield + finally: + lost = tuple( + f"{event}/{missing}" + for event in added + for missing in sorted(_unregister(client, name, event)) + ) + if lost: + pytest.fail(f"unregistering {name!r} dropped {list(lost)}: {_LOST_ENTRIES}") diff --git a/tests/e2e/logging/s3_reader.py b/tests/e2e/logging/s3_reader.py new file mode 100644 index 00000000000..8cc4409013e --- /dev/null +++ b/tests/e2e/logging/s3_reader.py @@ -0,0 +1,197 @@ +"""Read-back for the S3 logging tests against the real bucket the proxy writes to. + +Delivery is judged on what actually landed in S3: the proxy uploads with its own +AWS credentials exactly as in production (no local minio, no endpoint override), +and the tests fetch the object back with the official AWS SDK. Which bucket the +proxy is configured to write to is not discoverable over any proxy route, so the +test process is told through E2E_S3_LOG_BUCKET (plus E2E_S3_LOG_REGION, +E2E_S3_LOG_PATH when the deployment sets an s3_path prefix, and +E2E_S3_LOG_PROFILE to read back through a named AWS profile). A missing bucket is +a hard failure, never an empty result. + +An object's key ends in the id litellm assigned the call - the completion id for +a successful call, the litellm call id (the x-litellm-call-id response header) +for a failed one - so a test locates its own object exactly, without scanning +anyone else's, and deletes just that key on teardown. + +boto3 is the AWS SDK's own client and is deliberately not routed through +``e2e_http``: hand-signing SigV4 to reach S3 would be a worse test than using the +vendor client. Its s3 client is untyped, so it is confined to this module behind +the ``_S3Client`` protocol, its responses are narrowed here, and every payload is +validated into a pydantic model before a test sees it. +""" + +from __future__ import annotations + +import os +import time +from collections.abc import Iterator +from dataclasses import dataclass +from typing import Protocol, cast + +import boto3 +import pytest +from pydantic import BaseModel, ConfigDict, ValidationError + +from e2e_config import POLL_INTERVAL, POLL_TIMEOUT + +S3_LOG_BUCKET = os.environ.get("E2E_S3_LOG_BUCKET", "").strip() +S3_LOG_REGION = os.environ.get("E2E_S3_LOG_REGION", "us-east-1").strip() +#: The deployment's litellm_settings.s3_callback_params.s3_path, if it sets one. +S3_LOG_PATH = os.environ.get("E2E_S3_LOG_PATH", "").strip() +#: Named AWS profile for the read-back; empty uses boto3's default credential chain. +S3_LOG_PROFILE = os.environ.get("E2E_S3_LOG_PROFILE", "").strip() + + +class _S3Body(Protocol): + def read(self) -> bytes: ... + + +class _S3Client(Protocol): + """The three S3 calls this module makes.""" + + def list_objects_v2(self, **kwargs: str) -> object: ... + + def get_object(self, *, Bucket: str, Key: str) -> object: ... + + def delete_object(self, *, Bucket: str, Key: str) -> object: ... + + +class S3LogError(BaseModel): + """The error_information block of a failed call's payload.""" + + model_config = ConfigDict(extra="ignore") + + error_class: str = "" + error_code: str = "" + llm_provider: str = "" + + +class S3LogMessage(BaseModel): + model_config = ConfigDict(extra="ignore") + + role: str + content: str + + +class S3LogRecord(BaseModel): + """The StandardLoggingPayload as the proxy wrote it into the bucket. Only the + fields the delivery tests pin are modelled.""" + + model_config = ConfigDict(extra="ignore") + + id: str + status: str + call_type: str + model_group: str + response_cost: float + total_tokens: int + messages: tuple[S3LogMessage, ...] = () + error_information: S3LogError = S3LogError() + error_str: str | None = None + + +@dataclass(frozen=True, slots=True) +class S3LogObject: + key: str + record: S3LogRecord + + +@dataclass(frozen=True, slots=True) +class _ListPage: + keys: tuple[str, ...] + next_token: str | None + + +def _s3_client() -> _S3Client: + session = boto3.Session(profile_name=S3_LOG_PROFILE) if S3_LOG_PROFILE else boto3.Session() + client = session.client("s3", region_name=S3_LOG_REGION) # pyright: ignore[reportUnknownMemberType] # boto3 ships no types for the per-service clients; _S3Client pins the calls made here + return cast(_S3Client, client) + + +def _response(raw: object, call: str) -> dict[str, object]: + if not isinstance(raw, dict): + pytest.fail(f"S3 {call} returned {type(raw).__name__}, not a response mapping") + return cast(dict[str, object], raw) + + +def _list_page(raw: object) -> _ListPage: + listing = _response(raw, "list_objects_v2") + contents = listing.get("Contents", []) + if not isinstance(contents, list): + pytest.fail(f"S3 list_objects_v2 Contents was {type(contents).__name__}, not a list") + keys = tuple( + key + for entry in cast(list[object], contents) + if isinstance(entry, dict) and isinstance(key := cast(dict[str, object], entry).get("Key"), str) + ) + token = listing.get("NextContinuationToken") + truncated = listing.get("IsTruncated") is True + return _ListPage(keys=keys, next_token=token if truncated and isinstance(token, str) else None) + + +def _body(raw: object) -> bytes: + stream = _response(raw, "get_object").get("Body") + if not hasattr(stream, "read"): + pytest.fail(f"S3 get_object returned a Body of {type(stream).__name__}, which is not readable") + return cast(_S3Body, stream).read() + + +@dataclass(frozen=True, slots=True) +class S3LogsReader: + bucket: str + path: str + + def _keys(self, client: _S3Client) -> Iterator[str]: + page = _list_page(client.list_objects_v2(Bucket=self.bucket, Prefix=self.path)) + while True: + yield from page.keys + if page.next_token is None: + return + page = _list_page( + client.list_objects_v2(Bucket=self.bucket, Prefix=self.path, ContinuationToken=page.next_token) + ) + + def _record(self, client: _S3Client, key: str) -> S3LogRecord: + body = _body(client.get_object(Bucket=self.bucket, Key=key)) + try: + return S3LogRecord.model_validate_json(body) + except ValidationError as exc: + pytest.fail(f"the object at {key} is not a StandardLoggingPayload: {exc}") + + def objects_for_call(self, call_id: str) -> tuple[S3LogObject, ...]: + """Every object the proxy wrote for the call litellm identified as + ``call_id``. More than one is a duplicate-delivery bug, so this never + collapses to a single object.""" + client = _s3_client() + suffix = f"_{call_id}.json" + return tuple( + S3LogObject(key=key, record=self._record(client, key)) + for key in self._keys(client) + if key.endswith(suffix) + ) + + def poll_objects_for_call(self, call_id: str) -> tuple[S3LogObject, ...]: + """Poll until the call's object is readable - the integration batches + uploads behind a flush interval - and return whatever is there at the + deadline, so the caller's assertion, not a timeout, reports the gap.""" + deadline = time.monotonic() + POLL_TIMEOUT + while time.monotonic() < deadline: + found = self.objects_for_call(call_id) + if found: + return found + time.sleep(POLL_INTERVAL) + return self.objects_for_call(call_id) + + def delete(self, key: str) -> None: + _ = _s3_client().delete_object(Bucket=self.bucket, Key=key) + + +def build_s3_logs_reader() -> S3LogsReader: + if not S3_LOG_BUCKET: + pytest.fail( + "E2E_S3_LOG_BUCKET must name the bucket the proxy's s3_callback_params write to: " + "the s3 tests read delivery back out of the real bucket, and no proxy route reports " + "which one is configured; missing it is a hard failure, not a skip" + ) + return S3LogsReader(bucket=S3_LOG_BUCKET, path=S3_LOG_PATH) diff --git a/tests/e2e/logging/test_datadog_failure_log_e2e.py b/tests/e2e/logging/test_datadog_failure_log_e2e.py new file mode 100644 index 00000000000..beffd3f6ff9 --- /dev/null +++ b/tests/e2e/logging/test_datadog_failure_log_e2e.py @@ -0,0 +1,195 @@ +"""Live e2e: DataDog log delivery for calls that FAIL at the provider. + +Covers logging.datadog.failure.exports_metric: when a call dies upstream the +caller gets an error and no usage, so the DataDog event is the only record that +the request ever happened. An operator's failure-rate alert is built on it, so +the event has to arrive, has to say the call failed, and has to name what went +wrong - a success-only logging path, or one that swallows failures, silently +zeroes out that alert. + +The failure is real, not simulated: the test registers a deployment whose +upstream api_key is invalid, so OpenAI itself rejects the call and litellm's +failure path - not its success path - is what has to deliver. Both halves of the +contract are asserted: the recorded state (the proxy reports a datadog callback +registered for failure events) and the enforced behavior (the event at the real +DataDog intake, carrying the provider's error class, code and provider name, +cross-checked against the 401 the caller received, and carrying the prompt of +the very call that failed). +""" + +from __future__ import annotations + +import time +from collections.abc import Iterator + +import pytest +from pydantic import BaseModel, ConfigDict + +from callback_config import callback_enabled, registered_callbacks +from datadog_reader import DdLogEvent, DdLogsReader +from e2e_config import unique_marker +from e2e_http import StreamingResponse +from lifecycle import ResourceManager +from logging_client import INVALID_UPSTREAM_API_KEY, LoggingClient +from models import LiteLLMParamsBody + +pytestmark = pytest.mark.e2e + +#: The integration's name in litellm's callback settings. +DD_CALLBACK_NAME = "datadog" +#: A deployment litellm routes to OpenAI, so the invalid key is rejected by +#: OpenAI rather than by litellm's own request validation. +UPSTREAM_MODEL = "openai/gpt-4o-mini" +#: Present only in OpenAI's own rejection, never in the gateway's auth error, so +#: it tells an upstream 401 (the behavior under test) apart from a proxy 401. +UPSTREAM_REJECTION = "OpenAIException" + + +class _DdErrorInformation(BaseModel): + model_config = ConfigDict(extra="ignore") + + error_class: str + error_code: str + llm_provider: str + + +class _DdMessage(BaseModel): + model_config = ConfigDict(extra="ignore") + + role: str + content: str + + +class _DdFailurePayload(BaseModel): + """The fields of a failed call's StandardLoggingPayload the scenario pins.""" + + model_config = ConfigDict(extra="ignore") + + litellm_call_id: str + status: str + model_group: str + call_type: str + response_cost: float + total_tokens: int + error_str: str + error_information: _DdErrorInformation + messages: list[_DdMessage] = [] + + +@pytest.fixture +def datadog_failure_logging(client: LoggingClient) -> Iterator[None]: + """The proxy must ship failed calls to DataDog for the duration of the test. + A deployment that already does is left exactly as it is.""" + yield from callback_enabled(client, DD_CALLBACK_NAME, events=("failure",)) + + +def _first_upstream_rejection(client: LoggingClient, key: str, model: str, prompt: str) -> StreamingResponse: + """The first response that is the provider's rejection rather than the + gateway's own. A freshly created key can briefly 401 at the gateway until the + data plane's auth cache picks it up, and that 401 never reaches the provider, + so it would leave nothing for the integration to deliver; retry to a deadline + until the body carries the upstream error. + + A gateway 401 is itself logged as a failure (ProxyLogging.post_call_failure_hook + routes an auth_error on an llm api route into async_failure_handler), and that + event carries the same prompt, so a discarded attempt is indistinguishable from + the accepted one by prompt alone. The caller must therefore correlate on the + returned response's x-litellm-call-id, which is unique per attempt.""" + deadline = time.monotonic() + client.proxy.poll_timeout + while True: + outcome = client.chat_raw(key, model, prompt, max_tokens=16) + if UPSTREAM_REJECTION in outcome.body: + return outcome + if time.monotonic() >= deadline: + pytest.fail( + f"the call never reached the provider (status {outcome.status_code}): {outcome.body[:300]}" + ) + time.sleep(client.proxy.poll_interval) + + +class TestDataDogFailureLogDelivery: + @pytest.mark.covers("logging.datadog.failure.exports_metric", exercised_on=["chat_completions"]) + def test_chat_completions_failure_emits_one_log_event( + self, + client: LoggingClient, + dd_logs: DdLogsReader, + resources: ResourceManager, + datadog_failure_logging: None, + ) -> None: + """A /chat/completions call the provider rejects must reach the DataDog + logs intake as exactly one log event whose payload records the failure, + names the upstream error, and carries the prompt of the failed call.""" + assert DD_CALLBACK_NAME in registered_callbacks(client, "failure"), ( + "the proxy must report the datadog callback registered for failure events " + "before delivery can be asserted" + ) + + marker = unique_marker() + model_name = f"dd-failure-{marker}" + model_id = client.create_model( + model_name, LiteLLMParamsBody(model=UPSTREAM_MODEL, api_key=INVALID_UPSTREAM_API_KEY) + ) + resources.defer(lambda: client.delete_model(model_id)) + + key = client.key_with_alias(f"dd-failure-{marker}", models=[model_name]) + resources.defer(lambda: client.delete_key(key)) + + prompt = f"reply with one word {marker}" + outcome = _first_upstream_rejection(client, key, model_name, prompt) + assert outcome.status_code == 401, ( + f"an invalid upstream key must surface as the provider's 401, got {outcome.status_code}: " + f"{outcome.body[:300]}" + ) + call_id = outcome.call_id + assert call_id is not None, "the failed response must still carry x-litellm-call-id" + + events = dd_logs.poll_events_for_marker(call_id) + assert events, "no DataDog log event for the failed call reached the intake within the deadline" + assert len(events) == 1, ( + f"expected exactly ONE DataDog log event for call {call_id}, got {len(events)}" + ) + event: DdLogEvent = events[0] + assert "source:litellm" in event.tags, ( + f"the ingested event must carry the litellm source (shipped as ddsource), got tags {event.tags!r}" + ) + assert event.status != "ok", ( + "the integration ships a failure at DataDogStatus.ERROR, so the ingested event must not " + f"index at DataDog's ok severity - which is where a failure logged as a success lands; got " + f"{event.status!r}" + ) + + payload = _DdFailurePayload.model_validate(event.attributes) + assert payload.litellm_call_id == call_id, ( + f"the event must be the one for the attempt the caller accepted; searched {call_id}, " + f"got a payload for {payload.litellm_call_id}" + ) + assert payload.status == "failure", f"payload status must be failure, got {payload.status!r}" + assert payload.model_group == model_name, ( + f"payload model_group must be {model_name!r}, got {payload.model_group!r}" + ) + assert payload.call_type == "acompletion", ( + f"payload call_type must be acompletion, got {payload.call_type!r}" + ) + assert payload.error_information.error_code == "401", ( + f"the payload must carry the provider's status code, got {payload.error_information.error_code!r}" + ) + assert payload.error_information.error_class == "AuthenticationError", ( + f"the payload must name the upstream error class, got {payload.error_information.error_class!r}" + ) + assert payload.error_information.llm_provider == "openai", ( + f"the payload must name the provider that rejected the call, got " + f"{payload.error_information.llm_provider!r}" + ) + assert UPSTREAM_REJECTION in payload.error_str, ( + f"the payload's error_str must carry the provider's own message, got {payload.error_str!r}" + ) + assert payload.response_cost == 0, ( + "a failed call bills nothing, so a non-zero cost means the failure was accounted as if it " + f"had succeeded; got {payload.response_cost}" + ) + assert payload.total_tokens == 0, ( + f"a failed call consumes no tokens, got {payload.total_tokens}" + ) + assert any(marker in message.content for message in payload.messages), ( + f"the delivered event must carry the failed call's own prompt, got {payload.messages!r}" + ) diff --git a/tests/e2e/logging/test_s3_log_e2e.py b/tests/e2e/logging/test_s3_log_e2e.py new file mode 100644 index 00000000000..d58ae4fc1f4 --- /dev/null +++ b/tests/e2e/logging/test_s3_log_e2e.py @@ -0,0 +1,224 @@ +"""Live e2e: S3 object delivery for successful and for failed calls. + +Covers logging.s3.success.writes_object and logging.s3.failure.writes_object. +The S3 bucket is the audit trail: for a lot of deployments it, not the database, +is the record of what was asked and what it cost, and a compliance answer needs +the failed calls in it too, because a call that died upstream still left the +prompt with a provider. So the promise is an object per call, holding that +call's payload - and the integration batches uploads behind a flush interval and +swallows upload errors, which is exactly the shape of a bug that looks healthy +from the proxy and loses every record. + +Delivery is judged on what is in the bucket. The proxy uploads with its own AWS +credentials as in production, and the tests fetch the object back with the AWS +SDK (see s3_reader.py), so a dropped upload, a misaddressed bucket, or a payload +that lost the error fails here. Both halves of the contract are asserted for +each case: the recorded state (the proxy reports the s3_v2 callback registered +for that event) and the enforced behavior (the object in the bucket, keyed by +the id of the very call the caller made, with the cost cross-checked against +that response's x-litellm-response-cost header on the success path and the +provider's error preserved on the failure path). +""" + +from __future__ import annotations + +import time +from collections.abc import Iterator + +import pytest + +from callback_config import callback_enabled, registered_callbacks +from e2e_config import CHEAP_ANTHROPIC_MODEL, unique_marker +from e2e_http import StreamingResponse +from lifecycle import ResourceManager +from logging_client import ( + INVALID_UPSTREAM_API_KEY, + LoggingClient, + completion_response_id, + first_ok, +) +from models import LiteLLMParamsBody +from s3_reader import S3LogObject, S3LogsReader, build_s3_logs_reader + +pytestmark = pytest.mark.e2e + +#: The integration's name in litellm's callback settings. +S3_CALLBACK_NAME = "s3_v2" +#: A deployment litellm routes to OpenAI, so the invalid key is rejected by +#: OpenAI rather than by litellm's own request validation. +UPSTREAM_MODEL = "openai/gpt-4o-mini" +#: Present only in OpenAI's own rejection, never in the gateway's auth error, so +#: it tells an upstream 401 (the behavior under test) apart from a proxy 401. +UPSTREAM_REJECTION = "OpenAIException" + + +@pytest.fixture(scope="session") +def s3_logs() -> S3LogsReader: + """Read-back client for the real bucket the proxy's s3_callback_params name.""" + return build_s3_logs_reader() + + +@pytest.fixture(scope="module") +def s3_logging(client: LoggingClient) -> Iterator[None]: + """The proxy must ship both successful and failed calls to S3 for the + duration of this module. A deployment that already does is left as it is.""" + yield from callback_enabled(client, S3_CALLBACK_NAME, events=("success", "failure")) + + +def _first_upstream_rejection(client: LoggingClient, key: str, model: str, prompt: str) -> StreamingResponse: + """The first response that is the provider's rejection rather than the + gateway's own. A freshly created key can briefly 401 at the gateway until the + data plane's auth cache picks it up, and that 401 never reaches the provider, + so it would leave nothing for the integration to write; retry to a deadline + until the body carries the upstream error. + + A gateway 401 is itself logged as a failure and carries the same prompt, so a + discarded attempt is indistinguishable from the accepted one by prompt alone. + The caller must therefore correlate on the returned response's + x-litellm-call-id, which is unique per attempt and is what the object key + ends in.""" + deadline = time.monotonic() + client.proxy.poll_timeout + while True: + outcome = client.chat_raw(key, model, prompt, max_tokens=16) + if UPSTREAM_REJECTION in outcome.body: + return outcome + if time.monotonic() >= deadline: + pytest.fail( + f"the call never reached the provider (status {outcome.status_code}): {outcome.body[:300]}" + ) + time.sleep(client.proxy.poll_interval) + + +def _sole_object( + s3_logs: S3LogsReader, resources: ResourceManager, call_id: str +) -> S3LogObject: + """The one object the bucket holds for the call, queued for deletion the + moment it is found so a later assertion failure still cleans up.""" + written = s3_logs.poll_objects_for_call(call_id) + for entry in written: + resources.defer(lambda key=entry.key: s3_logs.delete(key)) + assert written, f"no S3 object for call {call_id} appeared in the bucket within the deadline" + assert len(written) == 1, ( + f"expected exactly ONE S3 object for the call, got {len(written)}: {[o.key for o in written]}" + ) + return written[0] + + +class TestS3LogDelivery: + @pytest.mark.covers("logging.s3.success.writes_object", exercised_on=["chat_completions"]) + def test_chat_completions_success_writes_one_object( + self, + client: LoggingClient, + s3_logs: S3LogsReader, + resources: ResourceManager, + s3_logging: None, + ) -> None: + """One successful /chat/completions call must leave exactly one object in + the bucket, keyed by that completion's id, holding the prompt, the token + counts and the same cost the caller was charged.""" + assert S3_CALLBACK_NAME in registered_callbacks(client, "success"), ( + "the proxy must report the s3_v2 callback registered for successful calls " + "before delivery can be asserted" + ) + + marker = unique_marker() + key = client.key_with_alias(f"s3-success-{marker}", models=[CHEAP_ANTHROPIC_MODEL]) + resources.defer(lambda: client.delete_key(key)) + + prompt = f"reply with one word {marker}" + outcome = first_ok( + client, lambda: client.chat_raw(key, CHEAP_ANTHROPIC_MODEL, prompt, max_tokens=16) + ) + assert outcome.response_cost is not None and outcome.response_cost > 0, ( + f"the response must report x-litellm-response-cost, got {outcome.response_cost!r}" + ) + completion_id = completion_response_id(outcome.body) + assert completion_id is not None, f"the completion must carry an id: {outcome.body[:300]}" + + record = _sole_object(s3_logs, resources, completion_id).record + assert record.id == completion_id, ( + f"the object must hold the payload of this completion, got id {record.id!r}" + ) + assert record.status == "success", f"payload status must be success, got {record.status!r}" + assert record.call_type == "acompletion", f"payload call_type must be acompletion, got {record.call_type!r}" + assert record.model_group == CHEAP_ANTHROPIC_MODEL, ( + f"payload model_group must be {CHEAP_ANTHROPIC_MODEL!r}, got {record.model_group!r}" + ) + assert record.total_tokens > 0, f"payload must count real tokens, got {record.total_tokens}" + assert record.response_cost == outcome.response_cost, ( + f"the stored cost {record.response_cost} must equal the cost the caller was charged " + f"{outcome.response_cost}" + ) + assert any(marker in message.content for message in record.messages), ( + f"the stored object must carry this call's own prompt, got {record.messages!r}" + ) + + @pytest.mark.covers("logging.s3.failure.writes_object", exercised_on=["chat_completions"]) + def test_chat_completions_failure_writes_one_object( + self, + client: LoggingClient, + s3_logs: S3LogsReader, + resources: ResourceManager, + s3_logging: None, + ) -> None: + """A /chat/completions call the provider rejects must still leave exactly + one object in the bucket, keyed by the litellm call id the caller was + handed, recording the failure and the upstream error rather than being + dropped with the response.""" + assert S3_CALLBACK_NAME in registered_callbacks(client, "failure"), ( + "the proxy must report the s3_v2 callback registered for failed calls " + "before delivery can be asserted" + ) + + marker = unique_marker() + model_name = f"s3-failure-{marker}" + model_id = client.create_model( + model_name, LiteLLMParamsBody(model=UPSTREAM_MODEL, api_key=INVALID_UPSTREAM_API_KEY) + ) + resources.defer(lambda: client.delete_model(model_id)) + + key = client.key_with_alias(f"s3-failure-{marker}", models=[model_name]) + resources.defer(lambda: client.delete_key(key)) + + prompt = f"reply with one word {marker}" + outcome = _first_upstream_rejection(client, key, model_name, prompt) + assert outcome.status_code == 401, ( + f"an invalid upstream key must surface as the provider's 401, got {outcome.status_code}: " + f"{outcome.body[:300]}" + ) + assert outcome.call_id is not None, "the failed response must still carry x-litellm-call-id" + + record = _sole_object(s3_logs, resources, outcome.call_id).record + assert record.id == outcome.call_id, ( + f"the object must hold the payload of this call, got id {record.id!r}" + ) + assert record.status == "failure", f"payload status must be failure, got {record.status!r}" + assert record.call_type == "acompletion", f"payload call_type must be acompletion, got {record.call_type!r}" + assert record.model_group == model_name, ( + f"payload model_group must be {model_name!r}, got {record.model_group!r}" + ) + assert record.error_information.error_code == "401", ( + f"the stored payload must carry the provider's status code, got " + f"{record.error_information.error_code!r}" + ) + assert record.error_information.error_class == "AuthenticationError", ( + f"the stored payload must name the upstream error class, got " + f"{record.error_information.error_class!r}" + ) + assert record.error_information.llm_provider == "openai", ( + f"the stored payload must name the provider that rejected the call, got " + f"{record.error_information.llm_provider!r}" + ) + assert record.error_str is not None and UPSTREAM_REJECTION in record.error_str, ( + f"the stored payload's error_str must carry the provider's own message, got {record.error_str!r}" + ) + assert record.response_cost == 0, ( + "a failed call bills nothing, so a non-zero cost means the failure was archived as if it " + f"had succeeded; got {record.response_cost}" + ) + assert record.total_tokens == 0, ( + f"a failed call consumes no tokens, got {record.total_tokens}" + ) + assert any(marker in message.content for message in record.messages), ( + f"the stored object must carry the failed call's own prompt, got {record.messages!r}" + )