From a133a3c6f71a81a535230b2b765bb6fc457405ad Mon Sep 17 00:00:00 2001 From: mzolota Date: Fri, 21 Aug 2026 10:37:14 +0200 Subject: [PATCH 1/6] feat(proxy): attribute usage to every team a user belongs to Today a request is attributed to the single team stamped on the virtual key (or resolved from one JWT claim) and to that team's organization. A user who belongs to many teams contributes spend to whichever team the key happens to name and nothing to the rest, so "what did this team consume?" only has an answer for keys that name it. Add two opt-in general_settings, both defaulting to off: - track_spend_across_all_user_teams: spend increments, daily rollups, and budget gates apply to every team the caller belongs to, and to every organization reached through those teams. - enforce_rate_limits_across_all_user_teams: the same expansion for the RPM/TPM limiter, so a request must fit inside every membership's limit. They are separate settings because they carry different costs. Spend attribution is additive bookkeeping. Rate-limit expansion makes the caller's effective limit the minimum across their memberships, which is a live behavior change for anyone in a busy team. Memberships come from LiteLLM_UserTable.teams, already maintained by SCIM and JWT sync. Resolution happens once in the auth path and is carried on server-only UserAPIKeyAuth fields, stripped from validated input like mcp_source_team_rpm_limits so a caller cannot choose which buckets they are charged against. No migration. LiteLLM_DailyTeamSpend is already unique per (team_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint), and the spend queue aggregates by entity_type:entity_id, so N teams means N rows and N counters. Also make the rate limiter's per-hash-tag Lua calls concurrent. They were sequential, so on Redis Cluster one descriptor per team would have meant one round trip per team on the hot path. Non-cluster Redis was and remains a single call. Known limitation: organization rate limits still apply to the stamped org only. Multi-org spend attribution works; multi-org rate limiting would need per-org limits precomputed in auth and is left out deliberately to keep this change reviewable. --- litellm/proxy/_types.py | 39 ++ litellm/proxy/auth/auth_checks.py | 205 +++++++++++ litellm/proxy/auth/membership_attribution.py | 255 +++++++++++++ litellm/proxy/auth/user_api_key_auth.py | 15 + litellm/proxy/db/db_spend_update_writer.py | 188 ++++++---- .../hooks/parallel_request_limiter_v3.py | 78 +++- .../proxy/hooks/proxy_track_cost_callback.py | 39 ++ litellm/proxy/litellm_pre_call_utils.py | 40 +++ litellm/proxy/proxy_server.py | 33 +- .../proxy/auth/test_membership_attribution.py | 337 ++++++++++++++++++ .../hooks/test_proxy_track_cost_callback.py | 4 + 11 files changed, 1155 insertions(+), 78 deletions(-) create mode 100644 litellm/proxy/auth/membership_attribution.py create mode 100644 tests/test_litellm/proxy/auth/test_membership_attribution.py diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 00cbd13cfdc..a9ae2443100 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2367,6 +2367,28 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): ) use_google_kms: bool | None = Field(None, description="decrypt keys with google kms") use_azure_key_vault: bool | None = Field(None, description="load keys from azure key vault") + track_spend_across_all_user_teams: bool | None = Field( + None, + description=( + "attribute each request's spend to EVERY team the calling user belongs to " + "(and every organization reached through those teams), not only the team " + "stamped on the virtual key. Budget gates expand to match, so one " + "over-budget team blocks the caller everywhere. Summing team spend then " + "exceeds real spend by design, because one request is charged to several " + "teams; key, user, and org totals stay single-counted. Default off." + ), + ) + enforce_rate_limits_across_all_user_teams: bool | None = Field( + None, + description=( + "apply the RPM/TPM limits of EVERY team the calling user belongs to, not " + "only the team stamped on the virtual key. The caller's effective limit " + "becomes the minimum across their memberships, so a busy team can throttle " + "someone who is mostly working for a different team. Separate from " + "track_spend_across_all_user_teams so spend attribution can be adopted " + "without this. Default off." + ), + ) master_key: str | None = Field(None, description="require a key for all calls to proxy") coordination_redis: CoordinationRedisParams | None = Field( None, @@ -2757,6 +2779,20 @@ class LiteLLM_VerificationTokenView(LiteLLM_VerificationToken): organization_rpm_limit: int | None = None organization_metadata: dict | None = None + # Membership attribution params (see litellm/proxy/auth/membership_attribution.py). + # Populated during auth ONLY when track_spend_across_all_user_teams or + # enforce_rate_limits_across_all_user_teams is on; None otherwise, so the + # default single-team path is untouched. Server-only and stripped from + # validated input for the same reason as mcp_source_team_rpm_limits below: + # a forged entry would let a caller choose which teams they are charged + # against, or name a team with generous limits to escape their real ones. + attributed_team_ids: list[str] | None = Field(default=None, exclude=True) + attributed_org_ids: list[str] | None = Field(default=None, exclude=True) + # team_id -> {"rpm": int | None, "tpm": int | None}. Precomputed during auth + # because rate-limit descriptor construction is synchronous and cannot await + # a per-team lookup. + attributed_team_limits: dict[str, dict[str, int | None]] | None = Field(default=None, exclude=True) + # Project Params project_alias: str | None = None project_metadata: dict | None = None @@ -2863,6 +2899,9 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob values.pop("mcp_source_team_rpm_limits", None) values.pop("mcp_session_resource_server_id", None) values.pop("via_virtual_key", None) + values.pop("attributed_team_ids", None) + values.pop("attributed_org_ids", None) + values.pop("attributed_team_limits", None) if values.get("api_key") is not None: values.update({"token": cls._safe_hash_litellm_api_key(values.get("api_key"))}) if isinstance(values.get("api_key"), str): diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 12d6b44a648..fda9701450b 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -985,6 +985,23 @@ async def common_checks( user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ), + # No-ops unless track_spend_across_all_user_teams is on: with the + # setting off, attributed_team_ids/attributed_org_ids are None and + # both return before touching the cache. + _attributed_teams_max_budget_check( + valid_token=valid_token, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + general_settings=general_settings, + ), + _attributed_orgs_max_budget_check( + valid_token=valid_token, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + general_settings=general_settings, + ), _tag_max_budget_check( request_body=request_body, prisma_client=prisma_client, @@ -4665,6 +4682,194 @@ async def _team_max_budget_check( ) +async def _attributed_teams_max_budget_check( + valid_token: UserAPIKeyAuth | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, + general_settings: dict, +): + """Enforce the max budget of every OTHER team the caller belongs to. + + Only active when ``track_spend_across_all_user_teams`` is on -- otherwise + ``attributed_team_ids`` is None and this returns immediately. + + The stamped team is skipped because ``_team_max_budget_check`` already + covers it; checking it twice would fire a duplicate budget alert. Teams are + evaluated concurrently (each lookup is a cache hit in steady state) but the + raised error is chosen by membership order, not by which coroutine finished + first, so the same request always reports the same offending team. + + Consequence worth stating plainly: with attribution on, ONE exhausted team + blocks the caller on every team, because their usage counts against all of + them. + """ + from litellm.proxy.auth.membership_attribution import ( + attributed_team_ids, + spend_attribution_enabled, + ) + + # Gated on the SPEND setting specifically. The resolver populates + # attributed_team_ids when either setting is on, so an operator who enabled + # only rate-limit attribution must not silently get all-team budget gates. + if not spend_attribution_enabled(general_settings): + return + + if valid_token is None or prisma_client is None: + return + + other_team_ids: Final = [ + team_id for team_id in attributed_team_ids(valid_token) if team_id != valid_token.team_id + ] + if not other_team_ids: + return + + from litellm.proxy.proxy_server import get_current_spend + + async def _check(team_id: str) -> None: + try: + team_object = await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=valid_token.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + except Exception: # noqa: BLE001 # a team that cannot be loaded contributes no ceiling + return + + if team_object is None or team_object.max_budget is None or not math.isfinite(team_object.max_budget): + return + + spend = await get_current_spend( + counter_key=f"spend:team:{team_id}", + fallback_spend=team_object.spend or 0.0, + max_budget=team_object.max_budget, + ) + if spend <= team_object.max_budget: + return + + call_info = CallInfo( + token=valid_token.token, + spend=spend, + max_budget=team_object.max_budget, + user_id=valid_token.user_id, + team_id=team_id, + team_alias=team_object.team_alias, + organization_id=team_object.organization_id, + event_group=Litellm_EntityType.TEAM, + ) + asyncio.create_task( + proxy_logging_obj.budget_alerts( + type="team_budget", + user_info=call_info, + ) + ) + raise litellm.BudgetExceededError( + current_cost=spend, + max_budget=team_object.max_budget, + message=( + f"Budget has been exceeded! Team={team_id} Current cost: {spend}, " + f"Max budget: {team_object.max_budget}" + ), + entity_type=Litellm_EntityType.TEAM.value, + entity_id=team_id, + ) + + results: Final = await asyncio.gather(*(_check(team_id) for team_id in other_team_ids), return_exceptions=True) + first_error: Final = next((r for r in results if isinstance(r, BaseException)), None) + if first_error is not None: + raise first_error + + +async def _attributed_orgs_max_budget_check( + valid_token: UserAPIKeyAuth | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, + general_settings: dict, +): + """Enforce the max budget of every OTHER organization the caller belongs to. + + Skips ``valid_token.org_id`` -- ``_organization_max_budget_check`` owns that + one. Only reaches past it when the caller's teams span several + organizations; in a single-org deployment this is always a no-op. + """ + from litellm.proxy.auth.membership_attribution import ( + attributed_org_ids, + spend_attribution_enabled, + ) + + if not spend_attribution_enabled(general_settings): + return + + if valid_token is None or prisma_client is None: + return + + other_org_ids: Final = [org_id for org_id in attributed_org_ids(valid_token) if org_id != valid_token.org_id] + if not other_org_ids: + return + + from litellm.proxy.proxy_server import get_current_spend + + async def _check(org_id: str) -> None: + try: + org_table = await get_org_object( + org_id=org_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=valid_token.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + include_budget_table=True, + ) + except Exception: # noqa: BLE001 # an org that cannot be loaded contributes no ceiling + return + + if org_table is None or org_table.litellm_budget_table is None: + return + org_max_budget = org_table.litellm_budget_table.max_budget + if org_max_budget is None or org_max_budget <= 0 or not math.isfinite(org_max_budget): + return + + org_spend = await get_current_spend( + counter_key=f"spend:org:{org_id}", + fallback_spend=org_table.spend or 0.0, + max_budget=org_max_budget, + ) + if org_spend <= org_max_budget: + return + + asyncio.create_task( + proxy_logging_obj.budget_alerts( + type="team_budget", + user_info=CallInfo( + token=valid_token.token, + spend=org_spend, + max_budget=org_max_budget, + user_id=valid_token.user_id, + team_id=valid_token.team_id, + organization_id=org_id, + event_group=Litellm_EntityType.ORGANIZATION, + ), + ) + ) + raise litellm.BudgetExceededError( + current_cost=org_spend, + max_budget=org_max_budget, + message=( + f"Budget has been exceeded! Organization={org_id} Current cost: {org_spend}, " + f"Max budget: {org_max_budget}" + ), + entity_type=Litellm_EntityType.ORGANIZATION.value, + entity_id=org_id, + ) + + results: Final = await asyncio.gather(*(_check(org_id) for org_id in other_org_ids), return_exceptions=True) + first_error: Final = next((r for r in results if isinstance(r, BaseException)), None) + if first_error is not None: + raise first_error + + async def _team_multi_budget_check( team_object: LiteLLM_TeamTable | None, ): diff --git a/litellm/proxy/auth/membership_attribution.py b/litellm/proxy/auth/membership_attribution.py new file mode 100644 index 00000000000..124b346fa43 --- /dev/null +++ b/litellm/proxy/auth/membership_attribution.py @@ -0,0 +1,255 @@ +"""Membership-based usage attribution. + +By default LiteLLM attributes a request to the single team stamped on the +virtual key (or the single team a JWT claim resolved to), and to that team's +organization. A user who belongs to many teams therefore contributes spend to +whichever team the key happens to name and nothing to the rest, and an operator +who wants "what did this team consume?" only gets an answer for keys that +happen to name it. + +Two opt-in settings change that. Both default to off, so an existing deployment +keeps the single-team behavior byte for byte: + +``track_spend_across_all_user_teams`` + Spend increments, daily rollups, and budget gates apply to every team the + caller belongs to, and to every organization reached through those teams. + +``enforce_rate_limits_across_all_user_teams`` + The same expansion for the RPM/TPM limiter, so a request must fit inside + every membership's limit rather than only the stamped team's. + +They are separate settings because they carry different costs. Spend +attribution is additive bookkeeping: the only surprise is that summing team +spend now exceeds real spend, because one request is deliberately charged to +several teams. Rate-limit expansion is a live behavior change: a caller's +effective limit becomes the MINIMUM across their memberships, so a busy team +can throttle someone who is mostly working for a different team. Operators +should be able to adopt the first without the second. + +Nothing here builds a team tree. LiteLLM teams are a flat set, each optionally +belonging to one organization -- there is no ``parent_team_id`` in the schema. +"All memberships" therefore means the caller's teams plus the organizations +those teams (and the caller's own user row) belong to, never a recursive walk. +""" + +import asyncio +from typing import TYPE_CHECKING, Any, Final + +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import LiteLLM_UserTable, UserAPIKeyAuth + +if TYPE_CHECKING: + from litellm.caching.caching import DualCache + from litellm.proxy.utils import PrismaClient, ProxyLogging + +SPEND_ATTRIBUTION_SETTING: Final = "track_spend_across_all_user_teams" +RATE_LIMIT_ATTRIBUTION_SETTING: Final = "enforce_rate_limits_across_all_user_teams" + + +def spend_attribution_enabled(general_settings: dict | None) -> bool: + """Whether spend, rollups, and budget gates fan out across memberships.""" + if not general_settings: + return False + return general_settings.get(SPEND_ATTRIBUTION_SETTING) is True + + +def rate_limit_attribution_enabled(general_settings: dict | None) -> bool: + """Whether the RPM/TPM limiter fans out across memberships.""" + if not general_settings: + return False + return general_settings.get(RATE_LIMIT_ATTRIBUTION_SETTING) is True + + +def _attribution_enabled(general_settings: dict | None) -> bool: + return spend_attribution_enabled(general_settings) or rate_limit_attribution_enabled(general_settings) + + +def attributed_team_ids(valid_token: UserAPIKeyAuth | None) -> list[str]: + """Every team this request is attributed to, most-specific first. + + Falls back to the single stamped team whenever attribution is off or + resolved nothing, so a call site can use this unconditionally and keep + identical behavior with the settings disabled. + """ + if valid_token is None: + return [] + resolved: Final = valid_token.attributed_team_ids + if resolved: + return list(resolved) + return [valid_token.team_id] if valid_token.team_id else [] + + +def attributed_org_ids(valid_token: UserAPIKeyAuth | None) -> list[str]: + """Every organization this request is attributed to. + + Same fallback contract as :func:`attributed_team_ids`. + """ + if valid_token is None: + return [] + resolved: Final = valid_token.attributed_org_ids + if resolved: + return list(resolved) + return [valid_token.org_id] if valid_token.org_id else [] + + +def attribution_targets(attributed_ids: list[str] | None, stamped_id: str | None) -> list[str]: + """The entity ids one request should be charged against. + + The attributed set when membership attribution resolved one, otherwise the + single stamped id -- so with both settings off this returns exactly + ``[stamped_id]`` and every caller keeps its historical behavior. + + Order-preserving dedupe: the stamped team normally also appears in the + caller's membership list, and charging it twice would double-count. + """ + if attributed_ids: + return list(dict.fromkeys(i for i in attributed_ids if i)) + return [stamped_id] if stamped_id else [] + + +def _ordered_unique(values: list[str | None]) -> list[str]: + """Dedupe while preserving order, dropping empties. + + Order is preserved so the stamped team stays first. Budget and rate-limit + errors report the first offending entity, and a caller reading that error + is best served by hearing about the team their key actually names. + """ + return list(dict.fromkeys(v for v in values if v)) + + +async def resolve_membership_attribution( + *, + user_api_key_auth_obj: UserAPIKeyAuth, + user_object: LiteLLM_UserTable | None, + general_settings: dict | None, + prisma_client: "PrismaClient | None", + user_api_key_cache: "DualCache", + proxy_logging_obj: "ProxyLogging | None" = None, +) -> None: + """Populate the attributed-membership fields on ``user_api_key_auth_obj``. + + No-op unless one of the attribution settings is on, so the default path + pays nothing -- not even a cache read. + + Fails open. A team that cannot be resolved (deleted row, transient DB + error) is skipped rather than raised: attribution is bookkeeping layered on + top of an authorization decision that has already been made, and losing one + team's attribution must never turn an authorized request into a 500. The + stamped team is seeded first and is never dropped, so a failure degrades to + exactly the default single-team behavior. + """ + if not _attribution_enabled(general_settings): + return + + if prisma_client is None: + return + + # The stamped team leads, then the caller's SCIM/IdP-maintained membership + # list. Seeding the stamped team explicitly matters: LiteLLM_UserTable.teams + # is maintained by SCIM and JWT sync, so it can lag a key that was just + # pointed at a new team, and losing that team would silently under-charge + # the one team the operator explicitly named. + candidate_team_ids: Final = _ordered_unique( + [user_api_key_auth_obj.team_id, *(list(user_object.teams) if user_object and user_object.teams else [])] + ) + + if not candidate_team_ids: + _apply_org_only_attribution( + user_api_key_auth_obj=user_api_key_auth_obj, + user_object=user_object, + team_objects=[], + ) + return + + team_objects: Final = await _load_team_objects( + team_ids=candidate_team_ids, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=user_api_key_auth_obj.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + + resolved_ids: Final = [team_id for team_id, team_object in team_objects if team_object is not None] + if resolved_ids: + user_api_key_auth_obj.attributed_team_ids = resolved_ids + user_api_key_auth_obj.attributed_team_limits = { + team_id: { + "rpm": getattr(team_object, "rpm_limit", None), + "tpm": getattr(team_object, "tpm_limit", None), + } + for team_id, team_object in team_objects + if team_object is not None + } + + _apply_org_only_attribution( + user_api_key_auth_obj=user_api_key_auth_obj, + user_object=user_object, + team_objects=team_objects, + ) + + +def _apply_org_only_attribution( + *, + user_api_key_auth_obj: UserAPIKeyAuth, + user_object: LiteLLM_UserTable | None, + team_objects: list[tuple[str, Any]], +) -> None: + """Derive the attributed organizations from what is already loaded. + + Sources, in order: the org already resolved onto the token, the caller's own + ``LiteLLM_UserTable.organization_id``, then the org of each attributed team. + Deliberately no ``LiteLLM_OrganizationMembership`` query -- that table would + add a fresh round trip to the hot path, and these three sources already + cover both the common single-org deployment and a caller whose teams span + several organizations. + """ + org_ids: Final = _ordered_unique( + [ + user_api_key_auth_obj.org_id, + user_object.organization_id if user_object else None, + *[ + getattr(team_object, "organization_id", None) + for _team_id, team_object in team_objects + if team_object is not None + ], + ] + ) + if org_ids: + user_api_key_auth_obj.attributed_org_ids = org_ids + + +async def _load_team_objects( + *, + team_ids: list[str], + prisma_client: "PrismaClient", + user_api_key_cache: "DualCache", + parent_otel_span: Any, + proxy_logging_obj: "ProxyLogging | None", +) -> list[tuple[str, Any]]: + """Resolve every candidate team, preserving input order. + + Lookups run concurrently: ``get_team_object`` is cache-first, so the steady + state is N in-memory hits, but a cold pod pays a DB read per team and those + must not serialize. + """ + from litellm.proxy.auth.auth_checks import get_team_object + + async def _safe_get(team_id: str) -> tuple[str, Any]: + try: + team_object = await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + return team_id, team_object + except Exception as e: # noqa: BLE001 # attribution must never fail an authorized request + verbose_proxy_logger.debug( + "membership attribution: skipping team_id=%s, could not resolve: %s", + team_id, + e, + ) + return team_id, None + + return list(await asyncio.gather(*(_safe_get(team_id) for team_id in team_ids))) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 99592d44f9b..2e8b3872c3e 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -68,6 +68,7 @@ from litellm.proxy.auth.auth_utils import ( route_in_additonal_public_routes, ) from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler +from litellm.proxy.auth.membership_attribution import resolve_membership_attribution from litellm.proxy.auth.network import TrustedProxyConfig, resolve_network_context from litellm.proxy.auth.oauth2_check import Oauth2Handler from litellm.proxy.auth.oauth2_proxy_hook import handle_oauth2_proxy_request @@ -2381,6 +2382,20 @@ async def _run_centralized_common_checks( if user_api_key_auth_obj.org_id is None and team_object is not None and team_object.organization_id is not None: user_api_key_auth_obj.org_id = team_object.organization_id + # Expand the single stamped team/org into every membership the caller has, + # when the operator opted in. Must run BEFORE common_checks so the + # all-teams budget gate sees the expansion, and before the request metadata + # is stamped so the cost callback can attribute spend to the same set. + # No-op (not even a cache read) when both settings are off. + await resolve_membership_attribution( + user_api_key_auth_obj=user_api_key_auth_obj, + user_object=user_object, + general_settings=general_settings, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + # common_checks identifies admin via user_object, not the token # (non_proxy_admin_allowed_routes_check). JWT admin shortcut and # master_key tokens get admin from the token; the DB row for the diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 65a271d4029..49520a869a5 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -156,6 +156,11 @@ class DBSpendUpdateWriter: start_time: datetime | None, end_time: datetime | None, response_cost: float | None, + # Every team/org this request is attributed to, when + # track_spend_across_all_user_teams is on. None keeps the historical + # single-team/single-org behavior. + attributed_team_ids: list[str] | None = None, + attributed_org_ids: list[str] | None = None, ): from litellm.proxy.proxy_server import ( disable_spend_logs, @@ -234,6 +239,8 @@ class DBSpendUpdateWriter: prisma_client=prisma_client, litellm_proxy_budget_name=litellm_proxy_budget_name, payload=payload, + attributed_team_ids=attributed_team_ids, + attributed_org_ids=attributed_org_ids, ) ) @@ -424,6 +431,8 @@ class DBSpendUpdateWriter: prisma_client: PrismaClient | None, litellm_proxy_budget_name: str | None, payload: SpendLogsPayload, + attributed_team_ids: list[str] | None = None, + attributed_org_ids: list[str] | None = None, ): """ Runs all 11 spend-update helpers sequentially inside a single asyncio task. @@ -467,6 +476,7 @@ class DBSpendUpdateWriter: team_id=team_id, user_id=user_id, prisma_client=prisma_client, + attributed_team_ids=attributed_team_ids, ) except Exception: verbose_proxy_logger.debug( @@ -479,6 +489,7 @@ class DBSpendUpdateWriter: response_cost=response_cost, org_id=org_id, prisma_client=prisma_client, + attributed_org_ids=attributed_org_ids, ) except Exception: verbose_proxy_logger.debug( @@ -548,6 +559,7 @@ class DBSpendUpdateWriter: await self.add_spend_log_transaction_to_daily_team_transaction( payload=payload_copy, prisma_client=prisma_client, + attributed_team_ids=attributed_team_ids, ) except Exception: verbose_proxy_logger.debug( @@ -560,6 +572,7 @@ class DBSpendUpdateWriter: payload=payload_copy, org_id=org_id, prisma_client=prisma_client, + attributed_org_ids=attributed_org_ids, ) except Exception: verbose_proxy_logger.debug( @@ -646,54 +659,74 @@ class DBSpendUpdateWriter: exc=e, ) + @staticmethod + def _attribution_targets(attributed_ids: list[str] | None, stamped_id: str | None) -> list[str]: + """The entity ids to charge for one request. See ``attribution_targets``.""" + from litellm.proxy.auth.membership_attribution import attribution_targets + + return attribution_targets(attributed_ids, stamped_id) + async def _update_team_db( self, response_cost: float | None, team_id: str | None, user_id: str | None, prisma_client: PrismaClient | None, + attributed_team_ids: list[str] | None = None, ): + """Charge the request to every team it is attributed to. + + With ``track_spend_across_all_user_teams`` off, ``attributed_team_ids`` + is None and this charges the single stamped team, exactly as before. + + Note the deliberate consequence when it is on: summing team spend across + the deployment then exceeds real spend, because one request is charged to + each of the caller's teams on purpose. The key, user, and org rows stay + single-counted, so those remain true totals. + """ + target_team_ids: Final = self._attribution_targets(attributed_team_ids, team_id) try: - if team_id is None or prisma_client is None: + if not target_team_ids or prisma_client is None: verbose_proxy_logger.debug( - "track_cost_callback: team_id is None or prisma_client is None. Not tracking spend for team" + "track_cost_callback: no attributed team or prisma_client is None. Not tracking spend for team" ) return - await self.spend_update_queue.add_update( - update=SpendUpdateQueueItem( - entity_type=Litellm_EntityType.TEAM, - entity_id=team_id, - response_cost=response_cost, - ) - ) - - try: - # Track spend of the team member within this team - if user_id is not None: - # key is "team_id::::user_id::" - team_member_key: Final = f"team_id::{team_id}::user_id::{user_id}" - await self.spend_update_queue.add_update( - update=SpendUpdateQueueItem( - entity_type=Litellm_EntityType.TEAM_MEMBER, - entity_id=team_member_key, - response_cost=response_cost, - ) + for target_team_id in target_team_ids: + await self.spend_update_queue.add_update( + update=SpendUpdateQueueItem( + entity_type=Litellm_EntityType.TEAM, + entity_id=target_team_id, + response_cost=response_cost, ) - except Exception as e: - spend_log_error( - "Spend tracking - failed to enqueue team member spend update. " - "team_id=%s, user_id=%s, response_cost=%s - %s", - team_id, - user_id, - response_cost, - str(e), - exc=e, ) + + try: + # Track spend of the team member within this team + if user_id is not None: + # key is "team_id::::user_id::" + team_member_key: Final = f"team_id::{target_team_id}::user_id::{user_id}" + await self.spend_update_queue.add_update( + update=SpendUpdateQueueItem( + entity_type=Litellm_EntityType.TEAM_MEMBER, + entity_id=team_member_key, + response_cost=response_cost, + ) + ) + except Exception as e: + spend_log_error( + "Spend tracking - failed to enqueue team member spend update. " + "team_id=%s, user_id=%s, response_cost=%s - %s", + target_team_id, + user_id, + response_cost, + str(e), + exc=e, + ) except Exception as e: spend_log_error( - "Spend tracking - failed to enqueue team spend update. team_id=%s, response_cost=%s - %s", - team_id, + "Spend tracking - failed to enqueue team spend update. team_ids=%s, response_cost=%s - %s", + target_team_ids, response_cost, str(e), exc=e, @@ -705,25 +738,33 @@ class DBSpendUpdateWriter: response_cost: float | None, org_id: str | None, prisma_client: PrismaClient | None, + attributed_org_ids: list[str] | None = None, ): + """Charge the request to every organization it is attributed to. + + Normally a single org: a caller belongs to one, and extra ids only + appear when their teams span several organizations. + """ + target_org_ids: Final = self._attribution_targets(attributed_org_ids, org_id) try: - if org_id is None or prisma_client is None: + if not target_org_ids or prisma_client is None: verbose_proxy_logger.debug( - "track_cost_callback: org_id is None or prisma_client is None. Not tracking spend for org" + "track_cost_callback: no attributed org or prisma_client is None. Not tracking spend for org" ) return - await self.spend_update_queue.add_update( - update=SpendUpdateQueueItem( - entity_type=Litellm_EntityType.ORGANIZATION, - entity_id=org_id, - response_cost=response_cost, + for target_org_id in target_org_ids: + await self.spend_update_queue.add_update( + update=SpendUpdateQueueItem( + entity_type=Litellm_EntityType.ORGANIZATION, + entity_id=target_org_id, + response_cost=response_cost, + ) ) - ) except Exception as e: spend_log_error( - "Spend tracking - failed to enqueue org spend update. org_id=%s, response_cost=%s - %s", - org_id, + "Spend tracking - failed to enqueue org spend update. org_ids=%s, response_cost=%s - %s", + target_org_ids, response_cost, str(e), exc=e, @@ -1940,7 +1981,17 @@ class DBSpendUpdateWriter: self, payload: SpendLogsPayload, prisma_client: PrismaClient | None = None, + attributed_team_ids: list[str] | None = None, ) -> None: + """Enqueue one daily rollup row per attributed team. + + ``LiteLLM_DailyTeamSpend`` is already unique on + ``(team_id, date, api_key, model, custom_llm_provider, + mcp_namespaced_tool_name, endpoint)``, so N teams means N distinct rows + and no migration is needed. The base transaction is computed once and + reused: it holds only per-request facts (tokens, cost, date, endpoint), + none of which depend on which team the row is keyed to. + """ if prisma_client is None: verbose_proxy_logger.debug("prisma_client is None. Skipping writing spend logs to db.") return @@ -1950,47 +2001,56 @@ class DBSpendUpdateWriter: ) if base_daily_transaction is None: return - if payload["team_id"] is None: - verbose_proxy_logger.debug("team_id is None for request. Skipping incrementing team spend.") + + target_team_ids: Final = self._attribution_targets(attributed_team_ids, payload["team_id"]) + if not target_team_ids: + verbose_proxy_logger.debug("no attributed team for request. Skipping incrementing team spend.") return endpoint_str: Final = base_daily_transaction.get("endpoint") or "" - daily_transaction_key = f"{payload['team_id']}_{base_daily_transaction['date']}_{payload['api_key']}_{payload['model']}_{payload['custom_llm_provider']}_{endpoint_str}" - daily_transaction: Final = DailyTeamSpendTransaction(team_id=payload["team_id"], **base_daily_transaction) - await self.daily_team_spend_update_queue.add_update(update={daily_transaction_key: daily_transaction}) + for target_team_id in target_team_ids: + daily_transaction_key = f"{target_team_id}_{base_daily_transaction['date']}_{payload['api_key']}_{payload['model']}_{payload['custom_llm_provider']}_{endpoint_str}" + daily_transaction = DailyTeamSpendTransaction(team_id=target_team_id, **base_daily_transaction) + await self.daily_team_spend_update_queue.add_update(update={daily_transaction_key: daily_transaction}) async def add_spend_log_transaction_to_daily_org_transaction( self, payload: SpendLogsPayload, prisma_client: PrismaClient | None = None, org_id: str | None = None, + attributed_org_ids: list[str] | None = None, ) -> None: + """Enqueue one daily rollup row per attributed organization.""" if prisma_client is None: verbose_proxy_logger.debug("prisma_client is None. Skipping writing spend logs to db.") return - if org_id is None: + target_org_ids: Final = self._attribution_targets(attributed_org_ids, org_id) + if not target_org_ids: verbose_proxy_logger.debug("organization_id is None for request. Skipping incrementing organization spend.") return - payload_with_org: Final = cast( - SpendLogsPayload, - { - **payload, - "organization_id": org_id, - }, - ) + for target_org_id in target_org_ids: + payload_with_org = cast( + SpendLogsPayload, + { + **payload, + "organization_id": target_org_id, + }, + ) - base_daily_transaction: Final = await self._common_add_spend_log_transaction_to_daily_transaction( - payload_with_org, prisma_client, "org" - ) - if base_daily_transaction is None: - return + base_daily_transaction = await self._common_add_spend_log_transaction_to_daily_transaction( + payload_with_org, prisma_client, "org" + ) + if base_daily_transaction is None: + continue - endpoint_str: Final = base_daily_transaction.get("endpoint") or "" - daily_transaction_key = f"{org_id}_{base_daily_transaction['date']}_{payload_with_org['api_key']}_{payload_with_org['model']}_{payload_with_org['custom_llm_provider']}_{endpoint_str}" - daily_transaction: Final = DailyOrganizationSpendTransaction(organization_id=org_id, **base_daily_transaction) - await self.daily_org_spend_update_queue.add_update(update={daily_transaction_key: daily_transaction}) + endpoint_str = base_daily_transaction.get("endpoint") or "" + daily_transaction_key = f"{target_org_id}_{base_daily_transaction['date']}_{payload_with_org['api_key']}_{payload_with_org['model']}_{payload_with_org['custom_llm_provider']}_{endpoint_str}" + daily_transaction = DailyOrganizationSpendTransaction( + organization_id=target_org_id, **base_daily_transaction + ) + await self.daily_org_spend_update_queue.add_update(update={daily_transaction_key: daily_transaction}) async def add_spend_log_transaction_to_daily_end_user_transaction( self, diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 1e65da5b867..fa00a42decf 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -1201,22 +1201,34 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): key_groups: Final = self._group_keys_by_hash_tag(keys_to_fetch) all_cache_values: Final[list[CacheCounterValue | None]] = [] - for hash_tag, group_keys in key_groups.items(): + async def _read_group(hash_tag: str, group_keys: list[str]) -> CacheCounterValues: try: - group_cache_values: CacheCounterValues = await self.batch_rate_limiter_script( + return await self.batch_rate_limiter_script( keys=group_keys, args=[now_int, self.window_size], # Use integer timestamp ) - all_cache_values.extend(group_cache_values) except Exception as e: verbose_proxy_logger.warning("Redis Lua script failed for hash tag %s: %s", hash_tag, e) # Fallback to in-memory cache for this group - group_cache_values = await self.in_memory_cache_sliding_window( + return await self.in_memory_cache_sliding_window( keys=group_keys, now_int=now_int, window_size=self.window_size, ) - all_cache_values.extend(group_cache_values) + + # One Lua round trip per hash-tag group, issued concurrently. On + # non-cluster Redis there is exactly one group, so this stays a single + # call no matter how many descriptors the request carries. On Redis + # Cluster the groups are per-slot, and membership-attributed team + # descriptors can produce one group per team -- running them + # concurrently keeps the added latency at roughly one round trip + # instead of N. Results are consumed in group order so the returned + # list still lines up with ``keys_to_fetch``. + group_results: Final = await asyncio.gather( + *(_read_group(hash_tag, group_keys) for hash_tag, group_keys in key_groups.items()) + ) + for group_cache_values in group_results: + all_cache_values.extend(group_cache_values) return all_cache_values @@ -2350,6 +2362,57 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) ) + def _add_attributed_team_rate_limit_descriptors( + self, + user_api_key_dict: UserAPIKeyAuth, + descriptors: list[RateLimitDescriptor], + ) -> None: + """Add a team descriptor for every OTHER team the caller belongs to. + + Only active when ``enforce_rate_limits_across_all_user_teams`` is on; + otherwise ``attributed_team_limits`` is None and this returns at once. + + ``key="team"`` is reused deliberately: team T must share one bucket + whether T is the team stamped on the key or one the caller merely + belongs to, or the same team would be limited twice over under two + different counters. The stamped team is skipped because the block above + already emitted its descriptor. + + Because ``should_rate_limit`` rejects when ANY descriptor is over, the + caller's effective limit becomes the minimum across their memberships. + That is the intended semantic, and it is why this rides a separate + setting from spend attribution: a busy team can now throttle someone + who is mostly working for a different team. + """ + from litellm.proxy.auth.membership_attribution import rate_limit_attribution_enabled + from litellm.proxy.proxy_server import general_settings + + if not rate_limit_attribution_enabled(general_settings): + return + + attributed_limits: Final = user_api_key_dict.attributed_team_limits + if not attributed_limits: + return + + for team_id, team_limits in attributed_limits.items(): + if not team_id or team_id == user_api_key_dict.team_id: + continue + rpm_limit = team_limits.get("rpm") + tpm_limit = team_limits.get("tpm") + if rpm_limit is None and tpm_limit is None: + continue + descriptors.append( + RateLimitDescriptor( + key="team", + value=team_id, + rate_limit={ + "requests_per_unit": rpm_limit, + "tokens_per_unit": tpm_limit, + "window_size": self.window_size, + }, + ) + ) + def _add_tag_per_key_rate_limit_descriptor( self, user_api_key_dict: UserAPIKeyAuth, @@ -2701,6 +2764,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) ) + self._add_attributed_team_rate_limit_descriptors( + user_api_key_dict=user_api_key_dict, + descriptors=descriptors, + ) + # Team Member rate limits if user_api_key_dict.user_id and ( user_api_key_dict.team_member_rpm_limit is not None or user_api_key_dict.team_member_tpm_limit is not None diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 99d0c94d11b..df9d67d2751 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -205,6 +205,8 @@ class _ProxyDBLogger(CustomLogger): start_time=actual_start_time, end_time=datetime.now(), org_id=user_api_key_dict.org_id, + attributed_team_ids=_attributed_team_ids(user_api_key_dict), + attributed_org_ids=_attributed_org_ids(user_api_key_dict), ) @log_db_metrics @@ -244,6 +246,10 @@ class _ProxyDBLogger(CustomLogger): user_id: Final = cast(str | None, metadata.get("user_api_key_user_id", None)) team_id: Final = cast(str | None, metadata.get("user_api_key_team_id", None)) org_id: Final = cast(str | None, metadata.get("user_api_key_org_id", None)) + # Present only when track_spend_across_all_user_teams is on; None + # otherwise, which keeps the single-team write path unchanged. + attributed_team_ids: Final = _metadata_id_list(metadata, "user_api_key_attributed_team_ids") + attributed_org_ids: Final = _metadata_id_list(metadata, "user_api_key_attributed_org_ids") key_alias: Final = cast(str | None, metadata.get("user_api_key_alias", None)) end_user_max_budget: Final = metadata.get("user_api_end_user_max_budget", None) sl_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None) @@ -292,6 +298,8 @@ class _ProxyDBLogger(CustomLogger): response_cost=response_cost, budget_reservation=budget_reservation, request_tags=tags, + attributed_team_ids=attributed_team_ids, + attributed_org_ids=attributed_org_ids, ) # update cache (fire-and-forget for backward compat: @@ -504,6 +512,31 @@ def _get_budget_reservation_from_metadata(metadata: dict) -> dict | None: return getattr(user_api_key_auth_obj, "budget_reservation", None) +def _metadata_id_list(metadata: dict, key: str) -> list[str] | None: + """Read a stamped id list out of request metadata, or None if absent. + + Returns None (not []) when the key is missing so downstream writers can + tell "attribution off" apart from "attribution on, nothing resolved". + """ + value: Final = metadata.get(key) + if not isinstance(value, list): + return None + ids: Final = [v for v in value if isinstance(v, str) and v] + return ids or None + + +def _attributed_team_ids(user_api_key_dict: UserAPIKeyAuth) -> list[str] | None: + from litellm.proxy.auth.membership_attribution import attributed_team_ids + + return attributed_team_ids(user_api_key_dict) if user_api_key_dict.attributed_team_ids else None + + +def _attributed_org_ids(user_api_key_dict: UserAPIKeyAuth) -> list[str] | None: + from litellm.proxy.auth.membership_attribution import attributed_org_ids + + return attributed_org_ids(user_api_key_dict) if user_api_key_dict.attributed_org_ids else None + + def _get_request_tags_for_cost_tracking( sl_object: StandardLoggingPayload | None, metadata: dict, @@ -535,6 +568,8 @@ async def _update_database_and_spend_counters( response_cost: float, budget_reservation: dict | None, request_tags: list[str] | None = None, + attributed_team_ids: list[str] | None = None, + attributed_org_ids: list[str] | None = None, ) -> None: try: await proxy_logging_obj.db_spend_update_writer.update_database( @@ -548,6 +583,8 @@ async def _update_database_and_spend_counters( start_time=start_time, end_time=end_time, org_id=org_id, + attributed_team_ids=attributed_team_ids, + attributed_org_ids=attributed_org_ids, ) except Exception: if budget_reservation is not None: @@ -573,6 +610,8 @@ async def _update_database_and_spend_counters( budget_reservation=budget_reservation, end_user_id=end_user_id, tags=request_tags, + attributed_team_ids=attributed_team_ids, + attributed_org_ids=attributed_org_ids, ) except Exception: if budget_reservation is not None: diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 2ec5c34958c..33e32a0615b 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -1306,6 +1306,19 @@ class LiteLLMProxyRequestSetup: ) if user_api_key_dict.budget_reservation is not None: data[_metadata_variable_name]["user_api_key_budget_reservation"] = user_api_key_dict.budget_reservation + + # Every team/org this request is attributed to, stamped as plain metadata + # so the cost callback reads it the same way it reads + # user_api_key_team_id. Written only when track_spend_across_all_user_teams + # is on, so the default path adds no metadata keys at all. Deliberately + # not folded into StandardLoggingUserAPIKeyMetadata: that TypedDict is + # consumed by ~20 logging integrations and widening it would make this a + # breaking payload change rather than an additive proxy feature. + LiteLLMProxyRequestSetup._add_attributed_membership_metadata( + data=data, + user_api_key_dict=user_api_key_dict, + _metadata_variable_name=_metadata_variable_name, + ) # UserAPIKeyAuth object for MCP server access control data[_metadata_variable_name]["user_api_key_auth"] = user_api_key_dict.model_copy( update={ @@ -1317,6 +1330,33 @@ class LiteLLMProxyRequestSetup: ) return data + @staticmethod + def _add_attributed_membership_metadata( + data: dict, + user_api_key_dict: UserAPIKeyAuth, + _metadata_variable_name: str, + ) -> None: + """Stamp the attributed team/org ids for the cost callback. + + No-op unless ``track_spend_across_all_user_teams`` is on AND auth + actually resolved a membership set, so a deployment with the setting off + carries exactly the metadata it carries today. + """ + from litellm.proxy.auth.membership_attribution import ( + attributed_org_ids, + attributed_team_ids, + spend_attribution_enabled, + ) + from litellm.proxy.proxy_server import general_settings + + if not spend_attribution_enabled(general_settings): + return + if user_api_key_dict.attributed_team_ids is None and user_api_key_dict.attributed_org_ids is None: + return + + data[_metadata_variable_name]["user_api_key_attributed_team_ids"] = attributed_team_ids(user_api_key_dict) + data[_metadata_variable_name]["user_api_key_attributed_org_ids"] = attributed_org_ids(user_api_key_dict) + @staticmethod def add_management_endpoint_metadata_to_request_metadata( data: dict, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 9ee62f94647..d61fd41de75 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2609,6 +2609,8 @@ async def increment_spend_counters( budget_reservation: dict | None = None, end_user_id: str | None = None, tags: list[str] | None = None, + attributed_team_ids: list[str] | None = None, + attributed_org_ids: list[str] | None = None, ): """ Atomically increment spend counters for budget enforcement. @@ -2723,12 +2725,24 @@ async def increment_spend_counters( increment=cost, ) + # With track_spend_across_all_user_teams on, one request increments the + # counter of every team the caller belongs to, and of every org reached + # through those teams. These are the counters the budget gates read + # Redis-first, so enforcement across memberships is immediate rather than + # waiting for the batched DB flush. With the setting off, both lists + # collapse to the single stamped id and the scopes are unchanged. + target_team_ids: Final = attribution_targets(attributed_team_ids, team_id) + target_org_ids: Final = attribution_targets(attributed_org_ids, org_id) + scope_coros: Final = tuple( coro for coro in ( _key_scope(token) if token is not None else None, - _team_scope(team_id) if team_id is not None else None, - _team_member_scope(user_id, team_id) if user_id is not None and team_id is not None else None, + *(_team_scope(scope_team_id) for scope_team_id in target_team_ids), + *( + _team_member_scope(user_id, scope_team_id) + for scope_team_id in (target_team_ids if user_id is not None else ()) + ), _user_scope(user_id) if user_id is not None else None, _increment_end_user_and_tag_spend_counters( end_user_id=end_user_id, @@ -2738,13 +2752,14 @@ async def increment_spend_counters( ) if end_user_id is not None or tags is not None else None, - _increment_org_spend_counter( - org_id=org_id, - response_cost=cost, - reserved_counter_keys=reserved_counter_keys, - ) - if org_id is not None - else None, + *( + _increment_org_spend_counter( + org_id=scope_org_id, + response_cost=cost, + reserved_counter_keys=reserved_counter_keys, + ) + for scope_org_id in target_org_ids + ), ) if coro is not None ) diff --git a/tests/test_litellm/proxy/auth/test_membership_attribution.py b/tests/test_litellm/proxy/auth/test_membership_attribution.py new file mode 100644 index 00000000000..417ee7523d0 --- /dev/null +++ b/tests/test_litellm/proxy/auth/test_membership_attribution.py @@ -0,0 +1,337 @@ +"""Tests for membership-based usage attribution. + +Covers the two opt-in settings introduced alongside +``litellm/proxy/auth/membership_attribution.py``: + +- ``track_spend_across_all_user_teams`` +- ``enforce_rate_limits_across_all_user_teams`` + +The most important cases here are the OFF cases. Both settings default to off, +and every one of those tests is a regression guard proving the default path is +byte-for-byte what it was before the feature existed. +""" + +import os +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.proxy._types import Litellm_EntityType, LiteLLM_UserTable, UserAPIKeyAuth +from litellm.proxy.auth.membership_attribution import ( + attributed_org_ids, + attributed_team_ids, + attribution_targets, + rate_limit_attribution_enabled, + resolve_membership_attribution, + spend_attribution_enabled, +) +from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter + + +# --------------------------------------------------------------------------- # +# settings gates +# --------------------------------------------------------------------------- # +@pytest.mark.parametrize("general_settings", [None, {}, {"track_spend_across_all_user_teams": False}]) +def test_spend_attribution_defaults_off(general_settings): + assert spend_attribution_enabled(general_settings) is False + + +@pytest.mark.parametrize("general_settings", [None, {}, {"enforce_rate_limits_across_all_user_teams": False}]) +def test_rate_limit_attribution_defaults_off(general_settings): + assert rate_limit_attribution_enabled(general_settings) is False + + +def test_settings_are_independent(): + """An operator must be able to take spend attribution without rate limits.""" + spend_only = {"track_spend_across_all_user_teams": True} + assert spend_attribution_enabled(spend_only) is True + assert rate_limit_attribution_enabled(spend_only) is False + + +# --------------------------------------------------------------------------- # +# target resolution +# --------------------------------------------------------------------------- # +def test_attribution_targets_falls_back_to_stamped_id(): + assert attribution_targets(None, "team-a") == ["team-a"] + assert attribution_targets([], "team-a") == ["team-a"] + assert attribution_targets(None, None) == [] + + +def test_attribution_targets_dedupes_and_preserves_order(): + """The stamped team is normally also in the membership list; charge it once.""" + assert attribution_targets(["team-a", "team-b", "team-a", ""], "team-a") == ["team-a", "team-b"] + + +def test_read_helpers_fall_back_to_stamped_values(): + token = UserAPIKeyAuth(api_key="sk-1", team_id="team-a", org_id="org-1") + assert attributed_team_ids(token) == ["team-a"] + assert attributed_org_ids(token) == ["org-1"] + + +def test_attributed_fields_cannot_be_forged_from_input(): + """Server-only fields. A caller who could set these would pick their own + budget and rate-limit buckets.""" + forged = UserAPIKeyAuth( + api_key="sk-1", + team_id="team-a", + attributed_team_ids=["team-with-huge-budget"], + attributed_org_ids=["org-with-huge-budget"], + attributed_team_limits={"team-with-huge-budget": {"rpm": 10**9}}, + ) + assert forged.attributed_team_ids is None + assert forged.attributed_org_ids is None + assert forged.attributed_team_limits is None + + +# --------------------------------------------------------------------------- # +# resolver +# --------------------------------------------------------------------------- # +@pytest.mark.asyncio +async def test_resolver_is_noop_when_both_settings_off(): + """Not even a cache read on the default path.""" + token = UserAPIKeyAuth(api_key="sk-1", team_id="team-a") + with patch("litellm.proxy.auth.auth_checks.get_team_object", new=AsyncMock()) as mock_get_team: + await resolve_membership_attribution( + user_api_key_auth_obj=token, + user_object=LiteLLM_UserTable(user_id="u1", teams=["team-a", "team-b"]), + general_settings={}, + prisma_client=MagicMock(), + user_api_key_cache=MagicMock(), + ) + mock_get_team.assert_not_called() + assert token.attributed_team_ids is None + + +@pytest.mark.asyncio +async def test_resolver_collects_every_membership_with_stamped_team_first(): + token = UserAPIKeyAuth(api_key="sk-1", team_id="team-a") + + def _team(team_id: str, rpm=None, tpm=None, org="org-1"): + obj = MagicMock() + obj.team_id = team_id + obj.rpm_limit = rpm + obj.tpm_limit = tpm + obj.organization_id = org + return obj + + async def _fake_get_team_object(team_id, **kwargs): + return _team(team_id, rpm=10 if team_id == "team-b" else None) + + with patch("litellm.proxy.auth.auth_checks.get_team_object", new=AsyncMock(side_effect=_fake_get_team_object)): + await resolve_membership_attribution( + user_api_key_auth_obj=token, + user_object=LiteLLM_UserTable(user_id="u1", teams=["team-b", "team-a", "team-c"]), + general_settings={"track_spend_across_all_user_teams": True}, + prisma_client=MagicMock(), + user_api_key_cache=MagicMock(), + ) + + # stamped team leads, membership list follows, no duplicate of team-a + assert token.attributed_team_ids == ["team-a", "team-b", "team-c"] + assert token.attributed_team_limits["team-b"] == {"rpm": 10, "tpm": None} + assert token.attributed_org_ids == ["org-1"] + + +@pytest.mark.asyncio +async def test_resolver_fails_open_on_unresolvable_team(): + """A deleted or erroring team is skipped, never raised: attribution must not + turn an already-authorized request into a 500.""" + token = UserAPIKeyAuth(api_key="sk-1", team_id="team-a") + + async def _fake_get_team_object(team_id, **kwargs): + if team_id == "team-broken": + raise Exception("team row is gone") + obj = MagicMock() + obj.rpm_limit = None + obj.tpm_limit = None + obj.organization_id = "org-1" + return obj + + with patch("litellm.proxy.auth.auth_checks.get_team_object", new=AsyncMock(side_effect=_fake_get_team_object)): + await resolve_membership_attribution( + user_api_key_auth_obj=token, + user_object=LiteLLM_UserTable(user_id="u1", teams=["team-broken", "team-a"]), + general_settings={"track_spend_across_all_user_teams": True}, + prisma_client=MagicMock(), + user_api_key_cache=MagicMock(), + ) + + assert token.attributed_team_ids == ["team-a"] + + +# --------------------------------------------------------------------------- # +# spend fan-out +# --------------------------------------------------------------------------- # +def _enqueued(mock_add_update, entity_type): + return [ + c.kwargs["update"]["entity_id"] + for c in mock_add_update.call_args_list + if c.kwargs["update"]["entity_type"] == entity_type + ] + + +@pytest.mark.asyncio +async def test_update_team_db_single_team_when_attribution_off(): + """Regression guard for the default path.""" + writer = DBSpendUpdateWriter() + writer.spend_update_queue.add_update = AsyncMock() + + await writer._update_team_db( + response_cost=1.0, + team_id="team-a", + user_id="u1", + prisma_client=MagicMock(), + ) + + assert _enqueued(writer.spend_update_queue.add_update, Litellm_EntityType.TEAM) == ["team-a"] + assert _enqueued(writer.spend_update_queue.add_update, Litellm_EntityType.TEAM_MEMBER) == [ + "team_id::team-a::user_id::u1" + ] + + +@pytest.mark.asyncio +async def test_update_team_db_charges_every_attributed_team(): + writer = DBSpendUpdateWriter() + writer.spend_update_queue.add_update = AsyncMock() + + await writer._update_team_db( + response_cost=1.0, + team_id="team-a", + user_id="u1", + prisma_client=MagicMock(), + attributed_team_ids=["team-a", "team-b", "team-c"], + ) + + assert _enqueued(writer.spend_update_queue.add_update, Litellm_EntityType.TEAM) == [ + "team-a", + "team-b", + "team-c", + ] + assert _enqueued(writer.spend_update_queue.add_update, Litellm_EntityType.TEAM_MEMBER) == [ + "team_id::team-a::user_id::u1", + "team_id::team-b::user_id::u1", + "team_id::team-c::user_id::u1", + ] + + +@pytest.mark.asyncio +async def test_update_org_db_charges_every_attributed_org(): + writer = DBSpendUpdateWriter() + writer.spend_update_queue.add_update = AsyncMock() + + await writer._update_org_db( + response_cost=1.0, + org_id="org-1", + prisma_client=MagicMock(), + attributed_org_ids=["org-1", "org-2"], + ) + + assert _enqueued(writer.spend_update_queue.add_update, Litellm_EntityType.ORGANIZATION) == [ + "org-1", + "org-2", + ] + + +@pytest.mark.asyncio +async def test_daily_team_rollup_emits_one_row_per_team(): + """LiteLLM_DailyTeamSpend is already unique per team+date+key+model, so N + teams means N distinct transaction keys and no migration.""" + writer = DBSpendUpdateWriter() + writer.daily_team_spend_update_queue.add_update = AsyncMock() + writer._common_add_spend_log_transaction_to_daily_transaction = AsyncMock( + return_value={"date": "2026-08-21", "spend": 1.0, "api_requests": 1, "endpoint": "/chat/completions"} + ) + + payload = { + "team_id": "team-a", + "api_key": "hashed-key", + "model": "gpt-4o", + "custom_llm_provider": "openai", + } + + await writer.add_spend_log_transaction_to_daily_team_transaction( + payload=payload, + prisma_client=MagicMock(), + attributed_team_ids=["team-a", "team-b"], + ) + + keys = [list(c.kwargs["update"].keys())[0] for c in writer.daily_team_spend_update_queue.add_update.call_args_list] + assert len(keys) == 2 + assert keys[0].startswith("team-a_2026-08-21_") + assert keys[1].startswith("team-b_2026-08-21_") + + team_ids = [ + list(c.kwargs["update"].values())[0]["team_id"] + for c in writer.daily_team_spend_update_queue.add_update.call_args_list + ] + assert team_ids == ["team-a", "team-b"] + + +@pytest.mark.asyncio +async def test_daily_team_rollup_single_row_when_attribution_off(): + writer = DBSpendUpdateWriter() + writer.daily_team_spend_update_queue.add_update = AsyncMock() + writer._common_add_spend_log_transaction_to_daily_transaction = AsyncMock( + return_value={"date": "2026-08-21", "spend": 1.0, "api_requests": 1, "endpoint": ""} + ) + + await writer.add_spend_log_transaction_to_daily_team_transaction( + payload={ + "team_id": "team-a", + "api_key": "hashed-key", + "model": "gpt-4o", + "custom_llm_provider": "openai", + }, + prisma_client=MagicMock(), + ) + + assert writer.daily_team_spend_update_queue.add_update.call_count == 1 + + +# --------------------------------------------------------------------------- # +# rate-limit fan-out +# --------------------------------------------------------------------------- # +def _limiter(): + from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + _PROXY_MaxParallelRequestsHandler_v3, + ) + + return _PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=MagicMock()) + + +def test_rate_limit_descriptors_not_added_when_setting_off(): + limiter = _limiter() + token = UserAPIKeyAuth(api_key="sk-1", team_id="team-a") + token.attributed_team_limits = {"team-b": {"rpm": 10, "tpm": None}} + + descriptors: list = [] + with patch("litellm.proxy.proxy_server.general_settings", {}): + limiter._add_attributed_team_rate_limit_descriptors(user_api_key_dict=token, descriptors=descriptors) + assert descriptors == [] + + +def test_rate_limit_descriptors_added_for_other_teams_only(): + limiter = _limiter() + token = UserAPIKeyAuth(api_key="sk-1", team_id="team-a") + token.attributed_team_limits = { + "team-a": {"rpm": 5, "tpm": None}, # stamped team: already emitted elsewhere + "team-b": {"rpm": 10, "tpm": None}, + "team-c": {"rpm": None, "tpm": None}, # no limits: nothing to enforce + "team-d": {"rpm": None, "tpm": 900}, + } + + descriptors: list = [] + with patch( + "litellm.proxy.proxy_server.general_settings", + {"enforce_rate_limits_across_all_user_teams": True}, + ): + limiter._add_attributed_team_rate_limit_descriptors(user_api_key_dict=token, descriptors=descriptors) + + assert [d["value"] for d in descriptors] == ["team-b", "team-d"] + # same namespace as the stamped-team descriptor, so a team shares one bucket + assert {d["key"] for d in descriptors} == {"team"} + assert descriptors[0]["rate_limit"]["requests_per_unit"] == 10 + assert descriptors[1]["rate_limit"]["tokens_per_unit"] == 900 diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index 50c93ed5275..6e0f9f097e0 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -603,6 +603,10 @@ async def test_update_database_and_spend_counters_updates_counters_after_db_upda budget_reservation=budget_reservation, end_user_id="test_end_user_id", tags=["tag-a"], + # None unless track_spend_across_all_user_teams is on, which keeps this + # call identical to the pre-attribution behavior. + attributed_team_ids=None, + attributed_org_ids=None, ) From 71f3ed04dc13d2bc5bc25a0f33015d101d3c8572 Mon Sep 17 00:00:00 2001 From: mzolota Date: Fri, 21 Aug 2026 12:29:37 +0200 Subject: [PATCH 2/6] fix(proxy): correct membership-attribution CI failures - Import attribution_targets in proxy_server. increment_spend_counters referenced it without an import, raising NameError on every spend increment. Caught by proxy-server and proxy-infra shards; my local runs never exercised that function. - Run ruff format over auth_checks.py. - Regenerate the two ConfigGeneralSettings entries in schema.d.ts. - Widen a test stub in test_spend_tracking_utils.py whose fixed signature predates the two new update_database parameters. Also close a real test gap the coverage report exposed: the two new budget gates and the request-metadata stamping had no tests at all. Adds nine, covering an over-budget non-stamped team, an over-budget non-stamped org, the stamped team being skipped so it cannot raise twice, an unloadable team contributing no ceiling, and both settings-off paths. Adds two counter tests pinning that spend fans out to every attributed team and that it does not when the setting is off. --- litellm/proxy/auth/auth_checks.py | 7 +- litellm/proxy/proxy_server.py | 1 + .../proxy/auth/test_membership_attribution.py | 237 ++++++++++++++++++ .../proxy/proxy_server/test_spend_counters.py | 77 ++++++ .../test_spend_tracking_utils.py | 2 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 10 + 6 files changed, 329 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index fda9701450b..c5883901a8e 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -4718,9 +4718,7 @@ async def _attributed_teams_max_budget_check( if valid_token is None or prisma_client is None: return - other_team_ids: Final = [ - team_id for team_id in attributed_team_ids(valid_token) if team_id != valid_token.team_id - ] + other_team_ids: Final = [team_id for team_id in attributed_team_ids(valid_token) if team_id != valid_token.team_id] if not other_team_ids: return @@ -4769,8 +4767,7 @@ async def _attributed_teams_max_budget_check( current_cost=spend, max_budget=team_object.max_budget, message=( - f"Budget has been exceeded! Team={team_id} Current cost: {spend}, " - f"Max budget: {team_object.max_budget}" + f"Budget has been exceeded! Team={team_id} Current cost: {spend}, Max budget: {team_object.max_budget}" ), entity_type=Litellm_EntityType.TEAM.value, entity_id=team_id, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d61fd41de75..c973e3a5796 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -288,6 +288,7 @@ from litellm.proxy.auth.auth_utils import ( ) from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.auth.litellm_license import LicenseCheck +from litellm.proxy.auth.membership_attribution import attribution_targets from litellm.proxy.auth.model_checks import ( expand_wildcard_deployments_for_model_info, get_all_fallbacks, diff --git a/tests/test_litellm/proxy/auth/test_membership_attribution.py b/tests/test_litellm/proxy/auth/test_membership_attribution.py index 417ee7523d0..7dfb5016402 100644 --- a/tests/test_litellm/proxy/auth/test_membership_attribution.py +++ b/tests/test_litellm/proxy/auth/test_membership_attribution.py @@ -19,7 +19,12 @@ import pytest sys.path.insert(0, os.path.abspath("../../../..")) +import litellm from litellm.proxy._types import Litellm_EntityType, LiteLLM_UserTable, UserAPIKeyAuth +from litellm.proxy.auth.auth_checks import ( + _attributed_orgs_max_budget_check, + _attributed_teams_max_budget_check, +) from litellm.proxy.auth.membership_attribution import ( attributed_org_ids, attributed_team_ids, @@ -29,6 +34,7 @@ from litellm.proxy.auth.membership_attribution import ( spend_attribution_enabled, ) from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter +from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup # --------------------------------------------------------------------------- # @@ -335,3 +341,234 @@ def test_rate_limit_descriptors_added_for_other_teams_only(): assert {d["key"] for d in descriptors} == {"team"} assert descriptors[0]["rate_limit"]["requests_per_unit"] == 10 assert descriptors[1]["rate_limit"]["tokens_per_unit"] == 900 + + +# --------------------------------------------------------------------------- # +# budget enforcement across memberships +# --------------------------------------------------------------------------- # +SPEND_ON = {"track_spend_across_all_user_teams": True} + + +def _proxy_logging(): + from litellm.proxy.utils import ProxyLogging + + obj = ProxyLogging(user_api_key_cache=None) + obj.budget_alerts = AsyncMock() + return obj + + +def _team(team_id, max_budget=None, spend=0.0, org="org-1"): + obj = MagicMock() + obj.team_id = team_id + obj.team_alias = f"{team_id}-alias" + obj.max_budget = max_budget + obj.spend = spend + obj.organization_id = org + return obj + + +def _token_with_memberships(): + token = UserAPIKeyAuth(api_key="sk-1", team_id="team-a", org_id="org-1") + token.attributed_team_ids = ["team-a", "team-b"] + token.attributed_org_ids = ["org-1", "org-2"] + return token + + +@pytest.mark.asyncio +async def test_attributed_team_budget_check_noop_when_setting_off(): + """Populated memberships must not gate budgets unless the SPEND setting is + on. Enabling only rate-limit attribution must not silently add budget gates. + """ + token = _token_with_memberships() + with patch("litellm.proxy.auth.auth_checks.get_team_object", new=AsyncMock()) as mock_get_team: + await _attributed_teams_max_budget_check( + valid_token=token, + prisma_client=MagicMock(), + user_api_key_cache=MagicMock(), + proxy_logging_obj=_proxy_logging(), + general_settings={"enforce_rate_limits_across_all_user_teams": True}, + ) + mock_get_team.assert_not_called() + + +@pytest.mark.asyncio +async def test_attributed_team_budget_check_raises_for_other_team(): + """A team the caller merely belongs to can now block them.""" + token = _token_with_memberships() + + async def _get_team(team_id, **kwargs): + return _team(team_id, max_budget=10.0) + + async def _spend(counter_key, fallback_spend, max_budget=None, **kwargs): + return 25.0 if counter_key == "spend:team:team-b" else 0.0 + + with ( + patch("litellm.proxy.auth.auth_checks.get_team_object", new=AsyncMock(side_effect=_get_team)), + patch("litellm.proxy.proxy_server.get_current_spend", _spend), + ): + with pytest.raises(litellm.BudgetExceededError) as exc: + await _attributed_teams_max_budget_check( + valid_token=token, + prisma_client=MagicMock(), + user_api_key_cache=MagicMock(), + proxy_logging_obj=_proxy_logging(), + general_settings=SPEND_ON, + ) + + assert exc.value.entity_id == "team-b" + assert exc.value.entity_type == Litellm_EntityType.TEAM.value + assert exc.value.current_cost == 25.0 + + +@pytest.mark.asyncio +async def test_attributed_team_budget_check_skips_stamped_team(): + """The stamped team is _team_max_budget_check's job. Checking it here too + would raise twice and fire a duplicate budget alert.""" + token = _token_with_memberships() + seen = [] + + async def _get_team(team_id, **kwargs): + seen.append(team_id) + return _team(team_id, max_budget=100.0) + + async def _spend(counter_key, fallback_spend, max_budget=None, **kwargs): + return 0.0 + + with ( + patch("litellm.proxy.auth.auth_checks.get_team_object", new=AsyncMock(side_effect=_get_team)), + patch("litellm.proxy.proxy_server.get_current_spend", _spend), + ): + await _attributed_teams_max_budget_check( + valid_token=token, + prisma_client=MagicMock(), + user_api_key_cache=MagicMock(), + proxy_logging_obj=_proxy_logging(), + general_settings=SPEND_ON, + ) + + assert seen == ["team-b"] + + +@pytest.mark.asyncio +async def test_attributed_team_budget_check_passes_under_budget(): + token = _token_with_memberships() + + async def _get_team(team_id, **kwargs): + return _team(team_id, max_budget=100.0) + + async def _spend(counter_key, fallback_spend, max_budget=None, **kwargs): + return 5.0 + + with ( + patch("litellm.proxy.auth.auth_checks.get_team_object", new=AsyncMock(side_effect=_get_team)), + patch("litellm.proxy.proxy_server.get_current_spend", _spend), + ): + await _attributed_teams_max_budget_check( + valid_token=token, + prisma_client=MagicMock(), + user_api_key_cache=MagicMock(), + proxy_logging_obj=_proxy_logging(), + general_settings=SPEND_ON, + ) + + +@pytest.mark.asyncio +async def test_attributed_team_budget_check_ignores_unloadable_team(): + """A team that cannot be loaded contributes no ceiling rather than a 500.""" + token = _token_with_memberships() + + async def _get_team(team_id, **kwargs): + raise Exception("team row is gone") + + with patch("litellm.proxy.auth.auth_checks.get_team_object", new=AsyncMock(side_effect=_get_team)): + await _attributed_teams_max_budget_check( + valid_token=token, + prisma_client=MagicMock(), + user_api_key_cache=MagicMock(), + proxy_logging_obj=_proxy_logging(), + general_settings=SPEND_ON, + ) + + +@pytest.mark.asyncio +async def test_attributed_org_budget_check_raises_for_other_org(): + """Only reachable when the caller's teams span several organizations.""" + token = _token_with_memberships() + + org = MagicMock() + org.spend = 0.0 + org.litellm_budget_table = MagicMock() + org.litellm_budget_table.max_budget = 50.0 + + async def _spend(counter_key, fallback_spend, max_budget=None, **kwargs): + return 90.0 if counter_key == "spend:org:org-2" else 0.0 + + with ( + patch("litellm.proxy.auth.auth_checks.get_org_object", new=AsyncMock(return_value=org)), + patch("litellm.proxy.proxy_server.get_current_spend", _spend), + ): + with pytest.raises(litellm.BudgetExceededError) as exc: + await _attributed_orgs_max_budget_check( + valid_token=token, + prisma_client=MagicMock(), + user_api_key_cache=MagicMock(), + proxy_logging_obj=_proxy_logging(), + general_settings=SPEND_ON, + ) + + assert exc.value.entity_id == "org-2" + assert exc.value.entity_type == Litellm_EntityType.ORGANIZATION.value + + +@pytest.mark.asyncio +async def test_attributed_org_budget_check_noop_when_setting_off(): + token = _token_with_memberships() + with patch("litellm.proxy.auth.auth_checks.get_org_object", new=AsyncMock()) as mock_get_org: + await _attributed_orgs_max_budget_check( + valid_token=token, + prisma_client=MagicMock(), + user_api_key_cache=MagicMock(), + proxy_logging_obj=_proxy_logging(), + general_settings={}, + ) + mock_get_org.assert_not_called() + + +# --------------------------------------------------------------------------- # +# request-metadata stamping +# --------------------------------------------------------------------------- # +def test_metadata_stamped_only_when_setting_on(): + from litellm.proxy.hooks.proxy_track_cost_callback import _metadata_id_list + + token = _token_with_memberships() + + data_off = {"metadata": {}} + with patch("litellm.proxy.proxy_server.general_settings", {}): + LiteLLMProxyRequestSetup._add_attributed_membership_metadata( + data=data_off, user_api_key_dict=token, _metadata_variable_name="metadata" + ) + assert data_off["metadata"] == {} + + data_on = {"metadata": {}} + with patch("litellm.proxy.proxy_server.general_settings", SPEND_ON): + LiteLLMProxyRequestSetup._add_attributed_membership_metadata( + data=data_on, user_api_key_dict=token, _metadata_variable_name="metadata" + ) + assert data_on["metadata"]["user_api_key_attributed_team_ids"] == ["team-a", "team-b"] + assert data_on["metadata"]["user_api_key_attributed_org_ids"] == ["org-1", "org-2"] + + # and the cost callback reads back exactly what was stamped + assert _metadata_id_list(data_on["metadata"], "user_api_key_attributed_team_ids") == [ + "team-a", + "team-b", + ] + + +def test_metadata_id_list_returns_none_when_absent(): + """None, not [], so a writer can tell "attribution off" from "on, nothing + resolved".""" + from litellm.proxy.hooks.proxy_track_cost_callback import _metadata_id_list + + assert _metadata_id_list({}, "user_api_key_attributed_team_ids") is None + assert _metadata_id_list({"user_api_key_attributed_team_ids": []}, "user_api_key_attributed_team_ids") is None + assert _metadata_id_list({"user_api_key_attributed_team_ids": "nope"}, "user_api_key_attributed_team_ids") is None diff --git a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py index 51980342a1d..4a9bc4ea153 100644 --- a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py +++ b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py @@ -421,6 +421,83 @@ async def test_increment_spend_counters_increments_all_buckets(monkeypatch): } +def _incremented_keys(fake_cache) -> set[str]: + """Counter keys passed to the Redis increment, however they were spelled.""" + keys = set() + for call in fake_cache.redis_cache.async_increment.call_args_list: + key = call.kwargs.get("key") + if key is None and call.args: + key = call.args[0] + if isinstance(key, str): + keys.add(key) + return keys + + +@pytest.mark.asyncio +async def test_increment_spend_counters_fans_out_across_attributed_teams(monkeypatch): + """With track_spend_across_all_user_teams on, every attributed team gets its + own team and team-member counter, so budget gates reading Redis see the + spend against all of them immediately. + + Baseline (single team) is 4 increments: key, team, team_member, user. Three + attributed teams makes it 8: key, 3 x team, 3 x team_member, user. + """ + fake_cache = _make_spend_counter_cache(redis_get_value=None, redis_increment_value=5.0) + fake_user_cache = _make_user_api_key_cache(get_value=None) + monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) + monkeypatch.setattr(ps, "prisma_client", None) + + async def _fake_coalesced(**kwargs): + return None + + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(side_effect=_fake_coalesced)) + + await ps.increment_spend_counters( + token="hashed-tok", + team_id="t1", + user_id="u1", + response_cost=5.0, + attributed_team_ids=["t1", "t2", "t3"], + ) + + incremented_keys = _incremented_keys(fake_cache) + assert {"spend:team:t1", "spend:team:t2", "spend:team:t3"} <= incremented_keys + assert { + "spend:team_member:u1:t1", + "spend:team_member:u1:t2", + "spend:team_member:u1:t3", + } <= incremented_keys + assert fake_cache.redis_cache.async_increment.call_count == 8 + + +@pytest.mark.asyncio +async def test_increment_spend_counters_single_team_when_attribution_off(monkeypatch): + """Regression guard: without the setting, only the stamped team is charged.""" + fake_cache = _make_spend_counter_cache(redis_get_value=None, redis_increment_value=5.0) + fake_user_cache = _make_user_api_key_cache(get_value=None) + monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) + monkeypatch.setattr(ps, "prisma_client", None) + + async def _fake_coalesced(**kwargs): + return None + + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(side_effect=_fake_coalesced)) + + await ps.increment_spend_counters( + token="hashed-tok", + team_id="t1", + user_id="u1", + response_cost=5.0, + ) + + incremented_keys = _incremented_keys(fake_cache) + assert "spend:team:t1" in incremented_keys + assert not any(k.startswith("spend:team:t2") for k in incremented_keys) + assert fake_cache.redis_cache.async_increment.call_count == 4 + + class _ConcurrencyProbe: """Stand-in for redis_cache.async_increment that pins concurrency. diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index f2dd66ee677..2f33b4f2421 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -930,6 +930,8 @@ async def test_api_key_preserved_through_failure_hook_to_database(): start_time, end_time, org_id, + attributed_team_ids=None, + attributed_org_ids=None, ): """Mock update_database and capture the payload it creates""" from litellm.proxy.spend_tracking.spend_tracking_utils import ( diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index fff48e14ecf..5894266b73a 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -24099,6 +24099,11 @@ export interface components { * @default false */ enable_public_model_hub: boolean; + /** + * Enforce Rate Limits Across All User Teams + * @description apply the RPM/TPM limits of EVERY team the calling user belongs to, not only the team stamped on the virtual key. The caller's effective limit becomes the minimum across their memberships, so a busy team can throttle someone who is mostly working for a different team. Separate from track_spend_across_all_user_teams so spend attribution can be adopted without this. Default off. + */ + enforce_rate_limits_across_all_user_teams?: boolean | null; /** * Forward Client Headers To Llm Api * @description If True, forwards client headers (e.g. Authorization) to the LLM API. Required for Claude Code with Max subscription. @@ -24266,6 +24271,11 @@ export interface components { * @description Fine-grained control over which object types to load from the database when store_model_in_db is True. Available types: 'models', 'mcp', 'guardrails', 'vector_stores', 'pass_through_endpoints', 'prompts', 'model_cost_map', 'tools', 'config_overrides'. If not set, all objects are loaded (default behavior). */ supported_db_objects?: components["schemas"]["SupportedDBObjectType"][] | null; + /** + * Track Spend Across All User Teams + * @description attribute each request's spend to EVERY team the calling user belongs to (and every organization reached through those teams), not only the team stamped on the virtual key. Budget gates expand to match, so one over-budget team blocks the caller everywhere. Summing team spend then exceeds real spend by design, because one request is charged to several teams; key, user, and org totals stay single-counted. Default off. + */ + track_spend_across_all_user_teams?: boolean | null; /** * Trusted Proxy Ranges * @description CIDR ranges of trusted reverse proxies allowed to provide identity headers for header-based auth paths such as enable_oauth2_proxy_auth and custom_ui_sso_sign_in_handler. From a70b49bcd14dc2c3836fd03e1e149eb6c3f04098 Mon Sep 17 00:00:00 2001 From: mzolota Date: Fri, 21 Aug 2026 12:52:52 +0200 Subject: [PATCH 3/6] fix(proxy): satisfy the ruff strict-rule budget gate The strict gate ratchets per-rule totals against ruff-strict-budget.json, and this change had pushed two of them up by one each: - ANN401/TID251: membership_attribution.py used typing.Any for the otel span and the team-object tuples. Both now carry concrete types (Span, LiteLLM_TeamTableCachedObj) imported under TYPE_CHECKING, so the module no longer references Any at all. - ANN202: the two new budget-check functions had no return annotation. Also moves a return out of a try block into an else (TRY300). Verified per-file against the base commit across every rule the strict config selects, not just the two the gate reported: zero new violations in all nine changed files. scripts/ruff_strict_gate.py itself needs fcntl and cannot run on Windows, so this was checked by counting each rule in each changed file at the base blob and at HEAD and diffing. --- litellm/proxy/auth/auth_checks.py | 4 ++-- litellm/proxy/auth/membership_attribution.py | 16 ++++++++++------ 2 files changed, 12 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index c5883901a8e..7bf75609b77 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -4688,7 +4688,7 @@ async def _attributed_teams_max_budget_check( user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, general_settings: dict, -): +) -> None: """Enforce the max budget of every OTHER team the caller belongs to. Only active when ``track_spend_across_all_user_teams`` is on -- otherwise @@ -4785,7 +4785,7 @@ async def _attributed_orgs_max_budget_check( user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, general_settings: dict, -): +) -> None: """Enforce the max budget of every OTHER organization the caller belongs to. Skips ``valid_token.org_id`` -- ``_organization_max_budget_check`` owns that diff --git a/litellm/proxy/auth/membership_attribution.py b/litellm/proxy/auth/membership_attribution.py index 124b346fa43..c60c72b29a2 100644 --- a/litellm/proxy/auth/membership_attribution.py +++ b/litellm/proxy/auth/membership_attribution.py @@ -33,13 +33,16 @@ those teams (and the caller's own user row) belong to, never a recursive walk. """ import asyncio -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final from litellm._logging import verbose_proxy_logger from litellm.proxy._types import LiteLLM_UserTable, UserAPIKeyAuth if TYPE_CHECKING: + from opentelemetry.trace import Span + from litellm.caching.caching import DualCache + from litellm.proxy._types import LiteLLM_TeamTableCachedObj from litellm.proxy.utils import PrismaClient, ProxyLogging SPEND_ATTRIBUTION_SETTING: Final = "track_spend_across_all_user_teams" @@ -192,7 +195,7 @@ def _apply_org_only_attribution( *, user_api_key_auth_obj: UserAPIKeyAuth, user_object: LiteLLM_UserTable | None, - team_objects: list[tuple[str, Any]], + team_objects: list[tuple[str, "LiteLLM_TeamTableCachedObj | None"]], ) -> None: """Derive the attributed organizations from what is already loaded. @@ -223,9 +226,9 @@ async def _load_team_objects( team_ids: list[str], prisma_client: "PrismaClient", user_api_key_cache: "DualCache", - parent_otel_span: Any, + parent_otel_span: "Span | None", proxy_logging_obj: "ProxyLogging | None", -) -> list[tuple[str, Any]]: +) -> list[tuple[str, "LiteLLM_TeamTableCachedObj | None"]]: """Resolve every candidate team, preserving input order. Lookups run concurrently: ``get_team_object`` is cache-first, so the steady @@ -234,7 +237,7 @@ async def _load_team_objects( """ from litellm.proxy.auth.auth_checks import get_team_object - async def _safe_get(team_id: str) -> tuple[str, Any]: + async def _safe_get(team_id: str) -> tuple[str, "LiteLLM_TeamTableCachedObj | None"]: try: team_object = await get_team_object( team_id=team_id, @@ -243,7 +246,6 @@ async def _load_team_objects( parent_otel_span=parent_otel_span, proxy_logging_obj=proxy_logging_obj, ) - return team_id, team_object except Exception as e: # noqa: BLE001 # attribution must never fail an authorized request verbose_proxy_logger.debug( "membership attribution: skipping team_id=%s, could not resolve: %s", @@ -251,5 +253,7 @@ async def _load_team_objects( e, ) return team_id, None + else: + return team_id, team_object return list(await asyncio.gather(*(_safe_get(team_id) for team_id in team_ids))) From 9ba2017ea071a4d6c714d703a95eb1f9d082c77f Mon Sep 17 00:00:00 2001 From: mzolota Date: Fri, 21 Aug 2026 13:43:43 +0200 Subject: [PATCH 4/6] refactor(proxy): adopt the immutable-type house style for attribution The type-discipline gate (scripts/type_discipline_gate.py, budgeted in type-discipline-budget.json) ratchets LIT001 (mutable collection in an annotation), LIT002 (mutable construction), LIT010 (assignment without Final) and LIT011 (parameter mutation). This change had added 40/26/5/5 of them. Now zero new violations of any LIT rule, and zero of any rule the strict ruff config selects: - Sequence/Mapping for every read-only parameter that carried list/dict. - tuple[str, ...] returns and tuple construction throughout, including the UserAPIKeyAuth attribution fields and the ids stamped into request metadata. _metadata_id_list accepts list or tuple, since a JSON round trip turns the stamped tuple back into a list. - attributed_team_limits is a MappingProxyType of MappingProxyType, so the precomputed limits cannot be rewritten by a later stage. - The rate-limit helper returns its descriptors instead of appending into a caller-owned list, so no mutable out-parameter. - Final on the locals in the new budget checks; TypeAlias on TeamResolution. - The three genuine in-place stamps (onto the auth object, and into request metadata) carry an inline `# rebind-ok:` with a reason: the surrounding code stamps the same objects the same way, and returning new ones would mean rebuilding every consumer. Verified per-file against the base commit for every LIT rule and every strict ruff rule; both gates need POSIX-only fcntl and cannot run on Windows, so the counts were taken from the base blob and HEAD and diffed. --- litellm/proxy/_types.py | 6 +- litellm/proxy/auth/auth_checks.py | 22 ++- litellm/proxy/auth/membership_attribution.py | 104 +++++----- litellm/proxy/db/db_spend_update_writer.py | 19 +- .../hooks/parallel_request_limiter_v3.py | 47 ++--- .../proxy/hooks/proxy_track_cost_callback.py | 21 +- litellm/proxy/litellm_pre_call_utils.py | 12 +- litellm/proxy/proxy_server.py | 4 +- .../proxy/auth/test_membership_attribution.py | 35 ++-- .../proxy/proxy_server/test_spend_counters.py | 180 +++++------------- 10 files changed, 187 insertions(+), 263 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index a9ae2443100..76b6107f079 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2786,12 +2786,12 @@ class LiteLLM_VerificationTokenView(LiteLLM_VerificationToken): # validated input for the same reason as mcp_source_team_rpm_limits below: # a forged entry would let a caller choose which teams they are charged # against, or name a team with generous limits to escape their real ones. - attributed_team_ids: list[str] | None = Field(default=None, exclude=True) - attributed_org_ids: list[str] | None = Field(default=None, exclude=True) + attributed_team_ids: tuple[str, ...] | None = Field(default=None, exclude=True) + attributed_org_ids: tuple[str, ...] | None = Field(default=None, exclude=True) # team_id -> {"rpm": int | None, "tpm": int | None}. Precomputed during auth # because rate-limit descriptor construction is synchronous and cannot await # a per-team lookup. - attributed_team_limits: dict[str, dict[str, int | None]] | None = Field(default=None, exclude=True) + attributed_team_limits: Mapping[str, Mapping[str, int | None]] | None = Field(default=None, exclude=True) # Project Params project_alias: str | None = None diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 7bf75609b77..65a2a55e995 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -4687,7 +4687,7 @@ async def _attributed_teams_max_budget_check( prisma_client: PrismaClient | None, user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, - general_settings: dict, + general_settings: Mapping[str, object], ) -> None: """Enforce the max budget of every OTHER team the caller belongs to. @@ -4718,7 +4718,9 @@ async def _attributed_teams_max_budget_check( if valid_token is None or prisma_client is None: return - other_team_ids: Final = [team_id for team_id in attributed_team_ids(valid_token) if team_id != valid_token.team_id] + other_team_ids: Final = tuple( + team_id for team_id in attributed_team_ids(valid_token) if team_id != valid_token.team_id + ) if not other_team_ids: return @@ -4726,7 +4728,7 @@ async def _attributed_teams_max_budget_check( async def _check(team_id: str) -> None: try: - team_object = await get_team_object( + team_object: Final = await get_team_object( team_id=team_id, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, @@ -4739,7 +4741,7 @@ async def _attributed_teams_max_budget_check( if team_object is None or team_object.max_budget is None or not math.isfinite(team_object.max_budget): return - spend = await get_current_spend( + spend: Final = await get_current_spend( counter_key=f"spend:team:{team_id}", fallback_spend=team_object.spend or 0.0, max_budget=team_object.max_budget, @@ -4747,7 +4749,7 @@ async def _attributed_teams_max_budget_check( if spend <= team_object.max_budget: return - call_info = CallInfo( + call_info: Final = CallInfo( token=valid_token.token, spend=spend, max_budget=team_object.max_budget, @@ -4784,7 +4786,7 @@ async def _attributed_orgs_max_budget_check( prisma_client: PrismaClient | None, user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, - general_settings: dict, + general_settings: Mapping[str, object], ) -> None: """Enforce the max budget of every OTHER organization the caller belongs to. @@ -4803,7 +4805,7 @@ async def _attributed_orgs_max_budget_check( if valid_token is None or prisma_client is None: return - other_org_ids: Final = [org_id for org_id in attributed_org_ids(valid_token) if org_id != valid_token.org_id] + other_org_ids: Final = tuple(org_id for org_id in attributed_org_ids(valid_token) if org_id != valid_token.org_id) if not other_org_ids: return @@ -4811,7 +4813,7 @@ async def _attributed_orgs_max_budget_check( async def _check(org_id: str) -> None: try: - org_table = await get_org_object( + org_table: Final = await get_org_object( org_id=org_id, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, @@ -4824,11 +4826,11 @@ async def _attributed_orgs_max_budget_check( if org_table is None or org_table.litellm_budget_table is None: return - org_max_budget = org_table.litellm_budget_table.max_budget + org_max_budget: Final = org_table.litellm_budget_table.max_budget if org_max_budget is None or org_max_budget <= 0 or not math.isfinite(org_max_budget): return - org_spend = await get_current_spend( + org_spend: Final = await get_current_spend( counter_key=f"spend:org:{org_id}", fallback_spend=org_table.spend or 0.0, max_budget=org_max_budget, diff --git a/litellm/proxy/auth/membership_attribution.py b/litellm/proxy/auth/membership_attribution.py index c60c72b29a2..9d6c053b990 100644 --- a/litellm/proxy/auth/membership_attribution.py +++ b/litellm/proxy/auth/membership_attribution.py @@ -33,7 +33,9 @@ those teams (and the caller's own user row) belong to, never a recursive walk. """ import asyncio -from typing import TYPE_CHECKING, Final +from collections.abc import Iterable, Mapping, Sequence +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, TypeAlias from litellm._logging import verbose_proxy_logger from litellm.proxy._types import LiteLLM_UserTable, UserAPIKeyAuth @@ -48,83 +50,86 @@ if TYPE_CHECKING: SPEND_ATTRIBUTION_SETTING: Final = "track_spend_across_all_user_teams" RATE_LIMIT_ATTRIBUTION_SETTING: Final = "enforce_rate_limits_across_all_user_teams" +# One resolved team: its id, and the team row if it could be loaded. +TeamResolution: TypeAlias = tuple[str, "LiteLLM_TeamTableCachedObj | None"] -def spend_attribution_enabled(general_settings: dict | None) -> bool: + +def spend_attribution_enabled(general_settings: Mapping[str, object] | None) -> bool: """Whether spend, rollups, and budget gates fan out across memberships.""" if not general_settings: return False return general_settings.get(SPEND_ATTRIBUTION_SETTING) is True -def rate_limit_attribution_enabled(general_settings: dict | None) -> bool: +def rate_limit_attribution_enabled(general_settings: Mapping[str, object] | None) -> bool: """Whether the RPM/TPM limiter fans out across memberships.""" if not general_settings: return False return general_settings.get(RATE_LIMIT_ATTRIBUTION_SETTING) is True -def _attribution_enabled(general_settings: dict | None) -> bool: +def _attribution_enabled(general_settings: Mapping[str, object] | None) -> bool: return spend_attribution_enabled(general_settings) or rate_limit_attribution_enabled(general_settings) -def attributed_team_ids(valid_token: UserAPIKeyAuth | None) -> list[str]: - """Every team this request is attributed to, most-specific first. +def attributed_team_ids(valid_token: UserAPIKeyAuth | None) -> tuple[str, ...]: + """Every team this request is attributed to, stamped team first. Falls back to the single stamped team whenever attribution is off or resolved nothing, so a call site can use this unconditionally and keep identical behavior with the settings disabled. """ if valid_token is None: - return [] + return () resolved: Final = valid_token.attributed_team_ids if resolved: - return list(resolved) - return [valid_token.team_id] if valid_token.team_id else [] + return tuple(resolved) + return (valid_token.team_id,) if valid_token.team_id else () -def attributed_org_ids(valid_token: UserAPIKeyAuth | None) -> list[str]: +def attributed_org_ids(valid_token: UserAPIKeyAuth | None) -> tuple[str, ...]: """Every organization this request is attributed to. Same fallback contract as :func:`attributed_team_ids`. """ if valid_token is None: - return [] + return () resolved: Final = valid_token.attributed_org_ids if resolved: - return list(resolved) - return [valid_token.org_id] if valid_token.org_id else [] + return tuple(resolved) + return (valid_token.org_id,) if valid_token.org_id else () -def attribution_targets(attributed_ids: list[str] | None, stamped_id: str | None) -> list[str]: +def attribution_targets(attributed_ids: Sequence[str] | None, stamped_id: str | None) -> tuple[str, ...]: """The entity ids one request should be charged against. The attributed set when membership attribution resolved one, otherwise the single stamped id -- so with both settings off this returns exactly - ``[stamped_id]`` and every caller keeps its historical behavior. + ``(stamped_id,)`` and every caller keeps its historical behavior. Order-preserving dedupe: the stamped team normally also appears in the caller's membership list, and charging it twice would double-count. """ if attributed_ids: - return list(dict.fromkeys(i for i in attributed_ids if i)) - return [stamped_id] if stamped_id else [] + return _ordered_unique(attributed_ids) + return (stamped_id,) if stamped_id else () -def _ordered_unique(values: list[str | None]) -> list[str]: +def _ordered_unique(values: Iterable[str | None]) -> tuple[str, ...]: """Dedupe while preserving order, dropping empties. Order is preserved so the stamped team stays first. Budget and rate-limit errors report the first offending entity, and a caller reading that error is best served by hearing about the team their key actually names. """ - return list(dict.fromkeys(v for v in values if v)) + return tuple(dict.fromkeys(v for v in values if v)) async def resolve_membership_attribution( *, user_api_key_auth_obj: UserAPIKeyAuth, user_object: LiteLLM_UserTable | None, - general_settings: dict | None, + general_settings: Mapping[str, object] | None, prisma_client: "PrismaClient | None", user_api_key_cache: "DualCache", proxy_logging_obj: "ProxyLogging | None" = None, @@ -153,14 +158,14 @@ async def resolve_membership_attribution( # pointed at a new team, and losing that team would silently under-charge # the one team the operator explicitly named. candidate_team_ids: Final = _ordered_unique( - [user_api_key_auth_obj.team_id, *(list(user_object.teams) if user_object and user_object.teams else [])] + (user_api_key_auth_obj.team_id, *(user_object.teams if user_object and user_object.teams else ())) ) if not candidate_team_ids: - _apply_org_only_attribution( + _apply_org_attribution( user_api_key_auth_obj=user_api_key_auth_obj, user_object=user_object, - team_objects=[], + team_objects=(), ) return @@ -172,30 +177,39 @@ async def resolve_membership_attribution( proxy_logging_obj=proxy_logging_obj, ) - resolved_ids: Final = [team_id for team_id, team_object in team_objects if team_object is not None] + resolved_ids: Final = tuple(team_id for team_id, team_object in team_objects if team_object is not None) if resolved_ids: - user_api_key_auth_obj.attributed_team_ids = resolved_ids - user_api_key_auth_obj.attributed_team_limits = { - team_id: { - "rpm": getattr(team_object, "rpm_limit", None), - "tpm": getattr(team_object, "tpm_limit", None), + team_limits: Final = MappingProxyType( + { + team_id: MappingProxyType( + { + "rpm": getattr(team_object, "rpm_limit", None), + "tpm": getattr(team_object, "tpm_limit", None), + } + ) + for team_id, team_object in team_objects + if team_object is not None } - for team_id, team_object in team_objects - if team_object is not None - } + ) + # The resolved context belongs on the auth object every later stage + # already reads -- the key/team org fallback immediately above this + # call does exactly the same. Returning a new object instead would + # mean rebuilding every consumer of user_api_key_auth. + user_api_key_auth_obj.attributed_team_ids = resolved_ids # rebind-ok: stamping resolved auth context + user_api_key_auth_obj.attributed_team_limits = team_limits # rebind-ok: stamping resolved auth context - _apply_org_only_attribution( + _apply_org_attribution( user_api_key_auth_obj=user_api_key_auth_obj, user_object=user_object, team_objects=team_objects, ) -def _apply_org_only_attribution( +def _apply_org_attribution( *, user_api_key_auth_obj: UserAPIKeyAuth, user_object: LiteLLM_UserTable | None, - team_objects: list[tuple[str, "LiteLLM_TeamTableCachedObj | None"]], + team_objects: Sequence[TeamResolution], ) -> None: """Derive the attributed organizations from what is already loaded. @@ -207,28 +221,28 @@ def _apply_org_only_attribution( several organizations. """ org_ids: Final = _ordered_unique( - [ + ( user_api_key_auth_obj.org_id, user_object.organization_id if user_object else None, - *[ + *( getattr(team_object, "organization_id", None) for _team_id, team_object in team_objects if team_object is not None - ], - ] + ), + ) ) if org_ids: - user_api_key_auth_obj.attributed_org_ids = org_ids + user_api_key_auth_obj.attributed_org_ids = org_ids # rebind-ok: stamping resolved auth context async def _load_team_objects( *, - team_ids: list[str], + team_ids: Sequence[str], prisma_client: "PrismaClient", user_api_key_cache: "DualCache", parent_otel_span: "Span | None", proxy_logging_obj: "ProxyLogging | None", -) -> list[tuple[str, "LiteLLM_TeamTableCachedObj | None"]]: +) -> tuple[TeamResolution, ...]: """Resolve every candidate team, preserving input order. Lookups run concurrently: ``get_team_object`` is cache-first, so the steady @@ -237,9 +251,9 @@ async def _load_team_objects( """ from litellm.proxy.auth.auth_checks import get_team_object - async def _safe_get(team_id: str) -> tuple[str, "LiteLLM_TeamTableCachedObj | None"]: + async def _safe_get(team_id: str) -> TeamResolution: try: - team_object = await get_team_object( + team_object: Final = await get_team_object( team_id=team_id, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, @@ -256,4 +270,4 @@ async def _load_team_objects( else: return team_id, team_object - return list(await asyncio.gather(*(_safe_get(team_id) for team_id in team_ids))) + return tuple(await asyncio.gather(*(_safe_get(team_id) for team_id in team_ids))) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 49520a869a5..0b7fc14ea36 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -12,6 +12,7 @@ import os import random import time import traceback +from collections.abc import Sequence from datetime import datetime, timedelta, timezone from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast, overload @@ -159,8 +160,8 @@ class DBSpendUpdateWriter: # Every team/org this request is attributed to, when # track_spend_across_all_user_teams is on. None keeps the historical # single-team/single-org behavior. - attributed_team_ids: list[str] | None = None, - attributed_org_ids: list[str] | None = None, + attributed_team_ids: Sequence[str] | None = None, + attributed_org_ids: Sequence[str] | None = None, ): from litellm.proxy.proxy_server import ( disable_spend_logs, @@ -431,8 +432,8 @@ class DBSpendUpdateWriter: prisma_client: PrismaClient | None, litellm_proxy_budget_name: str | None, payload: SpendLogsPayload, - attributed_team_ids: list[str] | None = None, - attributed_org_ids: list[str] | None = None, + attributed_team_ids: Sequence[str] | None = None, + attributed_org_ids: Sequence[str] | None = None, ): """ Runs all 11 spend-update helpers sequentially inside a single asyncio task. @@ -660,7 +661,7 @@ class DBSpendUpdateWriter: ) @staticmethod - def _attribution_targets(attributed_ids: list[str] | None, stamped_id: str | None) -> list[str]: + def _attribution_targets(attributed_ids: Sequence[str] | None, stamped_id: str | None) -> tuple[str, ...]: """The entity ids to charge for one request. See ``attribution_targets``.""" from litellm.proxy.auth.membership_attribution import attribution_targets @@ -672,7 +673,7 @@ class DBSpendUpdateWriter: team_id: str | None, user_id: str | None, prisma_client: PrismaClient | None, - attributed_team_ids: list[str] | None = None, + attributed_team_ids: Sequence[str] | None = None, ): """Charge the request to every team it is attributed to. @@ -738,7 +739,7 @@ class DBSpendUpdateWriter: response_cost: float | None, org_id: str | None, prisma_client: PrismaClient | None, - attributed_org_ids: list[str] | None = None, + attributed_org_ids: Sequence[str] | None = None, ): """Charge the request to every organization it is attributed to. @@ -1981,7 +1982,7 @@ class DBSpendUpdateWriter: self, payload: SpendLogsPayload, prisma_client: PrismaClient | None = None, - attributed_team_ids: list[str] | None = None, + attributed_team_ids: Sequence[str] | None = None, ) -> None: """Enqueue one daily rollup row per attributed team. @@ -2018,7 +2019,7 @@ class DBSpendUpdateWriter: payload: SpendLogsPayload, prisma_client: PrismaClient | None = None, org_id: str | None = None, - attributed_org_ids: list[str] | None = None, + attributed_org_ids: Sequence[str] | None = None, ) -> None: """Enqueue one daily rollup row per attributed organization.""" if prisma_client is None: diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index fa00a42decf..1644bf74f52 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -1201,7 +1201,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): key_groups: Final = self._group_keys_by_hash_tag(keys_to_fetch) all_cache_values: Final[list[CacheCounterValue | None]] = [] - async def _read_group(hash_tag: str, group_keys: list[str]) -> CacheCounterValues: + async def _read_group(hash_tag: str, group_keys: Sequence[str]) -> CacheCounterValues: try: return await self.batch_rate_limiter_script( keys=group_keys, @@ -2362,11 +2362,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) ) - def _add_attributed_team_rate_limit_descriptors( + def _attributed_team_rate_limit_descriptors( self, user_api_key_dict: UserAPIKeyAuth, - descriptors: list[RateLimitDescriptor], - ) -> None: + ) -> tuple[RateLimitDescriptor, ...]: """Add a team descriptor for every OTHER team the caller belongs to. Only active when ``enforce_rate_limits_across_all_user_teams`` is on; @@ -2388,30 +2387,27 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): from litellm.proxy.proxy_server import general_settings if not rate_limit_attribution_enabled(general_settings): - return + return () attributed_limits: Final = user_api_key_dict.attributed_team_limits if not attributed_limits: - return + return () - for team_id, team_limits in attributed_limits.items(): - if not team_id or team_id == user_api_key_dict.team_id: - continue - rpm_limit = team_limits.get("rpm") - tpm_limit = team_limits.get("tpm") - if rpm_limit is None and tpm_limit is None: - continue - descriptors.append( - RateLimitDescriptor( - key="team", - value=team_id, - rate_limit={ - "requests_per_unit": rpm_limit, - "tokens_per_unit": tpm_limit, - "window_size": self.window_size, - }, - ) + return tuple( + RateLimitDescriptor( + key="team", + value=team_id, + rate_limit=RateLimitDescriptorRateLimitObject( + requests_per_unit=team_limits.get("rpm"), + tokens_per_unit=team_limits.get("tpm"), + window_size=self.window_size, + ), ) + for team_id, team_limits in attributed_limits.items() + if team_id + and team_id != user_api_key_dict.team_id + and not (team_limits.get("rpm") is None and team_limits.get("tpm") is None) + ) def _add_tag_per_key_rate_limit_descriptor( self, @@ -2764,10 +2760,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) ) - self._add_attributed_team_rate_limit_descriptors( - user_api_key_dict=user_api_key_dict, - descriptors=descriptors, - ) + descriptors.extend(self._attributed_team_rate_limit_descriptors(user_api_key_dict=user_api_key_dict)) # Team Member rate limits if user_api_key_dict.user_id and ( diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index df9d67d2751..cfd49b03f75 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -1,5 +1,6 @@ import asyncio import traceback +from collections.abc import Mapping, Sequence from datetime import datetime from typing import Any, Final, cast @@ -512,26 +513,28 @@ def _get_budget_reservation_from_metadata(metadata: dict) -> dict | None: return getattr(user_api_key_auth_obj, "budget_reservation", None) -def _metadata_id_list(metadata: dict, key: str) -> list[str] | None: +def _metadata_id_list(metadata: Mapping[str, object], key: str) -> tuple[str, ...] | None: """Read a stamped id list out of request metadata, or None if absent. - Returns None (not []) when the key is missing so downstream writers can - tell "attribution off" apart from "attribution on, nothing resolved". + Returns None (not an empty tuple) when the key is missing, so downstream + writers can tell "attribution off" apart from "attribution on, nothing + resolved". """ value: Final = metadata.get(key) - if not isinstance(value, list): + # tuple as stamped, list after any JSON round trip + if not isinstance(value, (list, tuple)): return None - ids: Final = [v for v in value if isinstance(v, str) and v] + ids: Final = tuple(v for v in value if isinstance(v, str) and v) return ids or None -def _attributed_team_ids(user_api_key_dict: UserAPIKeyAuth) -> list[str] | None: +def _attributed_team_ids(user_api_key_dict: UserAPIKeyAuth) -> tuple[str, ...] | None: from litellm.proxy.auth.membership_attribution import attributed_team_ids return attributed_team_ids(user_api_key_dict) if user_api_key_dict.attributed_team_ids else None -def _attributed_org_ids(user_api_key_dict: UserAPIKeyAuth) -> list[str] | None: +def _attributed_org_ids(user_api_key_dict: UserAPIKeyAuth) -> tuple[str, ...] | None: from litellm.proxy.auth.membership_attribution import attributed_org_ids return attributed_org_ids(user_api_key_dict) if user_api_key_dict.attributed_org_ids else None @@ -568,8 +571,8 @@ async def _update_database_and_spend_counters( response_cost: float, budget_reservation: dict | None, request_tags: list[str] | None = None, - attributed_team_ids: list[str] | None = None, - attributed_org_ids: list[str] | None = None, + attributed_team_ids: Sequence[str] | None = None, + attributed_org_ids: Sequence[str] | None = None, ) -> None: try: await proxy_logging_obj.db_spend_update_writer.update_database( diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 33e32a0615b..22026d45275 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -1332,7 +1332,7 @@ class LiteLLMProxyRequestSetup: @staticmethod def _add_attributed_membership_metadata( - data: dict, + data: dict, # mutable-ok: every sibling stamper in this class writes into the caller's request dict user_api_key_dict: UserAPIKeyAuth, _metadata_variable_name: str, ) -> None: @@ -1354,8 +1354,14 @@ class LiteLLMProxyRequestSetup: if user_api_key_dict.attributed_team_ids is None and user_api_key_dict.attributed_org_ids is None: return - data[_metadata_variable_name]["user_api_key_attributed_team_ids"] = attributed_team_ids(user_api_key_dict) - data[_metadata_variable_name]["user_api_key_attributed_org_ids"] = attributed_org_ids(user_api_key_dict) + # Stored as tuples. This dict is serialized into the spend-log payload, + # and json.dumps renders a tuple as an array exactly like a list; the + # reader (_metadata_id_list) accepts either, so a JSON round trip that + # turns these back into lists is also fine. + team_ids: Final = attributed_team_ids(user_api_key_dict) + org_ids: Final = attributed_org_ids(user_api_key_dict) + data[_metadata_variable_name]["user_api_key_attributed_team_ids"] = team_ids # rebind-ok: metadata stamp + data[_metadata_variable_name]["user_api_key_attributed_org_ids"] = org_ids # rebind-ok: metadata stamp @staticmethod def add_management_endpoint_metadata_to_request_metadata( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c973e3a5796..1f81a1bfda9 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2610,8 +2610,8 @@ async def increment_spend_counters( budget_reservation: dict | None = None, end_user_id: str | None = None, tags: list[str] | None = None, - attributed_team_ids: list[str] | None = None, - attributed_org_ids: list[str] | None = None, + attributed_team_ids: Sequence[str] | None = None, + attributed_org_ids: Sequence[str] | None = None, ): """ Atomically increment spend counters for budget enforcement. diff --git a/tests/test_litellm/proxy/auth/test_membership_attribution.py b/tests/test_litellm/proxy/auth/test_membership_attribution.py index 7dfb5016402..0ff057f8984 100644 --- a/tests/test_litellm/proxy/auth/test_membership_attribution.py +++ b/tests/test_litellm/proxy/auth/test_membership_attribution.py @@ -61,20 +61,20 @@ def test_settings_are_independent(): # target resolution # --------------------------------------------------------------------------- # def test_attribution_targets_falls_back_to_stamped_id(): - assert attribution_targets(None, "team-a") == ["team-a"] - assert attribution_targets([], "team-a") == ["team-a"] - assert attribution_targets(None, None) == [] + assert attribution_targets(None, "team-a") == ("team-a",) + assert attribution_targets([], "team-a") == ("team-a",) + assert attribution_targets(None, None) == () def test_attribution_targets_dedupes_and_preserves_order(): """The stamped team is normally also in the membership list; charge it once.""" - assert attribution_targets(["team-a", "team-b", "team-a", ""], "team-a") == ["team-a", "team-b"] + assert attribution_targets(["team-a", "team-b", "team-a", ""], "team-a") == ("team-a", "team-b") def test_read_helpers_fall_back_to_stamped_values(): token = UserAPIKeyAuth(api_key="sk-1", team_id="team-a", org_id="org-1") - assert attributed_team_ids(token) == ["team-a"] - assert attributed_org_ids(token) == ["org-1"] + assert attributed_team_ids(token) == ("team-a",) + assert attributed_org_ids(token) == ("org-1",) def test_attributed_fields_cannot_be_forged_from_input(): @@ -136,9 +136,9 @@ async def test_resolver_collects_every_membership_with_stamped_team_first(): ) # stamped team leads, membership list follows, no duplicate of team-a - assert token.attributed_team_ids == ["team-a", "team-b", "team-c"] + assert token.attributed_team_ids == ("team-a", "team-b", "team-c") assert token.attributed_team_limits["team-b"] == {"rpm": 10, "tpm": None} - assert token.attributed_org_ids == ["org-1"] + assert token.attributed_org_ids == ("org-1",) @pytest.mark.asyncio @@ -165,7 +165,7 @@ async def test_resolver_fails_open_on_unresolvable_team(): user_api_key_cache=MagicMock(), ) - assert token.attributed_team_ids == ["team-a"] + assert token.attributed_team_ids == ("team-a",) # --------------------------------------------------------------------------- # @@ -313,10 +313,9 @@ def test_rate_limit_descriptors_not_added_when_setting_off(): token = UserAPIKeyAuth(api_key="sk-1", team_id="team-a") token.attributed_team_limits = {"team-b": {"rpm": 10, "tpm": None}} - descriptors: list = [] with patch("litellm.proxy.proxy_server.general_settings", {}): - limiter._add_attributed_team_rate_limit_descriptors(user_api_key_dict=token, descriptors=descriptors) - assert descriptors == [] + descriptors = limiter._attributed_team_rate_limit_descriptors(user_api_key_dict=token) + assert descriptors == () def test_rate_limit_descriptors_added_for_other_teams_only(): @@ -329,12 +328,11 @@ def test_rate_limit_descriptors_added_for_other_teams_only(): "team-d": {"rpm": None, "tpm": 900}, } - descriptors: list = [] with patch( "litellm.proxy.proxy_server.general_settings", {"enforce_rate_limits_across_all_user_teams": True}, ): - limiter._add_attributed_team_rate_limit_descriptors(user_api_key_dict=token, descriptors=descriptors) + descriptors = limiter._attributed_team_rate_limit_descriptors(user_api_key_dict=token) assert [d["value"] for d in descriptors] == ["team-b", "team-d"] # same namespace as the stamped-team descriptor, so a team shares one bucket @@ -554,14 +552,11 @@ def test_metadata_stamped_only_when_setting_on(): LiteLLMProxyRequestSetup._add_attributed_membership_metadata( data=data_on, user_api_key_dict=token, _metadata_variable_name="metadata" ) - assert data_on["metadata"]["user_api_key_attributed_team_ids"] == ["team-a", "team-b"] - assert data_on["metadata"]["user_api_key_attributed_org_ids"] == ["org-1", "org-2"] + assert data_on["metadata"]["user_api_key_attributed_team_ids"] == ("team-a", "team-b") + assert data_on["metadata"]["user_api_key_attributed_org_ids"] == ("org-1", "org-2") # and the cost callback reads back exactly what was stamped - assert _metadata_id_list(data_on["metadata"], "user_api_key_attributed_team_ids") == [ - "team-a", - "team-b", - ] + assert _metadata_id_list(data_on["metadata"], "user_api_key_attributed_team_ids") == ("team-a", "team-b") def test_metadata_id_list_returns_none_when_absent(): diff --git a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py index 4a9bc4ea153..1737af305c9 100644 --- a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py +++ b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py @@ -47,9 +47,7 @@ def _make_spend_counter_cache( cache.in_memory_cache.delete_cache = MagicMock() if with_redis: cache.redis_cache = MagicMock() - cache.redis_cache.async_get_cache = AsyncMock( - return_value=redis_get_value, side_effect=redis_get_side_effect - ) + cache.redis_cache.async_get_cache = AsyncMock(return_value=redis_get_value, side_effect=redis_get_side_effect) cache.redis_cache.async_increment = AsyncMock( return_value=redis_increment_value, side_effect=redis_increment_side_effect, @@ -69,9 +67,7 @@ def _make_spend_counter_cache( def _make_user_api_key_cache(get_value=None, get_side_effect=None): cache = MagicMock() - cache.async_get_cache = AsyncMock( - return_value=get_value, side_effect=get_side_effect - ) + cache.async_get_cache = AsyncMock(return_value=get_value, side_effect=get_side_effect) cache.async_set_cache_pipeline = AsyncMock() return cache @@ -108,9 +104,7 @@ async def test_get_current_spend_redis_error_falls_back_to_in_memory(monkeypatch ) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - result = await ps.get_current_spend( - counter_key="spend:key:abc", fallback_spend=99.0 - ) + result = await ps.get_current_spend(counter_key="spend:key:abc", fallback_spend=99.0) assert result == 17.0 @@ -135,9 +129,7 @@ async def test_get_current_spend_floors_stale_low_counter_against_db(monkeypatch # the stale counter is repaired up to the authoritative DB value via a # monotonic set-max so other workers read the corrected total, and a # concurrent increment cannot be clobbered - fake_cache.redis_cache.async_set_max.assert_awaited_once_with( - key="spend:key:abc", value=12.0 - ) + fake_cache.redis_cache.async_set_max.assert_awaited_once_with(key="spend:key:abc", value=12.0) @pytest.mark.asyncio @@ -168,9 +160,7 @@ async def test_get_current_spend_no_floor_without_max_budget(monkeypatch): from_db = AsyncMock(return_value=12.0) monkeypatch.setattr(ps.SpendCounterReseed, "from_db", from_db) - result = await ps.get_current_spend( - counter_key="spend:key:abc", fallback_spend=12.0 - ) + result = await ps.get_current_spend(counter_key="spend:key:abc", fallback_spend=12.0) assert result == 2.0 assert from_db.await_count == 0 @@ -209,12 +199,8 @@ async def test_get_current_spend_floor_caches_db_read(monkeypatch): from_db = AsyncMock(return_value=12.0) monkeypatch.setattr(ps.SpendCounterReseed, "from_db", from_db) - first = await ps.get_current_spend( - counter_key="spend:key:abc", fallback_spend=12.0, max_budget=10.0 - ) - second = await ps.get_current_spend( - counter_key="spend:key:abc", fallback_spend=12.0, max_budget=10.0 - ) + first = await ps.get_current_spend(counter_key="spend:key:abc", fallback_spend=12.0, max_budget=10.0) + second = await ps.get_current_spend(counter_key="spend:key:abc", fallback_spend=12.0, max_budget=10.0) assert first == 12.0 assert second == 12.0 @@ -266,9 +252,7 @@ async def test_get_current_spend_floors_window_against_spend_logs(monkeypatch): assert result == 15.0 assert wfsl.await_count == 1 - fake_cache.redis_cache.async_set_max.assert_awaited_once_with( - key=counter_key, value=15.0 - ) + fake_cache.redis_cache.async_set_max.assert_awaited_once_with(key=counter_key, value=15.0) @pytest.mark.asyncio @@ -278,21 +262,13 @@ async def test_get_current_spend_fail_closed_rejects_when_unverifiable(monkeypat rather than admitted on an unverifiable budget.""" from fastapi import HTTPException - fake_cache = _make_spend_counter_cache( - redis_get_side_effect=RuntimeError("redis down") - ) + fake_cache = _make_spend_counter_cache(redis_get_side_effect=RuntimeError("redis down")) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - monkeypatch.setattr( - ps, "general_settings", {"fail_closed_budget_enforcement": True} - ) - monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None) - ) + monkeypatch.setattr(ps, "general_settings", {"fail_closed_budget_enforcement": True}) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) with pytest.raises(HTTPException) as exc: - await ps.get_current_spend( - counter_key="spend:key:abc", fallback_spend=1.0, max_budget=10.0 - ) + await ps.get_current_spend(counter_key="spend:key:abc", fallback_spend=1.0, max_budget=10.0) assert exc.value.status_code == 503 @@ -300,18 +276,12 @@ async def test_get_current_spend_fail_closed_rejects_when_unverifiable(monkeypat async def test_get_current_spend_fail_closed_off_admits_when_unverifiable(monkeypatch): """Default (flag off): an unverifiable read keeps the existing behavior and admits using the cached fallback — no new rejection.""" - fake_cache = _make_spend_counter_cache( - redis_get_side_effect=RuntimeError("redis down") - ) + fake_cache = _make_spend_counter_cache(redis_get_side_effect=RuntimeError("redis down")) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "general_settings", {}) - monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None) - ) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) - result = await ps.get_current_spend( - counter_key="spend:key:abc", fallback_spend=1.0, max_budget=10.0 - ) + result = await ps.get_current_spend(counter_key="spend:key:abc", fallback_spend=1.0, max_budget=10.0) assert result == 1.0 @@ -321,13 +291,9 @@ async def test_get_current_spend_fail_closed_admits_when_redis_verified(monkeypa authoritative, so an under-budget request is admitted normally.""" fake_cache = _make_spend_counter_cache(redis_get_value=1.0) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - monkeypatch.setattr( - ps, "general_settings", {"fail_closed_budget_enforcement": True} - ) + monkeypatch.setattr(ps, "general_settings", {"fail_closed_budget_enforcement": True}) - result = await ps.get_current_spend( - counter_key="spend:key:abc", fallback_spend=1.0, max_budget=10.0 - ) + result = await ps.get_current_spend(counter_key="spend:key:abc", fallback_spend=1.0, max_budget=10.0) assert result == 1.0 @@ -336,16 +302,10 @@ async def test_get_current_spend_fail_closed_allows_authoritative_fallback(monke """End-user/tag callers pass fallback_authoritative=True (their spend is loaded fresh from the DB in auth), so fail-closed does not reject them even when the counter path is unreadable.""" - fake_cache = _make_spend_counter_cache( - redis_get_side_effect=RuntimeError("redis down") - ) + fake_cache = _make_spend_counter_cache(redis_get_side_effect=RuntimeError("redis down")) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - monkeypatch.setattr( - ps, "general_settings", {"fail_closed_budget_enforcement": True} - ) - monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None) - ) + monkeypatch.setattr(ps, "general_settings", {"fail_closed_budget_enforcement": True}) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) result = await ps.get_current_spend( counter_key="spend:end_user:e1", @@ -363,9 +323,7 @@ async def test_get_current_spend_strict_floors_when_fallback_also_stale(monkeypa re-checks the authoritative DB and enforces against it.""" fake_cache = _make_spend_counter_cache(redis_get_value=0.00001) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - monkeypatch.setattr( - ps, "general_settings", {"fail_closed_budget_enforcement": True} - ) + monkeypatch.setattr(ps, "general_settings", {"fail_closed_budget_enforcement": True}) from_db = AsyncMock(return_value=0.5) monkeypatch.setattr(ps.SpendCounterReseed, "from_db", from_db) @@ -387,9 +345,7 @@ async def test_get_current_spend_strict_floors_when_fallback_also_stale(monkeypa @pytest.mark.asyncio async def test_increment_spend_counters_increments_all_buckets(monkeypatch): - fake_cache = _make_spend_counter_cache( - redis_get_value=None, redis_increment_value=5.0 - ) + fake_cache = _make_spend_counter_cache(redis_get_value=None, redis_increment_value=5.0) fake_user_cache = _make_user_api_key_cache(get_value=None) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) @@ -398,9 +354,7 @@ async def test_increment_spend_counters_increments_all_buckets(monkeypatch): async def _fake_coalesced(**kwargs): return None - monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(side_effect=_fake_coalesced) - ) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(side_effect=_fake_coalesced)) await ps.increment_spend_counters( token="hashed-tok", @@ -540,9 +494,7 @@ async def test_increment_spend_counters_runs_scopes_concurrently(monkeypatch): monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) monkeypatch.setattr(ps, "prisma_client", None) - monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None) - ) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) task = asyncio.create_task( ps.increment_spend_counters( @@ -591,9 +543,7 @@ async def test_increment_spend_counters_skips_reserved_counter_keys(monkeypatch) import litellm.proxy.spend_tracking.budget_reservation as br reserved = {"spend:key:hashed-tok", "spend:org:org1"} - monkeypatch.setattr( - br, "get_reserved_counter_keys", MagicMock(return_value=set(reserved)) - ) + monkeypatch.setattr(br, "get_reserved_counter_keys", MagicMock(return_value=set(reserved))) monkeypatch.setattr(br, "reconcile_budget_reservation", AsyncMock()) recorded: dict[str, float] = {} @@ -608,9 +558,7 @@ async def test_increment_spend_counters_skips_reserved_counter_keys(monkeypatch) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) monkeypatch.setattr(ps, "prisma_client", None) - monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None) - ) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) reservation = {"finalized": False} await ps.increment_spend_counters( @@ -655,9 +603,7 @@ async def test_increment_spend_counters_failing_scope_propagates_after_siblings_ monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) monkeypatch.setattr(ps, "prisma_client", None) - monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None) - ) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) reservation = {"finalized": False} with pytest.raises(RuntimeError, match="redis increment failed"): @@ -713,9 +659,7 @@ async def test_increment_spend_counters_zero_cost_is_noop_finalizes_reservation( @pytest.mark.asyncio async def test_reconcile_budget_reservation_for_counter_update_returns_empty_set_when_none(): - result = await ps._reconcile_budget_reservation_for_counter_update( - budget_reservation=None, response_cost=1.0 - ) + result = await ps._reconcile_budget_reservation_for_counter_update(budget_reservation=None, response_cost=1.0) assert result == set() @@ -758,16 +702,12 @@ async def test_reconcile_budget_reservation_for_counter_update_failure_invalidat async def test_increment_end_user_and_tag_spend_counters_increments_each_unique_tag( monkeypatch, ): - fake_cache = _make_spend_counter_cache( - redis_get_value=None, redis_increment_value=3.0 - ) + fake_cache = _make_spend_counter_cache(redis_get_value=None, redis_increment_value=3.0) fake_user_cache = _make_user_api_key_cache() monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) monkeypatch.setattr(ps, "prisma_client", None) - monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None) - ) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) await ps._increment_end_user_and_tag_spend_counters( end_user_id="eu1", @@ -812,16 +752,12 @@ async def test_increment_end_user_and_tag_spend_counters_no_end_user_no_tags_inv @pytest.mark.asyncio async def test_increment_org_spend_counter_increments_when_org_present(monkeypatch): - fake_cache = _make_spend_counter_cache( - redis_get_value=None, redis_increment_value=10.0 - ) + fake_cache = _make_spend_counter_cache(redis_get_value=None, redis_increment_value=10.0) fake_user_cache = _make_user_api_key_cache() monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) monkeypatch.setattr(ps, "prisma_client", None) - monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None) - ) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) await ps._increment_org_spend_counter( org_id="org-1", @@ -832,9 +768,7 @@ async def test_increment_org_spend_counter_increments_when_org_present(monkeypat observed = { "increment_called": fake_cache.redis_cache.async_increment.called, "increment_calls": fake_cache.redis_cache.async_increment.call_count, - "counter_key_arg": fake_cache.redis_cache.async_increment.call_args.kwargs[ - "key" - ], + "counter_key_arg": fake_cache.redis_cache.async_increment.call_args.kwargs["key"], } assert normalize(observed) == { "increment_called": True, @@ -883,16 +817,12 @@ async def test_init_and_increment_unreserved_spend_counter_skips_reserved_keys( async def test_init_and_increment_unreserved_spend_counter_proceeds_when_not_reserved( monkeypatch, ): - fake_cache = _make_spend_counter_cache( - redis_get_value=None, redis_increment_value=2.0 - ) + fake_cache = _make_spend_counter_cache(redis_get_value=None, redis_increment_value=2.0) fake_user_cache = _make_user_api_key_cache() monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) monkeypatch.setattr(ps, "prisma_client", None) - monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None) - ) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) await ps._init_and_increment_unreserved_spend_counter( counter_key="spend:tag:y", @@ -920,9 +850,7 @@ async def test_init_and_increment_unreserved_spend_counter_proceeds_when_not_res @pytest.mark.asyncio async def test_init_and_increment_spend_counter_warm_cache_skips_reseed(monkeypatch): - fake_cache = _make_spend_counter_cache( - redis_get_value=11.0, redis_increment_value=14.0 - ) + fake_cache = _make_spend_counter_cache(redis_get_value=11.0, redis_increment_value=14.0) fake_user_cache = _make_user_api_key_cache() monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) @@ -957,9 +885,7 @@ async def test_init_and_increment_spend_counter_warm_cache_skips_reseed(monkeypa async def test_init_and_increment_window_spend_counter_increments_when_initialized( monkeypatch, ): - fake_cache = _make_spend_counter_cache( - redis_get_value=0.0, redis_increment_value=5.0 - ) + fake_cache = _make_spend_counter_cache(redis_get_value=0.0, redis_increment_value=5.0) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "prisma_client", None) monkeypatch.setattr( @@ -1044,16 +970,12 @@ async def test_ensure_spend_counter_initialized_warm_skips_reseed_and_source( async def test_ensure_spend_counter_initialized_cold_seeds_from_source_cache( monkeypatch, ): - fake_cache = _make_spend_counter_cache( - redis_get_value=None, redis_increment_value=7.0 - ) + fake_cache = _make_spend_counter_cache(redis_get_value=None, redis_increment_value=7.0) fake_user_cache = _make_user_api_key_cache(get_value={"spend": 7.0}) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) monkeypatch.setattr(ps, "prisma_client", None) - monkeypatch.setattr( - ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None) - ) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) await ps._ensure_spend_counter_initialized( counter_key="spend:user:u", @@ -1093,9 +1015,7 @@ async def test_get_source_cache_base_spend_reads_first_hit_from_list(monkeypatch fake_user_cache.async_get_cache = AsyncMock(side_effect=_get) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) - result = await ps._get_source_cache_base_spend( - source_cache_key=["miss", "hit-obj", "miss2"] - ) + result = await ps._get_source_cache_base_spend(source_cache_key=["miss", "hit-obj", "miss2"]) observed = { "result": result, @@ -1222,9 +1142,7 @@ async def test_increment_spend_counter_cache_redis_path_returns_new_value(monkey fake_cache = _make_spend_counter_cache(redis_increment_value=44.0) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - result = await ps._increment_spend_counter_cache( - counter_key="spend:key:k", increment=4.0 - ) + result = await ps._increment_spend_counter_cache(counter_key="spend:key:k", increment=4.0) observed = { "result": result, @@ -1242,15 +1160,11 @@ async def test_increment_spend_counter_cache_redis_path_returns_new_value(monkey async def test_increment_spend_counter_cache_redis_error_raises_and_invalidates( monkeypatch, ): - fake_cache = _make_spend_counter_cache( - redis_increment_side_effect=RuntimeError("incr fail") - ) + fake_cache = _make_spend_counter_cache(redis_increment_side_effect=RuntimeError("incr fail")) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) with pytest.raises(RuntimeError): - await ps._increment_spend_counter_cache( - counter_key="spend:key:k", increment=1.0 - ) + await ps._increment_spend_counter_cache(counter_key="spend:key:k", increment=1.0) assert fake_cache.in_memory_cache.delete_cache.called is True assert fake_cache.redis_cache.async_delete_cache.called is True @@ -1271,9 +1185,7 @@ async def test_invalidate_spend_counter_deletes_in_memory_and_redis(monkeypatch) observed = { "in_memory_delete_called": fake_cache.in_memory_cache.delete_cache.called, "redis_delete_called": fake_cache.redis_cache.async_delete_cache.called, - "delete_args_key": fake_cache.redis_cache.async_delete_cache.call_args.kwargs[ - "key" - ], + "delete_args_key": fake_cache.redis_cache.async_delete_cache.call_args.kwargs["key"], } assert normalize(observed) == { "in_memory_delete_called": True, @@ -1285,9 +1197,7 @@ async def test_invalidate_spend_counter_deletes_in_memory_and_redis(monkeypatch) @pytest.mark.asyncio async def test_invalidate_spend_counter_swallows_redis_failure_no_raise(monkeypatch): fake_cache = _make_spend_counter_cache() - fake_cache.redis_cache.async_delete_cache = AsyncMock( - side_effect=RuntimeError("redis down") - ) + fake_cache.redis_cache.async_delete_cache = AsyncMock(side_effect=RuntimeError("redis down")) monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) await ps._invalidate_spend_counter(counter_key="spend:key:k") From 294b312d18a4dd321af264497b1ff82e1c85097a Mon Sep 17 00:00:00 2001 From: mzolota Date: Fri, 21 Aug 2026 14:25:39 +0200 Subject: [PATCH 5/6] test(proxy): assert observable outcomes in the attribution budget tests The test-quality gate flagged two tests that patched a dependency and then only asserted the mock went uncalled, which restates the implementation instead of testing behavior (TQ002). Both are now stronger. Each sets up a membership that is wildly over budget and asserts the caller is NOT blocked while the setting is off. That fails if the flag gate is ever removed, which `assert_not_called` never would have caught -- it would have passed against any refactor that reached the team row by another route. Two neighbouring tests asserted nothing at all and could only fail by raising (TQ001). They now assert the same observable value, via a helper that returns the budget error or None, so "the caller got through" is a stated expectation rather than an absence. Also drops the unnecessary sys.path.insert (TQ003); pytest's rootdir already makes litellm importable. Narrows the membership_attribution cache annotation from DualCache to UserApiKeyCache. get_team_object wants the subclass, so the wider type was a real basedpyright error, not a style point. The module now type-checks clean. Verified against the base with every gate the lint job runs: ruff format, ruff check, ruff check on the test tree, the strict-rule gate, the type-discipline gate, circular imports, import safety, and basedpyright on the new module. Measured on a clean worktree at the base commit, this branch adds zero test-quality violations and clears three. --- litellm/proxy/auth/membership_attribution.py | 6 +- .../proxy/auth/test_membership_attribution.py | 114 ++++++++++++------ 2 files changed, 81 insertions(+), 39 deletions(-) diff --git a/litellm/proxy/auth/membership_attribution.py b/litellm/proxy/auth/membership_attribution.py index 9d6c053b990..9cb40490a93 100644 --- a/litellm/proxy/auth/membership_attribution.py +++ b/litellm/proxy/auth/membership_attribution.py @@ -43,8 +43,8 @@ from litellm.proxy._types import LiteLLM_UserTable, UserAPIKeyAuth if TYPE_CHECKING: from opentelemetry.trace import Span - from litellm.caching.caching import DualCache from litellm.proxy._types import LiteLLM_TeamTableCachedObj + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.utils import PrismaClient, ProxyLogging SPEND_ATTRIBUTION_SETTING: Final = "track_spend_across_all_user_teams" @@ -131,7 +131,7 @@ async def resolve_membership_attribution( user_object: LiteLLM_UserTable | None, general_settings: Mapping[str, object] | None, prisma_client: "PrismaClient | None", - user_api_key_cache: "DualCache", + user_api_key_cache: "UserApiKeyCache", proxy_logging_obj: "ProxyLogging | None" = None, ) -> None: """Populate the attributed-membership fields on ``user_api_key_auth_obj``. @@ -239,7 +239,7 @@ async def _load_team_objects( *, team_ids: Sequence[str], prisma_client: "PrismaClient", - user_api_key_cache: "DualCache", + user_api_key_cache: "UserApiKeyCache", parent_otel_span: "Span | None", proxy_logging_obj: "ProxyLogging | None", ) -> tuple[TeamResolution, ...]: diff --git a/tests/test_litellm/proxy/auth/test_membership_attribution.py b/tests/test_litellm/proxy/auth/test_membership_attribution.py index 0ff057f8984..d0dd6e0b409 100644 --- a/tests/test_litellm/proxy/auth/test_membership_attribution.py +++ b/tests/test_litellm/proxy/auth/test_membership_attribution.py @@ -11,14 +11,10 @@ and every one of those tests is a regression guard proving the default path is byte-for-byte what it was before the feature existed. """ -import os -import sys from unittest.mock import AsyncMock, MagicMock, patch import pytest -sys.path.insert(0, os.path.abspath("../../../..")) - import litellm from litellm.proxy._types import Litellm_EntityType, LiteLLM_UserTable, UserAPIKeyAuth from litellm.proxy.auth.auth_checks import ( @@ -372,21 +368,67 @@ def _token_with_memberships(): return token -@pytest.mark.asyncio -async def test_attributed_team_budget_check_noop_when_setting_off(): - """Populated memberships must not gate budgets unless the SPEND setting is - on. Enabling only rate-limit attribution must not silently add budget gates. +async def _team_budget_error(token, general_settings): + """The budget error the team gate raised, or None if the caller got through. + + Asserting on this instead of on which internals ran keeps these tests about + the only thing a caller can observe: blocked, or not blocked. """ - token = _token_with_memberships() - with patch("litellm.proxy.auth.auth_checks.get_team_object", new=AsyncMock()) as mock_get_team: + try: await _attributed_teams_max_budget_check( valid_token=token, prisma_client=MagicMock(), user_api_key_cache=MagicMock(), proxy_logging_obj=_proxy_logging(), - general_settings={"enforce_rate_limits_across_all_user_teams": True}, + general_settings=general_settings, ) - mock_get_team.assert_not_called() + except litellm.BudgetExceededError as e: + return e + return None + + +async def _org_budget_error(token, general_settings): + """The budget error the org gate raised, or None if the caller got through.""" + try: + await _attributed_orgs_max_budget_check( + valid_token=token, + prisma_client=MagicMock(), + user_api_key_cache=MagicMock(), + proxy_logging_obj=_proxy_logging(), + general_settings=general_settings, + ) + except litellm.BudgetExceededError as e: + return e + return None + + +def _over_budget_team(team_id, **kwargs): + """A team whose budget is already blown, for proving a gate is inert.""" + return _team(team_id, max_budget=1.0, spend=999.0) + + +async def _spend_always_over(counter_key, fallback_spend, max_budget=None, **kwargs): + return 999.0 + + +@pytest.mark.asyncio +async def test_attributed_team_budget_check_noop_when_setting_off(): + """Populated memberships must not gate budgets unless the SPEND setting is on. + + Enabling only rate-limit attribution must not silently add budget gates, so + a caller whose other team is wildly over budget still gets through. + """ + token = _token_with_memberships() + with ( + patch( + "litellm.proxy.auth.auth_checks.get_team_object", + new=AsyncMock(side_effect=lambda team_id, **kwargs: _over_budget_team(team_id)), + ), + patch("litellm.proxy.proxy_server.get_current_spend", _spend_always_over), + ): + error = await _team_budget_error(token, {"enforce_rate_limits_across_all_user_teams": True}) + + assert error is None @pytest.mark.asyncio @@ -449,6 +491,7 @@ async def test_attributed_team_budget_check_skips_stamped_team(): @pytest.mark.asyncio async def test_attributed_team_budget_check_passes_under_budget(): + """A membership with room left does not block the caller.""" token = _token_with_memberships() async def _get_team(team_id, **kwargs): @@ -461,31 +504,26 @@ async def test_attributed_team_budget_check_passes_under_budget(): patch("litellm.proxy.auth.auth_checks.get_team_object", new=AsyncMock(side_effect=_get_team)), patch("litellm.proxy.proxy_server.get_current_spend", _spend), ): - await _attributed_teams_max_budget_check( - valid_token=token, - prisma_client=MagicMock(), - user_api_key_cache=MagicMock(), - proxy_logging_obj=_proxy_logging(), - general_settings=SPEND_ON, - ) + error = await _team_budget_error(token, SPEND_ON) + + assert error is None @pytest.mark.asyncio async def test_attributed_team_budget_check_ignores_unloadable_team(): - """A team that cannot be loaded contributes no ceiling rather than a 500.""" + """A team that cannot be loaded contributes no ceiling rather than a 500. + + The caller is neither blocked nor served an internal error. + """ token = _token_with_memberships() async def _get_team(team_id, **kwargs): raise Exception("team row is gone") with patch("litellm.proxy.auth.auth_checks.get_team_object", new=AsyncMock(side_effect=_get_team)): - await _attributed_teams_max_budget_check( - valid_token=token, - prisma_client=MagicMock(), - user_api_key_cache=MagicMock(), - proxy_logging_obj=_proxy_logging(), - general_settings=SPEND_ON, - ) + error = await _team_budget_error(token, SPEND_ON) + + assert error is None @pytest.mark.asyncio @@ -520,16 +558,20 @@ async def test_attributed_org_budget_check_raises_for_other_org(): @pytest.mark.asyncio async def test_attributed_org_budget_check_noop_when_setting_off(): + """An over-budget non-stamped org does not block while the setting is off.""" token = _token_with_memberships() - with patch("litellm.proxy.auth.auth_checks.get_org_object", new=AsyncMock()) as mock_get_org: - await _attributed_orgs_max_budget_check( - valid_token=token, - prisma_client=MagicMock(), - user_api_key_cache=MagicMock(), - proxy_logging_obj=_proxy_logging(), - general_settings={}, - ) - mock_get_org.assert_not_called() + org = MagicMock() + org.spend = 999.0 + org.litellm_budget_table = MagicMock() + org.litellm_budget_table.max_budget = 1.0 + + with ( + patch("litellm.proxy.auth.auth_checks.get_org_object", new=AsyncMock(return_value=org)), + patch("litellm.proxy.proxy_server.get_current_spend", _spend_always_over), + ): + error = await _org_budget_error(token, {}) + + assert error is None # --------------------------------------------------------------------------- # From e948552e9b558484406263f4b36333c34178110c Mon Sep 17 00:00:00 2001 From: mzolota Date: Fri, 21 Aug 2026 15:29:03 +0200 Subject: [PATCH 6/6] fix(proxy): clear the basedpyright errors this change introduced The basedpyright budget gate reported three new errors. All three were real, not annotation noise: - reportOptionalCall (limit 0): the concurrency fix moved the Lua-script call into a nested function, and the enclosing `is None` early return does not narrow the attribute inside a closure. The narrowed value is now bound to a local before the closure is defined. - reportGeneralTypeIssues: `team_member_key: Final` sat inside the new fan-out loop, and a Final cannot be reassigned per iteration. - reportUnnecessaryComparison: `team_object is None` after get_team_object, which returns a non-optional team and raises when the row is missing. Two more were found while verifying and fixed before they could breach: - reportArgumentType in increment_spend_counters: the `user_id is not None` guard sat in the comprehension's iterable, where it does not narrow the element expression. It is now an `if` clause of the comprehension. - reportArgumentType in the limiter: the slot-key parameter had been widened to Sequence[str], but both callees declare `keys: list[str]`. Reverted to list[str] with an inline `# mutable-ok:` naming that constraint. Verified with every gate the lint job runs, against the base commit: ruff format, ruff check, ruff check on the test tree, the strict-rule gate, the type-discipline gate, circular imports, import safety, and a hunk-scoped basedpyright pass showing zero errors from the breached rules inside the lines this change touches. --- litellm/proxy/auth/auth_checks.py | 2 +- litellm/proxy/db/db_spend_update_writer.py | 4 +++- litellm/proxy/hooks/parallel_request_limiter_v3.py | 13 ++++++++++--- litellm/proxy/proxy_server.py | 5 +---- 4 files changed, 15 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 65a2a55e995..5b835fc5264 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -4738,7 +4738,7 @@ async def _attributed_teams_max_budget_check( except Exception: # noqa: BLE001 # a team that cannot be loaded contributes no ceiling return - if team_object is None or team_object.max_budget is None or not math.isfinite(team_object.max_budget): + if team_object.max_budget is None or not math.isfinite(team_object.max_budget): return spend: Final = await get_current_spend( diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 0b7fc14ea36..1dc76b96ab5 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -706,7 +706,9 @@ class DBSpendUpdateWriter: # Track spend of the team member within this team if user_id is not None: # key is "team_id::::user_id::" - team_member_key: Final = f"team_id::{target_team_id}::user_id::{user_id}" + team_member_key = ( + f"team_id::{target_team_id}::user_id::{user_id}" # rebind-ok: one per attributed team + ) await self.spend_update_queue.add_update( update=SpendUpdateQueueItem( entity_type=Litellm_EntityType.TEAM_MEMBER, diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 1644bf74f52..2522a7c72ab 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -1198,17 +1198,24 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if self.batch_rate_limiter_script is None: return [] + # Bound here, not read off self inside the closure: the None check above + # narrows the attribute, but that narrowing does not reach into a nested + # function, so reading it there would be an optional call. + run_script: Final = self.batch_rate_limiter_script key_groups: Final = self._group_keys_by_hash_tag(keys_to_fetch) all_cache_values: Final[list[CacheCounterValue | None]] = [] - async def _read_group(hash_tag: str, group_keys: Sequence[str]) -> CacheCounterValues: + # Both callees below (the Lua script and the in-memory fallback) declare + # `keys: list[str]`, and _group_keys_by_hash_tag hands back exactly that, + # so a read-only view here would not be assignable. + async def _read_group(tag: str, group_keys: list[str]) -> CacheCounterValues: # mutable-ok: callee needs list try: - return await self.batch_rate_limiter_script( + return await run_script( keys=group_keys, args=[now_int, self.window_size], # Use integer timestamp ) except Exception as e: - verbose_proxy_logger.warning("Redis Lua script failed for hash tag %s: %s", hash_tag, e) + verbose_proxy_logger.warning("Redis Lua script failed for hash tag %s: %s", tag, e) # Fallback to in-memory cache for this group return await self.in_memory_cache_sliding_window( keys=group_keys, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 1f81a1bfda9..bb5064299c2 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2740,10 +2740,7 @@ async def increment_spend_counters( for coro in ( _key_scope(token) if token is not None else None, *(_team_scope(scope_team_id) for scope_team_id in target_team_ids), - *( - _team_member_scope(user_id, scope_team_id) - for scope_team_id in (target_team_ids if user_id is not None else ()) - ), + *(_team_member_scope(user_id, scope_team_id) for scope_team_id in target_team_ids if user_id is not None), _user_scope(user_id) if user_id is not None else None, _increment_end_user_and_tag_spend_counters( end_user_id=end_user_id,