diff --git a/litellm/litellm_core_utils/cli_credential_lock.py b/litellm/litellm_core_utils/cli_credential_lock.py new file mode 100644 index 00000000000..d310936a2ee --- /dev/null +++ b/litellm/litellm_core_utils/cli_credential_lock.py @@ -0,0 +1,107 @@ +import ctypes.wintypes +import errno +import os +import stat +import sys +import time +from collections.abc import Generator +from contextlib import contextmanager +from hashlib import sha256 +from pathlib import Path +from typing import TYPE_CHECKING, Final, cast + +from filelock import Timeout + +if TYPE_CHECKING or sys.platform != "win32": + import fcntl + + +@contextmanager +def credential_lock(home: Path, timeout: float = 30) -> Generator[None, None, None]: + """Serialize credential changes on this host without writing inside the home directory.""" + if sys.platform == "win32": + with _windows_mutex(home, timeout): + yield + return + fd: Final = _open_posix_lock(home) + try: + deadline: Final = time.monotonic() + timeout + while True: + try: + fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB) + break + except OSError as error: + if error.errno not in (errno.EACCES, errno.EAGAIN, errno.EWOULDBLOCK): + raise + if time.monotonic() >= deadline: + raise Timeout(str(home)) from None + time.sleep(0.05) + try: + os.utime(fd, None) + if os.fstat(fd).st_nlink != 1: + raise OSError("The CLI lock file was removed while waiting") + yield + finally: + fcntl.flock(fd, fcntl.LOCK_UN) + finally: + os.close(fd) + + +def _open_posix_lock(home: Path) -> int: + if sys.platform == "win32": + raise OSError("POSIX lock files are unavailable on Windows") + directory: Final = Path("/tmp") / f"litellm-cli-{os.getuid()}" + directory.mkdir(mode=0o700, exist_ok=True) + directory_fd: Final = os.open(directory, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW) + try: + directory_stat: Final = os.fstat(directory_fd) + if directory_stat.st_uid != os.getuid() or stat.S_IMODE(directory_stat.st_mode) & 0o077: + raise PermissionError("The CLI lock directory must be private and owned by the current user") + identity: Final = sha256(str(home.resolve()).encode()).hexdigest() + fd: Final = os.open( + f"{identity}.lock", os.O_RDWR | os.O_CREAT | os.O_NOFOLLOW | os.O_NONBLOCK, 0o600, dir_fd=directory_fd + ) + try: + file_stat: Final = os.fstat(fd) + if not stat.S_ISREG(file_stat.st_mode) or file_stat.st_uid != os.getuid() or file_stat.st_nlink != 1: + raise PermissionError("The CLI lock must be a regular file owned only by the current user") + return fd + except OSError: + os.close(fd) + raise + finally: + os.close(directory_fd) + + +@contextmanager +def _windows_mutex(home: Path, timeout: float) -> Generator[None, None, None]: + kernel: Final = ctypes.WinDLL("kernel32", use_last_error=True) + create: Final = ctypes.WINFUNCTYPE( + ctypes.wintypes.HANDLE, ctypes.c_void_p, ctypes.wintypes.BOOL, ctypes.wintypes.LPCWSTR, use_last_error=True + )(("CreateMutexW", kernel)) + wait: Final = ctypes.WINFUNCTYPE( + ctypes.wintypes.DWORD, ctypes.wintypes.HANDLE, ctypes.wintypes.DWORD, use_last_error=True + )(("WaitForSingleObject", kernel)) + release: Final = ctypes.WINFUNCTYPE(ctypes.wintypes.BOOL, ctypes.wintypes.HANDLE, use_last_error=True)( + ("ReleaseMutex", kernel) + ) + close: Final = ctypes.WINFUNCTYPE(ctypes.wintypes.BOOL, ctypes.wintypes.HANDLE, use_last_error=True)( + ("CloseHandle", kernel) + ) + identity: Final = sha256(os.path.normcase(str(home.resolve())).encode()).hexdigest() + handle: Final = cast(int | None, create(None, False, f"Global\\litellm-cli-{identity}")) + if handle is None: + raise ctypes.WinError(ctypes.get_last_error()) + try: + result: Final = cast(int, wait(handle, max(0, int(timeout * 1000)))) + if result == 0x102: + raise Timeout(str(home)) + if result not in (0, 0x80): + raise ctypes.WinError(ctypes.get_last_error()) + try: + yield + finally: + if not release(handle) and sys.exc_info()[0] is None: + raise ctypes.WinError(ctypes.get_last_error()) + finally: + close(handle) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index d42c1c6b879..d0879b41728 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -1912,6 +1912,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): @@ -1926,6 +1927,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 e66504af47a..a98cc1d0a2e 100644 --- a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py +++ b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py @@ -563,6 +563,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 @@ -584,6 +585,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, @@ -598,6 +604,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 bb063f8f77a..76f41114db0 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -33819,6 +33819,22 @@ ], "title": "Resource" } + }, + { + "in": "query", + "name": "team", + "required": false, + "schema": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Team" + } } ], "responses": { @@ -34632,6 +34648,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 8c05264c6e6..1dd8f6dce65 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. @@ -697,6 +699,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/_cli_context.py b/litellm/proxy/client/cli/commands/_cli_context.py index 74c29653d16..5ccbb30392c 100644 --- a/litellm/proxy/client/cli/commands/_cli_context.py +++ b/litellm/proxy/client/cli/commands/_cli_context.py @@ -1,7 +1,7 @@ from typing import Final import click -from typing_extensions import ReadOnly, TypedDict +from typing_extensions import NotRequired, ReadOnly, TypedDict class CliContextValues(TypedDict): @@ -9,6 +9,7 @@ class CliContextValues(TypedDict): base_url: ReadOnly[str] api_key: ReadOnly[str | None] + api_key_from_token_file: ReadOnly[NotRequired[bool]] _UNSET_CLI_CONTEXT: Final[CliContextValues] = {"base_url": "", "api_key": None} diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index 98af32fa7aa..0bba43cd75a 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -3,16 +3,20 @@ import sys import time import webbrowser from collections.abc import Callable, Mapping, Sequence +from contextlib import AbstractContextManager +from pathlib import Path from typing import Any, Final, TypeVar from urllib.parse import urlencode import click import requests +from filelock import Timeout from rich.console import Console from rich.table import Table from typing_extensions import NotRequired, ReadOnly, TypedDict, assert_never from litellm.constants import CLI_JWT_EXPIRATION_HOURS +from litellm.litellm_core_utils.cli_credential_lock import credential_lock from litellm.litellm_core_utils.cli_keyring import ( DISABLE_KEYRING_ENV_VAR, SYSTEM_KEYRING, @@ -239,6 +243,10 @@ def _renewal_reader(vault: SecretVault) -> Callable[[], Mapping[str, object] | N return reload +def _credential_lock() -> AbstractContextManager[None]: + return credential_lock(Path.home()) + + def get_stored_api_key( expected_base_url: str | None = None, *, @@ -256,6 +264,26 @@ def get_stored_api_key( return None if expected_base_url is not None and token_data.get("base_url") != expected_base_url.rstrip("/"): return None + if is_cli_token_fresh(token_data) or not token_data.get("refresh_token"): + return _key_from_record({**token_data, "refresh_token": None}, vault) + try: + with _credential_lock(): + return _get_stored_api_key(expected_base_url, vault) + except (OSError, Timeout) as error: + _warn(f"Could not lock the saved login: {error}") + return None + + +def _get_stored_api_key(expected_base_url: str | None, vault: SecretVault) -> str | None: + token_data: Final = load_token(vault=vault) + if token_data is None: + return None + if expected_base_url is not None and token_data.get("base_url") != expected_base_url.rstrip("/"): + return None + return _key_from_record(token_data, vault) + + +def _key_from_record(token_data: Mapping[str, object], vault: SecretVault) -> str | None: return fresh_api_key( token_data, _renewal_saver(vault), @@ -619,7 +647,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, required_team_id: str | None = None +) -> CliAuthResult | None: """ Poll the server for authentication completion and handle team selection. @@ -639,6 +683,11 @@ def _poll_for_authentication(base_url: str, key_id: str, poll_secret: str) -> Cl team_details: Final = data.get("team_details") user_id = data.get("user_id") normalized_teams: Final[list[CliTeam]] = _normalize_teams(teams, team_details) + if ( + required_team_id is not None + and match_requested_team(normalized_teams, required_team_id) != required_team_id + ): + raise click.ClickException("The requested team is not available for this login") if not normalized_teams: click.echo("Warning: No teams available for selection.") return None @@ -649,6 +698,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 @@ -657,7 +707,7 @@ def _poll_for_authentication(base_url: str, key_id: str, poll_secret: str) -> Cl "api_key": jwt_with_team, "user_id": user_id, "teams": teams, - "team_id": None, # Set by server in JWT + "team_id": match_requested_team(normalized_teams, team), } click.echo("Team selection cancelled or JWT generation failed.") @@ -685,7 +735,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 +757,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.") @@ -828,10 +885,43 @@ def _finish_login(base_url: str, api_key: str, config_claude: bool, stored: Secr show_commands() -def _replace_stored_token(record: CliTokenData, http: Http, vault: SecretVault) -> SecretSave: +def _replace_stored_token( + record: CliTokenData, http: Http, vault: SecretVault, required_team_id: str | None = None +) -> SecretSave: + if required_team_id is not None and record.get("team_id") != required_team_id: + refused_revocation: Final = revoke_stored_credential(record, http) + if refused_revocation is not None: + click.echo( + f"Could not revoke the rejected login's refresh token on the proxy ({refused_revocation.reason}); " + "it expires on its own." + ) + raise click.ClickException("The login did not select the requested team; your saved login has not changed") + try: + with _credential_lock(): + return _persist_replacement(record, http, vault) + except (OSError, Timeout) as error: + return CredentialNotSaved(f"Could not lock the saved login: {error}") + + +def _persist_replacement(record: CliTokenData, http: Http, vault: SecretVault) -> SecretSave: previous: Final = load_token(vault=vault) + previous_secret: Final = vault.read() stored: Final = save_token(record, vault=vault) - if previous is None or isinstance(stored, CredentialNotSaved): + if previous is not None and isinstance(stored, CredentialNotRecorded): + restored: Final = ( + vault.write(previous_secret.blob) if isinstance(previous_secret, SecretFound) else previous_secret + ) + if isinstance(restored, SecretStored): + abandoned_revocation: Final = revoke_stored_credential(record, http) + if abandoned_revocation is not None: + click.echo( + "Could not revoke the abandoned login's refresh token " + f"on the proxy ({abandoned_revocation.reason}); " + "it expires on its own." + ) + return CredentialNotSaved("The replacement could not be recorded; your previous login was restored") + click.echo("Could not restore the previous login after the partial save; sign in again to repair it.") + if previous is None or isinstance(stored, (CredentialNotSaved, CredentialNotRecorded)): return stored revocation: Final = revoke_stored_credential(previous, http) if revocation is not None: @@ -842,14 +932,17 @@ 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, required_team_id: str | None = None +) -> bool: 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 - stored: Final = _replace_stored_token(pkce_token_record(base_url, credential), http, vault) + return False + stored: Final = _replace_stored_token(pkce_token_record(base_url, credential), http, vault, required_team_id) _finish_login(base_url, credential.access_token, config_claude, stored) + return not isinstance(stored, (CredentialNotSaved, CredentialNotRecorded)) @click.command(name="login") @@ -866,15 +959,33 @@ 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""" + login_to_proxy(ctx, config_claude, pkce, team) + + +def login_to_proxy( + ctx: click.Context, config_claude: bool, pkce: bool, team: str | None, required_team_id: str | None = None +) -> bool: from litellm.constants import LITELLM_CLI_SOURCE_IDENTIFIER ctx_obj: Final[CliContextObj] = ctx.obj @@ -888,8 +999,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)) - return + return _pkce_login(base_url, config_claude, context_secret_vault(ctx), team, required_team_id) cli_sso_flow: Final = _start_cli_sso_flow(base_url=base_url) key_id: Final = cli_sso_flow["login_id"] poll_secret: Final = cli_sso_flow["poll_secret"] @@ -919,7 +1029,13 @@ 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, required_team_id=required_team_id + ) + if required_team_id is not None + else _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"] @@ -937,31 +1053,33 @@ def login(ctx: click.Context, config_claude: bool, pkce: bool) -> None: "auth_header_name": "Authorization", "jwt_token": "", "timestamp": time.time(), + "team_id": auth_result["team_id"], }, requests.Session(), context_secret_vault(ctx), + required_team_id, ) _finish_login(base_url, api_key, config_claude, stored) - return + return not isinstance(stored, (CredentialNotSaved, CredentialNotRecorded)) else: click.echo("Authentication timed out. Please try again.") click.echo( "The proxy never reported the browser sign-in as finished. If you did complete it, " "check the proxy logs for /sso/callback errors and confirm SSO is configured on the proxy." ) - return + return False except KeyboardInterrupt: click.echo("\nAuthentication cancelled by user.") - return + return False except click.ClickException: # Login itself already succeeded; only the post-login step failed, so this # must not be relabelled as an authentication failure by the handler below. raise except Exception as e: click.echo(f"Authentication failed: {e}") - return + return False @click.command(name="logout") @@ -969,6 +1087,14 @@ def login(ctx: click.Context, config_claude: bool, pkce: bool) -> None: def logout(ctx: click.Context): """Logout and clear stored authentication""" vault: Final = context_secret_vault(ctx) + try: + with _credential_lock(): + _logout(vault) + except (OSError, Timeout) as error: + raise click.ClickException(f"Could not lock the saved login: {error}") from error + + +def _logout(vault: SecretVault) -> None: token_data: Final = load_token(vault=vault) revocation: Final = revoke_stored_credential(token_data, requests.Session()) if token_data is not None else None match revocation: @@ -1042,15 +1168,13 @@ def print_token(ctx: click.Context): click.echo(keychain_unreadable_notice(vault), err=True) sys.exit(1) + saved_base_url: Final = token_data.get("base_url") api_key: Final = ( ctx_obj.get("api_key") if issued_for_this_server and ctx_obj.get("api_key_from_token_file") - else fresh_api_key( - token_data, - _renewal_saver(vault), - requests.Session(), - reload=_renewal_reader(vault), - warn=_warn, + else get_stored_api_key( + saved_base_url if isinstance(saved_base_url, str) else None, + vault=vault, ) ) if not 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/client/cli/commands/teams.py b/litellm/proxy/client/cli/commands/teams.py index 1a941786f19..fe47254102e 100644 --- a/litellm/proxy/client/cli/commands/teams.py +++ b/litellm/proxy/client/cli/commands/teams.py @@ -12,6 +12,7 @@ from typing_extensions import ReadOnly, TypedDict from litellm.proxy.client import Client from ._cli_context import cli_context_values +from .auth import context_secret_vault, load_token, login_to_proxy class _TeamRow(TypedDict): @@ -141,6 +142,28 @@ def assign_key(ctx: click.Context, team_id: str | None): click.echo("No API key found. Please login first using 'litellm login'") raise click.Abort() + stored_token: Final = load_token(vault=context_secret_vault(ctx)) + if ( + stored_token is not None + and stored_token.get("base_url") == context["base_url"].rstrip("/") + and (context.get("api_key_from_token_file", False) or stored_token.get("key") == api_key) + and (stored_token.get("refresh_token") or not api_key.startswith("sk-")) + ): + if not context.get("api_key_from_token_file", False): + raise click.ClickException("Unset --api-key and LITELLM_PROXY_API_KEY to switch your saved CLI session") + click.echo("Signing in again to select the team for your CLI session") + saved: Final = login_to_proxy( + ctx, + config_claude=False, + pkce=bool(stored_token.get("refresh_token")), + team=team_id, + required_team_id=team_id, + ) + if not saved: + raise click.ClickException("CLI session team assignment did not complete") + click.echo(f"Successfully assigned CLI session to team: {team_id}" if team_id else "CLI session team selected") + return + try: # If no team_id provided, show teams and let user select if not team_id: 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 2943ff4b74a..03c2af2bf2c 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 @@ -2137,6 +2137,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 ecfabf33e87..532dafe32b9 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -48710,6 +48710,7 @@ export interface operations { response_type?: string | null; scope?: string | null; resource?: string | null; + team?: string | null; }; header?: never; path?: never; @@ -77885,6 +77886,7 @@ export interface operations { response_type?: string | null; scope?: string | null; resource?: string | null; + team?: string | null; }; header?: never; path: {