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

@@ -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.

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