From 5a9805c10248f2dae99c3e891706abe6a398f01a Mon Sep 17 00:00:00 2001 From: jesus Date: Tue, 8 Sep 2026 23:08:49 +0000 Subject: [PATCH 1/3] fix(auth): refresh lite session token team models from team row Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/user_api_key_auth.py | 33 +++++ .../proxy/auth/test_user_api_key_auth.py | 123 ++++++++++++++++++ 2 files changed, 156 insertions(+) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index d8679627bc3..53b382a0180 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1742,6 +1742,15 @@ async def _user_api_key_auth_builder( ): valid_token = ExperimentalUIJWTToken.get_key_object_from_ui_hash_key(api_key) + if valid_token is not None and valid_token.is_session_token: + valid_token = await _refresh_session_token_team_grants( + valid_token=valid_token, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + if ( valid_token is not None and isinstance(valid_token, UserAPIKeyAuth) @@ -2344,6 +2353,30 @@ def _team_obj_from_token(valid_token: UserAPIKeyAuth) -> LiteLLM_TeamTableCached ) +async def _refresh_session_token_team_grants( + valid_token: UserAPIKeyAuth, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None, + proxy_logging_obj: ProxyLogging, +) -> UserAPIKeyAuth: + if valid_token.team_id is None or valid_token.team_id == UI_TEAM_ID or prisma_client is None: + return valid_token + try: + team_obj: Final = await get_team_object( + team_id=valid_token.team_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 HTTPException: + return valid_token + return valid_token.model_copy( + update={"team_models": list(team_obj.models), "team_alias": team_obj.team_alias} # mutable-ok: auth model requires a fresh list + ) + + 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. 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 f1a269cc00a..22d3ccffb41 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 @@ -6111,6 +6111,129 @@ async def test_non_admin_cli_session_token_reaches_production_auth_path(monkeypa assert result.is_session_token is True +@pytest.mark.asyncio +async def test_cli_session_token_refreshes_team_grants_from_team_row(monkeypatch): + monkeypatch.delenv("EXPERIMENTAL_UI_LOGIN", raising=False) + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-salt-cli-test") + + from fastapi import Request + from starlette.datastructures import URL + + from litellm.proxy._types import LiteLLM_TeamTableCachedObj + from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken + + user_info = LiteLLM_UserTable( + user_id="internal-user-1", + user_role=LitellmUserRoles.INTERNAL_USER.value, + models=[], + ) + cli_token = ExperimentalUIJWTToken.get_cli_jwt_auth_token( + user_info, + team_id="team-abc", + team_alias="old-name", + team_models=["old-model"], + ) + fresh_team = LiteLLM_TeamTableCachedObj( + team_id="team-abc", + team_alias="renamed", + models=["new-model"], + ) + attrs = _proxy_attrs_for_db_lookup() + import litellm.proxy.proxy_server as _proxy_server_mod + + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for key, value in attrs.items(): + setattr(_proxy_server_mod, key, value) + request = Request(scope={"type": "http"}) + request._url = URL(url="/v1/models") + with ( + patch( # test-quality-ok: direct collaborator patch exercises the real auth builder + "litellm.proxy.auth.user_api_key_auth.get_team_object", + new_callable=AsyncMock, + return_value=fresh_team, + ), + patch( # test-quality-ok: direct collaborator patch exercises the real auth builder + "litellm.proxy.auth.user_api_key_auth.get_user_object", + new_callable=AsyncMock, + return_value=None, + ), + ): + result = await _user_api_key_auth_builder( + request=request, + api_key=f"Bearer {cli_token}", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={}, + ) + finally: + for key, value in originals.items(): + setattr(_proxy_server_mod, key, value) + + assert result.team_models == ["new-model"] + assert result.team_alias == "renamed" + + +@pytest.mark.asyncio +async def test_cli_session_token_keeps_minted_team_grants_when_team_lookup_is_404(monkeypatch): + monkeypatch.delenv("EXPERIMENTAL_UI_LOGIN", raising=False) + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-salt-cli-test") + + from fastapi import HTTPException, Request + from starlette.datastructures import URL + + from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken + + user_info = LiteLLM_UserTable( + user_id="internal-user-1", + user_role=LitellmUserRoles.INTERNAL_USER.value, + models=[], + ) + cli_token = ExperimentalUIJWTToken.get_cli_jwt_auth_token( + user_info, + team_id="team-abc", + team_alias="old-name", + team_models=["old-model"], + ) + attrs = _proxy_attrs_for_db_lookup() + import litellm.proxy.proxy_server as _proxy_server_mod + + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for key, value in attrs.items(): + setattr(_proxy_server_mod, key, value) + request = Request(scope={"type": "http"}) + request._url = URL(url="/v1/models") + with ( + patch( # test-quality-ok: direct collaborator patch exercises the real auth builder + "litellm.proxy.auth.user_api_key_auth.get_team_object", + new_callable=AsyncMock, + side_effect=HTTPException(status_code=404), + ), + patch( # test-quality-ok: direct collaborator patch exercises the real auth builder + "litellm.proxy.auth.user_api_key_auth.get_user_object", + new_callable=AsyncMock, + return_value=None, + ), + ): + result = await _user_api_key_auth_builder( + request=request, + api_key=f"Bearer {cli_token}", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={}, + ) + finally: + for key, value in originals.items(): + setattr(_proxy_server_mod, key, value) + + assert result.team_models == ["old-model"] + + @pytest.mark.asyncio async def test_cli_session_token_authenticates_when_jwt_auth_enabled_without_license(monkeypatch): """A lite login token is an encrypted (non-JWT) session blob. With From d1fdf64aa87e1148b65dcc672336abd7131ea713 Mon Sep 17 00:00:00 2001 From: jesus Date: Tue, 8 Sep 2026 23:24:42 +0000 Subject: [PATCH 2/3] style(auth): format session token grant refresh Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/user_api_key_auth.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 53b382a0180..b7b7c8d7f1c 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -2373,7 +2373,10 @@ async def _refresh_session_token_team_grants( except HTTPException: return valid_token return valid_token.model_copy( - update={"team_models": list(team_obj.models), "team_alias": team_obj.team_alias} # mutable-ok: auth model requires a fresh list + update={ # mutable-ok: model_copy requires a mutable update mapping + "team_models": list(team_obj.models), # mutable-ok: auth model requires a fresh list + "team_alias": team_obj.team_alias, + } ) From a443a68347695f51e7b0509683981bfa9286f321 Mon Sep 17 00:00:00 2001 From: jesus Date: Wed, 9 Sep 2026 00:04:31 +0000 Subject: [PATCH 3/3] refactor(auth): sync session token team grants from resolved team object Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/user_api_key_auth.py | 39 ++----------------------- 1 file changed, 3 insertions(+), 36 deletions(-) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index b7b7c8d7f1c..45a32379054 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1742,15 +1742,6 @@ async def _user_api_key_auth_builder( ): valid_token = ExperimentalUIJWTToken.get_key_object_from_ui_hash_key(api_key) - if valid_token is not None and valid_token.is_session_token: - valid_token = await _refresh_session_token_team_grants( - valid_token=valid_token, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - ) - if ( valid_token is not None and isinstance(valid_token, UserAPIKeyAuth) @@ -2214,6 +2205,9 @@ async def _user_api_key_auth_builder( # guardrails (or any other metadata) added after the key was cached # are picked up on subsequent requests without a cache eviction. valid_token.team_metadata = _team_obj.metadata + if valid_token.is_session_token: + valid_token.team_models = list(_team_obj.models) # mutable-ok: auth model requires a fresh list + valid_token.team_alias = _team_obj.team_alias else: valid_token.team_object_permission = None @@ -2353,33 +2347,6 @@ def _team_obj_from_token(valid_token: UserAPIKeyAuth) -> LiteLLM_TeamTableCached ) -async def _refresh_session_token_team_grants( - valid_token: UserAPIKeyAuth, - prisma_client: PrismaClient | None, - user_api_key_cache: UserApiKeyCache, - parent_otel_span: Span | None, - proxy_logging_obj: ProxyLogging, -) -> UserAPIKeyAuth: - if valid_token.team_id is None or valid_token.team_id == UI_TEAM_ID or prisma_client is None: - return valid_token - try: - team_obj: Final = await get_team_object( - team_id=valid_token.team_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 HTTPException: - return valid_token - return valid_token.model_copy( - update={ # mutable-ok: model_copy requires a mutable update mapping - "team_models": list(team_obj.models), # mutable-ok: auth model requires a fresh list - "team_alias": team_obj.team_alias, - } - ) - - 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.