diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index e7fee8d6eb2..cd8103abf5e 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/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, ) diff --git a/litellm/proxy/auth/auth_checks_organization.py b/litellm/proxy/auth/auth_checks_organization.py index 44c1d158cbe..b4caff9b8ee 100644 --- a/litellm/proxy/auth/auth_checks_organization.py +++ b/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} diff --git a/tests/proxy_behavior/management/test_team_update.py b/tests/proxy_behavior/management/test_team_update.py index 9cb2b0fecda..23ea89fa74d 100644 --- a/tests/proxy_behavior/management/test_team_update.py +++ b/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 diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index d623149ff6a..204e6a671e3 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/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(): """