From 948e5755eba9cb80e1239ecebfb717fcad9b2c36 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Sat, 5 Sep 2026 13:03:28 -0700 Subject: [PATCH] test(e2e): cover presidio post_call, tool_permission, and weave logging cells (#39279) * test(e2e): cover presidio post_call, tool_permission, and weave logging cells Five registry cells in Logging & Guardrails had no covering test. Each one now has a live scenario read back from the real destination: - guardrail.presidio.post_call.masks: an output-scoped Presidio guardrail anonymizes the PII the model repeats back. The prompt also asks for the address's local part, which Presidio does not mask, so one response proves the model saw the raw address (no pre-call masking) while the address itself comes back as - guardrail.tool_permission.pre_call.blocks / .allows: an allow-list of one tool. A request declaring an unlisted tool is rejected 400 naming it; a request declaring the permitted tool is served and carries x-litellm-applied-guardrails, so the allow half cannot pass by the guardrail never running - logging.niche_integrations.success.logs_spend / .failure.logs_spend: a key-scoped weave_otel callback delivers to the real Weave project, read back through Weave's query API. Success asserts exactly one call whose llm.response.cost equals the x-litellm-response-cost header; failure asserts one ERROR-status call naming the provider exception and carrying no cost Logging & Guardrails coverage goes 24/59 to 29/59. No registry rows are added. * test(e2e): make the tool-permission allow case deterministic and scope the Weave read-back Review follow-ups on the coverage PR. - the allow scenario forced the outcome to depend on whether the model felt like calling an optional tool, and checked for the tool name as a substring of the whole body, which a prose mention would satisfy. It now sends tool_choice="required" and asserts the parsed response carries exactly one tool call, for the permitted tool - the Weave read-back queried the newest 200 calls of a shared project and filtered client-side, so busy traffic could push the target out of the window and read as a delivery failure. The query now scopes server-side to the litellm_request op and to calls started after the request, and pages through the window with offset - the reader builds its results as tuples instead of accumulating into lists Also unblocks the lint gate: `basedpyright tests/e2e` runs only on PRs that touch tests/e2e, and it has been failing on staging for three FakeItem arguments in test_junit_properties.py. The stand-in now goes through one typed adapter that says why, so the gate is green without touching junit_properties.py itself. * test(e2e): scope the presidio post_call guardrail to email and phone Running the suite three times in a row caught a real flake: Presidio's broader recognizers sometimes claim the email's local part as an NRP entity, so the answer came back as `\n\n` and the assertion that the raw local part survives failed. That token is what tells output masking apart from input masking, so it has to survive. The post_call guardrail now registers pii_entities_config for EMAIL_ADDRESS and PHONE_NUMBER only, which is also the narrower thing the scenario means. Verified against the exact marker that failed, plus two others. * test(e2e): mark weave logging cells stage red * test(e2e): use per-test stage red skips for the weave logging cells --- tests/e2e/CONTRIBUTING.md | 4 +- tests/e2e/guardrails/guardrails_client.py | 71 ++++- .../guardrails/test_presidio_masking_e2e.py | 106 ++++++- .../test_tool_permission_guardrail_e2e.py | 167 +++++++++++ tests/e2e/logging/logging_client.py | 39 +++ tests/e2e/logging/test_weave_log_e2e.py | 192 ++++++++++++ tests/e2e/logging/weave_reader.py | 282 ++++++++++++++++++ tests/e2e/models.py | 2 + 8 files changed, 844 insertions(+), 19 deletions(-) create mode 100644 tests/e2e/guardrails/test_tool_permission_guardrail_e2e.py create mode 100644 tests/e2e/logging/test_weave_log_e2e.py create mode 100644 tests/e2e/logging/weave_reader.py diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index 871d6b3904c..b270feb820e 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -50,7 +50,9 @@ The suites run against a live proxy, so bring one up first by running the litell They also need a proxy whose bundled UI contains the change under test, so run the proxy from your branch (an editable install serves the UI your checkout builds) -Some suites need extra services the bare proxy does not start. The `logging/` OTEL trace-completeness tests read spans back from a jaeger query API at `http://localhost:16686` (override with `E2E_OTEL_QUERY_URL`); run a `jaegertracing/all-in-one` and point `PHOENIX_COLLECTOR_HTTP_ENDPOINT` at its OTLP ingest. The `mcp/` suite needs the deterministic upstream MCP server in `mcp_tests/mcp_e2e_upstream_server.py` reachable by the proxy +Some suites need extra services the bare proxy does not start. The `logging/` OTEL trace-completeness tests read spans back from a jaeger query API at `http://localhost:16686` (override with `E2E_OTEL_QUERY_URL`); run a `jaegertracing/all-in-one` and point `PHOENIX_COLLECTOR_HTTP_ENDPOINT` at its OTLP ingest. The `mcp/` suite needs the deterministic upstream MCP server in `mcp_tests/mcp_e2e_upstream_server.py` reachable by the proxy. The presidio guardrail tests need a running Presidio analyzer and anonymizer the proxy can reach, addressed by `PRESIDIO_ANALYZER_API_BASE` / `PRESIDIO_ANONYMIZER_API_BASE` + +A couple of logging destinations are configured on the proxy rather than by the test. The Weave tests scope their callback to the key they create, but litellm builds the `weave_otel` logger from `WANDB_API_KEY` and `WANDB_PROJECT_ID` before it applies the per-key vars, so the proxy needs both in its own environment or the key-scoped callback never initializes and nothing ships ### Record and replay diff --git a/tests/e2e/guardrails/guardrails_client.py b/tests/e2e/guardrails/guardrails_client.py index f03e70df84a..1f55a0f9a56 100644 --- a/tests/e2e/guardrails/guardrails_client.py +++ b/tests/e2e/guardrails/guardrails_client.py @@ -18,6 +18,7 @@ from models import ( ChatBody, ChatMessage, ChatResponse, + ChatTool, KeyGenerateBody, LiteLLMParamsBody, TeamDeleteBody, @@ -31,6 +32,8 @@ from proxy_client import ProxyClient from pydantic import BaseModel GuardrailMode = Literal["pre_call", "post_call", "during_call", "logging_only"] +PiiEntity = Literal["EMAIL_ADDRESS", "PHONE_NUMBER", "PERSON", "CREDIT_CARD", "US_SSN"] +PiiAction = Literal["MASK", "BLOCK"] BlockedWordAction = Literal["BLOCK", "MASK"] @@ -81,6 +84,27 @@ class PresidioParamsBody(GuardrailParamsBase): presidio_filter_scope: Literal["input", "output", "both"] | None = None presidio_language: str | None = None output_parse_pii: bool | None = None + pii_entities_config: dict[PiiEntity, PiiAction] | None = None + + +class ToolPermissionRuleBody(BaseModel): + """One tool_permission rule: a decision for the tool named by `tool_name`.""" + + id: str + tool_name: str + decision: Literal["allow", "deny"] + + +class ToolPermissionParamsBody(GuardrailParamsBase): + """Tool-permission guardrail params. `default_action="deny"` makes the rules + an allow-list, and `on_disallowed_action="block"` turns a disallowed tool into + a 400 instead of rewriting the request; "rewrite" is a different product + promise and belongs to its own scenario.""" + + guardrail: Literal["tool_permission"] = "tool_permission" + rules: list[ToolPermissionRuleBody] + default_action: Literal["allow", "deny"] = "deny" + on_disallowed_action: Literal["block", "rewrite"] = "block" GuardrailParamsBody = ( @@ -89,6 +113,7 @@ GuardrailParamsBody = ( | OpenAIModerationParamsBody | BlockCodeExecutionParamsBody | PresidioParamsBody + | ToolPermissionParamsBody ) @@ -200,9 +225,7 @@ class GuardrailsClient: self.proxy.transport.post( "/guardrails", headers=self.proxy.transport.master, - json=GuardrailCreateBody( - guardrail=GuardrailSpecBody(guardrail_name=name, litellm_params=params) - ), + json=GuardrailCreateBody(guardrail=GuardrailSpecBody(guardrail_name=name, litellm_params=params)), response_type=GuardrailCreateResponse, ) ).guardrail_id @@ -241,9 +264,7 @@ class GuardrailsClient: ) def create_key_in_team(self, team_id: str) -> str: - return self.proxy.generate_key( - KeyGenerateBody(team_id=team_id, user_id="e2e-guardrails-user") - ) + return self.proxy.generate_key(KeyGenerateBody(team_id=team_id, user_id="e2e-guardrails-user")) def chat( self, @@ -253,6 +274,7 @@ class GuardrailsClient: *, guardrails: list[str] | None = None, max_tokens: int = 16, + tools: list[ChatTool] | None = None, ) -> Result[ChatResponse]: """Drive a chat call, optionally opting into named guardrails for this request only (the per-request `guardrails` selector). With `guardrails` @@ -266,6 +288,35 @@ class GuardrailsClient: messages=[ChatMessage(role="user", content=text)], max_tokens=max_tokens, guardrails=guardrails, + tools=tools, + ), + ) + + def chat_raw( + self, + key: str, + model: str, + text: str, + *, + guardrails: list[str] | None = None, + max_tokens: int = 16, + tools: list[ChatTool] | None = None, + tool_choice: str | None = None, + ) -> StreamingResponse: + """Drive /chat/completions returning the raw HTTP outcome, for the + assertions a typed body cannot carry: the `x-litellm-applied-guardrails` + response header, which is how an ALLOW scenario proves the guardrail ran + rather than being absent.""" + return self.proxy.transport.send( + "/chat/completions", + headers=self.proxy.transport.bearer(key), + json=ChatBody( + model=model, + messages=[ChatMessage(role="user", content=text)], + max_tokens=max_tokens, + guardrails=guardrails, + tools=tools, + tool_choice=tool_choice, ), ) @@ -323,9 +374,7 @@ class GuardrailsClient: return self.proxy.transport.send( "/v1/responses", headers=self.proxy.transport.bearer(key), - json=_ResponsesGuardrailBody( - model=model, input=text, guardrails=guardrails - ), + json=_ResponsesGuardrailBody(model=model, input=text, guardrails=guardrails), ) def apply_guardrail(self, key: str, *, name: str, text: str) -> Result[ApplyGuardrailResponse]: @@ -349,9 +398,7 @@ class GuardrailsClient: if isinstance(last, Success): return time.sleep(POLL_INTERVAL) - raise AssertionError( - f"team {team_id!r} was created but /team/info never returned it: {last}" - ) + raise AssertionError(f"team {team_id!r} was created but /team/info never returned it: {last}") def build_client(proxy: ProxyClient) -> GuardrailsClient: diff --git a/tests/e2e/guardrails/test_presidio_masking_e2e.py b/tests/e2e/guardrails/test_presidio_masking_e2e.py index 6d927292975..c6d87473c21 100644 --- a/tests/e2e/guardrails/test_presidio_masking_e2e.py +++ b/tests/e2e/guardrails/test_presidio_masking_e2e.py @@ -6,11 +6,19 @@ messages BEFORE the model runs, so the model only ever sees placeholders like must come back with the placeholders echoed and the raw PII absent, on /chat/completions and on /v1/messages (Anthropic format). +post_call: the mirror hook. The request reaches the model unmasked and the +MODEL OUTPUT is what gets anonymized, so the caller never receives raw PII the +model repeated back. The two hooks are told apart behaviorally rather than by +configuration: the post_call prompt asks for a value derived from the raw email +(its local part, which is not itself an entity Presidio masks) alongside the +address itself, so the answer proves the model saw the raw address while the +address in the same response comes back as . + The analyzer/anonymizer endpoints come from PRESIDIO_ANALYZER_API_BASE / PRESIDIO_ANONYMIZER_API_BASE; missing env is a hard failure, never a skip. -Each guardrail registers with presidio_filter_scope="input" so only the -configured hook's callback exists (the default "both" adds a second post_call -output masker), and is deleted on teardown. +Each guardrail registers with an explicit presidio_filter_scope so only the +configured hook's callback exists (the default "both" registers input masking +AND a post_call output masker), and is deleted on teardown. """ from __future__ import annotations @@ -18,13 +26,14 @@ from __future__ import annotations import os import time from collections.abc import Callable +from typing import Literal import pytest from pydantic import BaseModel from e2e_config import unique_marker from e2e_http import Result, Success -from guardrails_client import GuardrailsClient, PresidioParamsBody +from guardrails_client import GuardrailMode, GuardrailsClient, PiiAction, PiiEntity, PresidioParamsBody from lifecycle import ResourceManager from models import AnthropicMessagesResponse, ChatResponse @@ -65,16 +74,20 @@ def _register_presidio( resources: ResourceManager, *, name: str, + mode: GuardrailMode = "pre_call", + filter_scope: Literal["input", "output", "both"] = "input", + entities: dict[PiiEntity, PiiAction] | None = None, ) -> None: analyzer, anonymizer = _presidio_bases() guardrail_id = client.register( name, PresidioParamsBody( - mode="pre_call", + mode=mode, default_on=False, presidio_analyzer_api_base=analyzer, presidio_anonymizer_api_base=anonymizer, - presidio_filter_scope="input", + presidio_filter_scope=filter_scope, + pii_entities_config=entities, ), ) resources.defer(lambda: client.delete_guardrail(guardrail_id)) @@ -182,3 +195,84 @@ class TestPresidioPreCallMasking: _messages_text, email=email, ) + + +#: Room for the model's reasoning tokens plus the three-line answer; a lower cap +#: truncates the response before the address it is supposed to mask. +_POST_CALL_MAX_TOKENS = 512 + +#: The post_call scenario masks these two entities and nothing else. Left +#: unscoped, Presidio's broader recognizers claim the local part too (a random +#: marker reads as an NRP), which would erase the very token that tells output +#: masking apart from input masking. +_POST_CALL_ENTITIES: dict[PiiEntity, PiiAction] = {"EMAIL_ADDRESS": "MASK", "PHONE_NUMBER": "MASK"} + + +def _post_call_prompt(marker: str, local_part: str) -> str: + """Ask for the local part and the full address in one answer. Presidio masks + an EMAIL_ADDRESS entity and a bare local part is not one, so the two land + differently in the same response and pin the hook point behaviorally.""" + return ( + f"{marker} My email address is {local_part}@example.com and my phone number is {FAKE_PHONE}. " + "Reply with exactly three lines and nothing else. " + "Line 1: the part of the email address before the @ sign. " + "Line 2: the full email address. " + "Line 3: the phone number." + ) + + +class TestPresidioPostCallMasking: + @pytest.mark.covers( + "guardrail.presidio.post_call.masks", + exercised_on=["chat_completions"], + ) + def test_post_call_masks_pii_in_model_output( + self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str + ) -> None: + """A guardrail scoped to the output must anonymize the PII the model + repeats back, so a caller (or a downstream log of the response) never + receives it, while the request itself reaches the model untouched. + + Both facts are asserted from one response: the local part comes back raw, + which is only possible if the model saw the real address, and the address + itself comes back as in the same answer. + """ + name = f"e2e-presidio-post-chat-{unique_marker()}" + _register_presidio( + client, + resources, + name=name, + mode="post_call", + filter_scope="output", + entities=_POST_CALL_ENTITIES, + ) + + local_part = f"e2euser{unique_marker()}" + email = f"{local_part}@example.com" + prompt = _post_call_prompt(unique_marker(), local_part) + + deadline = time.monotonic() + GUARDRAIL_PROPAGATION_DEADLINE_SECONDS + last = "" + while True: + result = client.chat(scoped_key, MODEL, prompt, guardrails=[name], max_tokens=_POST_CALL_MAX_TOKENS) + match result: + case Success(data=data): + last = _first_content(data) + if MASKED_EMAIL_TOKEN in last and email not in last: + assert local_part in last, ( + "the model must have seen the RAW address (it is asked for the local " + "part, which Presidio does not mask); the local part is missing, so " + f"this response cannot tell post_call masking from pre_call: {last[:300]!r}" + ) + assert MASKED_PHONE_TOKEN in last and FAKE_PHONE not in last, ( + f"the phone number in the model's answer must be masked too, got: {last[:300]!r}" + ) + return + case _: + last = f"" + if time.monotonic() >= deadline: + pytest.fail( + f"presidio post_call guardrail never masked the model's output within " + f"{GUARDRAIL_PROPAGATION_DEADLINE_SECONDS}s; last observation: {last[:300]!r}" + ) + time.sleep(GUARDRAIL_PROPAGATION_POLL_INTERVAL_SECONDS) diff --git a/tests/e2e/guardrails/test_tool_permission_guardrail_e2e.py b/tests/e2e/guardrails/test_tool_permission_guardrail_e2e.py new file mode 100644 index 00000000000..9ef3650625c --- /dev/null +++ b/tests/e2e/guardrails/test_tool_permission_guardrail_e2e.py @@ -0,0 +1,167 @@ +"""Live e2e: the tool_permission guardrail gates which tools a request may declare. + +The guardrail is registered `mode="pre_call"` with `default_action="deny"`, so its +rules are an allow-list applied to the tools the CALLER declares, before the model +runs. Two halves of one product promise: + +- blocks: a request declaring a tool outside the allow-list is rejected with a 400 + naming the denied tool, and never reaches the model +- allows: a request declaring only the permitted tool is served normally, comes + back with a real tool call for that tool, and carries an + `x-litellm-applied-guardrails` header naming the guardrail, which is what + separates "the guardrail ran and allowed it" from "the guardrail was never + attached". `tool_choice="required"` keeps the model from answering directly and + making the outcome depend on its mood + +No vendor API is involved: `tool_permission` is a built-in guardrail, so the +verdict comes from the proxy itself. +""" + +from __future__ import annotations + +from typing import Final + +import pytest + +from e2e_config import unique_marker +from e2e_http import StreamingResponse, UnknownApiError +from guardrails_client import ( + GuardrailsClient, + ToolPermissionParamsBody, + ToolPermissionRuleBody, + poll_until_blocked, +) +from lifecycle import ResourceManager +from models import ChatResponse, ChatTool, ChatToolFunction + +pytestmark = pytest.mark.e2e + +MODEL = "gemini-2.5-flash" + +#: The one tool the guardrail permits, and one it does not. Both are declared by +#: the caller in the request body; the guardrail reads them there. +ALLOWED_TOOL: Final = ChatTool( + function=ChatToolFunction( + name="get_weather", + description="Get the current weather for a city", + parameters={ + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + ) +) +DENIED_TOOL: Final = ChatTool( + function=ChatToolFunction( + name="delete_customer_database", + description="Permanently delete the customer database", + parameters={"type": "object", "properties": {}}, + ) +) + +TOOL_PROMPT: Final = "What is the weather in Paris right now?" + + +def _register_tool_permission(client: GuardrailsClient, resources: ResourceManager, *, name: str) -> None: + """Allow-list exactly one tool: everything else falls to `default_action=deny` + and, with `on_disallowed_action=block`, is rejected outright.""" + guardrail_id = client.register( + name, + ToolPermissionParamsBody( + mode="pre_call", + default_on=False, + default_action="deny", + on_disallowed_action="block", + rules=[ + ToolPermissionRuleBody( + id="allow-get-weather", + tool_name=ALLOWED_TOOL.function.name, + decision="allow", + ) + ], + ), + ) + resources.defer(lambda: client.delete_guardrail(guardrail_id)) + + +def _applied_guardrails(outcome: StreamingResponse) -> str: + return outcome.headers.get("x-litellm-applied-guardrails", "") + + +def _tool_call_names(response: ChatResponse) -> tuple[str, ...]: + return tuple( + call.function.name + for choice in response.choices + if choice.message + for call in choice.message.tool_calls or () + if call.function.name + ) + + +class TestToolPermissionPreCall: + @pytest.mark.covers("guardrail.tool_permission.pre_call.blocks", exercised_on=["chat_completions"]) + def test_pre_call_blocks_tool_outside_the_allow_list( + self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str + ) -> None: + """A request declaring a tool the guardrail does not permit must be + rejected with a 400 that names the denied tool. An unauthorized tool that + merely reaches the model is the whole failure mode this guardrail exists + to prevent, so a 200 here is a hard failure.""" + name = f"e2e-toolperm-block-{unique_marker()}" + _register_tool_permission(client, resources, name=name) + + result = poll_until_blocked( + lambda: client.chat( + scoped_key, + MODEL, + TOOL_PROMPT, + guardrails=[name], + max_tokens=128, + tools=[DENIED_TOOL], + ) + ) + + match result: + case UnknownApiError(status_code=status, body=body): + assert status == 400, f"expected the guardrail block status 400, got {status}: {body[:400]}" + assert DENIED_TOOL.function.name in body, ( + f"the block must name the denied tool so the caller can fix the request; got: {body[:400]}" + ) + assert "guardrail" in body.lower(), ( + f"the block body should identify itself as a guardrail verdict; got: {body[:400]}" + ) + case _: + pytest.fail(f"tool_permission let a tool outside the allow-list through; got {result}") + + @pytest.mark.covers("guardrail.tool_permission.pre_call.allows", exercised_on=["chat_completions"]) + def test_pre_call_allows_permitted_tool( + self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str + ) -> None: + """The mirror half: a request declaring only the permitted tool is served + and the model calls it. Without the header check a guardrail that never + attached would pass this test for the wrong reason, so the 200 alone is + not the contract.""" + name = f"e2e-toolperm-allow-{unique_marker()}" + _register_tool_permission(client, resources, name=name) + + outcome = client.chat_raw( + scoped_key, + MODEL, + TOOL_PROMPT, + guardrails=[name], + max_tokens=128, + tools=[ALLOWED_TOOL], + tool_choice="required", + ) + + assert outcome.ok, f"the permitted tool must be served, got {outcome.status_code}: {outcome.body[:400]}" + applied = _applied_guardrails(outcome) + assert name in applied, ( + "the allowed call must carry x-litellm-applied-guardrails naming the guardrail; " + f"without it the 200 only proves the guardrail never ran. Got {applied!r}" + ) + + called = _tool_call_names(ChatResponse.model_validate_json(outcome.body)) + assert called == (ALLOWED_TOOL.function.name,), ( + f"the served call must carry one tool call for the permitted tool, got {called!r}: {outcome.body[:400]}" + ) diff --git a/tests/e2e/logging/logging_client.py b/tests/e2e/logging/logging_client.py index f0f7ad7eaa4..66dfa233ec4 100644 --- a/tests/e2e/logging/logging_client.py +++ b/tests/e2e/logging/logging_client.py @@ -189,6 +189,45 @@ class LangfuseCreds: ) +@dataclass(frozen=True, slots=True) +class WeaveCreds: + """Weights & Biases Weave credentials for a key-scoped ``weave_otel`` callback. + + The proxy still needs WANDB_API_KEY / WANDB_PROJECT_ID in its own environment: + the weave_otel logger is constructed from those before the per-key vars are + applied, so a key-scoped callback on a proxy without them never initializes. + The per-key vars are what direct THIS key's spans at this project. + """ + + api_key: str + project_id: str + + def key_logging_metadata(self) -> KeyMetadata: + return KeyMetadata( + logging=[ + KeyLoggingCallback( + callback_name="weave_otel", + callback_type="success_and_failure", + callback_vars=KeyLoggingCallbackVars( + wandb_api_key=self.api_key, + weave_project_id=self.project_id, + ), + ) + ] + ) + + +def load_weave_creds() -> WeaveCreds: + api_key = os.getenv("WANDB_API_KEY") + project_id = (os.getenv("WEAVE_PROJECT_ID") or os.getenv("WANDB_PROJECT_ID") or "").strip() + if not (api_key and project_id): + pytest.fail( + "Weave e2e requires WANDB_API_KEY and WEAVE_PROJECT_ID (or WANDB_PROJECT_ID, " + "format /); missing credentials is a hard failure, not a skip" + ) + return WeaveCreds(api_key=api_key, project_id=project_id) + + def load_langfuse_creds() -> LangfuseCreds: public_key = os.getenv("LANGFUSE_PUBLIC_KEY") secret_key = os.getenv("LANGFUSE_SECRET_KEY") diff --git a/tests/e2e/logging/test_weave_log_e2e.py b/tests/e2e/logging/test_weave_log_e2e.py new file mode 100644 index 00000000000..dab5993c87a --- /dev/null +++ b/tests/e2e/logging/test_weave_log_e2e.py @@ -0,0 +1,192 @@ +"""Live e2e: key-scoped Weave (Weights & Biases) delivery, success and failure. + +Covers the two `logging.niche_integrations.*.logs_spend` cells with a real member +of that cohort. A key carrying a `weave_otel` callback in its logging metadata +must deliver its calls to the real Weave project, and each call must arrive +exactly once, carrying the same cost the response header reported: + +- success: one `litellm_request` call, OTEL status OK, `llm.response.cost` equal + to `x-litellm-response-cost`, and non-zero tokens +- failure: a provider-rejected call arrives too, as one call with OTEL status + ERROR naming the provider exception, and with no cost - a failed call that + silently never reaches the destination is an invisible outage, and a billed + one is worse + +Both halves assert the recorded state (the key's callback registration answers +success and the destination holds the call) and the enforced behavior (the +delivered payload's status and cost). Delivery is read back through Weave's own +query API; nothing is mocked. +""" + +from __future__ import annotations + +import time + +import pytest + +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, + WeaveCreds, + costs_agree, + first_ok, + load_weave_creds, +) +from models import LiteLLMParamsBody +from weave_reader import WeaveCall, WeaveReader, build_weave_reader + +pytestmark = pytest.mark.e2e + + +@pytest.fixture(scope="session") +def weave_creds() -> WeaveCreds: + return load_weave_creds() + + +@pytest.fixture(scope="session") +def weave_reader() -> WeaveReader: + return build_weave_reader() + + +#: How far before the request the Weave read-back window opens, to absorb clock +#: skew between this host and Weave. Without it a host running slightly fast +#: would filter out its own call. +_WINDOW_SKEW_SECONDS = 120.0 + + +def _window_start() -> float: + return time.time() - _WINDOW_SKEW_SECONDS + + +def _exactly_one(calls: tuple[WeaveCall, ...], *, marker: str, what: str) -> WeaveCall: + assert calls, f"no Weave call for the {what} (marker {marker}) reached the project within the deadline" + assert len(calls) == 1, ( + f"expected exactly ONE Weave call for the {what} (marker {marker}), got {len(calls)}: " + f"{[call.id for call in calls]} - more than one call for one request is the " + "duplicate-delivery bug" + ) + return calls[0] + + +WEAVE_STAGE_RED_REASON = ( + "stage red: product gap, key-scoped weave_otel spans are not delivered when the OTEL v2 callback is active" +) + + +class TestWeaveLogDelivery: + @pytest.mark.skip(reason=WEAVE_STAGE_RED_REASON) + @pytest.mark.covers("logging.niche_integrations.success.logs_spend", exercised_on=["chat_completions"]) + def test_chat_completions_delivers_one_call_with_spend( + self, + client: LoggingClient, + weave_creds: WeaveCreds, + weave_reader: WeaveReader, + resources: ResourceManager, + ) -> None: + alias = f"weave-key-{unique_marker()}" + key = client.key_with_alias( + alias, + models=[CHEAP_ANTHROPIC_MODEL], + metadata=weave_creds.key_logging_metadata(), + ) + resources.defer(lambda: client.delete_key(key)) + + marker = unique_marker() + since = _window_start() + outcome = first_ok( + client, + lambda: client.chat_raw(key, CHEAP_ANTHROPIC_MODEL, f"reply with one word {marker}", max_tokens=64), + ) + 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}" + ) + + call = _exactly_one( + weave_reader.poll_calls_matching(marker, since=since), marker=marker, what="successful call" + ) + + assert call.status_code == "OK", f"a successful call must land at OK span status, got {call.status_code!r}" + cost = call.response_cost + assert cost is not None and costs_agree(outcome.response_cost, cost), ( + f"the Weave call's llm.response.cost {cost!r} must agree with the header cost " + f"{outcome.response_cost} - a delivered span with the wrong cost is a silent " + "billing-attribution bug" + ) + assert call.total_tokens is not None and call.total_tokens > 0, ( + f"the delivered call must carry token usage, got {call.total_tokens!r}" + ) + + @pytest.mark.skip(reason=WEAVE_STAGE_RED_REASON) + @pytest.mark.covers("logging.niche_integrations.failure.logs_spend", exercised_on=["chat_completions"]) + def test_failed_chat_completions_delivers_one_error_call( + self, + client: LoggingClient, + weave_creds: WeaveCreds, + weave_reader: WeaveReader, + resources: ResourceManager, + ) -> None: + """A deployment with an invalid upstream key passes proxy auth and fails + at the provider, so exactly one provider failure exists for it. Proxy-side + 401s during key propagation never reach the provider and ship no payload, + which is what the retry loop below relies on.""" + model_name = f"weave-err-{unique_marker()}" + model_id = client.create_model( + model_name, + LiteLLMParamsBody(model="anthropic/claude-haiku-4-5", api_key=INVALID_UPSTREAM_API_KEY), + ) + resources.defer(lambda: client.delete_model(model_id)) + key = client.key_with_alias( + f"weave-err-key-{unique_marker()}", + models=[model_name], + metadata=weave_creds.key_logging_metadata(), + ) + resources.defer(lambda: client.delete_key(key)) + + marker = unique_marker() + since = _window_start() + outcome = _provoke_provider_failure(client, key, model_name, marker) + + call = _exactly_one(weave_reader.poll_calls_matching(marker, since=since), marker=marker, what="failed call") + + assert call.status_code == "ERROR", ( + f"a failed call must land at ERROR span status, got {call.status_code!r} - " + "Weave's own summary.weave.status reads success either way, which is exactly " + "why the span status is what this asserts on" + ) + error = call.error + assert error is not None and error.message is not None and "AnthropicException" in error.message, ( + f"the delivered call must carry the provider error, got {error!r}" + ) + assert not call.response_cost, f"a failed call must not be billed, got llm.response.cost={call.response_cost!r}" + assert outcome.status_code == 401, ( + f"an upstream auth failure must map to 401, got {outcome.status_code}: {outcome.body[:200]}" + ) + + +def _provoke_provider_failure(client: LoggingClient, key: str, model_name: str, marker: str) -> StreamingResponse: + """Send until the provider (not the proxy) is the one rejecting the call. + + A network failure between the test and the proxy is NOT retried: the request + may have been served, and a retry would double-log the failure payload and + falsely trip the exactly-one assertion. + """ + deadline = time.monotonic() + client.proxy.poll_timeout + while True: + outcome = client.chat_raw(key, model_name, f"trigger an upstream auth failure {marker}", max_tokens=16) + assert not outcome.ok, "the call must fail; the deployment's upstream key is invalid" + assert outcome.status_code != -1, ( + "network failure between the test and the proxy while provoking the provider failure; " + "retrying now could double-log the failure payload and falsely trip the exactly-one " + f"assertion - fix the rig connectivity first: {outcome.body[:200]}" + ) + if "AnthropicException" in outcome.body or time.monotonic() >= deadline: + break + time.sleep(client.proxy.poll_interval) + assert "AnthropicException" in outcome.body, ( + "never saw the upstream provider failure before the deadline; the key may still be " + f"propagating - last outcome {outcome.status_code}: {outcome.body[:200]}" + ) + return outcome diff --git a/tests/e2e/logging/weave_reader.py b/tests/e2e/logging/weave_reader.py new file mode 100644 index 00000000000..2f8f759d299 --- /dev/null +++ b/tests/e2e/logging/weave_reader.py @@ -0,0 +1,282 @@ +"""Read-back for the Weave (Weights & Biases) logging tests against the real +Weave project. + +The proxy ships OTEL spans to https://trace.wandb.ai/otel/v1/traces with the +``weave_otel`` callback, and the tests read the ingested calls back through +Weave's own query API (``POST /calls/stream_query``), which answers JSON Lines: +one JSON object per call, so the body is parsed line by line rather than as one +document. + +The project is shared with other traffic, so the read never relies on the target +being among the newest N calls: the query is scoped server-side to the +``litellm_request`` op and to calls that started after the test's own request, +and pages with ``offset`` until the window is exhausted. + +Weave's own ``summary.weave.status`` is a rollup that reads "success" even for a +span the exporter marked failed, so status comes from the OTEL span itself +(``attributes.otel_span.status.code``), and the shipped cost from +``attributes.otel_span.attributes.llm.response.cost`` - the StandardLogging +``response_cost``, which is what makes this a spend assertion rather than a +delivery ping. + +Missing configuration is a hard failure, never a skip. +""" + +from __future__ import annotations + +import base64 +import json +import os +import time +from dataclasses import dataclass +from itertools import count, takewhile +from typing import Final + +import pytest +from pydantic import BaseModel, ConfigDict, Field + +from e2e_config import POLL_INTERVAL, POLL_TIMEOUT +from e2e_http import URL, AuthHeaders, send + +_WEAVE_TRACE_API: Final = "https://trace.wandb.ai" + +#: The op every litellm LLM call lands under. The proxy also exports a root +#: server span ("Received Proxy Server Request") and management spans; only the +#: LLM call carries the usage and cost this suite asserts on. +LITELLM_REQUEST_OP: Final = "litellm_request" + +#: How long to keep re-reading after the first matching call before trusting the +#: exactly-one assertion. The OTEL batch exporter flushes on its own schedule, so +#: a duplicate export can surface well after the first one, and a duplicate IS +#: the bug being guarded against. +WEAVE_SETTLE_SECONDS: Final = 45.0 + +#: Rows per page. The query is already scoped to this run's time window, so this +#: only bounds one round trip, not what the read can see. +_PAGE_SIZE: Final = 500 + + +class _WeaveSortBy(BaseModel): + field: str + direction: str + + +class _WeaveOpFilter(BaseModel): + op_names: list[str] + + +class _WeaveGetField(BaseModel): + get_field: str = Field(serialization_alias="$getField") + + +class _WeaveLiteral(BaseModel): + literal: float = Field(serialization_alias="$literal") + + +class _WeaveGreaterThan(BaseModel): + gt: tuple[_WeaveGetField, _WeaveLiteral] = Field(serialization_alias="$gt") + + +class _WeaveQuery(BaseModel): + expr: _WeaveGreaterThan = Field(serialization_alias="$expr") + + +class _WeaveQueryBody(BaseModel): + project_id: str + filter: _WeaveOpFilter + query: _WeaveQuery + limit: int = _PAGE_SIZE + offset: int = 0 + sort_by: list[_WeaveSortBy] = [_WeaveSortBy(field="started_at", direction="asc")] + + +class _OtelStatus(BaseModel): + model_config = ConfigDict(extra="ignore") + + code: str | None = None + message: str | None = None + + +class _OtelError(BaseModel): + model_config = ConfigDict(extra="ignore") + + code: str | None = None + type: str | None = None + message: str | None = None + + +class _LlmResponse(BaseModel): + model_config = ConfigDict(extra="ignore") + + cost: float | None = None + + +class _LlmAttributes(BaseModel): + model_config = ConfigDict(extra="ignore") + + response: _LlmResponse | None = None + + +class _OtelSpanAttributes(BaseModel): + model_config = ConfigDict(extra="ignore") + + llm: _LlmAttributes | None = None + error: _OtelError | None = None + + +class _OtelSpan(BaseModel): + model_config = ConfigDict(extra="ignore") + + name: str | None = None + status: _OtelStatus | None = None + attributes: _OtelSpanAttributes | None = None + + +class _CallAttributes(BaseModel): + model_config = ConfigDict(extra="ignore") + + otel_span: _OtelSpan | None = None + + +class _Usage(BaseModel): + model_config = ConfigDict(extra="ignore") + + total_tokens: int | None = None + + +class _WeaveSummary(BaseModel): + model_config = ConfigDict(extra="ignore") + + usage: dict[str, _Usage] = {} + + +class WeaveCall(BaseModel): + """One ingested Weave call, reduced to what the scenarios assert on.""" + + model_config = ConfigDict(extra="ignore") + + id: str + op_name: str + started_at: str | None = None + inputs: dict[str, object] = {} + attributes: _CallAttributes | None = None + summary: _WeaveSummary | None = Field(default=None) + + @property + def op(self) -> str: + """The bare op name out of ``weave://///op/:``.""" + return self.op_name.split("/op/")[-1].split(":")[0] + + @property + def status_code(self) -> str | None: + """The OTEL span status, not Weave's own rollup (which reads "success" + even for a span the exporter marked ERROR).""" + span = self.attributes.otel_span if self.attributes else None + return span.status.code if span and span.status else None + + @property + def error(self) -> _OtelError | None: + span = self.attributes.otel_span if self.attributes else None + return span.attributes.error if span and span.attributes else None + + @property + def response_cost(self) -> float | None: + span = self.attributes.otel_span if self.attributes else None + llm = span.attributes.llm if span and span.attributes else None + return llm.response.cost if llm and llm.response else None + + @property + def total_tokens(self) -> int | None: + """Weave keys usage by model, so the total is summed across whatever + models the call reported.""" + if not self.summary or not self.summary.usage: + return None + totals = [usage.total_tokens for usage in self.summary.usage.values() if usage.total_tokens is not None] + return sum(totals) if totals else None + + def mentions(self, needle: str) -> bool: + return needle in json.dumps(self.inputs, default=str) + + +@dataclass(frozen=True, slots=True) +class WeaveReader: + project_id: str + api_key: str + + @property + def _headers(self) -> AuthHeaders: + """Weave authenticates with HTTP Basic as the fixed user ``api``.""" + token = base64.b64encode(f"api:{self.api_key}".encode()).decode() + return AuthHeaders(authorization=f"Basic {token}") + + def _query_body(self, *, since: float, offset: int, op: str) -> _WeaveQueryBody: + return _WeaveQueryBody( + project_id=self.project_id, + filter=_WeaveOpFilter(op_names=[f"weave:///{self.project_id}/op/{op}:*"]), + query=_WeaveQuery( + expr=_WeaveGreaterThan(gt=(_WeaveGetField(get_field="started_at"), _WeaveLiteral(literal=since))) + ), + offset=offset, + ) + + def _page(self, *, since: float, offset: int, op: str) -> tuple[WeaveCall, ...]: + outcome = send( + URL(f"{_WEAVE_TRACE_API}/calls/stream_query"), + headers=self._headers, + json=self._query_body(since=since, offset=offset, op=op), + ) + if not outcome.ok: + pytest.fail( + f"Weave calls query for project {self.project_id!r} failed " + f"({outcome.status_code}): {outcome.body[:300]}" + ) + return tuple(WeaveCall.model_validate_json(line) for line in outcome.body.splitlines() if line.strip()) + + def calls_matching(self, marker: str, *, since: float, op: str = LITELLM_REQUEST_OP) -> tuple[WeaveCall, ...]: + """Every call under ``op`` started after ``since`` whose inputs carry + ``marker``, paging until the window is exhausted. + + More than one is the duplicate-delivery bug, so this never collapses to a + single call. + """ + pages = tuple( + takewhile( + bool, + (self._page(since=since, offset=offset, op=op) for offset in count(0, _PAGE_SIZE)), + ) + ) + return tuple(call for page in pages for call in page if call.mentions(marker)) + + def poll_calls_matching(self, marker: str, *, since: float, op: str = LITELLM_REQUEST_OP) -> tuple[WeaveCall, ...]: + """Poll until the call is readable, then keep re-reading for + WEAVE_SETTLE_SECONDS so a duplicate exported by a later batch flush + cannot hide from the exactly-one assertion. A duplicate ends the settle + early, because more waiting cannot clear it.""" + deadline = time.monotonic() + POLL_TIMEOUT + while time.monotonic() < deadline: + calls = self.calls_matching(marker, since=since, op=op) + if calls: + return self._settled(marker, since=since, op=op, first=calls) + time.sleep(POLL_INTERVAL) + return () + + def _settled(self, marker: str, *, since: float, op: str, first: tuple[WeaveCall, ...]) -> tuple[WeaveCall, ...]: + """A transiently empty re-read never downgrades what was already seen.""" + settle_deadline = time.monotonic() + WEAVE_SETTLE_SECONDS + latest = first # rebind-ok: one settle window, re-read per poll interval + while time.monotonic() < settle_deadline and len(latest) <= 1: + time.sleep(POLL_INTERVAL) + latest = self.calls_matching(marker, since=since, op=op) or latest + return latest + + +def build_weave_reader() -> WeaveReader: + project_id = (os.environ.get("WEAVE_PROJECT_ID") or os.environ.get("WANDB_PROJECT_ID") or "").strip() + api_key = os.environ.get("WANDB_API_KEY", "").strip() + if not project_id or not api_key: + pytest.fail( + "Weave e2e requires WANDB_API_KEY and WEAVE_PROJECT_ID (or WANDB_PROJECT_ID, " + "format /): the test reads the proxy's weave_otel delivery " + "back from the real Weave project; missing credentials is a hard failure, not a skip" + ) + return WeaveReader(project_id=project_id, api_key=api_key) diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 5de49ead3ed..1379ecb4530 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -35,6 +35,8 @@ class KeyLoggingCallbackVars(BaseModel): langfuse_public_key: str | None = None langfuse_secret_key: str | None = None langfuse_host: str | None = None + wandb_api_key: str | None = None + weave_project_id: str | None = None class KeyLoggingCallback(BaseModel):