Merge 422246978d into 9071ca503e

litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py

@@ -1879,6 +1879,7 @@ async def authorize(
     response_type: str | None = None,
     scope: str | None = None,
     resource: str | None = None,
+    team: str | None = None,
 ):
     # Redirect to real OAuth provider with PKCE support
     if mcp_server_name is None and client_id and is_gateway_dcr_client_id(client_id):
@@ -1893,6 +1894,7 @@ async def authorize(
                 response_type=response_type,
                 session_user_id=_session_cookie_user_id(request),
                 lookup_consent_teams=lookup_consent_teams,
+                requested_team=team,
             )
         return aggregate_authorize(
             request=request,

litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py

@@ -509,6 +509,7 @@ async def native_client_authorize(
     response_type: str | None,
     session_user_id: str | None,
     lookup_consent_teams: LookupConsentTeams,
+    requested_team: str | None = None,
 ) -> Response:
     """The authorize verb for a native client that named the proxy API itself as its
     RFC 8707 ``resource``: the same client, redirect, PKCE, and sign-in checks as the
@@ -530,6 +531,11 @@ async def native_client_authorize(
     teams: Final = await lookup_consent_teams(session_user_id)
     if not isinstance(teams, tuple):
         return _consent_lookup_failure_response(teams)
+    selected_team_id: Final = (
+        next((t.team_id for t in teams if requested_team in (t.team_id, t.team_alias)), None)
+        if requested_team
+        else None
+    )
     handle: Final = secrets.token_urlsafe(24)
     flow: Final = _new_connect_flow(
         session_user_id=session_user_id,
@@ -544,6 +550,7 @@ async def native_client_authorize(
         client_origin=_origin_only(redirect_uri),
         user_id=session_user_id,
         teams=tuple((team.team_id, team.team_alias or team.team_id) for team in teams),
+        selected_team_id=selected_team_id,
         flow_handle=handle,
         complete_url=f"{base_url}/authorize/complete",
     )

litellm/proxy/_lazy_openapi_snapshot.json

@@ -24918,6 +24918,22 @@
                 ],
                 "title": "Resource"
               }
+            },
+            {
+              "in": "query",
+              "name": "team",
+              "required": false,
+              "schema": {
+                "anyOf": [
+                  {
+                    "type": "string"
+                  },
+                  {
+                    "type": "null"
+                  }
+                ],
+                "title": "Team"
+              }
             }
           ],
           "responses": {
@@ -25731,6 +25747,22 @@
                 ],
                 "title": "Resource"
               }
+            },
+            {
+              "in": "query",
+              "name": "team",
+              "required": false,
+              "schema": {
+                "anyOf": [
+                  {
+                    "type": "string"
+                  },
+                  {
+                    "type": "null"
+                  }
+                ],
+                "title": "Team"
+              }
             }
           ],
           "responses": {

litellm/proxy/client/cli/README.md

@@ -506,6 +506,8 @@ The token minted by `lite login` is a short-lived, per-session agent credential,
 
 The credential is short-lived by design (default 24h, configurable via `LITELLM_CLI_JWT_EXPIRATION_HOURS`); run `lite login` again to refresh it, which also re-reads your latest team and user settings. It does not appear in the Keys UI and cannot be rotated or revoked mid-session. `lite auth print-token` (usable as Claude Code's `apiKeyHelper`) prints it while it's still fresh and fails once it expires -- there is no silent renewal, so a long-running session needs a fresh `lite login` once a day. `lite claude`, `lite codex`, `lite opencode`, and `lite pi` work with it on a default deployment; `EXPERIMENTAL_UI_LOGIN` is not required. `lite login --pkce` is the exception to the daily re-login: it signs in through your system browser with OAuth authorization code and PKCE and stores a refresh token next to the key, so every `lite` command and `lite auth print-token` renew the key on their own shortly before it expires, `lite whoami` shows when the current key expires, and `lite logout` revokes the refresh token on the proxy (it needs a proxy that serves `/.well-known/litellm-cli-auth`; see [Browser sign-in with PKCE](https://docs.litellm.ai/docs/proxy/cli_sso#browser-sign-in-with-pkce)). When a renewal is refused, for example after a `lite logout` run from another copy of the credential, the command prints why on stderr and, once the key has run out, tells you to run `lite login --pkce` again. Only the holder can end a `--pkce` session early, with `lite logout`; an admin has no button for it, but every renewal re-reads the user on the proxy, so deactivating the user or removing them from the team makes the next renewal fail and the key runs out within `LITELLM_CLI_JWT_EXPIRATION_HOURS`. On a proxy with more than one worker or replica, configure Redis (`litellm_settings.cache` with Redis `cache_params`, or `general_settings.coordination_redis`) so a refresh token stays single-use and `lite logout` holds on every worker; without Redis each worker keeps its own record. If you need a long-lived, rotatable key that shows up in the Keys UI, create a dedicated virtual key in the dashboard and pass it via `--api-key` or `LITELLM_PROXY_API_KEY` instead.
 
+If you belong to several teams, `lite login` normally asks which one to attribute the credential to. Pass `--team <id-or-alias>` (or export `LITELLM_PROXY_TEAM`) to skip that pick: when the value matches one of your teams by id or alias it is used directly, otherwise the pick still runs. `LITELLM_PROXY_TEAM` applies to both the SSO and the `--pkce` flow, and `LITELLM_PROXY_LOGIN_PKCE=true` is equivalent to passing `--pkce`
+
 ### Route Every Claude Code Session Through the Proxy
 
 `lite claude` wraps a single invocation, but `lite up` goes further: it patches `~/.claude/settings.json`, Claude Code's own config file, so that every Claude Code session started afterward -- from any terminal, launched normally with just `claude`, no wrapper needed -- routes through your LiteLLM proxy. It sets `env.ANTHROPIC_BASE_URL` to the proxy URL, `env.ENABLE_TOOL_SEARCH` to `true` and `env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY` to `1` when those keys are missing, writes the key it resolved (your fresh `lite login`, or an explicit `--api-key`) into `env.ANTHROPIC_AUTH_TOKEN` as a static token, drops any stray `ANTHROPIC_API_KEY` or `apiKeyHelper` so nothing fights that token, and leaves every other setting in the file untouched. It backs up the original file before patching it. Nothing here writes an `apiKeyHelper`: Claude Code would spawn `lite` (and its keychain check) on every credential refresh, so the key is copied in instead and `lite up` restores the file when it stops.
@@ -657,6 +659,8 @@ The CLI respects the following environment variables:
 
 - `LITELLM_PROXY_URL`: Base URL of the proxy server
 - `LITELLM_PROXY_API_KEY`: API key for authentication
+- `LITELLM_PROXY_TEAM`: Default team (id or alias) for `lite login`, equivalent to `--team`
+- `LITELLM_PROXY_LOGIN_PKCE`: Set to `true` to make `lite login` use the PKCE browser flow, equivalent to `--pkce`
 
 `LITELLM_PROXY_URL` takes precedence over a `base_url` stored via `lite config set`, and the `--base-url` option overrides both. See the Configuration section for the full precedence order.
 

litellm/proxy/client/cli/commands/auth.py

@@ -619,7 +619,23 @@ def _get_cli_sso_poll_headers(poll_secret: str) -> dict[str, str]:
     return {"x-litellm-cli-poll-secret": poll_secret}
 
 
-def _poll_for_authentication(base_url: str, key_id: str, poll_secret: str) -> CliAuthResult | None:
+def match_requested_team(teams: Sequence[CliTeam], requested_team: str | None) -> str | None:
+    """The team_id of the first team whose id or alias equals ``requested_team``."""
+    if requested_team is None:
+        return None
+    return next(
+        (
+            team_id
+            for team in teams
+            if (team_id := team.get("team_id")) is not None and requested_team in (team_id, team.get("team_alias"))
+        ),
+        None,
+    )
+
+
+def _poll_for_authentication(
+    base_url: str, key_id: str, poll_secret: str, team: str | None = None
+) -> CliAuthResult | None:
     """
     Poll the server for authentication completion and handle team selection.
 
@@ -649,6 +665,7 @@ def _poll_for_authentication(base_url: str, key_id: str, poll_secret: str) -> Cl
             key_id=key_id,
             poll_secret=poll_secret,
             teams=normalized_teams,
+            requested_team=team,
         )
 
         # Use the team-specific JWT if selection succeeded
@@ -685,7 +702,11 @@ def _poll_for_authentication(base_url: str, key_id: str, poll_secret: str) -> Cl
 
 
 def _handle_team_selection_during_polling(
-    base_url: str, key_id: str, poll_secret: str, teams: list[CliTeam]
+    base_url: str,
+    key_id: str,
+    poll_secret: str,
+    teams: list[CliTeam],
+    requested_team: str | None = None,
 ) -> str | None:
     """
     Handle team selection and re-poll with selected team_id.
@@ -703,7 +724,10 @@ def _handle_team_selection_during_polling(
     click.echo("\n" + "=" * 60)
     click.echo("Select a team for your CLI session...")
 
-    team_id: Final = _render_and_prompt_for_team_selection(teams)
+    matched_team_id: Final = match_requested_team(teams, requested_team)
+    if requested_team is not None and matched_team_id is None:
+        click.echo(f"Team '{requested_team}' was not found among your teams; select one below.")
+    team_id: Final = matched_team_id or _render_and_prompt_for_team_selection(teams)
 
     if not team_id:
         click.echo("No team selected.")
@@ -842,9 +866,9 @@ def _replace_stored_token(record: CliTokenData, http: Http, vault: SecretVault)
     return stored
 
 
-def _pkce_login(base_url: str, config_claude: bool, vault: SecretVault) -> None:
+def _pkce_login(base_url: str, config_claude: bool, vault: SecretVault, team: str | None) -> None:
     http: Final = requests.Session()
-    credential: Final = run_pkce_login(base_url, http, echo=click.echo)
+    credential: Final = run_pkce_login(base_url, http, echo=click.echo, team=team)
     if isinstance(credential, PkceFailure):
         click.echo(f"Authentication failed: {credential.reason}")
         return
@@ -866,14 +890,26 @@ def _pkce_login(base_url: str, config_claude: bool, vault: SecretVault) -> None:
     "--pkce",
     is_flag=True,
     default=False,
+    envvar="LITELLM_PROXY_LOGIN_PKCE",
+    show_envvar=True,
     help=(
         "Sign in with OAuth authorization code + PKCE through your system browser (loopback redirect), "
         "with a refresh token that renews the key automatically. Requires a proxy that serves "
         "/.well-known/litellm-cli-auth."
     ),
 )
+@click.option(
+    "--team",
+    envvar="LITELLM_PROXY_TEAM",
+    show_envvar=True,
+    default=None,
+    help=(
+        "Team id or alias to attribute this login to. Skips the team pick when it matches one of your "
+        "teams; otherwise you pick as usual."
+    ),
+)
 @click.pass_context
-def login(ctx: click.Context, config_claude: bool, pkce: bool) -> None:
+def login(ctx: click.Context, config_claude: bool, pkce: bool, team: str | None) -> None:
     """Login to LiteLLM proxy using SSO authentication"""
     from litellm.constants import LITELLM_CLI_SOURCE_IDENTIFIER
 
@@ -888,7 +924,7 @@ def login(ctx: click.Context, config_claude: bool, pkce: bool) -> None:
 
     try:
         if pkce:
-            _pkce_login(base_url, config_claude, context_secret_vault(ctx))
+            _pkce_login(base_url, config_claude, context_secret_vault(ctx), team)
             return
         cli_sso_flow: Final = _start_cli_sso_flow(base_url=base_url)
         key_id: Final = cli_sso_flow["login_id"]
@@ -919,7 +955,9 @@ def login(ctx: click.Context, config_claude: bool, pkce: bool) -> None:
         # Poll for authentication completion
         click.echo("Waiting for authentication...")
 
-        auth_result: Final = _poll_for_authentication(base_url=base_url, key_id=key_id, poll_secret=poll_secret)
+        auth_result: Final = _poll_for_authentication(
+            base_url=base_url, key_id=key_id, poll_secret=poll_secret, team=team
+        )
 
         if auth_result:
             api_key: Final = auth_result["api_key"]

litellm/proxy/client/cli/commands/pkce_login.py

@@ -304,7 +304,14 @@ def pkce_pair() -> tuple[str, str]:
     return verifier, urlsafe_b64encode(digest).rstrip(b"=").decode("ascii")
 
 
-def authorize_url(contract: CliAuthContract, client_id: str, redirect_uri: str, state: str, code_challenge: str) -> str:
+def authorize_url(
+    contract: CliAuthContract,
+    client_id: str,
+    redirect_uri: str,
+    state: str,
+    code_challenge: str,
+    team: str | None = None,
+) -> str:
     query: Final = urlencode(
         _form(
             response_type="code",
@@ -314,6 +321,7 @@ def authorize_url(contract: CliAuthContract, client_id: str, redirect_uri: str,
             code_challenge=code_challenge,
             code_challenge_method="S256",
             resource=contract.resource,
+            **({"team": team} if team is not None else {}),
         )
     )
     return f"{contract.authorization_endpoint}?{query}"
@@ -440,6 +448,7 @@ def run_pkce_login(
     open_browser: Callable[[str], object] = webbrowser.open,
     echo: Callable[[str], None] = print,
     timeout_seconds: float = LOGIN_TIMEOUT_SECONDS,
+    team: str | None = None,
 ) -> PkceCredential | PkceFailure:
     contract: Final = discover_cli_auth(base_url, http)
     if isinstance(contract, PkceFailure):
@@ -450,7 +459,7 @@ def run_pkce_login(
         client_id: Final = register_client(contract, server.redirect_uri, http)
         if isinstance(client_id, PkceFailure):
             return client_id
-        url: Final = authorize_url(contract, client_id, server.redirect_uri, state, challenge)
+        url: Final = authorize_url(contract, client_id, server.redirect_uri, state, challenge, team=team)
         echo(f"Opening browser to: {url}")
         echo("Approve the sign-in in your browser. Waiting...")
         threading.Thread(target=open_browser, args=(url,), name="lite-login-browser", daemon=True).start()

litellm/proxy/common_utils/html_forms/native_client_consent.py

@@ -12,6 +12,7 @@ def render_native_client_consent_page(
     teams: Sequence[tuple[str, str]],
     flow_handle: str,
     complete_url: str,
+    selected_team_id: str | None = None,
 ) -> str:
     """The consent page a native client's sign-in lands on: who is signed in, which
     loopback client asked, which team the credential is attributed to, and an explicit
@@ -62,7 +63,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_field(teams, selected_team_id)}
 <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 +75,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_field(teams: Sequence[tuple[str, str]], selected_team_id: str | None = None) -> str:
     if not teams:
         return ""
     if len(teams) == 1:
@@ -84,7 +85,9 @@ def _team_field(teams: Sequence[tuple[str, str]]) -> str:
             f"<p>Requests are attributed to team <strong>{escape(team_label)}</strong>.</p>"
         )
     options: Final = "".join(
-        f'<option value="{escape(team_id)}">{escape(team_label)}</option>' for team_id, team_label in teams
+        f'<option value="{escape(team_id)}"{" selected" if team_id == selected_team_id else ""}>'
+        f"{escape(team_label)}</option>"
+        for team_id, team_label in teams
     )
     return (
         f'<label for="team_id">Attribute requests to team</label><select id="team_id" name="team_id">{options}</select>'

tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py

@@ -2108,6 +2108,26 @@ async def test_introspect_fails_closed_on_dead_user_and_503s_on_outage():
     assert (status, body["error"]) == (500, "server_error")
 
 
+@pytest.mark.asyncio
+async def test_native_authorize_preselects_the_requested_team_alias_on_the_consent_page():
+    client_id = (await _register([LOOPBACK_REDIRECT_URI]))["client_id"]
+    response = await _native_authorize(client_id, requested_team="Team A")
+    assert response.status_code == 200
+    body = response.body.decode()
+    assert '<option value="team-a" selected>Team A</option>' in body
+    assert '<option value="team-b">team-b</option>' in body
+
+
+@pytest.mark.asyncio
+async def test_native_authorize_with_an_unknown_requested_team_renders_a_plain_chooser():
+    client_id = (await _register([LOOPBACK_REDIRECT_URI]))["client_id"]
+    response = await _native_authorize(client_id, requested_team="no-such-team")
+    assert response.status_code == 200
+    body = response.body.decode()
+    assert "<select" in body
+    assert " selected" not in body
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
     "auth_type", [None, "none", "api_key", "bearer_token", "basic", "authorization", "token", "aws_sigv4"]

tests/test_litellm/proxy/auth/test_cli_auth.py

@@ -276,6 +276,7 @@ async def test_poll_for_authentication_team_selection_success(
             {"team_id": "1", "team_alias": None},
             {"team_id": "2", "team_alias": None},
         ],
+        requested_team=None,
     )
     click_mock.assert_not_called()
 
@@ -310,6 +311,7 @@ async def test_poll_for_authentication_team_selection_cancelled(
         key_id="key-123",
         poll_secret="poll-secret",
         teams=[{"team_id": "team-1", "team_alias": None}],
+        requested_team=None,
     )
     click_mock.assert_called_once()
     assert "Team selection cancelled" in click_mock.call_args[0][0]

tests/test_litellm/proxy/client/cli/test_auth_commands.py

@@ -2044,3 +2044,88 @@ class TestGetStoredApiKeyRefresh:
         assert captured.out == ""
         assert captured.err == "Could not renew the key: token request failed with 503: temporarily_unavailable\n"
         save.assert_not_called()
+
+
+class TestRequestedTeamLoginOption:
+    """`lite login --team` (or LITELLM_PROXY_TEAM) pre-picks a team in both login flows."""
+
+    _TEAMS = [
+        {"team_id": "team-alpha", "team_alias": "Alpha Team"},
+        {"team_id": "team-beta", "team_alias": "Beta Team"},
+    ]
+
+    def setup_method(self):
+        self.runner = CliRunner()
+
+    def test_match_requested_team_resolves_ids_and_aliases(self):
+        from litellm.proxy.client.cli.commands.auth import match_requested_team
+
+        assert match_requested_team(self._TEAMS, "team-beta") == "team-beta"
+        assert match_requested_team(self._TEAMS, "Alpha Team") == "team-alpha"
+
+    def test_match_requested_team_returns_none_without_a_match_or_a_request(self):
+        from litellm.proxy.client.cli.commands.auth import match_requested_team
+
+        assert match_requested_team(self._TEAMS, "no-such-team") is None
+        assert match_requested_team(self._TEAMS, None) is None
+        assert match_requested_team([], "team-alpha") is None
+
+    def test_matching_requested_team_skips_the_prompt_and_polls_for_that_team(self):
+        from litellm.proxy.client.cli.commands import auth
+
+        ready = Mock()
+        ready.status_code = 200
+        ready.json.return_value = {"status": "ready", "key": "jwt-beta"}
+
+        def prompt_must_not_run(*args, **kwargs):
+            raise AssertionError("the interactive team pick ran despite a matching --team")
+
+        with (
+            patch("requests.get", return_value=ready) as get,
+            patch("click.prompt", side_effect=prompt_must_not_run),
+        ):
+            jwt = auth._handle_team_selection_during_polling(
+                "https://test.example.com", "sess-1", "poll-secret", self._TEAMS, requested_team="Beta Team"
+            )
+
+        assert jwt == "jwt-beta"
+        assert "team_id=team-beta" in get.call_args.args[0]
+
+    def test_unknown_requested_team_warns_and_falls_back_to_the_prompt(self, capsys):
+        from litellm.proxy.client.cli.commands import auth
+
+        ready = Mock()
+        ready.status_code = 200
+        ready.json.return_value = {"status": "ready", "key": "jwt-alpha"}
+
+        with (
+            patch("requests.get", return_value=ready) as get,
+            patch("click.prompt", return_value="1") as prompt,
+        ):
+            jwt = auth._handle_team_selection_during_polling(
+                "https://test.example.com", "sess-1", "poll-secret", self._TEAMS, requested_team="no-such-team"
+            )
+
+        assert jwt == "jwt-alpha"
+        prompt.assert_called_once()
+        assert "team_id=team-alpha" in get.call_args.args[0]
+        assert "Team 'no-such-team' was not found among your teams; select one below." in capsys.readouterr().out
+
+    def test_env_vars_select_the_pkce_flow_and_pass_the_requested_team(self):
+        with (
+            patch(  # test-quality-ok: the click command exposes no injection seam for its login-flow dispatch
+                "litellm.proxy.client.cli.commands.auth._pkce_login"
+            ) as pkce_login,
+            patch(  # test-quality-ok: same; proves the SSO path stayed untouched
+                "litellm.proxy.client.cli.commands.auth._start_cli_sso_flow"
+            ) as sso_start,
+        ):
+            result = self.runner.invoke(
+                login,
+                obj={"base_url": PKCE_BASE_URL},
+                env={"LITELLM_PROXY_LOGIN_PKCE": "true", "LITELLM_PROXY_TEAM": "Beta Team"},
+            )
+
+        assert result.exit_code == 0, result.output
+        sso_start.assert_not_called()
+        assert pkce_login.call_args.args[3] == "Beta Team"

tests/test_litellm/proxy/client/cli/test_pkce_login.py

@@ -714,3 +714,13 @@ def test_revoke_stored_credential_revokes_only_pkce_records():
         "token_type_hint": "refresh_token",
         "client_id": "llm_dcrc_abc",
     }
+
+
+def test_authorize_url_includes_the_team_only_when_one_is_given():
+    with_team = authorize_url(
+        CONTRACT, "llm_dcrc_abc", "http://127.0.0.1:5/callback", "state-1", "challenge-1", team="Team A"
+    )
+    assert parse_qs(urlparse(with_team).query)["team"] == ["Team A"]
+
+    without_team = authorize_url(CONTRACT, "llm_dcrc_abc", "http://127.0.0.1:5/callback", "state-1", "challenge-1")
+    assert "team" not in parse_qs(urlparse(without_team).query)

tests/test_litellm/proxy/common_utils/html_forms/test_native_client_consent.py

@@ -65,3 +65,15 @@ def test_consent_page_promises_only_what_logout_can_deliver():
     assert f"expires within {CLI_JWT_EXPIRATION_HOURS} hours" in page
     assert "<code>lite logout</code> stops it from being renewed" in page
     assert "revoked" not in page
+
+
+def test_consent_page_marks_the_requested_team_selected_in_the_chooser():
+    page = _render(teams=(("team-a", "Team A"), ("team-b", "team-b")), selected_team_id="team-b")
+    assert '<option value="team-b" selected>team-b</option>' in page
+    assert '<option value="team-a">Team A</option>' in page
+
+
+def test_consent_page_marks_nothing_selected_without_a_matching_team():
+    for selected in (None, "team-elsewhere"):
+        page = _render(teams=(("team-a", "Team A"), ("team-b", "team-b")), selected_team_id=selected)
+        assert "selected" not in page

ui/litellm-dashboard/src/lib/http/schema.d.ts

@@ -41575,6 +41575,7 @@ export interface operations {
                 response_type?: string | null;
                 scope?: string | null;
                 resource?: string | null;
+                team?: string | null;
             };
             header?: never;
             path?: never;
@@ -68859,6 +68860,7 @@ export interface operations {
                 response_type?: string | null;
                 scope?: string | null;
                 resource?: string | null;
+                team?: string | null;
             };
             header?: never;
             path: {