mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-26 01:12:21 +00:00
test(e2e): pin end-user and tag attribution from Codex-style headers on /v1/responses (#43093)
* test(e2e): pin end-user and tag attribution from Codex-style headers on /v1/responses Codex CLI has no body field for the end user, so its config.toml http_headers attach x-litellm-customer-id or x-litellm-end-user-id plus x-litellm-tags to every /v1/responses call. The proxy already honors those headers on the Responses route, but nothing in the e2e stack pinned it. The new case sends that exact wire shape with each standard customer header and fails unless the spend row carries the end user, the tags, the aresponses call type, and a nonzero cost. * test(e2e): assert the customer's /customer/info total matches the header-attributed Responses row --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
parent
25fb7810c2
commit
57bead9842
4 changed files with 93 additions and 4 deletions
|
|
@ -49,6 +49,7 @@
|
|||
- {id: quota_management.spend_tracking.surface_consistency.matches_every_surface, module: quota_management, tier: P1, behavior: spend_tracking, variant: surface_consistency, assertions: [matches_every_surface], exercised_on: [chat_completions], source: "proxy/db/db_spend_update_writer.py", rationale: "One priced request lands the same response_cost on the spend log row, /key/info, /team/info, the usage export's /user/daily/activity/aggregated row, and the litellm_spend_metric Prometheus sample; each is a separate writer, so a rounding, dropped, or double-counted write on one drifts it from the rest (LIT-3620, LIT-5045)"}
|
||||
- {id: quota_management.spend_tracking.tags.attributes_spend, module: quota_management, tier: P1, behavior: spend_tracking, variant: tags, assertions: [attributes_spend], exercised_on: [chat_completions], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "Request tags round-trip to spend rows and tag rollups match tagged logs"}
|
||||
- {id: quota_management.spend_tracking.end_user.attributes_spend, module: quota_management, tier: P1, behavior: spend_tracking, variant: end_user, assertions: [attributes_spend], exercised_on: [chat_completions], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "user= attribution lands the end-user id on the spend row"}
|
||||
- {id: quota_management.spend_tracking.end_user.attributes_responses_header, module: quota_management, tier: P1, behavior: spend_tracking, variant: end_user, assertions: [attributes_responses_header], exercised_on: [responses], source: "proxy/auth/auth_utils.py", rationale: "A /v1/responses call carrying x-litellm-customer-id or x-litellm-end-user-id plus x-litellm-tags, the headers Codex CLI attaches through its config.toml http_headers because it has no body field for the end user, lands the end user and the tags on a costed aresponses spend row whose spend the customer's /customer/info total matches (LIT-8575)"}
|
||||
- {id: quota_management.spend_tracking.per_model.writes_own_rows, module: quota_management, tier: P2, behavior: spend_tracking, variant: per_model, assertions: [writes_own_rows], exercised_on: [chat_completions], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "Each model on a shared key gets its own spend row"}
|
||||
- {id: quota_management.spend_tracking.failure.writes_failure_row, module: quota_management, tier: P1, behavior: spend_tracking, variant: failure, assertions: [writes_failure_row], exercised_on: [chat_completions], source: "proxy/spend_tracking/spend_log_error_logger.py", rationale: "A failed call writes a failure-status spend row"}
|
||||
- {id: quota_management.spend_tracking.failure.writes_normalized_error, module: quota_management, tier: P1, behavior: spend_tracking, variant: failure, assertions: [writes_normalized_error], exercised_on: [chat_completions], source: "litellm_core_utils/error_normalization.py", rationale: "Failure rows carry a stable metadata.error_information.normalized_error key next to the unchanged error_message, so two upstream auth failures with different provider wording share one cluster key a dashboard can group by"}
|
||||
|
|
|
|||
|
|
@ -23,7 +23,7 @@ proxy + SpendLogs rows. Status: `covered` / `partial` / `gap`.
|
|||
| per-model / per-provider attribution | `test_spend_tracking_utils.py` | unit | covered | yes (`test_each_model_on_a_shared_key_gets_its_own_row`) |
|
||||
| field population (model/tokens/api_key/team/org) | `test_spend_tracking_utils.py` | unit | partial | yes (asserts real values) |
|
||||
| `request_tags` propagation | `test_db_spend_update_writer.py` | unit | partial | yes (`test_request_tags_round_trip`) |
|
||||
| `end_user` attribution | unit | unit | partial | yes (`test_end_user_spend_attributed_on_row`) |
|
||||
| `end_user` attribution | unit | unit | partial | yes (`test_end_user_spend_attributed_on_row`, `test_end_user_header_attributes_responses_row`) |
|
||||
|
||||
## Cost calculation by modality
|
||||
|
||||
|
|
@ -67,6 +67,7 @@ proxy + SpendLogs rows. Status: `covered` / `partial` / `gap`.
|
|||
| `test_request_tags_round_trip` | tags persist onto the row |
|
||||
| `test_tag_spend_matches_sum_of_tagged_logs` | `/spend/tags` SUM/COUNT == tagged rows |
|
||||
| `test_end_user_spend_attributed_on_row` | `end_user` attributed + costed |
|
||||
| `test_end_user_header_attributes_responses_row` | `x-litellm-customer-id` / `x-litellm-end-user-id` + `x-litellm-tags` headers on `/v1/responses` (the Codex CLI `http_headers` shape) attributed + tagged + costed, and `/customer/info` spend equals the row |
|
||||
| `test_each_model_on_a_shared_key_gets_its_own_row` | per-model/provider rows, correct model + cost, distinct request_ids matching response id |
|
||||
| `test_failure_call_writes_failure_status_row` | failed call -> `status=failure`, `spend=0` |
|
||||
| `test_spend_calculate_returns_nonzero_cost` | cost-map smoke (no batch wait) |
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@ from typing import Final
|
|||
|
||||
from e2e_config import unique_marker
|
||||
from e2e_http import (
|
||||
AuthHeaders,
|
||||
FileUploadForm,
|
||||
Headers,
|
||||
NoBody,
|
||||
|
|
@ -36,6 +37,7 @@ from models import (
|
|||
ChatMessage,
|
||||
ChatMetadata,
|
||||
ChatResponse,
|
||||
CustomerInfoParams,
|
||||
DateRangeParams,
|
||||
EmbedBody,
|
||||
EmbedResponse,
|
||||
|
|
@ -63,9 +65,10 @@ METRICS_PATH: Final = "/metrics/"
|
|||
|
||||
__all__ = [
|
||||
"BatchCreateBody",
|
||||
"BatchObject",
|
||||
"CallbackLogMetadata",
|
||||
"CallbackLogPayload",
|
||||
"BatchObject",
|
||||
"ClientAttributionHeaders",
|
||||
"DailyActivityKeyBreakdown",
|
||||
"FileObject",
|
||||
"ProbeResult",
|
||||
|
|
@ -80,6 +83,16 @@ __all__ = [
|
|||
]
|
||||
|
||||
|
||||
class ClientAttributionHeaders(AuthHeaders):
|
||||
"""The attribution headers a coding agent attaches to every call from its own
|
||||
config (Codex CLI's config.toml ``http_headers``, Claude Code's
|
||||
``ANTHROPIC_CUSTOM_HEADERS``) because it has no body field for the end user."""
|
||||
|
||||
x_litellm_customer_id: str | None = Field(default=None, alias="x-litellm-customer-id")
|
||||
x_litellm_end_user_id: str | None = Field(default=None, alias="x-litellm-end-user-id")
|
||||
x_litellm_tags: str | None = Field(default=None, alias="x-litellm-tags")
|
||||
|
||||
|
||||
class GeminiApiKeyHeaders(Headers):
|
||||
x_goog_api_key: str = Field(serialization_alias="x-goog-api-key")
|
||||
content_type: str = Field(default="application/json", serialization_alias="Content-Type")
|
||||
|
|
@ -222,6 +235,10 @@ class TeamInfoSpendResponse(BaseModel):
|
|||
team_info: TeamInfoSpend
|
||||
|
||||
|
||||
class CustomerSpendResponse(BaseModel):
|
||||
spend: float | None = None
|
||||
|
||||
|
||||
def _chat_body(
|
||||
model: str,
|
||||
content: str,
|
||||
|
|
@ -371,6 +388,31 @@ class SpendClient:
|
|||
)
|
||||
return outcome.result if isinstance(outcome, Converged) else outcome.last_result
|
||||
|
||||
def customer_spend(self, customer_id: str) -> float:
|
||||
"""0.0 until the spend writer has upserted the end-user row, which /customer/info 404s before."""
|
||||
looked_up: Final = self.proxy.transport.get(
|
||||
"/customer/info",
|
||||
headers=self.proxy.transport.master,
|
||||
params=CustomerInfoParams(end_user_id=customer_id),
|
||||
response_type=CustomerSpendResponse,
|
||||
)
|
||||
match looked_up:
|
||||
case Success(data=data):
|
||||
return data.spend or 0.0
|
||||
case _:
|
||||
return 0.0
|
||||
|
||||
def poll_customer_spend(self, customer_id: str, *, minimum: float = 0.0) -> float:
|
||||
outcome: Final = await_converged(
|
||||
lambda: self.customer_spend(customer_id),
|
||||
converged=lambda spend: spend > minimum,
|
||||
timeout=self.proxy.poll_timeout,
|
||||
interval=self.proxy.poll_interval,
|
||||
now=time.monotonic,
|
||||
sleep=time.sleep,
|
||||
)
|
||||
return outcome.result if isinstance(outcome, Converged) else outcome.last_result
|
||||
|
||||
def scrape_metrics(self) -> Mapping[str, ProbeResult]:
|
||||
"""GET /metrics/ on every replica in PROXY_REPLICA_URLS, keyed by replica. The
|
||||
counter is per pod, so the union of the replicas is the fleet's exposition; the
|
||||
|
|
@ -479,9 +521,12 @@ class SpendClient:
|
|||
)
|
||||
|
||||
def send_responses(self, key: str, model: str, content: str) -> StreamingResponse:
|
||||
return self.send_responses_with_headers(self.proxy.transport.bearer(key), model, content)
|
||||
|
||||
def send_responses_with_headers(self, headers: AuthHeaders, model: str, content: str) -> StreamingResponse:
|
||||
return self.proxy.transport.send(
|
||||
"/v1/responses",
|
||||
headers=self.proxy.transport.bearer(key),
|
||||
headers=headers,
|
||||
json=ResponsesBody(model=model, input=content),
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -24,7 +24,14 @@ import pytest
|
|||
from e2e_http import RateLimitedError, Success
|
||||
from lifecycle import ResourceManager
|
||||
from models import KeyGenerateBody, LiteLLMParamsBody, SpendLogs, SpendLogsParams
|
||||
from spend_e2e_client import SpendClient, SpendLogRow, is_ok, unique_marker, unwrap
|
||||
from spend_e2e_client import (
|
||||
ClientAttributionHeaders,
|
||||
SpendClient,
|
||||
SpendLogRow,
|
||||
is_ok,
|
||||
unique_marker,
|
||||
unwrap,
|
||||
)
|
||||
|
||||
pytestmark = pytest.mark.e2e
|
||||
|
||||
|
|
@ -425,6 +432,41 @@ def test_end_user_spend_attributed_on_row(
|
|||
assert (row.spend or 0) > 0, f"end-user row should cost > 0: {_summarize(rows)}"
|
||||
|
||||
|
||||
@pytest.mark.covers("quota_management.spend_tracking.end_user.attributes_responses_header")
|
||||
@pytest.mark.parametrize("header", ["x-litellm-customer-id", "x-litellm-end-user-id"])
|
||||
def test_end_user_header_attributes_responses_row(
|
||||
client: SpendClient, scoped_key: str, resources: ResourceManager, header: str
|
||||
) -> None:
|
||||
"""Codex CLI has no body field for the end user, so its config.toml http_headers
|
||||
attach the customer header (and x-litellm-tags) to every /v1/responses call.
|
||||
A regression that stops reading either header on the Responses route, drops the
|
||||
tags, costs the row at zero, or leaves the customer's own spend total behind the
|
||||
row fails here."""
|
||||
customer = resources.customer(f"e2e-codex-{unique_marker()}")
|
||||
tag = f"codex-{unique_marker()}"
|
||||
headers = ClientAttributionHeaders.model_validate(
|
||||
{"authorization": f"Bearer {scoped_key}", header: customer, "x-litellm-tags": tag}
|
||||
)
|
||||
sent = client.send_responses_with_headers(
|
||||
headers, "openai-responses-codex", f"one word {unique_marker()}"
|
||||
)
|
||||
assert sent.ok, f"/v1/responses failed with {sent.status_code}: {sent.body[:300]}"
|
||||
|
||||
rows = client.poll_logs_for_key(
|
||||
scoped_key, predicate=lambda rs: any(r.end_user == customer for r in rs)
|
||||
)
|
||||
row = _require_row(
|
||||
rows, lambda r: r.end_user == customer, f"attributed to end_user {customer!r} via {header}"
|
||||
)
|
||||
assert row.call_type == "aresponses", f"row is not a Responses row: {_summarize(rows)}"
|
||||
assert tag in (row.request_tags or []), f"tag {tag!r} missing from {row.request_tags}"
|
||||
assert (row.spend or 0) > 0, f"end-user row should cost > 0: {_summarize(rows)}"
|
||||
customer_total = client.poll_customer_spend(customer)
|
||||
assert _approx_equal(customer_total, row.spend or 0), (
|
||||
f"/customer/info spend {customer_total} != the row's {row.spend}: {_summarize(rows)}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.covers("quota_management.spend_tracking.per_model.writes_own_rows")
|
||||
def test_each_model_on_a_shared_key_gets_its_own_row(
|
||||
client: SpendClient, scoped_key: str
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue