From 1981160775fe34ca86cd283794cdf6b89f5019a1 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 29 Aug 2026 10:16:03 -0700 Subject: [PATCH 01/16] feat(budgets): add model access group budget table and shared types Adds the durable row that a model access group budget hangs off. Model access groups live only as free-text strings inside model_info.access_groups, so unlike tags there is no existing row to carry a budget_id. Foundation only: schema, migration, repository, entity type, spend transaction bucket, auth carrier field and registry cache keys. Nothing reads or writes these yet. --- .../migration.sql | 20 +++++++++++++++++++ litellm/constants.py | 1 + litellm/proxy/_types.py | 3 +++ .../proxy/common_utils/user_api_key_cache.py | 15 ++++++++++++++ litellm/repositories/table_repositories.py | 6 ++++++ schema.prisma | 12 +++++++++++ 6 files changed, 57 insertions(+) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260829000000_add_model_access_group_budget_table/migration.sql diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260829000000_add_model_access_group_budget_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260829000000_add_model_access_group_budget_table/migration.sql new file mode 100644 index 00000000000..62398da7f04 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260829000000_add_model_access_group_budget_table/migration.sql @@ -0,0 +1,20 @@ +-- CreateTable +CREATE TABLE IF NOT EXISTS "LiteLLM_ModelAccessGroupBudgetTable" ( + "access_group_name" TEXT NOT NULL, + "spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0, + "budget_id" TEXT, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "created_by" TEXT, + "updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updated_by" TEXT, + + CONSTRAINT "LiteLLM_ModelAccessGroupBudgetTable_pkey" PRIMARY KEY ("access_group_name") +); + +-- AddForeignKey +DO $$ +BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_ModelAccessGroupBudgetTable_budget_id_fkey') THEN + ALTER TABLE "LiteLLM_ModelAccessGroupBudgetTable" ADD CONSTRAINT "LiteLLM_ModelAccessGroupBudgetTable_budget_id_fkey" FOREIGN KEY ("budget_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE SET NULL ON UPDATE CASCADE; + END IF; +END $$; diff --git a/litellm/constants.py b/litellm/constants.py index fc88086805f..24340d7bd55 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1681,6 +1681,7 @@ DEFAULT_MCP_NAMESPACE_CSV_MAX_TOKENS: Final = 16 # Ceilings on the cached auth registries; larger tables fall back to per-row lookups # instead of holding an unbounded id set in every worker. TAG_REGISTRY_MAX_SIZE: Final = 5000 +ACCESS_GROUP_REGISTRY_MAX_SIZE: Final = 5000 END_USER_RESTRICTED_REGISTRY_MAX_SIZE: Final = 5000 # How long a failed registry load is remembered as "unusable", so a degraded Postgres # is not re-scanned on every request on top of the per-id lookups it falls back to. diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index f40b0632398..4d86d258788 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -241,6 +241,7 @@ class Litellm_EntityType(enum.Enum): PROJECT = "project" TAG = "tag" AGENT = "agent" + ACCESS_GROUP = "access_group" # global proxy level entity PROXY = "proxy" @@ -2886,6 +2887,7 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob ), ) budget_reservation: dict[str, Any] | None = Field(default=None, exclude=True) + matched_access_groups: list[str] | None = Field(default=None, exclude=True) budget_throttle_pct: float | None = Field(default=None, exclude=True) user: Any | None = None # Expanded user object when expand=user is used created_by_user: Any | None = None # Expanded created_by user when expand=user is used @@ -4918,6 +4920,7 @@ class DBSpendUpdateTransactions(TypedDict): org_list_transactions: dict[str, float] | None tag_list_transactions: dict[str, float] | None agent_list_transactions: dict[str, float] | None + access_group_list_transactions: dict[str, float] | None class SpendUpdateQueueItem(TypedDict, total=False): diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index b8df0105b7b..96152240181 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -185,6 +185,21 @@ def tag_registry_cache_key() -> str: return "tag_registry" +#: Cached under ``access_group_registry_cache_key`` when the table exceeds +#: ``ACCESS_GROUP_REGISTRY_MAX_SIZE``: registry unusable, fall back to the per-group lookup. +ACCESS_GROUP_REGISTRY_OVERFLOW_SENTINEL: Final = "__access_group_registry_overflow__" + + +def access_group_cache_key(access_group_name: str) -> str: + """Cache key one model access group budget row is stored under; shared so auth, spend tracking and the management endpoints cannot drift.""" + return f"access_group:{access_group_name}" + + +def access_group_registry_cache_key() -> str: + """Cache key for the set of model access group names that have a budget row.""" + return "access_group_registry" + + #: Cached under ``end_user_restricted_registry_cache_key`` when the restricted set exceeds #: ``END_USER_RESTRICTED_REGISTRY_MAX_SIZE``: registry unusable, fall back to the per-id fetch. END_USER_RESTRICTED_REGISTRY_OVERFLOW_SENTINEL: Final = "__end_user_restricted_registry_overflow__" diff --git a/litellm/repositories/table_repositories.py b/litellm/repositories/table_repositories.py index e02f652caf6..de13e5aeb6c 100644 --- a/litellm/repositories/table_repositories.py +++ b/litellm/repositories/table_repositories.py @@ -100,6 +100,12 @@ class TagRepository(PrismaTableRepository["prisma_models.LiteLLM_TagTable"]): table_name = "litellm_tagtable" +class ModelAccessGroupBudgetRepository( + PrismaTableRepository["prisma_models.LiteLLM_ModelAccessGroupBudgetTable"] +): + table_name = "litellm_modelaccessgroupbudgettable" + + class InvitationLinkRepository(PrismaTableRepository["prisma_models.LiteLLM_InvitationLink"]): table_name = "litellm_invitationlink" diff --git a/schema.prisma b/schema.prisma index 2bb850139a2..8ddbee63973 100644 --- a/schema.prisma +++ b/schema.prisma @@ -29,6 +29,7 @@ model LiteLLM_BudgetTable { keys LiteLLM_VerificationToken[] // multiple keys can have the same budget end_users LiteLLM_EndUserTable[] // multiple end-users can have the same budget tags LiteLLM_TagTable[] // multiple tags can have the same budget + model_access_groups LiteLLM_ModelAccessGroupBudgetTable[] // multiple model access groups can have the same budget team_membership LiteLLM_TeamMembership[] // budgets of Users within a Team organization_membership LiteLLM_OrganizationMembership[] // budgets of Users within a Organization } @@ -585,6 +586,17 @@ model LiteLLM_EndUserTable { blocked Boolean @default(false) } +model LiteLLM_ModelAccessGroupBudgetTable { + access_group_name String @id + spend Float @default(0.0) + budget_id String? + litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) + created_at DateTime @default(now()) @map("created_at") + created_by String? + updated_at DateTime @default(now()) @updatedAt + updated_by String? +} + // Track tags with budgets and spend model LiteLLM_TagTable { tag_name String @id From d2440639d5fb3f691bdc18f68079997eebcec300 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 29 Aug 2026 12:13:42 -0700 Subject: [PATCH 02/16] feat(budgets): enforce shared budgets on model access groups A model access group could gate which models a caller reaches but never how much that group of callers could spend in total. Capping a shared pool meant setting a per-entity budget on every key by hand, which caps each key separately and still leaves no way to read what the group cost. Spend is attributed to a group only when the group's name appears on an allowlist the caller was granted (key, team, team-member scope, project or org) and that group serves the requested model. Asking for a model that merely belongs to a group attributes nothing, because nothing about the caller named the group. Levels are unioned rather than ranked, so a team granted "*" whose member is scoped to one group still counts as gated by that group. Enforcement runs on both paths tags already use: a reservation counter on the pre-call path and a read-time max_budget check inside the existing concurrent budget gather, so the ceiling still holds under disable_budget_reservation. Adds LiteLLM_ModelAccessGroupBudgetTable, which is the only place a group is ever a row: the groups themselves stay free-text strings in model_info.access_groups, so a row exists only once someone gives that group a budget. GET, PUT and DELETE /access_group/{name}/budget manage it, and /access_group/{name}/info now carries the spend and budget alongside the models. --- basedpyright-code-budget.json | 2 +- litellm/constants.py | 3 +- .../internal_call_metadata.py | 7 + litellm/litellm_core_utils/litellm_logging.py | 44 +- litellm/proxy/_types.py | 6 +- litellm/proxy/auth/auth_checks.py | 401 +++++++++++ .../proxy/common_utils/reset_budget_job.py | 39 +- .../proxy/common_utils/user_api_key_cache.py | 24 +- litellm/proxy/db/db_spend_update_writer.py | 133 +++- .../redis_update_buffer.py | 9 + .../spend_update_queue.py | 4 + litellm/proxy/litellm_pre_call_utils.py | 5 + ...model_access_group_management_endpoints.py | 387 ++++++++++- .../pass_through_endpoints.py | 2 + .../spend_tracking/budget_reservation.py | 64 +- .../spend_tracking/spend_tracking_utils.py | 24 +- litellm/repositories/prisma_protocols.py | 3 + litellm/repositories/unit_of_work.py | 2 + .../model_management_endpoints.py | 35 +- .../types/proxy/model_access_group_budget.py | 19 + litellm/types/utils.py | 3 +- schema.prisma | 3 + .../test_litellm_logging.py | 54 ++ .../auth/test_model_access_group_budgets.py | 504 ++++++++++++++ .../common_utils/test_reset_budget_job.py | 110 ++- .../proxy/db/test_model_access_group_spend.py | 508 ++++++++++++++ .../test_access_group_management.py | 655 ++++++++++++++++++ .../proxy/test_budget_reservation.py | 111 ++- .../proxy/test_litellm_pre_call_utils.py | 36 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 228 +++++- 30 files changed, 3376 insertions(+), 49 deletions(-) create mode 100644 litellm/types/proxy/model_access_group_budget.py create mode 100644 tests/test_litellm/proxy/auth/test_model_access_group_budgets.py create mode 100644 tests/test_litellm/proxy/db/test_model_access_group_spend.py diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index ef88ae574fb..989ab76dc29 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -3,7 +3,7 @@ "limit": 17270 }, "reportArgumentType": { - "limit": 2539 + "limit": 2538 }, "reportAssignmentType": { "limit": 319 diff --git a/litellm/constants.py b/litellm/constants.py index 24340d7bd55..2139b12e024 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1681,7 +1681,8 @@ DEFAULT_MCP_NAMESPACE_CSV_MAX_TOKENS: Final = 16 # Ceilings on the cached auth registries; larger tables fall back to per-row lookups # instead of holding an unbounded id set in every worker. TAG_REGISTRY_MAX_SIZE: Final = 5000 -ACCESS_GROUP_REGISTRY_MAX_SIZE: Final = 5000 +DEFAULT_MODEL_ACCESS_GROUP_CACHE_TTL: Final = int(os.getenv("DEFAULT_MODEL_ACCESS_GROUP_CACHE_TTL", 600)) +MODEL_ACCESS_GROUP_REGISTRY_MAX_SIZE: Final = 5000 END_USER_RESTRICTED_REGISTRY_MAX_SIZE: Final = 5000 # How long a failed registry load is remembered as "unusable", so a degraded Postgres # is not re-scanned on every request on top of the per-id lookups it falls back to. diff --git a/litellm/litellm_core_utils/internal_call_metadata.py b/litellm/litellm_core_utils/internal_call_metadata.py index 34d5797a6d8..4d043701f40 100644 --- a/litellm/litellm_core_utils/internal_call_metadata.py +++ b/litellm/litellm_core_utils/internal_call_metadata.py @@ -25,6 +25,13 @@ from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, Inter BUDGET_RESERVATION_METADATA_KEYS: Final = frozenset({"user_api_key_budget_reservation"}) +MODEL_ACCESS_GROUP_METADATA_KEY: Final = "user_api_key_matched_model_access_groups" +"""Where auth records the model access groups that authorized the request, for the spend writer. + +The ``user_api_key`` prefix is load-bearing, not cosmetic: when a request carries both +``metadata`` and ``litellm_metadata``, ``get_litellm_metadata_from_kwargs`` returns the latter and +copies a key across only when ``user_api_key`` appears in its name.""" + _USER_API_KEY_AUTH_KEY: Final = "user_api_key_auth" FORWARDABLE_IDENTITY_METADATA_KEYS: Final = frozenset( diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index a8672b5c112..262c54ec535 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -64,7 +64,10 @@ from litellm.integrations.mlflow import MlflowLogger from litellm.integrations.sqs import SQSLogger from litellm.litellm_core_utils.core_helpers import is_expected_client_error, reconstruct_model_name from litellm.litellm_core_utils.get_litellm_params import get_litellm_params -from litellm.litellm_core_utils.internal_call_metadata import is_unbilled_non_inference_call +from litellm.litellm_core_utils.internal_call_metadata import ( + MODEL_ACCESS_GROUP_METADATA_KEY, + is_unbilled_non_inference_call, +) from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import ( cost_breakdown_with_guardrail, guardrail_information_cost, @@ -5049,6 +5052,42 @@ def is_valid_sha256_hash(value: str) -> bool: return bool(re.fullmatch(r"[a-fA-F0-9]{64}", value)) +def coerce_model_access_groups(value: object) -> tuple[str, ...]: + """Model access group names out of untrusted request metadata, deduped and order preserving.""" + if not isinstance(value, (list, tuple)): + return () + return tuple(dict.fromkeys(group for group in value if isinstance(group, str) and group)) + + +def _model_access_groups_on_auth_object(user_api_key_auth: object) -> object: + if isinstance(user_api_key_auth, Mapping): + return user_api_key_auth.get("matched_model_access_groups") + return getattr(user_api_key_auth, "matched_model_access_groups", None) + + +def _model_access_groups_from_metadata(metadata: Mapping[str, object]) -> tuple[str, ...]: + stamped: Final = coerce_model_access_groups(metadata.get(MODEL_ACCESS_GROUP_METADATA_KEY)) + if stamped: + return stamped + return coerce_model_access_groups(_model_access_groups_on_auth_object(metadata.get("user_api_key_auth"))) + + +def request_model_access_groups_from_litellm_params(litellm_params: Mapping[str, object]) -> tuple[str, ...]: + """Access groups the auth layer stamped onto this request, from whichever metadata field carries them. + + Detached internal sub-calls only inherit the identity keys, so the auth object is the + fallback there, exactly as _get_budget_reservation_from_metadata does for reservations. + """ + for metadata_variable_name in ("metadata", "litellm_metadata"): + metadata = litellm_params.get(metadata_variable_name) + if not isinstance(metadata, Mapping): + continue + model_access_groups = _model_access_groups_from_metadata(metadata) + if model_access_groups: + return model_access_groups + return () + + class StandardLoggingPayloadSetup: @staticmethod def cleanup_timestamps( @@ -5896,6 +5935,7 @@ def get_standard_logging_object_payload( request_tags: Final = StandardLoggingPayloadSetup._get_request_tags( litellm_params=litellm_params, proxy_server_request=proxy_server_request ) + request_model_access_groups: Final = request_model_access_groups_from_litellm_params(litellm_params) # cleanup timestamps ( @@ -6058,6 +6098,7 @@ def get_standard_logging_object_payload( prompt_tokens=usage_dict.get("prompt_tokens", 0), completion_tokens=usage_dict.get("completion_tokens", 0), request_tags=request_tags, + request_model_access_groups=request_model_access_groups, end_user=end_user_id, api_base=StandardLoggingPayloadSetup.strip_trailing_slash(litellm_params.get("api_base", "")) or "", model_group=_model_group, @@ -6269,6 +6310,7 @@ def create_dummy_standard_logging_payload() -> StandardLoggingPayload: cache_key=None, saved_cache_cost=saved_cache_cost, request_tags=[], + request_model_access_groups=(), end_user=None, requester_ip_address="127.0.0.1", messages=messages, diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 4d86d258788..95e634e8b2d 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -241,7 +241,7 @@ class Litellm_EntityType(enum.Enum): PROJECT = "project" TAG = "tag" AGENT = "agent" - ACCESS_GROUP = "access_group" + MODEL_ACCESS_GROUP = "model_access_group" # global proxy level entity PROXY = "proxy" @@ -2887,7 +2887,7 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob ), ) budget_reservation: dict[str, Any] | None = Field(default=None, exclude=True) - matched_access_groups: list[str] | None = Field(default=None, exclude=True) + matched_model_access_groups: list[str] | None = Field(default=None, exclude=True) budget_throttle_pct: float | None = Field(default=None, exclude=True) user: Any | None = None # Expanded user object when expand=user is used created_by_user: Any | None = None # Expanded created_by user when expand=user is used @@ -4920,7 +4920,7 @@ class DBSpendUpdateTransactions(TypedDict): org_list_transactions: dict[str, float] | None tag_list_transactions: dict[str, float] | None agent_list_transactions: dict[str, float] | None - access_group_list_transactions: dict[str, float] | None + model_access_group_list_transactions: ReadOnly[dict[str, float] | None] class SpendUpdateQueueItem(TypedDict, total=False): diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index c1f9407cdad..cd4d03e08e2 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -30,8 +30,10 @@ from litellm.constants import ( DEFAULT_ACCESS_GROUP_CACHE_TTL, DEFAULT_IN_MEMORY_TTL, DEFAULT_MAX_RECURSE_DEPTH, + DEFAULT_MODEL_ACCESS_GROUP_CACHE_TTL, EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE, END_USER_RESTRICTED_REGISTRY_MAX_SIZE, + MODEL_ACCESS_GROUP_REGISTRY_MAX_SIZE, REGISTRY_ERROR_NEGATIVE_CACHE_TTL, TAG_REGISTRY_MAX_SIZE, ) @@ -78,11 +80,15 @@ from litellm.proxy.common_utils.http_parsing_utils import ( from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time from litellm.proxy.common_utils.user_api_key_cache import ( END_USER_RESTRICTED_REGISTRY_OVERFLOW_SENTINEL, + MODEL_ACCESS_GROUP_REGISTRY_OVERFLOW_SENTINEL, TAG_REGISTRY_OVERFLOW_SENTINEL, UserApiKeyCache, end_user_cache_key, end_user_restricted_registry_cache_key, get_management_object_ttl, + model_access_group_cache_key, + model_access_group_registry_cache_key, + model_access_group_spend_counter_key, object_permission_cache_key, tag_cache_key, tag_registry_cache_key, @@ -107,12 +113,14 @@ from litellm.repositories.table_repositories import ( EndUserRepository, JWTKeyMappingRepository, ManagedVectorStoresRepository, + ModelAccessGroupBudgetRepository, TagRepository, TeamMembershipRepository, ) from litellm.repositories.team_repository import TeamRepository from litellm.repositories.user_repository import UserRepository from litellm.router import Router +from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget from litellm.utils import get_utc_datetime from .auth_checks_organization import ( @@ -251,6 +259,43 @@ def _end_user_table(repo: _PrismaTableHolder[_PrismaEndUserRow]) -> _PrismaAuthT return repo.table +class _PrismaMaxBudgetRow(Protocol): + @property + def max_budget(self) -> float | None: ... + + +class _PrismaModelAccessGroupBudgetRow(Protocol): + access_group_name: str + + @property + def spend(self) -> float | None: ... + + @property + def litellm_budget_table(self) -> _PrismaMaxBudgetRow | None: ... + + +def _model_access_group_budget_table( + repo: _PrismaTableHolder[_PrismaModelAccessGroupBudgetRow], +) -> _PrismaAuthTable[_PrismaModelAccessGroupBudgetRow]: + return repo.table + + +class _MemberModelScope(Protocol): + @property + def allowed_models(self) -> Sequence[str] | None: ... + + +class _TeamMembershipModelScope(Protocol): + @property + def litellm_budget_table(self) -> _MemberModelScope | None: ... + + +def _member_allowed_models(membership: _TeamMembershipModelScope) -> Sequence[str]: + """The member's own model scope, read through a narrowed view of the membership row.""" + budget_table: Final = membership.litellm_budget_table + return () if budget_table is None else (budget_table.allowed_models or ()) + + class _RawCacheRead(Protocol): async def async_get_cache(self, *, key: str) -> object: ... @@ -807,6 +852,7 @@ async def common_checks( 1.1. If project is blocked 2. If team can call model 2.2 If project can call model + 2.3 Which model access groups authorized this request 3. If team is in budget 3.0.2. If project is in budget 3.0.3. If project is over soft budget (alert only) @@ -925,6 +971,18 @@ async def common_checks( proxy_logging_obj=proxy_logging_obj, ) + # 2.3 Which model access groups authorized this request + matched_model_access_groups: Final = await stamp_matched_model_access_groups( + model=_model, + valid_token=valid_token, + team_object=team_object, + project_object=project_object, + llm_router=llm_router, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + # Run before apply_key_tags_pre_auth injects key metadata.tags into request_body. _reject_clientside_metadata_tags_check(general_settings, request_body, route) @@ -1004,6 +1062,13 @@ async def common_checks( proxy_logging_obj=proxy_logging_obj, valid_token=valid_token, ), + _model_access_group_max_budget_check( + matched_model_access_groups=matched_model_access_groups, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + if matched_model_access_groups + else None, _user_max_budget_check(), _check_team_member_budget( team_object=team_object, @@ -1444,6 +1509,7 @@ _REGISTRY_NOT_CACHED: Final = _RegistryNotCached() #: One lock per registry; module-level because the stampede to collapse is worker-wide. _TAG_REGISTRY_LOAD_LOCK: Final = asyncio.Lock() _END_USER_REGISTRY_LOAD_LOCK: Final = asyncio.Lock() +_MODEL_ACCESS_GROUP_REGISTRY_LOAD_LOCK: Final = asyncio.Lock() async def _cached_registry( @@ -1836,6 +1902,105 @@ async def _load_tag_registry( ) +async def _load_model_access_group_registry( + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, +) -> frozenset[str] | None: + """The set of model access group names that have a row in ``LiteLLM_ModelAccessGroupBudgetTable``.""" + + async def fetch_ids() -> tuple[str, ...]: + registry_rows: Final = await _model_access_group_budget_table( + ModelAccessGroupBudgetRepository(prisma_client) + ).find_many(take=MODEL_ACCESS_GROUP_REGISTRY_MAX_SIZE + 1) + return tuple(row.access_group_name for row in registry_rows) + + return await _load_bounded_registry( + cache_key=model_access_group_registry_cache_key(), + overflow_sentinel=MODEL_ACCESS_GROUP_REGISTRY_OVERFLOW_SENTINEL, + max_size=MODEL_ACCESS_GROUP_REGISTRY_MAX_SIZE, + load_lock=_MODEL_ACCESS_GROUP_REGISTRY_LOAD_LOCK, + fetch_ids=fetch_ids, + user_api_key_cache=user_api_key_cache, + ) + + +async def _fetch_uncached_model_access_group_budgets( + uncached_groups: Sequence[str], + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, +) -> tuple[tuple[str, ModelAccessGroupBudget], ...]: + """Budget rows for the groups a cache probe missed. + + No registry gate here, unlike the tag path: the names only ever come from + ``matched_model_access_groups``, which :func:`collect_matched_model_access_groups` already + intersected with the registry, so a name that has no row cannot reach this. + """ + if not uncached_groups: + return () + + try: + db_rows: Final = await _model_access_group_budget_table( + ModelAccessGroupBudgetRepository(prisma_client) + ).find_many( + where={"access_group_name": {"in": list(uncached_groups)}}, + include={"litellm_budget_table": True}, + ) + fetched: Final = tuple((row.access_group_name, _model_access_group_budget(row)) for row in db_rows) + for fetched_name, fetched_obj in fetched: + await user_api_key_cache.async_set_cache( + key=model_access_group_cache_key(fetched_name), + value=fetched_obj, + model_type=ModelAccessGroupBudget, + ttl=DEFAULT_MODEL_ACCESS_GROUP_CACHE_TTL, + ) + except Exception as e: # noqa: BLE001 # fail-safe: a budget fetch error must yield "no budget rows", never break auth + verbose_proxy_logger.debug("Error batch fetching model access group budgets from database: %s", e) + return () + else: + return fetched + + +def _model_access_group_budget(row: _PrismaModelAccessGroupBudgetRow) -> ModelAccessGroupBudget: + budget_table: Final = row.litellm_budget_table + return ModelAccessGroupBudget( + access_group_name=row.access_group_name, + spend=row.spend or 0.0, + max_budget=None if budget_table is None else budget_table.max_budget, + ) + + +@log_db_metrics +async def get_model_access_group_budgets_batch( + access_group_names: Sequence[str], + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, +) -> dict[str, ModelAccessGroupBudget]: + """Budget rows for the given model access groups, served from cache where possible. + + Shared by the two enforcement paths so they read one row per group per request: the + reservation counters when reservations are on, and :func:`_model_access_group_max_budget_check` + when ``disable_budget_reservation`` turns them off. + """ + if prisma_client is None or not access_group_names: + return {} + + probed: Final = [ + ( + group, + await user_api_key_cache.async_get_cache( + key=model_access_group_cache_key(group), model_type=ModelAccessGroupBudget + ), + ) + for group in access_group_names + ] + fetched: Final = await _fetch_uncached_model_access_group_budgets( + uncached_groups=tuple(group for group, budget in probed if budget is None), + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + return {group: budget for group, budget in (*probed, *fetched) if budget is not None} + + async def _fetch_uncached_tags( uncached_tags: Sequence[str], prisma_client: PrismaClient, @@ -3882,6 +4047,192 @@ def _resolve_key_models_for_auth_check(valid_token: UserAPIKeyAuth) -> list[str] return models +def _model_access_groups_serving_model( + model: str | Sequence[str], + llm_router: Router, + team_id: str | None, +) -> frozenset[str]: + """Every model access group whose deployments serve the requested model(s).""" + requested: Final = (model,) if isinstance(model, str) else tuple(model) + return frozenset( + group + for requested_model in requested + for group in llm_router.get_model_access_groups(model_name=requested_model, team_id=team_id) + ) + + +async def _team_member_granted_models( + valid_token: UserAPIKeyAuth, + team_object: LiteLLM_TeamTable | None, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> Sequence[str]: + """The member's own ``allowed_models`` scope; empty when the member is not narrowed below the team.""" + if team_object is None or valid_token.user_id is None: + return () + + team_membership: Final = await get_team_membership( + user_id=valid_token.user_id, + team_id=team_object.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + return () if team_membership is None else _member_allowed_models(team_membership) + + +async def _org_granted_models( + valid_token: UserAPIKeyAuth, + team_object: LiteLLM_TeamTable | None, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> Sequence[str]: + """The org allowlist reached through the key, or through its team when the key names no org.""" + org_id: Final = valid_token.org_id or (team_object.organization_id if team_object is not None else None) + if org_id is None: + return () + + try: + org_object: Final = await get_org_object( + org_id=org_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + except Exception as e: # noqa: BLE001 # fail-safe: attribution degrades to "no org grant", it must never break auth + verbose_proxy_logger.debug("access group attribution: org lookup failed: %s", e) + return () + return org_object.models if org_object is not None else () + + +async def _granted_model_lists( + valid_token: UserAPIKeyAuth, + team_object: LiteLLM_TeamTable | None, + project_object: LiteLLM_ProjectTableCachedObj | None, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> tuple[Sequence[str], ...]: + """One model allowlist per level that participates in authorizing the request.""" + return ( + _resolve_key_models_for_auth_check(valid_token=valid_token), + team_object.models if team_object is not None else (), + await _team_member_granted_models( + valid_token=valid_token, + team_object=team_object, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ), + project_object.models if project_object is not None else (), + await _org_granted_models( + valid_token=valid_token, + team_object=team_object, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ), + ) + + +async def collect_matched_model_access_groups( + model: str | Sequence[str] | None, + valid_token: UserAPIKeyAuth | None, + team_object: LiteLLM_TeamTable | None, + project_object: LiteLLM_ProjectTableCachedObj | None, + llm_router: Router | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> tuple[str, ...]: + """ + The budgeted model access groups that authorized this request, sorted and deduplicated. + + A group is charged only when its name appears on an allowlist the caller was granted -- key, + team, team-member scope, project or org -- *and* that group serves the requested model. Asking + for a model that merely belongs to a group attributes nothing, because nothing about the caller + named the group. + + Levels are unioned, never ranked: a team granted ``*`` whose member is scoped to one group is + still a caller gated by that group. An unrestricted allowlist (empty, ``*``) names no group and + so contributes nothing. + + The whole walk is gated on the budget registry, because collecting every match costs a full scan + of each allowlist where the plain access check stops at the first hit. An empty registry means no + group carries a budget, so there is nothing to attribute and no work worth doing. + """ + if model is None or valid_token is None or llm_router is None or prisma_client is None: + return () + + registry: Final = await _load_model_access_group_registry( + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + if registry is not None and not registry: + return () + + covering_groups: Final = _model_access_groups_serving_model( + model=model, + llm_router=llm_router, + team_id=valid_token.team_id, + ) + budgeted_groups: Final = covering_groups if registry is None else covering_groups & registry + if not budgeted_groups: + return () + + granted: Final = frozenset( + granted_model + for granted_models in await _granted_model_lists( + valid_token=valid_token, + team_object=team_object, + project_object=project_object, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + for granted_model in granted_models + ) + return tuple(sorted(budgeted_groups & granted)) + + +async def stamp_matched_model_access_groups( + model: str | Sequence[str] | None, + valid_token: UserAPIKeyAuth | None, + team_object: LiteLLM_TeamTable | None, + project_object: LiteLLM_ProjectTableCachedObj | None, + llm_router: Router | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> tuple[str, ...]: + """Record the groups that authorized this request on its auth object, for the post-call spend + writer and the reservation counters, and hand them back for the budget check.""" + if valid_token is None: + return () + + try: + matched: Final = await collect_matched_model_access_groups( + model=model, + valid_token=valid_token, + team_object=team_object, + project_object=project_object, + llm_router=llm_router, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + except Exception as e: # noqa: BLE001 # fail-safe: attribution is spend telemetry, it must never break auth + verbose_proxy_logger.debug("model access group attribution failed: %s", e) + return () + if not matched: + return () + matched_groups: Final = list(matched) # mutable-ok: the auth field is typed list[str] | None + valid_token.matched_model_access_groups = matched_groups # rebind-ok: request-scoped carrier for the writer + return matched + + async def can_key_call_model( model: str | list[str], llm_model_list: list | None, @@ -5256,6 +5607,56 @@ async def _tag_max_budget_check( ) +async def _model_access_group_max_budget_check( + matched_model_access_groups: Sequence[str], + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, +) -> None: + """Block the request when a model access group that authorized it is over its max budget. + + Only the groups auth already matched are charged and therefore only they are checked, so a + request that no budgeted group authorized costs nothing here. + + Like the tag check this is a plain read with no reservation, so concurrent requests can + overshoot the ceiling slightly. The reservation counters are the precise path; this one covers + the ``disable_budget_reservation`` case. + + Raises: + BudgetExceededError if a matched group is over its max budget. + """ + if prisma_client is None or not matched_model_access_groups: + return + + budgets: Final = await get_model_access_group_budgets_batch( + access_group_names=matched_model_access_groups, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + + from litellm.proxy.proxy_server import get_current_spend + + for group in matched_model_access_groups: + budget = budgets.get(group) + if budget is None or budget.max_budget is None: + continue + + group_spend = await get_current_spend( + counter_key=model_access_group_spend_counter_key(group), + fallback_spend=budget.spend, + max_budget=budget.max_budget, + fallback_authoritative=True, + ) + if group_spend <= budget.max_budget: + continue + raise litellm.BudgetExceededError( + current_cost=group_spend, + max_budget=budget.max_budget, + message=f"Budget has been exceeded! Model access group={group} Current cost: {group_spend}, Max budget: {budget.max_budget}", + entity_type=Litellm_EntityType.MODEL_ACCESS_GROUP.value, + entity_id=group, + ) + + def is_model_allowed_by_pattern(model: str, allowed_model_pattern: str) -> bool: """ Check if a model matches an allowed pattern. diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index d065b062517..55539ded02b 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -33,7 +33,11 @@ from litellm.proxy.common_utils.timezone_utils import ( compute_budget_reset_at, get_budget_reset_settings, ) -from litellm.proxy.common_utils.user_api_key_cache import tag_cache_key +from litellm.proxy.common_utils.user_api_key_cache import ( + model_access_group_cache_key, + model_access_group_spend_counter_key, + tag_cache_key, +) from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry from litellm.proxy.utils import PrismaClient, ProxyLogging @@ -41,6 +45,7 @@ from litellm.repositories.organization_repository import OrganizationRepository from litellm.repositories.prisma_protocols import SpendLinkedTable from litellm.repositories.table_repositories import ( EndUserRepository, + ModelAccessGroupBudgetRepository, TagRepository, TeamMembershipRepository, ) @@ -92,6 +97,11 @@ class _TagRow(_BudgetLinkedRow, Protocol): def tag_name(self) -> str: ... +class _ModelAccessGroupRow(_BudgetLinkedRow, Protocol): + @property + def access_group_name(self) -> str: ... + + class _EndUserRow(_BudgetLinkedRow, Protocol): @property def user_id(self) -> str: ... @@ -154,6 +164,14 @@ def _tag_cache_keys(row: _TagRow) -> tuple[str, ...]: return (tag_cache_key(row.tag_name),) +def _model_access_group_counter_key(row: _ModelAccessGroupRow) -> str: + return model_access_group_spend_counter_key(row.access_group_name) + + +def _model_access_group_cache_keys(row: _ModelAccessGroupRow) -> tuple[str, ...]: + return (model_access_group_cache_key(row.access_group_name),) + + def _budget_link_where( budget_ids: Sequence[str], extra: Mapping[str, object] = MappingProxyType({}), @@ -610,6 +628,11 @@ class ResetBudgetJob: where=_budget_link_where(budget_ids, _SPENT_ROWS_WHERE), log_subject="tags", ) + model_access_groups: Final[tuple[_ModelAccessGroupRow, ...]] = await self._fetch_linked_rows( + table=ModelAccessGroupBudgetRepository(self.prisma_client).table, + where=_budget_link_where(budget_ids, _SPENT_ROWS_WHERE), + log_subject="model access groups", + ) rollover_caps: Final[Mapping[str, float]] = MappingProxyType( { # mutable-ok: MappingProxyType wraps a one-shot dict comprehension b.budget_id: cap @@ -639,6 +662,10 @@ class ResetBudgetJob: *((_key_counter_key(row), _row_carried_spend(row, rollover_caps)) for row in keys), *((_org_counter_key(row), _row_carried_spend(row, rollover_caps)) for row in orgs), *((_tag_counter_key(row), _row_carried_spend(row, rollover_caps)) for row in tags), + *( + (_model_access_group_counter_key(row), _row_carried_spend(row, rollover_caps)) + for row in model_access_groups + ), ), rollover_caps=rollover_caps, cache_keys=( @@ -646,6 +673,7 @@ class ResetBudgetJob: *(key for row in keys for key in _key_cache_keys(row)), *(key for row in orgs for key in _org_cache_keys(row)), *(key for row in tags for key in _tag_cache_keys(row)), + *(key for row in model_access_groups for key in _model_access_group_cache_keys(row)), ), ) @@ -671,6 +699,7 @@ class ResetBudgetJob: _queue_budget_linked_resets(uow.keys, cascade, extra=_LINKED_KEYS_WHERE) _queue_budget_linked_resets(uow.organizations, cascade, extra=_SPENT_ROWS_WHERE) _queue_budget_linked_resets(uow.tags, cascade, extra=_SPENT_ROWS_WHERE) + _queue_budget_linked_resets(uow.model_access_groups, cascade, extra=_SPENT_ROWS_WHERE) _queue_enduser_resets(uow.endusers, cascade) for budget_id, budget_reset_at in cascade.budget_resets: uow.budgets.queue_window_advance(budget_id=budget_id, budget_reset_at=budget_reset_at) @@ -714,7 +743,8 @@ class ResetBudgetJob: async def reset_budget_for_litellm_budget_table(self) -> None: """ Resets the spend a budget tier gates (end users, team members, keys, - orgs, tags) and advances the tier's budget_reset_at, atomically. + orgs, tags, model access groups) and advances the tier's + budget_reset_at, atomically. Caches are invalidated only after the transaction commits, so a failed run cannot leave a zeroed counter in front of an un-reset DB row. @@ -745,8 +775,9 @@ class ResetBudgetJob: return _ChunkOutcome(fetched=len(cascade.budgets), advanced=advanced) case _BudgetCascadeFailed(cascade=cascade, error=error): verbose_proxy_logger.exception( - "Failed to reset the budget table cascade (team member, enduser, org and tag spend, plus " - "budget_reset_at); nothing was committed and the budgets stay due for the next run: %s", + "Failed to reset the budget table cascade (team member, enduser, org, tag and model access " + "group spend, plus budget_reset_at); nothing was committed and the budgets stay due for the " + "next run: %s", error, exc_info=error, ) diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index 96152240181..ae0b39135f1 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -185,19 +185,29 @@ def tag_registry_cache_key() -> str: return "tag_registry" -#: Cached under ``access_group_registry_cache_key`` when the table exceeds -#: ``ACCESS_GROUP_REGISTRY_MAX_SIZE``: registry unusable, fall back to the per-group lookup. -ACCESS_GROUP_REGISTRY_OVERFLOW_SENTINEL: Final = "__access_group_registry_overflow__" +#: Cached under ``model_access_group_registry_cache_key`` when the table exceeds +#: ``MODEL_ACCESS_GROUP_REGISTRY_MAX_SIZE``: registry unusable, fall back to the per-group lookup. +MODEL_ACCESS_GROUP_REGISTRY_OVERFLOW_SENTINEL: Final = "__model_access_group_registry_overflow__" -def access_group_cache_key(access_group_name: str) -> str: +def model_access_group_cache_key(access_group_name: str) -> str: """Cache key one model access group budget row is stored under; shared so auth, spend tracking and the management endpoints cannot drift.""" - return f"access_group:{access_group_name}" + return f"model_access_group:{access_group_name}" -def access_group_registry_cache_key() -> str: +def model_access_group_registry_cache_key() -> str: """Cache key for the set of model access group names that have a budget row.""" - return "access_group_registry" + return "model_access_group_registry" + + +def model_access_group_spend_counter_key(access_group_name: str) -> str: + """Spend counter key for one model access group; shared so its three owners cannot drift. + + The reservation path writes it, auth reads it to enforce ``max_budget``, and the reset job + clears it on rollover. A copy that drifts in any one of them silently resets or reads a + counter nobody else touches, which shows up as a budget that never trips or never resets. + """ + return f"spend:model_access_group:{access_group_name}" #: Cached under ``end_user_restricted_registry_cache_key`` when the restricted set exceeds diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 3f2777ff1f3..886394d66f9 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 Mapping, Sequence from datetime import datetime, timedelta, timezone from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast, overload @@ -23,6 +24,7 @@ from litellm.constants import ( DB_SPEND_UPDATE_JOB_NAME, INTERNAL_CALL_ORIGIN_METADATA_KEY, ) +from litellm.litellm_core_utils.litellm_logging import coerce_model_access_groups from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.proxy._types import ( DB_RETRY_SAFE_ERROR_TYPES, @@ -86,6 +88,7 @@ class _SpendBatch(Protocol): litellm_organizationtable: BatchTable litellm_tagtable: BatchTable litellm_agentstable: BatchTable + litellm_modelaccessgroupbudgettable: BatchTable class _SpendBatchManager(Protocol): @@ -123,6 +126,52 @@ def _get_llm_router(): return None +class _DeploymentLookup(Protocol): + def get_model_info(self, id: str) -> Mapping[str, object] | None: ... + + +def _served_model_access_groups( + router: _DeploymentLookup | None, + served_model_id: str | None, +) -> frozenset[str] | None: + """Access groups declared by the deployment that actually served the request. + + None when the served deployment cannot be identified, in which case the set + attributed at auth time stands unchanged. + """ + if router is None or not served_model_id: + return None + deployment: Final = router.get_model_info(id=served_model_id) + if deployment is None: + return None + model_info: Final = deployment.get("model_info") + if not isinstance(model_info, Mapping): + return None + declared: Final = model_info.get("access_groups") + if not isinstance(declared, (list, tuple)): + return frozenset() + return frozenset(group for group in declared if isinstance(group, str)) + + +def debitable_model_access_groups( + attributed: Sequence[str] | None, + served_model_id: str | None, + router: _DeploymentLookup | None, +) -> tuple[str, ...]: + """Groups to debit: the set attributed at auth time, narrowed to those the served model belongs to. + + The router may fall back to a model outside the pool auth reserved against, so the + attributed set is the hard upper bound: a group absent from it is never debited. + """ + ordered: Final = coerce_model_access_groups(attributed) + if not ordered: + return () + served: Final = _served_model_access_groups(router=router, served_model_id=served_model_id) + if served is None: + return ordered + return tuple(group for group in ordered if group in served) + + class DBSpendUpdateWriter: """ Module responsible for @@ -187,6 +236,7 @@ class DBSpendUpdateWriter: ## CREATE SPEND LOG PAYLOAD ## from litellm.proxy.spend_tracking.spend_tracking_utils import ( get_logging_payload, + get_request_model_access_groups, ) payload: Final = get_logging_payload( @@ -239,6 +289,7 @@ class DBSpendUpdateWriter: prisma_client=prisma_client, litellm_proxy_budget_name=litellm_proxy_budget_name, payload=payload, + request_model_access_groups=get_request_model_access_groups(kwargs), ) ) @@ -431,9 +482,10 @@ class DBSpendUpdateWriter: prisma_client: PrismaClient | None, litellm_proxy_budget_name: str | None, payload: SpendLogsPayload, + request_model_access_groups: Sequence[str] = (), ): """ - Runs all 11 spend-update helpers sequentially inside a single asyncio task. + Runs all 12 spend-update helpers sequentially inside a single asyncio task. Each helper is wrapped in try/except so one failure doesn't prevent the others. @@ -505,6 +557,20 @@ class DBSpendUpdateWriter: traceback.format_exc(), ) + try: + await self._update_model_access_group_db( + response_cost=response_cost, + request_model_access_groups=request_model_access_groups, + served_model_id=payload_copy.get("model_id"), + prisma_client=prisma_client, + router=_get_llm_router(), + ) + except Exception: + verbose_proxy_logger.debug( + "_batch_database_updates: _update_model_access_group_db failed: %s", + traceback.format_exc(), + ) + _agent_id_for_spend: Final = payload_copy.get("agent_id") try: await self._update_agent_db( @@ -814,6 +880,51 @@ class DBSpendUpdateWriter: ) raise e + async def _update_model_access_group_db( + self, + response_cost: float | None, + request_model_access_groups: Sequence[str] | None, + served_model_id: str | None, + prisma_client: PrismaClient | None, + router: _DeploymentLookup | None = None, + ): + """ + Update spend for every model access group this request is billed against. + + Args: + response_cost: Cost of the request, charged in full to each group + request_model_access_groups: Groups attributed at auth time, the upper bound on what may be debited + served_model_id: Deployment id actually served, used to narrow the attributed set + prisma_client: Prisma client instance + router: Deployment lookup used to re-resolve groups after a fallback + """ + try: + if prisma_client is None: + return + + for model_access_group in debitable_model_access_groups( + attributed=request_model_access_groups, + served_model_id=served_model_id, + router=router, + ): + await self.spend_update_queue.add_update( + update=SpendUpdateQueueItem( + entity_type=Litellm_EntityType.MODEL_ACCESS_GROUP, + entity_id=model_access_group, + response_cost=response_cost, + ) + ) + except Exception as e: + spend_log_error( + "Spend tracking - failed to enqueue model access group spend update. " + "model_access_groups=%s, response_cost=%s - %s", + request_model_access_groups, + response_cost, + str(e), + exc=e, + ) + raise e + async def _insert_spend_log_to_db( self, payload: dict | SpendLogsPayload, @@ -927,7 +1038,8 @@ class DBSpendUpdateWriter: if db_spend_update_transactions is not None: verbose_proxy_logger.info( "Spend tracking - committing spend updates from Redis to DB: " - "keys=%d, users=%d, teams=%d, orgs=%d, end_users=%d, team_members=%d, tags=%d, agents=%d", + "keys=%d, users=%d, teams=%d, orgs=%d, end_users=%d, team_members=%d, tags=%d, agents=%d, " + "model_access_groups=%d", len(db_spend_update_transactions.get("key_list_transactions") or {}), len(db_spend_update_transactions.get("user_list_transactions") or {}), len(db_spend_update_transactions.get("team_list_transactions") or {}), @@ -936,6 +1048,7 @@ class DBSpendUpdateWriter: len(db_spend_update_transactions.get("team_member_list_transactions") or {}), len(db_spend_update_transactions.get("tag_list_transactions") or {}), len(db_spend_update_transactions.get("agent_list_transactions") or {}), + len(db_spend_update_transactions.get("model_access_group_list_transactions") or {}), ) await self._commit_spend_updates_to_db( prisma_client=prisma_client, @@ -1433,6 +1546,20 @@ class DBSpendUpdateWriter: proxy_logging_obj=proxy_logging_obj, ) + ### UPDATE MODEL ACCESS GROUP TABLE ### + model_access_group_list_transactions: Final = db_spend_update_transactions.get( + "model_access_group_list_transactions" + ) + await DBSpendUpdateWriter._update_entity_spend_in_db( + entity_name="Model access group", + transactions=model_access_group_list_transactions, + table_accessor="litellm_modelaccessgroupbudgettable", + where_field="access_group_name", + n_retry_times=n_retry_times, + prisma_client=prisma_client, + proxy_logging_obj=proxy_logging_obj, + ) + ### UPDATE AGENT TABLE ### agent_list_transactions: Final = db_spend_update_transactions["agent_list_transactions"] await DBSpendUpdateWriter._update_entity_spend_in_db( @@ -1449,7 +1576,7 @@ class DBSpendUpdateWriter: async def _update_entity_spend_in_db( entity_name: str, transactions: dict[str, float] | None, - table_accessor: Literal["litellm_tagtable", "litellm_agentstable"], + table_accessor: Literal["litellm_tagtable", "litellm_agentstable", "litellm_modelaccessgroupbudgettable"], where_field: str, n_retry_times: int, prisma_client: PrismaClient, diff --git a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py index ad92902221a..39b2111fb22 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -65,6 +65,7 @@ _SpendTransactionField: TypeAlias = Literal[ "org_list_transactions", "tag_list_transactions", "agent_list_transactions", + "model_access_group_list_transactions", ] _SPEND_TRANSACTION_FIELDS: Final[tuple[_SpendTransactionField, ...]] = ( @@ -76,6 +77,7 @@ _SPEND_TRANSACTION_FIELDS: Final[tuple[_SpendTransactionField, ...]] = ( "org_list_transactions", "tag_list_transactions", "agent_list_transactions", + "model_access_group_list_transactions", ) _ValueT = TypeVar("_ValueT") @@ -397,6 +399,10 @@ class RedisUpdateBuffer: Litellm_EntityType.AGENT, db_spend_update_transactions.get("agent_list_transactions"), ), + ( + Litellm_EntityType.MODEL_ACCESS_GROUP, + db_spend_update_transactions.get("model_access_group_list_transactions"), + ), ] for entity_type, entities in entity_entries: if not entities: @@ -826,6 +832,9 @@ class RedisUpdateBuffer: org_list_transactions=_merged_entity_transactions(list_of_transactions, "org_list_transactions"), tag_list_transactions=_merged_entity_transactions(list_of_transactions, "tag_list_transactions"), agent_list_transactions=_merged_entity_transactions(list_of_transactions, "agent_list_transactions"), + model_access_group_list_transactions=_merged_entity_transactions( + list_of_transactions, "model_access_group_list_transactions" + ), ) async def _emit_new_item_added_to_redis_buffer_event( diff --git a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py index 57cb5e73b64..8c0076b10c1 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py +++ b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py @@ -139,6 +139,7 @@ class SpendUpdateQueue(BaseUpdateQueue): org_list_transactions={}, tag_list_transactions={}, agent_list_transactions={}, + model_access_group_list_transactions={}, ) # Map entity types to their corresponding transaction dictionary keys @@ -151,6 +152,7 @@ class SpendUpdateQueue(BaseUpdateQueue): Litellm_EntityType.ORGANIZATION: "org_list_transactions", Litellm_EntityType.TAG: "tag_list_transactions", Litellm_EntityType.AGENT: "agent_list_transactions", + Litellm_EntityType.MODEL_ACCESS_GROUP: "model_access_group_list_transactions", } for update in updates: @@ -190,6 +192,8 @@ class SpendUpdateQueue(BaseUpdateQueue): transactions_dict = db_spend_update_transactions["tag_list_transactions"] elif dict_key == "agent_list_transactions": transactions_dict = db_spend_update_transactions["agent_list_transactions"] + elif dict_key == "model_access_group_list_transactions": + transactions_dict = db_spend_update_transactions["model_access_group_list_transactions"] else: continue diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index b7688790e54..105145dfd20 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -29,6 +29,7 @@ from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( _request_blocked_callback_params, iter_client_callback_metadata_dicts, ) +from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.litellm_core_utils.url_utils import ( is_url_destination_allowed_by_host, @@ -1377,6 +1378,10 @@ 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 + if user_api_key_dict.matched_model_access_groups: + data[_metadata_variable_name][MODEL_ACCESS_GROUP_METADATA_KEY] = ( + user_api_key_dict.matched_model_access_groups + ) # UserAPIKeyAuth object for MCP server access control data[_metadata_variable_name]["user_api_key_auth"] = user_api_key_dict.model_copy( update={ diff --git a/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py b/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py index e1a7645e988..bf65da1621f 100644 --- a/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py @@ -2,18 +2,33 @@ Allow proxy admin to manage model access groups Endpoints here: -- POST /model_group/new - Create a new access group with multiple model names +- POST /access_group/new - Create a new access group with multiple model names +- GET /access_group/list - List every access group +- GET /access_group/{access_group}/info - Read one access group, including its budget +- PUT /access_group/{access_group}/update - Replace an access group's deployments +- DELETE /access_group/{access_group}/delete - Delete an access group and its budget +- GET /access_group/{access_group}/budget - Read an access group's shared budget and spend +- PUT /access_group/{access_group}/budget - Set or replace an access group's shared budget +- DELETE /access_group/{access_group}/budget - Clear an access group's shared budget """ import json from collections.abc import Mapping, Sequence +from datetime import datetime from typing import TYPE_CHECKING, Any, Final, Protocol from fastapi import APIRouter, Depends, HTTPException +from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.user_api_key_cache import ( + UserApiKeyCache, + model_access_group_cache_key, + model_access_group_registry_cache_key, +) +from litellm.proxy.management_endpoints.common_utils import validate_budget_duration # Clear cache and reload models to pick up the access group changes from litellm.proxy.management_endpoints.model_management_endpoints import ( @@ -22,10 +37,16 @@ from litellm.proxy.management_endpoints.model_management_endpoints import ( model_info_as_mapping, reload_serving_verdict, ) +from litellm.proxy.management_helpers.utils import handle_budget_for_entity from litellm.proxy.utils import PrismaClient from litellm.repositories.model_repository import ModelRepository +from litellm.repositories.table_repositories import ModelAccessGroupBudgetRepository from litellm.types.proxy.management_endpoints.model_management_endpoints import ( + AccessGroupBudget, + AccessGroupBudgetRequest, + AccessGroupBudgetResponse, AccessGroupInfo, + DeleteAccessGroupBudgetResponse, DeleteModelGroupResponse, ListAccessGroupsResponse, NewModelGroupRequest, @@ -36,7 +57,43 @@ from litellm.types.proxy.management_endpoints.model_management_endpoints import if TYPE_CHECKING: from litellm import Router -router: Final = APIRouter() +router: Final = APIRouter(tags=["model management"]) + +_AUTH_DEPENDENCIES: Final = (Depends(user_api_key_auth),) + + +class _ErrorDetail(TypedDict): + error: ReadOnly[str] + + +class _ModelAccessGroupWhere(TypedDict): + access_group_name: ReadOnly[str] + + +class _BudgetInclude(TypedDict): + litellm_budget_table: ReadOnly[bool] + + +class _ModelAccessGroupBudgetCreate(TypedDict): + access_group_name: ReadOnly[str] + budget_id: ReadOnly[str | None] + created_by: ReadOnly[str] + updated_by: ReadOnly[str] + + +class _ModelAccessGroupBudgetUpdate(TypedDict): + budget_id: ReadOnly[str | None] + updated_by: ReadOnly[str] + + +class _ModelAccessGroupBudgetUpsert(TypedDict): + create: ReadOnly[_ModelAccessGroupBudgetCreate] + update: ReadOnly[_ModelAccessGroupBudgetUpdate] + + +def _http_error(status_code: int, message: str) -> HTTPException: + detail: Final[_ErrorDetail] = {"error": message} + return HTTPException(status_code=status_code, detail=detail) class _DeploymentRow(Protocol): @@ -58,10 +115,140 @@ class _ModelTableClient(Protocol): async def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> object: ... +class _BudgetRow(Protocol): + @property + def budget_id(self) -> str: ... + + @property + def max_budget(self) -> float | None: ... + + @property + def soft_budget(self) -> float | None: ... + + @property + def budget_duration(self) -> str | None: ... + + @property + def budget_reset_at(self) -> datetime | None: ... + + +class _ModelAccessGroupBudgetRow(Protocol): + @property + def spend(self) -> float: ... + + @property + def budget_id(self) -> str | None: ... + + @property + def litellm_budget_table(self) -> _BudgetRow | None: ... + + +class _ModelAccessGroupBudgetTableClient(Protocol): + async def find_unique( + self, *, where: Mapping[str, object], include: Mapping[str, object] | None = None + ) -> _ModelAccessGroupBudgetRow | None: ... + + async def upsert( + self, + *, + where: Mapping[str, object], + data: Mapping[str, object], + include: Mapping[str, object] | None = None, + ) -> _ModelAccessGroupBudgetRow: ... + + async def delete(self, *, where: Mapping[str, object]) -> _ModelAccessGroupBudgetRow | None: ... + + def _model_table(prisma_client: PrismaClient) -> _ModelTableClient: return ModelRepository(prisma_client).table +def _model_access_group_budget_table(prisma_client: PrismaClient) -> _ModelAccessGroupBudgetTableClient: + return ModelAccessGroupBudgetRepository(prisma_client).table + + +def _prisma_client_or_500() -> PrismaClient: + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise _http_error(500, "Database not connected.") + return prisma_client + + +def _auth_cache() -> UserApiKeyCache: + from litellm.proxy.proxy_server import user_api_key_cache + + return user_api_key_cache + + +async def _evict_model_access_group_cache_keys(access_group: str, auth_cache: UserApiKeyCache) -> None: + """ + Every endpoint that writes an access group budget row must call this, or the budget stays + unenforced until the TTL expires: auth gates the feature on a cached registry of the groups + that have a budget row, read cache-first with no freshness check. + """ + from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( + evict_and_broadcast, + ) + + await evict_and_broadcast( + cache_keys=(model_access_group_cache_key(access_group), model_access_group_registry_cache_key()), + user_api_key_cache=auth_cache, + ) + + +async def _model_access_group_budget_row( + access_group: str, prisma_client: PrismaClient +) -> _ModelAccessGroupBudgetRow | None: + where: Final[_ModelAccessGroupWhere] = {"access_group_name": access_group} + include: Final[_BudgetInclude] = {"litellm_budget_table": True} + return await _model_access_group_budget_table(prisma_client).find_unique(where=where, include=include) + + +def _budget_or_none(row: _ModelAccessGroupBudgetRow | None) -> AccessGroupBudget | None: + budget: Final = row.litellm_budget_table if row is not None else None + if budget is None: + return None + return AccessGroupBudget( + budget_id=budget.budget_id, + max_budget=budget.max_budget, + soft_budget=budget.soft_budget, + budget_duration=budget.budget_duration, + budget_reset_at=budget.budget_reset_at, + ) + + +def _budget_response(access_group: str, row: _ModelAccessGroupBudgetRow | None) -> AccessGroupBudgetResponse: + return AccessGroupBudgetResponse( + access_group=access_group, + spend=row.spend if row is not None else 0.0, + budget=_budget_or_none(row), + ) + + +async def _delete_model_access_group_budget_row( + access_group: str, prisma_client: PrismaClient, auth_cache: UserApiKeyCache +) -> bool: + """ + Drop the group's budget row only, matching /tag/delete: the LiteLLM_BudgetTable row survives + because the link is ON DELETE SET NULL and a budget_id an admin passed in may be shared with + other entities. + + Evicts unconditionally: a group with no row of its own can still be sitting in the cached + registry, so skipping the eviction when nothing was deleted would leave that stale. + """ + where: Final[_ModelAccessGroupWhere] = {"access_group_name": access_group} + row: Final = await _model_access_group_budget_table(prisma_client).delete(where=where) + await _evict_model_access_group_cache_keys(access_group, auth_cache) + return row is not None + + +async def _raise_404_if_model_access_group_missing(access_group: str, prisma_client: PrismaClient) -> None: + access_groups_map: Final = await get_all_access_groups_from_db(prisma_client=prisma_client) + if access_group not in access_groups_map: + raise _http_error(404, f"Access group '{access_group}' not found") + + def validate_models_exist(model_names: Sequence[str], llm_router: "Router | None") -> tuple[bool, Sequence[str]]: """ Validate that all requested model names exist in the router. @@ -356,8 +543,7 @@ async def get_all_access_groups_from_db( @router.post( "/access_group/new", - tags=["model management"], - dependencies=[Depends(user_api_key_auth)], + dependencies=_AUTH_DEPENDENCIES, response_model=NewModelGroupResponse, ) async def create_model_group( @@ -503,8 +689,7 @@ async def create_model_group( @router.get( "/access_group/list", - tags=["model management"], - dependencies=[Depends(user_api_key_auth)], + dependencies=_AUTH_DEPENDENCIES, response_model=ListAccessGroupsResponse, ) async def list_access_groups( @@ -553,8 +738,7 @@ async def list_access_groups( @router.get( "/access_group/{access_group}/info", - tags=["model management"], - dependencies=[Depends(user_api_key_auth)], + dependencies=_AUTH_DEPENDENCIES, response_model=AccessGroupInfo, ) async def get_access_group_info( @@ -574,7 +758,7 @@ async def get_access_group_info( - access_group: str - The access group name (URL path parameter) Returns: - - AccessGroupInfo with the access group details + - AccessGroupInfo with the access group details, its shared budget and its spend Raises: - HTTPException 404: If access group not found @@ -596,7 +780,15 @@ async def get_access_group_info( detail={"error": f"Access group '{access_group}' not found"}, ) - return access_groups_map[access_group] + info: Final = access_groups_map[access_group] + budget_row: Final = await _model_access_group_budget_row(access_group, prisma_client) + return AccessGroupInfo( + access_group=info.access_group, + model_names=info.model_names, + deployment_count=info.deployment_count, + spend=budget_row.spend if budget_row is not None else 0.0, + budget=_budget_or_none(budget_row), + ) except HTTPException: raise @@ -610,8 +802,7 @@ async def get_access_group_info( @router.put( "/access_group/{access_group}/update", - tags=["model management"], - dependencies=[Depends(user_api_key_auth)], + dependencies=_AUTH_DEPENDENCIES, response_model=NewModelGroupResponse, ) async def update_access_group( @@ -765,13 +956,13 @@ async def update_access_group( @router.delete( "/access_group/{access_group}/delete", - tags=["model management"], - dependencies=[Depends(user_api_key_auth)], + dependencies=_AUTH_DEPENDENCIES, response_model=DeleteModelGroupResponse, ) async def delete_access_group( access_group: str, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + auth_cache: UserApiKeyCache = Depends(_auth_cache), ): """ Delete an access group. @@ -835,6 +1026,13 @@ async def delete_access_group( removed_pairs: Final = tuple(pair for pair in removed if pair is not None) models_updated: Final = len(removed_pairs) + # Budget last, deliberately: failing here strands a budget row for a group already on no + # deployment (clutter), where the reverse order can leave a live group enforcing nothing. + # The LiteLLM_BudgetTable row it linked is left alone, as /tag/delete leaves a tag's. + await _delete_model_access_group_budget_row( + access_group=access_group, prisma_client=prisma_client, auth_cache=auth_cache + ) + # Clear cache and reload models to pick up the access group changes live_before_reload: Final = live_model_ids_snapshot() reload_outcome: Final = await clear_cache() @@ -864,3 +1062,164 @@ async def delete_access_group( status_code=500, detail={"error": f"Failed to delete access group: {e}"}, ) + + +@router.get( + "/access_group/{access_group}/budget", + dependencies=_AUTH_DEPENDENCIES, + response_model=AccessGroupBudgetResponse, +) +async def get_access_group_budget( + access_group: str, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +) -> AccessGroupBudgetResponse: + """ + Get the shared budget of an access group, and the spend drawn against it. + + Example: + ```bash + curl -X GET 'http://localhost:4000/access_group/production-models/budget' \\ + -H 'Authorization: Bearer sk-1234' + ``` + + Parameters: + - access_group: str - The access group name (URL path parameter) + + Returns: + - AccessGroupBudgetResponse; budget is null when the group has no budget set + + Raises: + - HTTPException 404: If access group not found + """ + prisma_client: Final = _prisma_client_or_500() + await _raise_404_if_model_access_group_missing(access_group=access_group, prisma_client=prisma_client) + + return _budget_response( + access_group=access_group, + row=await _model_access_group_budget_row(access_group, prisma_client), + ) + + +@router.put( + "/access_group/{access_group}/budget", + dependencies=_AUTH_DEPENDENCIES, + response_model=AccessGroupBudgetResponse, +) +async def set_access_group_budget( + access_group: str, + data: AccessGroupBudgetRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + auth_cache: UserApiKeyCache = Depends(_auth_cache), +) -> AccessGroupBudgetResponse: + """ + Set or replace the shared budget of an access group. Idempotent. + + Every key that can reach a model in the group draws from this one budget. + + Example: + ```bash + curl -X PUT 'http://localhost:4000/access_group/production-models/budget' \\ + -H 'Authorization: Bearer sk-1234' \\ + -H 'Content-Type: application/json' \\ + -d '{ + "max_budget": 100.0, + "budget_duration": "30d" + }' + ``` + + Parameters: + - access_group: str - The access group name (URL path parameter) + - max_budget: Optional[float] - Requests fail once the group's shared spend exceeds this + - soft_budget: Optional[float] - Fires an alert when reached; requests still succeed + - budget_duration: Optional[str] - Frequency of resetting the group's spend (e.g. '30d') + - budget_id: Optional[str] - Link an existing budget instead of creating one + + Returns: + - AccessGroupBudgetResponse with the stored budget and current spend + + Raises: + - HTTPException 400: If no budget field is given, or budget_duration cannot be parsed + - HTTPException 404: If access group not found + """ + from litellm.proxy.proxy_server import litellm_proxy_admin_name + + prisma_client: Final = _prisma_client_or_500() + if not data.model_dump(exclude_none=True): + raise _http_error(400, "One of max_budget, soft_budget, budget_duration or budget_id is required") + validate_budget_duration(data.budget_duration) + await _raise_404_if_model_access_group_missing(access_group=access_group, prisma_client=prisma_client) + + existing_row: Final = await _model_access_group_budget_row(access_group, prisma_client) + budget_id: Final = await handle_budget_for_entity( + data=data, + existing_budget_id=existing_row.budget_id if existing_row is not None else None, + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + litellm_proxy_admin_name=litellm_proxy_admin_name, + ) + actor: Final = user_api_key_dict.user_id or litellm_proxy_admin_name + upsert_data: Final[_ModelAccessGroupBudgetUpsert] = { + "create": { + "access_group_name": access_group, + "budget_id": budget_id, + "created_by": actor, + "updated_by": actor, + }, + "update": {"budget_id": budget_id, "updated_by": actor}, + } + where: Final[_ModelAccessGroupWhere] = {"access_group_name": access_group} + include: Final[_BudgetInclude] = {"litellm_budget_table": True} + row: Final = await _model_access_group_budget_table(prisma_client).upsert( + where=where, data=upsert_data, include=include + ) + await _evict_model_access_group_cache_keys(access_group, auth_cache) + + verbose_proxy_logger.info("Set budget %s on access group '%s'", budget_id, access_group) + return _budget_response(access_group=access_group, row=row) + + +@router.delete( + "/access_group/{access_group}/budget", + dependencies=_AUTH_DEPENDENCIES, + response_model=DeleteAccessGroupBudgetResponse, +) +async def delete_access_group_budget( + access_group: str, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + auth_cache: UserApiKeyCache = Depends(_auth_cache), +) -> DeleteAccessGroupBudgetResponse: + """ + Clear the shared budget of an access group, leaving the group itself in place. + + Example: + ```bash + curl -X DELETE 'http://localhost:4000/access_group/production-models/budget' \\ + -H 'Authorization: Bearer sk-1234' + ``` + + Parameters: + - access_group: str - The access group name (URL path parameter) + + Returns: + - DeleteAccessGroupBudgetResponse; budget_deleted is false when there was nothing to clear + + Raises: + - HTTPException 404: If access group not found + """ + prisma_client: Final = _prisma_client_or_500() + await _raise_404_if_model_access_group_missing(access_group=access_group, prisma_client=prisma_client) + + budget_deleted: Final = await _delete_model_access_group_budget_row( + access_group=access_group, + prisma_client=prisma_client, + auth_cache=auth_cache, + ) + return DeleteAccessGroupBudgetResponse( + access_group=access_group, + budget_deleted=budget_deleted, + message=( + f"Budget for access group '{access_group}' deleted successfully" + if budget_deleted + else f"Access group '{access_group}' has no budget to delete" + ), + ) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 3d60f4f5f3a..09d3dedaafa 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -47,6 +47,7 @@ from litellm.litellm_core_utils.core_helpers import ( get_metadata_variable_name_from_kwargs, get_or_create_metadata_bucket, ) +from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.litellm_core_utils.safe_json_dumps import safe_dumps @@ -577,6 +578,7 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): _metadata["user_api_key"] = user_api_key_dict.api_key _metadata["litellm_parent_otel_span"] = user_api_key_dict.parent_otel_span _metadata["user_api_key_budget_reservation"] = user_api_key_dict.budget_reservation + _metadata[MODEL_ACCESS_GROUP_METADATA_KEY] = user_api_key_dict.matched_model_access_groups # The per-model budget counters are keyed off these. get_sanitized_user_information_from_key # returns StandardLoggingUserAPIKeyMetadata, which carries no budget field, so without this # the post-call increment finds nothing and every passthrough request goes untracked and diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 7cad3f0a022..46086711923 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -12,7 +12,6 @@ from fastapi import HTTPException, status import litellm from litellm._logging import verbose_proxy_logger -from litellm.caching import DualCache from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.litellm_core_utils.llm_cost_calc.tiered_pricing import select_tier_for_input, tier_rate from litellm.proxy._types import ( @@ -26,12 +25,16 @@ from litellm.proxy.auth.auth_utils import get_model_from_request from litellm.proxy.auth.budget_throttle import should_throttle_budget_exceeded from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.user_api_key_cache import ( + UserApiKeyCache, end_user_cache_key, + model_access_group_cache_key, + model_access_group_spend_counter_key, tag_cache_key, team_membership_reservation_cache_key, ) from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.router import Router +from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget @dataclass @@ -53,6 +56,7 @@ _COUNTER_ENTITY_TYPES: Final[Mapping[str, str]] = { "User": Litellm_EntityType.USER.value, "EndUser": Litellm_EntityType.END_USER.value, "Tag": Litellm_EntityType.TAG.value, + "Model access group": Litellm_EntityType.MODEL_ACCESS_GROUP.value, "Organization": Litellm_EntityType.ORGANIZATION.value, } @@ -158,7 +162,7 @@ async def reserve_budget_for_request( team_object: LiteLLM_TeamTable | None, user_object: LiteLLM_UserTable | None, prisma_client: PrismaClient | None, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, end_user_id: str | None = None, end_user_object: object = None, @@ -348,7 +352,7 @@ async def _get_budget_counters( team_object: LiteLLM_TeamTable | None, user_object: LiteLLM_UserTable | None, prisma_client: PrismaClient | None, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, end_user_id: str | None = None, end_user_object: object = None, @@ -437,6 +441,14 @@ async def _get_budget_counters( ) ) + counters.extend( + await _get_model_access_group_budget_counters( + valid_token=valid_token, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + ) + team_member_counter: Final = await _get_team_member_budget_counter( valid_token=valid_token, team_object=team_object, @@ -491,7 +503,7 @@ async def _get_end_user_budget_counter( async def _get_tag_budget_counters( request_body: dict, prisma_client: PrismaClient | None, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, ) -> list[_BudgetCounter]: from litellm.proxy.auth.auth_checks import get_tag_objects_batch @@ -530,6 +542,46 @@ async def _get_tag_budget_counters( return counters +async def _get_model_access_group_budget_counters( + valid_token: UserAPIKeyAuth, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, +) -> list[_BudgetCounter]: + """Reservation counters for the model access groups that authorized this request. + + The names come off the auth object rather than the request body: ``common_checks`` already + resolved which granted groups serve the requested model, and re-deriving that here would both + duplicate the walk and risk disagreeing with what the spend writer attributes. + """ + from litellm.proxy.auth.auth_checks import get_model_access_group_budgets_batch + + group_names: Final = tuple(dict.fromkeys(valid_token.matched_model_access_groups or ())) + if not group_names: + return [] + + budgets: Final = await get_model_access_group_budgets_batch( + access_group_names=group_names, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + candidates: Final = (_model_access_group_counter(group, budgets.get(group)) for group in group_names) + return [counter for counter in candidates if counter is not None] + + +def _model_access_group_counter(group: str, budget: ModelAccessGroupBudget | None) -> _BudgetCounter | None: + """A counter for one group, or nothing when the group carries no budget to reserve against.""" + if budget is None or budget.max_budget is None or budget.max_budget <= 0: + return None + return _BudgetCounter( + counter_key=model_access_group_spend_counter_key(group), + source_cache_key=model_access_group_cache_key(group), + max_budget=budget.max_budget, + fallback_spend=budget.spend, + entity_type="Model access group", + entity_id=group, + ) + + def _dedupe_tags(tags: list[str]) -> list[str]: seen: Final = set() deduped_tags: Final = [] @@ -545,7 +597,7 @@ async def _get_team_member_budget_counter( valid_token: UserAPIKeyAuth, team_object: LiteLLM_TeamTable | None, user_object: LiteLLM_UserTable | None, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, ) -> _BudgetCounter | None: if team_object is None or team_object.team_id is None or user_object is None or valid_token.user_id is None: return None @@ -588,7 +640,7 @@ async def _get_team_member_budget_counter( async def _get_org_budget_counter( valid_token: UserAPIKeyAuth, team_object: LiteLLM_TeamTable | None, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, ) -> _BudgetCounter | None: org_id: str | None = None if valid_token.org_id is not None: diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index a6b6375fd9c..d66f8c8e9d7 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -1,6 +1,7 @@ import os import re import secrets +from collections.abc import Mapping from datetime import datetime, timezone from datetime import datetime as dt from typing import Any, Final, Literal, cast @@ -23,7 +24,11 @@ from litellm.litellm_core_utils.core_helpers import ( reconstruct_model_name, ) from litellm.litellm_core_utils.internal_call_metadata import is_unbilled_non_inference_call -from litellm.litellm_core_utils.litellm_logging import is_valid_sha256_hash +from litellm.litellm_core_utils.litellm_logging import ( + coerce_model_access_groups, + is_valid_sha256_hash, + request_model_access_groups_from_litellm_params, +) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps, strip_null_bytes from litellm.proxy._types import SpendLogsMetadata, SpendLogsPayload from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error @@ -243,6 +248,23 @@ def _extract_usage_for_ocr_call(response_obj: Any, response_obj_dict: dict) -> d return {} +def get_request_model_access_groups(kwargs: Mapping[str, object] | None) -> tuple[str, ...]: + """Model access groups that authorized this request, as stamped onto request metadata at auth time.""" + if kwargs is None: + return () + + standard_logging_payload: Final = kwargs.get("standard_logging_object") + if isinstance(standard_logging_payload, Mapping): + from_payload: Final = coerce_model_access_groups(standard_logging_payload.get("request_model_access_groups")) + if from_payload: + return from_payload + + litellm_params: Final = kwargs.get("litellm_params") + if not isinstance(litellm_params, Mapping): + return () + return request_model_access_groups_from_litellm_params(litellm_params) + + def _sl_attribution_fallback( standard_logging_payload: StandardLoggingPayload | None, field: Literal["model_id", "model_group", "api_base", "custom_llm_provider"], diff --git a/litellm/repositories/prisma_protocols.py b/litellm/repositories/prisma_protocols.py index 2aa1b8e0e3f..d962934dfb1 100644 --- a/litellm/repositories/prisma_protocols.py +++ b/litellm/repositories/prisma_protocols.py @@ -144,4 +144,7 @@ class PrismaBatch(Protocol): @property def litellm_endusertable(self) -> BatchTable: ... + @property + def litellm_modelaccessgroupbudgettable(self) -> BatchTable: ... + async def commit(self) -> None: ... diff --git a/litellm/repositories/unit_of_work.py b/litellm/repositories/unit_of_work.py index eb11ebe3b9c..a497d0580db 100644 --- a/litellm/repositories/unit_of_work.py +++ b/litellm/repositories/unit_of_work.py @@ -118,6 +118,7 @@ class BudgetCascadeUnitOfWork: keys: LinkedSpendResetWrites organizations: LinkedSpendResetWrites tags: LinkedSpendResetWrites + model_access_groups: LinkedSpendResetWrites endusers: LinkedSpendResetWrites budgets: BudgetWindowWrites @@ -143,6 +144,7 @@ async def budget_cascade_unit_of_work( keys=LinkedSpendResetWrites(table=batch.litellm_verificationtoken), organizations=LinkedSpendResetWrites(table=batch.litellm_organizationtable), tags=LinkedSpendResetWrites(table=batch.litellm_tagtable), + model_access_groups=LinkedSpendResetWrites(table=batch.litellm_modelaccessgroupbudgettable), endusers=LinkedSpendResetWrites(table=batch.litellm_endusertable), budgets=BudgetWindowWrites(table=batch.litellm_budgettable), ) diff --git a/litellm/types/proxy/management_endpoints/model_management_endpoints.py b/litellm/types/proxy/management_endpoints/model_management_endpoints.py index 6e18787a224..9d4663631fe 100644 --- a/litellm/types/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/types/proxy/management_endpoints/model_management_endpoints.py @@ -1,6 +1,7 @@ +from datetime import datetime from typing import Any -from pydantic import BaseModel, Field +from pydantic import BaseModel, ConfigDict, Field from ...router import ModelGroupInfo @@ -53,10 +54,42 @@ class DeleteModelGroupResponse(BaseModel): message: str +class AccessGroupBudget(BaseModel): + budget_id: str + max_budget: float | None = None + soft_budget: float | None = None + budget_duration: str | None = None + budget_reset_at: datetime | None = None + + +class AccessGroupBudgetRequest(BaseModel): + budget_id: str | None = None # Link an existing budget instead of creating one + max_budget: float | None = Field(default=None, ge=0) + soft_budget: float | None = Field(default=None, ge=0) + budget_duration: str | None = None + + # rejects tpm_limit/rpm_limit/max_parallel_requests: those are not enforced per access group + model_config = ConfigDict(extra="forbid") + + +class AccessGroupBudgetResponse(BaseModel): + access_group: str + spend: float # Shared spend accrued by every key that can reach this access group + budget: AccessGroupBudget | None = None + + +class DeleteAccessGroupBudgetResponse(BaseModel): + access_group: str + budget_deleted: bool # False when the access group had no budget to begin with + message: str + + class AccessGroupInfo(BaseModel): access_group: str model_names: list[str] # List of model names in this access group deployment_count: int # Total number of deployments with this access group + spend: float | None = None # Only populated by /access_group/{access_group}/info + budget: AccessGroupBudget | None = None class ListAccessGroupsResponse(BaseModel): diff --git a/litellm/types/proxy/model_access_group_budget.py b/litellm/types/proxy/model_access_group_budget.py new file mode 100644 index 00000000000..cccbe92b5d6 --- /dev/null +++ b/litellm/types/proxy/model_access_group_budget.py @@ -0,0 +1,19 @@ +"""The model access group budget state auth and the spend reservation path share.""" + +from __future__ import annotations + +from pydantic import BaseModel + + +class ModelAccessGroupBudget(BaseModel): + """One model access group's budget, flattened out of its joined ``LiteLLM_ModelAccessGroupBudgetTable`` row. + + Both readers want only the recorded spend and the ceiling, and this sits on the per-request hot + path behind a cache, so the linked budget row is collapsed to ``max_budget`` rather than cached + whole. ``spend`` is the DB-recorded value, which lags the live counter and is only ever a + fallback for it. + """ + + access_group_name: str + spend: float = 0.0 + max_budget: float | None = None diff --git a/litellm/types/utils.py b/litellm/types/utils.py index f0319a7c664..abddd1aa7c3 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -40,7 +40,7 @@ from pydantic import ( field_serializer, field_validator, ) -from typing_extensions import ReadOnly, Required, TypedDict +from typing_extensions import NotRequired, ReadOnly, Required, TypedDict from litellm._logging import verbose_logger from litellm._uuid import uuid @@ -3258,6 +3258,7 @@ class StandardLoggingPayload(TypedDict): cache_key: str | None saved_cache_cost: float request_tags: list + request_model_access_groups: NotRequired[ReadOnly[Sequence[str]]] end_user: str | None requester_ip_address: str | None user_agent: str | None diff --git a/schema.prisma b/schema.prisma index 8ddbee63973..48d41edffbd 100644 --- a/schema.prisma +++ b/schema.prisma @@ -586,6 +586,9 @@ model LiteLLM_EndUserTable { blocked Boolean @default(false) } +// Budget and shared spend for a model access group. The groups themselves are not rows anywhere: +// they are free-text strings in LiteLLM_ProxyModelTable.model_info.access_groups, so a row here +// exists only once someone gives that group a budget. model LiteLLM_ModelAccessGroupBudgetTable { access_group_name String @id spend Float @default(0.0) diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 1193160c831..bc5ecccc10d 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -3727,6 +3727,60 @@ def test_get_standard_logging_object_payload_includes_litellm_call_id(logging_ob assert payload["litellm_call_id"] == call_id +def test_get_standard_logging_object_payload_carries_matched_access_groups(logging_obj): + """Access groups stamped at auth time reach the logging payload, so integrations see what a request billed.""" + from datetime import datetime + + from litellm.litellm_core_utils.litellm_logging import ( + get_standard_logging_object_payload, + ) + + now = datetime.now() + payload = get_standard_logging_object_payload( + kwargs={ + "model": "gpt-4o", + "messages": [], + "litellm_params": { + "metadata": { + "user_api_key_matched_model_access_groups": ["premium-pool", "shared-pool"] + }, + "proxy_server_request": {"body": {}}, + }, + }, + init_response_obj={}, + start_time=now, + end_time=now, + logging_obj=logging_obj, + status="success", + ) + + assert payload is not None + assert payload["request_model_access_groups"] == ("premium-pool", "shared-pool") + + +def test_get_standard_logging_object_payload_has_no_access_groups_when_unstamped( + logging_obj, +): + from datetime import datetime + + from litellm.litellm_core_utils.litellm_logging import ( + get_standard_logging_object_payload, + ) + + now = datetime.now() + payload = get_standard_logging_object_payload( + kwargs={"model": "gpt-4o", "messages": []}, + init_response_obj={}, + start_time=now, + end_time=now, + logging_obj=logging_obj, + status="success", + ) + + assert payload is not None + assert payload["request_model_access_groups"] == () + + def test_get_standard_logging_object_payload_preserves_absent_end_user_as_none(logging_obj): from datetime import datetime from typing import Final diff --git a/tests/test_litellm/proxy/auth/test_model_access_group_budgets.py b/tests/test_litellm/proxy/auth/test_model_access_group_budgets.py new file mode 100644 index 00000000000..2896f76eec6 --- /dev/null +++ b/tests/test_litellm/proxy/auth/test_model_access_group_budgets.py @@ -0,0 +1,504 @@ +""" +Which model access groups a request is charged to. + +A group is attributed only when its name appears on an allowlist the caller was granted, so the +group is what authorized the call. Asking for a model that merely belongs to a group attributes +nothing, and every level that can name a group (key, team, team-member scope, project, org) is +unioned rather than ranked. +""" + +from types import SimpleNamespace +from unittest.mock import patch + +import pytest + +import litellm +from litellm import Router +from litellm.proxy._types import ( + LiteLLM_BudgetTable, + Litellm_EntityType, + LiteLLM_OrganizationTable, + LiteLLM_ProjectTableCachedObj, + LiteLLM_TeamMembership, + LiteLLM_TeamTable, + UserAPIKeyAuth, +) +from litellm.proxy.auth.auth_checks import ( + _model_access_group_max_budget_check, + collect_matched_model_access_groups, + common_checks, + stamp_matched_model_access_groups, +) +from litellm.proxy.common_utils.reset_budget_job import _model_access_group_counter_key +from litellm.proxy.common_utils.user_api_key_cache import ( + UserApiKeyCache, + model_access_group_registry_cache_key, + model_access_group_spend_counter_key, + team_membership_reservation_cache_key, +) +from litellm.proxy.utils import ProxyLogging + +TEAM_ID = "team-1" +USER_ID = "user-1" +ORG_ID = "org-1" +BUDGETED_GROUPS = ("tier-a", "tier-b", "claude-tier") +MODEL_ACCESS_GROUP_COUNTER_KEY = model_access_group_spend_counter_key("tier-a") + +MODEL_LIST = [ + { + "model_name": "gpt-4o", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "k"}, + "model_info": {"access_groups": ["tier-a", "tier-b"]}, + }, + { + "model_name": "claude-sonnet", + "litellm_params": {"model": "anthropic/claude-sonnet", "api_key": "k"}, + "model_info": {"access_groups": ["claude-tier"]}, + }, +] + + +class _ExplodingPrismaClient: + """Every lookup in these tests is served from the injected cache; a real DB read is a bug.""" + + def __getattr__(self, name: str) -> object: + raise AssertionError(f"unexpected database access: {name}") + + +class _CountingRouter(Router): + """Counts access-group lookups, so a test can prove the registry gate skipped them.""" + + def __init__(self, **kwargs) -> None: + super().__init__(**kwargs) + self.access_group_lookups = 0 + + def get_model_access_groups(self, *args, **kwargs): + self.access_group_lookups += 1 + return super().get_model_access_groups(*args, **kwargs) + + +async def _cache( + budgeted_groups: tuple[str, ...] = BUDGETED_GROUPS, + member_allowed_models: tuple[str, ...] = (), + org_models: tuple[str, ...] = (), +) -> UserApiKeyCache: + cache = UserApiKeyCache() + await cache.async_set_cache(key=model_access_group_registry_cache_key(), value=budgeted_groups) + if member_allowed_models: + await cache.async_set_cache( + key=team_membership_reservation_cache_key(user_id=USER_ID, team_id=TEAM_ID), + value=LiteLLM_TeamMembership( + user_id=USER_ID, + team_id=TEAM_ID, + budget_id="member-budget", + litellm_budget_table=LiteLLM_BudgetTable(allowed_models=list(member_allowed_models)), + ), + model_type=LiteLLM_TeamMembership, + ) + if org_models: + await cache.async_set_cache( + key=f"org_id:{ORG_ID}", + value=LiteLLM_OrganizationTable( + organization_id=ORG_ID, + budget_id="org-budget", + models=list(org_models), + created_by=USER_ID, + updated_by=USER_ID, + ), + model_type=LiteLLM_OrganizationTable, + ) + return cache + + +async def _matched( + *, + model: str = "gpt-4o", + key_models: list[str] | None = None, + team_models: list[str] | None = None, + team_org_id: str | None = None, + project_models: list[str] | None = None, + valid_token: UserAPIKeyAuth | None = None, + cache: UserApiKeyCache | None = None, + llm_router: Router | None = None, +) -> tuple[str, ...]: + resolved_cache = cache if cache is not None else await _cache() + return await collect_matched_model_access_groups( + model=model, + valid_token=valid_token + if valid_token is not None + else UserAPIKeyAuth(api_key="hashed", models=key_models or [], team_id=TEAM_ID, user_id=USER_ID), + team_object=( + LiteLLM_TeamTable(team_id=TEAM_ID, models=team_models, organization_id=team_org_id) + if team_models is not None + else None + ), + project_object=( + LiteLLM_ProjectTableCachedObj(project_id="project-1", models=project_models) + if project_models is not None + else None + ), + llm_router=llm_router if llm_router is not None else Router(model_list=MODEL_LIST), + prisma_client=_ExplodingPrismaClient(), + user_api_key_cache=resolved_cache, + proxy_logging_obj=ProxyLogging(user_api_key_cache=resolved_cache), + ) + + +@pytest.mark.asyncio +async def test_group_named_on_the_key_is_attributed(): + assert await _matched(key_models=["tier-a"]) == ("tier-a",) + + +@pytest.mark.asyncio +async def test_model_granted_directly_on_the_key_attributes_nothing(): + assert await _matched(key_models=["gpt-4o"]) == () + + +@pytest.mark.asyncio +@pytest.mark.parametrize("key_models", [["*"], [], ["all-proxy-models"]]) +async def test_unrestricted_key_attributes_nothing(key_models: list[str]): + assert await _matched(key_models=key_models) == () + + +@pytest.mark.asyncio +async def test_group_that_does_not_serve_the_requested_model_is_not_attributed(): + assert await _matched(model="gpt-4o", key_models=["claude-tier"]) == () + + +@pytest.mark.asyncio +async def test_both_granted_groups_covering_the_model_are_attributed(): + assert await _matched(key_models=["tier-b", "tier-a"]) == ("tier-a", "tier-b") + + +@pytest.mark.asyncio +async def test_group_named_only_on_the_team_is_attributed(): + assert await _matched(key_models=[], team_models=["tier-a"]) == ("tier-a",) + + +@pytest.mark.asyncio +async def test_group_named_only_in_a_team_members_scope_is_attributed(): + assert await _matched( + model="claude-sonnet", + key_models=["*"], + team_models=["*"], + cache=await _cache(member_allowed_models=("claude-tier",)), + ) == ("claude-tier",) + + +@pytest.mark.asyncio +async def test_group_named_only_on_the_project_is_attributed(): + assert await _matched(key_models=["*"], project_models=["tier-b"]) == ("tier-b",) + + +@pytest.mark.asyncio +async def test_group_named_only_on_the_org_is_attributed(): + assert await _matched( + valid_token=UserAPIKeyAuth(api_key="hashed", models=["*"], user_id=USER_ID, org_id=ORG_ID), + cache=await _cache(org_models=("tier-a",)), + ) == ("tier-a",) + + +@pytest.mark.asyncio +async def test_group_named_on_the_teams_org_is_attributed_when_the_key_names_no_org(): + assert await _matched( + key_models=["*"], + team_models=["*"], + team_org_id=ORG_ID, + cache=await _cache(org_models=("tier-b",)), + ) == ("tier-b",) + + +@pytest.mark.asyncio +async def test_all_team_models_sentinel_on_the_key_resolves_to_the_teams_groups(): + assert await _matched( + valid_token=UserAPIKeyAuth( + api_key="hashed", + models=["all-team-models"], + team_models=["tier-a"], + team_id=TEAM_ID, + user_id=USER_ID, + ), + ) == ("tier-a",) + + +@pytest.mark.asyncio +async def test_group_without_a_budget_is_not_attributed(): + assert await _matched(key_models=["tier-a"], cache=await _cache(budgeted_groups=("tier-b",))) == () + + +@pytest.mark.asyncio +async def test_empty_registry_skips_the_access_group_matching_entirely(): + router = _CountingRouter(model_list=MODEL_LIST) + + assert await _matched(key_models=["tier-a"], cache=await _cache(budgeted_groups=()), llm_router=router) == () + assert router.access_group_lookups == 0 + + assert await _matched(key_models=["tier-a"], llm_router=router) == ("tier-a",) + assert router.access_group_lookups == 1 + + +@pytest.mark.asyncio +async def test_stamp_records_the_matched_groups_on_the_auth_object(): + cache = await _cache() + valid_token = UserAPIKeyAuth(api_key="hashed", models=["tier-a", "tier-b"], team_id=TEAM_ID, user_id=USER_ID) + + await stamp_matched_model_access_groups( + model="gpt-4o", + valid_token=valid_token, + team_object=None, + project_object=None, + llm_router=Router(model_list=MODEL_LIST), + prisma_client=_ExplodingPrismaClient(), + user_api_key_cache=cache, + proxy_logging_obj=ProxyLogging(user_api_key_cache=cache), + ) + + assert valid_token.matched_model_access_groups == ["tier-a", "tier-b"] + + +class _BrokenRouter(Router): + def get_model_access_groups(self, *args, **kwargs): + raise RuntimeError("access group store unavailable") + + +@pytest.mark.asyncio +async def test_stamp_does_not_break_auth_when_the_access_group_lookup_fails(): + cache = await _cache() + valid_token = UserAPIKeyAuth(api_key="hashed", models=["tier-a"], team_id=TEAM_ID, user_id=USER_ID) + + await stamp_matched_model_access_groups( + model="gpt-4o", + valid_token=valid_token, + team_object=None, + project_object=None, + llm_router=_BrokenRouter(model_list=MODEL_LIST), + prisma_client=_ExplodingPrismaClient(), + user_api_key_cache=cache, + proxy_logging_obj=ProxyLogging(user_api_key_cache=cache), + ) + + assert valid_token.matched_model_access_groups is None + + +@pytest.mark.asyncio +async def test_stamp_leaves_the_auth_object_untouched_when_nothing_matched(): + cache = await _cache() + valid_token = UserAPIKeyAuth(api_key="hashed", models=["gpt-4o"], team_id=TEAM_ID, user_id=USER_ID) + + await stamp_matched_model_access_groups( + model="gpt-4o", + valid_token=valid_token, + team_object=None, + project_object=None, + llm_router=Router(model_list=MODEL_LIST), + prisma_client=_ExplodingPrismaClient(), + user_api_key_cache=cache, + proxy_logging_obj=ProxyLogging(user_api_key_cache=cache), + ) + + assert valid_token.matched_model_access_groups is None + + +class _MagBudgetRow: + """One ``LiteLLM_ModelAccessGroupBudgetTable`` row as prisma hands it back.""" + + def __init__(self, access_group_name: str, spend: float = 0.0, max_budget: float | None = None) -> None: + self.access_group_name = access_group_name + self.spend = spend + self.litellm_budget_table = None if max_budget is None else SimpleNamespace(max_budget=max_budget) + + +class _RecordingPrismaClient: + """Serves budget rows and records which groups actually reached the database.""" + + def __init__(self, *rows: _MagBudgetRow) -> None: + self.rows = {row.access_group_name: row for row in rows} + self.batches: list[list[str]] = [] + self.db = SimpleNamespace( + litellm_modelaccessgroupbudgettable=SimpleNamespace(find_many=self._find_many) + ) + + async def _find_many(self, **kwargs): + requested = list(kwargs["where"]["access_group_name"]["in"]) + self.batches.append(requested) + return [self.rows[group] for group in requested if group in self.rows] + + +def _spend_reader(spend_by_counter_key: dict[str, float]): + """Stand-in for proxy_server.get_current_spend, recording every counter key it is asked for.""" + seen: list[str] = [] + + async def read(counter_key, fallback_spend, max_budget=None, **kwargs): + seen.append(counter_key) + return spend_by_counter_key.get(counter_key, fallback_spend) + + return read, seen + + +async def _enforce( + matched: tuple[str, ...], + *rows: _MagBudgetRow, + spend_by_counter_key: dict[str, float] | None = None, + prisma_client: object | None = None, + cache: UserApiKeyCache | None = None, +) -> list[str]: + read, seen = _spend_reader(spend_by_counter_key or {}) + with patch("litellm.proxy.proxy_server.get_current_spend", read): + await _model_access_group_max_budget_check( + matched_model_access_groups=matched, + prisma_client=prisma_client if prisma_client is not None else _RecordingPrismaClient(*rows), + user_api_key_cache=cache if cache is not None else UserApiKeyCache(), + ) + return seen + + +@pytest.mark.asyncio +async def test_group_under_its_max_budget_passes(): + assert await _enforce( + ("tier-a",), + _MagBudgetRow("tier-a", spend=4.0, max_budget=10.0), + spend_by_counter_key={MODEL_ACCESS_GROUP_COUNTER_KEY: 4.0}, + ) == [MODEL_ACCESS_GROUP_COUNTER_KEY] + + +@pytest.mark.asyncio +async def test_group_exactly_at_its_max_budget_passes(): + """The ceiling is inclusive, matching the tag check it mirrors; only spend strictly above it blocks.""" + await _enforce( + ("tier-a",), + _MagBudgetRow("tier-a", max_budget=10.0), + spend_by_counter_key={MODEL_ACCESS_GROUP_COUNTER_KEY: 10.0}, + ) + + +@pytest.mark.asyncio +async def test_group_over_its_max_budget_blocks_the_request_and_names_the_group(): + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _enforce( + ("tier-a",), + _MagBudgetRow("tier-a", max_budget=10.0), + spend_by_counter_key={MODEL_ACCESS_GROUP_COUNTER_KEY: 10.5}, + ) + + assert exc_info.value.entity_id == "tier-a" + assert exc_info.value.entity_type == Litellm_EntityType.MODEL_ACCESS_GROUP.value + assert exc_info.value.current_cost == 10.5 + assert exc_info.value.max_budget == 10.0 + assert "tier-a" in str(exc_info.value) + + +@pytest.mark.asyncio +async def test_group_with_a_row_but_no_budget_never_blocks(): + """An admin can register a group without a ceiling; that must not become an implicit zero budget.""" + assert ( + await _enforce( + ("tier-a",), + _MagBudgetRow("tier-a", spend=9999.0), + spend_by_counter_key={MODEL_ACCESS_GROUP_COUNTER_KEY: 9999.0}, + ) + == [] + ) + + +@pytest.mark.asyncio +async def test_a_cold_counter_falls_back_to_the_spend_recorded_on_the_row(): + """After a counter expires the DB row is the only record of the spend, so it has to be read.""" + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _enforce(("tier-a",), _MagBudgetRow("tier-a", spend=12.0, max_budget=10.0)) + + assert exc_info.value.current_cost == 12.0 + + +@pytest.mark.asyncio +async def test_an_over_budget_group_blocks_even_when_another_matched_group_is_fine(): + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _enforce( + ("tier-a", "tier-b"), + _MagBudgetRow("tier-a", max_budget=10.0), + _MagBudgetRow("tier-b", max_budget=1.0), + spend_by_counter_key={ + MODEL_ACCESS_GROUP_COUNTER_KEY: 1.0, + model_access_group_spend_counter_key("tier-b"): 5.0, + }, + ) + + assert exc_info.value.entity_id == "tier-b" + + +@pytest.mark.asyncio +async def test_request_that_matched_no_group_touches_neither_database_nor_counters(): + assert await _enforce((), prisma_client=_ExplodingPrismaClient()) == [] + + +@pytest.mark.asyncio +async def test_budget_check_reads_the_counter_key_the_reset_job_clears(): + """Reads and resets must agree, or a rollover clears a counter nobody reads.""" + reset_job_key = _model_access_group_counter_key(SimpleNamespace(access_group_name="tier-a")) + + assert await _enforce(("tier-a",), _MagBudgetRow("tier-a", max_budget=10.0)) == [reset_job_key] + + +@pytest.mark.asyncio +async def test_a_second_request_serves_the_budget_row_from_cache(): + cache = UserApiKeyCache() + prisma_client = _RecordingPrismaClient(_MagBudgetRow("tier-a", max_budget=10.0)) + + await _enforce(("tier-a",), prisma_client=prisma_client, cache=cache) + await _enforce(("tier-a",), prisma_client=prisma_client, cache=cache) + + assert prisma_client.batches == [["tier-a"]] + + +@pytest.mark.asyncio +async def test_a_database_error_does_not_block_the_request(): + class _FailingPrismaClient: + def __init__(self) -> None: + self.db = SimpleNamespace( + litellm_modelaccessgroupbudgettable=SimpleNamespace(find_many=self._boom) + ) + + async def _boom(self, **kwargs): + raise RuntimeError("database unavailable") + + assert await _enforce(("tier-a",), prisma_client=_FailingPrismaClient()) == [] + + +async def _common_checks_with_over_budget_group(*, skip_budget_checks: bool) -> bool: + cache = await _cache() + prisma_client = _RecordingPrismaClient(_MagBudgetRow("tier-a", max_budget=1.0)) + read, _ = _spend_reader({MODEL_ACCESS_GROUP_COUNTER_KEY: 99.0}) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma_client), + patch("litellm.proxy.proxy_server.user_api_key_cache", cache), + patch("litellm.proxy.proxy_server.get_current_spend", read), + patch("litellm.proxy.auth.auth_checks._is_api_route_allowed", return_value=True), + ): + return await common_checks( + request_body={"model": "gpt-4o", "messages": []}, + team_object=None, + user_object=None, + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route="/v1/chat/completions", + llm_router=Router(model_list=MODEL_LIST), + proxy_logging_obj=ProxyLogging(user_api_key_cache=cache), + valid_token=UserAPIKeyAuth(api_key="hashed", models=["tier-a"], user_id=USER_ID), + request=SimpleNamespace(method="POST", headers={}, query_params={}, url=SimpleNamespace(path="/v1/chat/completions")), + skip_budget_checks=skip_budget_checks, + ) + + +@pytest.mark.asyncio +async def test_common_checks_blocks_a_request_whose_group_is_over_budget(): + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _common_checks_with_over_budget_group(skip_budget_checks=False) + + assert exc_info.value.entity_id == "tier-a" + + +@pytest.mark.asyncio +async def test_free_model_routes_skip_the_model_access_group_budget_check(): + """skip_budget_checks is how free models stay free; it has to cover this budget too.""" + assert await _common_checks_with_over_budget_group(skip_budget_checks=True) is True diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index 5d3afd95a55..eb713b6f3c2 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py @@ -77,6 +77,7 @@ class MockBatcher: self.litellm_teammembership = _Table("team_membership", self) self.litellm_organizationtable = _Table("org", self) self.litellm_tagtable = _Table("tag", self) + self.litellm_modelaccessgroupbudgettable = _Table("model_access_group", self) self.litellm_endusertable = _Table("enduser", self) async def commit(self): @@ -91,6 +92,7 @@ class MockDB: self.litellm_endusertable = MockTable() self.litellm_organizationtable = MockTable() self.litellm_tagtable = MockTable() + self.litellm_modelaccessgroupbudgettable = MockTable() self.batch_calls: List[Dict[str, Any]] = [] self.batchers: List[MockBatcher] = [] @@ -521,6 +523,7 @@ _LINKED_TABLE_CASES = [ ), ("org", {"budget_id": {"in": ["7d-budget-tier"]}, "spend": {"gt": 0}}), ("tag", {"budget_id": {"in": ["7d-budget-tier"]}, "spend": {"gt": 0}}), + ("model_access_group", {"budget_id": {"in": ["7d-budget-tier"]}, "spend": {"gt": 0}}), ] @@ -1299,13 +1302,19 @@ _INVALIDATION_CASES = [ "spend:tag:tenant-42", {"tag:tenant-42"}, ), + ( + "litellm_modelaccessgroupbudgettable", + type("AccessGroup", (), {"access_group_name": "gpt-4-group"}), + "spend:model_access_group:gpt-4-group", + {"model_access_group:gpt-4-group"}, + ), ] @pytest.mark.parametrize( "table_attr, linked_row, counter_key, cache_keys", _INVALIDATION_CASES, - ids=["team_membership", "key", "org", "tag"], + ids=["team_membership", "key", "org", "tag", "model_access_group"], ) def test_budget_table_reset_invalidates_counters_and_management_cache( reset_budget_job, mock_prisma_client, monkeypatch, table_attr, linked_row, counter_key, cache_keys @@ -1359,6 +1368,102 @@ def test_budget_table_reset_commits_even_when_cache_eviction_fails(reset_budget_ assert mock_prisma_client.db.batchers[0].committed is True +# --------------------------------------------------------------------------- +# Model access group budgets ride the same cascade +# --------------------------------------------------------------------------- + + +def _model_access_group_row(name: str = "gpt-4-group", spend: float = 12.0, budget_id: str = "budget-1"): + """A LiteLLM_ModelAccessGroupBudgetTable row, shaped like prisma hands it back.""" + return type("AccessGroup", (), {"access_group_name": name, "spend": spend, "budget_id": budget_id}) + + +def test_access_group_reset_only_matches_rows_that_have_spend(reset_budget_job, mock_prisma_client, monkeypatch): + """Both the read and the write are filtered to spend > 0 on the due tiers. + + A group sitting at spend 0 has nothing to reset, and a group hanging off a + tier that is not due yet must not be swept along: both are excluded by the + filter, not by anything downstream. + """ + _make_counter_invalidation_job(monkeypatch) + mock_prisma_client.data["budget"] = [_budget_row(budget_id="budget-due", budget_duration="7d")] + mock_prisma_client.db.litellm_modelaccessgroupbudgettable.set_find_many_results( + [_model_access_group_row(budget_id="budget-due")] + ) + + asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) + + expected_where = {"budget_id": {"in": ["budget-due"]}, "spend": {"gt": 0}} + assert mock_prisma_client.db.litellm_modelaccessgroupbudgettable.find_many_calls == [{"where": expected_where}] + writes = _batch_writes(mock_prisma_client, "model_access_group", op="update_many") + assert len(writes) == 1 + assert writes[0]["where"] == expected_where + assert writes[0]["data"] == {"spend": 0} + + +def test_access_groups_are_untouched_when_no_budget_is_due(reset_budget_job, mock_prisma_client, monkeypatch): + """No due tier means the group table is never read, written or evicted.""" + counter_cache = _make_counter_invalidation_job(monkeypatch) + mock_prisma_client.db.litellm_modelaccessgroupbudgettable.set_find_many_results([_model_access_group_row()]) + + asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) + + assert mock_prisma_client.db.litellm_modelaccessgroupbudgettable.find_many_calls == [] + assert _batch_writes(mock_prisma_client, "model_access_group") == [] + counter_cache.in_memory_cache.set_cache.assert_not_called() + counter_cache.user_api_key_cache.async_delete_cache.assert_not_awaited() + + +def test_budget_table_reset_invalidates_every_access_group_not_just_the_first( + reset_budget_job, mock_prisma_client, monkeypatch +): + """When several groups share the expiring tier, all of them are evicted.""" + counter_cache = _make_counter_invalidation_job(monkeypatch) + mock_prisma_client.data["budget"] = [_budget_row(budget_id="budget-1")] + mock_prisma_client.db.litellm_modelaccessgroupbudgettable.set_find_many_results( + [_model_access_group_row(name=name) for name in ("group-a", "group-b", "group-c")] + ) + + asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) + + deleted = {call.kwargs.get("key") for call in counter_cache.user_api_key_cache.async_delete_cache.await_args_list} + assert deleted == {"model_access_group:group-a", "model_access_group:group-b", "model_access_group:group-c"} + for name in ("group-a", "group-b", "group-c"): + counter_cache.in_memory_cache.set_cache.assert_any_call(key=f"spend:model_access_group:{name}", value=0.0, ttl=60) + + +def test_budget_cascade_carries_access_group_overage_when_rollover_enabled( + rollover_enabled, reset_budget_job, mock_prisma_client, monkeypatch +): + """A group 5 over the tier cap keeps a spend of 5 in the next window, the + same way a tag or a team member does: over-cap rows are decremented by the + cap, the rest are zeroed, and the counter is seeded with the carried spend.""" + counter_cache = _make_counter_invalidation_job(monkeypatch) + mock_prisma_client.data["budget"] = [_budget_row(budget_id="budget-roll", budget_duration="7d", max_budget=10.0)] + mock_prisma_client.db.litellm_modelaccessgroupbudgettable.set_find_many_results( + [_model_access_group_row(spend=15.0, budget_id="budget-roll")] + ) + + asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) + + writes = _batch_writes(mock_prisma_client, "model_access_group") + assert { + "table": "model_access_group", + "op": "update_many", + "where": {"budget_id": "budget-roll", "spend": {"gt": 10.0}}, + "data": {"spend": {"decrement": 10.0}}, + } in writes + assert { + "table": "model_access_group", + "op": "update_many", + "where": {"budget_id": "budget-roll", "spend": {"gt": 0, "lte": 10.0}}, + "data": {"spend": 0}, + } in writes + assert _replay_spend_writes(writes, 15.0) == 5.0 + assert _replay_spend_writes(writes, 8.0) == 0 + counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:model_access_group:gpt-4-group", value=5.0, ttl=60) + + # --------------------------------------------------------------------------- # Atomicity of the budget-table cascade (LIT-5138) # --------------------------------------------------------------------------- @@ -1471,6 +1576,7 @@ def test_budget_cascade_writes_land_in_a_single_transaction(reset_budget_job, mo ("key", "update_many"), ("org", "update_many"), ("tag", "update_many"), + ("model_access_group", "update_many"), ("enduser", "update_many"), ("budget", "update_many"), } @@ -1506,7 +1612,7 @@ def test_failed_cascade_is_logged_as_a_cascade_failure(monkeypatch): assert mock_exception.call_count == 1 message = mock_exception.call_args.args[0] assert "cascade" in message - for mentioned in ("team member", "enduser", "org", "tag", "budget_reset_at"): + for mentioned in ("team member", "enduser", "org", "tag", "model access group", "budget_reset_at"): assert mentioned in message, f"failure log should mention {mentioned}: {message}" diff --git a/tests/test_litellm/proxy/db/test_model_access_group_spend.py b/tests/test_litellm/proxy/db/test_model_access_group_spend.py new file mode 100644 index 00000000000..79dffc29e7e --- /dev/null +++ b/tests/test_litellm/proxy/db/test_model_access_group_spend.py @@ -0,0 +1,508 @@ +"""Spend accumulation for model access group budgets.""" + +import asyncio +from collections.abc import Mapping, Sequence + +import pytest + +from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY +from litellm.proxy._types import DBSpendUpdateTransactions, Litellm_EntityType, SpendUpdateQueueItem +from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter, debitable_model_access_groups +from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import DailySpendUpdateQueue +from litellm.proxy.db.db_transaction_queue.redis_update_buffer import RedisUpdateBuffer +from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue +from litellm.proxy.spend_tracking.spend_tracking_utils import get_request_model_access_groups + + +class _FakeRouter: + """Deployment lookup returning the access groups each deployment declares.""" + + def __init__(self, deployments: Mapping[str, Sequence[str] | None]) -> None: + self._deployments = deployments + + def get_model_info(self, id: str) -> dict | None: + if id not in self._deployments: + return None + declared = self._deployments[id] + model_info: dict = {"id": id} + if declared is not None: + model_info["access_groups"] = list(declared) + return {"model_name": "some-model", "model_info": model_info} + + +class _FakeBatchTable: + def __init__(self) -> None: + self.calls: list[tuple[dict, dict]] = [] + + def update_many(self, where: dict, data: dict) -> None: + self.calls.append((where, data)) + + +class _FakeBatcher: + def __init__(self) -> None: + self.tables: dict[str, _FakeBatchTable] = {} + + def __getattr__(self, name: str) -> _FakeBatchTable: + return self.tables.setdefault(name, _FakeBatchTable()) + + +class _FakeBatchManager: + def __init__(self, batcher: _FakeBatcher) -> None: + self._batcher = batcher + + async def __aenter__(self) -> _FakeBatcher: + return self._batcher + + async def __aexit__(self, *exc_info: object) -> bool: + return False + + +class _FakeTransaction: + def __init__(self, batcher: _FakeBatcher) -> None: + self._batcher = batcher + + def batch_(self) -> _FakeBatchManager: + return _FakeBatchManager(self._batcher) + + async def __aenter__(self) -> "_FakeTransaction": + return self + + async def __aexit__(self, *exc_info: object) -> bool: + return False + + +class _FakeDb: + def __init__(self, batcher: _FakeBatcher) -> None: + self._batcher = batcher + + def tx(self, timeout: object = None) -> _FakeTransaction: + return _FakeTransaction(self._batcher) + + +class _FakePrismaClient: + def __init__(self) -> None: + self.batcher = _FakeBatcher() + self.db = _FakeDb(self.batcher) + + +def _empty_transactions(**overrides: dict[str, float]) -> DBSpendUpdateTransactions: + return DBSpendUpdateTransactions( + user_list_transactions=overrides.get("user_list_transactions", {}), + end_user_list_transactions=overrides.get("end_user_list_transactions", {}), + key_list_transactions=overrides.get("key_list_transactions", {}), + team_list_transactions=overrides.get("team_list_transactions", {}), + team_member_list_transactions=overrides.get("team_member_list_transactions", {}), + org_list_transactions=overrides.get("org_list_transactions", {}), + tag_list_transactions=overrides.get("tag_list_transactions", {}), + agent_list_transactions=overrides.get("agent_list_transactions", {}), + model_access_group_list_transactions=overrides.get("model_access_group_list_transactions", {}), + ) + + +async def _drain(queue: SpendUpdateQueue) -> list[SpendUpdateQueueItem]: + return await queue.flush_all_updates_from_in_memory_queue() + + +# --- enqueue --------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_single_matched_group_enqueues_one_item_with_full_cost(): + writer = DBSpendUpdateWriter() + + await writer._update_model_access_group_db( + response_cost=0.42, + request_model_access_groups=["premium-pool"], + served_model_id="deployment-1", + prisma_client=object(), + router=_FakeRouter({"deployment-1": ["premium-pool"]}), + ) + + updates = await _drain(writer.spend_update_queue) + assert updates == [ + SpendUpdateQueueItem( + entity_type=Litellm_EntityType.MODEL_ACCESS_GROUP, + entity_id="premium-pool", + response_cost=0.42, + ) + ] + + +@pytest.mark.asyncio +async def test_every_matched_group_is_charged_the_full_cost_not_a_split(): + writer = DBSpendUpdateWriter() + + await writer._update_model_access_group_db( + response_cost=0.30, + request_model_access_groups=["pool-a", "pool-b", "pool-c"], + served_model_id="deployment-1", + prisma_client=object(), + router=_FakeRouter({"deployment-1": ["pool-a", "pool-b", "pool-c"]}), + ) + + updates = await _drain(writer.spend_update_queue) + assert [update["entity_id"] for update in updates] == ["pool-a", "pool-b", "pool-c"] + assert [update["response_cost"] for update in updates] == [0.30, 0.30, 0.30] + assert {update["entity_type"] for update in updates} == {Litellm_EntityType.MODEL_ACCESS_GROUP} + + +@pytest.mark.parametrize("attributed", [None, [], ()]) +@pytest.mark.asyncio +async def test_no_attributed_groups_enqueues_nothing(attributed): + writer = DBSpendUpdateWriter() + + await writer._update_model_access_group_db( + response_cost=1.0, + request_model_access_groups=attributed, + served_model_id="deployment-1", + prisma_client=object(), + router=_FakeRouter({"deployment-1": ["premium-pool"]}), + ) + + assert await _drain(writer.spend_update_queue) == [] + + +@pytest.mark.asyncio +async def test_no_prisma_client_enqueues_nothing(): + writer = DBSpendUpdateWriter() + + await writer._update_model_access_group_db( + response_cost=1.0, + request_model_access_groups=["premium-pool"], + served_model_id="deployment-1", + prisma_client=None, + router=_FakeRouter({"deployment-1": ["premium-pool"]}), + ) + + assert await _drain(writer.spend_update_queue) == [] + + +@pytest.mark.asyncio +async def test_group_outside_the_attributed_set_is_never_debited(): + """The served deployment also sits in a pool auth never attributed; that pool stays untouched.""" + writer = DBSpendUpdateWriter() + + await writer._update_model_access_group_db( + response_cost=0.10, + request_model_access_groups=["premium-pool"], + served_model_id="deployment-1", + prisma_client=object(), + router=_FakeRouter({"deployment-1": ["premium-pool", "unattributed-pool"]}), + ) + + updates = await _drain(writer.spend_update_queue) + assert [update["entity_id"] for update in updates] == ["premium-pool"] + + +# --- fallback guard -------------------------------------------------------- + + +def test_fallback_to_a_model_in_another_pool_debits_nothing(): + assert ( + debitable_model_access_groups( + attributed=["premium-pool"], + served_model_id="fallback-deployment", + router=_FakeRouter({"fallback-deployment": ["cheap-pool"]}), + ) + == () + ) + + +def test_fallback_to_a_model_in_no_pool_debits_nothing(): + assert ( + debitable_model_access_groups( + attributed=["premium-pool"], + served_model_id="fallback-deployment", + router=_FakeRouter({"fallback-deployment": None}), + ) + == () + ) + + +def test_attributed_set_stands_when_the_served_deployment_is_unknown(): + assert debitable_model_access_groups( + attributed=["premium-pool"], + served_model_id="not-in-router", + router=_FakeRouter({"deployment-1": ["premium-pool"]}), + ) == ("premium-pool",) + + +def test_attributed_set_stands_without_a_router(): + assert debitable_model_access_groups( + attributed=["premium-pool", "premium-pool"], + served_model_id="deployment-1", + router=None, + ) == ("premium-pool",) + + +def test_partial_overlap_keeps_only_the_intersection(): + assert debitable_model_access_groups( + attributed=["pool-a", "pool-b"], + served_model_id="deployment-1", + router=_FakeRouter({"deployment-1": ["pool-b", "pool-c"]}), + ) == ("pool-b",) + + +def test_only_real_group_names_ever_become_entity_ids(): + """Whatever shape the attributed set arrives in, an empty or non-string name never reaches the queue.""" + assert debitable_model_access_groups( + attributed=["pool-a", "", "pool-a", None, 7], + served_model_id=None, + router=None, + ) == ("pool-a",) + + +# --- metadata extraction --------------------------------------------------- + + +def test_access_groups_read_from_request_metadata(): + kwargs = {"litellm_params": {"metadata": {MODEL_ACCESS_GROUP_METADATA_KEY: ["pool-a", "pool-b", "pool-a"]}}} + assert get_request_model_access_groups(kwargs) == ("pool-a", "pool-b") + + +def test_access_groups_read_from_litellm_metadata(): + kwargs = {"litellm_params": {"litellm_metadata": {MODEL_ACCESS_GROUP_METADATA_KEY: ["pool-a"]}}} + assert get_request_model_access_groups(kwargs) == ("pool-a",) + + +def test_standard_logging_payload_wins_over_metadata(): + kwargs = { + "litellm_params": {"metadata": {MODEL_ACCESS_GROUP_METADATA_KEY: ["from-metadata"]}}, + "standard_logging_object": {"request_model_access_groups": ["from-payload"]}, + } + assert get_request_model_access_groups(kwargs) == ("from-payload",) + + +def test_metadata_is_used_when_the_logging_payload_carries_no_groups(): + kwargs = { + "litellm_params": {"metadata": {MODEL_ACCESS_GROUP_METADATA_KEY: ["from-metadata"]}}, + "standard_logging_object": {"request_model_access_groups": []}, + } + assert get_request_model_access_groups(kwargs) == ("from-metadata",) + + +@pytest.mark.parametrize("stamped", ["pool-a", 7, {"pool-a": 1}]) +def test_non_list_access_group_metadata_is_ignored(stamped): + kwargs = {"litellm_params": {"metadata": {MODEL_ACCESS_GROUP_METADATA_KEY: stamped}}} + assert get_request_model_access_groups(kwargs) == () + + +def test_key_absent_from_metadata_yields_no_groups(): + """The chat path only stamps the key when something matched, so absent must mean nothing to debit.""" + kwargs = {"litellm_params": {"metadata": {"user_api_key_user_id": "u-1"}}} + assert get_request_model_access_groups(kwargs) == () + + +def test_explicit_none_yields_no_groups(): + """The pass-through path stamps the key unconditionally, so it can be present and None.""" + kwargs = {"litellm_params": {"metadata": {MODEL_ACCESS_GROUP_METADATA_KEY: None}}} + assert get_request_model_access_groups(kwargs) == () + + +@pytest.mark.parametrize( + "metadata", + [ + {"user_api_key_user_id": "u-1"}, + {MODEL_ACCESS_GROUP_METADATA_KEY: None}, + ], + ids=["key-absent", "key-present-but-none"], +) +@pytest.mark.asyncio +async def test_neither_absent_nor_none_metadata_debits_anything(metadata): + writer = DBSpendUpdateWriter() + + await writer._update_model_access_group_db( + response_cost=0.5, + request_model_access_groups=get_request_model_access_groups({"litellm_params": {"metadata": metadata}}), + served_model_id="deployment-1", + prisma_client=object(), + router=_FakeRouter({"deployment-1": ["premium-pool"]}), + ) + + assert await _drain(writer.spend_update_queue) == [] + + +def test_detached_sub_call_falls_back_to_the_auth_object(): + """Sub-calls inherit only the identity keys, so the groups come off user_api_key_auth there.""" + + class _Auth: + matched_model_access_groups = ["premium-pool"] + + kwargs = {"litellm_params": {"metadata": {"user_api_key_auth": _Auth()}}} + assert get_request_model_access_groups(kwargs) == ("premium-pool",) + + +def test_stamped_metadata_wins_over_the_auth_object(): + class _Auth: + matched_model_access_groups = ["stale-pool"] + + kwargs = { + "litellm_params": { + "metadata": { + MODEL_ACCESS_GROUP_METADATA_KEY: ["fresh-pool"], + "user_api_key_auth": _Auth(), + } + } + } + assert get_request_model_access_groups(kwargs) == ("fresh-pool",) + + +def test_auth_object_without_matched_groups_yields_no_groups(): + class _Auth: + matched_model_access_groups = None + + kwargs = {"litellm_params": {"metadata": {"user_api_key_auth": _Auth()}}} + assert get_request_model_access_groups(kwargs) == () + + +def test_non_string_entries_are_dropped(): + kwargs = {"litellm_params": {"metadata": {MODEL_ACCESS_GROUP_METADATA_KEY: ["pool-a", None, "", 3]}}} + assert get_request_model_access_groups(kwargs) == ("pool-a",) + + +def test_missing_metadata_yields_no_groups(): + assert get_request_model_access_groups(None) == () + assert get_request_model_access_groups({}) == () + assert get_request_model_access_groups({"litellm_params": {}}) == () + + +# --- queue bucketing and redis round trip ---------------------------------- + + +def test_access_group_updates_aggregate_into_their_own_bucket(): + queue = SpendUpdateQueue() + + transactions = queue.get_aggregated_db_spend_update_transactions( + [ + SpendUpdateQueueItem( + entity_type=Litellm_EntityType.MODEL_ACCESS_GROUP, entity_id="pool-a", response_cost=0.1 + ), + SpendUpdateQueueItem( + entity_type=Litellm_EntityType.MODEL_ACCESS_GROUP, entity_id="pool-a", response_cost=0.2 + ), + SpendUpdateQueueItem( + entity_type=Litellm_EntityType.MODEL_ACCESS_GROUP, entity_id="pool-b", response_cost=0.5 + ), + SpendUpdateQueueItem(entity_type=Litellm_EntityType.TAG, entity_id="pool-a", response_cost=9.0), + ] + ) + + assert transactions["model_access_group_list_transactions"] == {"pool-a": pytest.approx(0.3), "pool-b": 0.5} + assert transactions["tag_list_transactions"] == {"pool-a": 9.0} + + +def test_access_group_transactions_survive_the_redis_buffer_merge(): + merged = RedisUpdateBuffer._combine_list_of_transactions( + [ + _empty_transactions(model_access_group_list_transactions={"pool-a": 0.25}), + _empty_transactions(model_access_group_list_transactions={"pool-a": 0.25, "pool-b": 1.0}), + ] + ) + + assert merged["model_access_group_list_transactions"] == {"pool-a": 0.5, "pool-b": 1.0} + + +@pytest.mark.asyncio +async def test_redis_buffer_requeues_access_group_transactions_as_queue_items(): + queue = SpendUpdateQueue() + daily_queue = DailySpendUpdateQueue() + + await RedisUpdateBuffer._restore_spend_updates_to_in_memory_queues( + db_spend_update_transactions=_empty_transactions(model_access_group_list_transactions={"pool-a": 0.75}), + daily_spend_update_transactions=None, + daily_team_spend_update_transactions=None, + daily_org_spend_update_transactions=None, + daily_end_user_spend_update_transactions=None, + daily_agent_spend_update_transactions=None, + spend_update_queue=queue, + daily_spend_update_queue=daily_queue, + daily_team_spend_update_queue=daily_queue, + daily_org_spend_update_queue=daily_queue, + daily_end_user_spend_update_queue=daily_queue, + daily_agent_spend_update_queue=daily_queue, + ) + + updates = await _drain(queue) + assert updates == [ + SpendUpdateQueueItem( + entity_type=Litellm_EntityType.MODEL_ACCESS_GROUP, + entity_id="pool-a", + response_cost=0.75, + ) + ] + + +# --- flush to postgres ----------------------------------------------------- + + +@pytest.mark.asyncio +async def test_commit_increments_spend_on_the_model_access_group_budget_table(): + prisma_client = _FakePrismaClient() + + await DBSpendUpdateWriter()._commit_spend_updates_to_db( + prisma_client=prisma_client, + n_retry_times=0, + proxy_logging_obj=None, + db_spend_update_transactions=_empty_transactions( + model_access_group_list_transactions={"pool-b": 0.5, "pool-a": 0.25} + ), + ) + + assert prisma_client.batcher.tables["litellm_modelaccessgroupbudgettable"].calls == [ + ({"access_group_name": "pool-a"}, {"spend": {"increment": 0.25}}), + ({"access_group_name": "pool-b"}, {"spend": {"increment": 0.5}}), + ] + assert "litellm_tagtable" not in prisma_client.batcher.tables + + +# --- end-to-end through the batched fan-out -------------------------------- + + +@pytest.mark.asyncio +async def test_batch_database_updates_enqueues_access_group_spend(): + writer = DBSpendUpdateWriter() + + await writer._batch_database_updates( + response_cost=0.15, + user_id=None, + hashed_token=None, + team_id=None, + org_id=None, + end_user_id=None, + prisma_client=object(), + litellm_proxy_budget_name=None, + payload={"model_id": "deployment-1", "spend": 0.15}, + request_model_access_groups=("pool-a", "pool-b"), + ) + await asyncio.sleep(0) + + access_group_updates = [ + update + for update in await _drain(writer.spend_update_queue) + if update["entity_type"] is Litellm_EntityType.MODEL_ACCESS_GROUP + ] + assert [(update["entity_id"], update["response_cost"]) for update in access_group_updates] == [ + ("pool-a", 0.15), + ("pool-b", 0.15), + ] + + +@pytest.mark.asyncio +async def test_batch_database_updates_enqueues_nothing_without_access_groups(): + writer = DBSpendUpdateWriter() + + await writer._batch_database_updates( + response_cost=0.15, + user_id=None, + hashed_token=None, + team_id=None, + org_id=None, + end_user_id=None, + prisma_client=object(), + litellm_proxy_budget_name=None, + payload={"model_id": "deployment-1", "spend": 0.15}, + ) + await asyncio.sleep(0) + + updates = await _drain(writer.spend_update_queue) + assert [update for update in updates if update["entity_type"] is Litellm_EntityType.MODEL_ACCESS_GROUP] == [] diff --git a/tests/test_litellm/proxy/management_endpoints/test_access_group_management.py b/tests/test_litellm/proxy/management_endpoints/test_access_group_management.py index db0557cfbf0..74171774fc6 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_access_group_management.py +++ b/tests/test_litellm/proxy/management_endpoints/test_access_group_management.py @@ -2,6 +2,10 @@ Test access group management endpoints """ +from contextlib import contextmanager +from dataclasses import dataclass +from datetime import datetime +from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -449,6 +453,7 @@ async def test_delete_access_group_ignores_models_that_were_already_dead(): mock_prisma = MagicMock() mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[deploy_broken]) mock_prisma.db.litellm_proxymodeltable.update = AsyncMock() + mock_prisma.db.litellm_modelaccessgroupbudgettable.delete = AsyncMock(return_value=None) from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( @@ -468,6 +473,7 @@ async def test_delete_access_group_ignores_models_that_were_already_dead(): response = await delete_access_group( access_group="doomed-group", user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), + auth_cache=_FakeAuthCache(), ) assert response.models_updated == 1 @@ -568,3 +574,652 @@ async def test_create_access_group_model_missing_everywhere_still_400s(): assert exc_info.value.status_code == 400 assert model_name in str(exc_info.value.detail) + +@dataclass +class _FakeBudgetRow: + budget_id: str + max_budget: float | None = None + soft_budget: float | None = None + budget_duration: str | None = None + budget_reset_at: datetime | None = None + + +@dataclass +class _FakeAccessGroupBudgetRow: + access_group_name: str + budget_id: str | None = None + spend: float = 0.0 + litellm_budget_table: _FakeBudgetRow | None = None + + +@dataclass +class _FakeDeployment: + model_id: str + model_name: str + model_info: dict + + +class _FakeBudgetTable: + """Stands in for litellm_budgettable so a test can see whether a budget row was created, + updated in place, or left orphaned.""" + + def __init__(self, journal: list[str]) -> None: + self.journal = journal + self.rows: dict[str, _FakeBudgetRow] = {} + self.create_calls: list[dict] = [] + self.update_calls: list[tuple[str, dict]] = [] + self.deleted_ids: list[str] = [] + self._sequence = 0 + + async def create(self, data, include=None): + self._sequence += 1 + budget_id = str(data.get("budget_id") or f"budget-{self._sequence}") + row = _FakeBudgetRow( + budget_id=budget_id, + max_budget=data.get("max_budget"), + soft_budget=data.get("soft_budget"), + budget_duration=data.get("budget_duration"), + ) + self.rows[budget_id] = row + self.create_calls.append(dict(data)) + self.journal.append(f"budget_table.create:{budget_id}") + return row + + async def update(self, where, data, include=None): + budget_id = where["budget_id"] + self.update_calls.append((budget_id, dict(data))) + self.journal.append(f"budget_table.update:{budget_id}") + row = self.rows.get(budget_id) + if row is None: + return None + for field_name in ("max_budget", "soft_budget", "budget_duration"): + if data.get(field_name) is not None: + setattr(row, field_name, data[field_name]) + return row + + async def delete(self, where, include=None): + budget_id = where["budget_id"] + self.journal.append(f"budget_table.delete:{budget_id}") + self.deleted_ids.append(budget_id) + return self.rows.pop(budget_id, None) + + +class _FakeAccessGroupBudgetTable: + """Stands in for litellm_modelaccessgroupbudgettable, resolving `include` against the fake + budget table the way prisma resolves the relation.""" + + def __init__(self, journal: list[str], budget_table: _FakeBudgetTable) -> None: + self.journal = journal + self.budget_table = budget_table + self.rows: dict[str, _FakeAccessGroupBudgetRow] = {} + self.upsert_calls: list[dict] = [] + + def _resolve(self, row, include): + if row is None: + return None + row.litellm_budget_table = ( + self.budget_table.rows.get(row.budget_id) if include and row.budget_id is not None else None + ) + return row + + async def find_unique(self, where, include=None): + return self._resolve(self.rows.get(where["access_group_name"]), include) + + async def upsert(self, where, data, include=None): + access_group_name = where["access_group_name"] + self.upsert_calls.append(dict(data)) + self.journal.append(f"access_group_budget.upsert:{access_group_name}") + existing = self.rows.get(access_group_name) + payload = data["update"] if existing is not None else data["create"] + row = existing or _FakeAccessGroupBudgetRow(access_group_name=access_group_name) + row.budget_id = payload.get("budget_id") + self.rows[access_group_name] = row + return self._resolve(row, include) + + async def delete(self, where, include=None): + access_group_name = where["access_group_name"] + self.journal.append(f"access_group_budget.delete:{access_group_name}") + return self.rows.pop(access_group_name, None) + + +class _FakeModelTable: + def __init__(self, journal: list[str], deployments) -> None: + self.journal = journal + self.deployments = list(deployments) + self.updates: list[tuple[dict, dict]] = [] + + async def find_many(self, where=None, **kwargs): + return list(self.deployments) + + async def find_unique(self, where, include=None): + return next((d for d in self.deployments if d.model_id == where["model_id"]), None) + + async def update(self, where, data, include=None): + self.journal.append(f"model_table.update:{where['model_id']}") + self.updates.append((dict(where), dict(data))) + return None + + +class _FakePrismaClient: + def __init__(self, journal: list[str], deployments=()) -> None: + self.budget_table = _FakeBudgetTable(journal) + self.access_group_budget_table = _FakeAccessGroupBudgetTable(journal, self.budget_table) + self.model_table = _FakeModelTable(journal, deployments) + self.db = SimpleNamespace( + litellm_budgettable=self.budget_table, + litellm_modelaccessgroupbudgettable=self.access_group_budget_table, + litellm_proxymodeltable=self.model_table, + ) + + def jsonify_object(self, data): + return dict(data) + + +class _FakeAuthCache: + """Spy for the auth cache the endpoints evict through. Injected into the endpoint rather than + patched over the proxy_server global, so dropping the eviction call fails a test.""" + + def __init__(self, journal: list[str] | None = None) -> None: + self.journal = journal if journal is not None else [] + self.deleted_keys: list[str] = [] + + async def async_delete_cache(self, key): + self.deleted_keys.append(key) + self.journal.append(f"auth_cache.delete:{key}") + + +def _admin(): + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + + return UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + + +def _deployment(model_id="deploy-1", model_name="gpt-4o", access_groups=("prod-models",)): + return _FakeDeployment( + model_id=model_id, + model_name=model_name, + model_info={"access_groups": list(access_groups)}, + ) + + +def _seed_budget(prisma, access_group, spend=0.0, budget_id="budget-seed", **budget_fields): + prisma.budget_table.rows[budget_id] = _FakeBudgetRow(budget_id=budget_id, **budget_fields) + prisma.access_group_budget_table.rows[access_group] = _FakeAccessGroupBudgetRow( + access_group_name=access_group, + budget_id=budget_id, + spend=spend, + ) + + +@contextmanager +def _proxy(prisma): + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + yield + + +def _eviction_journal(access_group): + """Both auth cache keys, in the order a write path has to evict them.""" + from litellm.proxy.common_utils.user_api_key_cache import ( + model_access_group_cache_key, + model_access_group_registry_cache_key, + ) + + return [ + f"auth_cache.delete:{model_access_group_cache_key(access_group)}", + f"auth_cache.delete:{model_access_group_registry_cache_key()}", + ] + + +def _assert_evicted_after_write(journal, access_group, write_entry): + """Exactly the two keys, in order, after the DB write. Deliberately not a tail slice: what + has to hold is that the eviction follows the write, not that nothing follows the eviction.""" + evictions = [entry for entry in journal if entry.startswith("auth_cache.delete:")] + assert evictions == _eviction_journal(access_group) + assert journal.index(write_entry) < journal.index(evictions[0]) + + +@pytest.mark.asyncio +async def test_put_access_group_budget_creates_the_row_and_its_budget(): + """First PUT has to create both halves: the budget row it links, and the access group row + that carries the link and the shared spend.""" + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + set_access_group_budget, + ) + from litellm.types.proxy.management_endpoints.model_management_endpoints import ( + AccessGroupBudgetRequest, + ) + + prisma = _FakePrismaClient([], deployments=[_deployment()]) + cache = _FakeAuthCache() + + with _proxy(prisma): + response = await set_access_group_budget( + access_group="prod-models", + data=AccessGroupBudgetRequest(max_budget=100.0, soft_budget=80.0, budget_duration="30d"), + user_api_key_dict=_admin(), + auth_cache=cache, + ) + + assert response.access_group == "prod-models" + assert response.spend == 0.0 + assert response.budget is not None + assert response.budget.max_budget == 100.0 + assert response.budget.soft_budget == 80.0 + assert response.budget.budget_duration == "30d" + assert len(prisma.budget_table.create_calls) == 1 + assert prisma.access_group_budget_table.rows["prod-models"].budget_id == response.budget.budget_id + + +@pytest.mark.asyncio +async def test_second_put_replaces_the_budget_instead_of_creating_another(): + """PUT is idempotent: a second call must update the budget already linked to the group, + not leave a second budget row (and a second group row) behind.""" + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + set_access_group_budget, + ) + from litellm.types.proxy.management_endpoints.model_management_endpoints import ( + AccessGroupBudgetRequest, + ) + + prisma = _FakePrismaClient([], deployments=[_deployment()]) + cache = _FakeAuthCache() + + with _proxy(prisma): + first = await set_access_group_budget( + access_group="prod-models", + data=AccessGroupBudgetRequest(max_budget=100.0), + user_api_key_dict=_admin(), + auth_cache=cache, + ) + second = await set_access_group_budget( + access_group="prod-models", + data=AccessGroupBudgetRequest(max_budget=250.0), + user_api_key_dict=_admin(), + auth_cache=cache, + ) + + assert first.budget is not None and second.budget is not None + assert second.budget.budget_id == first.budget.budget_id + assert second.budget.max_budget == 250.0 + assert len(prisma.budget_table.create_calls) == 1 + assert len(prisma.budget_table.rows) == 1 + assert len(prisma.access_group_budget_table.rows) == 1 + assert prisma.budget_table.update_calls[-1][0] == first.budget.budget_id + + +@pytest.mark.asyncio +async def test_put_access_group_budget_links_an_existing_budget_without_creating_one(): + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + set_access_group_budget, + ) + from litellm.types.proxy.management_endpoints.model_management_endpoints import ( + AccessGroupBudgetRequest, + ) + + prisma = _FakePrismaClient([], deployments=[_deployment()]) + prisma.budget_table.rows["shared-budget"] = _FakeBudgetRow(budget_id="shared-budget", max_budget=7.0) + cache = _FakeAuthCache() + + with _proxy(prisma): + response = await set_access_group_budget( + access_group="prod-models", + data=AccessGroupBudgetRequest(budget_id="shared-budget"), + user_api_key_dict=_admin(), + auth_cache=cache, + ) + + assert prisma.budget_table.create_calls == [] + assert response.budget is not None + assert response.budget.budget_id == "shared-budget" + assert response.budget.max_budget == 7.0 + assert prisma.access_group_budget_table.rows["prod-models"].budget_id == "shared-budget" + + +@pytest.mark.asyncio +async def test_put_access_group_budget_rejects_an_empty_body(): + """An empty PUT would register the group as budgeted while enforcing nothing.""" + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + set_access_group_budget, + ) + from litellm.types.proxy.management_endpoints.model_management_endpoints import ( + AccessGroupBudgetRequest, + ) + + prisma = _FakePrismaClient([], deployments=[_deployment()]) + cache = _FakeAuthCache() + + with _proxy(prisma), pytest.raises(HTTPException) as exc_info: + await set_access_group_budget( + access_group="prod-models", + data=AccessGroupBudgetRequest(), + user_api_key_dict=_admin(), + auth_cache=cache, + ) + + assert exc_info.value.status_code == 400 + assert prisma.access_group_budget_table.rows == {} + assert cache.deleted_keys == [] + + +@pytest.mark.asyncio +async def test_put_access_group_budget_rejects_an_unparseable_duration(): + """An unparseable duration can only be discovered by the reset job, long after the write.""" + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + set_access_group_budget, + ) + from litellm.types.proxy.management_endpoints.model_management_endpoints import ( + AccessGroupBudgetRequest, + ) + + prisma = _FakePrismaClient([], deployments=[_deployment()]) + cache = _FakeAuthCache() + + with _proxy(prisma), pytest.raises(HTTPException) as exc_info: + await set_access_group_budget( + access_group="prod-models", + data=AccessGroupBudgetRequest(max_budget=10.0, budget_duration="every other tuesday"), + user_api_key_dict=_admin(), + auth_cache=cache, + ) + + assert exc_info.value.status_code == 400 + assert prisma.budget_table.create_calls == [] + assert prisma.access_group_budget_table.rows == {} + + +def test_access_group_budget_request_rejects_rate_limit_fields(): + """tpm/rpm/max_parallel_requests are not enforced per access group, so accepting them would + promise rate limiting that never happens.""" + from pydantic import ValidationError + + from litellm.types.proxy.management_endpoints.model_management_endpoints import ( + AccessGroupBudgetRequest, + ) + + for unsupported in ({"tpm_limit": 10}, {"rpm_limit": 10}, {"max_parallel_requests": 10}): + with pytest.raises(ValidationError): + AccessGroupBudgetRequest(max_budget=1.0, **unsupported) + + +@pytest.mark.asyncio +async def test_get_access_group_budget_returns_the_budget_and_the_shared_spend(): + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + get_access_group_budget, + ) + + prisma = _FakePrismaClient([], deployments=[_deployment()]) + _seed_budget(prisma, "prod-models", spend=42.5, max_budget=100.0, budget_duration="30d") + + with _proxy(prisma): + response = await get_access_group_budget(access_group="prod-models", user_api_key_dict=_admin()) + + assert response.access_group == "prod-models" + assert response.spend == 42.5 + assert response.budget is not None + assert response.budget.max_budget == 100.0 + assert response.budget.budget_duration == "30d" + + +@pytest.mark.asyncio +async def test_get_access_group_budget_on_a_budgetless_group_is_200_not_404(): + """A real group that simply has no budget is not an error; only an unknown group is.""" + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + get_access_group_budget, + ) + + prisma = _FakePrismaClient([], deployments=[_deployment()]) + + with _proxy(prisma): + response = await get_access_group_budget(access_group="prod-models", user_api_key_dict=_admin()) + + assert response.spend == 0.0 + assert response.budget is None + + +@pytest.mark.asyncio +async def test_access_group_budget_routes_404_on_an_unknown_group(): + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + delete_access_group_budget, + get_access_group_budget, + set_access_group_budget, + ) + from litellm.types.proxy.management_endpoints.model_management_endpoints import ( + AccessGroupBudgetRequest, + ) + + prisma = _FakePrismaClient([], deployments=[_deployment()]) + cache = _FakeAuthCache() + admin = _admin() + + calls = ( + lambda: get_access_group_budget(access_group="ghost-group", user_api_key_dict=admin), + lambda: set_access_group_budget( + access_group="ghost-group", + data=AccessGroupBudgetRequest(max_budget=1.0), + user_api_key_dict=admin, + auth_cache=cache, + ), + lambda: delete_access_group_budget( + access_group="ghost-group", user_api_key_dict=admin, auth_cache=cache + ), + ) + + with _proxy(prisma): + for make_call in calls: + with pytest.raises(HTTPException) as exc_info: + await make_call() + assert exc_info.value.status_code == 404 + + assert prisma.budget_table.create_calls == [] + assert prisma.access_group_budget_table.rows == {} + + +@pytest.mark.asyncio +async def test_delete_access_group_budget_drops_the_row_and_spares_the_shared_budget(): + """The group row goes; the LiteLLM_BudgetTable row it linked survives, as /tag/delete leaves a + tag's. That row can be shared, so deleting it would be data loss for whatever else points at it.""" + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + delete_access_group_budget, + ) + + prisma = _FakePrismaClient([], deployments=[_deployment()]) + _seed_budget(prisma, "prod-models", spend=12.0, max_budget=100.0) + cache = _FakeAuthCache() + + with _proxy(prisma): + response = await delete_access_group_budget( + access_group="prod-models", user_api_key_dict=_admin(), auth_cache=cache + ) + + assert response.budget_deleted is True + assert prisma.access_group_budget_table.rows == {} + assert prisma.budget_table.deleted_ids == [] + assert prisma.budget_table.rows["budget-seed"].max_budget == 100.0 + + +@pytest.mark.asyncio +async def test_delete_access_group_budget_on_a_budgetless_group_still_evicts(): + """budget_deleted is False, but the group can still be sitting in the cached registry of + budgeted groups, so the eviction has to run whether or not a row was there to drop.""" + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + delete_access_group_budget, + ) + + journal: list[str] = [] + prisma = _FakePrismaClient(journal, deployments=[_deployment()]) + cache = _FakeAuthCache(journal) + + with _proxy(prisma): + response = await delete_access_group_budget( + access_group="prod-models", user_api_key_dict=_admin(), auth_cache=cache + ) + + assert response.budget_deleted is False + assert prisma.budget_table.deleted_ids == [] + _assert_evicted_after_write(journal, "prod-models", "access_group_budget.delete:prod-models") + + +@pytest.mark.asyncio +async def test_deleting_the_access_group_strips_deployments_before_dropping_the_budget(): + """Ordering is the point: stripping first means a failure leaves an unreachable budget row, + while the reverse leaves a live group whose enforcement silently vanished. The shared + LiteLLM_BudgetTable row survives here too.""" + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + delete_access_group, + ) + + journal: list[str] = [] + prisma = _FakePrismaClient(journal, deployments=[_deployment()]) + _seed_budget(prisma, "prod-models", spend=3.0, max_budget=100.0) + cache = _FakeAuthCache() + + never_served_router = MagicMock() + never_served_router.get_model_ids.return_value = [] + with ( + _proxy(prisma), + patch("litellm.proxy.proxy_server.llm_router", never_served_router), + patch( + "litellm.proxy.management_endpoints.model_access_group_management_endpoints.clear_cache", + new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)), + ), + ): + response = await delete_access_group( + access_group="prod-models", user_api_key_dict=_admin(), auth_cache=cache + ) + + assert response.models_updated == 1 + assert prisma.access_group_budget_table.rows == {} + assert prisma.budget_table.deleted_ids == [] + assert prisma.budget_table.rows["budget-seed"].max_budget == 100.0 + assert journal.index("model_table.update:deploy-1") < journal.index("access_group_budget.delete:prod-models") + + +@pytest.mark.asyncio +async def test_access_group_info_surfaces_the_budget_and_spend(): + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + get_access_group_info, + ) + + prisma = _FakePrismaClient([], deployments=[_deployment()]) + _seed_budget(prisma, "prod-models", spend=9.5, max_budget=100.0, soft_budget=50.0) + + with _proxy(prisma): + info = await get_access_group_info(access_group="prod-models", user_api_key_dict=_admin()) + + assert info.model_names == ["gpt-4o"] + assert info.spend == 9.5 + assert info.budget is not None + assert info.budget.max_budget == 100.0 + assert info.budget.soft_budget == 50.0 + + +@pytest.mark.asyncio +async def test_put_access_group_budget_evicts_both_auth_cache_keys(): + """Auth reads the per-group row and the registry of budgeted groups cache-first with no + freshness check, so a PUT that skips either eviction returns 200 and enforces nothing until + the TTL expires. Both keys, after the write.""" + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + set_access_group_budget, + ) + from litellm.types.proxy.management_endpoints.model_management_endpoints import ( + AccessGroupBudgetRequest, + ) + + journal: list[str] = [] + prisma = _FakePrismaClient(journal, deployments=[_deployment()]) + cache = _FakeAuthCache(journal) + + with _proxy(prisma): + await set_access_group_budget( + access_group="prod-models", + data=AccessGroupBudgetRequest(max_budget=100.0), + user_api_key_dict=_admin(), + auth_cache=cache, + ) + + _assert_evicted_after_write(journal, "prod-models", "access_group_budget.upsert:prod-models") + + +@pytest.mark.asyncio +async def test_delete_access_group_budget_evicts_both_auth_cache_keys(): + """Clearing a budget has the same window as setting one: until both keys are dropped, auth + keeps enforcing the budget that is already gone.""" + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + delete_access_group_budget, + ) + + journal: list[str] = [] + prisma = _FakePrismaClient(journal, deployments=[_deployment()]) + _seed_budget(prisma, "prod-models", spend=12.0, max_budget=100.0) + cache = _FakeAuthCache(journal) + + with _proxy(prisma): + await delete_access_group_budget( + access_group="prod-models", user_api_key_dict=_admin(), auth_cache=cache + ) + + _assert_evicted_after_write(journal, "prod-models", "access_group_budget.delete:prod-models") + + +@pytest.mark.asyncio +async def test_deleting_the_access_group_evicts_both_auth_cache_keys(): + """The group-delete cascade drops the budget row too, so it owes the same two evictions.""" + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + delete_access_group, + ) + + journal: list[str] = [] + prisma = _FakePrismaClient(journal, deployments=[_deployment()]) + _seed_budget(prisma, "prod-models", spend=3.0, max_budget=100.0) + cache = _FakeAuthCache(journal) + + never_served_router = MagicMock() + never_served_router.get_model_ids.return_value = [] + with ( + _proxy(prisma), + patch("litellm.proxy.proxy_server.llm_router", never_served_router), + patch( + "litellm.proxy.management_endpoints.model_access_group_management_endpoints.clear_cache", + new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)), + ), + ): + await delete_access_group(access_group="prod-models", user_api_key_dict=_admin(), auth_cache=cache) + + _assert_evicted_after_write(journal, "prod-models", "access_group_budget.delete:prod-models") + + +@pytest.mark.asyncio +async def test_deleting_an_access_group_that_never_had_a_budget_still_evicts(): + """The cascade's delete finds no row and reports nothing dropped, but the group can still be + sitting in the cached registry of budgeted groups, so both keys have to go regardless.""" + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + delete_access_group, + ) + + journal: list[str] = [] + prisma = _FakePrismaClient(journal, deployments=[_deployment()]) + cache = _FakeAuthCache(journal) + + never_served_router = MagicMock() + never_served_router.get_model_ids.return_value = [] + with ( + _proxy(prisma), + patch("litellm.proxy.proxy_server.llm_router", never_served_router), + patch( + "litellm.proxy.management_endpoints.model_access_group_management_endpoints.clear_cache", + new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)), + ), + ): + response = await delete_access_group( + access_group="prod-models", user_api_key_dict=_admin(), auth_cache=cache + ) + + assert response.models_updated == 1 + assert prisma.access_group_budget_table.rows == {} + _assert_evicted_after_write(journal, "prod-models", "access_group_budget.delete:prod-models") diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 38a346e7fb7..1b28566366c 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -2,6 +2,7 @@ import asyncio import threading from collections.abc import Mapping from datetime import datetime, timedelta, timezone +from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -10,15 +11,16 @@ from fastapi import HTTPException import litellm from litellm.caching.dual_cache import DualCache from litellm.constants import STREAM_SSE_KEEPALIVE_PING_BYTES -from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( - AnthropicMessagesStreamingResponse, -) from litellm.llms.anthropic.experimental_pass_through.messages.agentic_streaming_iterator import ( AgenticAnthropicStreamingIterator, ) +from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( + AnthropicMessagesStreamingResponse, +) from litellm.proxy._types import ( LiteLLM_BudgetTable, LiteLLM_EndUserTable, + Litellm_EntityType, LiteLLM_OrganizationTable, LiteLLM_TagTable, LiteLLM_TeamMembership, @@ -27,9 +29,15 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.proxy.common_utils.reset_budget_job import _model_access_group_counter_key +from litellm.proxy.common_utils.user_api_key_cache import ( + UserApiKeyCache, + model_access_group_cache_key, +) from litellm.proxy.spend_tracking.budget_reservation import ( TOKENIZE_OFF_EVENT_LOOP_MIN_CHARS, _approximate_input_size, + _get_model_access_group_budget_counters, estimate_request_max_cost, get_budget_window_start, invalidate_budget_reservation_counters, @@ -2962,3 +2970,100 @@ async def test_small_prompt_is_tokenized_inline(spend_counter_state): assert reservation is not None assert threads == [threading.main_thread()] + + +class _ModelAccessGroupBudgetPrisma: + """Serves ``LiteLLM_ModelAccessGroupBudgetTable`` rows, recording what reached the database.""" + + def __init__(self, **max_budget_by_group) -> None: + self.rows = { + group: SimpleNamespace( + access_group_name=group, + spend=7.0, + litellm_budget_table=None if max_budget is None else SimpleNamespace(max_budget=max_budget), + ) + for group, max_budget in max_budget_by_group.items() + } + self.batches = [] + self.db = SimpleNamespace( + litellm_modelaccessgroupbudgettable=SimpleNamespace(find_many=self._find_many) + ) + + async def _find_many(self, **kwargs): + requested = list(kwargs["where"]["access_group_name"]["in"]) + self.batches.append(requested) + return [self.rows[group] for group in requested if group in self.rows] + + +async def _model_access_group_counters(matched, **max_budget_by_group): + return await _get_model_access_group_budget_counters( + valid_token=UserAPIKeyAuth(api_key="hashed", matched_model_access_groups=matched), + prisma_client=_ModelAccessGroupBudgetPrisma(**max_budget_by_group), + user_api_key_cache=UserApiKeyCache(), + ) + + +@pytest.mark.asyncio +async def test_model_access_group_with_a_budget_reserves_against_the_reset_jobs_counter_key(): + counters = await _model_access_group_counters(["premium"], premium=25.0) + + assert len(counters) == 1 + counter = counters[0] + assert counter.counter_key == _model_access_group_counter_key(SimpleNamespace(access_group_name="premium")) + assert counter.source_cache_key == model_access_group_cache_key("premium") + assert counter.max_budget == 25.0 + assert counter.fallback_spend == 7.0 + assert counter.entity_type == "Model access group" + assert counter.entity_id == "premium" + + +@pytest.mark.asyncio +async def test_model_access_group_without_a_budget_reserves_nothing(): + assert await _model_access_group_counters(["premium"], premium=None) == [] + + +@pytest.mark.asyncio +async def test_model_access_group_with_a_zero_budget_reserves_nothing(): + """Zero is how a budget is cleared, not a ceiling that blocks every request.""" + assert await _model_access_group_counters(["premium"], premium=0.0) == [] + + +@pytest.mark.asyncio +async def test_model_access_group_counters_come_from_the_auth_object(): + """Auth already resolved which granted groups serve the model; re-deriving it here would drift.""" + assert await _model_access_group_counters(None, premium=25.0) == [] + + +@pytest.mark.asyncio +async def test_repeated_model_access_group_reserves_once(): + counters = await _model_access_group_counters(["premium", "premium"], premium=25.0) + + assert [counter.entity_id for counter in counters] == ["premium"] + + +@pytest.mark.asyncio +async def test_model_access_group_counter_blocks_a_request_over_the_group_budget(spend_counter_state): + """End to end through the reservation path, which is what runs when reservations are enabled.""" + counter_cache, key_cache = spend_counter_state + prisma_client = _ModelAccessGroupBudgetPrisma(premium=1.0) + valid_token = UserAPIKeyAuth(api_key="hashed", token="tok", matched_model_access_groups=["premium"]) + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=0.5, + ): + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=None, + user_object=None, + prisma_client=prisma_client, + user_api_key_cache=key_cache, + proxy_logging_obj=ProxyLogging(user_api_key_cache=key_cache), + ) + + assert exc_info.value.entity_id == "premium" + assert exc_info.value.entity_type == Litellm_EntityType.MODEL_ACCESS_GROUP.value diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index 9fb31d2a6db..260db842370 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -33,6 +33,8 @@ from litellm.proxy.litellm_pre_call_utils import ( check_if_token_is_service_account, clean_headers, ) +from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs +from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY from litellm.litellm_core_utils.get_provider_specific_headers import ( ProviderSpecificHeaderUtils, ) @@ -7681,3 +7683,37 @@ async def test_add_litellm_data_to_request_keeps_litellm_metadata_on_litellm_met ) assert updated["litellm_metadata"]["trace_id"] == "abc" + + +def _stamp_model_access_groups(matched_model_access_groups, metadata_variable_name="metadata"): + user_api_key_dict = UserAPIKeyAuth(api_key="hashed-key") + user_api_key_dict.matched_model_access_groups = matched_model_access_groups + return LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata( + data={metadata_variable_name: {}}, + user_api_key_dict=user_api_key_dict, + _metadata_variable_name=metadata_variable_name, + )[metadata_variable_name] + + +def test_matched_model_access_groups_are_stamped_into_request_metadata(): + """The post-call spend writer reads the groups off request metadata, not off UserAPIKeyAuth.""" + stamped = _stamp_model_access_groups(["tier-a", "tier-b"]) + + assert stamped[MODEL_ACCESS_GROUP_METADATA_KEY] == ["tier-a", "tier-b"] + assert MODEL_ACCESS_GROUP_METADATA_KEY not in _stamp_model_access_groups(None) + + +def test_stamped_model_access_groups_survive_the_litellm_metadata_merge(): + """ + The key must keep its ``user_api_key`` prefix: when a request carries both metadata dicts, + get_litellm_metadata_from_kwargs returns litellm_metadata and copies a key over from metadata + only when that substring is in its name, so an unprefixed key is silently dropped. + """ + kwargs = { + "litellm_params": { + "metadata": _stamp_model_access_groups(["tier-a"]), + "litellm_metadata": {"trace_id": "abc"}, + } + } + + assert get_litellm_metadata_from_kwargs(kwargs)[MODEL_ACCESS_GROUP_METADATA_KEY] == ["tier-a"] diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index eaa05ddc005..7b23b03fb3b 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -497,6 +497,91 @@ export interface paths { patch?: never; trace?: never; }; + "/access_group/{access_group}/budget": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Get Access Group Budget + * @description Get the shared budget of an access group, and the spend drawn against it. + * + * Example: + * ```bash + * curl -X GET 'http://localhost:4000/access_group/production-models/budget' \ + * -H 'Authorization: Bearer sk-1234' + * ``` + * + * Parameters: + * - access_group: str - The access group name (URL path parameter) + * + * Returns: + * - AccessGroupBudgetResponse; budget is null when the group has no budget set + * + * Raises: + * - HTTPException 404: If access group not found + */ + get: operations["get_access_group_budget_access_group__access_group__budget_get"]; + /** + * Set Access Group Budget + * @description Set or replace the shared budget of an access group. Idempotent. + * + * Every key that can reach a model in the group draws from this one budget. + * + * Example: + * ```bash + * curl -X PUT 'http://localhost:4000/access_group/production-models/budget' \ + * -H 'Authorization: Bearer sk-1234' \ + * -H 'Content-Type: application/json' \ + * -d '{ + * "max_budget": 100.0, + * "budget_duration": "30d" + * }' + * ``` + * + * Parameters: + * - access_group: str - The access group name (URL path parameter) + * - max_budget: Optional[float] - Requests fail once the group's shared spend exceeds this + * - soft_budget: Optional[float] - Fires an alert when reached; requests still succeed + * - budget_duration: Optional[str] - Frequency of resetting the group's spend (e.g. '30d') + * - budget_id: Optional[str] - Link an existing budget instead of creating one + * + * Returns: + * - AccessGroupBudgetResponse with the stored budget and current spend + * + * Raises: + * - HTTPException 400: If no budget field is given, or budget_duration cannot be parsed + * - HTTPException 404: If access group not found + */ + put: operations["set_access_group_budget_access_group__access_group__budget_put"]; + post?: never; + /** + * Delete Access Group Budget + * @description Clear the shared budget of an access group, leaving the group itself in place. + * + * Example: + * ```bash + * curl -X DELETE 'http://localhost:4000/access_group/production-models/budget' \ + * -H 'Authorization: Bearer sk-1234' + * ``` + * + * Parameters: + * - access_group: str - The access group name (URL path parameter) + * + * Returns: + * - DeleteAccessGroupBudgetResponse; budget_deleted is false when there was nothing to clear + * + * Raises: + * - HTTPException 404: If access group not found + */ + delete: operations["delete_access_group_budget_access_group__access_group__budget_delete"]; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/access_group/{access_group}/delete": { parameters: { query?: never; @@ -555,7 +640,7 @@ export interface paths { * - access_group: str - The access group name (URL path parameter) * * Returns: - * - AccessGroupInfo with the access group details + * - AccessGroupInfo with the access group details, its shared budget and its spend * * Raises: * - HTTPException 404: If access group not found @@ -22182,6 +22267,38 @@ export interface components { */ type: "restricted_sso_group"; }; + /** AccessGroupBudget */ + AccessGroupBudget: { + /** Budget Duration */ + budget_duration?: string | null; + /** Budget Id */ + budget_id: string; + /** Budget Reset At */ + budget_reset_at?: string | null; + /** Max Budget */ + max_budget?: number | null; + /** Soft Budget */ + soft_budget?: number | null; + }; + /** AccessGroupBudgetRequest */ + AccessGroupBudgetRequest: { + /** Budget Duration */ + budget_duration?: string | null; + /** Budget Id */ + budget_id?: string | null; + /** Max Budget */ + max_budget?: number | null; + /** Soft Budget */ + soft_budget?: number | null; + }; + /** AccessGroupBudgetResponse */ + AccessGroupBudgetResponse: { + /** Access Group */ + access_group: string; + budget?: components["schemas"]["AccessGroupBudget"] | null; + /** Spend */ + spend: number; + }; /** AccessGroupCreateRequest */ AccessGroupCreateRequest: { /** Access Agent Ids */ @@ -22203,10 +22320,13 @@ export interface components { AccessGroupInfo: { /** Access Group */ access_group: string; + budget?: components["schemas"]["AccessGroupBudget"] | null; /** Deployment Count */ deployment_count: number; /** Model Names */ model_names: string[]; + /** Spend */ + spend?: number | null; }; /** AccessGroupResponse */ AccessGroupResponse: { @@ -26004,6 +26124,15 @@ export interface components { [key: string]: unknown; }; }; + /** DeleteAccessGroupBudgetResponse */ + DeleteAccessGroupBudgetResponse: { + /** Access Group */ + access_group: string; + /** Budget Deleted */ + budget_deleted: boolean; + /** Message */ + message: string; + }; /** * DeleteCustomerRequest * @description Delete multiple Customers @@ -39024,6 +39153,103 @@ export interface operations { }; }; }; + get_access_group_budget_access_group__access_group__budget_get: { + parameters: { + query?: never; + header?: never; + path: { + access_group: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["AccessGroupBudgetResponse"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + set_access_group_budget_access_group__access_group__budget_put: { + parameters: { + query?: never; + header?: never; + path: { + access_group: string; + }; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["AccessGroupBudgetRequest"]; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["AccessGroupBudgetResponse"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + delete_access_group_budget_access_group__access_group__budget_delete: { + parameters: { + query?: never; + header?: never; + path: { + access_group: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["DeleteAccessGroupBudgetResponse"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; delete_access_group_access_group__access_group__delete_delete: { parameters: { query?: never; From 183a782e57111d5656d49f600437f43588c7a39b Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" <41898282+github-actions[bot]@users.noreply.github.com> Date: Sat, 29 Aug 2026 19:15:11 +0000 Subject: [PATCH 03/16] chore: sync schema.prisma copies from root --- .../litellm_proxy_extras/schema.prisma | 15 +++++++++++++++ litellm/proxy/schema.prisma | 15 +++++++++++++++ 2 files changed, 30 insertions(+) diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 2bb850139a2..48d41edffbd 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -29,6 +29,7 @@ model LiteLLM_BudgetTable { keys LiteLLM_VerificationToken[] // multiple keys can have the same budget end_users LiteLLM_EndUserTable[] // multiple end-users can have the same budget tags LiteLLM_TagTable[] // multiple tags can have the same budget + model_access_groups LiteLLM_ModelAccessGroupBudgetTable[] // multiple model access groups can have the same budget team_membership LiteLLM_TeamMembership[] // budgets of Users within a Team organization_membership LiteLLM_OrganizationMembership[] // budgets of Users within a Organization } @@ -585,6 +586,20 @@ model LiteLLM_EndUserTable { blocked Boolean @default(false) } +// Budget and shared spend for a model access group. The groups themselves are not rows anywhere: +// they are free-text strings in LiteLLM_ProxyModelTable.model_info.access_groups, so a row here +// exists only once someone gives that group a budget. +model LiteLLM_ModelAccessGroupBudgetTable { + access_group_name String @id + spend Float @default(0.0) + budget_id String? + litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) + created_at DateTime @default(now()) @map("created_at") + created_by String? + updated_at DateTime @default(now()) @updatedAt + updated_by String? +} + // Track tags with budgets and spend model LiteLLM_TagTable { tag_name String @id diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 2bb850139a2..48d41edffbd 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -29,6 +29,7 @@ model LiteLLM_BudgetTable { keys LiteLLM_VerificationToken[] // multiple keys can have the same budget end_users LiteLLM_EndUserTable[] // multiple end-users can have the same budget tags LiteLLM_TagTable[] // multiple tags can have the same budget + model_access_groups LiteLLM_ModelAccessGroupBudgetTable[] // multiple model access groups can have the same budget team_membership LiteLLM_TeamMembership[] // budgets of Users within a Team organization_membership LiteLLM_OrganizationMembership[] // budgets of Users within a Organization } @@ -585,6 +586,20 @@ model LiteLLM_EndUserTable { blocked Boolean @default(false) } +// Budget and shared spend for a model access group. The groups themselves are not rows anywhere: +// they are free-text strings in LiteLLM_ProxyModelTable.model_info.access_groups, so a row here +// exists only once someone gives that group a budget. +model LiteLLM_ModelAccessGroupBudgetTable { + access_group_name String @id + spend Float @default(0.0) + budget_id String? + litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) + created_at DateTime @default(now()) @map("created_at") + created_by String? + updated_at DateTime @default(now()) @updatedAt + updated_by String? +} + // Track tags with budgets and spend model LiteLLM_TagTable { tag_name String @id From 8e1d1f1ef0daf20dc227cf51ce02a44eb0621207 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 29 Aug 2026 12:49:31 -0700 Subject: [PATCH 04/16] fix(budgets): write the model access group spend counter after each call Two problems, both caught in review. The new table only landed in the root schema.prisma. Client generation reads litellm/proxy/schema.prisma and packaging reads the copy under litellm-proxy-extras, so the generated client had no litellm_modelaccessgroupbudgettable and every budget read and write against it would have failed at runtime. The root is the source of truth; both copies are now byte-identical to it. Nothing incremented spend:model_access_group:{group} after a call. Only the reservation path ever wrote it, so with disable_budget_reservation the read-time check was reading a counter nobody maintained and falling back to the row's spend, which is cached for up to DEFAULT_MODEL_ACCESS_GROUP_CACHE_TTL. A caller could run well past the pool inside that window, which is precisely the case the read-time check exists to cover. increment_spend_counters now takes the matched groups and charges them through _init_and_increment_unreserved_spend_counter, so a group already covered by a reservation is skipped rather than counted twice. The cost callback sources the names with get_request_model_access_groups, the same reader the spend writer uses. --- .../proxy/common_utils/user_api_key_cache.py | 9 +- .../proxy/hooks/proxy_track_cost_callback.py | 6 + litellm/proxy/proxy_server.py | 37 ++++++ .../hooks/test_proxy_track_cost_callback.py | 89 +++++++++++++- .../proxy/test_budget_reservation.py | 116 ++++++++++++++++++ 5 files changed, 249 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index ae0b39135f1..589d8fe68d1 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -201,11 +201,12 @@ def model_access_group_registry_cache_key() -> str: def model_access_group_spend_counter_key(access_group_name: str) -> str: - """Spend counter key for one model access group; shared so its three owners cannot drift. + """Spend counter key for one model access group; shared so its four owners cannot drift. - The reservation path writes it, auth reads it to enforce ``max_budget``, and the reset job - clears it on rollover. A copy that drifts in any one of them silently resets or reads a - counter nobody else touches, which shows up as a budget that never trips or never resets. + The reservation path writes it up front, the cost callback writes it after the call, auth + reads it to enforce ``max_budget``, and the reset job clears it on rollover. A copy that + drifts in any one of them silently resets or reads a counter nobody else touches, which shows + up as a budget that never trips or never resets. """ return f"spend:model_access_group:{access_group_name}" diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index f593e94b36f..ad857b8c3f5 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 Sequence from datetime import datetime from typing import TYPE_CHECKING, Any, Final, cast @@ -27,6 +28,7 @@ from litellm.proxy.spend_tracking.spend_log_error_logger import ( ) from litellm.proxy.spend_tracking.spend_tracking_utils import ( _sanitize_error_information_for_spend_logs, + get_request_model_access_groups, ) from litellm.proxy.utils import ProxyUpdateSpend from litellm.types.utils import ( @@ -258,6 +260,7 @@ class _ProxyDBLogger(CustomLogger): sl_object=sl_object, metadata=metadata, ) + model_access_groups: Final = get_request_model_access_groups(kwargs) if response_cost is not None: user_api_key: Final = metadata.get("user_api_key", None) @@ -296,6 +299,7 @@ class _ProxyDBLogger(CustomLogger): response_cost=response_cost, budget_reservation=budget_reservation, request_tags=tags, + model_access_groups=model_access_groups, ) # update cache (fire-and-forget for backward compat: @@ -572,6 +576,7 @@ async def _update_database_and_spend_counters( response_cost: float, budget_reservation: dict | None, request_tags: list[str] | None = None, + model_access_groups: Sequence[str] | None = None, ) -> None: try: await proxy_logging_obj.db_spend_update_writer.update_database( @@ -610,6 +615,7 @@ async def _update_database_and_spend_counters( budget_reservation=budget_reservation, end_user_id=end_user_id, tags=request_tags, + model_access_groups=model_access_groups, ) except Exception: if budget_reservation is not None: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index af0aa9743bc..accf03b33ab 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -382,6 +382,8 @@ from litellm.proxy.common_utils.user_api_key_cache import ( UserApiKeyCache, end_user_cache_key, get_management_object_ttl, + model_access_group_cache_key, + model_access_group_spend_counter_key, tag_cache_key, ) from litellm.proxy.config_resolvers import resolve_fields @@ -2648,6 +2650,7 @@ async def increment_spend_counters( budget_reservation: dict | None = None, end_user_id: str | None = None, tags: list[str] | None = None, + model_access_groups: Sequence[str] | None = None, ): """ Atomically increment spend counters for budget enforcement. @@ -2777,6 +2780,13 @@ async def increment_spend_counters( ) if end_user_id is not None or tags is not None else None, + _increment_model_access_group_spend_counters( + model_access_groups=model_access_groups, + response_cost=cost, + reserved_counter_keys=reserved_counter_keys, + ) + if model_access_groups + else None, _increment_org_spend_counter( org_id=org_id, response_cost=cost, @@ -2865,6 +2875,33 @@ async def _increment_end_user_and_tag_spend_counters( ) +async def _increment_model_access_group_spend_counters( + model_access_groups: Sequence[object], + response_cost: float, + reserved_counter_keys: set[str], +) -> None: + """Charge the model access groups that authorized this request. + + Without this the counter auth reads is written only by the reservation path, so + ``disable_budget_reservation`` would leave ``_model_access_group_max_budget_check`` enforcing + against the DB row's spend, which lags by up to the cache TTL. + + Typed ``object`` rather than ``str`` because the names reach the cost callback out of request + metadata, which the coercion upstream filters to a list but not to strings. A non-string that + slipped through would build a counter key nothing else ever reads. + """ + unique_groups: Final = tuple( + dict.fromkeys(group for group in model_access_groups if group and isinstance(group, str)) + ) + for group in unique_groups: + await _init_and_increment_unreserved_spend_counter( + counter_key=model_access_group_spend_counter_key(group), + source_cache_key=model_access_group_cache_key(group), + increment=response_cost, + reserved_counter_keys=reserved_counter_keys, + ) + + async def _increment_org_spend_counter( org_id: str | None, response_cost: float, 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 ca517474a5c..e923a28f371 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 @@ -1,14 +1,14 @@ -import pytest - - from datetime import datetime from unittest.mock import AsyncMock, MagicMock, patch +import pytest + +from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.hooks.proxy_track_cost_callback import ( - _ProxyDBLogger, _get_budget_reservation_from_metadata, + _ProxyDBLogger, _should_track_cost_callback, _update_database_and_spend_counters, ) @@ -586,6 +586,7 @@ async def test_update_database_and_spend_counters_updates_counters_after_db_upda response_cost=0.2, budget_reservation=budget_reservation, request_tags=["tag-a"], + model_access_groups=("premium",), ) proxy_logging_obj.db_spend_update_writer.update_database.assert_awaited_once() @@ -598,6 +599,7 @@ 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"], + model_access_groups=("premium",), ) @@ -1875,3 +1877,82 @@ async def test_track_cost_callback_logs_unauthenticated_pass_through_request( assert mock_proxy_logging.db_spend_update_writer.update_database.await_count == ( 1 if expect_spend_log else 0 ) + + +@pytest.mark.asyncio +async def test_track_cost_callback_charges_the_model_access_groups_auth_stamped(): + """Auth stamps the matched groups onto request metadata; the callback has to carry them through. + + Without this hop nothing writes ``spend:model_access_group:*`` on the normal path, so with + reservations disabled the budget check reads a counter no one maintains. + """ + logger = _ProxyDBLogger() + kwargs = { + "model": "gpt-4", + "call_type": "acompletion", + "litellm_params": { + "metadata": { + "user_api_key": "hashed-key", + "user_api_key_user_id": "user-1", + MODEL_ACCESS_GROUP_METADATA_KEY: ["premium", "starter"], + }, + }, + "standard_logging_object": {"response_cost": 0.25, "request_tags": None}, + "stream": False, + } + + with ( + patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging, + patch( + "litellm.proxy.proxy_server.increment_spend_counters", new_callable=AsyncMock + ) as mock_increment_spend_counters, + patch("litellm.proxy.proxy_server.update_cache", new_callable=AsyncMock), + ): + mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() + mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock() + await logger._PROXY_track_cost_callback( + kwargs=kwargs, + completion_response=None, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + mock_increment_spend_counters.assert_awaited_once() + assert mock_increment_spend_counters.await_args.kwargs["model_access_groups"] == ( + "premium", + "starter", + ) + + +@pytest.mark.asyncio +async def test_track_cost_callback_charges_no_model_access_group_when_none_were_stamped(): + """A request no budgeted group authorized must not debit anything.""" + logger = _ProxyDBLogger() + kwargs = { + "model": "gpt-4", + "call_type": "acompletion", + "litellm_params": { + "metadata": {"user_api_key": "hashed-key", "user_api_key_user_id": "user-1"}, + }, + "standard_logging_object": {"response_cost": 0.25, "request_tags": None}, + "stream": False, + } + + with ( + patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging, + patch( + "litellm.proxy.proxy_server.increment_spend_counters", new_callable=AsyncMock + ) as mock_increment_spend_counters, + patch("litellm.proxy.proxy_server.update_cache", new_callable=AsyncMock), + ): + mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() + mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock() + await logger._PROXY_track_cost_callback( + kwargs=kwargs, + completion_response=None, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + mock_increment_spend_counters.assert_awaited_once() + assert mock_increment_spend_counters.await_args.kwargs["model_access_groups"] == () diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 1b28566366c..65c8c248a27 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -33,6 +33,7 @@ from litellm.proxy.common_utils.reset_budget_job import _model_access_group_coun from litellm.proxy.common_utils.user_api_key_cache import ( UserApiKeyCache, model_access_group_cache_key, + model_access_group_spend_counter_key, ) from litellm.proxy.spend_tracking.budget_reservation import ( TOKENIZE_OFF_EVENT_LOOP_MIN_CHARS, @@ -47,6 +48,7 @@ from litellm.proxy.spend_tracking.budget_reservation import ( ) from litellm.proxy.utils import ProxyLogging from litellm.router import Router +from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget @pytest.fixture() @@ -3067,3 +3069,117 @@ async def test_model_access_group_counter_blocks_a_request_over_the_group_budget assert exc_info.value.entity_id == "premium" assert exc_info.value.entity_type == Litellm_EntityType.MODEL_ACCESS_GROUP.value + + +async def _cache_model_access_group_budget(key_cache, group, spend, max_budget=None): + await key_cache.async_set_cache( + key=model_access_group_cache_key(group), + value=ModelAccessGroupBudget(access_group_name=group, spend=spend, max_budget=max_budget), + model_type=ModelAccessGroupBudget, + ) + + +async def _reserve_for_model_access_groups(key_cache, groups, estimate): + """Reserve against the given groups, whose rows are already cached, so nothing hits the DB.""" + with patch( + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=estimate, + ): + return await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=UserAPIKeyAuth( + api_key="hashed", token="tok-mag-counter", matched_model_access_groups=list(groups) + ), + team_object=None, + user_object=None, + prisma_client=_ModelAccessGroupBudgetPrisma(), + user_api_key_cache=key_cache, + proxy_logging_obj=ProxyLogging(user_api_key_cache=key_cache), + ) + + +@pytest.mark.asyncio +async def test_model_access_group_counter_accumulates_across_calls_without_a_reservation(spend_counter_state): + """With reservations disabled nothing writes the counter up front, so the cost callback must. + + Otherwise the read-time budget check enforces against the DB row's spend, which the cache + holds for the full TTL, and a caller runs past the ceiling for that whole window. + """ + counter_cache, key_cache = spend_counter_state + await _cache_model_access_group_budget(key_cache, "premium", spend=1.0, max_budget=25.0) + + from litellm.proxy.proxy_server import increment_spend_counters + + counter_key = model_access_group_spend_counter_key("premium") + + await increment_spend_counters( + token=None, team_id=None, user_id=None, response_cost=0.25, model_access_groups=["premium"] + ) + assert counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(1.25) + + await increment_spend_counters( + token=None, team_id=None, user_id=None, response_cost=0.75, model_access_groups=["premium", "premium", ""] + ) + assert counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(2.0) + assert counter_cache.in_memory_cache.get_cache(key=model_access_group_spend_counter_key("")) is None + + +@pytest.mark.asyncio +async def test_reserved_model_access_group_is_not_charged_twice(spend_counter_state): + """The reservation already wrote this counter, so the post-call pass has to skip it.""" + counter_cache, key_cache = spend_counter_state + await _cache_model_access_group_budget(key_cache, "premium", spend=1.0, max_budget=25.0) + + reservation = await _reserve_for_model_access_groups(key_cache, ["premium"], estimate=0.6) + counter_key = model_access_group_spend_counter_key("premium") + assert counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(1.6) + + from litellm.proxy.proxy_server import increment_spend_counters + + await increment_spend_counters( + token=None, + team_id=None, + user_id=None, + response_cost=0.2, + budget_reservation=reservation, + model_access_groups=["premium"], + ) + + # 1.0 recorded + the reservation reconciled down to the 0.2 actually spent. A second + # increment would land at 1.4. + assert counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(1.2) + + +@pytest.mark.asyncio +async def test_unreserved_model_access_group_is_charged_alongside_a_reserved_one(spend_counter_state): + """A budgetless group reserves nothing, so only the post-call pass can charge it. + + Both groups authorized the request and both get debited, each exactly once, whether or not + the reservation path happened to hold a counter for them. + """ + counter_cache, key_cache = spend_counter_state + await _cache_model_access_group_budget(key_cache, "premium", spend=1.0, max_budget=25.0) + await _cache_model_access_group_budget(key_cache, "starter", spend=4.0) + + reservation = await _reserve_for_model_access_groups(key_cache, ["premium", "starter"], estimate=0.6) + assert [entry["entity_id"] for entry in reservation["entries"]] == ["premium"] + + from litellm.proxy.proxy_server import increment_spend_counters + + await increment_spend_counters( + token=None, + team_id=None, + user_id=None, + response_cost=0.2, + budget_reservation=reservation, + model_access_groups=["premium", "starter", "starter", "premium"], + ) + + assert counter_cache.in_memory_cache.get_cache( + key=model_access_group_spend_counter_key("premium") + ) == pytest.approx(1.2) + assert counter_cache.in_memory_cache.get_cache( + key=model_access_group_spend_counter_key("starter") + ) == pytest.approx(4.2) From fb9e76fab3c325ad36dfa6f7a5ff956a0efda462 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 29 Aug 2026 14:00:12 -0700 Subject: [PATCH 05/16] chore(budgets): refresh the lazy OpenAPI snapshot and allowlist the model access group budget routes --- litellm/proxy/_lazy_openapi_snapshot.json | 335 +++++++++++++++++- .../endpointaudit/coverage_allowlist.txt | 3 + 2 files changed, 337 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index c1e89f8aa75..37608bc3d4b 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -461,6 +461,145 @@ "access_groups": { "components": { "schemas": { + "AccessGroupBudget": { + "properties": { + "budget_duration": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Budget Duration" + }, + "budget_id": { + "title": "Budget Id", + "type": "string" + }, + "budget_reset_at": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Budget Reset At" + }, + "max_budget": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Max Budget" + }, + "soft_budget": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Soft Budget" + } + }, + "required": [ + "budget_id" + ], + "title": "AccessGroupBudget", + "type": "object" + }, + "AccessGroupBudgetRequest": { + "additionalProperties": false, + "properties": { + "budget_duration": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Budget Duration" + }, + "budget_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Budget Id" + }, + "max_budget": { + "anyOf": [ + { + "minimum": 0.0, + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Max Budget" + }, + "soft_budget": { + "anyOf": [ + { + "minimum": 0.0, + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Soft Budget" + } + }, + "title": "AccessGroupBudgetRequest", + "type": "object" + }, + "AccessGroupBudgetResponse": { + "properties": { + "access_group": { + "title": "Access Group", + "type": "string" + }, + "budget": { + "anyOf": [ + { + "$ref": "#/components/schemas/AccessGroupBudget" + }, + { + "type": "null" + } + ] + }, + "spend": { + "title": "Spend", + "type": "number" + } + }, + "required": [ + "access_group", + "spend" + ], + "title": "AccessGroupBudgetResponse", + "type": "object" + }, "AccessGroupCreateRequest": { "properties": { "access_agent_ids": { @@ -561,6 +700,16 @@ "title": "Access Group", "type": "string" }, + "budget": { + "anyOf": [ + { + "$ref": "#/components/schemas/AccessGroupBudget" + }, + { + "type": "null" + } + ] + }, "deployment_count": { "title": "Deployment Count", "type": "integer" @@ -571,6 +720,17 @@ }, "title": "Model Names", "type": "array" + }, + "spend": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Spend" } }, "required": [ @@ -782,6 +942,29 @@ "title": "AccessGroupUpdateRequest", "type": "object" }, + "DeleteAccessGroupBudgetResponse": { + "properties": { + "access_group": { + "title": "Access Group", + "type": "string" + }, + "budget_deleted": { + "title": "Budget Deleted", + "type": "boolean" + }, + "message": { + "title": "Message", + "type": "string" + } + }, + "required": [ + "access_group", + "budget_deleted", + "message" + ], + "title": "DeleteAccessGroupBudgetResponse", + "type": "object" + }, "DeleteModelGroupResponse": { "properties": { "access_group": { @@ -1072,6 +1255,156 @@ ] } }, + "/access_group/{access_group}/budget": { + "delete": { + "description": "Clear the shared budget of an access group, leaving the group itself in place.\n\nExample:\n```bash\ncurl -X DELETE 'http://localhost:4000/access_group/production-models/budget' \\\n -H 'Authorization: Bearer sk-1234'\n```\n\nParameters:\n- access_group: str - The access group name (URL path parameter)\n\nReturns:\n- DeleteAccessGroupBudgetResponse; budget_deleted is false when there was nothing to clear\n\nRaises:\n- HTTPException 404: If access group not found", + "operationId": "delete_access_group_budget_access_group__access_group__budget_delete", + "parameters": [ + { + "in": "path", + "name": "access_group", + "required": true, + "schema": { + "title": "Access Group", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/DeleteAccessGroupBudgetResponse" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Delete Access Group Budget", + "tags": [ + "access_groups" + ] + }, + "get": { + "description": "Get the shared budget of an access group, and the spend drawn against it.\n\nExample:\n```bash\ncurl -X GET 'http://localhost:4000/access_group/production-models/budget' \\\n -H 'Authorization: Bearer sk-1234'\n```\n\nParameters:\n- access_group: str - The access group name (URL path parameter)\n\nReturns:\n- AccessGroupBudgetResponse; budget is null when the group has no budget set\n\nRaises:\n- HTTPException 404: If access group not found", + "operationId": "get_access_group_budget_access_group__access_group__budget_get", + "parameters": [ + { + "in": "path", + "name": "access_group", + "required": true, + "schema": { + "title": "Access Group", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AccessGroupBudgetResponse" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Get Access Group Budget", + "tags": [ + "access_groups" + ] + }, + "put": { + "description": "Set or replace the shared budget of an access group. Idempotent.\n\nEvery key that can reach a model in the group draws from this one budget.\n\nExample:\n```bash\ncurl -X PUT 'http://localhost:4000/access_group/production-models/budget' \\\n -H 'Authorization: Bearer sk-1234' \\\n -H 'Content-Type: application/json' \\\n -d '{\n \"max_budget\": 100.0,\n \"budget_duration\": \"30d\"\n }'\n```\n\nParameters:\n- access_group: str - The access group name (URL path parameter)\n- max_budget: Optional[float] - Requests fail once the group's shared spend exceeds this\n- soft_budget: Optional[float] - Fires an alert when reached; requests still succeed\n- budget_duration: Optional[str] - Frequency of resetting the group's spend (e.g. '30d')\n- budget_id: Optional[str] - Link an existing budget instead of creating one\n\nReturns:\n- AccessGroupBudgetResponse with the stored budget and current spend\n\nRaises:\n- HTTPException 400: If no budget field is given, or budget_duration cannot be parsed\n- HTTPException 404: If access group not found", + "operationId": "set_access_group_budget_access_group__access_group__budget_put", + "parameters": [ + { + "in": "path", + "name": "access_group", + "required": true, + "schema": { + "title": "Access Group", + "type": "string" + } + } + ], + "requestBody": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AccessGroupBudgetRequest" + } + } + }, + "required": true + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AccessGroupBudgetResponse" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Set Access Group Budget", + "tags": [ + "access_groups" + ] + } + }, "/access_group/{access_group}/delete": { "delete": { "description": "Delete an access group.\n\nRemoves the access group from all deployments that have it.\n\nExample:\n```bash\ncurl -X DELETE 'http://localhost:4000/access_group/production-models/delete' \\\n -H 'Authorization: Bearer sk-1234'\n```\n\nParameters:\n- access_group: str - The access group name (URL path parameter)\n\nReturns:\n- DeleteModelGroupResponse with deletion details\n\nRaises:\n- HTTPException 404: If access group not found", @@ -1122,7 +1455,7 @@ }, "/access_group/{access_group}/info": { "get": { - "description": "Get information about a specific access group.\n\nExample:\n```bash\ncurl -X GET 'http://localhost:4000/access_group/production-models/info' \\\n -H 'Authorization: Bearer sk-1234'\n```\n\nParameters:\n- access_group: str - The access group name (URL path parameter)\n\nReturns:\n- AccessGroupInfo with the access group details\n\nRaises:\n- HTTPException 404: If access group not found", + "description": "Get information about a specific access group.\n\nExample:\n```bash\ncurl -X GET 'http://localhost:4000/access_group/production-models/info' \\\n -H 'Authorization: Bearer sk-1234'\n```\n\nParameters:\n- access_group: str - The access group name (URL path parameter)\n\nReturns:\n- AccessGroupInfo with the access group details, its shared budget and its spend\n\nRaises:\n- HTTPException 404: If access group not found", "operationId": "get_access_group_info_access_group__access_group__info_get", "parameters": [ { diff --git a/terraform/provider/tools/endpointaudit/coverage_allowlist.txt b/terraform/provider/tools/endpointaudit/coverage_allowlist.txt index d10be89b90c..052962e078e 100644 --- a/terraform/provider/tools/endpointaudit/coverage_allowlist.txt +++ b/terraform/provider/tools/endpointaudit/coverage_allowlist.txt @@ -126,3 +126,6 @@ POST /customer/delete # known gap: litellm_customer GET /team/{team_id}/callback # known gap: team callback resource POST /team/{team_id}/callback # known gap: team callback resource DELETE /team/{team_id}/callback/{callback_name} # known gap: team callback resource +GET /access_group/{access_group}/budget # known gap: budget attributes on litellm_access_group +PUT /access_group/{access_group}/budget # known gap: budget attributes on litellm_access_group +DELETE /access_group/{access_group}/budget # known gap: budget attributes on litellm_access_group From 06d0665f5026e85a7f7f45626a168718f6122c6a Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 29 Aug 2026 14:00:56 -0700 Subject: [PATCH 06/16] fix(budgets): drop the bespoke access group cache TTL and format the repository The cache TTL for a group's budget row was a new DEFAULT_MODEL_ACCESS_GROUP_CACHE_TTL env var defaulting to 600 seconds, which nobody asked for and which the docs gate rightly rejected as undocumented. Every other management object cached in auth_checks, tags included, already reads get_management_object_ttl, so it honors general_settings user_api_key_cache_ttl and falls back to the shared default. Group budgets now do the same, which drops a constant, drops an env var, and makes the row expire on the same operator knob as keys and teams. Also formats ModelAccessGroupBudgetRepository, and teaches the FakeBatch double in the unit of work tests about the new table. That double is read while the cascade unit of work is constructed rather than inside the block, so three tests that never mention access groups were failing at the async with. The new test alongside it walks the dataclass fields, so the next dependent added to the cascade is covered without anyone remembering to update a list. --- litellm/constants.py | 1 - litellm/proxy/auth/auth_checks.py | 3 +- litellm/repositories/table_repositories.py | 4 +- .../repositories/test_unit_of_work.py | 38 +++++++++++++++++++ 4 files changed, 40 insertions(+), 6 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 2139b12e024..6ee4f151bc0 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1681,7 +1681,6 @@ DEFAULT_MCP_NAMESPACE_CSV_MAX_TOKENS: Final = 16 # Ceilings on the cached auth registries; larger tables fall back to per-row lookups # instead of holding an unbounded id set in every worker. TAG_REGISTRY_MAX_SIZE: Final = 5000 -DEFAULT_MODEL_ACCESS_GROUP_CACHE_TTL: Final = int(os.getenv("DEFAULT_MODEL_ACCESS_GROUP_CACHE_TTL", 600)) MODEL_ACCESS_GROUP_REGISTRY_MAX_SIZE: Final = 5000 END_USER_RESTRICTED_REGISTRY_MAX_SIZE: Final = 5000 # How long a failed registry load is remembered as "unusable", so a degraded Postgres diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index cd4d03e08e2..19e4f158144 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -30,7 +30,6 @@ from litellm.constants import ( DEFAULT_ACCESS_GROUP_CACHE_TTL, DEFAULT_IN_MEMORY_TTL, DEFAULT_MAX_RECURSE_DEPTH, - DEFAULT_MODEL_ACCESS_GROUP_CACHE_TTL, EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE, END_USER_RESTRICTED_REGISTRY_MAX_SIZE, MODEL_ACCESS_GROUP_REGISTRY_MAX_SIZE, @@ -1951,7 +1950,7 @@ async def _fetch_uncached_model_access_group_budgets( key=model_access_group_cache_key(fetched_name), value=fetched_obj, model_type=ModelAccessGroupBudget, - ttl=DEFAULT_MODEL_ACCESS_GROUP_CACHE_TTL, + ttl=get_management_object_ttl(user_api_key_cache), ) except Exception as e: # noqa: BLE001 # fail-safe: a budget fetch error must yield "no budget rows", never break auth verbose_proxy_logger.debug("Error batch fetching model access group budgets from database: %s", e) diff --git a/litellm/repositories/table_repositories.py b/litellm/repositories/table_repositories.py index de13e5aeb6c..fac142bb017 100644 --- a/litellm/repositories/table_repositories.py +++ b/litellm/repositories/table_repositories.py @@ -100,9 +100,7 @@ class TagRepository(PrismaTableRepository["prisma_models.LiteLLM_TagTable"]): table_name = "litellm_tagtable" -class ModelAccessGroupBudgetRepository( - PrismaTableRepository["prisma_models.LiteLLM_ModelAccessGroupBudgetTable"] -): +class ModelAccessGroupBudgetRepository(PrismaTableRepository["prisma_models.LiteLLM_ModelAccessGroupBudgetTable"]): table_name = "litellm_modelaccessgroupbudgettable" diff --git a/tests/test_litellm/repositories/test_unit_of_work.py b/tests/test_litellm/repositories/test_unit_of_work.py index 1ebfd917e36..1a76b537e95 100644 --- a/tests/test_litellm/repositories/test_unit_of_work.py +++ b/tests/test_litellm/repositories/test_unit_of_work.py @@ -1,9 +1,11 @@ +from dataclasses import fields from datetime import datetime, timezone from typing import Any, Dict, List, Mapping, Tuple import pytest from litellm.repositories.unit_of_work import ( + LinkedSpendResetWrites, budget_cascade_unit_of_work, spend_reset_unit_of_work, ) @@ -32,6 +34,7 @@ class FakeBatch: self.litellm_teammembership = FakeBatchTable("litellm_teammembership", self.calls) self.litellm_organizationtable = FakeBatchTable("litellm_organizationtable", self.calls) self.litellm_tagtable = FakeBatchTable("litellm_tagtable", self.calls) + self.litellm_modelaccessgroupbudgettable = FakeBatchTable("litellm_modelaccessgroupbudgettable", self.calls) self.litellm_endusertable = FakeBatchTable("litellm_endusertable", self.calls) async def commit(self) -> None: @@ -90,6 +93,7 @@ async def test_budget_cascade_dependents_and_window_advance_share_one_batch(): uow.keys.queue_spend_zero(where=linked) uow.organizations.queue_spend_zero(where=linked) uow.tags.queue_spend_zero(where=linked) + uow.model_access_groups.queue_spend_zero(where=linked) uow.endusers.queue_spend_zero(where={"user_id": {"in": ["enduser-1"]}}) uow.budgets.queue_window_advance(budget_id="budget-1", budget_reset_at=reset_at) assert batch.commit_count == 0 @@ -100,6 +104,7 @@ async def test_budget_cascade_dependents_and_window_advance_share_one_batch(): ("litellm_verificationtoken.update_many", linked, {"spend": 0}), ("litellm_organizationtable.update_many", linked, {"spend": 0}), ("litellm_tagtable.update_many", linked, {"spend": 0}), + ("litellm_modelaccessgroupbudgettable.update_many", linked, {"spend": 0}), ("litellm_endusertable.update_many", {"user_id": {"in": ["enduser-1"]}}, {"spend": 0}), ("litellm_budgettable.update_many", {"budget_id": "budget-1"}, {"budget_reset_at": reset_at}), ] @@ -117,6 +122,39 @@ async def test_budget_window_advance_tolerates_a_tier_deleted_mid_chunk(): assert [call[0] for call in batch.calls] == ["litellm_budgettable.update_many"] +async def test_every_cascade_dependent_writes_to_its_own_table_on_the_one_batch(): + """Walks the dataclass instead of naming tables, so a dependent added to + BudgetCascadeUnitOfWork later cannot go uncovered. + + The named test above only proves the tables it lists, and an unbound + dependent surfaces as an AttributeError from whichever tests happen to + open a cascade. This pins the real contract: every field writes, each to a + distinct table, all on the same batch. + """ + reset_at = datetime(2026, 8, 3, 12, 0, tzinfo=timezone.utc) + batches: List[FakeBatch] = [] + + def _new_batch() -> FakeBatch: + # Fresh per call like db.batch_(), unlike the `lambda: batch` above: a + # second transaction would otherwise alias onto the first and hide. + batches.append(FakeBatch()) + return batches[-1] + + async with budget_cascade_unit_of_work(_new_batch) as uow: + writes = [getattr(uow, field.name) for field in fields(uow)] + for write in writes: + if isinstance(write, LinkedSpendResetWrites): + write.queue_spend_zero(where={"budget_id": "budget-1"}) + else: + write.queue_window_advance(budget_id="budget-1", budget_reset_at=reset_at) + + assert len(batches) == 1, "the cascade must open exactly one transaction" + batch = batches[0] + assert len(batch.calls) == len(writes), "a dependent bound to a batch of its own would not land here" + assert len({call[0] for call in batch.calls}) == len(writes), "two dependents share one table" + assert batch.commit_count == 1 + + async def test_budget_cascade_raising_inside_block_skips_commit(): """A failure part-way through must leave budget_reset_at where it was, so the tier is still due on the next tick.""" From 6b2e7f8a1f5e5d6bcf83793a6248dcad5c7d0e51 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 29 Aug 2026 14:10:58 -0700 Subject: [PATCH 07/16] refactor(budgets): declare route dependencies with Annotated instead of argument defaults --- ...model_access_group_management_endpoints.py | 22 +++++++++---------- .../test_access_group_management.py | 22 ++++++------------- 2 files changed, 17 insertions(+), 27 deletions(-) diff --git a/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py b/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py index bf65da1621f..dabb9334b16 100644 --- a/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py @@ -15,7 +15,7 @@ Endpoints here: import json from collections.abc import Mapping, Sequence from datetime import datetime -from typing import TYPE_CHECKING, Any, Final, Protocol +from typing import TYPE_CHECKING, Annotated, Any, Final, Protocol from fastapi import APIRouter, Depends, HTTPException from typing_extensions import ReadOnly, TypedDict @@ -548,7 +548,7 @@ async def get_all_access_groups_from_db( ) async def create_model_group( data: NewModelGroupRequest, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], ): """ Create a new access group containing multiple model names. @@ -693,7 +693,7 @@ async def create_model_group( response_model=ListAccessGroupsResponse, ) async def list_access_groups( - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], ): """ List all access groups. @@ -743,7 +743,7 @@ async def list_access_groups( ) async def get_access_group_info( access_group: str, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], ): """ Get information about a specific access group. @@ -808,7 +808,7 @@ async def get_access_group_info( async def update_access_group( access_group: str, data: UpdateModelGroupRequest, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], ): """ Update an access group's model names. @@ -961,8 +961,8 @@ async def update_access_group( ) async def delete_access_group( access_group: str, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - auth_cache: UserApiKeyCache = Depends(_auth_cache), + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], + auth_cache: Annotated[UserApiKeyCache, Depends(_auth_cache)], ): """ Delete an access group. @@ -1071,7 +1071,6 @@ async def delete_access_group( ) async def get_access_group_budget( access_group: str, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ) -> AccessGroupBudgetResponse: """ Get the shared budget of an access group, and the spend drawn against it. @@ -1108,8 +1107,8 @@ async def get_access_group_budget( async def set_access_group_budget( access_group: str, data: AccessGroupBudgetRequest, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - auth_cache: UserApiKeyCache = Depends(_auth_cache), + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], + auth_cache: Annotated[UserApiKeyCache, Depends(_auth_cache)], ) -> AccessGroupBudgetResponse: """ Set or replace the shared budget of an access group. Idempotent. @@ -1185,8 +1184,7 @@ async def set_access_group_budget( ) async def delete_access_group_budget( access_group: str, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - auth_cache: UserApiKeyCache = Depends(_auth_cache), + auth_cache: Annotated[UserApiKeyCache, Depends(_auth_cache)], ) -> DeleteAccessGroupBudgetResponse: """ Clear the shared budget of an access group, leaving the group itself in place. diff --git a/tests/test_litellm/proxy/management_endpoints/test_access_group_management.py b/tests/test_litellm/proxy/management_endpoints/test_access_group_management.py index 74171774fc6..a0d8ea48770 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_access_group_management.py +++ b/tests/test_litellm/proxy/management_endpoints/test_access_group_management.py @@ -955,7 +955,7 @@ async def test_get_access_group_budget_returns_the_budget_and_the_shared_spend() _seed_budget(prisma, "prod-models", spend=42.5, max_budget=100.0, budget_duration="30d") with _proxy(prisma): - response = await get_access_group_budget(access_group="prod-models", user_api_key_dict=_admin()) + response = await get_access_group_budget(access_group="prod-models") assert response.access_group == "prod-models" assert response.spend == 42.5 @@ -974,7 +974,7 @@ async def test_get_access_group_budget_on_a_budgetless_group_is_200_not_404(): prisma = _FakePrismaClient([], deployments=[_deployment()]) with _proxy(prisma): - response = await get_access_group_budget(access_group="prod-models", user_api_key_dict=_admin()) + response = await get_access_group_budget(access_group="prod-models") assert response.spend == 0.0 assert response.budget is None @@ -998,16 +998,14 @@ async def test_access_group_budget_routes_404_on_an_unknown_group(): admin = _admin() calls = ( - lambda: get_access_group_budget(access_group="ghost-group", user_api_key_dict=admin), + lambda: get_access_group_budget(access_group="ghost-group"), lambda: set_access_group_budget( access_group="ghost-group", data=AccessGroupBudgetRequest(max_budget=1.0), user_api_key_dict=admin, auth_cache=cache, ), - lambda: delete_access_group_budget( - access_group="ghost-group", user_api_key_dict=admin, auth_cache=cache - ), + lambda: delete_access_group_budget(access_group="ghost-group", auth_cache=cache), ) with _proxy(prisma): @@ -1033,9 +1031,7 @@ async def test_delete_access_group_budget_drops_the_row_and_spares_the_shared_bu cache = _FakeAuthCache() with _proxy(prisma): - response = await delete_access_group_budget( - access_group="prod-models", user_api_key_dict=_admin(), auth_cache=cache - ) + response = await delete_access_group_budget(access_group="prod-models", auth_cache=cache) assert response.budget_deleted is True assert prisma.access_group_budget_table.rows == {} @@ -1056,9 +1052,7 @@ async def test_delete_access_group_budget_on_a_budgetless_group_still_evicts(): cache = _FakeAuthCache(journal) with _proxy(prisma): - response = await delete_access_group_budget( - access_group="prod-models", user_api_key_dict=_admin(), auth_cache=cache - ) + response = await delete_access_group_budget(access_group="prod-models", auth_cache=cache) assert response.budget_deleted is False assert prisma.budget_table.deleted_ids == [] @@ -1160,9 +1154,7 @@ async def test_delete_access_group_budget_evicts_both_auth_cache_keys(): cache = _FakeAuthCache(journal) with _proxy(prisma): - await delete_access_group_budget( - access_group="prod-models", user_api_key_dict=_admin(), auth_cache=cache - ) + await delete_access_group_budget(access_group="prod-models", auth_cache=cache) _assert_evicted_after_write(journal, "prod-models", "access_group_budget.delete:prod-models") From 7a2f6c4cc21cbe66d620e401b662dd4d79eef7fb Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 29 Aug 2026 14:17:45 -0700 Subject: [PATCH 08/16] fix(budgets): let the model access group spend helper swallow its own failure The helper already logs through spend_log_error, so re-raising only for _batch_database_updates to catch and log again was double reporting. Isolate inside the helper instead and drop the caller's wrapper. --- litellm/proxy/db/db_spend_update_writer.py | 27 ++++++++-------------- 1 file changed, 10 insertions(+), 17 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 886394d66f9..4d005f0aaac 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -485,7 +485,7 @@ class DBSpendUpdateWriter: request_model_access_groups: Sequence[str] = (), ): """ - Runs all 12 spend-update helpers sequentially inside a single asyncio task. + Runs all 13 spend-update helpers sequentially inside a single asyncio task. Each helper is wrapped in try/except so one failure doesn't prevent the others. @@ -557,19 +557,13 @@ class DBSpendUpdateWriter: traceback.format_exc(), ) - try: - await self._update_model_access_group_db( - response_cost=response_cost, - request_model_access_groups=request_model_access_groups, - served_model_id=payload_copy.get("model_id"), - prisma_client=prisma_client, - router=_get_llm_router(), - ) - except Exception: - verbose_proxy_logger.debug( - "_batch_database_updates: _update_model_access_group_db failed: %s", - traceback.format_exc(), - ) + await self._update_model_access_group_db( + response_cost=response_cost, + request_model_access_groups=request_model_access_groups, + served_model_id=payload_copy.get("model_id"), + prisma_client=prisma_client, + router=_get_llm_router(), + ) _agent_id_for_spend: Final = payload_copy.get("agent_id") try: @@ -887,7 +881,7 @@ class DBSpendUpdateWriter: served_model_id: str | None, prisma_client: PrismaClient | None, router: _DeploymentLookup | None = None, - ): + ) -> None: """ Update spend for every model access group this request is billed against. @@ -914,7 +908,7 @@ class DBSpendUpdateWriter: response_cost=response_cost, ) ) - except Exception as e: + except Exception as e: # noqa: BLE001 # isolation: a helper failure must not stop the batch spend_log_error( "Spend tracking - failed to enqueue model access group spend update. " "model_access_groups=%s, response_cost=%s - %s", @@ -923,7 +917,6 @@ class DBSpendUpdateWriter: str(e), exc=e, ) - raise e async def _insert_spend_log_to_db( self, From acf3ed7d9b8baebd03f9f8af8cf9ef53ba8cb0ba Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 29 Aug 2026 14:25:56 -0700 Subject: [PATCH 09/16] fix(budgets): narrow model access group spend counters to the served deployment The database writer already intersects the auth-matched groups with the ones the served deployment declares, but the live spend counters got the unnarrowed set. A caller granted two pools that both cover a model group debited both counters while only one row moved, so the in-memory ceiling could block a pool its persisted spend never touched. Narrow once at the callback so both consumers read the same set. --- litellm/proxy/db/db_spend_update_writer.py | 8 +- .../proxy/hooks/proxy_track_cost_callback.py | 10 ++- .../hooks/test_proxy_track_cost_callback.py | 80 +++++++++++++++++++ 3 files changed, 93 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 4d005f0aaac..dc18a57c10c 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -112,7 +112,7 @@ def _spend_update_tx(prisma_client: PrismaClient) -> _SpendTransactionManager: return tx -def _get_llm_router(): +def get_llm_router(): """The proxy's router, or None outside a running proxy. Injected rather than imported where it is used, so the savings computation stays @@ -371,7 +371,7 @@ class DBSpendUpdateWriter: routing_decision=metadata.get("routing_decision"), usage_object=usage_object_raw if isinstance(usage_object_raw, dict) else None, model_id=payload.get("model_id"), - llm_router=_get_llm_router, + llm_router=get_llm_router, cost_breakdown=metadata.get("cost_breakdown"), recorded_autorouter_savings=metadata.get("autorouter_savings"), ) @@ -562,7 +562,7 @@ class DBSpendUpdateWriter: request_model_access_groups=request_model_access_groups, served_model_id=payload_copy.get("model_id"), prisma_client=prisma_client, - router=_get_llm_router(), + router=get_llm_router(), ) _agent_id_for_spend: Final = payload_copy.get("agent_id") @@ -2004,7 +2004,7 @@ class DBSpendUpdateWriter: gateway_injected_cache=marks_gateway_injection(_metadata, payload.get("model_id")), routing_decision=_metadata.get("routing_decision"), model_id=payload.get("model_id"), - llm_router=_get_llm_router, + llm_router=get_llm_router, usage_object=usage_obj, cost_breakdown=_metadata.get("cost_breakdown"), recorded_autorouter_savings=_metadata.get("autorouter_savings"), diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index ad857b8c3f5..d0e4f548c99 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -21,6 +21,10 @@ from litellm.proxy.auth.auth_checks import ( log_db_metrics, ) from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy.db.db_spend_update_writer import ( + debitable_model_access_groups, + get_llm_router, +) from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup from litellm.proxy.spend_tracking.spend_log_error_logger import ( should_suppress_spend_log_tracebacks, @@ -260,7 +264,11 @@ class _ProxyDBLogger(CustomLogger): sl_object=sl_object, metadata=metadata, ) - model_access_groups: Final = get_request_model_access_groups(kwargs) + model_access_groups: Final = debitable_model_access_groups( + attributed=get_request_model_access_groups(kwargs), + served_model_id=sl_object.get("model_id") if sl_object is not None else None, + router=get_llm_router(), + ) if response_cost is not None: user_api_key: Final = metadata.get("user_api_key", None) 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 e923a28f371..c163656219f 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 @@ -1956,3 +1956,83 @@ async def test_track_cost_callback_charges_no_model_access_group_when_none_were_ mock_increment_spend_counters.assert_awaited_once() assert mock_increment_spend_counters.await_args.kwargs["model_access_groups"] == () + + +class _FakeDeploymentLookup: + """Deployment lookup returning the access groups each deployment declares.""" + + def __init__(self, deployments): + self._deployments = deployments + + def get_model_info(self, id): + if id not in self._deployments: + return None + return {"model_name": "premium-haiku", "model_info": {"id": id, "access_groups": list(self._deployments[id])}} + + +def _model_access_group_kwargs(granted, served_model_id): + return { + "call_type": "acompletion", + "model": "premium-haiku", + "litellm_call_id": "test-call-id", + "litellm_params": { + "metadata": { + "user_api_key": "hashed-key", + "user_api_key_user_id": "u-1", + MODEL_ACCESS_GROUP_METADATA_KEY: list(granted), + } + }, + "stream": False, + "standard_logging_object": {"response_cost": 0.25, "model_id": served_model_id}, + } + + +async def _run_callback_capturing_groups(kwargs, deployments): + logger = _ProxyDBLogger() + with ( + patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging, + patch( + "litellm.proxy.hooks.proxy_track_cost_callback._update_database_and_spend_counters", + new=AsyncMock(), + ) as mock_update, + patch("litellm.proxy.proxy_server.llm_router", new=_FakeDeploymentLookup(deployments)), + ): + mock_proxy_logging.failed_tracking_alert = AsyncMock() + mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock() + mock_proxy_logging.db_spend_update_writer = MagicMock() + mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() + + await logger._PROXY_track_cost_callback( + kwargs=kwargs, + completion_response=None, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + return mock_update.await_args.kwargs["model_access_groups"] + + +@pytest.mark.asyncio +async def test_spend_counters_only_debit_the_group_the_served_deployment_belongs_to(): + """A caller granted two pools that both cover the model group only draws down the pool that served. + + The database writer already narrows by served deployment, so passing the unnarrowed set to the + live counters let one request block a pool the persisted spend never debited. + """ + debited = await _run_callback_capturing_groups( + kwargs=_model_access_group_kwargs(granted=["premium", "tier0"], served_model_id="deployment-premium"), + deployments={"deployment-premium": ["premium"], "deployment-tier0": ["tier0"]}, + ) + + assert debited == ("premium",) + + +@pytest.mark.asyncio +async def test_spend_counters_keep_every_granted_group_when_the_deployment_is_unknown(): + """An unidentifiable deployment leaves the auth-time set standing, so nothing silently stops billing.""" + debited = await _run_callback_capturing_groups( + kwargs=_model_access_group_kwargs(granted=["premium", "tier0"], served_model_id="deployment-gone"), + deployments={"deployment-premium": ["premium"]}, + ) + + assert debited == ("premium", "tier0") From d7c0bc1e6d2439ab44c3e4906c861d8926b7e101 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 29 Aug 2026 14:44:00 -0700 Subject: [PATCH 10/16] test(budgets): clear the test-quality violations this branch added The four model access group callback tests now share one helper, so nine patches of proxy_server internals become three, and both mock-echo assertions go with them. The delete_access_group tests share a context manager for the same reason. test_group_exactly_at_its_max_budget_passes gained the assertion it was missing: it now proves the group reached the spend comparison, which a group skipped for a missing budget row would not. The route-allowed patch beside it was dead, so it is gone. What is left is suppressed with the collaborator each one cannot inject. --- .../auth/test_model_access_group_budgets.py | 23 ++- .../hooks/test_proxy_track_cost_callback.py | 143 ++++++------------ .../test_access_group_management.py | 57 ++++--- .../proxy/test_budget_reservation.py | 4 +- 4 files changed, 90 insertions(+), 137 deletions(-) diff --git a/tests/test_litellm/proxy/auth/test_model_access_group_budgets.py b/tests/test_litellm/proxy/auth/test_model_access_group_budgets.py index 2896f76eec6..4396dae202b 100644 --- a/tests/test_litellm/proxy/auth/test_model_access_group_budgets.py +++ b/tests/test_litellm/proxy/auth/test_model_access_group_budgets.py @@ -343,7 +343,9 @@ async def _enforce( cache: UserApiKeyCache | None = None, ) -> list[str]: read, seen = _spend_reader(spend_by_counter_key or {}) - with patch("litellm.proxy.proxy_server.get_current_spend", read): + # The check takes its client and cache as arguments, injected just below. get_current_spend is the + # one collaborator it reaches by a lazy `from litellm.proxy.proxy_server import`, with no parameter. + with patch("litellm.proxy.proxy_server.get_current_spend", read): # test-quality-ok: get_current_spend is lazily imported inside _model_access_group_max_budget_check and has no injection point await _model_access_group_max_budget_check( matched_model_access_groups=matched, prisma_client=prisma_client if prisma_client is not None else _RecordingPrismaClient(*rows), @@ -363,12 +365,16 @@ async def test_group_under_its_max_budget_passes(): @pytest.mark.asyncio async def test_group_exactly_at_its_max_budget_passes(): - """The ceiling is inclusive, matching the tag check it mirrors; only spend strictly above it blocks.""" - await _enforce( + """The ceiling is inclusive, matching the tag check it mirrors; only spend strictly above it blocks. + + Asserting the counter was read is what keeps this honest: a group that got skipped entirely, + because its row never arrived or carried no budget, would also not raise. + """ + assert await _enforce( ("tier-a",), _MagBudgetRow("tier-a", max_budget=10.0), spend_by_counter_key={MODEL_ACCESS_GROUP_COUNTER_KEY: 10.0}, - ) + ) == [MODEL_ACCESS_GROUP_COUNTER_KEY] @pytest.mark.asyncio @@ -469,10 +475,11 @@ async def _common_checks_with_over_budget_group(*, skip_budget_checks: bool) -> read, _ = _spend_reader({MODEL_ACCESS_GROUP_COUNTER_KEY: 99.0}) with ( - patch("litellm.proxy.proxy_server.prisma_client", prisma_client), - patch("litellm.proxy.proxy_server.user_api_key_cache", cache), - patch("litellm.proxy.proxy_server.get_current_spend", read), - patch("litellm.proxy.auth.auth_checks._is_api_route_allowed", return_value=True), + # common_checks resolves all three off the proxy_server module at call time; its signature + # has no client, cache or spend-reader parameter to pass them through instead. + patch("litellm.proxy.proxy_server.prisma_client", prisma_client), # test-quality-ok: common_checks lazily imports prisma_client from proxy_server and takes no client parameter + patch("litellm.proxy.proxy_server.user_api_key_cache", cache), # test-quality-ok: common_checks lazily imports user_api_key_cache from proxy_server and takes no cache parameter + patch("litellm.proxy.proxy_server.get_current_spend", read), # test-quality-ok: get_current_spend is lazily imported inside the budget check and has no injection point ): return await common_checks( request_body={"model": "gpt-4o", "messages": []}, 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 c163656219f..5160dd3431e 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 @@ -1879,85 +1879,6 @@ async def test_track_cost_callback_logs_unauthenticated_pass_through_request( ) -@pytest.mark.asyncio -async def test_track_cost_callback_charges_the_model_access_groups_auth_stamped(): - """Auth stamps the matched groups onto request metadata; the callback has to carry them through. - - Without this hop nothing writes ``spend:model_access_group:*`` on the normal path, so with - reservations disabled the budget check reads a counter no one maintains. - """ - logger = _ProxyDBLogger() - kwargs = { - "model": "gpt-4", - "call_type": "acompletion", - "litellm_params": { - "metadata": { - "user_api_key": "hashed-key", - "user_api_key_user_id": "user-1", - MODEL_ACCESS_GROUP_METADATA_KEY: ["premium", "starter"], - }, - }, - "standard_logging_object": {"response_cost": 0.25, "request_tags": None}, - "stream": False, - } - - with ( - patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging, - patch( - "litellm.proxy.proxy_server.increment_spend_counters", new_callable=AsyncMock - ) as mock_increment_spend_counters, - patch("litellm.proxy.proxy_server.update_cache", new_callable=AsyncMock), - ): - mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() - mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock() - await logger._PROXY_track_cost_callback( - kwargs=kwargs, - completion_response=None, - start_time=datetime.now(), - end_time=datetime.now(), - ) - - mock_increment_spend_counters.assert_awaited_once() - assert mock_increment_spend_counters.await_args.kwargs["model_access_groups"] == ( - "premium", - "starter", - ) - - -@pytest.mark.asyncio -async def test_track_cost_callback_charges_no_model_access_group_when_none_were_stamped(): - """A request no budgeted group authorized must not debit anything.""" - logger = _ProxyDBLogger() - kwargs = { - "model": "gpt-4", - "call_type": "acompletion", - "litellm_params": { - "metadata": {"user_api_key": "hashed-key", "user_api_key_user_id": "user-1"}, - }, - "standard_logging_object": {"response_cost": 0.25, "request_tags": None}, - "stream": False, - } - - with ( - patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging, - patch( - "litellm.proxy.proxy_server.increment_spend_counters", new_callable=AsyncMock - ) as mock_increment_spend_counters, - patch("litellm.proxy.proxy_server.update_cache", new_callable=AsyncMock), - ): - mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() - mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock() - await logger._PROXY_track_cost_callback( - kwargs=kwargs, - completion_response=None, - start_time=datetime.now(), - end_time=datetime.now(), - ) - - mock_increment_spend_counters.assert_awaited_once() - assert mock_increment_spend_counters.await_args.kwargs["model_access_groups"] == () - - class _FakeDeploymentLookup: """Deployment lookup returning the access groups each deployment declares.""" @@ -1970,32 +1891,38 @@ class _FakeDeploymentLookup: return {"model_name": "premium-haiku", "model_info": {"id": id, "access_groups": list(self._deployments[id])}} -def _model_access_group_kwargs(granted, served_model_id): +def _model_access_group_kwargs(granted, served_model_id=None): + metadata = {"user_api_key": "hashed-key", "user_api_key_user_id": "user-1"} + if granted is not None: + metadata[MODEL_ACCESS_GROUP_METADATA_KEY] = list(granted) return { "call_type": "acompletion", "model": "premium-haiku", "litellm_call_id": "test-call-id", - "litellm_params": { - "metadata": { - "user_api_key": "hashed-key", - "user_api_key_user_id": "u-1", - MODEL_ACCESS_GROUP_METADATA_KEY: list(granted), - } - }, + "litellm_params": {"metadata": metadata}, "stream": False, - "standard_logging_object": {"response_cost": 0.25, "model_id": served_model_id}, + "standard_logging_object": {"response_cost": 0.25, "request_tags": None, "model_id": served_model_id}, } -async def _run_callback_capturing_groups(kwargs, deployments): +async def _groups_charged_by_the_callback(kwargs, deployments=None): + """The groups the callback hands the spend counters for one request. + + The callback resolves ``proxy_logging_obj`` and the router by importing them off + ``proxy_server`` inside its own body, so there is no seam to inject either through. + """ logger = _ProxyDBLogger() with ( - patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging, - patch( + patch( # test-quality-ok: callback imports proxy_logging_obj off proxy_server in its body, no seam + "litellm.proxy.proxy_server.proxy_logging_obj" + ) as mock_proxy_logging, + patch( # test-quality-ok: the arguments to this call are the boundary under test "litellm.proxy.hooks.proxy_track_cost_callback._update_database_and_spend_counters", new=AsyncMock(), ) as mock_update, - patch("litellm.proxy.proxy_server.llm_router", new=_FakeDeploymentLookup(deployments)), + patch( # test-quality-ok: llm_router is a proxy_server global the callback reads lazily, no seam + "litellm.proxy.proxy_server.llm_router", new=_FakeDeploymentLookup(deployments or {}) + ), ): mock_proxy_logging.failed_tracking_alert = AsyncMock() mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock() @@ -2012,6 +1939,30 @@ async def _run_callback_capturing_groups(kwargs, deployments): return mock_update.await_args.kwargs["model_access_groups"] +@pytest.mark.asyncio +async def test_track_cost_callback_charges_the_model_access_groups_auth_stamped(): + """Auth stamps the matched groups onto request metadata; the callback has to carry them through. + + Without this hop nothing writes ``spend:model_access_group:*`` on the normal path, so with + reservations disabled the budget check reads a counter no one maintains. + """ + charged = await _groups_charged_by_the_callback( + kwargs=_model_access_group_kwargs(granted=["premium", "starter"]), + ) + + assert charged == ("premium", "starter") + + +@pytest.mark.asyncio +async def test_track_cost_callback_charges_no_model_access_group_when_none_were_stamped(): + """A request no budgeted group authorized must not debit anything.""" + charged = await _groups_charged_by_the_callback( + kwargs=_model_access_group_kwargs(granted=None), + ) + + assert charged == () + + @pytest.mark.asyncio async def test_spend_counters_only_debit_the_group_the_served_deployment_belongs_to(): """A caller granted two pools that both cover the model group only draws down the pool that served. @@ -2019,20 +1970,20 @@ async def test_spend_counters_only_debit_the_group_the_served_deployment_belongs The database writer already narrows by served deployment, so passing the unnarrowed set to the live counters let one request block a pool the persisted spend never debited. """ - debited = await _run_callback_capturing_groups( + charged = await _groups_charged_by_the_callback( kwargs=_model_access_group_kwargs(granted=["premium", "tier0"], served_model_id="deployment-premium"), deployments={"deployment-premium": ["premium"], "deployment-tier0": ["tier0"]}, ) - assert debited == ("premium",) + assert charged == ("premium",) @pytest.mark.asyncio async def test_spend_counters_keep_every_granted_group_when_the_deployment_is_unknown(): """An unidentifiable deployment leaves the auth-time set standing, so nothing silently stops billing.""" - debited = await _run_callback_capturing_groups( + charged = await _groups_charged_by_the_callback( kwargs=_model_access_group_kwargs(granted=["premium", "tier0"], served_model_id="deployment-gone"), deployments={"deployment-premium": ["premium"]}, ) - assert debited == ("premium", "tier0") + assert charged == ("premium", "tier0") diff --git a/tests/test_litellm/proxy/management_endpoints/test_access_group_management.py b/tests/test_litellm/proxy/management_endpoints/test_access_group_management.py index a0d8ea48770..de5fc96c7c3 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_access_group_management.py +++ b/tests/test_litellm/proxy/management_endpoints/test_access_group_management.py @@ -753,7 +753,29 @@ def _seed_budget(prisma, access_group, spend=0.0, budget_id="budget-seed", **bud @contextmanager def _proxy(prisma): - with patch("litellm.proxy.proxy_server.prisma_client", prisma): + with patch( # test-quality-ok: the endpoints import proxy_server.prisma_client themselves; no parameter to inject + "litellm.proxy.proxy_server.prisma_client", prisma + ): + yield + + +@contextmanager +def _proxy_with_stubbed_reload(prisma): + """delete_access_group finishes by reloading the router and judging what it serves afterwards. + Both collaborators it reaches for there are module globals it imports itself, so a fake can only + get in by patching them; auth_cache and prisma are the ones with a real seam.""" + never_served_router = MagicMock() + never_served_router.get_model_ids.return_value = [] + with ( + _proxy(prisma), + patch( # test-quality-ok: live_model_ids_snapshot() reads the llm_router global; the endpoint takes no router + "litellm.proxy.proxy_server.llm_router", never_served_router + ), + patch( # test-quality-ok: the endpoint calls its module-level clear_cache import; there is no parameter for it + "litellm.proxy.management_endpoints.model_access_group_management_endpoints.clear_cache", + new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)), + ), + ): yield @@ -1073,16 +1095,7 @@ async def test_deleting_the_access_group_strips_deployments_before_dropping_the_ _seed_budget(prisma, "prod-models", spend=3.0, max_budget=100.0) cache = _FakeAuthCache() - never_served_router = MagicMock() - never_served_router.get_model_ids.return_value = [] - with ( - _proxy(prisma), - patch("litellm.proxy.proxy_server.llm_router", never_served_router), - patch( - "litellm.proxy.management_endpoints.model_access_group_management_endpoints.clear_cache", - new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)), - ), - ): + with _proxy_with_stubbed_reload(prisma): response = await delete_access_group( access_group="prod-models", user_api_key_dict=_admin(), auth_cache=cache ) @@ -1171,16 +1184,7 @@ async def test_deleting_the_access_group_evicts_both_auth_cache_keys(): _seed_budget(prisma, "prod-models", spend=3.0, max_budget=100.0) cache = _FakeAuthCache(journal) - never_served_router = MagicMock() - never_served_router.get_model_ids.return_value = [] - with ( - _proxy(prisma), - patch("litellm.proxy.proxy_server.llm_router", never_served_router), - patch( - "litellm.proxy.management_endpoints.model_access_group_management_endpoints.clear_cache", - new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)), - ), - ): + with _proxy_with_stubbed_reload(prisma): await delete_access_group(access_group="prod-models", user_api_key_dict=_admin(), auth_cache=cache) _assert_evicted_after_write(journal, "prod-models", "access_group_budget.delete:prod-models") @@ -1198,16 +1202,7 @@ async def test_deleting_an_access_group_that_never_had_a_budget_still_evicts(): prisma = _FakePrismaClient(journal, deployments=[_deployment()]) cache = _FakeAuthCache(journal) - never_served_router = MagicMock() - never_served_router.get_model_ids.return_value = [] - with ( - _proxy(prisma), - patch("litellm.proxy.proxy_server.llm_router", never_served_router), - patch( - "litellm.proxy.management_endpoints.model_access_group_management_endpoints.clear_cache", - new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)), - ), - ): + with _proxy_with_stubbed_reload(prisma): response = await delete_access_group( access_group="prod-models", user_api_key_dict=_admin(), auth_cache=cache ) diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 65c8c248a27..95067929ac1 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -3050,7 +3050,7 @@ async def test_model_access_group_counter_blocks_a_request_over_the_group_budget prisma_client = _ModelAccessGroupBudgetPrisma(premium=1.0) valid_token = UserAPIKeyAuth(api_key="hashed", token="tok", matched_model_access_groups=["premium"]) - with patch( + with patch( # test-quality-ok: reserve_budget_for_request takes no estimator, so pinning the estimate needs this attribute "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", return_value=0.5, ): @@ -3081,7 +3081,7 @@ async def _cache_model_access_group_budget(key_cache, group, spend, max_budget=N async def _reserve_for_model_access_groups(key_cache, groups, estimate): """Reserve against the given groups, whose rows are already cached, so nothing hits the DB.""" - with patch( + with patch( # test-quality-ok: reserve_budget_for_request takes no estimator, so pinning the estimate needs this attribute "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", return_value=estimate, ): From 40a7fe922177429bff66d352c2806f1fc4085b75 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 29 Aug 2026 15:20:10 -0700 Subject: [PATCH 11/16] fix(budgets): make the model access group ceiling exclusive A pool whose recorded spend has reached max_budget has nothing left to give, so the next request is refused rather than admitted. This departs from the tag check it otherwise mirrors and matches where keys and organizations already draw the line. A non-positive budget now means no budget here too, so the read-time check and the reservation path agree on what counts as unbudgeted. --- litellm/proxy/auth/auth_checks.py | 9 +++- .../auth/test_model_access_group_budgets.py | 42 ++++++++++++++++--- 2 files changed, 44 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 19e4f158144..f6345eae06d 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -5620,6 +5620,11 @@ async def _model_access_group_max_budget_check( overshoot the ceiling slightly. The reservation counters are the precise path; this one covers the ``disable_budget_reservation`` case. + The ceiling is exclusive, unlike the tag check it otherwise mirrors: a pool whose recorded + spend has reached ``max_budget`` has nothing left to give, so the next request is refused. + Keys and organizations already draw the line there. A non-positive budget means no budget, + matching what the reservation path treats as unbudgeted. + Raises: BudgetExceededError if a matched group is over its max budget. """ @@ -5636,7 +5641,7 @@ async def _model_access_group_max_budget_check( for group in matched_model_access_groups: budget = budgets.get(group) - if budget is None or budget.max_budget is None: + if budget is None or budget.max_budget is None or budget.max_budget <= 0: continue group_spend = await get_current_spend( @@ -5645,7 +5650,7 @@ async def _model_access_group_max_budget_check( max_budget=budget.max_budget, fallback_authoritative=True, ) - if group_spend <= budget.max_budget: + if group_spend < budget.max_budget: continue raise litellm.BudgetExceededError( current_cost=group_spend, diff --git a/tests/test_litellm/proxy/auth/test_model_access_group_budgets.py b/tests/test_litellm/proxy/auth/test_model_access_group_budgets.py index 4396dae202b..fb82d8708fd 100644 --- a/tests/test_litellm/proxy/auth/test_model_access_group_budgets.py +++ b/tests/test_litellm/proxy/auth/test_model_access_group_budgets.py @@ -364,19 +364,51 @@ async def test_group_under_its_max_budget_passes(): @pytest.mark.asyncio -async def test_group_exactly_at_its_max_budget_passes(): - """The ceiling is inclusive, matching the tag check it mirrors; only spend strictly above it blocks. +async def test_group_exactly_at_its_max_budget_blocks_the_request(): + """A pool whose spend has reached the ceiling has nothing left, so the next request is refused. - Asserting the counter was read is what keeps this honest: a group that got skipped entirely, - because its row never arrived or carried no budget, would also not raise. + This is where the check departs from the tag one it otherwise mirrors, and it matches where + keys and organizations already draw the line. """ + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _enforce( + ("tier-a",), + _MagBudgetRow("tier-a", max_budget=10.0), + spend_by_counter_key={MODEL_ACCESS_GROUP_COUNTER_KEY: 10.0}, + ) + + assert exc_info.value.entity_id == "tier-a" + assert exc_info.value.current_cost == 10.0 + + +@pytest.mark.asyncio +async def test_group_just_under_its_max_budget_passes(): + """Asserting the counter was read is what keeps this honest: a group that got skipped entirely, + because its row never arrived or carried no budget, would also not raise.""" assert await _enforce( ("tier-a",), _MagBudgetRow("tier-a", max_budget=10.0), - spend_by_counter_key={MODEL_ACCESS_GROUP_COUNTER_KEY: 10.0}, + spend_by_counter_key={MODEL_ACCESS_GROUP_COUNTER_KEY: 9.99}, ) == [MODEL_ACCESS_GROUP_COUNTER_KEY] +@pytest.mark.asyncio +async def test_a_non_positive_budget_means_no_budget(): + """The reservation path treats max_budget <= 0 as unbudgeted, so the read-time check must agree. + + Without this the exclusive ceiling would turn a zero into a total freeze on one path and a + no-op on the other. + """ + assert ( + await _enforce( + ("tier-a",), + _MagBudgetRow("tier-a", max_budget=0.0), + spend_by_counter_key={MODEL_ACCESS_GROUP_COUNTER_KEY: 5.0}, + ) + == [] + ) + + @pytest.mark.asyncio async def test_group_over_its_max_budget_blocks_the_request_and_names_the_group(): with pytest.raises(litellm.BudgetExceededError) as exc_info: From e263c09e4fd2cecf32a3e146e79ccdf4cc31a7d8 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 29 Aug 2026 15:46:49 -0700 Subject: [PATCH 12/16] test(e2e): cover model access group budgets against a live proxy Four cases in tests/e2e/quota_management/budgets, driving real OpenAI calls through a group whose shared pool is drained to exhaustion: the spender key stays blocked, a key that spent nothing of its own is blocked by the same pool, a sibling group with no budget keeps serving, and the budget read reports the spend drawn against the group. Adds set/get/delete access group budget to BudgetClient and the four matching rows to the coverage registry. --- tests/e2e/CLAUDE.md | 8 +- .../coverage_registry/quota_management.yaml | 4 + .../quota_management/budgets/budget_client.py | 65 +++++++ .../test_model_access_group_budget_e2e.py | 161 ++++++++++++++++++ 4 files changed, 234 insertions(+), 4 deletions(-) create mode 100644 tests/e2e/quota_management/budgets/test_model_access_group_budget_e2e.py diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index cf912ddaa25..14b2b4e3299 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -177,15 +177,15 @@ quota_management... behavior : ratelimit | budget | spend_tracking variant : rpm | tpm | priority_generous | priority_strict key | internal_user | end_user | organization | team | team_member | tag - | model_max | soft | key_multi_window | team_multi_window - | fallback | spend_counter + | model_access_group | model_max | soft | key_multi_window + | team_multi_window | fallback | spend_counter chat_completions | stream | messages_bridge | embeddings | cache_hit | key_rollup | concurrent_burst | tags | end_user | per_model | failure | spend_calculate | pagination assertion : blocks_over_limit | resets_after_window | headers_report_remaining | picks_under_tpm | blocks_then_resets | resets_windows_independently | alerts_without_blocking - | isolates_per_model | isolates_per_member | enforced_across_keys | routes_to_fallback - | reseed_matches_db | logs_cost | zero_cost + | isolates_per_model | isolates_per_member | isolates_per_group | enforced_across_keys + | routes_to_fallback | reseed_matches_db | reports_spend | logs_cost | zero_cost | matches_sum_of_logs | loses_no_spend | attributes_spend | writes_own_rows | writes_failure_row | returns_cost | keeps_total e.g. quota_management.ratelimit.rpm.blocks_over_limit exercised_on=[chat_completions, messages] diff --git a/tests/e2e/coverage_registry/quota_management.yaml b/tests/e2e/coverage_registry/quota_management.yaml index 42a075681e0..d0afcaca848 100644 --- a/tests/e2e/coverage_registry/quota_management.yaml +++ b/tests/e2e/coverage_registry/quota_management.yaml @@ -20,6 +20,10 @@ - {id: quota_management.budget.organization.blocks_over_limit, module: quota_management, tier: P1, behavior: budget, variant: organization, assertions: [blocks_over_limit], exercised_on: [chat_completions], source: "proxy/auth/auth_checks.py", rationale: "An organization's max_budget blocks keys under its teams"} - {id: quota_management.budget.team_member.blocks_over_limit, module: quota_management, tier: P1, behavior: budget, variant: team_member, assertions: [blocks_over_limit], exercised_on: [chat_completions], source: "proxy/auth/auth_checks.py", rationale: "A member's per-team budget blocks independently of the team budget"} - {id: quota_management.budget.team_member.isolates_per_member, module: quota_management, tier: P1, behavior: budget, variant: team_member, assertions: [isolates_per_member], exercised_on: [chat_completions], source: "proxy/auth/auth_checks.py", rationale: "One team member's exhausted per-team budget does not block a different member on the same team"} +- {id: quota_management.budget.model_access_group.blocks_over_limit, module: quota_management, tier: P1, behavior: budget, variant: model_access_group, assertions: [blocks_over_limit], exercised_on: [chat_completions], source: "proxy/auth/auth_checks.py", rationale: "A model access group's shared max_budget blocks further calls to deployments in the group once the pool is spent"} +- {id: quota_management.budget.model_access_group.enforced_across_keys, module: quota_management, tier: P1, behavior: budget, variant: model_access_group, assertions: [enforced_across_keys], exercised_on: [chat_completions], source: "proxy/auth/auth_checks.py", rationale: "The pool is shared, so a key that spent nothing of its own is blocked once another key granted the same group drained it"} +- {id: quota_management.budget.model_access_group.isolates_per_group, module: quota_management, tier: P1, behavior: budget, variant: model_access_group, assertions: [isolates_per_group], exercised_on: [chat_completions], source: "proxy/db/db_spend_update_writer.py", rationale: "A request is charged only to the granted groups that serve the model it called, so an exhausted group never blocks a sibling group"} +- {id: quota_management.budget.model_access_group.reports_spend, module: quota_management, tier: P2, behavior: budget, variant: model_access_group, assertions: [reports_spend], exercised_on: [chat_completions], source: "proxy/management_endpoints/model_access_group_management_endpoints.py", rationale: "GET /access_group/{name}/budget reports the pool and the spend drawn against it, so an admin can see why calls are being refused"} - {id: quota_management.budget.tag.blocks_over_limit, module: quota_management, tier: P1, behavior: budget, variant: tag, assertions: [blocks_over_limit], exercised_on: [chat_completions], source: "router_strategy/budget_limiter.py", rationale: "Proxy-level tag budgets block tagged requests at the cap"} - {id: quota_management.budget.end_user_model_max.blocks_over_limit, module: quota_management, tier: P1, behavior: budget, variant: end_user_model_max, assertions: [blocks_over_limit], exercised_on: [chat_completions], source: "budget_management_endpoints.py", fail_before_fix: proven, rationale: "A per-model rpm_limit on an end-user budget is accepted and stored but never enforced; only key-attached budgets honour it"} - {id: quota_management.budget.model_max.isolates_per_model, module: quota_management, tier: P1, behavior: budget, variant: model_max, assertions: [isolates_per_model], exercised_on: [chat_completions], source: "proxy/hooks/model_max_budget_limiter.py", rationale: "model_max_budget caps one model without touching a sibling's budget"} diff --git a/tests/e2e/quota_management/budgets/budget_client.py b/tests/e2e/quota_management/budgets/budget_client.py index 543d5f959e0..087dc8ca522 100644 --- a/tests/e2e/quota_management/budgets/budget_client.py +++ b/tests/e2e/quota_management/budgets/budget_client.py @@ -151,6 +151,28 @@ class TagDeleteBody(BaseModel): name: str +class AccessGroupBudgetBody(BaseModel): + max_budget: float | None = None + soft_budget: float | None = None + budget_duration: str | None = None + + +class AccessGroupBudgetView(BaseModel): + budget_id: str + max_budget: float | None = None + soft_budget: float | None = None + budget_duration: str | None = None + + +class AccessGroupBudgetResponse(BaseModel): + """GET/PUT /access_group/{name}/budget: the group's shared pool and the spend + every key that can reach the group has drawn against it.""" + + access_group: str + spend: float + budget: AccessGroupBudgetView | None = None + + class BudgetNewBody(BaseModel): max_budget: float | None = None soft_budget: float | None = None @@ -514,6 +536,49 @@ class BudgetClient: response_type=NoBody, ) + # ---- model access group --------------------------------------------- + + def set_access_group_budget( + self, + access_group: str, + *, + max_budget: float | None = None, + soft_budget: float | None = None, + budget_duration: str | None = None, + ) -> AccessGroupBudgetResponse: + """Give a model access group one shared budget. Every key that can reach a + deployment in the group draws from it.""" + return unwrap( + self.proxy.transport.put( + f"/access_group/{access_group}/budget", + headers=self.proxy.transport.master, + json=AccessGroupBudgetBody( + max_budget=max_budget, + soft_budget=soft_budget, + budget_duration=budget_duration, + ), + response_type=AccessGroupBudgetResponse, + ) + ) + + def access_group_budget(self, access_group: str) -> AccessGroupBudgetResponse: + return unwrap( + self.proxy.transport.get( + f"/access_group/{access_group}/budget", + headers=self.proxy.transport.master, + params=NoBody(), + response_type=AccessGroupBudgetResponse, + ) + ) + + def delete_access_group_budget(self, access_group: str) -> None: + _ = self.proxy.transport.delete( + f"/access_group/{access_group}/budget", + headers=self.proxy.transport.master, + json=NoBody(), + response_type=NoBody, + ) + # ---- budget table --------------------------------------------------- def create_budget( diff --git a/tests/e2e/quota_management/budgets/test_model_access_group_budget_e2e.py b/tests/e2e/quota_management/budgets/test_model_access_group_budget_e2e.py new file mode 100644 index 00000000000..9c927a31216 --- /dev/null +++ b/tests/e2e/quota_management/budgets/test_model_access_group_budget_e2e.py @@ -0,0 +1,161 @@ +"""Live e2e: one shared budget across every key that can reach a model access group. + +A model access group is a free-text label on a deployment (`model_info.access_groups`), +and a key is granted the group by name. The budget hangs off the group, not the key, so +the interesting behaviors are the ones a per-key budget cannot produce: a key that has +spent nothing of its own is refused once somebody else drained the pool, and draining one +group leaves a second group untouched, because a request is only charged to the groups +the caller was granted that also serve the model being called. +""" + +from __future__ import annotations + +import os +import time +from collections.abc import Iterator +from dataclasses import dataclass +from typing import Final + +import pytest + +from budget_client import BudgetClient, is_budget_block +from e2e_config import unique_marker +from e2e_http import StreamingResponse, require_successful_call +from lifecycle import ResourceManager +from models import KeyGenerateBody, LiteLLMParamsBody, ModelInfoBody, ModelNewBody + +pytestmark = pytest.mark.e2e + +BACKEND: Final = "openai/gpt-5.4-nano" +TINY_BUDGET: Final = 5e-6 +MAX_TOKENS: Final = 16 +DRAIN_TIMEOUT_SECONDS: Final = 180 + + +@dataclass(frozen=True, slots=True) +class DrainedPool: + """A model access group whose shared budget has been spent to exhaustion, the + deployment inside it, the key that did the spending, and a second group holding + its own deployment that was never given a budget at all.""" + + access_group: str + model: str + spender_key: str + free_access_group: str + free_model: str + + +def _provider_key(env_var: str) -> str: + return os.environ.get(env_var) or f"os.environ/{env_var}" + + +def _grouped_model(model_name: str, access_group: str) -> ModelNewBody: + return ModelNewBody( + model_name=model_name, + litellm_params=LiteLLMParamsBody(model=BACKEND, api_key=_provider_key("OPENAI_API_KEY")), + model_info=ModelInfoBody(access_groups=[access_group]), + ) + + +def _call(client: BudgetClient, key: str, model: str) -> StreamingResponse: + return client.chat(key, model, f"hi {unique_marker()}", max_tokens=MAX_TOKENS) + + +def _drain(client: BudgetClient, key: str, model: str, access_group: str) -> None: + """Spend the group's pool until the proxy refuses the next request. The first call + lands under the cap and the block comes from the spend it recorded, so this needs at + least one round trip through the spend writer, not just one request.""" + deadline: Final = time.monotonic() + DRAIN_TIMEOUT_SECONDS + while time.monotonic() < deadline: + result = _call(client, key, model) + if is_budget_block(result): + return + require_successful_call(result) + time.sleep(1) + pytest.fail(f"budget on model access group {access_group!r} never blocked a request") + + +@pytest.fixture(scope="module") +def drained(client: BudgetClient) -> Iterator[DrainedPool]: + marker: Final = unique_marker() + pool: Final = DrainedPool( + access_group=f"e2e-mag-budget-{marker}", + model=f"e2e-mag-budgeted-{marker}", + spender_key=client.proxy.generate_key(KeyGenerateBody(models=[f"e2e-mag-budget-{marker}"])), + free_access_group=f"e2e-mag-free-{marker}", + free_model=f"e2e-mag-unbudgeted-{marker}", + ) + created: Final = ( + client.proxy.register_model(_grouped_model(pool.model, pool.access_group)), + client.proxy.register_model(_grouped_model(pool.free_model, pool.free_access_group)), + ) + try: + client.set_access_group_budget(pool.access_group, max_budget=TINY_BUDGET) + _drain(client, pool.spender_key, pool.model, pool.access_group) + yield pool + finally: + client.delete_access_group_budget(pool.access_group) + client.proxy.delete_key(pool.spender_key) + for model_id in created: + client.proxy.delete_model(model_id) + + +class TestModelAccessGroupBudget: + @pytest.mark.covers("quota_management.budget.model_access_group.blocks_over_limit") + def test_the_key_that_drained_the_pool_stays_blocked( + self, client: BudgetClient, drained: DrainedPool + ) -> None: + result = _call(client, drained.spender_key, drained.model) + assert is_budget_block(result), ( + f"an exhausted pool served {drained.model!r} again: {result.status_code} {result.body[:300]}" + ) + assert drained.access_group in result.body, ( + f"the block did not name the group that caused it: {result.body[:300]}" + ) + + @pytest.mark.covers("quota_management.budget.model_access_group.enforced_across_keys") + def test_a_key_that_spent_nothing_is_blocked_by_the_shared_pool( + self, client: BudgetClient, resources: ResourceManager, drained: DrainedPool + ) -> None: + newcomer = resources.key(models=[drained.access_group]) + + result = _call(client, newcomer, drained.model) + + assert is_budget_block(result), ( + "a freshly minted key with no spend of its own was served by an exhausted " + f"shared pool: {result.status_code} {result.body[:300]}" + ) + + @pytest.mark.covers("quota_management.budget.model_access_group.isolates_per_group") + def test_a_drained_group_does_not_block_a_different_group( + self, client: BudgetClient, resources: ResourceManager, drained: DrainedPool + ) -> None: + other = resources.key(models=[drained.free_access_group]) + + result = _call(client, other, drained.free_model) + + assert not is_budget_block(result), ( + f"{drained.free_access_group!r} has no budget of its own but was blocked by " + f"{drained.access_group!r}'s exhausted pool: {result.body[:300]}" + ) + require_successful_call(result) + + @pytest.mark.covers("quota_management.budget.model_access_group.reports_spend") + def test_the_budget_read_reports_the_spend_drawn_against_the_pool( + self, client: BudgetClient, drained: DrainedPool + ) -> None: + """Enforcement runs off a live counter while the group's row is written by the + batched spend writer, so the recorded spend an admin reads lands a beat after the + block. Poll for it: what matters is that it arrives and matches the pool.""" + deadline = time.monotonic() + client.proxy.poll_timeout + reported = client.access_group_budget(drained.access_group) + while reported.spend < TINY_BUDGET and time.monotonic() < deadline: + time.sleep(client.proxy.poll_interval) + reported = client.access_group_budget(drained.access_group) + + assert reported.budget is not None, "the group lost the budget that just blocked it" + assert reported.budget.max_budget == TINY_BUDGET + assert reported.spend >= TINY_BUDGET, ( + f"the pool blocked at {TINY_BUDGET} but only {reported.spend} was ever recorded " + f"against the group within {client.proxy.poll_timeout}s" + ) From 3e2999f29fb3adbe621e890387c305995791eb41 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 29 Aug 2026 16:05:57 -0700 Subject: [PATCH 13/16] fix(proxy): run SMTP send_email off the event loop with a connection timeout (#38473) * fix(proxy): run SMTP send_email off the event loop with a connection timeout Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): format utils.py and update _create_smtp_connection tests for timeout Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): keep malformed SMTP_TIMEOUT inside the email error boundary Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore: retrigger ci Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * ci: exclude misaligned circleci coverage flag from merged codecov report Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore: retrigger ci for codecov and benchmarks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * ci: disable carryforward for the circleci codecov flag Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * ci: exclude carried-forward coverage from the codecov patch status Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * ci: stop carrying forward the dead circleci codecov flag Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- codecov.yaml | 2 + litellm/proxy/utils.py | 67 +++++++++++++------ tests/test_litellm/proxy/test_proxy_utils.py | 7 +- .../proxy/utils/prisma_and_spend/conftest.py | 5 ++ .../utils/prisma_and_spend/test_send_email.py | 56 +++++++++++++--- 5 files changed, 103 insertions(+), 34 deletions(-) diff --git a/codecov.yaml b/codecov.yaml index bc0b3604329..4d93c18f3ac 100644 --- a/codecov.yaml +++ b/codecov.yaml @@ -25,6 +25,8 @@ flag_management: carryforward: false - name: proxy-db-schema-migration carryforward: false + - name: circleci + carryforward: false component_management: individual_components: diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index eab56c31c39..cf56fc0b1dd 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -6031,10 +6031,42 @@ def _should_use_smtp_ssl(smtp_port: int) -> bool: return os.getenv("SMTP_USE_SSL", "False") == "True" or smtp_port == 465 -def _create_smtp_connection(smtp_host: str, smtp_port: int) -> smtplib.SMTP: +def _create_smtp_connection(smtp_host: str, smtp_port: int, timeout: float) -> smtplib.SMTP: if _should_use_smtp_ssl(smtp_port=smtp_port): - return smtplib.SMTP_SSL(host=smtp_host, port=smtp_port, context=ssl.create_default_context()) - return smtplib.SMTP(host=smtp_host, port=smtp_port) + return smtplib.SMTP_SSL(host=smtp_host, port=smtp_port, context=ssl.create_default_context(), timeout=timeout) + return smtplib.SMTP(host=smtp_host, port=smtp_port, timeout=timeout) + + +def _send_smtp_message( + email_message: MIMEMultipart, + smtp_host: str, + smtp_port: int, + smtp_username: str | None, + smtp_password: str | None, + sender_email: str, + receiver_email: str, + timeout: float, +) -> None: + using_ssl: Final = _should_use_smtp_ssl(smtp_port=smtp_port) + with _create_smtp_connection( + smtp_host=smtp_host, + smtp_port=smtp_port, + timeout=timeout, + ) as server: + if not using_ssl and os.getenv("SMTP_TLS", "True") != "False": + server.starttls(context=ssl.create_default_context()) + + if smtp_username and smtp_password: + server.login( + user=smtp_username, + password=smtp_password, + ) + + server.send_message( + msg=email_message, + from_addr=sender_email, + to_addrs=receiver_email, + ) async def send_email( @@ -6080,27 +6112,18 @@ async def send_email( email_message.attach(MIMEText(html, "html")) try: - using_ssl: Final = _should_use_smtp_ssl(smtp_port=smtp_port) - with _create_smtp_connection( + smtp_timeout: Final = float(os.getenv("SMTP_TIMEOUT", "30")) + await asyncio.to_thread( + _send_smtp_message, + email_message=email_message, smtp_host=smtp_host, smtp_port=smtp_port, - ) as server: - if not using_ssl and os.getenv("SMTP_TLS", "True") != "False": - server.starttls(context=ssl.create_default_context()) - - # Login to your email account only if smtp_username and smtp_password are provided - if smtp_username and smtp_password: - server.login( - user=smtp_username, - password=smtp_password, - ) - - # Send the email - server.send_message( - msg=email_message, - from_addr=sender_email, - to_addrs=receiver_email, - ) + smtp_username=smtp_username, + smtp_password=smtp_password, + sender_email=sender_email, + receiver_email=receiver_email, + timeout=smtp_timeout, + ) except Exception as e: verbose_proxy_logger.exception("An error occurred while sending the email:" + str(e)) diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index 6920cc0dae3..dcaad968663 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -1266,13 +1266,14 @@ class TestCreateSmtpConnection: patch("smtplib.SMTP_SSL") as mock_smtp_ssl, patch("smtplib.SMTP") as mock_smtp, ): - result = _create_smtp_connection(smtp_host="mail.example.com", smtp_port=465) + result = _create_smtp_connection(smtp_host="mail.example.com", smtp_port=465, timeout=30.0) mock_smtp.assert_not_called() assert result is mock_smtp_ssl.return_value _, kwargs = mock_smtp_ssl.call_args assert kwargs["host"] == "mail.example.com" assert kwargs["port"] == 465 + assert kwargs["timeout"] == 30.0 context = kwargs["context"] assert isinstance(context, ssl.SSLContext) assert context.verify_mode == ssl.CERT_REQUIRED @@ -1286,11 +1287,11 @@ class TestCreateSmtpConnection: patch("smtplib.SMTP_SSL") as mock_smtp_ssl, patch("smtplib.SMTP") as mock_smtp, ): - result = _create_smtp_connection(smtp_host="mail.example.com", smtp_port=587) + result = _create_smtp_connection(smtp_host="mail.example.com", smtp_port=587, timeout=30.0) mock_smtp_ssl.assert_not_called() assert result is mock_smtp.return_value - mock_smtp.assert_called_once_with(host="mail.example.com", port=587) + mock_smtp.assert_called_once_with(host="mail.example.com", port=587, timeout=30.0) class TestSendEmailStartTls: diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/conftest.py b/tests/test_litellm/proxy/utils/prisma_and_spend/conftest.py index 19abcb5d66d..fce51c9296c 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/conftest.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/conftest.py @@ -14,6 +14,7 @@ from __future__ import annotations import asyncio import sys +import threading from dataclasses import dataclass, field from email.message import EmailMessage from pathlib import Path @@ -320,6 +321,7 @@ class _SentMessage: body: Optional[str] starttls_called: bool login_args: Optional[tuple] + thread_ident: int @dataclass @@ -328,6 +330,7 @@ class InMemorySMTP: sent: List[_SentMessage] = field(default_factory=list) raise_on_send: Optional[Exception] = None + connection_kwargs: List[Dict[str, Any]] = field(default_factory=list) def server_factory(self) -> Callable[..., Any]: outer = self @@ -370,10 +373,12 @@ class InMemorySMTP: body=body, starttls_called=self._starttls_called, login_args=self._login_args, + thread_ident=threading.get_ident(), ) ) def _factory(*args: Any, **kwargs: Any) -> _Conn: + outer.connection_kwargs.append(dict(kwargs)) return _Conn() return _factory diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_send_email.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_send_email.py index 739e942de52..0e8aba0a03b 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_send_email.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_send_email.py @@ -6,6 +6,7 @@ Symbols pinned here: from __future__ import annotations +import threading from typing import Any import pytest @@ -51,9 +52,7 @@ async def test_send_email_dispatches_via_smtp(in_memory_smtp: Any) -> None: @pytest.mark.asyncio -async def test_send_email_starttls_uses_ssl( - in_memory_smtp: Any, monkeypatch: pytest.MonkeyPatch -) -> None: +async def test_send_email_starttls_uses_ssl(in_memory_smtp: Any, monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("SMTP_USE_SSL", "True") await send_email( receiver_email="to@invalid", @@ -82,9 +81,7 @@ async def test_send_email_error_missing_sender_email( ) -> None: monkeypatch.delenv("SMTP_SENDER_EMAIL", raising=False) with pytest.raises(ValueError, match="SMTP_SENDER_EMAIL"): - await send_email( - receiver_email="x@y", subject="s", html="

h

" - ) + await send_email(receiver_email="x@y", subject="s", html="

h

") @pytest.mark.asyncio @@ -105,6 +102,49 @@ async def test_send_email_error_missing_html() -> None: await send_email(receiver_email="x@y", subject="s", html=None) +@pytest.mark.asyncio +async def test_send_email_sets_connection_timeout(in_memory_smtp: Any) -> None: + await send_email( + receiver_email="to@invalid", + subject="Hi", + html="

x

", + ) + assert in_memory_smtp.connection_kwargs[0].get("timeout") == 30.0 + + +@pytest.mark.asyncio +async def test_send_email_timeout_env_override(in_memory_smtp: Any, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("SMTP_TIMEOUT", "5") + monkeypatch.setenv("SMTP_USE_SSL", "True") + await send_email( + receiver_email="to@invalid", + subject="Hi", + html="

x

", + ) + assert in_memory_smtp.connection_kwargs[0].get("timeout") == 5.0 + + +@pytest.mark.asyncio +async def test_send_email_malformed_timeout_is_swallowed(in_memory_smtp: Any, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("SMTP_TIMEOUT", "30s") + await send_email( + receiver_email="to@invalid", + subject="Hi", + html="

x

", + ) + assert in_memory_smtp.sent == [] + + +@pytest.mark.asyncio +async def test_send_email_runs_off_event_loop_thread(in_memory_smtp: Any) -> None: + await send_email( + receiver_email="to@invalid", + subject="Hi", + html="

x

", + ) + assert in_memory_smtp.sent[0].thread_ident != threading.get_ident() + + @pytest.mark.asyncio async def test_send_email_smtp_failure_is_swallowed( in_memory_smtp: Any, @@ -113,7 +153,5 @@ async def test_send_email_smtp_failure_is_swallowed( does not raise so a failing email never blocks the proxy. """ in_memory_smtp.raise_on_send = RuntimeError("smtp boom") - await send_email( - receiver_email="to@invalid", subject="Hi", html="

x

" - ) + await send_email(receiver_email="to@invalid", subject="Hi", html="

x

") assert in_memory_smtp.sent == [] From 2c2f4bbfed769985bd0341b242e9d6b811fb9719 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 4 Aug 2026 16:56:37 -0700 Subject: [PATCH 14/16] feat(proxy): add LiteLLM_BudgetWindowSpend table for per-window budget spend Multi-window budgets (budget_limits on keys/teams) currently keep window spend only in cache. Every cold or expired counter recomputes the window by aggregating LiteLLM_SpendLogs, which has no usable index for that query and saturates the DB on large tables (#35766). This adds a LiteLLM_BudgetWindowSpend table holding one row per configured window, keyed (entity_type, entity_id, window_duration), with window_start identifying the period the spend belongs to. Follow-up PRs maintain these rows from the spend update writer and move window budget enforcement reads onto them. --- .../migration.sql | 13 +++++++++++++ .../litellm_proxy_extras/schema.prisma | 12 ++++++++++++ litellm/proxy/schema.prisma | 12 ++++++++++++ schema.prisma | 12 ++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 +- 5 files changed, 50 insertions(+), 1 deletion(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260804162853_add_budget_window_spend_table/migration.sql diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260804162853_add_budget_window_spend_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260804162853_add_budget_window_spend_table/migration.sql new file mode 100644 index 00000000000..45cc927a328 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260804162853_add_budget_window_spend_table/migration.sql @@ -0,0 +1,13 @@ +-- CreateTable +CREATE TABLE IF NOT EXISTS "LiteLLM_BudgetWindowSpend" ( + "entity_type" TEXT NOT NULL, + "entity_id" TEXT NOT NULL, + "window_duration" TEXT NOT NULL, + "window_start" TIMESTAMP(3) NOT NULL, + "spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updated_at" TIMESTAMP(3) NOT NULL, + + CONSTRAINT "LiteLLM_BudgetWindowSpend_pkey" PRIMARY KEY ("entity_type","entity_id","window_duration") +); + diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 48d41edffbd..60223265211 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -664,6 +664,18 @@ model LiteLLM_SpendLogs { @@index([session_id]) } +model LiteLLM_BudgetWindowSpend { + entity_type String + entity_id String + window_duration String + window_start DateTime + spend Float @default(0.0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + + @@id([entity_type, entity_id, window_duration]) +} + // View spend, model, api_key per request model LiteLLM_ErrorLogs { request_id String @id @default(uuid()) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 48d41edffbd..60223265211 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -664,6 +664,18 @@ model LiteLLM_SpendLogs { @@index([session_id]) } +model LiteLLM_BudgetWindowSpend { + entity_type String + entity_id String + window_duration String + window_start DateTime + spend Float @default(0.0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + + @@id([entity_type, entity_id, window_duration]) +} + // View spend, model, api_key per request model LiteLLM_ErrorLogs { request_id String @id @default(uuid()) diff --git a/schema.prisma b/schema.prisma index 48d41edffbd..60223265211 100644 --- a/schema.prisma +++ b/schema.prisma @@ -664,6 +664,18 @@ model LiteLLM_SpendLogs { @@index([session_id]) } +model LiteLLM_BudgetWindowSpend { + entity_type String + entity_id String + window_duration String + window_start DateTime + spend Float @default(0.0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + + @@id([entity_type, entity_id, window_duration]) +} + // View spend, model, api_key per request model LiteLLM_ErrorLogs { request_id String @id @default(uuid()) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index c762183dec4..37e97f98e63 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -26193,7 +26193,7 @@ export interface components { * @description Default role assigned to new users created * @default internal_user_viewer */ - user_role: ("proxy_admin" | "proxy_admin_viewer" | "internal_user" | "internal_user_viewer") | null; + user_role: ("internal_user" | "internal_user_viewer" | "proxy_admin" | "proxy_admin_viewer") | null; }; /** * DefaultTeamSSOParams From 882927673555a392c7fd7bbaa9355df0da1fae82 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 29 Aug 2026 11:36:59 -0700 Subject: [PATCH 15/16] chore(migrations): drop the generated comment from the budget window spend migration --- .../20260804162853_add_budget_window_spend_table/migration.sql | 1 - 1 file changed, 1 deletion(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260804162853_add_budget_window_spend_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260804162853_add_budget_window_spend_table/migration.sql index 45cc927a328..c3018006adb 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260804162853_add_budget_window_spend_table/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260804162853_add_budget_window_spend_table/migration.sql @@ -1,4 +1,3 @@ --- CreateTable CREATE TABLE IF NOT EXISTS "LiteLLM_BudgetWindowSpend" ( "entity_type" TEXT NOT NULL, "entity_id" TEXT NOT NULL, From 4c8838696a7135718acf9d126de765d295240456 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 29 Aug 2026 12:27:49 -0700 Subject: [PATCH 16/16] chore(ui): drop unrelated schema.d.ts enum reorder from the window spend schema branch --- ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 37e97f98e63..c762183dec4 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -26193,7 +26193,7 @@ export interface components { * @description Default role assigned to new users created * @default internal_user_viewer */ - user_role: ("internal_user" | "internal_user_viewer" | "proxy_admin" | "proxy_admin_viewer") | null; + user_role: ("proxy_admin" | "proxy_admin_viewer" | "internal_user" | "internal_user_viewer") | null; }; /** * DefaultTeamSSOParams