Merge pull request #32560 from BerriAI/litellm_/org-admins-team-budgets-1d4b26

fix(proxy): resolve team org from team_id so org admins can update team budgets

litellm/proxy/auth/auth_checks.py

@@ -98,7 +98,10 @@ from litellm.repositories.user_repository import UserRepository
 from litellm.router import Router
 from litellm.utils import get_utc_datetime
 
-from .auth_checks_organization import organization_role_based_access_check
+from .auth_checks_organization import (
+    add_team_org_context_to_request_body,
+    organization_role_based_access_check,
+)
 from .auth_utils import get_model_from_request
 
 if TYPE_CHECKING:
@@ -707,10 +710,28 @@ async def common_checks(
     # 10 [OPTIONAL] Organization RBAC checks
     organization_role_based_access_check(user_object=user_object, route=route, request_body=request_body)
 
+    async def _fetch_team_org_id(team_id: str) -> Optional[str]:
+        try:
+            team = await get_team_object(
+                team_id=team_id,
+                prisma_client=prisma_client,
+                user_api_key_cache=user_api_key_cache,
+                proxy_logging_obj=proxy_logging_obj,
+            )
+        except HTTPException:
+            return None
+        return team.organization_id
+
+    request_body_for_route_check = await add_team_org_context_to_request_body(
+        route=route,
+        request_body=request_body,
+        fetch_team_org_id=_fetch_team_org_id,
+    )
+
     _is_route_allowed = _is_api_route_allowed(
         route=route,
         request=request,
-        request_data=request_body,
+        request_data=request_body_for_route_check,
         valid_token=valid_token,
         user_obj=user_object,
     )

litellm/proxy/auth/auth_checks_organization.py

@@ -2,7 +2,7 @@
 Auth Checks for Organizations
 """
 
-from typing import Dict, List, Optional, Tuple
+from typing import Awaitable, Callable, Dict, List, Optional, Tuple
 
 from fastapi import status
 
@@ -170,3 +170,33 @@ def _user_is_org_admin(
 
     # User must be admin of ALL requested orgs, not just any one
     return all(org_id in admin_org_ids for org_id in candidate_org_ids)
+
+
+TEAM_ORG_CONTEXT_ROUTES = frozenset({"/team/update"})
+
+
+async def add_team_org_context_to_request_body(
+    route: str,
+    request_body: dict,
+    fetch_team_org_id: Callable[[str], Awaitable[Optional[str]]],
+) -> dict:
+    """
+    Return a copy of request_body with organization_id resolved from the target
+    team when the route identifies the team by team_id and the caller did not
+    pass organization_id. This lets an org admin of the team's own org reach the
+    org-scoped branch of the route gate (which keys off organization_id) without
+    the client having to send it. Returns request_body unchanged when it does
+    not apply, so callers that already pass organization_id and non-team routes
+    are untouched.
+    """
+    if route not in TEAM_ORG_CONTEXT_ROUTES:
+        return request_body
+    if request_body.get("organization_id"):
+        return request_body
+    team_id = request_body.get("team_id")
+    if not isinstance(team_id, str) or not team_id:
+        return request_body
+    org_id = await fetch_team_org_id(team_id)
+    if not org_id:
+        return request_body
+    return {**request_body, "organization_id": org_id}

tests/proxy_behavior/management/test_team_update.py

@@ -105,27 +105,35 @@ async def test_team_update_authz_matrix(
         assert row.team_alias != MARKER_ALIAS, "denied but team mutated"
 
 
-async def test_team_update_requires_proxy_admin_without_org_context(
+async def test_team_update_org_admin_resolved_from_team_without_org_context(
     proxy_client, prisma, scratch, world
 ):
-    """With no organization_id in the body the route gate has no org context
-    and falls back to proxy-admin-only: an org admin of the team's own org
-    is 401, PROXY_ADMIN is 200."""
+    """With no organization_id in the body the route gate resolves the target
+    team's org from team_id, so an org admin of the team's own org is allowed
+    (200), same as PROXY_ADMIN. A team admin of that same team stays denied
+    (401): the resolution grants org admins access, not team admins."""
     await _seed_target(prisma, world, "alpha", scratch.prefix)
 
-    denied = await proxy_client.post(
+    allowed_org_admin = await proxy_client.post(
         "/team/update",
         headers={"Authorization": f"Bearer {world.keys[Actor.ORG_ADMIN].cleartext}"},
         json={"team_id": scratch.prefix, "team_alias": MARKER_ALIAS},
     )
-    assert denied.status_code == 401, denied.text
+    assert allowed_org_admin.status_code == 200, allowed_org_admin.text
 
-    allowed = await proxy_client.post(
+    allowed_proxy_admin = await proxy_client.post(
         "/team/update",
         headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
         json={"team_id": scratch.prefix, "team_alias": MARKER_ALIAS},
     )
-    assert allowed.status_code == 200, allowed.text
+    assert allowed_proxy_admin.status_code == 200, allowed_proxy_admin.text
+
+    denied_team_admin = await proxy_client.post(
+        "/team/update",
+        headers={"Authorization": f"Bearer {world.keys[Actor.TEAM_ADMIN].cleartext}"},
+        json={"team_id": scratch.prefix, "team_alias": MARKER_ALIAS},
+    )
+    assert denied_team_admin.status_code == 401, denied_team_admin.text
 
 
 # Relocation gate — moving a team to a different org. The scratch team starts

tests/test_litellm/proxy/auth/test_route_checks.py

@@ -2595,6 +2595,135 @@ def test_org_admin_of_multiple_orgs_can_operate_on_both():
     assert _user_is_org_admin({"organizations": ["org-A", "org-B"]}, user_obj) is True
 
 
+# ── LIT-4221: /team/update org-context resolution from team_id ────────────────
+from litellm.proxy.auth.auth_checks_organization import (
+    add_team_org_context_to_request_body,
+)
+
+
+@pytest.mark.asyncio
+async def test_add_team_org_context_resolves_org_from_team():
+    """For /team/update with only team_id, the target team's org is resolved and
+    injected so the org-admin route gate can see it. This is what lets an org
+    admin update a team budget from the Hub UI, which sends team_id, not
+    organization_id (LIT-4221)."""
+
+    async def fetch(team_id: str):
+        assert team_id == "team-1"
+        return "org-1"
+
+    out = await add_team_org_context_to_request_body(
+        route="/team/update",
+        request_body={"team_id": "team-1", "max_budget": 42},
+        fetch_team_org_id=fetch,
+    )
+    assert out == {"team_id": "team-1", "max_budget": 42, "organization_id": "org-1"}
+
+
+@pytest.mark.asyncio
+async def test_add_team_org_context_noop_when_org_id_already_present():
+    """If the caller already passed organization_id, no lookup happens and the
+    body is returned unchanged."""
+
+    async def fetch(team_id: str):
+        raise AssertionError("must not resolve when organization_id is present")
+
+    body = {"team_id": "team-1", "organization_id": "org-explicit"}
+    out = await add_team_org_context_to_request_body(
+        route="/team/update", request_body=body, fetch_team_org_id=fetch
+    )
+    assert out == body
+
+
+@pytest.mark.asyncio
+async def test_add_team_org_context_noop_for_other_routes():
+    """Only /team/update opts into org resolution; other routes are untouched."""
+
+    async def fetch(team_id: str):
+        raise AssertionError("must not resolve for a non-opted-in route")
+
+    body = {"team_id": "team-1"}
+    out = await add_team_org_context_to_request_body(
+        route="/team/delete", request_body=body, fetch_team_org_id=fetch
+    )
+    assert out == body
+
+
+@pytest.mark.asyncio
+async def test_add_team_org_context_noop_when_team_has_no_org():
+    """A standalone team (no org) resolves to None, so nothing is injected and
+    the org-admin branch stays unreachable (no blanket access)."""
+
+    async def fetch(team_id: str):
+        return None
+
+    body = {"team_id": "team-1"}
+    out = await add_team_org_context_to_request_body(
+        route="/team/update", request_body=body, fetch_team_org_id=fetch
+    )
+    assert out == body
+
+
+def test_team_update_gate_allows_org_admin_with_resolved_org():
+    """Post-resolution (organization_id present), an org admin of that org clears
+    the gate for /team/update."""
+    user_obj = _make_org_admin_user("org-1")
+    valid_token = UserAPIKeyAuth(user_id="org-admin-user", user_role=LitellmUserRoles.INTERNAL_USER.value)
+    request = MagicMock(spec=Request)
+    request.method = "POST"
+    request.query_params = {}
+
+    RouteChecks.non_proxy_admin_allowed_routes_check(
+        user_obj=user_obj,
+        _user_role=LitellmUserRoles.INTERNAL_USER.value,
+        route="/team/update",
+        request=request,
+        valid_token=valid_token,
+        request_data={"team_id": "team-1", "organization_id": "org-1"},
+    )
+
+
+def test_team_update_gate_rejects_without_org_context():
+    """Without organization_id (i.e. resolution found no org, or a non-org-admin),
+    the gate still rejects /team/update — the fix adds no blanket allow. Guards
+    against re-widening the route (e.g. dropping it into self_managed_routes)."""
+    user_obj = _make_org_admin_user("org-1")
+    valid_token = UserAPIKeyAuth(user_id="org-admin-user", user_role=LitellmUserRoles.INTERNAL_USER.value)
+    request = MagicMock(spec=Request)
+    request.method = "POST"
+    request.query_params = {}
+
+    with pytest.raises(Exception):
+        RouteChecks.non_proxy_admin_allowed_routes_check(
+            user_obj=user_obj,
+            _user_role=LitellmUserRoles.INTERNAL_USER.value,
+            route="/team/update",
+            request=request,
+            valid_token=valid_token,
+            request_data={"team_id": "team-1", "max_budget": 42},
+        )
+
+
+def test_team_update_gate_rejects_cross_org_admin_with_resolved_org():
+    """Even after the target team's org is resolved, an org admin of a DIFFERENT
+    org is rejected at the gate (no cross-org escalation)."""
+    user_obj = _make_org_admin_user("org-1")
+    valid_token = UserAPIKeyAuth(user_id="org-admin-user", user_role=LitellmUserRoles.INTERNAL_USER.value)
+    request = MagicMock(spec=Request)
+    request.method = "POST"
+    request.query_params = {}
+
+    with pytest.raises(Exception):
+        RouteChecks.non_proxy_admin_allowed_routes_check(
+            user_obj=user_obj,
+            _user_role=LitellmUserRoles.INTERNAL_USER.value,
+            route="/team/update",
+            request=request,
+            valid_token=valid_token,
+            request_data={"team_id": "team-1", "organization_id": "org-2"},
+        )
+
+
 @pytest.mark.asyncio
 async def test_initialize_pass_through_registers_wildcard_for_auth_subpath():
     """