Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions tests/e2e/CLAUDE.md
Original file line number Diff line number Diff line change
Expand Up @@ -134,9 +134,9 @@ quota_management.<behavior>.<variant>.<assertion>
<budget> key | internal_user | end_user | organization | team_member | tag
| model_max | soft | key_multi_window | team_multi_window
| fallback | spend_counter
<spend_tracking> chat_completions | stream | embeddings | cache_hit | key_rollup
| concurrent_burst | tags | end_user | per_model | failure
| spend_calculate | pagination
<spend_tracking> 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 | routes_to_fallback | reseed_matches_db | logs_cost | zero_cost
Expand Down
1 change: 1 addition & 0 deletions tests/e2e/coverage_registry/quota_management.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,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"}
Expand Down
5 changes: 5 additions & 0 deletions tests/e2e/docker-compose.yml
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,11 @@ configs:
model: openai/text-embedding-3-small
api_key: os.environ/OPENAI_API_KEY

- model_name: openai-responses-codex
litellm_params:
model: openai/gpt-5.3-codex
api_key: os.environ/OPENAI_API_KEY

# v2 auto-router with the LLM complexity classifier. SIMPLE stays on the
# openai backend; every higher tier routes to the anthropic backend, so the
# served deployment (read back from the spend log's model) reveals whether
Expand Down
4 changes: 4 additions & 0 deletions tests/e2e/e2e_gateway.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
unwrap,
)
from models import (
AnthropicMessagesBody,
ChatBody,
ChatResponse,
CustomerDeleteBody,
Expand Down Expand Up @@ -227,6 +228,9 @@ def chat(self, key: str, body: ChatBody) -> Result[ChatResponse]:
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",
Expand Down
1 change: 1 addition & 0 deletions tests/e2e/quota_management/spend_tracking/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@ def _driver_params(provider_model: str, env_var: str) -> LiteLLMParamsBody:
("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"),
)


Expand Down
14 changes: 14 additions & 0 deletions tests/e2e/quota_management/spend_tracking/spend_e2e_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@
)
from e2e_gateway import Gateway, build_gateway
from models import (
AnthropicMessagesBody,
ChatBody,
ChatMessage,
ChatMetadata,
Expand Down Expand Up @@ -119,6 +120,19 @@ def chat_stream(
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.gateway.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.gateway.embed(key, EmbedBody(model=model, input=content))

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down Expand Up @@ -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")

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Proof in PR description does not cover this test

The PR's "Screenshots / Proof of Fix" section describes and demonstrates MCP key-access tests (tests/e2e/mcp/test_mcp_key_access_e2e.py) — a completely different suite. The command shown (python -m pytest tests/e2e/mcp/ -v) and the collected items (2 MCP tests) have no relation to this new spend-tracking test for the /v1/messages bridge. No evidence is supplied that test_streaming_messages_via_responses_bridge_tracks_spend passes against a live proxy, which is the claimed proof-of-fix for LIT-4547.

Rule Used: What: Ensure that any PR claiming to fix an issue ... (source)

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 bridged_costed_row(rows: list[SpendLogRow]) -> bool:
return any(
(r.spend or 0) > 0 and "anthropic_messages" in (r.call_type or "")
for r in rows
)

rows = client.poll_logs_for_key(scoped_key, predicate=bridged_costed_row)
costed = [r for r in rows if (r.spend or 0) > 0]
assert len(costed) == 1, (
f"expected exactly one costed row for the bridged stream, saw {_summarize(rows)}"
)

row = costed[0]
assert "anthropic_messages" in (row.call_type or ""), (
f"row was not billed as a /v1/messages call (call_type {row.call_type!r}); "
f"the bridge must keep the messages billing identity: {_summarize(rows)}"
)
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)}"
)

Comment thread
greptile-apps[bot] marked this conversation as resolved.

@pytest.mark.covers("quota_management.spend_tracking.embeddings.logs_cost")
def test_embedding_writes_nonzero_spend_row(
client: SpendClient, scoped_key: str
Expand Down
Loading