diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 385f39e02a4..2293a35cce6 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2755,6 +2755,16 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob "user id." ), ) + via_ui_session_blob: bool = Field( + default=False, + exclude=True, + description=( + "Server-only marker set exclusively where an Admin UI session blob is decrypted with the " + "proxy's own ui_hash_key, via post-construction assignment. Stripped from validated input " + "for the same reason as via_virtual_key. The blob carries no key row, so via_virtual_key " + "cannot speak for it, and a proxy-admin session returns before that marker is ever set." + ), + ) budget_reservation: dict[str, Any] | None = Field(default=None, exclude=True) budget_throttle_pct: float | None = Field(default=None, exclude=True) user: Any | None = None # Expanded user object when expand=user is used @@ -2781,6 +2791,7 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob values.pop("mcp_admitted_user_subject", None) values.pop("mcp_source_team_rpm_limits", None) values.pop("via_virtual_key", None) + values.pop("via_ui_session_blob", None) if values.get("api_key") is not None: values.update({"token": cls._safe_hash_litellm_api_key(values.get("api_key"))}) if isinstance(values.get("api_key"), str): diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 39e1c14a6e6..0dee7e48f9d 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -26,6 +26,7 @@ from litellm.constants import ( GLOBAL_PROXY_SPEND_CACHE_KEY, LITELLM_PROXY_BUDGET_NAME, LITELLM_PROXY_MASTER_KEY_ALIAS, + UI_SESSION_TOKEN_TEAM_ID, ) from litellm.integrations.otel.model.config import is_otel_v2_enabled from litellm.integrations.otel.runtime import phase_span, seed_request_identity @@ -1597,6 +1598,11 @@ async def _user_api_key_auth_builder( and get_secret_bool("EXPERIMENTAL_UI_LOGIN") is not False ): valid_token = ExperimentalUIJWTToken.get_key_object_from_ui_hash_key(api_key) + if valid_token is not None: + # Decryption with ui_hash_key is itself the provenance proof, and it is the + # only one this token gets: a proxy-admin session returns below, before the + # virtual-key paths that would mark it. + valid_token.via_ui_session_blob = True if ( valid_token is not None @@ -2163,6 +2169,26 @@ def _team_obj_from_token(valid_token: UserAPIKeyAuth) -> LiteLLM_TeamTableCached ) +def _is_ui_session_token(valid_token: UserAPIKeyAuth) -> bool: + """An Admin UI session key, whose reserved team is absent by design rather than deleted. + + The id alone cannot earn the exemption. The synthesized team's empty ``models`` reads as + every model, so an identity that merely names the reserved team is widened, not just waved + through, and several auth paths take ``team_id`` straight from data the proxy did not mint: + JWT claims, an OAuth2 introspection response, a custom auth handler's return value. + + So ask where the credential came from rather than trying to name every producer it did not + come from. Both markers are set only by post-construction assignment at a boundary the proxy + itself validated, and both are stripped from validated input, so no claim, header or handler + return can carry one in. A database-minted session key is marked as a virtual key. The + ``EXPERIMENTAL_UI_LOGIN`` blob has no key row to be marked by, and a proxy-admin one returns + before the virtual-key paths, so it is marked where it is decrypted instead. + """ + return valid_token.team_id == UI_SESSION_TOKEN_TEAM_ID and ( + valid_token.via_virtual_key or valid_token.via_ui_session_blob + ) + + 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. @@ -2272,7 +2298,9 @@ async def _run_centralized_common_checks( ) fetch_coros: Final = [] - if user_api_key_auth_obj.team_id is not None: + if _is_ui_session_token(user_api_key_auth_obj): + fetch_coros.append(_safe_fetch("team", _team_from_token(user_api_key_auth_obj))) + elif user_api_key_auth_obj.team_id is not None: fetch_coros.append( _safe_fetch( "team", @@ -2489,6 +2517,12 @@ async def _run_centralized_common_checks( ) +async def _team_from_token(valid_token: UserAPIKeyAuth) -> LiteLLM_TeamTableCachedObj: + """Coroutine wrapper over ``_team_obj_from_token`` for the gather below, used + where the team is known not to be resolvable from the database.""" + return _team_obj_from_token(valid_token) + + async def _noop_none() -> None: """Sentinel coroutine for asyncio.gather when a fetch is unnecessary (e.g. token has no team_id). Keeps the result tuple positional.""" diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 7e190e8b19d..60425a8220a 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -841,6 +841,21 @@ def _enforce_upperbound_key_params( ) +def reject_reserved_ui_session_team_id(team_id: str | None) -> None: + """``/team/new`` already refuses this id, so the row it names can never exist and a key + carrying it resolves its team from the token instead of the database. That exemption + belongs to Admin UI sessions alone: an ordinary key carrying the id inherits it, and with + it the skip of its owner's user-level model check. The UI login flow mints its session key + through ``generate_key_helper_fn`` directly, so reserving the id here does not reach it. + """ + if team_id == UI_SESSION_TOKEN_TEAM_ID: + raise HTTPException( + status_code=400, + detail=f"team_id '{UI_SESSION_TOKEN_TEAM_ID}' is reserved for LiteLLM UI dashboard " + "sessions and cannot be assigned to a key. Please use a different team id.", + ) + + async def _common_key_generation_helper( data: GenerateKeyRequest, user_api_key_dict: UserAPIKeyAuth, @@ -854,6 +869,8 @@ async def _common_key_generation_helper( prisma_client, ) + reject_reserved_ui_session_team_id(data.team_id) + common_key_access_checks( user_api_key_dict=user_api_key_dict, data=data, @@ -2770,6 +2787,8 @@ async def update_key_fn( ) try: + reject_reserved_ui_session_team_id(data.team_id) + # Validate budget values are not negative and are finite numbers if data.max_budget is not None and (not math.isfinite(data.max_budget) or data.max_budget < 0): raise HTTPException( 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 eea556b9a0d..3c6a26b7b05 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 @@ -4472,6 +4472,218 @@ async def test_centralized_common_checks_absent_team_refused_despite_db_unavaila setattr(_proxy_server_mod, k, v) +@pytest.mark.parametrize("marker", ["via_virtual_key", "via_ui_session_blob"]) +@pytest.mark.asyncio +async def test_centralized_common_checks_ui_session_team_is_not_treated_as_deleted(marker): + """Every Admin UI session token is stamped with the reserved + ``litellm-dashboard`` team id, which never has a row, so the absent-team + refusal read the ordinary UI case as a deleted team and hard 404'd every + dashboard request. The sentinel resolves to the token-derived team without a + lookup that can only fail. + + One row per way the proxy mints a UI session: a database-backed key row, and the + ``EXPERIMENTAL_UI_LOGIN`` blob, which has no row and so cannot be a virtual key.""" + import litellm.proxy.proxy_server as _proxy_server_mod + from fastapi import Request + from starlette.datastructures import URL + + from litellm.constants import UI_SESSION_TOKEN_TEAM_ID + from litellm.proxy.auth.user_api_key_auth import TeamNotFoundError + + token = UserAPIKeyAuth( + api_key="sk-test", + token="hashed-sk-test", + team_id=UI_SESSION_TOKEN_TEAM_ID, + models=[], + team_models=[], + ) + setattr(token, marker, True) + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + request._body = json.dumps({"model": "gpt-4.1"}).encode() + + 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( + "litellm.proxy.auth.user_api_key_auth.get_team_object", + new_callable=AsyncMock, + side_effect=TeamNotFoundError(team_id=UI_SESSION_TOKEN_TEAM_ID), + ) as mock_get_team, + patch( + "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-4.1"}, + route="/chat/completions", + ) + mock_get_team.assert_not_awaited() + mock_checks.assert_awaited_once() + assert mock_checks.call_args.kwargs["team_object"].team_id == UI_SESSION_TOKEN_TEAM_ID + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + +@pytest.mark.parametrize( + "identity", + [ + pytest.param( + {"token": "hashed-sk-test", "jwt_claims": {"team_id": "litellm-dashboard"}}, + id="jwt_mapped_to_a_virtual_key", + ), + pytest.param( + {"token": None, "jwt_claims": {"team_id": "litellm-dashboard"}}, + id="standard_jwt_identity", + ), + # The shape Oauth2Handler.check_oauth2_token returns: team_id read straight out of + # the introspection response, api_key set so a token hash exists, and no jwt_claims. + pytest.param({"api_key": "oauth2-access-token"}, id="oauth2_introspection_identity"), + pytest.param({"token": "handler-supplied"}, id="custom_auth_handler_return"), + pytest.param({"token": None, "jwt_claims": None}, id="tokenless_identity"), + ], +) +@pytest.mark.asyncio +async def test_centralized_common_checks_non_ui_identity_claiming_ui_team_still_refused(identity): + """Several auth paths take ``team_id`` from data the proxy did not mint, so any of them + could name the reserved UI team and inherit the exemption if the id alone earned it. The + synthesized team's empty ``models`` reads as every model, so that is a widening, not just + a bypass. + + None of these carry a provenance marker, and none can: both markers are stripped from + validated input, which the marker-forging test below pins separately.""" + import litellm.proxy.proxy_server as _proxy_server_mod + from fastapi import Request + from starlette.datastructures import URL + + from litellm.constants import UI_SESSION_TOKEN_TEAM_ID + from litellm.proxy.auth.user_api_key_auth import TeamNotFoundError + + token = UserAPIKeyAuth( + **{ + "api_key": None, + "team_id": UI_SESSION_TOKEN_TEAM_ID, + "models": [], + "team_models": [], + **identity, + } + ) + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + request._body = json.dumps({"model": "gpt-4.1"}).encode() + + 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( + "litellm.proxy.auth.user_api_key_auth.get_team_object", + new_callable=AsyncMock, + side_effect=TeamNotFoundError(team_id=UI_SESSION_TOKEN_TEAM_ID), + ) as mock_get_team, + patch( + "litellm.proxy.auth.user_api_key_auth.common_checks", + new_callable=AsyncMock, + ) as mock_checks, + pytest.raises(TeamNotFoundError), + ): + await _run_centralized_common_checks( + user_api_key_auth_obj=token, + request=request, + request_data={"model": "gpt-4.1"}, + route="/chat/completions", + ) + mock_get_team.assert_awaited_once() + mock_checks.assert_not_awaited() + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + +@pytest.mark.asyncio +async def test_ui_session_blob_is_marked_where_it_is_decrypted(monkeypatch): + """The blob carries no key row, so nothing downstream can mark it, and a proxy-admin + session returns from the builder before the virtual-key paths run. Losing the marker + at the decrypt boundary therefore costs the exemption outright and puts every Admin UI + request under ``EXPERIMENTAL_UI_LOGIN`` back on the 404 this PR exists to fix.""" + import litellm.proxy.proxy_server as _proxy_server_mod + from fastapi import Request + from starlette.datastructures import URL + + from litellm.constants import UI_SESSION_TOKEN_TEAM_ID + from litellm.proxy._types import LiteLLM_UserTable, LitellmUserRoles + from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken + from litellm.proxy.auth.user_api_key_auth import ( + _is_ui_session_token, + _user_api_key_auth_builder, + ) + + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-test-salt-ui-session-blob") + + admin = LiteLLM_UserTable( + user_id="ui-admin", + user_role=LitellmUserRoles.PROXY_ADMIN.value, + models=[], + ) + blob = ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token(admin) + assert not blob.startswith("sk-"), "a sk- credential would take the key-row path instead" + + attrs = _proxy_server_attrs_for_custom_auth(user_custom_auth=None) + originals = {attr: getattr(_proxy_server_mod, attr, None) for attr in attrs} + try: + for attr, val in attrs.items(): + setattr(_proxy_server_mod, attr, val) + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + + result = await _user_api_key_auth_builder( + request=request, + api_key=f"Bearer {blob}", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={}, + ) + + assert result.team_id == UI_SESSION_TOKEN_TEAM_ID + assert result.via_virtual_key is False, "the blob has no key row to be marked by" + assert result.via_ui_session_blob is True + assert _is_ui_session_token(result) is True + finally: + for attr, val in originals.items(): + setattr(_proxy_server_mod, attr, val) + + +@pytest.mark.parametrize("marker", ["via_virtual_key", "via_ui_session_blob"]) +def test_ui_session_provenance_markers_cannot_be_forged_from_validated_input(marker): + """The exemption rests entirely on these two markers, so the whole guard collapses if a + JWT claim splat, a custom auth handler's return value, or key metadata can set one. They + are server-only: settable by assignment, stripped from anything validated.""" + from litellm.constants import UI_SESSION_TOKEN_TEAM_ID + from litellm.proxy.auth.user_api_key_auth import _is_ui_session_token + + forged = UserAPIKeyAuth.model_validate( + {"team_id": UI_SESSION_TOKEN_TEAM_ID, "api_key": "sk-forged", marker: True} + ) + assert getattr(forged, marker) is False + assert _is_ui_session_token(forged) is False + + minted = UserAPIKeyAuth(team_id=UI_SESSION_TOKEN_TEAM_ID, api_key="sk-minted") + assert _is_ui_session_token(minted) is False + setattr(minted, marker, True) + assert _is_ui_session_token(minted) is True + + @pytest.mark.asyncio async def test_centralized_common_checks_unreadable_team_keeps_db_unavailable_optout(): """The counterpart: an unreadable team leaves the grant unknown rather than diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index bdf09a95e4b..cfacced95b9 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -16444,3 +16444,91 @@ async def test_regenerate_key_repoints_live_membership_not_the_key_row_it_read( assert await _authorized_models_for_key( access_groups, new_token_hash, ["ag-revoked-since", "ag-attached-since"] ) == ["attached-model"] + + +@pytest.mark.parametrize("team_id", [None, "a-real-team"]) +def test_reject_reserved_ui_session_team_id_allows_ordinary_team_ids(team_id): + from litellm.proxy.management_endpoints.key_management_endpoints import ( + reject_reserved_ui_session_team_id, + ) + + reject_reserved_ui_session_team_id(team_id) + + +def test_reject_reserved_ui_session_team_id_refuses_the_sentinel(): + """A key carrying the reserved id resolves its team from the token rather than the + database, which is the exemption Admin UI sessions need and nothing else may borrow: + it also skips the owner's user-level model check. ``/team/new`` already reserves this + id, so reserving it for keys closes the same door on the other side.""" + from fastapi import HTTPException + + from litellm.constants import UI_SESSION_TOKEN_TEAM_ID + from litellm.proxy.management_endpoints.key_management_endpoints import ( + reject_reserved_ui_session_team_id, + ) + + with pytest.raises(HTTPException) as exc: + reject_reserved_ui_session_team_id(UI_SESSION_TOKEN_TEAM_ID) + assert exc.value.status_code == 400 + assert "reserved" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_key_generation_refuses_the_reserved_ui_team_even_for_a_proxy_admin(): + """Only a proxy admin can reach this at all, since key_generation_check refuses a + non-admin naming a team with no row. Admin authority is not the question: the id is + reserved for sessions the UI mints, and the UI mints those through + generate_key_helper_fn directly, so it never passes through here.""" + from fastapi import HTTPException + + from litellm.constants import UI_SESSION_TOKEN_TEAM_ID + from litellm.proxy._types import GenerateKeyRequest, LitellmUserRoles + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _common_key_generation_helper, + ) + + admin = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + user_id="admin-user", + api_key="sk-admin", + ) + + with pytest.raises(HTTPException) as exc: + await _common_key_generation_helper( + data=GenerateKeyRequest(team_id=UI_SESSION_TOKEN_TEAM_ID), + user_api_key_dict=admin, + litellm_changed_by=None, + team_table=None, + ) + assert exc.value.status_code == 400 + assert UI_SESSION_TOKEN_TEAM_ID in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_key_update_refuses_moving_a_key_onto_the_reserved_ui_team(): + """Reserving the id only at creation would leave the same exemption one /key/update + away, so the update path carries the guard too. It runs before the key is even read, + so no database access is needed to refuse it.""" + from fastapi import Request + + from litellm.constants import UI_SESSION_TOKEN_TEAM_ID + from litellm.proxy._types import LitellmUserRoles, ProxyException, UpdateKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import update_key_fn + + admin = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + user_id="admin-user", + api_key="sk-admin", + ) + request = Request(scope={"type": "http", "headers": [], "method": "POST"}) + + # the endpoint's own handler re-raises an HTTPException as a ProxyException + with pytest.raises(ProxyException) as exc: + await update_key_fn( + request=request, + data=UpdateKeyRequest(key="sk-existing", team_id=UI_SESSION_TOKEN_TEAM_ID), + user_api_key_dict=admin, + litellm_changed_by=None, + ) + assert exc.value.code == "400" + assert UI_SESSION_TOKEN_TEAM_ID in str(exc.value.message)