Merge 2090047bd3 into 49affa7c01

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

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)

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

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