From 0e4fc7c86d6ab90c765ba28765536953ed1c86c3 Mon Sep 17 00:00:00 2001 From: jesus Date: Tue, 8 Sep 2026 17:01:55 +0000 Subject: [PATCH] fix(auth): backfill organization_alias from the team's org for JWT and key auth Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/user_api_key_auth.py | 31 +++ .../proxy/auth/test_user_api_key_auth.py | 176 ++++++++++++++++++ 2 files changed, 207 insertions(+) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index b39b1f330b3..227a578b07a 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -54,6 +54,7 @@ from litellm.proxy.auth.auth_checks import ( get_end_user_object, get_jwt_key_mapping_object, get_object_permission, + get_org_object, get_project_object, get_team_object, get_user_object, @@ -2344,6 +2345,27 @@ def _team_obj_from_token(valid_token: UserAPIKeyAuth) -> LiteLLM_TeamTableCached ) +async def _resolve_org_alias( + org_id: str, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None, + proxy_logging_obj: ProxyLogging, +) -> str | None: + try: + org_object: Final = await get_org_object( + org_id=org_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + except Exception as e: # noqa: BLE001 # fail-safe: attribution must never break auth + verbose_proxy_logger.debug("org alias attribution: org lookup failed: %s", e) + return None + return org_object.organization_alias if org_object is not None else None + + def _token_can_vouch_for_team(valid_token: UserAPIKeyAuth, lookup_error: BaseException) -> bool: """Whether the token's own team fields may stand in for a team that failed to resolve, without widening access. @@ -2601,6 +2623,15 @@ async def _run_centralized_common_checks( if user_api_key_auth_obj.org_id is None and team_object is not None and team_object.organization_id is not None: user_api_key_auth_obj.org_id = team_object.organization_id + if user_api_key_auth_obj.org_id is not None and user_api_key_auth_obj.organization_alias is None: + user_api_key_auth_obj.organization_alias = await _resolve_org_alias( + org_id=user_api_key_auth_obj.org_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + # common_checks identifies admin via user_object, not the token # (non_proxy_admin_allowed_routes_check). JWT admin shortcut and # master_key tokens get admin from the token; the DB row for the diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 541aeabcbcd..f4be8587f6e 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -5288,6 +5288,182 @@ async def test_centralized_common_checks_backfills_org_id_from_team(key_org_id, setattr(_proxy_server_mod, k, v) +@pytest.mark.asyncio +async def test_centralized_common_checks_backfills_org_alias_from_team_org(): + import litellm.proxy.proxy_server as _proxy_server_mod + from fastapi import Request + from starlette.datastructures import URL + + from litellm.proxy._types import LiteLLM_OrganizationTable, LiteLLM_TeamTableCachedObj + + token = UserAPIKeyAuth(api_key="sk-test", user_id="u", team_id="t1") + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + + fetched_team = LiteLLM_TeamTableCachedObj(team_id="t1", organization_id="org-1") + fetched_org = LiteLLM_OrganizationTable( + organization_id="org-1", + organization_alias="acme-org", + budget_id="budget-1", + created_by="test", + updated_by="test", + ) + + attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None) + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + with ( + patch( # test-quality-ok: centralized checks read this auth module seam + "litellm.proxy.auth.user_api_key_auth.get_team_object", + new_callable=AsyncMock, + return_value=fetched_team, + ), + patch( # test-quality-ok: centralized checks read this auth module seam + "litellm.proxy.auth.user_api_key_auth.get_org_object", + new_callable=AsyncMock, + return_value=fetched_org, + ) as mock_get_org, + patch( # test-quality-ok: centralized checks read this auth module seam + "litellm.proxy.auth.user_api_key_auth.common_checks", + new_callable=AsyncMock, + ) as mock_checks, + ): + await _run_centralized_common_checks( + user_api_key_auth_obj=token, + request=request, + request_data={"model": "gpt-4o"}, + route="/chat/completions", + ) + + mock_checks.assert_awaited_once() + assert token.organization_alias == "acme-org" + assert mock_checks.await_args.kwargs["valid_token"].organization_alias == "acme-org" + mock_get_org.assert_awaited_once_with( + org_id="org-1", + prisma_client=None, + user_api_key_cache=attrs["user_api_key_cache"], + parent_otel_span=None, + proxy_logging_obj=attrs["proxy_logging_obj"], + ) + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "org_id,organization_alias,team_org_id,expected_alias", + [ + ("org-pinned", "pinned", "org-from-team", "pinned"), + (None, None, None, None), + ], +) +async def test_centralized_common_checks_preserves_or_omits_org_alias_without_org_lookup( + org_id, organization_alias, team_org_id, expected_alias +): + import litellm.proxy.proxy_server as _proxy_server_mod + from fastapi import Request + from starlette.datastructures import URL + + from litellm.proxy._types import LiteLLM_TeamTableCachedObj + + token = UserAPIKeyAuth( + api_key="sk-test", + user_id="u", + team_id="t1", + org_id=org_id, + organization_alias=organization_alias, + ) + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + + fetched_team = LiteLLM_TeamTableCachedObj(team_id="t1", organization_id=team_org_id) + + attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None) + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + with ( + patch( # test-quality-ok: centralized checks read this auth module seam + "litellm.proxy.auth.user_api_key_auth.get_team_object", + new_callable=AsyncMock, + return_value=fetched_team, + ), + patch( # test-quality-ok: centralized checks read this auth module seam + "litellm.proxy.auth.user_api_key_auth.get_org_object", + new_callable=AsyncMock, + ) as mock_get_org, + patch( # test-quality-ok: centralized checks read this auth module seam + "litellm.proxy.auth.user_api_key_auth.common_checks", + new_callable=AsyncMock, + ), + ): + await _run_centralized_common_checks( + user_api_key_auth_obj=token, + request=request, + request_data={"model": "gpt-4o"}, + route="/chat/completions", + ) + + assert token.organization_alias == expected_alias + mock_get_org.assert_not_awaited() + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + +@pytest.mark.asyncio +async def test_centralized_common_checks_org_alias_lookup_failure_is_swallowed(): + import litellm.proxy.proxy_server as _proxy_server_mod + from fastapi import Request + from starlette.datastructures import URL + + from litellm.proxy._types import LiteLLM_TeamTableCachedObj + + token = UserAPIKeyAuth(api_key="sk-test", user_id="u", team_id="t1") + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + + fetched_team = LiteLLM_TeamTableCachedObj(team_id="t1", organization_id="org-1") + + attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None) + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + with ( + patch( # test-quality-ok: centralized checks read this auth module seam + "litellm.proxy.auth.user_api_key_auth.get_team_object", + new_callable=AsyncMock, + return_value=fetched_team, + ), + patch( # test-quality-ok: centralized checks read this auth module seam + "litellm.proxy.auth.user_api_key_auth.get_org_object", + new_callable=AsyncMock, + side_effect=Exception("db down"), + ), + patch( # test-quality-ok: centralized checks read this auth module seam + "litellm.proxy.auth.user_api_key_auth.common_checks", + new_callable=AsyncMock, + ) as mock_checks, + ): + await _run_centralized_common_checks( + user_api_key_auth_obj=token, + request=request, + request_data={"model": "gpt-4o"}, + route="/chat/completions", + ) + + mock_checks.assert_awaited_once() + assert token.organization_alias is None + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + @pytest.mark.asyncio async def test_cli_session_token_org_backfilled_from_team(monkeypatch): """LIT-4688 root cause: CLI session tokens (from /sso/cli/poll) are minted