[Infra] Merge dev branch

litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py

@@ -323,6 +323,14 @@ async def authorize_with_server(
         )
 
     parsed = urlparse(redirect_uri)
+    if parsed.scheme not in ("http", "https"):
+        raise HTTPException(
+            status_code=400,
+            detail={
+                "error": "invalid_redirect_uri",
+                "message": "redirect_uri must use http or https scheme",
+            },
+        )
     base_url = urlunparse(parsed._replace(query=""))
     request_base_url = get_request_base_url(request)
     encoded_state = encode_state_with_base_url(

litellm/proxy/auth/auth_checks.py

@@ -626,11 +626,17 @@ async def common_checks(  # noqa: PLR0915
             and user_object.max_budget is not None
         ):
             user_budget = user_object.max_budget
-            if user_budget < user_object.spend:
+            from litellm.proxy.proxy_server import get_current_spend
+
+            user_spend = await get_current_spend(
+                counter_key=f"spend:user:{user_object.user_id}",
+                fallback_spend=user_object.spend or 0.0,
+            )
+            if user_spend >= user_budget:
                 raise litellm.BudgetExceededError(
-                    current_cost=user_object.spend,
+                    current_cost=user_spend,
                     max_budget=user_budget,
-                    message=f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_object.spend}, Budget={user_budget}",
+                    message=f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_spend}, Budget={user_budget}",
                 )
 
         ## 4.2 check team member budget, if team key
@@ -3665,12 +3671,20 @@ async def _organization_max_budget_check(
     if org_max_budget is None or org_max_budget <= 0:
         return
 
+    # Read spend from cross-pod counter (Redis-first) or cached object (fallback)
+    from litellm.proxy.proxy_server import get_current_spend
+
+    org_spend = await get_current_spend(
+        counter_key=f"spend:org:{org_id}",
+        fallback_spend=org_table.spend or 0.0,
+    )
+
     # Check if organization spend exceeds max budget
-    if org_table.spend >= org_max_budget:
+    if org_spend >= org_max_budget:
         # Trigger budget alert
         call_info = CallInfo(
             token=valid_token.token,
-            spend=org_table.spend,
+            spend=org_spend,
             max_budget=org_max_budget,
             user_id=valid_token.user_id,
             team_id=valid_token.team_id,
@@ -3686,9 +3700,9 @@ async def _organization_max_budget_check(
         )
 
         raise litellm.BudgetExceededError(
-            current_cost=org_table.spend,
+            current_cost=org_spend,
             max_budget=org_max_budget,
-            message=f"Budget has been exceeded! Organization={org_id} Current cost: {org_table.spend}, Max budget: {org_max_budget}",
+            message=f"Budget has been exceeded! Organization={org_id} Current cost: {org_spend}, Max budget: {org_max_budget}",
         )
 
 

litellm/proxy/hooks/max_budget_limiter.py

@@ -21,21 +21,31 @@ async def async_pre_call_hook(
     ):
         try:
             verbose_proxy_logger.debug("Inside Max Budget Limiter Pre-Call Hook")
-            cache_key = f"{user_api_key_dict.user_id}_user_api_key_user_id"
-            user_row = await cache.async_get_cache(
-                cache_key, parent_otel_span=user_api_key_dict.parent_otel_span
-            )
-            if user_row is None:  # value not yet cached
-                return
-            max_budget = user_row["max_budget"]
-            curr_spend = user_row["spend"]
+            max_budget = user_api_key_dict.user_max_budget
+            user_id = user_api_key_dict.user_id
 
-            if max_budget is None:
+            if max_budget is None or user_id is None:
                 return
 
-            if curr_spend is None:
+            # Personal budget applies only to non-team requests, matching
+            # the explicit team-key exemption in common_checks section 4.1.
+            if user_api_key_dict.team_id is not None:
                 return
 
+            from litellm.proxy.proxy_server import get_current_spend
+
+            curr_spend = await get_current_spend(
+                counter_key=f"spend:user:{user_id}",
+                fallback_spend=user_api_key_dict.user_spend or 0.0,
+            )
+
+            verbose_proxy_logger.debug(
+                "MaxBudgetLimiter: user_id=%s, spend=%.6f, max=%.6f",
+                user_id,
+                curr_spend,
+                max_budget,
+            )
+
             # CHECK IF REQUEST ALLOWED
             if curr_spend >= max_budget:
                 raise HTTPException(status_code=429, detail="Max budget limit reached.")

litellm/proxy/hooks/proxy_track_cost_callback.py

@@ -213,6 +213,7 @@ async def _PROXY_track_cost_callback(
                         team_id=team_id,
                         user_id=user_id,
                         response_cost=response_cost,
+                        org_id=org_id,
                     )
 
                     # update cache (fire-and-forget for backward compat:

litellm/proxy/management_endpoints/mcp_management_endpoints.py

@@ -1336,18 +1336,24 @@ async def add_session_mcp_server(
 
         return _redact_mcp_credentials(temp_record)
 
-    def _get_cached_temporary_mcp_server_or_404(server_id: str) -> MCPServer:
+    def _get_cached_temporary_mcp_server_or_404(
+        server_id: str, request: Optional[Request] = None
+    ) -> MCPServer:
         server = get_cached_temporary_mcp_server(server_id)
         if server is None:
             # Fall back to real DB/config server (e.g. for the user-side OAuth flow
             # which calls these endpoints with a real server_id, not a temp session id).
             from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
                 global_mcp_server_manager,
             )
+            from litellm.proxy.auth.ip_address_utils import IPAddressUtils
 
+            client_ip = IPAddressUtils.get_mcp_client_ip(request) if request else None
             server = global_mcp_server_manager.get_mcp_server_by_id(
                 server_id
-            ) or global_mcp_server_manager.get_mcp_server_by_name(server_id)
+            ) or global_mcp_server_manager.get_mcp_server_by_name(
+                server_id, client_ip=client_ip
+            )
         if server is None:
             raise HTTPException(
                 status_code=status.HTTP_404_NOT_FOUND,
@@ -1358,10 +1364,12 @@ def _get_cached_temporary_mcp_server_or_404(server_id: str) -> MCPServer:
     @router.get(
         "/server/oauth/{server_id}/authorize",
         include_in_schema=False,
+        dependencies=[Depends(user_api_key_auth)],
     )
     async def mcp_authorize(
         request: Request,
         server_id: str,
+        user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
         client_id: Optional[str] = None,
         redirect_uri: str = Query(...),
         state: str = "",
@@ -1370,7 +1378,7 @@ async def mcp_authorize(
         response_type: Optional[str] = None,
         scope: Optional[str] = None,
     ):
-        mcp_server = _get_cached_temporary_mcp_server_or_404(server_id)
+        mcp_server = _get_cached_temporary_mcp_server_or_404(server_id, request=request)
         # Use the server's stored client_id when the caller doesn't supply one
         resolved_client_id = mcp_server.client_id or client_id or ""
         if not resolved_client_id:
@@ -1399,10 +1407,12 @@ async def mcp_authorize(
     @router.post(
         "/server/oauth/{server_id}/token",
         include_in_schema=False,
+        dependencies=[Depends(user_api_key_auth)],
     )
     async def mcp_token(
         request: Request,
         server_id: str,
+        user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
         grant_type: str = Form(...),
         code: Optional[str] = Form(None),
         redirect_uri: Optional[str] = Form(None),
@@ -1412,7 +1422,7 @@ async def mcp_token(
         refresh_token: Optional[str] = Form(None),
         scope: Optional[str] = Form(None),
     ):
-        mcp_server = _get_cached_temporary_mcp_server_or_404(server_id)
+        mcp_server = _get_cached_temporary_mcp_server_or_404(server_id, request=request)
         resolved_client_id = mcp_server.client_id or client_id or ""
         if not resolved_client_id:
             raise HTTPException(
@@ -1441,9 +1451,14 @@ async def mcp_token(
     @router.post(
         "/server/oauth/{server_id}/register",
         include_in_schema=False,
+        dependencies=[Depends(user_api_key_auth)],
     )
-    async def mcp_register(request: Request, server_id: str):
-        mcp_server = _get_cached_temporary_mcp_server_or_404(server_id)
+    async def mcp_register(
+        request: Request,
+        server_id: str,
+        user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
+    ):
+        mcp_server = _get_cached_temporary_mcp_server_or_404(server_id, request=request)
         request_data = await _read_request_body(request=request)
         data: dict = {**request_data}
 

litellm/proxy/proxy_server.py

@@ -1795,6 +1795,7 @@ async def increment_spend_counters(
     team_id: Optional[str],
     user_id: Optional[str],
     response_cost: Optional[float],
+    org_id: Optional[str] = None,
 ):
     """
     Atomically increment spend counters for budget enforcement.
@@ -1881,6 +1882,20 @@ async def increment_spend_counters(
             increment=response_cost,
         )
 
+    if user_id is not None:
+        await _init_and_increment_spend_counter(
+            counter_key=f"spend:user:{user_id}",
+            source_cache_key=user_id,
+            increment=response_cost,
+        )
+
+    if org_id is not None:
+        await _init_and_increment_spend_counter(
+            counter_key=f"spend:org:{org_id}",
+            source_cache_key=f"org_id:{org_id}",
+            increment=response_cost,
+        )
+
 
 async def _init_and_increment_spend_counter(
     counter_key: str,

tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py

@@ -1486,7 +1486,7 @@ async def test_mcp_authorize_proxies_to_discoverable_endpoint(self):
             )
 
         assert result is authorize_response
-        get_server.assert_called_once_with("server-1")
+        get_server.assert_called_once_with("server-1", request=request)
         authorize_mock.assert_awaited_once_with(
             request=request,
             mcp_server=server,
@@ -1533,7 +1533,7 @@ async def test_mcp_token_proxies_to_exchange_endpoint(self):
             )
 
         assert result is exchange_response
-        get_server.assert_called_once_with("server-1")
+        get_server.assert_called_once_with("server-1", request=request)
         exchange_mock.assert_awaited_once_with(
             request=request,
             mcp_server=server,
@@ -1581,7 +1581,7 @@ async def test_mcp_token_proxies_refresh_token_grant(self):
             )
 
         assert result is exchange_response
-        get_server.assert_called_once_with("server-1")
+        get_server.assert_called_once_with("server-1", request=request)
         exchange_mock.assert_awaited_once_with(
             request=request,
             mcp_server=server,
@@ -1628,7 +1628,7 @@ async def test_mcp_register_proxies_request_body(self):
             result = await mcp_register(request=request, server_id="server-1")
 
         assert result is register_response
-        get_server.assert_called_once_with("server-1")
+        get_server.assert_called_once_with("server-1", request=request)
         read_body.assert_awaited_once_with(request=request)
         register_mock.assert_awaited_once_with(
             request=request,