From 3503694ef1af0ace681447c46ab5d732fd01fb89 Mon Sep 17 00:00:00 2001
From: Jason Kidd <j.kidd@f5.com>
Date: Wed, 8 Apr 2026 12:51:47 -0700
Subject: [PATCH] fix: add missing prisma_client guard to
 _check_model_max_budget for consistency with main budget-check block

---
 litellm/proxy/auth/user_api_key_auth.py       |  4 ++++
 .../auth/test_auth_model_budget_bypass.py     | 20 +++++++++++++++++++
 2 files changed, 24 insertions(+)

diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py
index fce27bb5cb5..9fb1c2771ab 100644
--- a/litellm/proxy/auth/user_api_key_auth.py
+++ b/litellm/proxy/auth/user_api_key_auth.py
@@ -1084,6 +1084,7 @@ async def _user_api_key_auth_builder(  # noqa: PLR0915
                 request_data=request_data,
                 route=route,
                 model_max_budget_limiter=model_max_budget_limiter,
+                prisma_client=prisma_client,
             )
 
             return valid_token
@@ -1170,6 +1171,7 @@ async def _user_api_key_auth_builder(  # noqa: PLR0915
                 request_data=request_data,
                 route=route,
                 model_max_budget_limiter=model_max_budget_limiter,
+                prisma_client=prisma_client,
             )
 
             return _user_api_key_obj
@@ -1806,6 +1808,7 @@ async def _check_model_max_budget(
     request_data: dict,
     route: str,
     model_max_budget_limiter: "litellm.proxy.hooks.model_max_budget_limiter._PROXY_VirtualKeyModelMaxBudgetLimiter",
+    prisma_client: Optional["PrismaClient"] = None,
 ) -> None:
     """
     Run per-model budget checks for both key-level and end-user-level budgets.
@@ -1826,6 +1829,7 @@ async def _check_model_max_budget(
         max_budget_per_model is not None
         and isinstance(max_budget_per_model, dict)
         and len(max_budget_per_model) > 0
+        and prisma_client is not None
         and valid_token.token is not None
     ):
         await model_max_budget_limiter.is_key_within_model_budget(
diff --git a/tests/test_litellm/proxy/auth/test_auth_model_budget_bypass.py b/tests/test_litellm/proxy/auth/test_auth_model_budget_bypass.py
index 1d94d6ba86e..5dd20a85062 100644
--- a/tests/test_litellm/proxy/auth/test_auth_model_budget_bypass.py
+++ b/tests/test_litellm/proxy/auth/test_auth_model_budget_bypass.py
@@ -156,6 +156,7 @@ class TestCheckModelMaxBudget:
             request_data={"model": "gpt-4"},
             route="/chat/completions",
             model_max_budget_limiter=limiter,
+            prisma_client=MagicMock(),
         )
         limiter.is_key_within_model_budget.assert_awaited_once_with(
             user_api_key_dict=token,
@@ -238,6 +239,7 @@ class TestCheckModelMaxBudget:
                 request_data={"model": "gpt-4"},
                 route="/chat/completions",
                 model_max_budget_limiter=limiter,
+                prisma_client=MagicMock(),
             )
 
     @pytest.mark.asyncio
@@ -258,10 +260,28 @@ class TestCheckModelMaxBudget:
             request_data={"model": "gpt-4"},
             route="/chat/completions",
             model_max_budget_limiter=limiter,
+            prisma_client=MagicMock(),
         )
         limiter.is_key_within_model_budget.assert_awaited_once()
         limiter.is_end_user_within_model_budget.assert_awaited_once()
 
+    @pytest.mark.asyncio
+    async def test_should_skip_key_budget_when_prisma_client_is_none(self):
+        """Key-level check requires prisma_client to be non-None."""
+        limiter = MagicMock(spec=_PROXY_VirtualKeyModelMaxBudgetLimiter)
+        token = UserAPIKeyAuth(
+            token="hashed-key",
+            model_max_budget={"gpt-4": {"budget_limit": 10.0, "time_period": "1d"}},
+        )
+        await _check_model_max_budget(
+            valid_token=token,
+            request_data={"model": "gpt-4"},
+            route="/chat/completions",
+            model_max_budget_limiter=limiter,
+            prisma_client=None,
+        )
+        limiter.is_key_within_model_budget.assert_not_called()
+
     @pytest.mark.asyncio
     async def test_should_skip_when_model_max_budget_empty(self):
         """Empty model_max_budget dict should be treated as no budget."""