feat(proxy): support YAML model group budgets and deployment attribution

cookbook/model_access_group_budgets/README.md

@@ -0,0 +1,71 @@
+# Track a shared monthly model budget
+
+This guide configures a USD 300 calendar-month soft budget on existing Model Access Groups. It requires PostgreSQL and a proxy with the deployment attribution option. It does not filter by provider or grant model access.
+
+## Configure the deployments
+
+Add the same group and `access_group_spend_attribution: deployment` to each deployment whose spend should count. Include every deployment of a model, including fallback deployments. Keep your existing provider parameters.
+
+```yaml
+model_list:
+  - model_name: budgeted-chat
+    litellm_params:
+      model: openai/gpt-6.1-sol
+      api_key: os.environ/OPENAI_API_KEY
+    model_info:
+      access_groups: [shared-pool]
+      access_group_spend_attribution: deployment
+
+litellm_settings:
+  timezone: Europe/Madrid
+  budget_reset_time: "00:00"
+
+general_settings:
+  database_url: os.environ/DATABASE_URL
+  alerting: [webhook]
+  alert_types: [budget_alerts]
+```
+
+Set `WEBHOOK_URL` to your alert receiver and reload the configuration using your usual deployment process. The timezone and reset time are proxy-wide settings, so check their effect on other budgets first.
+
+Set this option on the models selected by a smart router. A router's classifier or embedding call counts only if its own deployment belongs to the group and enables deployment attribution. Setting the option only on the outer router is insufficient.
+
+Leave the option unset, or set it to `permission`, to retain the existing behavior. In that mode, only groups explicitly granted by an applicable caller allowlist are charged. Deployment attribution changes accounting, not authorization or hard-budget enforcement. Existing permission-based hard limits still apply.
+
+## Create the budget
+
+Use an administrative key against your own proxy.
+
+```bash
+curl --fail-with-body -X PUT "$LITELLM_URL/access_group/shared-pool/budget" \
+  -H "Authorization: Bearer $LITELLM_ADMIN_KEY" \
+  -H 'Content-Type: application/json' \
+  -d '{"soft_budget":300,"budget_duration":"1mo"}'
+```
+
+For a new pool, the response has `max_budget: null` and a `budget_reset_at` on the first of the next month in the configured timezone. A soft budget sends alerts without rejecting requests. Other key, team, organization or deployment limits remain effective. If reusing a pool that already has a hard budget, remove its old budget first or explicitly clear that limit through the budget API. Do not assume setting a soft limit clears an existing hard limit.
+
+The group can come from YAML, PostgreSQL, or both. PostgreSQL remains required for budgets and recorded spend. This does not move YAML models into PostgreSQL or rewrite the YAML file.
+
+## Verify the result
+
+```bash
+curl --fail-with-body "$LITELLM_URL/access_group/list" \
+  -H "Authorization: Bearer $LITELLM_ADMIN_KEY"
+curl --fail-with-body "$LITELLM_URL/access_group/shared-pool/budget" \
+  -H "Authorization: Bearer $LITELLM_ADMIN_KEY"
+```
+
+Call a configured model using an unrestricted or wildcard key, then read the budget again after the spend queue flushes. Confirm that the increment matches the priced request. Multiple applicable groups each receive the full cost, but duplicate membership in one group is counted once. There is no extra charge for a permission grant to the same group.
+
+Crossing the soft limit triggers the existing budget alert delivery mechanism after the cost callback updates the counters. Configure `budget_alerts` and a supported destination. Delivery uses the existing alert TTL and is best effort. Concurrent workers can send duplicate alerts, and delivery is not an exactly-once notification service.
+
+## Changes and operational limits
+
+Budget rows and recorded spend survive proxy restarts. Removing the last YAML member does not erase the budget. The group remains listed with no deployments, and its budget can still be read or deleted. Manage YAML membership in YAML. The group membership mutation endpoints still manage PostgreSQL deployments and reject replacement or deletion of groups containing YAML members. Budget deletion remains available.
+
+Attribution resolves the actual deployment ID against the serving router at cost-callback time. The resolved groups are shared by the counter and queued database update. A later YAML reload cannot change that queued charge. A reload *before* the callback can change membership, and removing a deployment before its callback can lose deployment-only attribution. Drain in-flight requests before removing or changing those deployments. Keep configuration consistent across replicas and spend workers. This feature does not retroactively reclassify historical spend.
+
+Only costs reported by the existing LiteLLM cost pipeline can be counted. Streaming uses its final or recovered usage event. Image generation needs supported pricing. Asynchronous video, batches, background responses and provider pass-through require their existing completion-cost tracking and a resolvable deployment ID. Status polls without a charge must not be treated as new inference. This example does not establish coverage for every provider, endpoint or plugin. Inspect spend logs and validate the specific route before relying on its totals.
+
+The existing spend queue, atomic database increments, event deduplication and reset job remain in use. Queue flushes and resets are asynchronous. This is an operational budget, not an invoice ledger with event-time month boundaries, guaranteed delivery during crashes, or retroactive recovery of missing provider costs. Use the existing spend-worker durability controls and reconcile against provider billing where needed.

litellm/proxy/auth/auth_checks.py

@@ -15,6 +15,7 @@ import re
 import time
 from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
 from dataclasses import dataclass
+from datetime import datetime
 from functools import partial
 from itertools import chain
 from types import MappingProxyType
@@ -356,6 +357,12 @@ def _end_user_table(repo: _PrismaTableHolder[_PrismaEndUserRow]) -> _PrismaAuthT
 
 
 class _PrismaMaxBudgetRow(Protocol):
+    @property
+    def soft_budget(self) -> float | None: ...
+
+    @property
+    def budget_reset_at(self) -> datetime | None: ...
+
     @property
     def max_budget(self) -> float | None: ...
 
@@ -2415,6 +2422,8 @@ def _model_access_group_budget(row: _PrismaModelAccessGroupBudgetRow) -> ModelAc
         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,
+        soft_budget=None if budget_table is None else budget_table.soft_budget,
+        budget_reset_at=None if budget_table is None else budget_table.budget_reset_at,
     )
 
 

litellm/proxy/db/db_spend_update_writer.py

@@ -458,15 +458,13 @@ def debitable_model_access_groups(
     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.
-    """
+    """Use the served deployment for opt-in accounting, otherwise narrow the auth grant."""
     ordered: Final = coerce_model_access_groups(attributed)
-    if not ordered:
-        return ()
+    deployment: Final = router.get_model_info(id=served_model_id) if router is not None and served_model_id else None
+    model_info: Final = deployment.get("model_info") if deployment is not None else None
     served: Final = _served_model_access_groups(router=router, served_model_id=served_model_id)
+    if isinstance(model_info, Mapping) and model_info.get("access_group_spend_attribution") == "deployment":
+        return tuple(sorted(served or ()))
     if served is None:
         return ordered
     return tuple(group for group in ordered if group in served)
@@ -513,6 +511,7 @@ class DBSpendUpdateWriter:
         end_time: datetime,
         response_cost: float | None,
         project_id: str | None = None,
+        resolved_model_access_groups: Sequence[str] | None = None,
     ) -> bool:
         """Record the request's spend, answering whether its cost still needs charging.
 
@@ -604,6 +603,7 @@ class DBSpendUpdateWriter:
                     litellm_proxy_budget_name=litellm_proxy_budget_name,
                     payload=payload,
                     request_model_access_groups=get_request_model_access_groups(kwargs),
+                    resolved_model_access_groups=resolved_model_access_groups,
                 )
             )
 
@@ -1023,6 +1023,7 @@ class DBSpendUpdateWriter:
         payload: SpendLogsPayload,
         request_model_access_groups: Sequence[str] = (),
         project_id: str | None = None,
+        resolved_model_access_groups: Sequence[str] | None = None,
     ):
         """
         Runs all 13 spend-update helpers sequentially inside a single asyncio task.
@@ -1116,6 +1117,7 @@ class DBSpendUpdateWriter:
             served_model_id=payload_copy.get("model_id"),
             prisma_client=prisma_client,
             router=get_llm_router(),
+            resolved_model_access_groups=resolved_model_access_groups,
         )
 
         _agent_id_for_spend: Final = payload_copy.get("agent_id")
@@ -1479,6 +1481,7 @@ class DBSpendUpdateWriter:
         served_model_id: str | None,
         prisma_client: PrismaClient | None,
         router: _DeploymentLookup | None = None,
+        resolved_model_access_groups: Sequence[str] | None = None,
     ) -> None:
         """
         Update spend for every model access group this request is billed against.
@@ -1494,11 +1497,16 @@ class DBSpendUpdateWriter:
             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,
-            ):
+            groups: Final = (
+                tuple(dict.fromkeys(resolved_model_access_groups))
+                if resolved_model_access_groups is not None
+                else debitable_model_access_groups(
+                    attributed=request_model_access_groups,
+                    served_model_id=served_model_id,
+                    router=router,
+                )
+            )
+            for model_access_group in groups:
                 await self.spend_update_queue.add_update(
                     update=SpendUpdateQueueItem(
                         entity_type=Litellm_EntityType.MODEL_ACCESS_GROUP,

litellm/proxy/hooks/proxy_track_cost_callback.py

@@ -24,9 +24,11 @@ from litellm.llms.azure.ptu_capacity import is_str_keyed_mapping
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.auth.auth_checks import (
     get_key_object,
+    get_model_access_group_budgets_batch,
     get_team_object,
 )
 from litellm.proxy.auth.route_checks import RouteChecks
+from litellm.proxy.common_utils.user_api_key_cache import model_access_group_spend_counter_key
 from litellm.proxy.db.db_lookup_gate import DBLookupDeadlineExceeded
 from litellm.proxy.db.db_spend_update_writer import (
     DBSpendUpdateWriter,
@@ -56,6 +58,8 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import (  # noqa: F401  #
     should_store_prompts_and_responses_in_spend_logs,
 )
 from litellm.proxy.utils import ProxyUpdateSpend
+from litellm.types.integrations.slack_alerting import CallInfo, Litellm_EntityType
+from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget
 from litellm.types.utils import (
     CallTypes,
     LiteLLMBatch,
@@ -465,6 +469,8 @@ class _ProxyDBLogger(CustomLogger):
                     if not charged:
                         return
 
+                    await alert_model_access_group_soft_budgets(model_access_groups, proxy_logging_obj)
+
                     # update cache (fire-and-forget for backward compat:
                     # cached object fields, soft budget alerts, etc.)
                     asyncio.create_task(
@@ -864,6 +870,7 @@ async def _update_database_and_spend_counters_in_batch(
             end_time=end_time,
             org_id=org_id,
             project_id=project_id,
+            resolved_model_access_groups=model_access_groups,
         )
     except Exception:
         if budget_reservation is not None:
@@ -963,3 +970,41 @@ async def _invalidate_budget_reservation_counters(
     await invalidate_budget_reservation_counters(
         budget_reservation=budget_reservation,
     )
+
+
+async def alert_model_access_group_soft_budgets(groups: Sequence[str], proxy_logging_obj: "ProxyLogging") -> None:
+    from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
+
+    if not groups or prisma_client is None or not proxy_logging_obj.alerting:
+        return
+    budgets: Final = await get_model_access_group_budgets_batch(
+        access_group_names=groups, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache
+    )
+    await asyncio.gather(
+        *(_alert_model_access_group_soft_budget(budget, proxy_logging_obj) for budget in budgets.values()),
+        return_exceptions=True,
+    )
+
+
+async def _alert_model_access_group_soft_budget(
+    budget: ModelAccessGroupBudget, proxy_logging_obj: "ProxyLogging"
+) -> None:
+    from litellm.proxy.proxy_server import get_current_spend
+
+    if budget.soft_budget is None:
+        return
+    spend: Final = await get_current_spend(
+        counter_key=model_access_group_spend_counter_key(budget.access_group_name), fallback_spend=budget.spend
+    )
+    if spend < budget.soft_budget:
+        return
+    await proxy_logging_obj.budget_alerts(
+        type="soft_budget",
+        user_info=CallInfo(
+            spend=spend,
+            soft_budget=budget.soft_budget,
+            token=f"model_access_group:{budget.access_group_name}:{budget.budget_reset_at}",
+            key_alias=budget.access_group_name,
+            event_group=Litellm_EntityType.MODEL_ACCESS_GROUP,
+        ),
+    )

litellm/proxy/management_endpoints/model_access_group_management_endpoints.py

@@ -19,9 +19,11 @@ from types import MappingProxyType
 from typing import TYPE_CHECKING, Annotated, Any, Final, Protocol
 
 from fastapi import APIRouter, Depends, HTTPException
+from pydantic import BaseModel, ConfigDict
 from typing_extensions import ReadOnly, TypedDict
 
 from litellm._logging import verbose_proxy_logger
+from litellm.litellm_core_utils.litellm_logging import coerce_model_access_groups
 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 (
@@ -273,9 +275,67 @@ async def _delete_model_access_group_budget_row(
     return row is not None
 
 
+class _GroupDeployment(BaseModel):
+    model_config = ConfigDict(frozen=True)
+    model_name: str
+    model_info: Mapping[str, object]
+
+
+def access_groups_from_deployments(deployments: Sequence[_GroupDeployment]) -> Mapping[str, AccessGroupInfo]:
+    memberships: Final = tuple(_deployment_groups(deployment) for deployment in deployments)
+    names: Final = frozenset[str]().union(*(groups for _, groups in memberships))
+    return {
+        group: AccessGroupInfo(
+            access_group=group,
+            model_names=sorted({name for name, groups in memberships if group in groups}),
+            deployment_count=sum(group in groups for _, groups in memberships),
+        )
+        for group in names
+    }
+
+
+def _deployment_groups(deployment: _GroupDeployment) -> tuple[str, frozenset[str]]:
+    declared: Final = deployment.model_info.get("access_groups")
+    return deployment.model_name, frozenset(coerce_model_access_groups(declared))
+
+
+async def get_all_access_groups(prisma_client: PrismaClient) -> Mapping[str, AccessGroupInfo]:
+    from litellm.proxy.proxy_server import llm_router
+
+    rows: Final = await _model_table(prisma_client).find_many()
+    database_ids: Final = frozenset(row.model_id for row in rows)
+    database: Final = tuple(
+        _GroupDeployment(model_name=row.model_name, model_info=dict(model_info_as_mapping(row.model_info) or {}))
+        for row in rows
+    )
+    live: Final = llm_router.get_model_list() or () if llm_router is not None else ()
+    configured: Final = tuple(_GroupDeployment.model_validate(deployment) for deployment in live)
+    config_only: Final = tuple(
+        deployment
+        for deployment in configured
+        if deployment.model_info.get("id") not in database_ids and deployment.model_info.get("db_model") is not True
+    )
+    return access_groups_from_deployments((*database, *config_only))
+
+
+async def _reject_config_membership_mutation(
+    access_group: str, database_groups: Mapping[str, AccessGroupInfo], prisma_client: PrismaClient
+) -> None:
+    all_groups: Final = await get_all_access_groups(prisma_client)
+    discovered: Final = all_groups.get(access_group)
+    database: Final = database_groups.get(access_group)
+    if discovered is not None and discovered.deployment_count > (database.deployment_count if database else 0):
+        raise _http_error(
+            400, "This group contains config deployments. Edit their access_groups in config.yaml instead."
+        )
+
+
 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:
+    access_groups_map: Final = await get_all_access_groups(prisma_client=prisma_client)
+    if (
+        access_group not in access_groups_map
+        and await _model_access_group_budget_row(access_group, prisma_client) is None
+    ):
         raise _http_error(404, f"Access group '{access_group}' not found")
 
 
@@ -529,41 +589,14 @@ def remove_access_group_from_deployment(model_info: dict[str, Any], access_group
 
 async def get_all_access_groups_from_db(
     prisma_client: PrismaClient,
-) -> dict[str, AccessGroupInfo]:
-    """
-    Get all access groups from the database.
-
-    Returns:
-        Dict[str, AccessGroupInfo]: Dictionary mapping access_group name to info
-    """
-    # Get all deployments
-    deployments: Final = await ModelRepository(prisma_client).table.find_many()
-
-    # Build access group map
-    model_names_by_group: Final[dict[str, list[str]]] = {}
-
-    for deployment in deployments:
-        model_info = deployment.model_info or {}
-        access_groups = model_info.get(  # pyright: ignore[reportAttributeAccessIssue]  # Json reads back as a dict
-            "access_groups", []
+) -> Mapping[str, AccessGroupInfo]:
+    deployments: Final = await _model_table(prisma_client).find_many()
+    return access_groups_from_deployments(
+        tuple(
+            _GroupDeployment(model_name=row.model_name, model_info=dict(model_info_as_mapping(row.model_info) or {}))
+            for row in deployments
         )
-        model_name = deployment.model_name
-
-        for access_group in access_groups:
-            if access_group not in model_names_by_group:
-                model_names_by_group[access_group] = []
-
-            model_names_by_group[access_group].append(model_name)
-
-    # Convert to AccessGroupInfo objects
-    return {
-        access_group: AccessGroupInfo(
-            access_group=access_group,
-            model_names=sorted(frozenset(model_names)),
-            deployment_count=len(model_names),
-        )
-        for access_group, model_names in model_names_by_group.items()
-    }
+    )
 
 
 @router.post(
@@ -653,7 +686,7 @@ async def create_model_group(
 
     try:
         # Check if access group already exists
-        existing_access_groups: Final = await get_all_access_groups_from_db(prisma_client=prisma_client)
+        existing_access_groups: Final = await get_all_access_groups(prisma_client=prisma_client)
 
         if data.access_group in existing_access_groups:
             raise HTTPException(
@@ -744,11 +777,20 @@ async def list_access_groups(
         )
 
     try:
-        access_groups_map: Final = await get_all_access_groups_from_db(prisma_client=prisma_client)
+        access_groups_map: Final = await get_all_access_groups(prisma_client=prisma_client)
         budget_rows: Final = await _model_access_group_budget_rows(prisma_client)
 
         access_groups_list: Final = sorted(
-            (_with_budget(info, budget_rows.get(info.access_group)) for info in access_groups_map.values()),
+            (
+                _with_budget(info, budget_rows.get(info.access_group))
+                for info in {
+                    **{
+                        name: AccessGroupInfo(access_group=name, model_names=[], deployment_count=0)
+                        for name in budget_rows
+                    },
+                    **access_groups_map,
+                }.values()
+            ),
             key=lambda group: group.access_group,
         )
 
@@ -798,17 +840,20 @@ async def get_access_group_info(
         )
 
     try:
-        access_groups_map: Final = await get_all_access_groups_from_db(prisma_client=prisma_client)
+        access_groups_map: Final = await get_all_access_groups(prisma_client=prisma_client)
 
-        if access_group not in access_groups_map:
+        budget_row: Final = await _model_access_group_budget_row(access_group, prisma_client)
+        if access_group not in access_groups_map and budget_row is None:
             raise HTTPException(
                 status_code=404,
                 detail={"error": f"Access group '{access_group}' not found"},
             )
 
         return _with_budget(
-            access_groups_map[access_group],
-            await _model_access_group_budget_row(access_group, prisma_client),
+            access_groups_map.get(
+                access_group, AccessGroupInfo(access_group=access_group, model_names=[], deployment_count=0)
+            ),
+            budget_row,
         )
 
     except HTTPException:
@@ -884,6 +929,7 @@ async def update_access_group(
     # Validation: Check if access group exists
     try:
         access_groups_map: Final = await get_all_access_groups_from_db(prisma_client=prisma_client)
+        await _reject_config_membership_mutation(access_group, access_groups_map, prisma_client)
         if access_group not in access_groups_map:
             raise HTTPException(
                 status_code=404,
@@ -1018,6 +1064,7 @@ async def delete_access_group(
     # Validation: Check if access group exists
     try:
         access_groups_map: Final = await get_all_access_groups_from_db(prisma_client=prisma_client)
+        await _reject_config_membership_mutation(access_group, access_groups_map, prisma_client)
         if access_group not in access_groups_map:
             raise HTTPException(
                 status_code=404,

litellm/proxy/management_helpers/utils.py

@@ -246,6 +246,11 @@ async def handle_budget_for_entity(
             _budget: Final[_PrismaBudgetRecord] = await BudgetRepository(prisma_client).table.create(
                 data={
                     **new_budget_data,
+                    **(
+                        {"budget_reset_at": get_budget_reset_time(budget_row.budget_duration)}
+                        if budget_row.budget_duration is not None
+                        else {}
+                    ),
                     "created_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
                     "updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
                 }

litellm/proxy/proxy_server.py

@@ -3069,8 +3069,9 @@ def cost_tracking():
         from litellm.integrations.shadow_eval_logger import ShadowEvalLogger
 
         spend_event_producer = build_spend_event_producer(CollectorSettings(), fallback=run_spend_event)
-        litellm.logging_callback_manager.add_litellm_callback(ProxyDBLogger(spend_event_producer))
-        litellm.logging_callback_manager.add_litellm_async_success_callback(ProxyDBLogger(spend_event_producer))
+        spend_logger: Final = ProxyDBLogger(spend_event_producer)
+        litellm.logging_callback_manager.add_litellm_callback(spend_logger)
+        litellm.logging_callback_manager.add_litellm_async_success_callback(spend_logger)
         litellm.logging_callback_manager.add_litellm_callback(ShadowEvalLogger())
 
 

litellm/types/proxy/model_access_group_budget.py

@@ -2,6 +2,8 @@
 
 from __future__ import annotations
 
+from datetime import datetime
+
 from litellm.types.llms.base import LiteLLMBaseModel
 
 
@@ -17,3 +19,6 @@ class ModelAccessGroupBudget(LiteLLMBaseModel):
     access_group_name: str
     spend: float = 0.0
     max_budget: float | None = None
+
+    soft_budget: float | None = None
+    budget_reset_at: datetime | None = None

litellm/types/router.py

@@ -275,6 +275,7 @@ class ModelInfo(MirroredPricingParams):
     # the model_name that can be used by the team when making LLM calls
     team_public_model_name: str | None = None
     member_auto_router: bool = False
+    access_group_spend_attribution: Literal["permission", "deployment"] = "permission"
 
     # admin-toggled pause flag; mirrors LiteLLM_ProxyModelTable.blocked
     blocked: bool | None = None

tests/integration/authorization/test_access_group_budget.py

@@ -429,3 +429,195 @@ def test_reservation_refuses_a_concurrent_call_while_an_in_flight_call_holds_the
             wire,
             f"Budget has been exceeded! Model access group={group.name} Current cost: {_MAX_BUDGET}, Max budget: {_MAX_BUDGET}",
         )
+
+
+def test_yaml_deployment_accounting_includes_router_subcalls_and_survives_restart(tmp_path: Path) -> None:
+    group: Final = f"yaml-pool-{uuid.uuid4().hex}"
+    target: Final = f"target-{uuid.uuid4().hex}"
+    judge_name: Final = f"judge-{uuid.uuid4().hex}"
+    smart: Final = f"smart-{uuid.uuid4().hex}"
+
+    def classifier(request: Request) -> Reply:
+        response: Final = _billed(request)
+        if request.method == "GET":
+            return response
+        body: Final = json.loads(response.body)
+        return Reply(
+            body=json.dumps(
+                {
+                    **body,
+                    "choices": [
+                        {
+                            "index": 0,
+                            "message": {"role": "assistant", "content": '{"tier":"COMPLEX"}'},
+                            "finish_reason": "stop",
+                        }
+                    ],
+                }
+            ).encode()
+        )
+
+    def receive_alert(request: Request) -> Reply:
+        assert request.method == "POST"
+        return Reply(body=b"{}")
+
+    with (
+        wire_server(_billed) as upstream,
+        wire_server(classifier) as judge,
+        wire_server(receive_alert) as alerts,
+        httpx.Client() as client,
+    ):
+        seed: Final = Gateway(client, "sk-local-group-budget-test", upstream.url)
+        deployments: Final = [
+            {
+                "model_name": name,
+                "litellm_params": {
+                    "model": f"openai/{_PROVIDER_MODEL}",
+                    "api_base": f"{url}/v1",
+                    "api_key": "synthetic",
+                    "input_cost_per_token": _INPUT_RATE,
+                    "output_cost_per_token": _OUTPUT_RATE,
+                },
+                "model_info": {"id": name, "access_groups": [group], "access_group_spend_attribution": "deployment"},
+            }
+            for name, url in ((target, upstream.url), (judge_name, judge.url))
+        ]
+        config: Final = {
+            "model_list": [
+                *deployments,
+                {
+                    "model_name": smart,
+                    "litellm_params": {
+                        "model": "auto_router/complexity_router",
+                        "complexity_router_config": {
+                            "classifier_type": "llm",
+                            "classifier_llm_config": {"model": judge_name, "timeout_ms": 20000},
+                            "tiers": {"SIMPLE": target, "MEDIUM": target, "COMPLEX": target, "REASONING": target},
+                        },
+                    },
+                },
+            ],
+            "general_settings": {
+                "database_url": "os.environ/DATABASE_URL",
+                "master_key": seed.key,
+                "store_model_in_db": True,
+                "alerting": ["webhook"],
+                "alert_types": ["budget_alerts"],
+            },
+            "litellm_settings": {"timezone": "UTC"},
+        }
+        path: Final = tmp_path / "yaml-groups.yaml"
+        path.write_text(yaml.safe_dump(config))
+        with owned_proxy(seed, tmp_path, {"WEBHOOK_URL": alerts.url}, config=path) as candidate:
+            created: Final = candidate.request(
+                "PUT", f"/access_group/{group}/budget", {"soft_budget": _CALL_COST / 2, "budget_duration": "1mo"}
+            )
+            assert created.status_code == 200, created.text
+            budget: Final = created.json()["budget"]
+            assert budget["max_budget"] is None
+            assert budget["budget_reset_at"][8:10] == "01"
+            with candidate.scenario() as scenario:
+                key: Final = scenario.key(models=["*"])
+                for stream in (False, True):
+                    response: Final = candidate.request(
+                        "POST",
+                        "/v1/chat/completions",
+                        {
+                            "model": smart,
+                            "stream": stream,
+                            "messages": [{"role": "user", "content": f"route {uuid.uuid4().hex}"}],
+                        },
+                        key=key,
+                    )
+                    assert response.status_code == 200, response.text
+                    assert "billed" in response.text
+
+                def call(index: int) -> int:
+                    result: Final = candidate.request(
+                        "POST",
+                        "/v1/chat/completions",
+                        {
+                            "model": target,
+                            "messages": [{"role": "user", "content": str(index)}],
+                        },
+                        key=key,
+                    )
+                    assert result.status_code == 200, result.text
+                    return result.status_code
+
+                with ThreadPoolExecutor(max_workers=4) as executor:
+                    assert tuple(executor.map(call, range(4))) == (200, 200, 200, 200)
+                assert len(_provider_calls(upstream)) == 6
+                assert len(_provider_calls(judge)) == 2
+                spent: Final = eventually(
+                    lambda: candidate.request("GET", f"/access_group/{group}/budget").json()["spend"],
+                    lambda value: value >= 8 * _CALL_COST - 1e-10,
+                    seconds=70,
+                )
+                assert spent == pytest.approx(8 * _CALL_COST)
+                notifications: Final = tuple(json.loads(request.body) for request in alerts.drain())
+                assert any(
+                    event.get("event") == "soft_budget_crossed"
+                    and event.get("key_alias") == group
+                    and event.get("event_group") == "model_access_group"
+                    for event in notifications
+                ), notifications
+        with owned_proxy(seed, tmp_path, {}, config=path) as restarted:
+            persisted: Final = restarted.request("GET", f"/access_group/{group}/budget")
+            assert persisted.status_code == 200, persisted.text
+            assert persisted.json()["spend"] == pytest.approx(8 * _CALL_COST)
+            assert persisted.json()["budget"]["budget_id"] == budget["budget_id"]
+            assert restarted.request("DELETE", f"/access_group/{group}/budget").status_code == 200
+            assert restarted.request("POST", "/budget/delete", {"id": budget["budget_id"]}).status_code == 200
+
+
+def test_yaml_image_group_counts_identical_generations_once_each(tmp_path: Path) -> None:
+    group: Final = f"image-pool-{uuid.uuid4().hex}"
+    name: Final = f"images-{uuid.uuid4().hex}"
+    price: Final = 0.25
+
+    def image(request: Request) -> Reply:
+        if request.method == "GET":
+            return Reply(body=b'{"object":"list","data":[]}')
+        assert request.target == "/images/generations"
+        body: Final = json.loads(request.body)
+        assert "access_group_spend_attribution" not in body
+        return Reply(body=b'{"created":1700000000,"data":[{"b64_json":"aW1n"}]}')
+
+    with wire_server(image) as upstream, httpx.Client() as client:
+        seed: Final = Gateway(client, "sk-local-image-budget-test", upstream.url)
+        config: Final = {
+            "model_list": [
+                {
+                    "model_name": name,
+                    "litellm_params": {
+                        "model": "openai/dall-e-3",
+                        "api_base": upstream.url,
+                        "api_key": "synthetic",
+                        "output_cost_per_image": price,
+                    },
+                    "model_info": {"access_groups": [group], "access_group_spend_attribution": "deployment"},
+                }
+            ],
+            "general_settings": {"database_url": "os.environ/DATABASE_URL", "master_key": seed.key},
+        }
+        path: Final = tmp_path / "image-group.yaml"
+        path.write_text(yaml.safe_dump(config))
+        with owned_proxy(seed, tmp_path, {}, config=path) as candidate, candidate.scenario() as scenario:
+            created: Final = candidate.request("PUT", f"/access_group/{group}/budget", {"soft_budget": 0.1})
+            assert created.status_code == 200, created.text
+            budget_id: Final = created.json()["budget"]["budget_id"]
+            scenario.cleanups.callback(scenario.delete_budget, budget_id)
+            scenario.cleanups.callback(candidate.request, "DELETE", f"/access_group/{group}/budget")
+            key: Final = scenario.key(models=["*"])
+            body: Final = {"model": name, "prompt": "a scripted sea otter", "size": "1024x1024", "n": 1}
+            for _ in range(2):
+                response: Final = candidate.request("POST", "/v1/images/generations", body, key=key)
+                assert response.status_code == 200, response.text
+            spent: Final = eventually(
+                lambda: candidate.request("GET", f"/access_group/{group}/budget").json()["spend"],
+                lambda value: value >= 2 * price,
+                seconds=70,
+            )
+            assert spent == pytest.approx(2 * price)
+            assert len(tuple(request for request in upstream.drain() if request.method == "POST")) == 2

tests/unit/proxy/auth/test_model_access_group_budgets.py

@@ -311,7 +311,7 @@ class _MagBudgetRow:
     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)
+        self.litellm_budget_table = None if max_budget is None else SimpleNamespace(max_budget=max_budget, soft_budget=None, budget_reset_at=None)
 
 
 class _RecordingPrismaClient:

tests/unit/proxy/db/test_model_access_group_spend.py

@@ -511,3 +511,47 @@ async def test_batch_database_updates_enqueues_nothing_without_access_groups():
 
     updates = await _drain(writer.spend_update_queue)
     assert [update for update in updates if update["entity_type"] is Litellm_EntityType.MODEL_ACCESS_GROUP] == []
+
+
+@pytest.mark.parametrize("attributed", [(), ("pool",), ("wrong-pool",), ("pool", "pool")])
+def test_deployment_accounting_charges_actual_groups_independently_of_grants(attributed):
+    from typing import Final
+    from litellm import Router
+
+    router: Final = Router(
+        model_list=[
+            {
+                "model_name": "selected-by-smart-router",
+                "litellm_params": {"model": "openai/group-budget-probe", "api_key": "synthetic"},
+                "model_info": {
+                    "id": "actual",
+                    "access_groups": ["pool", "pool", "other"],
+                    "access_group_spend_attribution": "deployment",
+                },
+            }
+        ]
+    )
+    assert debitable_model_access_groups(attributed, "actual", router) == ("other", "pool")
+    assert debitable_model_access_groups((), "unknown", router) == ()
+
+
+@pytest.mark.asyncio
+async def test_resolved_group_charge_is_not_recomputed_after_yaml_reload():
+    from typing import Final
+
+    writer: Final = DBSpendUpdateWriter()
+    await writer._update_model_access_group_db(
+        response_cost=0.42,
+        request_model_access_groups=(),
+        served_model_id="deployment-1",
+        prisma_client=object(),
+        router=_FakeRouter({"deployment-1": ["replacement-group"]}),
+        resolved_model_access_groups=("original-group", "original-group"),
+    )
+    assert await _drain(writer.spend_update_queue) == [
+        SpendUpdateQueueItem(
+            entity_type=Litellm_EntityType.MODEL_ACCESS_GROUP,
+            entity_id="original-group",
+            response_cost=0.42,
+        )
+    ]

tests/unit/proxy/management_endpoints/test_access_group_management.py

@@ -74,7 +74,6 @@ async def test_create_duplicate_access_group_fails():
         patch("litellm.proxy.proxy_server.llm_router", mock_router),
         patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
     ):
-
         # Should raise 409 Conflict
         with pytest.raises(HTTPException) as exc_info:
             await create_model_group(data=request_data, user_api_key_dict=mock_user)
@@ -103,9 +102,7 @@ async def test_create_access_group_with_model_ids_tags_only_specific_deployments
 
     mock_prisma = MagicMock()
     mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[])
-    mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(
-        return_value=deploy_a
-    )
+    mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=deploy_a)
     mock_prisma.db.litellm_proxymodeltable.update = AsyncMock()
 
     mock_user = UserAPIKeyAuth(
@@ -126,15 +123,11 @@ async def test_create_access_group_with_model_ids_tags_only_specific_deployments
             new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)),
         ),
     ):
-        response = await create_model_group(
-            data=request_data, user_api_key_dict=mock_user
-        )
+        response = await create_model_group(data=request_data, user_api_key_dict=mock_user)
 
     assert response.models_updated == 1
     assert response.model_ids == ["deploy-A"]
-    mock_prisma.db.litellm_proxymodeltable.find_unique.assert_called_once_with(
-        where={"model_id": "deploy-A"}
-    )
+    mock_prisma.db.litellm_proxymodeltable.find_unique.assert_called_once_with(where={"model_id": "deploy-A"})
     assert mock_prisma.db.litellm_proxymodeltable.update.call_count == 1
     update_call = mock_prisma.db.litellm_proxymodeltable.update.call_args
     assert update_call.kwargs["where"] == {"model_id": "deploy-A"}
@@ -169,9 +162,7 @@ async def test_create_access_group_with_model_names_tags_all_deployments():
     )
 
     mock_prisma = MagicMock()
-    mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(
-        side_effect=[[], [deploy_a, deploy_b, deploy_c]]
-    )
+    mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(side_effect=[[], [deploy_a, deploy_b, deploy_c]])
     mock_prisma.db.litellm_proxymodeltable.update = AsyncMock()
 
     mock_user = UserAPIKeyAuth(
@@ -179,9 +170,7 @@ async def test_create_access_group_with_model_names_tags_all_deployments():
         user_role=LitellmUserRoles.PROXY_ADMIN,
     )
 
-    request_data = NewModelGroupRequest(
-        access_group="production-models", model_names=["gpt-4o"]
-    )
+    request_data = NewModelGroupRequest(access_group="production-models", model_names=["gpt-4o"])
 
     with (
         patch("litellm.proxy.proxy_server.llm_router", mock_router),
@@ -191,9 +180,7 @@ async def test_create_access_group_with_model_names_tags_all_deployments():
             new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)),
         ),
     ):
-        response = await create_model_group(
-            data=request_data, user_api_key_dict=mock_user
-        )
+        response = await create_model_group(data=request_data, user_api_key_dict=mock_user)
 
     assert response.models_updated == 3
     assert response.model_names == ["gpt-4o"]
@@ -217,9 +204,7 @@ async def test_create_access_group_model_ids_takes_priority_over_model_names():
 
     mock_prisma = MagicMock()
     mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[])
-    mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(
-        return_value=deploy_a
-    )
+    mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=deploy_a)
     mock_prisma.db.litellm_proxymodeltable.update = AsyncMock()
 
     mock_user = UserAPIKeyAuth(
@@ -241,14 +226,10 @@ async def test_create_access_group_model_ids_takes_priority_over_model_names():
             new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)),
         ),
     ):
-        response = await create_model_group(
-            data=request_data, user_api_key_dict=mock_user
-        )
+        response = await create_model_group(data=request_data, user_api_key_dict=mock_user)
 
     assert response.models_updated == 1
-    mock_prisma.db.litellm_proxymodeltable.find_unique.assert_called_once_with(
-        where={"model_id": "deploy-A"}
-    )
+    mock_prisma.db.litellm_proxymodeltable.find_unique.assert_called_once_with(where={"model_id": "deploy-A"})
 
 
 @pytest.mark.asyncio
@@ -367,7 +348,6 @@ async def test_create_access_group_surfaces_dropped_models():
     assert "deploy-A" in str(exc_info.value.detail)
 
 
-
 @pytest.mark.asyncio
 async def test_create_access_group_trusts_reload_snapshot_over_post_lock_fresh_read():
     """A concurrent reconcile sampled after the lock is released must not make this
@@ -396,9 +376,7 @@ async def test_create_access_group_trusts_reload_snapshot_over_post_lock_fresh_r
         patch(
             "litellm.proxy.management_endpoints.model_access_group_management_endpoints.clear_cache",
             new=AsyncMock(
-                return_value=ReconcileOutcome(
-                    still_desired=frozenset({"deploy-A"}), live_after=frozenset({"deploy-A"})
-                )
+                return_value=ReconcileOutcome(still_desired=frozenset({"deploy-A"}), live_after=frozenset({"deploy-A"}))
             ),
         ),
     ):
@@ -575,6 +553,7 @@ 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
@@ -619,6 +598,7 @@ class _FakeBudgetTable:
             max_budget=data.get("max_budget"),
             soft_budget=data.get("soft_budget"),
             budget_duration=data.get("budget_duration"),
+            budget_reset_at=data.get("budget_reset_at"),
         )
         self.rows[budget_id] = row
         self.create_calls.append(dict(data))
@@ -1128,9 +1108,7 @@ async def test_deleting_the_access_group_strips_deployments_before_dropping_the_
     cache = _FakeAuthCache()
 
     with _proxy_with_stubbed_reload(prisma):
-        response = await delete_access_group(
-            access_group="prod-models", user_api_key_dict=_admin(), auth_cache=cache
-        )
+        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 == {}
@@ -1282,10 +1260,115 @@ async def test_deleting_an_access_group_that_never_had_a_budget_still_evicts():
     cache = _FakeAuthCache(journal)
 
     with _proxy_with_stubbed_reload(prisma):
-        response = await delete_access_group(
-            access_group="prod-models", user_api_key_dict=_admin(), auth_cache=cache
-        )
+        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")
+
+
+@pytest.mark.asyncio
+async def test_yaml_group_budget_survives_router_replacement_and_removal():
+    from typing import Final
+
+    from litellm.proxy.management_endpoints.model_access_group_management_endpoints import (
+        delete_access_group_budget,
+        get_access_group_budget,
+        list_access_groups,
+        set_access_group_budget,
+    )
+    from litellm.types.proxy.management_endpoints.model_management_endpoints import AccessGroupBudgetRequest
+
+    prisma: Final = _FakePrismaClient([])
+    cache: Final = _FakeAuthCache([])
+    models: Final = [
+        {
+            "model_name": "yaml-model",
+            "litellm_params": {"model": "openai/group-budget-probe", "api_key": "synthetic"},
+            "model_info": {"id": "yaml-id", "access_groups": ["yaml-pool", "yaml-pool"]},
+        }
+    ]
+    with _proxy(prisma), patch("litellm.proxy.proxy_server.llm_router", Router(model_list=models)):
+        created: Final = await set_access_group_budget(
+            "yaml-pool", AccessGroupBudgetRequest(soft_budget=300, budget_duration="1mo"), _admin(), cache
+        )
+        listed: Final = await list_access_groups(_admin())
+        assert [(g.access_group, g.model_names, g.deployment_count) for g in listed.access_groups] == [
+            ("yaml-pool", ["yaml-model"], 1)
+        ]
+        assert created.budget.soft_budget == 300
+        assert created.budget.max_budget is None
+    with _proxy(prisma), patch("litellm.proxy.proxy_server.llm_router", Router(model_list=models)):
+        assert await get_access_group_budget("yaml-pool") == created
+    with _proxy(prisma), patch("litellm.proxy.proxy_server.llm_router", Router(model_list=[])):
+        assert await get_access_group_budget("yaml-pool") == created
+        remaining: Final = await list_access_groups(_admin())
+        assert [(g.access_group, g.model_names, g.deployment_count) for g in remaining.access_groups] == [
+            ("yaml-pool", [], 0)
+        ]
+        assert (await delete_access_group_budget("yaml-pool", cache)).budget_deleted is True
+
+
+@pytest.mark.asyncio
+async def test_discovery_counts_database_deployment_once_and_uses_current_database_membership():
+    from typing import Final
+    from litellm.proxy.management_endpoints.model_access_group_management_endpoints import get_all_access_groups
+
+    prisma: Final = _FakePrismaClient([], deployments=[_deployment("db-id", "same-name", ("pool",))])
+    router: Final = Router(
+        model_list=[
+            {
+                "model_name": "same-name",
+                "litellm_params": {"model": "openai/probe", "api_key": "synthetic"},
+                "model_info": {"id": "db-id", "db_model": True, "access_groups": ["stale"]},
+            },
+            {
+                "model_name": "same-name",
+                "litellm_params": {"model": "openai/probe", "api_key": "synthetic"},
+                "model_info": {"id": "yaml-id", "access_groups": ["pool"]},
+            },
+        ]
+    )
+    with patch("litellm.proxy.proxy_server.llm_router", router):
+        groups: Final = await get_all_access_groups(prisma)
+    assert tuple(groups) == ("pool",)
+    assert groups["pool"].model_names == ["same-name"]
+    assert groups["pool"].deployment_count == 2
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("operation", ["update", "delete"])
+@pytest.mark.parametrize("with_database_member", [False, True])
+async def test_membership_mutations_preserve_yaml_groups_and_their_budget(operation, with_database_member):
+    from typing import Final
+
+    from fastapi import HTTPException
+    from litellm.proxy.management_endpoints.model_access_group_management_endpoints import (
+        delete_access_group,
+        update_access_group,
+    )
+    from litellm.types.proxy.management_endpoints.model_management_endpoints import UpdateModelGroupRequest
+
+    prisma: Final = _FakePrismaClient([], deployments=[_deployment()] if with_database_member else [])
+    _seed_budget(prisma, "prod-models", spend=12.0, soft_budget=300.0)
+    router: Final = Router(
+        model_list=[
+            {
+                "model_name": "yaml-model",
+                "litellm_params": {"model": "openai/probe", "api_key": "synthetic"},
+                "model_info": {"id": "yaml-id", "access_groups": ["prod-models"]},
+            }
+        ]
+    )
+    with _proxy(prisma), patch("litellm.proxy.proxy_server.llm_router", router):
+        mutation: Final = (
+            delete_access_group("prod-models", _admin(), _FakeAuthCache())
+            if operation == "delete"
+            else update_access_group("prod-models", UpdateModelGroupRequest(model_names=["yaml-model"]), _admin())
+        )
+        with pytest.raises(HTTPException) as error:
+            await mutation
+    assert error.value.status_code == 400
+    assert "config.yaml" in str(error.value.detail)
+    assert prisma.access_group_budget_table.rows["prod-models"].spend == 12.0
+    assert prisma.model_table.updates == []

ui/litellm-dashboard/src/lib/http/schema.d.ts

@@ -50021,6 +50021,12 @@ export interface components {
         };
         /** ModelInfo */
         litellm__types__router__ModelInfo: {
+            /**
+             * Access Group Spend Attribution
+             * @default permission
+             * @enum {string}
+             */
+            access_group_spend_attribution: "permission" | "deployment";
             /** Access Windows */
             access_windows?: components["schemas"]["ModelAccessWindow"][] | null;
             /** Allow Fail Open */