From 23df4f3b84c341ae39f2bc1ec442d93c133ab526 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Wed, 28 Oct 2026 22:12:15 -0700
Subject: [PATCH] fix(proxy): refuse a stale gateway team pick and require an
 explicit choice

---
 .../anthropic_endpoints/gateway_endpoints.py  | 17 +++---
 .../html_forms/native_client_consent.py       |  7 ++-
 litellm/proxy/management_endpoints/ui_sso.py  |  7 ++-
 .../test_gateway_endpoints.py                 | 56 +++++++++++++------
 .../proxy/management_endpoints/test_ui_sso.py | 10 ++--
 5 files changed, 64 insertions(+), 33 deletions(-)

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'<input type="hidden" name="team_id" value="{escape(team_id)}">'
             f"<p>Requests are attributed to team <strong>{escape(team_label)}</strong>.</p>"
         )
+    placeholder: Final = '<option value="" disabled selected>Choose a team</option>' if require_choice else ""
     options: Final = "".join(
         f'<option value="{escape(team_id)}">{escape(team_label)}</option>' for team_id, team_label in teams
     )
+    required: Final = " required" if require_choice else ""
     return (
-        f'<label for="team_id">Attribute requests to team</label><select id="team_id" name="team_id">{options}</select>'
+        '<label for="team_id">Attribute requests to team</label>'
+        f'<select id="team_id" name="team_id"{required}>{placeholder}{options}</select>'
     )
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(
           <p>{instructions}</p>
           <form method="post" action="{escaped_verify_url}">
             <input type="hidden" name="browser_complete_token" value="{escaped_browser_complete_token}" />
-            {team_select_field(team_choices)}
+            {team_select_field(team_choices, require_choice=True)}
             <label for="user_code">Verification code</label>
             <input id="user_code" name="user_code" autocomplete="one-time-code"{user_code_value_attr} required autofocus />
             <button type="submit">Continue</button>
@@ -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 <eu>"), ("team-claude", None))
+                for team_id, alias in (("team-ops", "Ops <eu>"), ("team-claude", None), ("team-ops-us", "Ops <eu>"))
             ]
         )
         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 (
-            '<select id="team_id" name="team_id">'
-            '<option value="team-ops">Ops &lt;eu&gt;</option>'
+            '<select id="team_id" name="team_id" required>'
+            '<option value="" disabled selected>Choose a team</option>'
+            '<option value="team-ops">Ops &lt;eu&gt; (team-ops)</option>'
             '<option value="team-claude">team-claude</option>'
+            '<option value="team-ops-us">Ops &lt;eu&gt; (team-ops-us)</option>'
             "</select>"
         ) in page