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

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()

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():

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):