Fix bugs that bypasses per-team member budget limit

litellm/proxy/auth/auth_checks.py

@@ -905,6 +905,63 @@ async def get_default_end_user_budget(
         return None
 
 
+@log_db_metrics
+async def get_team_member_default_budget(
+    budget_id: str,
+    prisma_client: Optional[PrismaClient],
+    user_api_key_cache: DualCache,
+) -> Optional[LiteLLM_BudgetTable]:
+    """
+    Fetches the team-level default per-member budget referenced by team.metadata["team_member_budget_id"].
+
+    This budget is applied to team members whose TeamMembership row has no
+    linked budget. Results are cached for performance.
+
+    Args:
+        budget_id: The budget_id pulled from team.metadata["team_member_budget_id"]
+        prisma_client: Database client instance
+        user_api_key_cache: Cache for storing/retrieving budget data
+
+    Returns:
+        LiteLLM_BudgetTable if found, None otherwise
+    """
+    if prisma_client is None:
+        return None
+
+    cache_key = f"team_member_default_budget:{budget_id}"
+
+    cached_budget = await user_api_key_cache.async_get_cache(key=cache_key)
+    if isinstance(cached_budget, LiteLLM_BudgetTable):
+        return cached_budget
+    if isinstance(cached_budget, dict):
+        return LiteLLM_BudgetTable(**cached_budget)
+
+    try:
+        budget_record = await prisma_client.db.litellm_budgettable.find_unique(
+            where={"budget_id": budget_id}
+        )
+
+        if budget_record is None:
+            verbose_proxy_logger.warning(
+                f"Team-default member budget not found in database: {budget_id}"
+            )
+            return None
+
+        await user_api_key_cache.async_set_cache(
+            key=cache_key,
+            value=budget_record.dict(),
+            ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
+        )
+
+        return LiteLLM_BudgetTable(**budget_record.dict())
+
+    except Exception:
+        verbose_proxy_logger.exception(
+            f"Error fetching team-default member budget {budget_id}"
+        )
+        return None
+
+
 async def _apply_default_budget_to_end_user(
     end_user_obj: LiteLLM_EndUserTable,
     prisma_client: PrismaClient,
@@ -3230,13 +3287,31 @@ async def _check_team_member_budget(
             proxy_logging_obj=proxy_logging_obj,
         )
 
+        # Per-member override wins; otherwise fall back to the team-level
+        # default configured via team.metadata["team_member_budget_id"].
+        team_member_budget: Optional[float] = None
         if (
             team_membership is not None
             and team_membership.litellm_budget_table is not None
-            and team_membership.litellm_budget_table.max_budget is not None
         ):
             team_member_budget = team_membership.litellm_budget_table.max_budget
-            team_member_spend = team_membership.spend or 0.0
+        else:
+            default_budget_id = (team_object.metadata or {}).get(
+                "team_member_budget_id"
+            )
+            if isinstance(default_budget_id, str):
+                default_budget = await get_team_member_default_budget(
+                    budget_id=default_budget_id,
+                    prisma_client=prisma_client,
+                    user_api_key_cache=user_api_key_cache,
+                )
+                if default_budget is not None:
+                    team_member_budget = default_budget.max_budget
+
+        if team_member_budget is not None:
+            team_member_spend = (
+                team_membership.spend if team_membership is not None else 0.0
+            ) or 0.0
 
             # Read from cross-pod counter (Redis-first) if available
             from litellm.proxy.proxy_server import get_current_spend

litellm/proxy/management_endpoints/team_endpoints.py

@@ -302,14 +302,15 @@ async def backfill_team_member_budget_entries(
         prisma_client: PrismaClient,
     ) -> None:
         """
-        Create team_memberships entries for existing members that don't have one.
-
-        Called after team_member_budget is set/updated on a team to ensure
-        members who joined before the budget was configured also get budget
-        enforcement.
-
-        Only creates missing entries — does not touch existing memberships
-        (which may carry individual per-member budgets).
+        Ensure every team member has a TeamMembership row linked to the
+        team_member_budget.
+
+        Called after team_member_budget is set/updated on a team. Creates
+        rows for members who don't have one, and populates budget_id on
+        existing rows where it is NULL. Rows with a non-NULL budget_id
+        are left untouched, which preserves per-member overrides but also
+        means rows pointing to a prior team-default budget_id are not
+        migrated to the new one.
         """
         if not members_with_roles:
             return
@@ -347,6 +348,21 @@ async def backfill_team_member_budget_entries(
                 _sanitize_for_log(team_member_budget_id),
             )
 
+        # Heal existing membership rows that predate the team_member_budget
+        # configuration: populate budget_id where it is currently NULL.
+        # Rows with an explicit budget_id (per-member override) are left alone.
+        updated = await prisma_client.db.litellm_teammembership.update_many(
+            where={"team_id": team_id, "budget_id": None},
+            data={"budget_id": team_member_budget_id},
+        )
+        if updated:
+            verbose_proxy_logger.info(
+                "Populated budget_id on %d existing team_memberships for team %s with budget %s",
+                updated,
+                _sanitize_for_log(team_id),
+                _sanitize_for_log(team_member_budget_id),
+            )
+
 
 def _get_default_team_param(field: str) -> Any:
     """

litellm/proxy/proxy_server.py

@@ -1908,34 +1908,102 @@ async def increment_spend_counters(
         )
 
 
+async def _reseed_spend_from_db(counter_key: str) -> float:
+    """
+    Read the authoritative spend for a missing counter from the DB. The
+    counter_key prefix encodes the table to query:
+
+        spend:key:{token}                 -> LiteLLM_VerificationToken.spend
+        spend:team:{team_id}              -> LiteLLM_TeamTable.spend
+        spend:team_member:{uid}:{tid}     -> LiteLLM_TeamMembership.spend
+        spend:user:{user_id}              -> LiteLLM_UserTable.spend
+        spend:org:{org_id}                -> LiteLLM_OrganizationTable.spend
+
+    Returns 0.0 if prisma is unavailable, the row is missing, or the
+    key format is unrecognized. On failure, logs and returns 0.0 rather
+    than raising so the caller can still record the current increment.
+    """
+    if prisma_client is None:
+        return 0.0
+    # Per-window counters (spend:*:window:{duration}) share prefixes with
+    # primary counters but don't correspond to a DB row; their ambiguity
+    # would otherwise be silently parsed as a regular counter and miss.
+    if ":window:" in counter_key:
+        return 0.0
+    try:
+        if counter_key.startswith("spend:key:"):
+            token = counter_key[len("spend:key:") :]
+            row = await prisma_client.db.litellm_verificationtoken.find_unique(
+                where={"token": token}
+            )
+        elif counter_key.startswith("spend:team_member:"):
+            suffix = counter_key[len("spend:team_member:") :]
+            if ":" not in suffix:
+                return 0.0
+            user_id, team_id = suffix.rsplit(":", 1)
+            row = await prisma_client.db.litellm_teammembership.find_unique(
+                where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}}
+            )
+        elif counter_key.startswith("spend:team:"):
+            team_id = counter_key[len("spend:team:") :]
+            row = await prisma_client.db.litellm_teamtable.find_unique(
+                where={"team_id": team_id}
+            )
+        elif counter_key.startswith("spend:user:"):
+            user_id = counter_key[len("spend:user:") :]
+            row = await prisma_client.db.litellm_usertable.find_unique(
+                where={"user_id": user_id}
+            )
+        elif counter_key.startswith("spend:org:"):
+            org_id = counter_key[len("spend:org:") :]
+            row = await prisma_client.db.litellm_organizationtable.find_unique(
+                where={"organization_id": org_id}
+            )
+        else:
+            return 0.0
+    except Exception:
+        verbose_proxy_logger.exception(
+            "Failed to reseed spend counter %s from DB", counter_key
+        )
+        return 0.0
+    if row is None:
+        return 0.0
+    return float(getattr(row, "spend", 0.0) or 0.0)
+
+
 async def _init_and_increment_spend_counter(
     counter_key: str,
     source_cache_key: str,
     increment: float,
 ):
     """
-    Initialize counter from cached object's DB-loaded spend if not yet set,
-    then atomically increment in both in-memory and Redis.
+    Initialize counter from the authoritative DB spend value if not yet
+    set, then atomically increment in both in-memory and Redis.
 
     On first access per pod:
-    1. Check spend_counter_cache (in-memory -> Redis via DualCache for init check)
-    2. If not found anywhere, read base spend from user_api_key_cache (DB-loaded object)
+    1. Check spend_counter_cache (in-memory -> Redis via DualCache)
+    2. If not found, reseed from the DB (`_reseed_spend_from_db`). Falls
+       back to the cached object's `.spend` via user_api_key_cache only
+       if prisma is unavailable, since that value can lag the flusher.
     3. Seed counter via async_increment_cache (not async_set_cache) to avoid a
        check-then-set race: if two pods cold-start simultaneously, both may see
-       the counter as absent and seed it. Using increment instead of set means
-       the worst case is over-counting (conservative — blocks slightly early)
-       rather than under-counting (would allow overspend).
+       the counter as absent and seed it. Using increment means the worst case
+       is over-counting (conservative, blocks slightly early) rather than
+       under-counting (would allow overspend).
     4. Increment atomically (both in-memory + Redis)
     """
     current = await spend_counter_cache.async_get_cache(key=counter_key)
     if current is None:
-        source = await user_api_key_cache.async_get_cache(key=source_cache_key)
-        base_spend = 0.0
-        if source is not None:
-            if isinstance(source, dict):
-                base_spend = source.get("spend", 0.0) or 0.0
-            else:
-                base_spend = getattr(source, "spend", 0.0) or 0.0
+        base_spend = await _reseed_spend_from_db(counter_key)
+        if prisma_client is None:
+            # Best-effort fallback when prisma is unavailable (tests or
+            # early-startup paths). May be stale but avoids resetting to 0.
+            source = await user_api_key_cache.async_get_cache(key=source_cache_key)
+            if source is not None:
+                if isinstance(source, dict):
+                    base_spend = source.get("spend", 0.0) or 0.0
+                else:
+                    base_spend = getattr(source, "spend", 0.0) or 0.0
         if base_spend > 0:
             await spend_counter_cache.async_increment_cache(
                 key=counter_key, value=base_spend