Add support for 0 cost models

litellm/proxy/auth/auth_checks.py

@@ -74,6 +74,75 @@ db_cache_expiry = DEFAULT_IN_MEMORY_TTL  # refresh every 5s
 all_routes = LiteLLMRoutes.openai_routes.value + LiteLLMRoutes.management_routes.value
 
 
+def _is_model_cost_zero(
+    model: Optional[Union[str, List[str]]], llm_router: Optional[Router]
+) -> bool:
+    """
+    Check if a model has zero cost (no configured pricing).
+    
+    Uses the router's get_model_group_info method to get pricing information.
+    
+    Args:
+        model: The model name or list of model names
+        llm_router: The LiteLLM router instance
+    
+    Returns:
+        bool: True if all costs for the model are zero, False otherwise
+    """
+    if model is None or llm_router is None:
+        return False
+    
+    # Handle list of models
+    model_list = [model] if isinstance(model, str) else model
+    
+    for model_name in model_list:
+        try:
+            # Use router's get_model_group_info method directly for better reliability
+            model_group_info = llm_router.get_model_group_info(model_group=model_name)
+            
+            if model_group_info is None:
+                # Model not found or no pricing info available
+                # Conservative approach: assume it has cost
+                verbose_proxy_logger.debug(
+                    f"No model group info found for {model_name}, assuming it has cost"
+                )
+                return False
+            
+            # Check costs for this model
+            # Only allow bypass if BOTH costs are explicitly set to 0 (not None)
+            input_cost = model_group_info.input_cost_per_token
+            output_cost = model_group_info.output_cost_per_token
+            
+            # If costs are not explicitly configured (None), assume it has cost
+            if input_cost is None or output_cost is None:
+                verbose_proxy_logger.debug(
+                    f"Model {model_name} has undefined cost (input: {input_cost}, output: {output_cost}), assuming it has cost"
+                )
+                return False
+            
+            # If either cost is non-zero, return False
+            if input_cost > 0 or output_cost > 0:
+                verbose_proxy_logger.debug(
+                    f"Model {model_name} has non-zero cost (input: {input_cost}, output: {output_cost})"
+                )
+                return False
+            
+            # This model has zero cost explicitly configured
+            verbose_proxy_logger.debug(
+                f"Model {model_name} has zero cost explicitly configured (input: {input_cost}, output: {output_cost})"
+            )
+            
+        except Exception as e:
+            # If we can't determine the cost, assume it has cost (conservative approach)
+            verbose_proxy_logger.debug(
+                f"Error checking cost for model {model_name}: {str(e)}, assuming it has cost"
+            )
+            return False
+    
+    # All models checked have zero cost
+    return True
+
+
 async def common_checks(
     request_body: dict,
     team_object: Optional[LiteLLM_TeamTable],
@@ -86,6 +155,7 @@ async def common_checks(
     proxy_logging_obj: ProxyLogging,
     valid_token: Optional[UserAPIKeyAuth],
     request: Request,
+    skip_budget_checks: bool = False,
 ) -> bool:
     """
     Common checks across jwt + key-based auth.
@@ -137,64 +207,66 @@ async def common_checks(
             user_object=user_object,
         )
 
-    # 3. If team is in budget
-    await _team_max_budget_check(
-        team_object=team_object,
-        proxy_logging_obj=proxy_logging_obj,
-        valid_token=valid_token,
-    )
+    # If this is a free model, skip all budget checks
+    if not skip_budget_checks:
+        # 3. If team is in budget
+        await _team_max_budget_check(
+            team_object=team_object,
+            proxy_logging_obj=proxy_logging_obj,
+            valid_token=valid_token,
+        )
 
-    # 3.1. If organization is in budget
-    await _organization_max_budget_check(
-        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,
-    )
+        # 3.1. If organization is in budget
+        await _organization_max_budget_check(
+            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,
+        )
 
-    await _tag_max_budget_check(
-        request_body=request_body,
-        prisma_client=prisma_client,
-        user_api_key_cache=user_api_key_cache,
-        proxy_logging_obj=proxy_logging_obj,
-        valid_token=valid_token,
-    )
+        await _tag_max_budget_check(
+            request_body=request_body,
+            prisma_client=prisma_client,
+            user_api_key_cache=user_api_key_cache,
+            proxy_logging_obj=proxy_logging_obj,
+            valid_token=valid_token,
+        )
 
-    # 4. If user is in budget
-    ## 4.1 check personal budget, if personal key
-    if (
-        (team_object is None or team_object.team_id is None)
-        and user_object is not None
-        and user_object.max_budget is not None
-    ):
-        user_budget = user_object.max_budget
-        if user_budget < user_object.spend:
-            raise litellm.BudgetExceededError(
-                current_cost=user_object.spend,
-                max_budget=user_budget,
-                message=f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_object.spend}, Budget={user_budget}",
-            )
+        # 4. If user is in budget
+        ## 4.1 check personal budget, if personal key
+        if (
+            (team_object is None or team_object.team_id is None)
+            and user_object is not None
+            and user_object.max_budget is not None
+        ):
+            user_budget = user_object.max_budget
+            if user_budget < user_object.spend:
+                raise litellm.BudgetExceededError(
+                    current_cost=user_object.spend,
+                    max_budget=user_budget,
+                    message=f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_object.spend}, Budget={user_budget}",
+                )
 
-    ## 4.2 check team member budget, if team key
-    await _check_team_member_budget(
-        team_object=team_object,
-        user_object=user_object,
-        valid_token=valid_token,
-        prisma_client=prisma_client,
-        user_api_key_cache=user_api_key_cache,
-        proxy_logging_obj=proxy_logging_obj,
-    )
+        ## 4.2 check team member budget, if team key
+        await _check_team_member_budget(
+            team_object=team_object,
+            user_object=user_object,
+            valid_token=valid_token,
+            prisma_client=prisma_client,
+            user_api_key_cache=user_api_key_cache,
+            proxy_logging_obj=proxy_logging_obj,
+        )
 
-    # 5. If end_user ('user' passed to /chat/completions, /embeddings endpoint) is in budget
-    if end_user_object is not None and end_user_object.litellm_budget_table is not None:
-        end_user_budget = end_user_object.litellm_budget_table.max_budget
-        if end_user_budget is not None and end_user_object.spend > end_user_budget:
-            raise litellm.BudgetExceededError(
-                current_cost=end_user_object.spend,
-                max_budget=end_user_budget,
-                message=f"ExceededBudget: End User={end_user_object.user_id} over budget. Spend={end_user_object.spend}, Budget={end_user_budget}",
-            )
+        # 5. If end_user ('user' passed to /chat/completions, /embeddings endpoint) is in budget
+        if end_user_object is not None and end_user_object.litellm_budget_table is not None:
+            end_user_budget = end_user_object.litellm_budget_table.max_budget
+            if end_user_budget is not None and end_user_object.spend > end_user_budget:
+                raise litellm.BudgetExceededError(
+                    current_cost=end_user_object.spend,
+                    max_budget=end_user_budget,
+                    message=f"ExceededBudget: End User={end_user_object.user_id} over budget. Spend={end_user_object.spend}, Budget={end_user_budget}",
+                )
 
     # 6. [OPTIONAL] If 'enforce_user_param' enabled - did developer pass in 'user' param for openai endpoints
     if (
@@ -237,6 +309,7 @@ async def common_checks(
     # 7. [OPTIONAL] If 'litellm.max_budget' is set (>0), is proxy under budget
     if (
         litellm.max_budget > 0
+        and not skip_budget_checks
         and global_proxy_spend is not None
         # only run global budget checks for OpenAI routes
         # Reason - the Admin UI should continue working if the proxy crosses it's global budget

litellm/proxy/auth/user_api_key_auth.py

@@ -586,6 +586,21 @@ async def _user_api_key_auth_builder(  # noqa: PLR0915
                     if team_object is not None
                     else None,
                 )
+                
+                # Check if model has zero cost - if so, skip all budget checks
+                model = get_model_from_request(request_data, route)
+                skip_budget_checks = False
+                if model is not None and llm_router is not None:
+                    from litellm.proxy.auth.auth_checks import _is_model_cost_zero
+                    
+                    skip_budget_checks = _is_model_cost_zero(
+                        model=model, llm_router=llm_router
+                    )
+                    if skip_budget_checks:
+                        verbose_proxy_logger.info(
+                            f"Skipping all budget checks for zero-cost model: {model}"
+                        )
+                
                 # run through common checks
                 _ = await common_checks(
                     request=request,
@@ -599,6 +614,7 @@ async def _user_api_key_auth_builder(  # noqa: PLR0915
                     llm_router=llm_router,
                     proxy_logging_obj=proxy_logging_obj,
                     valid_token=valid_token,
+                    skip_budget_checks=skip_budget_checks,
                 )
 
                 # return UserAPIKeyAuth object
@@ -990,8 +1006,22 @@ async def _user_api_key_auth_builder(  # noqa: PLR0915
                     )
                     user_obj = None
 
+            # Check 2a. Check if model has zero cost - if so, skip all budget checks
+            model = get_model_from_request(request_data, route)
+            skip_budget_checks = False
+            if model is not None and llm_router is not None:
+                from litellm.proxy.auth.auth_checks import _is_model_cost_zero
+                
+                skip_budget_checks = _is_model_cost_zero(
+                    model=model, llm_router=llm_router
+                )
+                if skip_budget_checks:
+                    verbose_proxy_logger.info(
+                        f"Skipping all budget checks for zero-cost model: {model}"
+                    )
+
             # Check 3. Check if user is in their team budget
-            if valid_token.team_member_spend is not None:
+            if not skip_budget_checks and valid_token.team_member_spend is not None:
                 if prisma_client is not None:
                     _cache_key = f"{valid_token.team_id}_{valid_token.user_id}"
 
@@ -1055,46 +1085,47 @@ async def _user_api_key_auth_builder(  # noqa: PLR0915
                         param=abbreviate_api_key(api_key=api_key),
                     )
 
-            # Check 4. Token Spend is under budget
-            if RouteChecks.is_llm_api_route(route=route):
-                await _virtual_key_max_budget_check(
+            if not skip_budget_checks:
+                # Check 4. Token Spend is under budget
+                if RouteChecks.is_llm_api_route(route=route):
+                    await _virtual_key_max_budget_check(
+                        valid_token=valid_token,
+                        proxy_logging_obj=proxy_logging_obj,
+                        user_obj=user_obj,
+                    )
+
+                # Check 5. Max Budget Alert Check
+                await _virtual_key_max_budget_alert_check(
                     valid_token=valid_token,
                     proxy_logging_obj=proxy_logging_obj,
                     user_obj=user_obj,
                 )
 
-            # Check 5. Max Budget Alert Check
-            await _virtual_key_max_budget_alert_check(
-                valid_token=valid_token,
-                proxy_logging_obj=proxy_logging_obj,
-                user_obj=user_obj,
-            )
-
-            # Check 6. Soft Budget Check
-            await _virtual_key_soft_budget_check(
-                valid_token=valid_token,
-                proxy_logging_obj=proxy_logging_obj,
-                user_obj=user_obj,
-            )
-
-            # Check 5. Token Model Spend is under Model budget
-            max_budget_per_model = valid_token.model_max_budget
-            current_model = request_data.get("model", None)
-
-            if (
-                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 current_model is not None
-                and valid_token.token is not None
-            ):
-                ## GET THE SPEND FOR THIS MODEL
-                await model_max_budget_limiter.is_key_within_model_budget(
-                    user_api_key_dict=valid_token,
-                    model=current_model,
+                # Check 6. Soft Budget Check
+                await _virtual_key_soft_budget_check(
+                    valid_token=valid_token,
+                    proxy_logging_obj=proxy_logging_obj,
+                    user_obj=user_obj,
                 )
 
+                # Check 5. Token Model Spend is under Model budget
+                max_budget_per_model = valid_token.model_max_budget
+                current_model = request_data.get("model", None)
+
+                if (
+                    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 current_model is not None
+                    and valid_token.token is not None
+                ):
+                    ## GET THE SPEND FOR THIS MODEL
+                    await model_max_budget_limiter.is_key_within_model_budget(
+                        user_api_key_dict=valid_token,
+                        model=current_model,
+                    )
+
             # Check 6: Additional Common Checks across jwt + key auth
             if valid_token.team_id is not None:
                 _team_obj: Optional[LiteLLM_TeamTable] = LiteLLM_TeamTable(
@@ -1162,6 +1193,7 @@ async def _user_api_key_auth_builder(  # noqa: PLR0915
                 llm_router=llm_router,
                 proxy_logging_obj=proxy_logging_obj,
                 valid_token=valid_token,
+                skip_budget_checks=skip_budget_checks,
             )
             # Token passed all checks
             if valid_token is None:

tests/proxy_unit_tests/test_zero_cost_model_budget_bypass.py

@@ -0,0 +1,590 @@
+"""
+Tests for zero-cost model budget bypass functionality.
+
+When a user exceeds their budget, the system should still allow requests
+to models with zero cost (e.g., on-premises models).
+"""
+
+import asyncio
+from typing import Optional
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+import litellm
+from litellm.caching.caching import DualCache
+from litellm.proxy._types import (
+    LiteLLM_BudgetTable,
+    LiteLLM_EndUserTable,
+    LiteLLM_TeamMembership,
+    LiteLLM_TeamTable,
+    LiteLLM_UserTable,
+    UserAPIKeyAuth,
+)
+from litellm.proxy.auth.auth_checks import (
+    _check_team_member_budget,
+    _is_model_cost_zero,
+    _team_max_budget_check,
+    common_checks,
+)
+from litellm.proxy.utils import ProxyLogging
+from litellm.router import Router
+from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
+
+
+@pytest.fixture
+def mock_router_with_zero_cost_model():
+    """Create a mock router with a zero-cost model."""
+    router = Router(
+        model_list=[
+            {
+                "model_name": "on-prem-model",
+                "litellm_params": {
+                    "model": "ollama/llama2",
+                    "api_base": "http://localhost:11434",
+                    "input_cost_per_token": 0.0,
+                    "output_cost_per_token": 0.0,
+                },
+                "model_info": {
+                    "id": "on-prem-model-id",
+                    "input_cost_per_token": 0.0,
+                    "output_cost_per_token": 0.0,
+                },
+            },
+            {
+                "model_name": "cloud-model",
+                "litellm_params": {
+                    "model": "gpt-3.5-turbo",
+                    "api_key": "sk-test",
+                },
+                "model_info": {
+                    "id": "cloud-model-id",
+                },
+            },
+        ]
+    )
+    return router
+
+
+@pytest.fixture
+def mock_router_with_paid_model():
+    """Create a mock router with only paid models."""
+    router = Router(
+        model_list=[
+            {
+                "model_name": "cloud-model",
+                "litellm_params": {
+                    "model": "gpt-3.5-turbo",
+                    "api_key": "sk-test",
+                },
+                "model_info": {
+                    "id": "cloud-model-id",
+                },
+            }
+        ]
+    )
+    return router
+
+
+@pytest.fixture
+def mock_proxy_logging():
+    """Create a mock ProxyLogging instance."""
+    proxy_logging = ProxyLogging(user_api_key_cache=None)
+    
+    async def mock_budget_alerts(*args, **kwargs):
+        pass
+    
+    proxy_logging.budget_alerts = mock_budget_alerts
+    return proxy_logging
+
+
+class TestIsModelCostZero:
+    """Tests for _is_model_cost_zero helper function."""
+
+    def test_zero_cost_model_in_router(self, mock_router_with_zero_cost_model):
+        """Test that a zero-cost model in router is correctly identified."""
+        result = _is_model_cost_zero(
+            model="on-prem-model", llm_router=mock_router_with_zero_cost_model
+        )
+        assert result is True
+
+    def test_paid_model_in_router(self, mock_router_with_zero_cost_model):
+        """Test that a paid model is correctly identified as non-zero cost."""
+        with patch("litellm.get_model_info") as mock_get_model_info:
+            # Mock the return value for gpt-3.5-turbo
+            mock_get_model_info.return_value = {
+                "input_cost_per_token": 0.0000015,
+                "output_cost_per_token": 0.000002,
+            }
+            result = _is_model_cost_zero(
+                model="cloud-model", llm_router=mock_router_with_zero_cost_model
+            )
+            assert result is False
+
+    def test_none_model(self, mock_router_with_zero_cost_model):
+        """Test that None model returns False."""
+        result = _is_model_cost_zero(
+            model=None, llm_router=mock_router_with_zero_cost_model
+        )
+        assert result is False
+
+    def test_none_router(self):
+        """Test that None router returns False."""
+        result = _is_model_cost_zero(model="some-model", llm_router=None)
+        assert result is False
+
+    def test_list_of_zero_cost_models(self, mock_router_with_zero_cost_model):
+        """Test that a list of zero-cost models returns True."""
+        result = _is_model_cost_zero(
+            model=["on-prem-model"], llm_router=mock_router_with_zero_cost_model
+        )
+        assert result is True
+
+    def test_mixed_cost_models(self, mock_router_with_zero_cost_model):
+        """Test that a list with mixed cost models returns False."""
+        with patch("litellm.get_model_info") as mock_get_model_info:
+            mock_get_model_info.return_value = {
+                "input_cost_per_token": 0.0000015,
+                "output_cost_per_token": 0.000002,
+            }
+            result = _is_model_cost_zero(
+                model=["on-prem-model", "cloud-model"],
+                llm_router=mock_router_with_zero_cost_model,
+            )
+            assert result is False
+
+
+class TestUserBudgetBypass:
+    """Tests for user budget bypass with zero-cost models."""
+
+    @pytest.mark.asyncio
+    async def test_user_over_budget_with_zero_cost_model_allowed(
+        self, mock_router_with_zero_cost_model, mock_proxy_logging
+    ):
+        """Test that user over budget can still use zero-cost models."""
+        user_object = LiteLLM_UserTable(
+            user_id="test-user",
+            spend=100.0,
+            max_budget=50.0,
+        )
+
+        request_body = {"model": "on-prem-model"}
+
+        # Should not raise BudgetExceededError
+        result = await common_checks(
+            request_body=request_body,
+            team_object=None,
+            user_object=user_object,
+            end_user_object=None,
+            global_proxy_spend=None,
+            general_settings={},
+            route="/v1/chat/completions",
+            llm_router=mock_router_with_zero_cost_model,
+            proxy_logging_obj=mock_proxy_logging,
+            valid_token=UserAPIKeyAuth(
+                token="test-token",
+                user_id="test-user",
+            ),
+            request=MagicMock(),
+            skip_budget_checks=True,  # This is set by user_api_key_auth for zero-cost models
+        )
+        assert result is True
+
+    @pytest.mark.asyncio
+    async def test_user_over_budget_with_paid_model_blocked(
+        self, mock_router_with_zero_cost_model, mock_proxy_logging
+    ):
+        """Test that user over budget cannot use paid models."""
+        user_object = LiteLLM_UserTable(
+            user_id="test-user",
+            spend=100.0,
+            max_budget=50.0,
+        )
+
+        request_body = {"model": "cloud-model"}
+
+        with patch("litellm.get_model_info") as mock_get_model_info:
+            mock_get_model_info.return_value = {
+                "input_cost_per_token": 0.0000015,
+                "output_cost_per_token": 0.000002,
+            }
+            with pytest.raises(litellm.BudgetExceededError) as exc_info:
+                await common_checks(
+                    request_body=request_body,
+                    team_object=None,
+                    user_object=user_object,
+                    end_user_object=None,
+                    global_proxy_spend=None,
+                    general_settings={},
+                    route="/v1/chat/completions",
+                    llm_router=mock_router_with_zero_cost_model,
+                    proxy_logging_obj=mock_proxy_logging,
+                    valid_token=UserAPIKeyAuth(
+                        token="test-token",
+                        user_id="test-user",
+                    ),
+                    request=MagicMock(),
+                )
+
+            assert exc_info.value.current_cost == 100.0
+            assert exc_info.value.max_budget == 50.0
+            assert "test-user" in str(exc_info.value)
+
+
+class TestEndUserBudgetBypass:
+    """Tests for end user budget bypass with zero-cost models."""
+
+    @pytest.mark.asyncio
+    async def test_end_user_over_budget_with_zero_cost_model_allowed(
+        self, mock_router_with_zero_cost_model, mock_proxy_logging
+    ):
+        """Test that end user over budget can still use zero-cost models."""
+        end_user_budget = LiteLLM_BudgetTable(max_budget=20.0)
+        end_user_object = LiteLLM_EndUserTable(
+            user_id="end-user-123",
+            spend=50.0,
+            litellm_budget_table=end_user_budget,
+            blocked=False,
+        )
+
+        request_body = {"model": "on-prem-model", "user": "end-user-123"}
+
+        # In the real flow, skip_budget_checks would be set to True for zero-cost models
+        result = await common_checks(
+            request_body=request_body,
+            team_object=None,
+            user_object=None,
+            end_user_object=end_user_object,
+            global_proxy_spend=None,
+            general_settings={},
+            route="/v1/chat/completions",
+            llm_router=mock_router_with_zero_cost_model,
+            proxy_logging_obj=mock_proxy_logging,
+            valid_token=UserAPIKeyAuth(
+                token="test-token",
+            ),
+            request=MagicMock(),
+            skip_budget_checks=True,  # This is set by user_api_key_auth for zero-cost models
+        )
+        assert result is True
+
+    @pytest.mark.asyncio
+    async def test_end_user_over_budget_with_paid_model_blocked(
+        self, mock_router_with_zero_cost_model, mock_proxy_logging
+    ):
+        """Test that end user over budget cannot use paid models."""
+        end_user_budget = LiteLLM_BudgetTable(max_budget=20.0)
+        end_user_object = LiteLLM_EndUserTable(
+            user_id="end-user-123",
+            spend=50.0,
+            litellm_budget_table=end_user_budget,
+            blocked=False,
+        )
+
+        request_body = {"model": "cloud-model", "user": "end-user-123"}
+
+        with patch("litellm.get_model_info") as mock_get_model_info:
+            mock_get_model_info.return_value = {
+                "input_cost_per_token": 0.0000015,
+                "output_cost_per_token": 0.000002,
+            }
+            with pytest.raises(litellm.BudgetExceededError) as exc_info:
+                await common_checks(
+                    request_body=request_body,
+                    team_object=None,
+                    user_object=None,
+                    end_user_object=end_user_object,
+                    global_proxy_spend=None,
+                    general_settings={},
+                    route="/v1/chat/completions",
+                    llm_router=mock_router_with_zero_cost_model,
+                    proxy_logging_obj=mock_proxy_logging,
+                    valid_token=UserAPIKeyAuth(
+                        token="test-token",
+                    ),
+                    request=MagicMock(),
+                )
+
+            assert exc_info.value.current_cost == 50.0
+            assert exc_info.value.max_budget == 20.0
+            assert "end-user-123" in str(exc_info.value)
+
+
+class TestTeamBudgetBypass:
+    """Tests for team budget bypass with zero-cost models."""
+
+    @pytest.mark.asyncio
+    async def test_team_over_budget_with_zero_cost_model_allowed(
+        self, mock_router_with_zero_cost_model, mock_proxy_logging
+    ):
+        """Test that team over budget can still use zero-cost models."""
+        team_object = LiteLLM_TeamTable(
+            team_id="test-team",
+            spend=150.0,
+            max_budget=100.0,
+        )
+
+        valid_token = UserAPIKeyAuth(
+            token="test-token",
+            team_id="test-team",
+        )
+
+        request_body = {"model": "on-prem-model"}
+
+        # In the real flow, skip_budget_checks would be set to True for zero-cost models
+        result = await common_checks(
+            request_body=request_body,
+            team_object=team_object,
+            user_object=None,
+            end_user_object=None,
+            global_proxy_spend=None,
+            general_settings={},
+            route="/v1/chat/completions",
+            llm_router=mock_router_with_zero_cost_model,
+            proxy_logging_obj=mock_proxy_logging,
+            valid_token=valid_token,
+            request=MagicMock(),
+            skip_budget_checks=True,  # This is set by user_api_key_auth for zero-cost models
+        )
+        assert result is True
+
+    @pytest.mark.asyncio
+    async def test_team_over_budget_with_paid_model_blocked(
+        self, mock_router_with_zero_cost_model, mock_proxy_logging
+    ):
+        """Test that team over budget cannot use paid models."""
+        team_object = LiteLLM_TeamTable(
+            team_id="test-team",
+            spend=150.0,
+            max_budget=100.0,
+        )
+
+        valid_token = UserAPIKeyAuth(
+            token="test-token",
+            team_id="test-team",
+        )
+
+        request_body = {"model": "cloud-model"}
+
+        with patch("litellm.get_model_info") as mock_get_model_info:
+            mock_get_model_info.return_value = {
+                "input_cost_per_token": 0.0000015,
+                "output_cost_per_token": 0.000002,
+            }
+            with pytest.raises(litellm.BudgetExceededError) as exc_info:
+                await common_checks(
+                    request_body=request_body,
+                    team_object=team_object,
+                    user_object=None,
+                    end_user_object=None,
+                    global_proxy_spend=None,
+                    general_settings={},
+                    route="/v1/chat/completions",
+                    llm_router=mock_router_with_zero_cost_model,
+                    proxy_logging_obj=mock_proxy_logging,
+                    valid_token=valid_token,
+                    request=MagicMock(),
+                )
+
+            assert exc_info.value.current_cost == 150.0
+            assert exc_info.value.max_budget == 100.0
+            assert "test-team" in str(exc_info.value)
+
+
+class TestTeamMemberBudgetBypass:
+    """Tests for team member budget bypass with zero-cost models."""
+
+    @pytest.mark.asyncio
+    async def test_team_member_over_budget_with_zero_cost_model_allowed(
+        self, mock_router_with_zero_cost_model, mock_proxy_logging
+    ):
+        """Test that team member over budget can still use zero-cost models."""
+        team_object = LiteLLM_TeamTable(
+            team_id="test-team",
+        )
+
+        user_object = LiteLLM_UserTable(
+            user_id="test-user",
+        )
+
+        valid_token = UserAPIKeyAuth(
+            token="test-token",
+            user_id="test-user",
+            team_id="test-team",
+        )
+
+        member_budget = LiteLLM_BudgetTable(max_budget=30.0)
+        team_membership = LiteLLM_TeamMembership(
+            user_id="test-user",
+            team_id="test-team",
+            spend=60.0,
+            litellm_budget_table=member_budget,
+        )
+
+        request_body = {"model": "on-prem-model"}
+
+        # Mock get_team_membership
+        with patch(
+            "litellm.proxy.auth.auth_checks.get_team_membership"
+        ) as mock_get_membership:
+            mock_get_membership.return_value = team_membership
+
+            # In the real flow, skip_budget_checks would be set to True for zero-cost models
+            result = await common_checks(
+                request_body=request_body,
+                team_object=team_object,
+                user_object=user_object,
+                end_user_object=None,
+                global_proxy_spend=None,
+                general_settings={},
+                route="/v1/chat/completions",
+                llm_router=mock_router_with_zero_cost_model,
+                proxy_logging_obj=mock_proxy_logging,
+                valid_token=valid_token,
+                request=MagicMock(),
+                skip_budget_checks=True,  # This is set by user_api_key_auth for zero-cost models
+            )
+            assert result is True
+
+    @pytest.mark.asyncio
+    async def test_team_member_over_budget_with_paid_model_blocked(
+        self, mock_router_with_zero_cost_model, mock_proxy_logging
+    ):
+        """Test that team member over budget cannot use paid models."""
+        team_object = LiteLLM_TeamTable(
+            team_id="test-team",
+        )
+
+        user_object = LiteLLM_UserTable(
+            user_id="test-user",
+        )
+
+        valid_token = UserAPIKeyAuth(
+            token="test-token",
+            user_id="test-user",
+            team_id="test-team",
+        )
+
+        member_budget = LiteLLM_BudgetTable(max_budget=30.0)
+        team_membership = LiteLLM_TeamMembership(
+            user_id="test-user",
+            team_id="test-team",
+            spend=60.0,
+            litellm_budget_table=member_budget,
+        )
+
+        request_body = {"model": "cloud-model"}
+
+        with patch(
+            "litellm.proxy.auth.auth_checks.get_team_membership"
+        ) as mock_get_membership:
+            mock_get_membership.return_value = team_membership
+
+            with patch("litellm.get_model_info") as mock_get_model_info:
+                mock_get_model_info.return_value = {
+                    "input_cost_per_token": 0.0000015,
+                    "output_cost_per_token": 0.000002,
+                }
+                with pytest.raises(litellm.BudgetExceededError) as exc_info:
+                    await common_checks(
+                        request_body=request_body,
+                        team_object=team_object,
+                        user_object=user_object,
+                        end_user_object=None,
+                        global_proxy_spend=None,
+                        general_settings={},
+                        route="/v1/chat/completions",
+                        llm_router=mock_router_with_zero_cost_model,
+                        proxy_logging_obj=mock_proxy_logging,
+                        valid_token=valid_token,
+                        request=MagicMock(),
+                    )
+
+                assert exc_info.value.current_cost == 60.0
+                assert exc_info.value.max_budget == 30.0
+                assert "test-user" in str(exc_info.value)
+                assert "test-team" in str(exc_info.value)
+
+
+class TestEdgeCases:
+    """Tests for edge cases and error handling."""
+
+    def test_model_not_in_router(self, mock_router_with_zero_cost_model):
+        """Test behavior when model is not found in router."""
+        with patch("litellm.get_model_info") as mock_get_model_info:
+            # Simulate model not found
+            mock_get_model_info.side_effect = Exception("Model not found")
+            result = _is_model_cost_zero(
+                model="nonexistent-model", llm_router=mock_router_with_zero_cost_model
+            )
+            # Should return False (conservative approach)
+            assert result is False
+
+    @pytest.mark.asyncio
+    async def test_user_under_budget_with_paid_model_allowed(
+        self, mock_router_with_zero_cost_model, mock_proxy_logging
+    ):
+        """Test that user under budget can use paid models normally."""
+        user_object = LiteLLM_UserTable(
+            user_id="test-user",
+            spend=30.0,
+            max_budget=100.0,
+        )
+
+        request_body = {"model": "cloud-model"}
+
+        with patch("litellm.get_model_info") as mock_get_model_info:
+            mock_get_model_info.return_value = {
+                "input_cost_per_token": 0.0000015,
+                "output_cost_per_token": 0.000002,
+            }
+            # Should not raise BudgetExceededError
+            result = await common_checks(
+                request_body=request_body,
+                team_object=None,
+                user_object=user_object,
+                end_user_object=None,
+                global_proxy_spend=None,
+                general_settings={},
+                route="/v1/chat/completions",
+                llm_router=mock_router_with_zero_cost_model,
+                proxy_logging_obj=mock_proxy_logging,
+                valid_token=UserAPIKeyAuth(
+                    token="test-token",
+                    user_id="test-user",
+                ),
+                request=MagicMock(),
+            )
+            assert result is True
+
+    @pytest.mark.asyncio
+    async def test_user_under_budget_with_zero_cost_model_allowed(
+        self, mock_router_with_zero_cost_model, mock_proxy_logging
+    ):
+        """Test that user under budget can use zero-cost models normally."""
+        user_object = LiteLLM_UserTable(
+            user_id="test-user",
+            spend=30.0,
+            max_budget=100.0,
+        )
+
+        request_body = {"model": "on-prem-model"}
+
+        # Should not raise BudgetExceededError
+        result = await common_checks(
+            request_body=request_body,
+            team_object=None,
+            user_object=user_object,
+            end_user_object=None,
+            global_proxy_spend=None,
+            general_settings={},
+            route="/v1/chat/completions",
+            llm_router=mock_router_with_zero_cost_model,
+            proxy_logging_obj=mock_proxy_logging,
+            valid_token=UserAPIKeyAuth(
+                token="test-token",
+                user_id="test-user",
+            ),
+            request=MagicMock(),
+        )
+        assert result is True