diff --git a/litellm/proxy/anthropic_endpoints/gateway_endpoints.py b/litellm/proxy/anthropic_endpoints/gateway_endpoints.py index 8c937e4cbd4..6c58f2edc02 100644 --- a/litellm/proxy/anthropic_endpoints/gateway_endpoints.py +++ b/litellm/proxy/anthropic_endpoints/gateway_endpoints.py @@ -261,12 +261,6 @@ async def device_authorization(request: Request) -> JSONResponse: return JSONResponse(content=body.model_dump(exclude_none=True)) -def _session_team_id(chosen_team_id: str | None, teams: tuple[str, ...]) -> str | None: - if chosen_team_id in teams: - return chosen_team_id - return teams[0] if teams else None - - def _validate_login(flow: Mapping[str, object]) -> _GatewayLogin | _OAuthError: from litellm.proxy.management_endpoints.ui_sso import selected_cli_sso_team_detail @@ -278,8 +272,15 @@ def _validate_login(flow: Mapping[str, object]) -> _GatewayLogin | _OAuthError: status_code=400, error="invalid_grant", description="The login session is malformed; sign in again" ) - chosen_team_id: Final = flow.get(CLI_SSO_CHOSEN_TEAM_ID) - team_id: Final = _session_team_id(chosen_team_id if isinstance(chosen_team_id, str) else None, session_data.teams) + chosen: Final = flow.get(CLI_SSO_CHOSEN_TEAM_ID) + chosen_team_id: Final = chosen if isinstance(chosen, str) else None + if chosen_team_id is not None and chosen_team_id not in session_data.teams: + return _OAuthError( + status_code=400, + error="invalid_grant", + description="The team picked at sign-in is no longer one of your teams; sign in again", + ) + team_id: Final = chosen_team_id or next(iter(session_data.teams), None) selected_team: Final = selected_cli_sso_team_detail(team_details=session_data.team_details, team_id=team_id) if selected_team is None: return _OAuthError( diff --git a/litellm/proxy/common_utils/html_forms/native_client_consent.py b/litellm/proxy/common_utils/html_forms/native_client_consent.py index 4c7e7a5fedd..aeadc93221c 100644 --- a/litellm/proxy/common_utils/html_forms/native_client_consent.py +++ b/litellm/proxy/common_utils/html_forms/native_client_consent.py @@ -74,7 +74,7 @@ button {{ flex: 1; padding: 10px; border-radius: 6px; font-size: 15px; cursor: p """ -def team_select_field(teams: Sequence[tuple[str, str]]) -> str: +def team_select_field(teams: Sequence[tuple[str, str]], *, require_choice: bool = False) -> str: if not teams: return "" if len(teams) == 1: @@ -83,9 +83,12 @@ def team_select_field(teams: Sequence[tuple[str, str]]) -> str: f'' f"

Requests are attributed to team {escape(team_label)}.

" ) + placeholder: Final = '' if require_choice else "" options: Final = "".join( f'' for team_id, team_label in teams ) + required: Final = " required" if require_choice else "" return ( - f'' + '' + f'' ) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index f347a742ae3..790556e5b7b 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -16,6 +16,7 @@ import json import os import re import secrets +from collections import Counter from collections.abc import Awaitable, Callable, Mapping, Sequence from copy import deepcopy from html import escape @@ -748,7 +749,7 @@ def _render_cli_sso_verification_page(

{instructions}

- {team_select_field(team_choices)} + {team_select_field(team_choices, require_choice=True)} @@ -2456,9 +2457,11 @@ def _browser_team_choices( ) -> tuple[tuple[str, str], ...]: if not flow.browser_team_choice: return () - return tuple( + labeled: Final = tuple( (detail.team_id, detail.team_alias or detail.team_id) for detail in team_details if detail.team_id is not None ) + label_counts: Final = Counter(label for _, label in labeled) + return tuple((team_id, label if label_counts[label] == 1 else f"{label} ({team_id})") for team_id, label in labeled) def selected_cli_sso_team_detail(team_details: object, team_id: str | None) -> CliSsoTeamDetail | None: diff --git a/tests/unit/proxy/anthropic_endpoints/test_gateway_endpoints.py b/tests/unit/proxy/anthropic_endpoints/test_gateway_endpoints.py index f0893a01f6e..3eab7a4d7e5 100644 --- a/tests/unit/proxy/anthropic_endpoints/test_gateway_endpoints.py +++ b/tests/unit/proxy/anthropic_endpoints/test_gateway_endpoints.py @@ -22,6 +22,7 @@ from fastapi.testclient import TestClient from litellm.caching.dual_cache import DualCache from litellm.proxy._types import ProxyException from litellm.proxy.anthropic_endpoints import gateway_endpoints +from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken from litellm.proxy.management_endpoints.ui_sso import ( _get_cli_sso_flow_cache_key, _hash_cli_sso_secret, @@ -272,6 +273,33 @@ def _login_id(device_code: str) -> str: return device_code.partition(".")[0] +def _request_token_after_browser_choice( + client: TestClient, cache: DualCache, chosen_team_id: str | None +) -> httpx.Response: + device_code: Final = _start_device_flow(client) + _complete_flow( + cache, + device_code, + session_data={ + **_COMPLETED_SESSION, + "teams": ["team-a", "team-b"], + "team_details": [ + *_COMPLETED_SESSION["team_details"], + { + "team_id": "team-b", + "team_alias": "Team B", + "team_models": ["claude-opus-4-5"], + "team_model_aliases": None, + }, + ], + }, + ) + flow_key: Final = _get_cli_sso_flow_cache_key(_login_id(device_code)) + flow: Final = cache.get_cache(key=flow_key) + cache.set_cache(key=flow_key, value={**flow, "chosen_team_id": chosen_team_id}, ttl=600) + return _request_token(client, device_code) + + def _complete_flow(cache: DualCache, device_code: str, session_data: Mapping[str, object] = _COMPLETED_SESSION) -> None: key: Final = _get_cli_sso_flow_cache_key(_login_id(device_code)) flow: Final = cache.get_cache(key=key) @@ -405,28 +433,22 @@ def test_token_success_mints_bearer_and_is_single_use(): def test_token_mints_for_the_team_chosen_in_the_browser( chosen_team_id: str | None, minted_team_id: str, minted_alias: str ): - two_team_session: Final = { - **_COMPLETED_SESSION, - "teams": ["team-a", "team-b"], - "team_details": [ - *_COMPLETED_SESSION["team_details"], - {"team_id": "team-b", "team_alias": "Team B", "team_models": ["claude-opus-4-5"], "team_model_aliases": None}, - ], - } with _gateway_env() as (client, cache): - device_code: Final = _start_device_flow(client) - _complete_flow(cache, device_code, session_data=two_team_session) - flow_key: Final = _get_cli_sso_flow_cache_key(_login_id(device_code)) - flow: Final = cache.get_cache(key=flow_key) - cache.set_cache(key=flow_key, value={**flow, "chosen_team_id": chosen_team_id}, ttl=600) - with patch(_MINT, return_value="sk-litellm-session-token") as mint: - resp: Final = _request_token(client, device_code) + resp: Final = _request_token_after_browser_choice(client, cache, chosen_team_id) + session: Final = ExperimentalUIJWTToken.get_key_object_from_ui_hash_key(resp.json()["access_token"]) assert resp.status_code == 200 - assert mint.call_args.kwargs["team_id"] == minted_team_id - assert mint.call_args.kwargs["team_alias"] == minted_alias + assert session is not None + assert (session.team_id, session.team_alias) == (minted_team_id, minted_alias) assert _opened_refresh(resp.json()["refresh_token"]).team_id == minted_team_id +def test_token_refuses_a_browser_choice_that_is_no_longer_one_of_the_session_teams(): + with _gateway_env() as (client, cache): + resp: Final = _request_token_after_browser_choice(client, cache, "team-gone") + assert resp.status_code == 400 + assert resp.json()["error"] == "invalid_grant" + + def test_token_teamless_user_mints_without_a_team(): with _gateway_env() as (client, cache): device_code = _start_device_flow(client) diff --git a/tests/unit/proxy/management_endpoints/test_ui_sso.py b/tests/unit/proxy/management_endpoints/test_ui_sso.py index 406b6b368a2..a87064961fe 100644 --- a/tests/unit/proxy/management_endpoints/test_ui_sso.py +++ b/tests/unit/proxy/management_endpoints/test_ui_sso.py @@ -4022,14 +4022,14 @@ class TestCLIKeyRegenerationFlow: mock_prisma.db.litellm_teamtable.find_many = AsyncMock( return_value=[ MagicMock(model_dump=lambda team_id=team_id, alias=alias: {"team_id": team_id, "team_alias": alias}) - for team_id, alias in (("team-ops", "Ops "), ("team-claude", None)) + for team_id, alias in (("team-ops", "Ops "), ("team-claude", None), ("team-ops-us", "Ops ")) ] ) with ( patch( "litellm.proxy.management_endpoints.ui_sso.get_user_info_from_db", return_value=LiteLLM_UserTable( - user_id="dev-1", user_role="internal_user", teams=["team-claude", "team-ops"] + user_id="dev-1", user_role="internal_user", teams=["team-claude", "team-ops", "team-ops-us"] ), ), patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), @@ -4047,9 +4047,11 @@ class TestCLIKeyRegenerationFlow: assert 'name="team_id"' not in page return assert ( - '' + '' + '' '' + '' "" ) in page