[Infra] Merge dev branch - #26197
Conversation
[Fix] MCP broker OAuth endpoint access controls
Brings user personal budget and organization budget enforcement in line with the existing key and team patterns, which already read spend from the atomic cross-pod Redis counter.
…lignment [Fix] Align user and org budget spend checks with atomic counter pattern
Greptile SummaryThis PR extends the cross-pod Redis spend-counter infrastructure to cover user and org budget tracking, migrates the
Confidence Score: 4/5Safe to merge after confirming the MCP OAuth endpoint auth addition won't break existing clients. The spend-counter additions are well-structured and consistent with existing patterns. The redirect_uri scheme validation is a clean security improvement. One P1 concern remains: adding user_api_key_auth to three previously-public MCP OAuth endpoints is a backwards-incompatible behaviour change that may break MCP clients attempting to initiate the OAuth flow without a pre-existing LiteLLM key. litellm/proxy/management_endpoints/mcp_management_endpoints.py — the auth addition to OAuth proxy endpoints needs confirmation that no existing MCP client relies on the previously-public behaviour.
|
| Filename | Overview |
|---|---|
| litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py | Adds redirect_uri scheme validation (http/https only) to prevent open-redirect via non-HTTP schemes — clean security hardening. |
| litellm/proxy/auth/auth_checks.py | Replaces DB-object spend read with cross-pod Redis counter (get_current_spend) for user and org budget checks, and tightens comparison from < to >=. |
| litellm/proxy/hooks/max_budget_limiter.py | Replaces stale cache-row lookup with user_api_key_dict fields + get_current_spend counter; correctly exempts team-key requests from personal budget enforcement. |
| litellm/proxy/hooks/proxy_track_cost_callback.py | Passes org_id to increment_spend_counters so org spend is now tracked in the cross-pod Redis counter — minor plumbing fix. |
| litellm/proxy/management_endpoints/mcp_management_endpoints.py | Adds user_api_key_auth to OAuth proxy endpoints (authorize/token/register) and passes request for IP extraction; the auth addition is a backwards-incompatible change that may break MCP clients initiating the OAuth flow. |
| litellm/proxy/proxy_server.py | Extends increment_spend_counters to atomically track spend:user:{user_id} and spend:org:{org_id} Redis counters; counter is incremented for all requests including team-key ones. |
| tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py | Updates mock assertions to expect request=request kwarg in _get_cached_temporary_mcp_server_or_404 calls — correct reflection of the signature change. |
Sequence Diagram
sequenceDiagram
participant Client as MCP Client
participant Proxy as LiteLLM Proxy
participant Redis as Redis (spend_counter_cache)
participant DB as Database
Note over Client,DB: Budget enforcement flow (new counters)
Client->>Proxy: API Request (with LiteLLM key)
Proxy->>Redis: get_current_spend(spend:user:{id})
Redis-->>Proxy: current spend (or fallback from DB cache)
Proxy->>Proxy: user_spend >= user_budget?
alt Over budget
Proxy-->>Client: 429 / BudgetExceededError
else Under budget
Proxy->>DB: Forward request
DB-->>Proxy: Response + cost
Proxy->>Redis: increment spend:key:{token}
Proxy->>Redis: increment spend:team:{team_id}
Proxy->>Redis: increment spend:user:{user_id}
Proxy->>Redis: increment spend:org:{org_id}
Proxy-->>Client: Response
end
Note over Client,DB: MCP OAuth proxy flow (now requires auth)
Client->>Proxy: POST /server/oauth/{id}/register + LiteLLM key
Proxy->>Proxy: user_api_key_auth check
Proxy->>DB: register client with external MCP server
Proxy-->>Client: client credentials
Reviews (2): Last reviewed commit: "Merge remote-tracking branch 'origin/lit..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| 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}", | ||
| ) |
There was a problem hiding this comment.
> → >= on user budget check
The original condition user_budget < user_object.spend blocked requests only when spend strictly exceeds the budget. The new condition user_spend >= user_budget also blocks when spend equals the budget exactly. A user who has spent precisely their budget will now be rejected on the next request, whereas previously they could still make one more call. This is likely intentional stricter enforcement, but it's worth documenting in a comment since it's a visible behavior change for anyone sitting exactly at their limit.
| @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), |
There was a problem hiding this comment.
Depends is redundant with parameter-level Depends
FastAPI deduplicates by callable per request, so having user_api_key_auth in both dependencies=[…] on the route decorator and as an injected parameter means the check only runs once anyway. The decorator entry adds no benefit and can be removed. The same pattern repeats on the /token (line 1410) and /register (line 1454) routes.
| @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), | |
| @router.get( | |
| "/server/oauth/{server_id}/authorize", | |
| include_in_schema=False, | |
| ) | |
| async def mcp_authorize( | |
| request: Request, | |
| server_id: str, | |
| user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), |
GitHub Actions
Inactive
GitHub Actions
Inactive
GitHub Actions
Inactive
GitHub Actions
Inactive
GitHub Actions
Inactive
| @router.get( | ||
| "/server/oauth/{server_id}/authorize", | ||
| include_in_schema=False, | ||
| dependencies=[Depends(user_api_key_auth)], | ||
| ) |
There was a problem hiding this comment.
Adding user_api_key_auth to /authorize, /token, and /register is a backwards-incompatible change (per the repo's policy against breaking changes without a flag). In a standard OAuth 2.0 / MCP OAuth proxy flow, an MCP client (e.g. Claude Desktop) calls the /register dynamic-client-registration endpoint and the /authorize redirect before it possesses a LiteLLM API key — obtaining one is the goal of the flow. Requiring user_api_key_auth on these three endpoints means any MCP client that does not already hold a LiteLLM key will receive a 401 and cannot complete registration or authorization.
If the intent is that only pre-authenticated LiteLLM users can initiate the OAuth proxy flow (i.e. the MCP client always carries the LiteLLM key as a bearer token), this deserves a comment explaining the design and ideally a feature flag (litellm.mcp_oauth_require_auth) so operators who relied on the old public-endpoint behaviour can opt out.
Rule Used: What: avoid backwards-incompatible changes without... (source)
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
[Infra] Merge dev branch
2 participants
Relevant issues
Pre-Submission checklist
Please complete all items before asking a LiteLLM maintainer to review your PR
tests/test_litellm/directory, Adding at least 1 test is a hard requirement - see detailsmake test-unit@greptileaiand received a Confidence Score of at least 4/5 before requesting a maintainer reviewDelays in PR merge?
If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).
CI (LiteLLM team)
Branch creation CI run
Link:
CI run for the last commit
Link:
Merge / cherry-pick CI run
Links:
Screenshots / Proof of Fix
Type
🆕 New Feature
🐛 Bug Fix
🧹 Refactoring
📖 Documentation
🚄 Infrastructure
✅ Test
Changes