diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 1fa78275085..47f3c74d7f1 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -146,9 +146,9 @@ quota_management... key | internal_user | end_user | organization | team | team_member | tag | model_max | soft | key_multi_window | team_multi_window | fallback | spend_counter - chat_completions | stream | embeddings | cache_hit | key_rollup - | concurrent_burst | tags | end_user | per_model | failure - | spend_calculate | pagination + chat_completions | stream | messages_bridge | embeddings + | cache_hit | key_rollup | concurrent_burst | tags | end_user + | per_model | failure | spend_calculate | pagination assertion : blocks_over_limit | resets_after_window | headers_report_remaining | picks_under_tpm | blocks_then_resets | resets_windows_independently | alerts_without_blocking | isolates_per_model | isolates_per_member | enforced_across_keys | routes_to_fallback diff --git a/tests/e2e/coverage_registry/quota_management.yaml b/tests/e2e/coverage_registry/quota_management.yaml index 633351ca97c..7e71f2bd3d1 100644 --- a/tests/e2e/coverage_registry/quota_management.yaml +++ b/tests/e2e/coverage_registry/quota_management.yaml @@ -29,6 +29,7 @@ - {id: quota_management.budget.spend_counter.reseed_matches_db, module: quota_management, tier: P2, behavior: budget, variant: spend_counter, assertions: [reseed_matches_db], exercised_on: [chat_completions], source: "proxy/spend_tracking/budget_reservation.py", rationale: "Concurrent cold-counter reseeds keep the enforcement counter equal to DB spend (#26829)"} - {id: quota_management.spend_tracking.chat_completions.logs_cost, module: quota_management, tier: P0, behavior: spend_tracking, variant: chat_completions, assertions: [logs_cost], exercised_on: [chat_completions], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "A paid chat call writes a nonzero spend row"} - {id: quota_management.spend_tracking.stream.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: stream, assertions: [logs_cost], exercised_on: [chat_completions], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "Streaming responses aggregate token counts into a spend row"} +- {id: quota_management.spend_tracking.messages_bridge.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: messages_bridge, assertions: [logs_cost], exercised_on: [messages], source: "llms/anthropic/experimental_pass_through/responses_adapters/handler.py", rationale: "A streaming /v1/messages request served by an openai-provider model is bridged through the anthropic-messages -> Responses adapter and must aggregate the consumed SSE stream into one spend row with nonzero cost and token counts, attributed to custom_llm_provider openai under call_type anthropic_messages"} - {id: quota_management.spend_tracking.embeddings.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: embeddings, assertions: [logs_cost], exercised_on: [embeddings], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "Embedding calls write nonzero spend rows"} - {id: quota_management.spend_tracking.cache_hit.zero_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: cache_hit, assertions: [zero_cost], exercised_on: [chat_completions], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "A response-cache hit logs at zero cost with the cache-hit marker"} - {id: quota_management.spend_tracking.key_rollup.matches_sum_of_logs, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_rollup, assertions: [matches_sum_of_logs], exercised_on: [chat_completions], source: "proxy/db/db_spend_update_writer.py", rationale: "A key's rolled-up spend equals the sum of its log rows"} diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index c466d415d0e..7eb86046375 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -232,6 +232,9 @@ class ProxyClient: def chat_stream(self, key: str, body: ChatBody) -> StreamingResponse: return self.transport.stream("/chat/completions", headers=self.transport.bearer(key), json=body) + def messages_stream(self, key: str, body: AnthropicMessagesBody) -> StreamingResponse: + return self.transport.stream("/v1/messages", headers=self.transport.bearer(key), json=body) + def embed(self, key: str, body: EmbedBody) -> Result[EmbedResponse]: return self.transport.post( "/embeddings", diff --git a/tests/e2e/quota_management/spend_tracking/conftest.py b/tests/e2e/quota_management/spend_tracking/conftest.py index c31e6b3c090..0597c9af400 100644 --- a/tests/e2e/quota_management/spend_tracking/conftest.py +++ b/tests/e2e/quota_management/spend_tracking/conftest.py @@ -35,6 +35,7 @@ DRIVER_MODELS: tuple[tuple[str, str, str], ...] = ( ("gemini-2.5-flash", "gemini/gemini-2.5-flash", "GEMINI_API_KEY"), ("claude-haiku-4-5", "anthropic/claude-haiku-4-5", "ANTHROPIC_API_KEY"), ("openai-text-embedding-3-small", "openai/text-embedding-3-small", "OPENAI_API_KEY"), + ("openai-responses-codex", "openai/gpt-5.3-codex", "OPENAI_API_KEY"), ) diff --git a/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py b/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py index 29ca5eb2ce6..0d49869aa91 100644 --- a/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py +++ b/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py @@ -29,6 +29,7 @@ from e2e_http import ( ) from proxy_client import ProxyClient from models import ( + AnthropicMessagesBody, ChatBody, ChatMessage, ChatMetadata, @@ -119,6 +120,19 @@ class SpendClient: key, _chat_body(model, content, max_tokens=max_tokens, stream=True) ) + def messages_stream( + self, key: str, model: str, content: str, *, max_tokens: int + ) -> StreamingResponse: + return self.proxy.messages_stream( + key, + AnthropicMessagesBody( + model=model, + messages=[ChatMessage(role="user", content=content)], + max_tokens=max_tokens, + stream=True, + ), + ) + def embed(self, key: str, model: str, content: str) -> Result[EmbedResponse]: return self.proxy.embed(key, EmbedBody(model=model, input=content)) diff --git a/tests/e2e/quota_management/spend_tracking/test_spend_tracking_e2e.py b/tests/e2e/quota_management/spend_tracking/test_spend_tracking_e2e.py index 465046e89af..d43d8e94898 100644 --- a/tests/e2e/quota_management/spend_tracking/test_spend_tracking_e2e.py +++ b/tests/e2e/quota_management/spend_tracking/test_spend_tracking_e2e.py @@ -41,6 +41,8 @@ def _summarize(rows: list[SpendLogRow]) -> list[dict[str, object]]: "spend", "status", "cache_hit", + "call_type", + "custom_llm_provider", "prompt_tokens", "completion_tokens", "total_tokens", @@ -122,6 +124,75 @@ def test_streaming_chat_completion_tracks_spend( assert (row.total_tokens or 0) == prompt + completion +@pytest.mark.covers("quota_management.spend_tracking.messages_bridge.logs_cost") +def test_streaming_messages_via_responses_bridge_tracks_spend( + client: SpendClient, scoped_key: str +) -> None: + """A streaming anthropic-format /v1/messages request served by an openai-provider + model is bridged through litellm's anthropic-messages -> Responses adapter, and + consuming the whole SSE stream writes exactly one costed spend row. + + The deployment is a Responses-only OpenAI model (gpt-5.3-codex, exposed only on + /v1/responses), so a served call could not have taken the chat-completions bridge: + that path would 404 at OpenAI on an endpoint the model does not have. The row + proving the Responses path carries custom_llm_provider "openai" (the openai + backend served it) under a call_type that keeps the /v1/messages billing identity + (never a chat call_type), with nonzero cost and prompt/completion tokens that the + bridge must aggregate out of the consumed stream. + """ + result = client.messages_stream( + scoped_key, + "openai-responses-codex", + f"reply with exactly one word {unique_marker()}", + max_tokens=64, + ) + assert ( + result.ok + ), f"bridged /v1/messages stream failed (status {result.status_code}): {result.body[:300]}" + assert result.is_streaming, ( + f"expected an SSE stream from /v1/messages, got content-type " + f"{result.content_type!r}" + ) + assert result.chunks > 0, "no SSE events were consumed from the /v1/messages stream" + assert ( + result.stream_error is None + ), f"the /v1/messages stream carried an error event: {result.stream_error}" + + def is_bridged_costed(row: SpendLogRow) -> bool: + return (row.spend or 0) > 0 and "anthropic_messages" in (row.call_type or "") + + rows = client.poll_logs_for_key( + scoped_key, predicate=lambda rs: any(is_bridged_costed(r) for r in rs) + ) + costed = [r for r in rows if (r.spend or 0) > 0] + bridged = [r for r in costed if is_bridged_costed(r)] + assert bridged == costed, ( + f"a costed row was not billed as a /v1/messages call (wrong call_type); " + f"the bridge must keep the messages billing identity: {_summarize(rows)}" + ) + assert len(bridged) == 1, ( + f"expected exactly one costed row for the bridged stream, saw {_summarize(rows)}" + ) + + row = bridged[0] + assert row.custom_llm_provider == "openai", ( + f"bridged row not attributed to the openai Responses backend " + f"(custom_llm_provider {row.custom_llm_provider!r}): {_summarize(rows)}" + ) + assert "codex" in (row.model or ""), ( + f"row model {row.model!r} is not the Responses-only codex deployment" + ) + + prompt = row.prompt_tokens or 0 + completion = row.completion_tokens or 0 + assert ( + prompt > 0 and completion > 0 + ), f"bridged stream tokens not tracked: {_summarize(rows)}" + assert (row.total_tokens or 0) == prompt + completion, ( + f"token arithmetic broken on the bridged row: {_summarize(rows)}" + ) + + @pytest.mark.covers("quota_management.spend_tracking.embeddings.logs_cost") def test_embedding_writes_nonzero_spend_row( client: SpendClient, scoped_key: str