diff --git a/litellm/proxy/management_endpoints/key_budget_resolver.py b/litellm/proxy/management_endpoints/key_budget_resolver.py index ea33c9a1676..ef78617aec5 100644 --- a/litellm/proxy/management_endpoints/key_budget_resolver.py +++ b/litellm/proxy/management_endpoints/key_budget_resolver.py @@ -155,6 +155,14 @@ _CUSTOM_AUTH_SKIPS_CHECKS_NOTE: Final = KeyBudgetNote( "at request time for the requests it authenticates; only the reservation layer still applies" ), ) +_ENTITY_UNAVAILABLE_NOTE: Final = KeyBudgetNote( + code="entity_unavailable", + severity="warning", + text=( + "this entity could not be read, so its budgets are unknown rather than unset; nothing on this " + "scope can be ruled out until it can be read again" + ), +) _USER_ON_TEAM_KEY_NOTE: Final = KeyBudgetNote( code="user_budget_not_applied_to_team_key", severity="info", @@ -279,6 +287,19 @@ class KeyBudgetResolverDeps: match_model_budget_key: ModelBudgetKeyMatcher = field(default=_match_model_budget_key) +@dataclass(frozen=True, slots=True) +class _Unavailable: + """An entity that could not be read, which is not the same as an entity that is not there.""" + + +_UNAVAILABLE: Final = _Unavailable() + + +@dataclass(frozen=True, slots=True) +class _UnknownSpend: + """Stands in for the counter of an entity that could not be read, so no number is invented for it.""" + + @dataclass(frozen=True, slots=True) class _CounterSpend: counter_key: str @@ -308,7 +329,7 @@ class _EndUserModelSpend: budget_config: BudgetConfig -_SpendSource = _CounterSpend | _RecordedSpend | _KeyModelSpend | _EndUserModelSpend +_SpendSource = _CounterSpend | _RecordedSpend | _KeyModelSpend | _EndUserModelSpend | _UnknownSpend @dataclass(frozen=True, slots=True) @@ -480,16 +501,17 @@ class _KeyBudgetContext: request_models: tuple[str, ...] match_model_budget_key: ModelBudgetKeyMatcher general_settings: Mapping[str, object] - proxy: _ProxyBudget | None + proxy: _ProxyBudget | _Unavailable | None team: LiteLLM_TeamTable | None user: LiteLLM_UserTable | None project: LiteLLM_ProjectTableCachedObj | None organization: LiteLLM_OrganizationTable | None team_member: TeamMemberBudget | None - tags: tuple[tuple[str, LiteLLM_TagTable | None], ...] + tags: tuple[tuple[str, LiteLLM_TagTable | _Unavailable | None], ...] end_user: LiteLLM_EndUserTable | None default_end_user_budget: LiteLLM_BudgetTable | None budget_meta: Mapping[str, _BudgetRowMeta] + unresolved: Mapping[BudgetScope, str | None] async def resolve_key_budgets( @@ -553,6 +575,8 @@ async def _read_spend(plan: _PlannedBudget, deps: KeyBudgetResolverDeps) -> _Spe source: Final = plan.spend_source try: match source: + case _UnknownSpend(): + return _SpendReading(value=None, state="unavailable") case _RecordedSpend(): return _SpendReading(value=source.value, state="live") case _CounterSpend(): @@ -611,6 +635,22 @@ def _entry_notes( return (*tightened, *plan.notes, *proxy_notes) +def _status(plan: _PlannedBudget, spend: float | None, exceeded: bool) -> BudgetStatus: + """ + A row nobody could evaluate reports ``unknown``, never ``unlimited``. + + Reading it as unset is the failure this endpoint exists to prevent: an unreadable entity and an + unreadable counter both leave a budget that may well be the one blocking the key. + """ + if isinstance(plan.spend_source, _UnknownSpend): + return "unknown" + if plan.max_budget is None: + return "unlimited" + if spend is None: + return "unknown" + return "exceeded" if exceeded else "ok" + + def _to_entry( plan: _PlannedBudget, reading: _SpendReading, @@ -624,7 +664,7 @@ def _to_entry( and spend is not None and (spend >= plan.max_budget if comparison == ">=" else spend > plan.max_budget) ) - status: Final[BudgetStatus] = "unlimited" if plan.max_budget is None else ("exceeded" if exceeded else "ok") + status: Final[BudgetStatus] = _status(plan=plan, spend=spend, exceeded=exceeded) return KeyBudgetEntry( scope=plan.scope, entity_type=_ENTITY_TYPE_BY_SCOPE[plan.scope], @@ -651,7 +691,7 @@ async def _load_context( deps: KeyBudgetResolverDeps, ) -> _KeyBudgetContext: token_inputs: Final = _token_budget_inputs(valid_token) - proxy, team, user, project, tags, end_user = await asyncio.gather( + loaded_proxy, loaded_team, loaded_user, loaded_project, tags, loaded_end_user = await asyncio.gather( _load_proxy_budget(deps), _load_team(valid_token.team_id, deps), _load_user(valid_token.user_id, deps), @@ -659,11 +699,17 @@ async def _load_context( _load_tags(token_inputs.tags, deps), _load_end_user(end_user_id, deps), ) - organization, team_member, default_end_user_budget = await asyncio.gather( - _load_organization(resolve_budget_org_id(valid_token=valid_token, team_object=team), deps), + team: Final = _resolved(loaded_team) + org_id: Final = resolve_budget_org_id(valid_token=valid_token, team_object=team) + loaded_organization, loaded_team_member, loaded_default_end_user_budget = await asyncio.gather( + _load_organization(org_id, deps), _load_team_member(valid_token=valid_token, team=team, deps=deps), _load_default_end_user_budget(deps), ) + project: Final = _resolved(loaded_project) + organization: Final = _resolved(loaded_organization) + end_user: Final = _resolved(loaded_end_user) + default_end_user_budget: Final = _resolved(loaded_default_end_user_budget) budget_ids: Final = _referenced_budget_ids( organization=organization, project=project, @@ -683,23 +729,90 @@ async def _load_context( request_models=_request_models(token_inputs.models), match_model_budget_key=deps.match_model_budget_key, general_settings=deps.general_settings, - proxy=proxy, + proxy=loaded_proxy, team=team, - user=user, + user=_resolved(loaded_user), project=project, organization=organization, - team_member=team_member, + team_member=_resolved(loaded_team_member), tags=tags, end_user=end_user, default_end_user_budget=default_end_user_budget, budget_meta=await _load_budget_meta(budget_ids, deps), + unresolved=_unresolved_scopes( + valid_token=valid_token, + end_user_id=end_user_id, + org_id=org_id, + loaded_team=loaded_team, + loaded_user=loaded_user, + loaded_project=loaded_project, + loaded_organization=loaded_organization, + loaded_team_member=loaded_team_member, + loaded_end_user=loaded_end_user, + loaded_default_end_user_budget=loaded_default_end_user_budget, + ), ) +def _resolved(value: _T | _Unavailable | None) -> _T | None: + return None if isinstance(value, _Unavailable) else value + + +def _unresolved_scopes( + valid_token: UserAPIKeyAuth, + end_user_id: str | None, + org_id: str | None, + loaded_team: LiteLLM_TeamTable | _Unavailable | None, + loaded_user: LiteLLM_UserTable | _Unavailable | None, + loaded_project: LiteLLM_ProjectTableCachedObj | _Unavailable | None, + loaded_organization: LiteLLM_OrganizationTable | _Unavailable | None, + loaded_team_member: TeamMemberBudget | _Unavailable | None, + loaded_end_user: LiteLLM_EndUserTable | _Unavailable | None, + loaded_default_end_user_budget: LiteLLM_BudgetTable | _Unavailable | None, +) -> Mapping[BudgetScope, str | None]: + """ + The scopes whose budgets could not be read, mapped to the entity they belong to. + + A failed lookup reaches the planners as an absent entity, which would drop the scope from the + report or leave it looking unconfigured, so each one is recorded here and reported on its own row. + A team carries three scopes and is also where an organization is inherited from, so losing it + leaves all four unknown. + """ + team_unavailable: Final = isinstance(loaded_team, _Unavailable) + end_user_object: Final = _resolved(loaded_end_user) + candidates: Final[tuple[tuple[BudgetScope, bool, str | None], ...]] = ( + ("team", team_unavailable, valid_token.team_id), + ("team_window", team_unavailable, valid_token.team_id), + ( + "team_member", + isinstance(loaded_team_member, _Unavailable) or (team_unavailable and valid_token.user_id is not None), + f"{valid_token.user_id}:{valid_token.team_id}", + ), + ("user", isinstance(loaded_user, _Unavailable), valid_token.user_id), + ( + "organization", + isinstance(loaded_organization, _Unavailable) or (team_unavailable and valid_token.org_id is None), + org_id, + ), + ("project", isinstance(loaded_project, _Unavailable), valid_token.project_id), + ( + "end_user", + isinstance(loaded_end_user, _Unavailable) + or ( + end_user_id is not None + and isinstance(loaded_default_end_user_budget, _Unavailable) + and (end_user_object is None or end_user_object.budget_id is None) + ), + end_user_id, + ), + ) + return MappingProxyType({scope: entity_id for scope, unavailable, entity_id in candidates if unavailable}) + + def _referenced_budget_ids( organization: LiteLLM_OrganizationTable | None, project: LiteLLM_ProjectTableCachedObj | None, - tags: Sequence[tuple[str, LiteLLM_TagTable | None]], + tags: Sequence[tuple[str, LiteLLM_TagTable | _Unavailable | None]], end_user: LiteLLM_EndUserTable | None, default_end_user_budget: LiteLLM_BudgetTable | None, key_budget_id: str | None, @@ -710,7 +823,7 @@ def _referenced_budget_ids( project.budget_id if project is not None else None, end_user.budget_id if end_user is not None else None, default_end_user_budget.budget_id if default_end_user_budget is not None else None, - *(tag.budget_id for _, tag in tags if tag is not None), + *(tag.budget_id for _, tag in tags if isinstance(tag, LiteLLM_TagTable)), ) return frozenset(budget_id for budget_id in candidates if budget_id is not None) @@ -740,12 +853,12 @@ async def _load_budget_meta( ) -async def _load_proxy_budget(deps: KeyBudgetResolverDeps) -> _ProxyBudget | None: +async def _load_proxy_budget(deps: KeyBudgetResolverDeps) -> _ProxyBudget | _Unavailable | None: try: row: Final = await UserRepository(deps.prisma_client).find_by_id(LITELLM_PROXY_BUDGET_NAME) except Exception: # noqa: BLE001 # every entity load degrades to "unknown" rather than failing the report verbose_proxy_logger.exception("Unable to load the proxy budget row") - return None + return _UNAVAILABLE if row is None: return None return _ProxyBudget( @@ -755,7 +868,7 @@ async def _load_proxy_budget(deps: KeyBudgetResolverDeps) -> _ProxyBudget | None ) -async def _load_team(team_id: str | None, deps: KeyBudgetResolverDeps) -> LiteLLM_TeamTable | None: +async def _load_team(team_id: str | None, deps: KeyBudgetResolverDeps) -> LiteLLM_TeamTable | _Unavailable | None: if team_id is None: return None try: @@ -767,10 +880,10 @@ async def _load_team(team_id: str | None, deps: KeyBudgetResolverDeps) -> LiteLL ) except Exception: # noqa: BLE001 # see _load_proxy_budget verbose_proxy_logger.exception("Unable to load team %s for budget resolution", team_id) - return None + return _UNAVAILABLE -async def _load_user(user_id: str | None, deps: KeyBudgetResolverDeps) -> LiteLLM_UserTable | None: +async def _load_user(user_id: str | None, deps: KeyBudgetResolverDeps) -> LiteLLM_UserTable | _Unavailable | None: if user_id is None: return None try: @@ -783,10 +896,12 @@ async def _load_user(user_id: str | None, deps: KeyBudgetResolverDeps) -> LiteLL ) except Exception: # noqa: BLE001 # see _load_proxy_budget verbose_proxy_logger.exception("Unable to load user %s for budget resolution", user_id) - return None + return _UNAVAILABLE -async def _load_project(project_id: str | None, deps: KeyBudgetResolverDeps) -> LiteLLM_ProjectTableCachedObj | None: +async def _load_project( + project_id: str | None, deps: KeyBudgetResolverDeps +) -> LiteLLM_ProjectTableCachedObj | _Unavailable | None: if project_id is None: return None try: @@ -798,10 +913,12 @@ async def _load_project(project_id: str | None, deps: KeyBudgetResolverDeps) -> ) except Exception: # noqa: BLE001 # see _load_proxy_budget verbose_proxy_logger.exception("Unable to load project %s for budget resolution", project_id) - return None + return _UNAVAILABLE -async def _load_organization(org_id: str | None, deps: KeyBudgetResolverDeps) -> LiteLLM_OrganizationTable | None: +async def _load_organization( + org_id: str | None, deps: KeyBudgetResolverDeps +) -> LiteLLM_OrganizationTable | _Unavailable | None: if org_id is None: return None try: @@ -814,13 +931,13 @@ async def _load_organization(org_id: str | None, deps: KeyBudgetResolverDeps) -> ) except Exception: # noqa: BLE001 # see _load_proxy_budget verbose_proxy_logger.exception("Unable to load organization %s for budget resolution", org_id) - return None + return _UNAVAILABLE async def _load_tags( tag_names: Sequence[str], deps: KeyBudgetResolverDeps, -) -> tuple[tuple[str, LiteLLM_TagTable | None], ...]: +) -> tuple[tuple[str, LiteLLM_TagTable | _Unavailable | None], ...]: if not tag_names: return () try: @@ -832,11 +949,13 @@ async def _load_tags( ) except Exception: # noqa: BLE001 # see _load_proxy_budget verbose_proxy_logger.exception("Unable to load tags %s for budget resolution", tag_names) - return tuple((tag_name, None) for tag_name in tag_names) + return tuple((tag_name, _UNAVAILABLE) for tag_name in tag_names) return tuple((tag_name, tag_objects.get(tag_name)) for tag_name in tag_names) -async def _load_end_user(end_user_id: str | None, deps: KeyBudgetResolverDeps) -> LiteLLM_EndUserTable | None: +async def _load_end_user( + end_user_id: str | None, deps: KeyBudgetResolverDeps +) -> LiteLLM_EndUserTable | _Unavailable | None: if end_user_id is None: return None try: @@ -848,10 +967,10 @@ async def _load_end_user(end_user_id: str | None, deps: KeyBudgetResolverDeps) - ) except Exception: # noqa: BLE001 # see _load_proxy_budget verbose_proxy_logger.exception("Unable to load end user %s for budget resolution", end_user_id) - return None + return _UNAVAILABLE -async def _load_default_end_user_budget(deps: KeyBudgetResolverDeps) -> LiteLLM_BudgetTable | None: +async def _load_default_end_user_budget(deps: KeyBudgetResolverDeps) -> LiteLLM_BudgetTable | _Unavailable | None: try: return await get_default_end_user_budget( prisma_client=deps.prisma_client, @@ -859,14 +978,14 @@ async def _load_default_end_user_budget(deps: KeyBudgetResolverDeps) -> LiteLLM_ ) except Exception: # noqa: BLE001 # see _load_proxy_budget verbose_proxy_logger.exception("Unable to load the default end user budget") - return None + return _UNAVAILABLE async def _load_team_member( valid_token: UserAPIKeyAuth, team: LiteLLM_TeamTable | None, deps: KeyBudgetResolverDeps, -) -> TeamMemberBudget | None: +) -> TeamMemberBudget | _Unavailable | None: if team is None or valid_token.user_id is None: return None try: @@ -879,7 +998,7 @@ async def _load_team_member( ) except Exception: # noqa: BLE001 # see _load_proxy_budget verbose_proxy_logger.exception("Unable to resolve the team-member budget for team %s", team.team_id) - return None + return _UNAVAILABLE def _plan_budgets(context: _KeyBudgetContext) -> tuple[_PlannedBudget, ...]: @@ -896,9 +1015,42 @@ def _plan_budgets(context: _KeyBudgetContext) -> tuple[_PlannedBudget, ...]: *_plan_project(context), *_plan_tags(context), *_plan_end_user(context), + *_unresolved_plans(context), ) +_UNRESOLVED_COMPARISON: Final[Mapping[BudgetScope, BudgetComparison]] = MappingProxyType( + { + "team": ">", + "team_window": ">=", + "team_member": ">=", + "user": ">=", + "organization": ">=", + "project": ">", + "tag": ">", + "end_user": ">", + } +) + + +def _unresolved_plan(scope: BudgetScope, entity_id: str | None) -> _PlannedBudget: + return _PlannedBudget( + scope=scope, + entity_id=entity_id, + entity_label=None, + enforcement="hard", + max_budget=None, + comparison=_UNRESOLVED_COMPARISON[scope], + source="unavailable", + spend_source=_UnknownSpend(), + notes=(_ENTITY_UNAVAILABLE_NOTE,), + ) + + +def _unresolved_plans(context: _KeyBudgetContext) -> tuple[_PlannedBudget, ...]: + return tuple(_unresolved_plan(scope, entity_id) for scope, entity_id in context.unresolved.items()) + + def _budget_meta(context: _KeyBudgetContext, budget_id: str | None) -> _BudgetRowMeta: if budget_id is None: return _BudgetRowMeta(budget_duration=None, budget_reset_at=None) @@ -907,20 +1059,37 @@ def _budget_meta(context: _KeyBudgetContext, budget_id: str | None) -> _BudgetRo def _plan_proxy(context: _KeyBudgetContext) -> tuple[_PlannedBudget, ...]: proxy: Final = context.proxy - return ( - _PlannedBudget( - scope="proxy", - entity_id=None, - entity_label=None, - enforcement="hard", - max_budget=litellm.max_budget if litellm.max_budget > 0 else None, - comparison=">", - source="litellm_settings.max_budget", - spend_source=_RecordedSpend(proxy.spend if proxy is not None else 0.0), - budget_duration=proxy.budget_duration if proxy is not None else None, - budget_reset_at=proxy.budget_reset_at if proxy is not None else None, - ), - ) + max_budget: Final = litellm.max_budget if litellm.max_budget > 0 else None + match proxy: + case _Unavailable(): + return ( + _PlannedBudget( + scope="proxy", + entity_id=None, + entity_label=None, + enforcement="hard", + max_budget=max_budget, + comparison=">", + source="litellm_settings.max_budget", + spend_source=_UnknownSpend(), + notes=(_ENTITY_UNAVAILABLE_NOTE,), + ), + ) + case _: + return ( + _PlannedBudget( + scope="proxy", + entity_id=None, + entity_label=None, + enforcement="hard", + max_budget=max_budget, + comparison=">", + source="litellm_settings.max_budget", + spend_source=_RecordedSpend(proxy.spend if proxy is not None else 0.0), + budget_duration=proxy.budget_duration if proxy is not None else None, + budget_reset_at=proxy.budget_reset_at if proxy is not None else None, + ), + ) def _key_max_budget_source(context: _KeyBudgetContext) -> str: @@ -1207,35 +1376,36 @@ def _plan_project(context: _KeyBudgetContext) -> tuple[_PlannedBudget, ...]: def _plan_tags(context: _KeyBudgetContext) -> tuple[_PlannedBudget, ...]: - return tuple( - _PlannedBudget( - scope="tag", - entity_id=tag_name, - entity_label=None, - enforcement="hard", - max_budget=( - tag.litellm_budget_table.max_budget - if tag is not None and tag.litellm_budget_table is not None - else None - ), - comparison=">", - source=f"budget_table:{tag.budget_id}" if tag is not None and tag.budget_id else "tag.budget_id", - spend_source=_CounterSpend( - counter_key=tag_spend_counter(tag_name), - fallback_spend=(tag.spend or 0.0) if tag is not None else 0.0, - fallback_authoritative=True, - ), - budget_duration=_budget_meta(context, tag.budget_id if tag is not None else None).budget_duration, - budget_reset_at=_budget_meta(context, tag.budget_id if tag is not None else None).budget_reset_at, - notes=(_TAG_NOTE,), - ) - for tag_name, tag in context.tags + return tuple(_plan_tag(context, tag_name, tag) for tag_name, tag in context.tags) + + +def _plan_tag(context: _KeyBudgetContext, tag_name: str, tag: LiteLLM_TagTable | _Unavailable | None) -> _PlannedBudget: + if isinstance(tag, _Unavailable): + return _unresolved_plan("tag", tag_name) + return _PlannedBudget( + scope="tag", + entity_id=tag_name, + entity_label=None, + enforcement="hard", + max_budget=( + tag.litellm_budget_table.max_budget if tag is not None and tag.litellm_budget_table is not None else None + ), + comparison=">", + source=f"budget_table:{tag.budget_id}" if tag is not None and tag.budget_id else "tag.budget_id", + spend_source=_CounterSpend( + counter_key=tag_spend_counter(tag_name), + fallback_spend=(tag.spend or 0.0) if tag is not None else 0.0, + fallback_authoritative=True, + ), + budget_duration=_budget_meta(context, tag.budget_id if tag is not None else None).budget_duration, + budget_reset_at=_budget_meta(context, tag.budget_id if tag is not None else None).budget_reset_at, + notes=(_TAG_NOTE,), ) def _plan_end_user(context: _KeyBudgetContext) -> tuple[_PlannedBudget, ...]: end_user_id: Final = context.end_user_id - if end_user_id is None: + if end_user_id is None or "end_user" in context.unresolved: return () end_user: Final = context.end_user budget: Final = end_user.litellm_budget_table if end_user is not None else context.default_end_user_budget diff --git a/litellm/types/proxy/management_endpoints/key_management_endpoints.py b/litellm/types/proxy/management_endpoints/key_management_endpoints.py index 2a2b1bc5182..ca3d87f0fac 100644 --- a/litellm/types/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/types/proxy/management_endpoints/key_management_endpoints.py @@ -130,13 +130,14 @@ BudgetEnforcement = Literal["hard", "soft", "throttled"] BudgetComparison = Literal[">=", ">"] -BudgetStatus = Literal["unlimited", "ok", "exceeded"] +BudgetStatus = Literal["unlimited", "ok", "exceeded", "unknown"] BudgetNoteCode = Literal[ "alert_only", "custom_auth_may_override_end_user_cap", "custom_auth_skips_read_time_checks", "end_user_route_only", + "entity_unavailable", "per_model_counters", "project_spend_not_tracked", "request_tags_add_budgets", @@ -171,7 +172,13 @@ class KeyBudgetNote(BaseModel): class KeyBudgetEntry(BaseModel): - """One budget that can gate requests made with a key, with its live spend.""" + """ + One budget that can gate requests made with a key, with its live spend. + + ``status`` is ``unknown`` when the row could not be evaluated, either because the entity behind it + was unreadable or because its spend was, and it is never ``unlimited`` in that case: an unreadable + scope is not a scope the reader may rule out. + """ scope: BudgetScope entity_type: Litellm_EntityType diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index d3d0ee7f752..5dbd531d65b 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -17393,12 +17393,110 @@ async def test_key_budgets_call_an_unreadable_counter_unreadable_rather_than_rep assert key_entry.spend is None assert key_entry.spend_state == "unavailable" assert key_entry.remaining is None + assert key_entry.status == "unknown", "a cap whose spend nobody could read is not a cap known to be under" assert "unavailable" not in [note.code for note in key_entry.notes], "a failed read is not a caveat" project_entry = next(e for e in budgets if e.scope == "project" and e.enforcement == "hard") assert project_entry.spend_state == "live", "the recorded-spend scopes do not go through the counter" +_BUDGETS_LOOKUP_BY_SCOPE = ( + ("team", "get_team_object"), + ("user", "get_user_object"), + ("project", "get_project_object"), + ("organization", "get_org_object"), + ("end_user", "get_end_user_object"), + ("team_member", "resolve_team_member_budget"), + ("tag", "get_tag_objects_batch"), +) + + +async def _budgets_with_failed_lookup(lookup): + async def _explode(*args, **kwargs): + raise RuntimeError(f"{lookup} is down") + + with _budgets_world(**_fully_populated_world()): + with patch(f"{_BUDGETS_RESOLVER}.{lookup}", _explode): + return await resolve_key_budgets( + valid_token=_budgets_token(), + end_user_id="end-user-budgets", + deps=_budgets_deps(), + ) + + +@pytest.mark.parametrize("scope, lookup", _BUDGETS_LOOKUP_BY_SCOPE) +@pytest.mark.asyncio +async def test_key_budgets_report_a_scope_it_could_not_read_instead_of_dropping_it(scope, lookup): + """ + A lookup that fails reaches the planners as an absent entity, which is the one lie this endpoint + cannot tell: the scope would either vanish from the report or sit there reading unlimited, and + either way the reader rules out the budget that may be the one blocking the key. + """ + budgets = await _budgets_with_failed_lookup(lookup) + + entry = next(e for e in budgets if e.scope == scope) + assert entry.status == "unknown", scope + assert entry.max_budget is None, scope + assert entry.spend is None, scope + assert entry.spend_state == "unavailable", scope + assert "entity_unavailable" in [note.code for note in entry.notes], scope + + +@pytest.mark.parametrize("scope, lookup", _BUDGETS_LOOKUP_BY_SCOPE) +@pytest.mark.asyncio +async def test_key_budgets_name_the_entity_of_a_scope_it_could_not_read(scope, lookup): + """An unknown row with no entity on it cannot be chased down, which is the whole point of reporting it.""" + budgets = await _budgets_with_failed_lookup(lookup) + + entity_id = next(e for e in budgets if e.scope == scope).entity_id + assert entity_id == { + "team": "team-budgets", + "user": "user-budgets", + "project": "project-budgets", + "organization": "org-budgets", + "end_user": "end-user-budgets", + "team_member": "user-budgets:team-budgets", + "tag": "prod", + }[scope], scope + + +@pytest.mark.asyncio +async def test_key_budgets_keep_every_scope_a_team_carries_when_the_team_cannot_be_read(): + """A team is three scopes plus the organization it is inherited from, so one failed read hides four budgets.""" + budgets = await _budgets_with_failed_lookup("get_team_object") + + unknown = {entry.scope for entry in budgets if entry.status == "unknown"} + assert {"team", "team_window", "team_member", "organization"} <= unknown + + +@pytest.mark.asyncio +async def test_key_budgets_report_the_proxy_budget_as_unknown_when_its_spend_row_cannot_be_read(): + """The global cap comes from config, so the row still names it, but nobody may read it as unspent.""" + repository = MagicMock() + repository.return_value.find_by_id = AsyncMock(side_effect=RuntimeError("postgres is down")) + + with _budgets_world(**_fully_populated_world()): + with patch(f"{_BUDGETS_RESOLVER}.UserRepository", repository), patch("litellm.max_budget", 1000.0): + budgets = await resolve_key_budgets( + valid_token=_budgets_token(), end_user_id=None, deps=_budgets_deps() + ) + + entry = next(e for e in budgets if e.scope == "proxy") + assert entry.max_budget == 1000.0 + assert entry.spend is None + assert entry.status == "unknown" + assert "entity_unavailable" in [note.code for note in entry.notes] + + +@pytest.mark.asyncio +async def test_key_budgets_do_not_duplicate_a_scope_they_could_not_read(): + """The unknown row replaces the scope's rows; emitting both would put two verdicts on one budget.""" + budgets = await _budgets_with_failed_lookup("get_end_user_object") + + assert len([entry for entry in budgets if entry.scope == "end_user"]) == 1 + assert not [entry for entry in budgets if entry.scope == "end_user_model"] + + @pytest.mark.asyncio async def test_key_budgets_say_a_throttling_key_throttles_rather_than_calling_it_a_blocker(): """`hard` plus a note meant every client that skipped the prose reported a denial that never happened.""" @@ -17523,6 +17621,8 @@ def test_key_budgets_classify_every_note_code_and_leave_none_to_a_default(): "per_model_counters": "warning", "project_spend_not_tracked": "warning", "request_tags_add_budgets": "warning", + # `status` carries it too, but only as a value no client built before this endpoint can know + "entity_unavailable": "warning", } diff --git a/ui/litellm-dashboard/src/components/templates/KeyBudgetsTableColumns.test.ts b/ui/litellm-dashboard/src/components/templates/KeyBudgetsTableColumns.test.ts index bad06814347..d5e6cb77851 100644 --- a/ui/litellm-dashboard/src/components/templates/KeyBudgetsTableColumns.test.ts +++ b/ui/litellm-dashboard/src/components/templates/KeyBudgetsTableColumns.test.ts @@ -216,7 +216,7 @@ describe("custom auth skipping read-time checks", () => { { ...UNCONFIGURED_BUDGET, notes: [CUSTOM_AUTH_SKIPS_NOTE] }, ]; expect(rows.map(cannotTrip)).toStrictEqual([false, false, false]); - expect(rows.map(rowRank)).toStrictEqual([0, 2, 3]); + expect(rows.map(rowRank)).toStrictEqual([0, 3, 4]); expect(isBlockingRow(rows[0])).toBe(true); }); @@ -271,6 +271,43 @@ describe("isBlockingRow", () => { }); }); +const ENTITY_UNAVAILABLE_NOTE = { + code: "entity_unavailable", + severity: "warning", + text: noteText("entity_unavailable"), +} as const; + +// The scope applies to the key, but the lookup behind it failed: no cap and no spend, which is the +// same shape as a scope with nothing configured and must never be ranked or read as one. +const UNRESOLVED: KeyBudgetEntry = { + ...UNCONFIGURED_BUDGET, + scope: "team", + entity_id: "team-budgets", + spend: null, + spend_state: "unavailable", + source: "unavailable", + status: "unknown", + notes: [ENTITY_UNAVAILABLE_NOTE], +}; + +describe("a scope the server could not resolve", () => { + it("is not dead, since the budget it hides can be the one denying every request", () => { + expect(cannotTrip(UNRESOLVED)).toBe(false); + }); + + it("is not blamed for the denial either, because no number on it supports that", () => { + expect(isBlockingRow(UNRESOLVED)).toBe(false); + expect(budgetThresholdRule(UNRESOLVED)).toBeNull(); + }); + + it("outranks every row known to be under its limit, and never sinks to an unlimited one", () => { + expect(rowRank(UNRESOLVED)).toBeGreaterThan(rowRank(BLOCKING)); + expect(rowRank(UNRESOLVED)).toBeLessThan(rowRank(HEALTHY)); + expect(rowRank(UNRESOLVED)).toBeLessThan(rowRank(UNCONFIGURED_BUDGET)); + expect(rowRank(UNRESOLVED)).toBeLessThan(rowRank(INERT)); + }); +}); + describe("rowRank", () => { it("ranks a blocking budget above an alert-only one, and both above healthy and unlimited", () => { expect(rowRank(BLOCKING)).toBeLessThan(rowRank(ALERT_ONLY)); diff --git a/ui/litellm-dashboard/src/components/templates/KeyBudgetsTableColumns.tsx b/ui/litellm-dashboard/src/components/templates/KeyBudgetsTableColumns.tsx index 67debc4594d..ef7aba56a20 100644 --- a/ui/litellm-dashboard/src/components/templates/KeyBudgetsTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/templates/KeyBudgetsTableColumns.tsx @@ -54,6 +54,7 @@ const CODE_KILLS_ROW: Readonly> = { custom_auth_may_override_end_user_cap: false, custom_auth_skips_read_time_checks: false, end_user_route_only: false, + entity_unavailable: false, per_model_counters: false, project_spend_not_tracked: true, request_tags_add_budgets: false, @@ -78,13 +79,19 @@ const canDeny = (entry: KeyBudgetEntry): boolean => !isAlertOnly(entry) && !cann export const isBlockingRow = (entry: KeyBudgetEntry): boolean => entry.status === "exceeded" && canDeny(entry); -/** Ascending relevance to "what stopped my request", so a row that cannot answer it sorts last. */ +/** + * Ascending relevance to "what stopped my request", so a row that cannot answer it sorts last. + * + * A budget nobody could read outranks one known to be under its limit: it is still a candidate for + * the denial being chased, and it is the only row on the table that needs someone to go look. + */ export const rowRank = (entry: KeyBudgetEntry): number => { - if (entry.status === "unlimited") return 3; - if (cannotTrip(entry)) return 4; + if (entry.status === "unlimited") return 4; + if (cannotTrip(entry)) return 5; if (isBlockingRow(entry)) return 0; if (entry.status === "exceeded") return 1; - return 2; + if (entry.status === "unknown") return 2; + return 3; }; /** @@ -116,6 +123,8 @@ export const budgetThresholdRule = (entry: KeyBudgetEntry): string | null => { const statusPresentation = (entry: KeyBudgetEntry): { tone: StatusTone; label: string } => { if (entry.status === "unlimited") return { tone: "neutral", label: "Unlimited" }; + // Before the exceeded branch: a row nobody could read has no numbers to be within or over. + if (entry.status === "unknown") return { tone: "warning", label: "Unknown" }; if (cannotTrip(entry)) return { tone: "neutral", label: "Cannot trip" }; if (entry.status !== "exceeded") return { tone: "success", label: "Within budget" }; if (isAlertOnly(entry)) return { tone: "warning", label: "Exceeded (alert only)" }; @@ -185,6 +194,12 @@ function EnforcementCell({ entry }: { entry: KeyBudgetEntry }) { return ; } +/** What sits beside an unreadable spend: the cap when there is one, and never "Unlimited" without one. */ +const unreadableSpendLimit = (entry: KeyBudgetEntry): string => { + if (entry.max_budget != null) return `of $${formatNumberWithCommas(entry.max_budget, 2)}`; + return entry.status === "unknown" ? "· of an unknown limit" : "· Unlimited"; +}; + function SpendCell({ entry }: { entry: KeyBudgetEntry }) { const rule = budgetThresholdRule(entry); return ( @@ -194,9 +209,7 @@ function SpendCell({ entry }: { entry: KeyBudgetEntry }) { ) : ( Unknown{" "} - - {entry.max_budget == null ? "· Unlimited" : `of $${formatNumberWithCommas(entry.max_budget, 2)}`} - + {unreadableSpendLimit(entry)} )} {rule && {rule}} @@ -205,7 +218,9 @@ function SpendCell({ entry }: { entry: KeyBudgetEntry }) { } function RemainingCell({ entry }: { entry: KeyBudgetEntry }) { - const unlimited = entry.max_budget == null; + // A missing cap reads as unlimited only when the row knows there is none; on an unresolved row it + // means nobody could read one, and printing "Unlimited" there is the answer this table must not give. + const unlimited = entry.max_budget == null && entry.status !== "unknown"; return ( { // compared against. Only a failed read has no number at all. expect(budget.spend_state === "unavailable").toBe(budget.spend === null); } - expect(budget.status === "unlimited").toBe(budget.max_budget === null); + // A row nobody could resolve has no cap to call unlimited and no spend to compare, so it is the + // one case where a null cap is not an unlimited budget. + if (budget.status === "unknown") { + expect(budget.spend).toBeNull(); + } else { + expect(budget.status === "unlimited").toBe(budget.max_budget === null); + } if (budget.spend == null || budget.max_budget == null) { expect(budget.remaining).toBeNull(); } else { @@ -349,13 +355,15 @@ describe("KeyInfoView Budgets tab", () => { spend_state: "unavailable", remaining: null, source: "budget_table:b-tag", - status: "ok", + status: "unknown", }; mockBudgets([unreadable]); const panel = await renderAndOpenBudgetsTab(); const row = rowFor(panel, "prod"); expect(row).toHaveTextContent("Unknown of $1,000.00"); + expect(cellUnder(panel, row, "Status")).toHaveTextContent("Unknown"); + expect(within(row).queryByText("Within budget")).not.toBeInTheDocument(); expect(row).not.toHaveTextContent("$0.00"); // The meter is the part that lies loudest: drawn at 0% it reads as untouched headroom. expect(within(row).queryByRole("meter")).not.toBeInTheDocument(); @@ -363,6 +371,44 @@ describe("KeyInfoView Budgets tab", () => { expect(row).toHaveTextContent("Blocks at ≥ $1,000.00"); }); + // A lookup that fails leaves the scope unresolved: no cap, no spend, and a warning note. The trap + // is that it looks exactly like a scope with nothing configured, which is the one reading that lets + // someone rule out the budget that is blocking them. + const UNRESOLVED_TEAM: KeyBudgetEntry = { + ...UNCONFIGURED_BUDGET, + scope: "team", + entity_type: "team", + entity_id: "team-123", + entity_label: null, + spend: null, + spend_state: "unavailable", + comparison: ">", + source: "unavailable", + status: "unknown", + notes: [{ code: "entity_unavailable", severity: "warning", text: noteText("entity_unavailable") }], + }; + + it("never reads a scope it could not resolve as an unlimited budget", async () => { + mockBudgets([UNRESOLVED_TEAM, KEY_UNLIMITED]); + const panel = await renderAndOpenBudgetsTab(); + + const row = rowFor(panel, "team-123"); + expect(cellUnder(panel, row, "Status")).toHaveTextContent("Unknown"); + expect(within(row).queryByText("Unlimited")).not.toBeInTheDocument(); + expect(within(row).queryByText("Within budget")).not.toBeInTheDocument(); + expect(within(row).queryByText("Cannot trip")).not.toBeInTheDocument(); + expect(cellUnder(panel, row, "Remaining")).toHaveTextContent("-"); + expect(row).toHaveTextContent(noteText("entity_unavailable")); + }); + + it("sorts a budget it could not read above the ones it could, since only that row needs chasing", async () => { + mockBudgets([KEY_UNLIMITED, USER_WITHIN_BUDGET, UNRESOLVED_TEAM]); + const panel = await renderAndOpenBudgetsTab(); + + const [, ...dataRows] = panel.getAllByRole("row"); + expect(dataRows[0]).toHaveTextContent("team-123"); + }); + // The resolver reports a cold counter as the 0.0 it will be enforced as, and always attaches the // per-model note to a `no_counter` reading, so a fixture without both is one the server cannot produce. const COLD_PER_MODEL: KeyBudgetEntry = { diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 000280b2108..5180ec9d5c2 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -26360,6 +26360,10 @@ export interface components { /** * KeyBudgetEntry * @description One budget that can gate requests made with a key, with its live spend. + * + * ``status`` is ``unknown`` when the row could not be evaluated, either because the entity behind it + * was unreadable or because its spend was, and it is never ``unlimited`` in that case: an unreadable + * scope is not a scope the reader may rule out. */ KeyBudgetEntry: { /** Budget Duration */ @@ -26408,7 +26412,7 @@ export interface components { * Status * @enum {string} */ - status: "unlimited" | "ok" | "exceeded"; + status: "unlimited" | "ok" | "exceeded" | "unknown"; /** Window Start */ window_start?: string | null; }; @@ -26428,7 +26432,7 @@ export interface components { * Code * @enum {string} */ - code: "alert_only" | "custom_auth_may_override_end_user_cap" | "custom_auth_skips_read_time_checks" | "end_user_route_only" | "per_model_counters" | "project_spend_not_tracked" | "request_tags_add_budgets" | "reservation_blocks_at_limit" | "rolling_window" | "throttled_instead_of_blocked" | "user_budget_not_applied_to_team_key"; + code: "alert_only" | "custom_auth_may_override_end_user_cap" | "custom_auth_skips_read_time_checks" | "end_user_route_only" | "entity_unavailable" | "per_model_counters" | "project_spend_not_tracked" | "request_tags_add_budgets" | "reservation_blocks_at_limit" | "rolling_window" | "throttled_instead_of_blocked" | "user_budget_not_applied_to_team_key"; /** * Severity * @enum {string}