From 474e60cdc555e5c46d3e77bd3c8577c27865edd9 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Thu, 29 Oct 2026 01:58:23 -0700
Subject: [PATCH] fix(gateway): list the browser-picked team first so an older
 replica mid-deploy issues it

---
 litellm/proxy/management_endpoints/ui_sso.py       | 14 +++++++++++++-
 .../anthropic_endpoints/test_gateway_endpoints.py  |  5 ++++-
 .../unit/proxy/management_endpoints/test_ui_sso.py |  7 ++++---
 3 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py
index 790556e5b7b..a980b9c33ff 100644
--- a/litellm/proxy/management_endpoints/ui_sso.py
+++ b/litellm/proxy/management_endpoints/ui_sso.py
@@ -846,7 +846,19 @@ async def cli_sso_complete(request: Request, login_id: str):
     _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},
+        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()
diff --git a/tests/unit/proxy/anthropic_endpoints/test_gateway_endpoints.py b/tests/unit/proxy/anthropic_endpoints/test_gateway_endpoints.py
index 3eab7a4d7e5..8d4571c2f01 100644
--- a/tests/unit/proxy/anthropic_endpoints/test_gateway_endpoints.py
+++ b/tests/unit/proxy/anthropic_endpoints/test_gateway_endpoints.py
@@ -446,7 +446,10 @@ def test_token_refuses_a_browser_choice_that_is_no_longer_one_of_the_session_tea
     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"
+    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():
diff --git a/tests/unit/proxy/management_endpoints/test_ui_sso.py b/tests/unit/proxy/management_endpoints/test_ui_sso.py
index a87064961fe..cfb64a85359 100644
--- a/tests/unit/proxy/management_endpoints/test_ui_sso.py
+++ b/tests/unit/proxy/management_endpoints/test_ui_sso.py
@@ -4057,12 +4057,12 @@ class TestCLIKeyRegenerationFlow:
 
     @pytest.mark.asyncio
     @pytest.mark.parametrize(
-        ("posted_team_id", "chosen_team_id"),
-        [("team-claude", "team-claude"), ("", None)],
+        ("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
+        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,
@@ -4094,6 +4094,7 @@ class TestCLIKeyRegenerationFlow:
         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):