From e8f0735d8a935cf4bf6ec97c5ef6b8d846667b95 Mon Sep 17 00:00:00 2001
From: simonp9 <349053457+simonp9@users.noreply.github.com>
Date: Thu, 29 Oct 2026 03:47:27 -0700
Subject: [PATCH] feat(proxy): let a Claude Code gateway user choose the
 session team (#49001)

* feat(proxy): let a Claude Code gateway user choose the session team

The gateway's browser verification page now lists the developer's teams, and the device code grant mints the session for the picked team. A team the developer is not on is refused with a 400, and a form without a team keeps the first-team default.

* fix(proxy): refuse a stale gateway team pick and require an explicit choice

* fix(gateway): list the browser-picked team first so an older replica mid-deploy issues it

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
---
 .../anthropic_endpoints/gateway_endpoints.py  |  17 ++-
 .../html_forms/native_client_consent.py       |   9 +-
 litellm/proxy/management_endpoints/ui_sso.py  |  68 ++++++++-
 .../test_claude_code_gateway_oauth.py         |  91 ++++++++++++
 .../test_gateway_endpoints.py                 |  55 ++++++++
 .../proxy/management_endpoints/test_ui_sso.py | 132 ++++++++++++++++++
 6 files changed, 364 insertions(+), 8 deletions(-)

diff --git a/litellm/proxy/anthropic_endpoints/gateway_endpoints.py b/litellm/proxy/anthropic_endpoints/gateway_endpoints.py
index 1911cb574ef..6c58f2edc02 100644
--- a/litellm/proxy/anthropic_endpoints/gateway_endpoints.py
+++ b/litellm/proxy/anthropic_endpoints/gateway_endpoints.py
@@ -45,7 +45,11 @@ from litellm.proxy.common_utils.http_parsing_utils import (  # noqa: F401  # leg
     safe_set_request_parsed_body,
 )
 from litellm.proxy.management_endpoints.sso_helper_utils import CLI_SSO_SESSIONS_TARGET
-from litellm.proxy.management_endpoints.ui_sso import CliSsoTeamDetail
+from litellm.proxy.management_endpoints.ui_sso import (
+    CLI_SSO_BROWSER_TEAM_CHOICE,
+    CLI_SSO_CHOSEN_TEAM_ID,
+    CliSsoTeamDetail,
+)
 from litellm.proxy.mcp_server.gateway_dcr_flow import (
     PROXY_API_AUDIENCE,
     MintedProxyCredential,
@@ -235,6 +239,7 @@ async def device_authorization(request: Request) -> JSONResponse:
         "sso_complete": False,
         "user_code_verified": False,
         "session_data": None,
+        CLI_SSO_BROWSER_TEAM_CHOICE: True,
     }
     _set_cli_sso_flow(login_id=login_id, cache=cli_sso_session_cache, flow=flow)
 
@@ -267,7 +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"
         )
 
-    team_id: Final = session_data.teams[0] if session_data.teams else None
+    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 dac92c4e787..aeadc93221c 100644
--- a/litellm/proxy/common_utils/html_forms/native_client_consent.py
+++ b/litellm/proxy/common_utils/html_forms/native_client_consent.py
@@ -62,7 +62,7 @@ button {{ flex: 1; padding: 10px; border-radius: 6px; font-size: 15px; cursor: p
 <p>Approving issues it a personal credential that expires within {CLI_JWT_EXPIRATION_HOURS} hours. <code>lite logout</code> stops it from being renewed. Only approve if you started this sign-in yourself.</p>
 <form method="post" action="{escape(complete_url)}">
 <input type="hidden" name="flow" value="{escape(flow_handle)}">
-{_team_field(teams)}
+{team_select_field(teams)}
 <div class="actions">
 <button type="submit" name="decision" value="deny" class="deny">Deny</button>
 <button type="submit" name="decision" value="approve" class="approve">Approve</button>
@@ -74,7 +74,7 @@ button {{ flex: 1; padding: 10px; border-radius: 6px; font-size: 15px; cursor: p
 """
 
 
-def _team_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_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 22c4f27d94a..a980b9c33ff 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
@@ -42,7 +43,7 @@ if TYPE_CHECKING:
 import jwt
 from fastapi import APIRouter, Depends, Header, HTTPException, Request, Response, status
 from fastapi.responses import RedirectResponse
-from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError, field_validator
+from pydantic import AliasPath, BaseModel, ConfigDict, Field, TypeAdapter, ValidationError, field_validator
 
 import litellm
 from litellm._internal_context import with_service_target
@@ -106,6 +107,7 @@ from litellm.proxy.common_utils.html_forms.default_credentials_hint import shoul
 from litellm.proxy.common_utils.html_forms.jwt_display_template import (
     jwt_display_template,
 )
+from litellm.proxy.common_utils.html_forms.native_client_consent import team_select_field
 from litellm.proxy.common_utils.html_forms.ui_login import build_ui_login_form
 from litellm.proxy.common_utils.user_api_key_cache import AUTH_OBJECTS_TARGET, UserApiKeyCache
 from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
@@ -184,6 +186,8 @@ _CLI_SSO_START_RATE_LIMIT_WINDOW_SECONDS: Final = 60
 _CLI_SSO_START_RATE_LIMIT_MAX_ATTEMPTS: Final = 30
 _CLI_SSO_USER_CODE_ALPHABET: Final = "ABCDEFGHJKLMNPQRSTUVWXYZ23456789"
 _CLI_SSO_LOGIN_ID_RE: Final = re.compile(r"^cli-[A-Za-z0-9_-]{12,124}$")
+CLI_SSO_BROWSER_TEAM_CHOICE: Final = "browser_team_choice"
+CLI_SSO_CHOSEN_TEAM_ID: Final = "chosen_team_id"
 _CLI_SSO_USER_CODE_RE = re.compile(rf"^[{_CLI_SSO_USER_CODE_ALPHABET}]{{4}}-[{_CLI_SSO_USER_CODE_ALPHABET}]{{4}}$")
 _CLI_SSO_SCALAR_TYPES: Final = (str, int, float, bool)
 _CLI_SSO_DEST_KEY_RE: Final = re.compile(r"^[A-Za-z0-9_.-]+$")
@@ -670,6 +674,7 @@ def _render_cli_sso_verification_page(
     verify_url: str,
     browser_complete_token: str,
     prefill_user_code: str | None = None,
+    team_choices: Sequence[tuple[str, str]] = (),
 ) -> str:
     escaped_verify_url: Final = escape(verify_url, quote=True)
     escaped_browser_complete_token: Final = escape(browser_complete_token, quote=True)
@@ -716,6 +721,15 @@ def _render_cli_sso_verification_page(
             letter-spacing: 0.08em;
             text-transform: uppercase;
           }}
+          select {{
+            box-sizing: border-box;
+            width: 100%;
+            padding: 12px;
+            margin-bottom: 18px;
+            border: 1px solid #cbd5e1;
+            border-radius: 6px;
+            font-size: 16px;
+          }}
           button {{
             margin-top: 16px;
             width: 100%;
@@ -735,6 +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, 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>
@@ -823,13 +838,46 @@ async def cli_sso_complete(request: Request, login_id: str):
     ):
         raise HTTPException(status_code=400, detail="Invalid verification code")
 
-    flow["user_code_verified"] = True
-    _set_cli_sso_flow(login_id=login_id, cache=cli_sso_session_cache, flow=flow)
+    team_choice_flow: Final = _CliSsoTeamChoiceFlow.model_validate(flow)
+    chosen_team_id: Final = _browser_chosen_team_id(team_choice_flow, form_values)
+    if chosen_team_id is not None and chosen_team_id not in team_choice_flow.teams:
+        raise HTTPException(status_code=400, detail="Choose one of the teams you belong to")
+
+    _set_cli_sso_flow(
+        login_id=login_id,
+        cache=cli_sso_session_cache,
+        flow={
+            **flow,
+            "user_code_verified": True,
+            CLI_SSO_CHOSEN_TEAM_ID: chosen_team_id,
+            "session_data": (
+                flow["session_data"]
+                if chosen_team_id is None
+                else {
+                    **flow["session_data"],
+                    "teams": [chosen_team_id, *(team for team in team_choice_flow.teams if team != chosen_team_id)],
+                }
+            ),
+        },
+    )
 
     html_content: Final = render_cli_sso_success_page()
     return HTMLResponse(content=html_content, status_code=200)
 
 
+class _CliSsoTeamChoiceFlow(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    browser_team_choice: bool = Field(default=False, validation_alias=CLI_SSO_BROWSER_TEAM_CHOICE)
+    teams: tuple[str, ...] = Field(default=(), validation_alias=AliasPath("session_data", "teams"))
+
+
+def _browser_chosen_team_id(flow: _CliSsoTeamChoiceFlow, form_values: Mapping[str, Sequence[str]]) -> str | None:
+    if not flow.browser_team_choice:
+        return None
+    return (form_values.get("team_id") or [""])[0] or None
+
+
 def normalize_email(email: str | None) -> str | None:
     """
     Normalize email address to lowercase for consistent storage and comparison.
@@ -2416,6 +2464,18 @@ def _cli_sso_session_teams(team_details: Sequence[CliSsoTeamDetail]) -> list[str
     return [detail.team_id for detail in team_details if detail.team_id is not None]
 
 
+def _browser_team_choices(
+    flow: _CliSsoTeamChoiceFlow, team_details: Sequence[CliSsoTeamDetail]
+) -> tuple[tuple[str, str], ...]:
+    if not flow.browser_team_choice:
+        return ()
+    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:
     """``None`` means the team's grants are unknown. An empty grant is a real value meaning unrestricted,
     so an unknown one must not be minted as empty."""
@@ -2484,6 +2544,7 @@ async def _complete_cli_sso_callback_session(
         request.scope.get("litellm_microsoft_interactive_subject"), resolved_user_id, prisma_client
     )
     resolved_teams: Final = _cli_sso_session_teams(team_details)
+    team_choices: Final = _browser_team_choices(_CliSsoTeamChoiceFlow.model_validate(flow), team_details)
     attribution_metadata: Final = build_cli_sso_attribution_metadata(result=result)
     if attribution_metadata:
         await _persist_cli_sso_user_metadata(
@@ -2521,6 +2582,7 @@ async def _complete_cli_sso_callback_session(
             verify_url=verify_url,
             browser_complete_token=browser_complete_token,
             prefill_user_code=prefill_user_code,
+            team_choices=team_choices,
         ),
         status_code=200,
     )
diff --git a/tests/integration/authorization/test_claude_code_gateway_oauth.py b/tests/integration/authorization/test_claude_code_gateway_oauth.py
index a80eeacad46..4dedbdcd8e4 100644
--- a/tests/integration/authorization/test_claude_code_gateway_oauth.py
+++ b/tests/integration/authorization/test_claude_code_gateway_oauth.py
@@ -545,3 +545,94 @@ def test_device_authorization_polls_complete_sso_and_mints_a_model_restricted_be
             assert set(token_form) == {"code", "code_verifier", "grant_type", "redirect_uri"}, token_form
             assert userinfo_request.method == "GET" and urlsplit(userinfo_request.target).path == "/userinfo"
             assert user_id == _USER_ID
+
+
+_TEAM_OPTION: Final = re.compile(r'<option value="([^"]+)">')
+
+
+def _device_code_grant(gateway: Gateway, device_code: str) -> httpx.Response:
+    return gateway.client.post(
+        "/claude_code_gateway/oauth/token",
+        data={
+            "grant_type": "urn:ietf:params:oauth:grant-type:device_code",
+            "device_code": device_code,
+            "client_id": _CLIENT_ID,
+        },
+    )
+
+
+@pytest.mark.timeout(300)
+def test_device_flow_mints_the_session_for_the_team_the_developer_picks_in_the_browser(
+    enabled_gateway: Gateway, identity_provider: Wire
+) -> None:
+    with enabled_gateway.scenario() as scenario:
+        user_id: Final = scenario.user(user_id=_USER_ID, user_email=_USER_EMAIL, user_role="internal_user")
+        team_models: Final = {
+            scenario.team(models=[model]): model
+            for model in (
+                scenario.model(model=f"anthropic/{_BACKEND}", api_key=_API_KEY),
+                scenario.model(model=f"anthropic/{_BACKEND}", api_key=_API_KEY),
+            )
+        }
+        for team_id in team_models:
+            enabled_gateway.post(
+                "/team/member_add", {"team_id": team_id, "member": {"role": "user", "user_id": user_id}}
+            )
+
+        authorization_response: Final = enabled_gateway.client.post(
+            "/claude_code_gateway/oauth/device_authorization",
+            data={"client_id": _CLIENT_ID},
+        )
+        assert authorization_response.status_code == 200, authorization_response.text
+        authorization: Final = _DeviceAuthorization.model_validate_json(authorization_response.content)
+        login_id: Final = authorization.device_code.partition(".")[0]
+        verification_response: Final = enabled_gateway.client.get(
+            authorization.verification_uri, follow_redirects=False
+        )
+        assert verification_response.status_code == 303, verification_response.text
+        identity_provider_response: Final = enabled_gateway.client.get(
+            verification_response.headers["location"], follow_redirects=False
+        )
+        assert identity_provider_response.status_code == 302, identity_provider_response.text
+        callback_response: Final = enabled_gateway.client.get(
+            identity_provider_response.headers["location"], follow_redirects=False
+        )
+        assert callback_response.status_code == 200, callback_response.text
+        offered_teams: Final = tuple(_TEAM_OPTION.findall(callback_response.text))
+        assert sorted(offered_teams) == sorted(team_models), callback_response.text
+        browser_token_match: Final = _BROWSER_TOKEN.search(callback_response.text)
+        assert browser_token_match is not None, callback_response.text
+        completion_form: Final = {
+            "user_code": authorization.user_code,
+            "browser_complete_token": html.unescape(browser_token_match.group(1)),
+        }
+
+        foreign_team: Final = enabled_gateway.client.post(
+            f"/sso/cli/complete/{login_id}",
+            data={**completion_form, "team_id": f"integration-{uuid.uuid4().hex}"},
+        )
+        assert foreign_team.status_code == 400, foreign_team.text
+        assert _JSON_OBJECT.validate_json(foreign_team.content) == {
+            "detail": "Choose one of the teams you belong to"
+        }, foreign_team.text
+        pending: Final = _device_code_grant(enabled_gateway, authorization.device_code)
+        assert pending.status_code == 400, pending.text
+        assert _JSON_OBJECT.validate_json(pending.content) == {"error": "authorization_pending"}, pending.text
+
+        chosen_team: Final = offered_teams[-1]
+        chosen: Final = enabled_gateway.client.post(
+            f"/sso/cli/complete/{login_id}", data={**completion_form, "team_id": chosen_team}
+        )
+        assert chosen.status_code == 200, chosen.text
+        token_response: Final = _device_code_grant(enabled_gateway, authorization.device_code)
+        assert token_response.status_code == 200, token_response.text
+        access_token: Final = _AccessToken.model_validate_json(token_response.content)
+        assert (access_token.user_id, access_token.team_id) == (user_id, chosen_team), token_response.text
+
+        with Anthropic(
+            base_url=f"{str(enabled_gateway.client.base_url).rstrip('/')}/claude_code_gateway",
+            auth_token=access_token.access_token,
+        ) as sdk:
+            served_models: Final = {model.id for model in sdk.models.list()}
+        assert served_models == {team_models[chosen_team]}
+        assert len(identity_provider.drain()) == 3
diff --git a/tests/unit/proxy/anthropic_endpoints/test_gateway_endpoints.py b/tests/unit/proxy/anthropic_endpoints/test_gateway_endpoints.py
index 0f064622f2d..8d4571c2f01 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)
@@ -397,6 +425,33 @@ def test_token_success_mints_bearer_and_is_single_use():
     assert replay.json()["error"] == "expired_token"
 
 
+@pytest.mark.parametrize(
+    ("chosen_team_id", "minted_team_id", "minted_alias"),
+    [("team-b", "team-b", "Team B"), (None, "team-a", "Team A")],
+    ids=["chosen_in_browser", "no_choice_recorded"],
+)
+def test_token_mints_for_the_team_chosen_in_the_browser(
+    chosen_team_id: str | None, minted_team_id: str, minted_alias: str
+):
+    with _gateway_env() as (client, cache):
+        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 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",
+        "error_description": "The team picked at sign-in is no longer one of your teams; sign in again",
+    }
+
+
 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 4135e9be6f6..cfb64a85359 100644
--- a/tests/unit/proxy/management_endpoints/test_ui_sso.py
+++ b/tests/unit/proxy/management_endpoints/test_ui_sso.py
@@ -3996,6 +3996,138 @@ class TestCLIKeyRegenerationFlow:
                 in result.body.decode()
             )
 
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("browser_team_choice", [True, False])
+    async def test_cli_sso_callback_offers_a_team_picker_only_when_the_flow_chooses_the_team_in_the_browser(
+        self, browser_team_choice: bool
+    ):
+        from litellm.proxy.management_endpoints.ui_sso import (
+            CLI_SSO_BROWSER_TEAM_CHOICE,
+            cli_sso_callback,
+        )
+
+        mock_request: Final = MagicMock(spec=Request)
+        mock_request.scope = {}
+        mock_request.base_url = "http://internal-proxy.local/"
+        mock_cache: Final = MagicMock(redis_cache=None)
+        mock_cache.get_cache.return_value = {
+            "poll_secret_hash": "poll-secret-hash",
+            "user_code_hash": "user-code-hash",
+            "sso_complete": False,
+            "user_code_verified": False,
+            "session_data": None,
+            **({CLI_SSO_BROWSER_TEAM_CHOICE: True} if browser_team_choice else {}),
+        }
+        mock_prisma: Final = MagicMock()
+        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), ("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", "team-ops-us"]
+                ),
+            ),
+            patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
+            patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache),
+            patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache),
+        ):
+            result: Final = await cli_sso_callback(
+                request=mock_request,
+                key="cli-session-4567890",
+                result=CustomOpenID(id="dev-1", email="dev-1@example.com", team_ids=[]),
+            )
+
+        page: Final = result.body.decode()
+        if not browser_team_choice:
+            assert 'name="team_id"' not in page
+            return
+        assert (
+            '<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
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        ("posted_team_id", "chosen_team_id", "stored_teams"),
+        [("team-claude", "team-claude", ["team-claude", "team-ops"]), ("", None, ["team-ops", "team-claude"])],
+        ids=["chosen", "not_posted"],
+    )
+    async def test_cli_sso_complete_records_the_team_chosen_in_the_browser(
+        self, posted_team_id: str, chosen_team_id: str | None, stored_teams: list[str]
+    ):
+        from litellm.proxy.management_endpoints.ui_sso import (
+            CLI_SSO_BROWSER_TEAM_CHOICE,
+            CLI_SSO_CHOSEN_TEAM_ID,
+            _hash_cli_sso_secret,
+            _normalize_cli_sso_user_code,
+            cli_sso_complete,
+        )
+
+        mock_request: Final = MagicMock(spec=Request)
+        mock_request.body = AsyncMock(
+            return_value=f"user_code=ABCD-EFGH&browser_complete_token=browser-token&team_id={posted_team_id}".encode()
+        )
+        mock_cache: Final = MagicMock(redis_cache=None)
+        mock_cache.get_cache.return_value = {
+            "poll_secret_hash": _hash_cli_sso_secret("poll-secret"),
+            "user_code_hash": _hash_cli_sso_secret(_normalize_cli_sso_user_code("ABCD-EFGH")),
+            "browser_complete_token_hash": _hash_cli_sso_secret("browser-token"),
+            "sso_complete": True,
+            "user_code_verified": False,
+            "session_data": {"user_id": "dev-1", "teams": ["team-ops", "team-claude"]},
+            CLI_SSO_BROWSER_TEAM_CHOICE: True,
+        }
+
+        with patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache):
+            result: Final = await cli_sso_complete(request=mock_request, login_id="cli-session-4567890")
+
+        assert result.status_code == 200
+        stored: Final = mock_cache.set_cache.call_args.kwargs["value"]
+        assert stored["user_code_verified"] is True
+        assert stored[CLI_SSO_CHOSEN_TEAM_ID] == chosen_team_id
+        assert stored["session_data"]["teams"] == stored_teams
+
+    @pytest.mark.asyncio
+    async def test_cli_sso_complete_refuses_a_team_the_user_is_not_on(self):
+        from litellm.proxy.management_endpoints.ui_sso import (
+            CLI_SSO_BROWSER_TEAM_CHOICE,
+            _hash_cli_sso_secret,
+            _normalize_cli_sso_user_code,
+            cli_sso_complete,
+        )
+
+        mock_request: Final = MagicMock(spec=Request)
+        mock_request.body = AsyncMock(
+            return_value=b"user_code=ABCD-EFGH&browser_complete_token=browser-token&team_id=team-finance"
+        )
+        mock_cache: Final = MagicMock(redis_cache=None)
+        mock_cache.get_cache.return_value = {
+            "poll_secret_hash": _hash_cli_sso_secret("poll-secret"),
+            "user_code_hash": _hash_cli_sso_secret(_normalize_cli_sso_user_code("ABCD-EFGH")),
+            "browser_complete_token_hash": _hash_cli_sso_secret("browser-token"),
+            "sso_complete": True,
+            "user_code_verified": False,
+            "session_data": {"user_id": "dev-1", "teams": ["team-ops", "team-claude"]},
+            CLI_SSO_BROWSER_TEAM_CHOICE: True,
+        }
+
+        with patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache):
+            with pytest.raises(HTTPException) as exc_info:
+                await cli_sso_complete(request=mock_request, login_id="cli-session-4567890")
+
+        assert exc_info.value.status_code == 400
+        assert exc_info.value.detail == "Choose one of the teams you belong to"
+        mock_cache.set_cache.assert_not_called()
+
     @pytest.mark.asyncio
     async def test_cli_poll_key_returns_teams_for_selection(self):
         """Test CLI poll endpoint returns teams for user selection when multiple teams exist"""