diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index bafe33d0a6b..a68580b29da 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/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, diff --git a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py index f3fdd54b39d..f4e1b330003 100644 --- a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py +++ b/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", ) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index cb7a18cd107..03d543152ac 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/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": { diff --git a/litellm/proxy/client/cli/README.md b/litellm/proxy/client/cli/README.md index 2071576a943..8314e0db282 100644 --- a/litellm/proxy/client/cli/README.md +++ b/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 ` (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. diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index 98af32fa7aa..5c2ce9eb53c 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/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"] diff --git a/litellm/proxy/client/cli/commands/pkce_login.py b/litellm/proxy/client/cli/commands/pkce_login.py index 93f5cfea21b..c53d45dc0c8 100644 --- a/litellm/proxy/client/cli/commands/pkce_login.py +++ b/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() 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..bee56970f1b 100644 --- a/litellm/proxy/common_utils/html_forms/native_client_consent.py +++ b/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

Approving issues it a personal credential that expires within {CLI_JWT_EXPIRATION_HOURS} hours. lite logout stops it from being renewed. Only approve if you started this sign-in yourself.

-{_team_field(teams)} +{_team_field(teams, selected_team_id)}
@@ -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"

Requests are attributed to team {escape(team_label)}.

" ) options: Final = "".join( - f'' for team_id, team_label in teams + f'" + for team_id, team_label in teams ) return ( f'' diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py index 7c80ee77cd7..95daa27ebb8 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py +++ b/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 '' in body + assert '' 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 "lite logout 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 '' in page + assert '' 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 diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 7ed8df6815d..bd1ada76e64 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/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: {