diff --git a/litellm/integrations/shadow_eval_logger.py b/litellm/integrations/shadow_eval_logger.py index 66295ea2319..eaaa6cfa088 100644 --- a/litellm/integrations/shadow_eval_logger.py +++ b/litellm/integrations/shadow_eval_logger.py @@ -16,7 +16,11 @@ Flow per sampled request (all in a detached background task, zero added latency) 3. Write a ``LiteLLM_ShadowEvalVerdict`` row and bump the job's counters. Shadow and judge calls carry ``shadow_eval_internal`` metadata so this logger -ignores its own traffic and cannot recurse. +ignores its own traffic and cannot recurse. They also carry the shadowed key's +identity metadata, so the provider spend they incur is attributed to that key +(and its team/org/user) and counts against every budget that key is subject to, +including the global proxy budget. A job additionally stops itself once its judge +spend reaches a multiple of the estimate quoted when it started. """ import asyncio @@ -26,6 +30,7 @@ import random import re from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass +from datetime import datetime, timezone from types import MappingProxyType from typing import TYPE_CHECKING, Final, TypeAlias @@ -34,6 +39,8 @@ from pydantic import BaseModel, TypeAdapter import litellm from litellm._logging import verbose_logger from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.internal_call_metadata import sanitized_forwardable_call_metadata +from litellm.types.utils import SHADOW_EVAL_JUDGE_CALL_ORIGIN, SHADOW_EVAL_ROUTER_CALL_ORIGIN if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient @@ -65,6 +72,14 @@ JUDGE_MAX_OUTPUT_TOKENS: Final = 500 _SEEN_FLUSH_INTERVAL_SECONDS: Final = 10.0 +# A job stops sampling once its judge spend reaches this multiple of the estimate shown +# at start. The headroom absorbs an estimate that undershot the real traffic mix without +# letting the job run away; the floor keeps a cent-sized estimate from stopping a job on +# its first verdict. cost_actual is read from the job cache, so overshoot is bounded by +# _JOB_CACHE_TTL_SECONDS worth of judge calls. +_SPEND_CAP_MULTIPLIER: Final = 1.5 +_SPEND_CAP_FLOOR_USD: Final = 1.0 + _EMPTY_METADATA: Final[Mapping[str, object]] = MappingProxyType({}) PAIRWISE_JUDGE_SYSTEM_PROMPT: Final = """You are an impartial quality judge comparing two responses to the same conversation. @@ -156,6 +171,21 @@ class ActiveShadowEvalJob: shadow_percentage: float judge_model: str status: str + cost_estimate: float | None = None + cost_actual: float = 0.0 + + +def _job_is_over_spend_cap(job: ActiveShadowEvalJob) -> bool: + """Whether the job has spent past what its start-time estimate justifies. + + Budgets bound what the *key* may spend; this bounds what a single eval may spend + even under a generous budget, so a bad estimate or a traffic spike cannot quietly + turn a "$3 eval" into a much larger bill. A job with no estimate (older rows, + unpriced judge model) has nothing to be a multiple of, so it is uncapped. + """ + if job.cost_estimate is None or job.cost_estimate <= 0.0: + return False + return job.cost_actual >= max(job.cost_estimate * _SPEND_CAP_MULTIPLIER, _SPEND_CAP_FLOOR_USD) _JobCache: TypeAlias = "dict[str, tuple[float, ActiveShadowEvalJob | None]]" # mutable-ok: TTL cache @@ -208,6 +238,9 @@ class ShadowEvalLogger(CustomLogger): job: Final = await self._get_active_job(api_key_hash) if job is None: return + if _job_is_over_spend_cap(job): + await self._stop_job_over_spend_cap(job) + return request_id: Final = payload.get("id") or "" if not request_id: return @@ -238,6 +271,7 @@ class ShadowEvalLogger(CustomLogger): model_parameters=MappingProxyType( dict(payload.get("model_parameters") or {}) # mutable-ok: frozen snapshot ), + parent_metadata=MappingProxyType(dict(request_metadata)), # mutable-ok: frozen snapshot ) ) task.add_done_callback(lambda _: setattr(self, "_inflight_shadow_tasks", self._inflight_shadow_tasks - 1)) @@ -269,6 +303,8 @@ class ShadowEvalLogger(CustomLogger): shadow_percentage=float(record.shadow_percentage), judge_model=str(record.judge_model), status=str(record.status), + cost_estimate=float(record.cost_estimate) if record.cost_estimate is not None else None, + cost_actual=float(record.cost_actual or 0.0), ) if record is not None else None @@ -279,6 +315,38 @@ class ShadowEvalLogger(CustomLogger): self._job_cache[api_key_hash] = (now, job) return job + async def _stop_job_over_spend_cap(self, job: ActiveShadowEvalJob) -> None: + """Flip a runaway job to completed, keeping the verdicts it already produced. + + Guarded on the job still being pending/running so two pods breaching the cap on + the same job cannot resurrect one an admin stopped in between. + """ + prisma: Final = self._prisma_provider() + if prisma is None: + return + self._job_cache = { # mutable-ok: TTL cache + k: v for k, v in self._job_cache.items() if v[1] is None or v[1].id != job.id + } + verbose_logger.info( + "shadow_eval: stopping job %s, spend $%.4f reached the cap for its $%.4f estimate", + job.id, + job.cost_actual, + job.cost_estimate or 0.0, + ) + try: + await prisma.db.litellm_shadowevaljob.update_many( + where={ # mutable-ok: Prisma filter + "id": job.id, + "status": {"in": ["pending", "running"]}, # mutable-ok: Prisma filter + }, + data={ # mutable-ok: Prisma payload + "status": "completed", + "completed_at": datetime.now(timezone.utc), + }, + ) + except Exception as e: # noqa: BLE001 # logging hooks must never fail the request + verbose_logger.debug("shadow_eval: failed to stop over-cap job %s: %s", job.id, e) + async def _flush_seen_counts(self) -> None: prisma: Final = self._prisma_provider() if prisma is None: @@ -304,6 +372,7 @@ class ShadowEvalLogger(CustomLogger): response_obj: object, real_model: str, model_parameters: Mapping[str, object], + parent_metadata: Mapping[str, object], ) -> None: """Detached background task: shadow call -> blind judge -> verdict row.""" prisma: Final = self._prisma_provider() @@ -312,7 +381,7 @@ class ShadowEvalLogger(CustomLogger): if not real_text or not messages: return - shadow: Final = await self._call_router_shadow(job.router_name, messages, model_parameters) + shadow: Final = await self._call_router_shadow(job.router_name, messages, model_parameters, parent_metadata) if shadow is None: await self._bump_failed(job.id) return @@ -323,6 +392,7 @@ class ShadowEvalLogger(CustomLogger): messages=messages, real_text=real_text, shadow_text=shadow_text, + parent_metadata=parent_metadata, ) if verdict is None: await self._bump_failed(job.id) @@ -373,16 +443,23 @@ class ShadowEvalLogger(CustomLogger): verbose_logger.debug("shadow_eval: failed_count increment failed: %s", e) async def _call_router_shadow( - self, router_name: str, messages: Sequence[Mapping[str, object]], model_parameters: Mapping[str, object] + self, + router_name: str, + messages: Sequence[Mapping[str, object]], + model_parameters: Mapping[str, object], + parent_metadata: Mapping[str, object], ) -> tuple[str, str, str | None, int | None] | None: """Send the prompt through the auto-router; return (text, model, tier, completion_tokens).""" router: Final = self._router_provider() if router is None: verbose_logger.debug("shadow_eval: no router available") return None - # The router's pre-routing hook writes its routing decision into this - # metadata dict; read it back after the call for tier attribution. - shadow_metadata: Final[dict[str, object]] = { # mutable-ok: router writes back + # Carries the shadowed key's identity so this call's real provider spend is + # attributed to, and budget-checked against, the key whose traffic it copies. + # The router's pre-routing hook also writes its routing decision into this dict; + # read it back after the call for tier attribution. + attribution: Final = sanitized_forwardable_call_metadata(parent_metadata, SHADOW_EVAL_ROUTER_CALL_ORIGIN) + shadow_metadata: Final[dict[str, object]] = attribution | { # mutable-ok: router writes back SHADOW_EVAL_INTERNAL_MARKER: True } shadow_params: Final = { # mutable-ok: splatted as kwargs @@ -417,6 +494,7 @@ class ShadowEvalLogger(CustomLogger): messages: Sequence[Mapping[str, object]], real_text: str, shadow_text: str, + parent_metadata: Mapping[str, object], ) -> tuple[str, float, str, float] | None: """Blind pairwise judge. Returns (preference, confidence, reasoning, cost).""" real_is_a: Final = random.random() < 0.5 @@ -434,6 +512,9 @@ class ShadowEvalLogger(CustomLogger): f"Response B:\n{response_b[:_MAX_JUDGE_CHARS]}\n\n" "Which response is better?" ) + judge_metadata: Final = sanitized_forwardable_call_metadata(parent_metadata, SHADOW_EVAL_JUDGE_CALL_ORIGIN) | { + SHADOW_EVAL_INTERNAL_MARKER: True + } try: response: Final = await litellm.acompletion( model=judge_model, @@ -443,7 +524,7 @@ class ShadowEvalLogger(CustomLogger): ], temperature=0, max_tokens=JUDGE_MAX_OUTPUT_TOKENS, - metadata={SHADOW_EVAL_INTERNAL_MARKER: True}, # mutable-ok: SDK metadata + metadata=judge_metadata, ) except Exception as e: # noqa: BLE001 # judge outages are a counted failure, not a crash verbose_logger.debug("shadow_eval: judge call failed: %s", e) diff --git a/litellm/litellm_core_utils/internal_call_metadata.py b/litellm/litellm_core_utils/internal_call_metadata.py new file mode 100644 index 00000000000..fa97aea8ff2 --- /dev/null +++ b/litellm/litellm_core_utils/internal_call_metadata.py @@ -0,0 +1,107 @@ +"""Metadata a request forwards to the internal LLM sub-calls it triggers. + +Internal features (the auto-router's classifier and its keyword embeddings, shadow +eval's shadow + judge calls) bill real provider spend that nobody typed a prompt for. +That spend must land on the same key/team/org/user as the request that caused it, so +the sub-call has to carry the caller's identity metadata: the proxy's cost callback +reads ``user_api_key`` / ``user_api_key_team_id`` / ... straight off the call's +metadata, and drops the spend log entirely when they are missing. + +Two things must never be forwarded as-is: + +* ``user_api_key_budget_reservation`` (and the same reservation nested inside + ``user_api_key_auth``) belongs to the parent completion, not to the sub-call. If a + sub-call's cost callback sees it, that callback finalizes the reservation and the + parent's own callback then skips incrementing the key/team budget counters, losing + the parent's spend. ``user_api_key_auth`` itself is kept, sanitized, because model + access-group filtering needs it. +* The origin of the sub-call. ``INTERNAL_CALL_ORIGIN_METADATA_KEY`` stamps which + feature made the call so a spend-log row can say it is not traffic the caller sent. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Final + +from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY +from litellm.types.utils import InternalCallOrigin + +BUDGET_RESERVATION_METADATA_KEYS: Final = frozenset({"user_api_key_budget_reservation"}) + +_USER_API_KEY_AUTH_KEY: Final = "user_api_key_auth" + +# The caller-identity subset a detached sub-call needs to be attributed and +# budget-checked exactly like the request that spawned it. Everything else on the +# parent's metadata (routing decisions, guardrail state, the standard logging object, +# the proxy request body) describes the parent call and would be a lie on a sub-call +# that runs after it has already returned. +FORWARDABLE_IDENTITY_METADATA_KEYS: Final = frozenset( + { + "user_api_key", + "user_api_key_hash", + "user_api_key_alias", + "user_api_key_team_id", + "user_api_key_org_id", + "user_api_key_user_id", + "user_api_key_end_user_id", + _USER_API_KEY_AUTH_KEY, + } +) + + +def sanitize_user_api_key_auth(auth: object) -> object: + """Copy of the auth object with its budget reservation removed. + + The proxy's cost callback falls back to reading the reservation from inside the + auth object when the top-level key is absent, so forwarding it unsanitized + re-creates the double-finalization that stripping the top-level key prevents. + """ + if isinstance(auth, dict): + return {k: v for k, v in auth.items() if k != "budget_reservation"} # mutable-ok: SDK metadata value + reservation: Final[object] = getattr(auth, "budget_reservation", None) + model_copy: Final[object] = getattr(auth, "model_copy", None) + if reservation is not None and callable(model_copy): + return model_copy(update={"budget_reservation": None}) # mutable-ok: pydantic update payload + return auth + + +def _sanitized(parent_metadata: Mapping[str, object]) -> dict[str, object]: # mutable-ok: SDK metadata kwarg + return { # mutable-ok: SDK metadata kwarg + k: sanitize_user_api_key_auth(v) if k == _USER_API_KEY_AUTH_KEY else v + for k, v in parent_metadata.items() + if k not in BUDGET_RESERVATION_METADATA_KEYS + } + + +def forwarded_internal_call_metadata( + parent_metadata: Mapping[str, object] | None, + call_origin: InternalCallOrigin, +) -> dict[str, object]: # mutable-ok: SDK metadata kwarg + """Parent metadata, minus its budget reservation, stamped with the sub-call's origin. + + For sub-calls made *inside* the parent request (the auto-router classifier and its + embeddings), where the parent's full context still describes the call being made. + A parent with no metadata at all had nothing to attribute in the first place, so + the sub-call is left unstamped rather than carrying a lone origin marker. + """ + if not parent_metadata: + return {} # mutable-ok: SDK metadata kwarg + return _sanitized(parent_metadata) | { # mutable-ok: SDK metadata kwarg + INTERNAL_CALL_ORIGIN_METADATA_KEY: call_origin + } + + +def sanitized_forwardable_call_metadata( + parent_metadata: Mapping[str, object], + call_origin: InternalCallOrigin, +) -> dict[str, object]: # mutable-ok: SDK metadata kwarg + """Just the caller's identity, stamped with the sub-call's origin. + + For sub-calls detached from the parent request (shadow eval), which outlive it and + must not inherit per-request state such as its routing decision or logging payload. + The origin stamp is unconditional here: these calls are always internal, whether or + not there was an identity to forward. + """ + identity: Final = {k: v for k, v in parent_metadata.items() if k in FORWARDABLE_IDENTITY_METADATA_KEYS} + return _sanitized(identity) | {INTERNAL_CALL_ORIGIN_METADATA_KEY: call_origin} # mutable-ok: SDK metadata kwarg diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index f20dc8dc76e..ba5b013b6fe 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -598,6 +598,8 @@ class _ShadowEvalJobRow(BaseModel): id: str status: ShadowEvalStatus router_name: str + api_key_id: str + team_id: str | None = None shadow_percentage: float request_count: int completed_count: int @@ -617,6 +619,8 @@ def _job_to_response(record: object, results: ShadowEvalResult | None) -> GetSha job_id=row.id, status=row.status, router_name=row.router_name, + api_key_id=row.api_key_id, + team_id=row.team_id, shadow_percentage=row.shadow_percentage, request_count=row.request_count, completed_count=row.completed_count, diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index f6ced0bb9d2..ffcb0fc471e 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -26,8 +26,9 @@ from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, cast from pydantic import BaseModel, create_model from litellm._logging import verbose_router_logger -from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY, RETURN_RAW_MODEL_NAME_METADATA_KEY +from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.internal_call_metadata import forwarded_internal_call_metadata from litellm.llms.base_llm.base_utils import type_to_response_format_param from litellm.types.utils import ( AUTOROUTER_CLASSIFIER_CALL_ORIGIN, @@ -172,38 +173,8 @@ def _append_custom_keywords(base_keywords: list[str], custom_keywords: list[str] return [*base_keywords, *deduped_custom.values()] -# Metadata keys that carry only the parent request's budget reservation state. These -# must not reach internal sub-calls (classifier, embedding): the reservation belongs to -# the routed completion being decided on, not to the sub-call itself, and forwarding it -# would let the sub-call's cost callback finalize the reservation, causing the routed -# completion's callback to skip incrementing key/team budget counters. -# -# Note: user_api_key_auth itself is intentionally kept; it is required by -# _filter_deployments_by_model_access_groups to scope embedding/classifier model -# selection to the caller's authorized access groups. It is forwarded as a sanitized -# copy with its budget_reservation sub-field removed, because the proxy cost callback -# (_get_budget_reservation_from_metadata) falls back to reading the reservation from -# inside the auth object when the top-level key is absent; forwarding it unsanitized -# would re-create the exact double-finalization this stripping exists to prevent. -_BUDGET_RESERVATION_METADATA_KEYS: Final = frozenset({"user_api_key_budget_reservation"}) - - -def _sanitize_user_api_key_auth(auth: Any) -> Any: - if isinstance(auth, dict): - return {k: v for k, v in auth.items() if k != "budget_reservation"} - if getattr(auth, "budget_reservation", None) is not None and hasattr(auth, "model_copy"): - return auth.model_copy(update={"budget_reservation": None}) - return auth - - -def _classifier_call_metadata(metadata: dict[str, Any] | None) -> dict[str, Any]: - if not metadata: - return {} - return { - k: _sanitize_user_api_key_auth(v) if k == "user_api_key_auth" else v - for k, v in metadata.items() - if k not in _BUDGET_RESERVATION_METADATA_KEYS - } | {INTERNAL_CALL_ORIGIN_METADATA_KEY: AUTOROUTER_CLASSIFIER_CALL_ORIGIN} +def _classifier_call_metadata(metadata: Mapping[str, Any] | None) -> dict[str, Any]: + return forwarded_internal_call_metadata(metadata, AUTOROUTER_CLASSIFIER_CALL_ORIGIN) def _parent_session_kwargs(request_kwargs: Mapping[str, Any] | None) -> Mapping[str, Any]: diff --git a/litellm/types/management_endpoints/auto_router_endpoints.py b/litellm/types/management_endpoints/auto_router_endpoints.py index 313a97a88e0..3d04a48c3cf 100644 --- a/litellm/types/management_endpoints/auto_router_endpoints.py +++ b/litellm/types/management_endpoints/auto_router_endpoints.py @@ -158,7 +158,12 @@ DEFAULT_SHADOW_EVAL_JUDGE_MODEL: Final[str] = "anthropic/claude-sonnet-5" class StartShadowEvalRequest(BaseModel): """Start shadowing a deployment's traffic through an auto-router for comparison.""" - api_key_id: str = Field(description="The hashed virtual key whose traffic will be shadowed") + api_key_id: str = Field( + description=( + "The hashed virtual key whose traffic will be shadowed. Shadow evaluation runs ONLY on this " + "key's traffic; requests made with any other key are not sampled." + ) + ) router_name: str = Field(description="The auto-router config to shadow requests through") shadow_percentage: float = Field( ge=0.1, @@ -213,6 +218,8 @@ class GetShadowEvalJobResponse(BaseModel): job_id: str status: ShadowEvalStatus router_name: str + api_key_id: str = Field(description="The hashed virtual key whose traffic this job evaluates, and only that key's") + team_id: str | None = Field(default=None, description="Team the shadowed key belongs to, when it has one") shadow_percentage: float request_count: int = Field(description="Total requests observed on the shadowed key since the job started") completed_count: int = Field(description="Verdicts written so far") diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 18cf9461648..4321685c2c3 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2781,11 +2781,13 @@ RoutingDecisionCause = Literal[ ] -InternalCallOrigin = Literal["autorouter_classifier"] +InternalCallOrigin = Literal["autorouter_classifier", "shadow_eval_router", "shadow_eval_judge"] """Which internal litellm feature originated a billed sub-call, so a spend log row records that it is not traffic the caller sent.""" AUTOROUTER_CLASSIFIER_CALL_ORIGIN: Final[InternalCallOrigin] = "autorouter_classifier" +SHADOW_EVAL_ROUTER_CALL_ORIGIN: Final[InternalCallOrigin] = "shadow_eval_router" +SHADOW_EVAL_JUDGE_CALL_ORIGIN: Final[InternalCallOrigin] = "shadow_eval_judge" class StandardLoggingRoutingDecision(TypedDict, total=False): diff --git a/tests/test_litellm/integrations/test_shadow_eval_logger.py b/tests/test_litellm/integrations/test_shadow_eval_logger.py index 319834a3c5a..47e5aa294cc 100644 --- a/tests/test_litellm/integrations/test_shadow_eval_logger.py +++ b/tests/test_litellm/integrations/test_shadow_eval_logger.py @@ -6,6 +6,7 @@ from unittest.mock import AsyncMock, MagicMock import pytest +from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY from litellm.integrations.shadow_eval_logger import ( JUDGE_MAX_OUTPUT_TOKENS, SHADOW_EVAL_INTERNAL_MARKER, @@ -15,6 +16,7 @@ from litellm.integrations.shadow_eval_logger import ( _sample_hits, _unmask_preference, ) +from litellm.types.utils import SHADOW_EVAL_JUDGE_CALL_ORIGIN, SHADOW_EVAL_ROUTER_CALL_ORIGIN class TestSampling: @@ -97,7 +99,7 @@ class TestJudgeOutputBudget: logger = ShadowEvalLogger(router_provider=lambda: MagicMock(), prisma_provider=lambda: MagicMock()) verdict = await logger._call_judge( - "gpt-4o-mini", [{"role": "user", "content": "hi"}], "real text", "shadow text" + "gpt-4o-mini", [{"role": "user", "content": "hi"}], "real text", "shadow text", {} ) assert verdict is not None @@ -188,6 +190,7 @@ class TestCallRouterShadowForwardsParameters: "claude-auto", [{"role": "user", "content": "hi"}], {"temperature": 0.2, "tools": [{"type": "function"}], "max_tokens": 500}, + {}, ) _, kwargs = router.acompletion.call_args @@ -204,13 +207,16 @@ class TestCallRouterShadowForwardsParameters: "claude-auto", [{"role": "user", "content": "hi"}], {"stream": True, "metadata": {"user_api_key_hash": "leaked"}, "temperature": 0.5}, + {}, ) _, kwargs = router.acompletion.call_args assert "stream" not in kwargs assert kwargs["temperature"] == 0.5 - # metadata must stay the logger's own internal marker, not the caller's - assert kwargs["metadata"] == {SHADOW_EVAL_INTERNAL_MARKER: True} + # model_parameters carries a stale copy of the parent's metadata; the shadow call + # builds its own from the parent metadata argument instead. + assert "leaked" not in str(kwargs["metadata"]) + assert kwargs["metadata"][SHADOW_EVAL_INTERNAL_MARKER] is True @pytest.mark.asyncio @@ -290,6 +296,7 @@ class TestStoppedJobCannotBeReactivated: response_obj={"choices": [{"message": {"content": "real text"}}]}, real_model="gpt-4o", model_parameters={}, + parent_metadata={}, ) prisma.db.litellm_shadowevaljob.update_many.assert_awaited_once() @@ -326,12 +333,274 @@ class TestVerdictWriteAccumulatesCost: response_obj={"choices": [{"message": {"content": "real text"}}]}, real_model="gpt-4o", model_parameters={}, + parent_metadata={}, ) _, call_kwargs = prisma.db.litellm_shadowevaljob.update_many.call_args assert call_kwargs["data"]["cost_actual"] == {"increment": 0.05} +_PARENT_METADATA = { + "user_api_key": "hashed-key", + "user_api_key_alias": "team-prod-key", + "user_api_key_team_id": "team-1", + "user_api_key_org_id": "org-1", + "user_api_key_user_id": "user-1", + "user_api_key_end_user_id": "end-user-1", + "user_api_key_budget_reservation": {"reserved_cost": 2.5, "finalized": False}, + "user_api_key_auth": {"models": ["gpt-4o"], "budget_reservation": {"reserved_cost": 2.5}}, + "routing_decision": {"tier": "SIMPLE"}, + "standard_logging_object": {"id": "parent-req"}, +} + +_IDENTITY_KEYS = ( + "user_api_key", + "user_api_key_alias", + "user_api_key_team_id", + "user_api_key_org_id", + "user_api_key_user_id", + "user_api_key_end_user_id", +) + + +def _assert_attributed(metadata, expected_origin): + """Every assertion the proxy's cost callback depends on to bill an internal sub-call.""" + for key in _IDENTITY_KEYS: + assert metadata[key] == _PARENT_METADATA[key], f"{key} must reach the sub-call for spend attribution" + assert metadata[INTERNAL_CALL_ORIGIN_METADATA_KEY] == expected_origin + # The parent completion owns this reservation. A sub-call carrying it would let the + # sub-call's cost callback finalize it, so the parent's own callback skips + # incrementing the key/team counters and the parent's spend is lost. + assert "user_api_key_budget_reservation" not in metadata + assert "budget_reservation" not in metadata["user_api_key_auth"] + assert metadata["user_api_key_auth"]["models"] == ["gpt-4o"] + # Per-request state of a request that already finished says nothing true about this one. + assert "routing_decision" not in metadata + assert "standard_logging_object" not in metadata + + +@pytest.mark.asyncio +class TestSubCallsAreAttributedToTheShadowedKey: + """Shadow and judge calls bill real provider spend nobody typed a prompt for. + + Without the caller's identity on their metadata, _PROXY_track_cost_callback drops + the spend log and no budget counter moves, so an admin-enabled eval spends against + a customer's key invisibly and past every configured limit. + """ + + async def test_judge_call_carries_the_key_identity_without_its_reservation(self, monkeypatch: pytest.MonkeyPatch): + import litellm as litellm_module + + acompletion = AsyncMock( + return_value={ + "choices": [{"message": {"content": '{"preference": "A", "confidence": 0.8, "reasoning": "clearer"}'}}] + } + ) + monkeypatch.setattr(litellm_module, "acompletion", acompletion) + logger = ShadowEvalLogger(router_provider=lambda: MagicMock(), prisma_provider=lambda: MagicMock()) + + verdict = await logger._call_judge( + "gpt-4o-mini", [{"role": "user", "content": "hi"}], "real text", "shadow text", _PARENT_METADATA + ) + + assert verdict is not None + metadata = acompletion.call_args.kwargs["metadata"] + _assert_attributed(metadata, SHADOW_EVAL_JUDGE_CALL_ORIGIN) + assert metadata[SHADOW_EVAL_INTERNAL_MARKER] is True + + async def test_shadow_router_call_carries_the_key_identity_and_the_recursion_guard(self): + router = MagicMock() + router.acompletion = AsyncMock(return_value={"choices": [{"message": {"content": "shadow reply"}}]}) + logger = ShadowEvalLogger(router_provider=lambda: router, prisma_provider=lambda: MagicMock()) + + result = await logger._call_router_shadow( + "claude-auto", [{"role": "user", "content": "hi"}], {}, _PARENT_METADATA + ) + + assert result is not None + metadata = router.acompletion.call_args.kwargs["metadata"] + _assert_attributed(metadata, SHADOW_EVAL_ROUTER_CALL_ORIGIN) + # Without the marker the shadow call's own success event would shadow itself. + assert metadata[SHADOW_EVAL_INTERNAL_MARKER] is True + + async def test_shadow_metadata_stays_writable_for_the_routing_decision_read_back(self): + """Tier attribution reads routing_decision back out of the dict the router was given.""" + router = MagicMock() + + async def acompletion(**kwargs): + kwargs["metadata"]["routing_decision"] = {"tier_label": "COMPLEX", "routed_model": "o1"} + return {"choices": [{"message": {"content": "shadow reply"}}]} + + router.acompletion = acompletion + logger = ShadowEvalLogger(router_provider=lambda: router, prisma_provider=lambda: MagicMock()) + + result = await logger._call_router_shadow( + "claude-auto", [{"role": "user", "content": "hi"}], {}, _PARENT_METADATA + ) + + assert result is not None + assert result[2] == "COMPLEX" + + async def test_parent_metadata_reaches_both_legs_from_the_success_hook(self): + """The hook is the only place the parent's metadata exists; a break here is invisible + to unit tests of the two legs in isolation.""" + prisma = MagicMock() + prisma.db.litellm_shadowevalverdict.create = AsyncMock() + prisma.db.litellm_shadowevaljob.update_many = AsyncMock() + job = ActiveShadowEvalJob(id="j1", router_name="r", shadow_percentage=100.0, judge_model="m", status="running") + logger, _, _ = _logger_with_mocks(job) + logger._prisma_provider = lambda: prisma + logger._call_router_shadow = AsyncMock(return_value=("shadow text", "shadow-model", "SIMPLE", 10)) + logger._call_judge = AsyncMock(return_value=("real", 0.9, "clearer", 0.01)) + + await logger.async_log_success_event( + { + "standard_logging_object": { + "id": "req-1", + "model": "gpt-4o", + "call_type": "acompletion", + "metadata": {"user_api_key_hash": "key-hash"}, + }, + "litellm_params": {"metadata": dict(_PARENT_METADATA)}, + "messages": [{"role": "user", "content": "hi"}], + }, + {"choices": [{"message": {"content": "real text"}}]}, + None, + None, + ) + await asyncio.sleep(0.05) + + assert logger._call_router_shadow.await_args.args[3]["user_api_key"] == "hashed-key" + assert logger._call_judge.await_args.kwargs["parent_metadata"]["user_api_key"] == "hashed-key" + + +@pytest.mark.asyncio +class TestPerJobSpendCap: + """Budgets bound the key; the cap bounds a single eval, so a bad estimate or a + traffic spike cannot turn a quoted eval into a much larger bill.""" + + @staticmethod + def _success_kwargs(): + return { + "standard_logging_object": { + "id": "req-1", + "model": "gpt-4o", + "call_type": "acompletion", + "metadata": {"user_api_key_hash": "key-hash"}, + }, + "litellm_params": {"metadata": {}}, + "messages": [{"role": "user", "content": "hi"}], + } + + async def test_job_over_the_cap_stops_sampling_and_completes_the_job(self): + job = ActiveShadowEvalJob( + id="j1", + router_name="r", + shadow_percentage=100.0, + judge_model="m", + status="running", + cost_estimate=10.0, + cost_actual=15.0, + ) + logger, prisma, router = _logger_with_mocks(job) + prisma.db.litellm_shadowevaljob.update_many = AsyncMock() + logger._run_shadow_eval = AsyncMock() + + await logger.async_log_success_event(self._success_kwargs(), MagicMock(), None, None) + await asyncio.sleep(0) + + logger._run_shadow_eval.assert_not_awaited() + router.acompletion.assert_not_called() + prisma.db.litellm_shadowevaljob.update_many.assert_awaited_once() + call_kwargs = prisma.db.litellm_shadowevaljob.update_many.call_args.kwargs + assert call_kwargs["where"]["id"] == "j1" + # Guarded so a pod breaching the cap cannot resurrect a job an admin already stopped. + assert set(call_kwargs["where"]["status"]["in"]) == {"pending", "running"} + assert call_kwargs["data"]["status"] == "completed" + assert call_kwargs["data"]["completed_at"] is not None + + async def test_job_just_under_the_cap_keeps_evaluating(self): + job = ActiveShadowEvalJob( + id="j1", + router_name="r", + shadow_percentage=100.0, + judge_model="m", + status="running", + cost_estimate=10.0, + cost_actual=14.99, + ) + logger, prisma, _ = _logger_with_mocks(job) + prisma.db.litellm_shadowevaljob.update_many = AsyncMock() + logger._run_shadow_eval = AsyncMock() + + await logger.async_log_success_event(self._success_kwargs(), MagicMock(), None, None) + await asyncio.sleep(0.01) + + logger._run_shadow_eval.assert_awaited_once() + prisma.db.litellm_shadowevaljob.update_many.assert_not_awaited() + + async def test_job_without_an_estimate_is_uncapped(self): + """A missing estimate is no multiple to compare against; treating it as 0 would + stop every such job on its first request instead.""" + job = ActiveShadowEvalJob( + id="j1", + router_name="r", + shadow_percentage=100.0, + judge_model="m", + status="running", + cost_estimate=None, + cost_actual=500.0, + ) + logger, prisma, _ = _logger_with_mocks(job) + prisma.db.litellm_shadowevaljob.update_many = AsyncMock() + logger._run_shadow_eval = AsyncMock() + + await logger.async_log_success_event(self._success_kwargs(), MagicMock(), None, None) + await asyncio.sleep(0.01) + + logger._run_shadow_eval.assert_awaited_once() + prisma.db.litellm_shadowevaljob.update_many.assert_not_awaited() + + async def test_a_cent_sized_estimate_is_not_stopped_by_its_first_verdict(self): + """1.5x a $0.01 estimate is $0.015; without the floor a single judge call ends the job.""" + job = ActiveShadowEvalJob( + id="j1", + router_name="r", + shadow_percentage=100.0, + judge_model="m", + status="running", + cost_estimate=0.01, + cost_actual=0.08, + ) + logger, prisma, _ = _logger_with_mocks(job) + prisma.db.litellm_shadowevaljob.update_many = AsyncMock() + logger._run_shadow_eval = AsyncMock() + + await logger.async_log_success_event(self._success_kwargs(), MagicMock(), None, None) + await asyncio.sleep(0.01) + + logger._run_shadow_eval.assert_awaited_once() + + async def test_stopping_evicts_the_cached_job_so_later_requests_do_not_rewrite_it(self): + job = ActiveShadowEvalJob( + id="j1", + router_name="r", + shadow_percentage=100.0, + judge_model="m", + status="running", + cost_estimate=10.0, + cost_actual=15.0, + ) + logger, prisma, _ = _logger_with_mocks(job) + prisma.db.litellm_shadowevaljob.update_many = AsyncMock() + prisma.db.litellm_shadowevaljob.find_first = AsyncMock(return_value=None) + + await logger.async_log_success_event(self._success_kwargs(), MagicMock(), None, None) + await logger.async_log_success_event(self._success_kwargs(), MagicMock(), None, None) + + assert prisma.db.litellm_shadowevaljob.update_many.await_count == 1 + + class TestExtractResponseText: def test_dict_response(self): resp = {"choices": [{"message": {"content": "hello"}}]} diff --git a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py index 357a2354cce..e0b2d551928 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py @@ -466,6 +466,49 @@ class TestAutoRouterBenchmarks: assert response.groups[0].tier_turns == expected +class TestShadowEvalJobResponseNamesWhatIsBeingEvaluated: + """A shadow eval only samples one key's traffic. An admin running several jobs cannot + tell which key a win rate belongs to unless the job response says so.""" + + @staticmethod + def _record(**overrides: object): + from datetime import datetime, timezone + + fields = { + "id": "job-1", + "status": "running", + "router_name": "claude-auto", + "api_key_id": "hashed-key-abc", + "team_id": "team-7", + "shadow_percentage": 10.0, + "request_count": 100, + "completed_count": 9, + "failed_count": 1, + "cost_estimate": 3.0, + "cost_actual": 0.42, + "created_at": datetime(2026, 8, 1, tzinfo=timezone.utc), + "completed_at": None, + **overrides, + } + return type("Row", (), fields)() + + def test_response_reports_the_shadowed_key_and_its_team(self): + from litellm.proxy.management_endpoints.auto_router_endpoints import _job_to_response + + response = _job_to_response(self._record(), None) + + assert response.api_key_id == "hashed-key-abc" + assert response.team_id == "team-7" + + def test_a_keyless_team_still_produces_a_response(self): + from litellm.proxy.management_endpoints.auto_router_endpoints import _job_to_response + + response = _job_to_response(self._record(team_id=None), None) + + assert response.api_key_id == "hashed-key-abc" + assert response.team_id is None + + class TestJudgeCostEstimate: """The upfront estimate must price the judge's real output budget, not a stale hardcoded one.""" diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 63eb71c7790..80c373b34b6 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -25329,6 +25329,11 @@ export interface components { * @description Status and, once available, results of a shadow-eval job. */ GetShadowEvalJobResponse: { + /** + * Api Key Id + * @description The hashed virtual key whose traffic this job evaluates, and only that key's + */ + api_key_id: string; /** Completed At */ completed_at?: string | null; /** @@ -25369,6 +25374,11 @@ export interface components { * @enum {string} */ status: "pending" | "running" | "completed" | "failed"; + /** + * Team Id + * @description Team the shadowed key belongs to, when it has one + */ + team_id?: string | null; }; /** * GetTeamMemberPermissionsResponse @@ -32768,7 +32778,7 @@ export interface components { StartShadowEvalRequest: { /** * Api Key Id - * @description The hashed virtual key whose traffic will be shadowed + * @description The hashed virtual key whose traffic will be shadowed. Shadow evaluation runs ONLY on this key's traffic; requests made with any other key are not sampled. */ api_key_id: string; /**