Merge f7e47bd108 into fb74957ddd

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)

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,

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",
     )

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": {

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.
@@ -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.
 

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}

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:

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

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

@@ -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 '<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

@@ -1,15 +1,21 @@
 import json
 import os
 import stat
+import sys
 import time
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import replace
 from pathlib import Path
+from threading import Event
 from unittest.mock import Mock, patch
 
 
 import pytest
 from click.testing import CliRunner
+from filelock import Timeout
 
 from litellm.constants import CLI_JWT_EXPIRATION_HOURS
+from litellm.litellm_core_utils.cli_credential_lock import _windows_mutex, credential_lock
 from litellm.litellm_core_utils.cli_keyring import (
     DISABLE_KEYRING_ENV_VAR,
     KeyringDisabled,
@@ -27,12 +33,19 @@ from litellm.proxy.client.cli import cli
 from litellm.proxy.client.cli.commands import claude_settings as claude_settings_module
 from litellm.proxy.client.cli.commands.claude_settings import SettingsFileOwner
 from litellm.proxy.client.cli.commands.auth import (
+    _replace_stored_token,
     get_stored_api_key,
+    load_token,
     login,
     logout,
     print_token,
     whoami,
 )
+from litellm.proxy.client.cli.commands.pkce_login import PkceFailure, RevocationUnavailable
+from tests.test_litellm_rust.support.child_interpreter import run_child_interpreter
+
+if sys.platform != "win32":
+    import fcntl
 
 
 @pytest.fixture
@@ -2044,3 +2057,634 @@ 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"
+
+
+@pytest.mark.parametrize("pkce,refreshed", [(False, False), (True, False), (True, True)])
+def test_assign_key_replaces_saved_session_and_retains_login_protocol(isolated_home, monkeypatch, pkce, refreshed):
+    monkeypatch.setenv(DISABLE_KEYRING_ENV_VAR, "1")
+    record = (
+        _pkce_record(key="session-old", expires_at=time.time() + 3600, team_id="team-a")
+        if pkce
+        else {
+            "base_url": PKCE_BASE_URL,
+            "key": "session-old",
+            "user_id": "u1",
+            "user_role": "cli",
+            "timestamp": time.time(),
+        }
+    )
+    save_cli_token(CliTokenRecord(**record))
+    _FakeSession.instances.clear()
+    with (
+        patch(
+            "litellm.proxy.client.cli.main.get_stored_api_key",
+            return_value="session-refreshed" if refreshed else "session-old",
+        ),
+        patch("litellm.proxy.client.cli.commands.teams.Client") as client,
+        patch(
+            "litellm.proxy.client.cli.commands.auth.run_pkce_login",
+            return_value=replace(_pkce_credential(), access_token="session-new"),
+        ) as run_pkce,
+        patch(
+            "litellm.proxy.client.cli.commands.auth._start_cli_sso_flow",
+            return_value={
+                "login_id": "login-id",
+                "poll_secret": "poll-secret",
+                "user_code": "ABCD-EFGH",
+            },
+        ),
+        patch(
+            "litellm.proxy.client.cli.commands.auth._poll_for_authentication",
+            return_value={
+                "api_key": "session-new",
+                "user_id": "u1",
+                "teams": ["team-a", "team-b"],
+                "team_id": "team-b",
+            },
+        ) as poll,
+        patch("litellm.proxy.client.cli.commands.auth.webbrowser.open"),
+        patch("litellm.proxy.client.cli.commands.auth.requests.Session", _FakeSession),
+        patch("litellm.proxy.client.cli.interface.show_commands"),
+    ):
+        result = CliRunner().invoke(cli, ["--base-url", PKCE_BASE_URL, "teams", "assign-key", "--team-id", "team-b"])
+
+    assert result.exit_code == 0, result.output
+    saved = load_token()
+    assert saved is not None
+    assert (saved["key"], saved["team_id"], saved["user_id"]) == ("session-new", "team-b", "u1")
+    assert "Successfully assigned CLI session to team: team-b" in result.output
+    client.return_value.keys.update.assert_not_called()
+    if pkce:
+        assert saved["refresh_token"] == "llm_srefresh_fresh"
+        assert run_pkce.call_args.kwargs["team"] == "team-b"
+        poll.assert_not_called()
+        assert (
+            next(session for session in _FakeSession.instances if session.posts).posts[0][1]["token"]
+            == "llm_srefresh_old"
+        )
+    else:
+        assert "refresh_token" not in saved
+        assert poll.call_args.kwargs["team"] == "team-b"
+        run_pkce.assert_not_called()
+
+
+@pytest.mark.parametrize("outcome", ["denied", "wrong-team", "storage-failure"])
+def test_assign_key_does_not_replace_saved_session_when_login_cannot_be_used(isolated_home, monkeypatch, outcome):
+    monkeypatch.setenv(DISABLE_KEYRING_ENV_VAR, "1")
+    save_cli_token(CliTokenRecord(**_pkce_record(key="session-old", expires_at=time.time() + 3600, team_id="team-a")))
+    before = load_token()
+    _FakeSession.instances.clear()
+    credential = (
+        PkceFailure("access_denied")
+        if outcome == "denied"
+        else replace(
+            _pkce_credential(), access_token="session-new", team_id="team-c" if outcome == "wrong-team" else "team-b"
+        )
+    )
+    with (
+        patch("litellm.proxy.client.cli.commands.teams.Client"),
+        patch("litellm.proxy.client.cli.commands.auth.run_pkce_login", return_value=credential),
+        patch(
+            "litellm.proxy.client.cli.commands.auth.save_token", return_value=CredentialNotSaved("read-only")
+        ) as save,
+        patch("litellm.proxy.client.cli.commands.auth.requests.Session", _FakeSession),
+        patch("litellm.proxy.client.cli.interface.show_commands"),
+    ):
+        result = CliRunner().invoke(cli, ["--base-url", PKCE_BASE_URL, "teams", "assign-key", "--team-id", "team-b"])
+
+    assert result.exit_code != 0, result.output
+    assert load_token() == before
+    assert "Successfully assigned" not in result.output
+    if outcome != "storage-failure":
+        save.assert_not_called()
+    posts = [post for session in _FakeSession.instances for post in session.posts]
+    assert posts == (
+        [
+            (
+                f"{PKCE_BASE_URL}/revoke",
+                {"token": "llm_srefresh_fresh", "token_type_hint": "refresh_token", "client_id": "llm_dcrc_abc"},
+            )
+        ]
+        if outcome == "wrong-team"
+        else []
+    )
+
+
+def test_assign_key_keeps_explicit_virtual_key_update_with_saved_session(isolated_home, monkeypatch):
+    monkeypatch.setenv(DISABLE_KEYRING_ENV_VAR, "1")
+    save_cli_token(CliTokenRecord(**_pkce_record(key="session-old", expires_at=time.time() + 3600)))
+    before = load_token()
+    with (
+        patch("litellm.proxy.client.cli.commands.teams.Client") as client,
+        patch("litellm.proxy.client.cli.commands.auth.run_pkce_login") as run_pkce,
+    ):
+        client.return_value.teams.list.return_value = []
+        result = CliRunner().invoke(
+            cli,
+            [
+                "--base-url",
+                PKCE_BASE_URL,
+                "--api-key",
+                "sk-virtual-key",
+                "teams",
+                "assign-key",
+                "--team-id",
+                "team-b",
+            ],
+        )
+
+    assert result.exit_code == 0, result.output
+    client.return_value.keys.update.assert_called_once_with(key="sk-virtual-key", team_id="team-b")
+    run_pkce.assert_not_called()
+    assert load_token() == before
+
+
+@pytest.mark.parametrize("restore_fails", [False, True])
+def test_assign_key_handles_keychain_update_without_metadata(isolated_home, secret_vault_factory, restore_fails):
+    vault = secret_vault_factory()
+    save_cli_token(
+        CliTokenRecord(
+            **_pkce_record(key="session-old", expires_at=time.time() + 3600, team_id="team-a", client_id="old-client")
+        ),
+        vault=vault,
+    )
+    before = load_token(vault=vault)
+    metadata_before = (isolated_home / ".litellm" / "token.json").read_bytes()
+    _FakeSession.instances.clear()
+
+    def refuse_metadata(*_):
+        if restore_fails:
+            vault.writable = False
+        raise OSError("read-only")
+
+    with (
+        patch("litellm.proxy.client.cli.commands.teams.Client"),
+        patch(
+            "litellm.proxy.client.cli.commands.auth.run_pkce_login",
+            return_value=replace(_pkce_credential(), access_token="session-new", team_id="team-b"),
+        ),
+        patch("litellm.litellm_core_utils.cli_token_utils.commit_staged_json", side_effect=refuse_metadata),
+        patch("litellm.proxy.client.cli.commands.auth.requests.Session", _FakeSession),
+    ):
+        result = CliRunner().invoke(
+            cli,
+            ["--base-url", PKCE_BASE_URL, "teams", "assign-key", "--team-id", "team-b"],
+            obj={"secret_vault": vault},
+        )
+
+    assert result.exit_code != 0, result.output
+    assert (isolated_home / ".litellm" / "token.json").read_bytes() == metadata_before
+    if restore_fails:
+        assert load_token(vault=vault)["key"] == "session-new"
+        assert "Could not restore the previous login" in result.output
+    else:
+        assert load_token(vault=vault) == before
+        assert "previous login was restored" in result.output
+    assert "Successfully assigned" not in result.output
+    assert [post for session in _FakeSession.instances for post in session.posts] == (
+        []
+        if restore_fails
+        else [
+            (
+                f"{PKCE_BASE_URL}/revoke",
+                {"token": "llm_srefresh_fresh", "token_type_hint": "refresh_token", "client_id": "llm_dcrc_abc"},
+            )
+        ]
+    )
+
+
+@pytest.mark.parametrize("reader", ["stored-key", "print-token"])
+def test_partial_replacement_preserves_a_concurrently_refreshed_login(isolated_home, secret_vault_factory, reader):
+    vault = secret_vault_factory()
+    save_cli_token(CliTokenRecord(**_pkce_record(team_id="team-a")), vault=vault)
+    refreshing = Event()
+    replacement_staged = Event()
+    rotation_saved = Event()
+    replace_file = os.replace
+    http = Mock()
+
+    def refresh_response(*_, **__):
+        refreshing.set()
+        replacement_staged.wait(1)
+        return _FakeHttpResponse(200, {**PKCE_TOKEN_RESPONSE, "team_id": "team-a"})
+
+    def commit_metadata(source, target):
+        if json.loads(Path(source).read_text())["team_id"] == "team-b":
+            replacement_staged.set()
+            assert rotation_saved.wait(5), "renewal did not persist its rotated credential"
+            raise OSError("cannot replace team metadata")
+        replace_file(source, target)
+        rotation_saved.set()
+
+    def read_key():
+        if reader == "stored-key":
+            return get_stored_api_key(PKCE_BASE_URL, vault=vault)
+        result = CliRunner().invoke(print_token, obj={"base_url": PKCE_BASE_URL, "secret_vault": vault})
+        assert result.exit_code == 0, result.output
+        return result.output.strip()
+
+    http.post.side_effect = refresh_response
+    with (
+        patch("litellm.proxy.client.cli.commands.auth.requests.Session", return_value=http),
+        patch("litellm.litellm_core_utils.private_json.os.replace", side_effect=commit_metadata),
+        ThreadPoolExecutor(max_workers=2) as executor,
+    ):
+        renewal = executor.submit(read_key)
+        assert refreshing.wait(5), "renewal did not reach the proxy"
+        replacement = executor.submit(
+            _replace_stored_token,
+            _pkce_record(key="session-new", refresh_token="llm_srefresh_new", team_id="team-b"),
+            _FakeSession(),
+            vault,
+            "team-b",
+        )
+        assert renewal.result(timeout=10) == "sk-cli-rotated"
+        assert isinstance(replacement.result(timeout=10), CredentialNotSaved)
+
+    saved = load_token(vault=vault)
+    assert (saved["key"], saved["refresh_token"], saved["team_id"]) == (
+        "sk-cli-rotated",
+        "llm_srefresh_rotated",
+        "team-a",
+    )
+
+
+@pytest.mark.parametrize("operation", ["logout", "renewal", "replacement"])
+def test_saved_login_works_with_read_only_home(isolated_home, secret_vault_factory, operation):
+    vault = secret_vault_factory()
+    save_cli_token(CliTokenRecord(**_pkce_record(team_id="team-a")), vault=vault)
+    http = _FakeSession()
+    isolated_home.chmod(0o500)
+    try:
+        with patch("litellm.proxy.client.cli.commands.auth.requests.Session") as session:
+            session.return_value.post.return_value = _FakeHttpResponse(200, PKCE_TOKEN_RESPONSE)
+            if operation == "logout":
+                result = CliRunner().invoke(logout, obj={"secret_vault": vault})
+                assert result.exit_code == 0, result.output
+                assert load_token(vault=vault) is None
+            elif operation == "renewal":
+                assert get_stored_api_key(PKCE_BASE_URL, vault=vault) == "sk-cli-rotated"
+                assert load_token(vault=vault)["refresh_token"] == "llm_srefresh_rotated"
+            else:
+                outcome = _replace_stored_token(_pkce_record(team_id="team-b"), http, vault, "team-b")
+                assert isinstance(outcome, SecretStored)
+                assert load_token(vault=vault)["team_id"] == "team-b"
+    finally:
+        isolated_home.chmod(0o700)
+
+
+def test_reading_a_fresh_login_does_not_need_a_writable_lock(isolated_home, secret_vault_factory, capsys):
+    vault = secret_vault_factory()
+    failing_boundary = "ctypes.WinDLL" if sys.platform == "win32" else "os.open"
+    with patch(failing_boundary, side_effect=PermissionError("lock unavailable")):
+        assert get_stored_api_key(PKCE_BASE_URL, vault=vault) is None
+    save_cli_token(CliTokenRecord(**_pkce_record(expires_at=time.time() + 3600)), vault=vault)
+
+    with patch(failing_boundary, side_effect=PermissionError("lock unavailable")):
+        assert get_stored_api_key(PKCE_BASE_URL, vault=vault) == "sk-cli-old"
+    assert capsys.readouterr().err == ""
+
+
+def test_lock_failure_leaves_saved_login_and_refresh_token_untouched(isolated_home, secret_vault_factory):
+    vault = secret_vault_factory()
+    save_cli_token(CliTokenRecord(**_pkce_record(team_id="team-a")), vault=vault)
+    before = load_token(vault=vault)
+    http = _FakeSession()
+
+    failing_boundary = "ctypes.WinDLL" if sys.platform == "win32" else "os.open"
+    with patch(failing_boundary, side_effect=PermissionError("lock unavailable")):
+        assert get_stored_api_key(PKCE_BASE_URL, vault=vault) is None
+        outcome = _replace_stored_token(_pkce_record(team_id="team-b"), http, vault, "team-b")
+        result = CliRunner().invoke(logout, obj={"secret_vault": vault})
+
+    assert isinstance(outcome, CredentialNotSaved)
+    assert "Could not lock the saved login" in outcome.detail
+    assert result.exit_code == 1, result.output
+    assert "Could not lock the saved login" in result.output
+    assert load_token(vault=vault) == before
+    assert http.posts == []
+
+
+@pytest.mark.parametrize("read_only", ["home", "config"])
+def test_credential_lock_serializes_processes_without_writable_directories(isolated_home, read_only):
+    config = isolated_home / ".litellm"
+    config.mkdir()
+    directory = isolated_home if read_only == "home" else config
+    directory.chmod(0o500)
+    child = """
+from pathlib import Path
+from filelock import Timeout
+from litellm.litellm_core_utils.cli_credential_lock import credential_lock
+try:
+    with credential_lock(Path.home(), timeout=0):
+        print('acquired')
+except Timeout:
+    print('blocked')
+"""
+    env = {**os.environ, "TMPDIR": str(isolated_home / "different-temporary-directory")}
+
+    def interrupted_writer():
+        with credential_lock(isolated_home):
+            blocked = run_child_interpreter(child, env=env, timeout=60)
+            assert blocked.returncode == 0, blocked.stderr
+            assert blocked.stdout.strip() == "blocked"
+            raise RuntimeError("release on error")
+
+    try:
+        with pytest.raises(RuntimeError, match="release on error"):
+            interrupted_writer()
+        released = run_child_interpreter(child, env=env, timeout=60)
+        assert released.returncode == 0, released.stderr
+        assert released.stdout.strip() == "acquired"
+    finally:
+        directory.chmod(0o700)
+
+
+@pytest.fixture
+def windows_mutex_api():
+    api = {
+        "CreateMutexW": Mock(return_value=123),
+        "WaitForSingleObject": Mock(return_value=0),
+        "ReleaseMutex": Mock(return_value=True),
+        "CloseHandle": Mock(return_value=True),
+    }
+    with (
+        patch("ctypes.WinDLL", create=True),
+        patch("ctypes.WINFUNCTYPE", return_value=lambda binding: api[binding[0]], create=True),
+        patch("ctypes.get_last_error", return_value=5, create=True),
+        patch("ctypes.WinError", side_effect=lambda code=0: OSError(code, "Win32 error"), create=True),
+    ):
+        yield api
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="POSIX filesystem locking")
+def test_saved_login_operations_do_not_lock_network_home_directories(isolated_home, secret_vault_factory):
+    vault = secret_vault_factory()
+    save_cli_token(CliTokenRecord(**_pkce_record(team_id="team-a")), vault=vault)
+    flock = fcntl.flock
+
+    def network_flock(fd, operation):
+        if stat.S_ISDIR(os.fstat(fd).st_mode):
+            raise OSError(9, "network filesystem requires a write-open regular file")
+        return flock(fd, operation)
+
+    with (
+        patch("fcntl.flock", side_effect=network_flock),
+        patch("litellm.proxy.client.cli.commands.auth.requests.Session") as session,
+    ):
+        session.return_value.post.return_value = _FakeHttpResponse(200, PKCE_TOKEN_RESPONSE)
+        assert get_stored_api_key(PKCE_BASE_URL, vault=vault) == "sk-cli-rotated"
+        outcome = _replace_stored_token(_pkce_record(team_id="team-b"), _FakeSession(), vault, "team-b")
+        assert isinstance(outcome, SecretStored)
+        assert load_token(vault=vault)["team_id"] == "team-b"
+        result = CliRunner().invoke(logout, obj={"secret_vault": vault})
+        assert result.exit_code == 0, result.output
+        assert load_token(vault=vault) is None
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="POSIX filesystem locking")
+@pytest.mark.parametrize("unsafe", ["directory-permissions", "directory-owner", "file-owner", "hardlink"])
+def test_credential_lock_refuses_unsafe_filesystem_state(isolated_home, unsafe):
+    fstat = os.fstat
+
+    def unsafe_stat(fd):
+        original = fstat(fd)
+        directory = stat.S_ISDIR(original.st_mode)
+        if unsafe == "directory-permissions" and directory:
+            return os.stat_result((original.st_mode | 0o020, *original[1:]))
+        if (unsafe == "directory-owner" and directory) or (unsafe == "file-owner" and not directory):
+            return os.stat_result((*original[:4], original.st_uid + 1, *original[5:]))
+        if unsafe == "hardlink" and not directory:
+            return os.stat_result((*original[:3], 2, *original[4:]))
+        return original
+
+    with patch("os.fstat", side_effect=unsafe_stat), pytest.raises(PermissionError, match="CLI lock"):
+        with credential_lock(isolated_home):
+            pytest.fail("unsafe lock entered the credential operation")
+
+
+@pytest.mark.parametrize("wait_result", [0, 0x80])
+@pytest.mark.parametrize("release_result", [True, False])
+def test_windows_mutex_preserves_body_errors_and_closes_handle(
+    tmp_path, windows_mutex_api, wait_result, release_result
+):
+    windows_mutex_api["WaitForSingleObject"].return_value = wait_result
+    windows_mutex_api["ReleaseMutex"].return_value = release_result
+
+    with pytest.raises(ValueError, match="body failed"), _windows_mutex(tmp_path, timeout=0):
+        raise ValueError("body failed")
+
+    windows_mutex_api["ReleaseMutex"].assert_called_once_with(123)
+    windows_mutex_api["CloseHandle"].assert_called_once_with(123)
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="POSIX filesystem locking")
+def test_credential_lock_refuses_a_file_removed_while_waiting(isolated_home):
+    fstat = os.fstat
+    link_counts = iter((1, 0))
+
+    def removed_stat(fd):
+        original = fstat(fd)
+        if stat.S_ISREG(original.st_mode):
+            return os.stat_result((*original[:3], next(link_counts), *original[4:]))
+        return original
+
+    with patch("os.fstat", side_effect=removed_stat), patch("fcntl.flock", wraps=fcntl.flock) as flock:
+        with pytest.raises(OSError, match="removed while waiting"):
+            with credential_lock(isolated_home):
+                pytest.fail("removed lock entered the credential operation")
+
+    with pytest.raises(OSError, match="Bad file descriptor"):
+        os.fstat(flock.call_args.args[0])
+
+
+@pytest.mark.parametrize(
+    ("handle", "wait_result", "error_type", "error_code"),
+    [(None, 0, OSError, 5), (123, 0xFFFFFFFF, OSError, 5), (123, 0x102, Timeout, None)],
+)
+def test_windows_mutex_acquisition_failure_preserves_error_and_closes_handle(
+    tmp_path, windows_mutex_api, handle, wait_result, error_type, error_code
+):
+    windows_mutex_api["CreateMutexW"].return_value = handle
+    windows_mutex_api["WaitForSingleObject"].return_value = wait_result
+
+    with pytest.raises(error_type) as raised, _windows_mutex(tmp_path, timeout=0):
+        pytest.fail("failed acquisition entered the critical section")
+
+    assert raised.value.errno == error_code
+    windows_mutex_api["ReleaseMutex"].assert_not_called()
+    assert windows_mutex_api["CloseHandle"].call_count == (0 if handle is None else 1)
+
+
+def test_windows_mutex_release_failure_preserves_error_and_closes_handle(tmp_path, windows_mutex_api):
+    windows_mutex_api["ReleaseMutex"].return_value = False
+
+    with pytest.raises(OSError, match="Win32 error") as raised, _windows_mutex(tmp_path, timeout=0):
+        pass
+
+    assert raised.value.errno == 5
+    windows_mutex_api["CloseHandle"].assert_called_once_with(123)
+
+
+def test_assign_key_warns_when_a_rejected_login_cannot_be_revoked(isolated_home, monkeypatch):
+    monkeypatch.setenv(DISABLE_KEYRING_ENV_VAR, "1")
+    save_cli_token(CliTokenRecord(**_pkce_record(key="session-old", expires_at=time.time() + 3600)))
+    before = load_token()
+    with (
+        patch(
+            "litellm.proxy.client.cli.commands.auth.run_pkce_login",
+            return_value=replace(_pkce_credential(), team_id="team-c"),
+        ),
+        patch(
+            "litellm.proxy.client.cli.commands.auth.revoke_stored_credential",
+            return_value=RevocationUnavailable("connection unavailable"),
+        ) as revoke,
+    ):
+        result = CliRunner().invoke(cli, ["--base-url", PKCE_BASE_URL, "teams", "assign-key", "--team-id", "team-b"])
+
+    assert result.exit_code != 0
+    assert load_token() == before
+    assert "Could not revoke the rejected login's refresh token" in result.output
+    assert "connection unavailable" in result.output
+    assert "your saved login has not changed" in result.output
+    assert revoke.call_args.args[0]["refresh_token"] == "llm_srefresh_fresh"
+
+
+@pytest.mark.parametrize("multiple_teams", [False, True])
+def test_assign_key_refuses_unavailable_team_in_sso_poll(isolated_home, monkeypatch, multiple_teams):
+    monkeypatch.setenv(DISABLE_KEYRING_ENV_VAR, "1")
+    save_cli_token(CliTokenRecord(base_url=PKCE_BASE_URL, key="session-old", timestamp=time.time()))
+    before = load_token()
+    with (
+        patch(
+            "litellm.proxy.client.cli.commands.auth._start_cli_sso_flow",
+            return_value={"login_id": "login-id", "poll_secret": "poll-secret", "user_code": "ABCD-EFGH"},
+        ),
+        patch("litellm.proxy.client.cli.commands.auth.webbrowser.open"),
+        patch("litellm.proxy.client.cli.commands.auth.requests.get") as get,
+    ):
+        get.return_value.status_code = 200
+        get.return_value.json.return_value = {
+            "status": "ready",
+            "requires_team_selection": multiple_teams,
+            "teams": ["team-c", "team-d"] if multiple_teams else ["team-c"],
+            "key": "session-wrong-team",
+            "team_id": "team-c",
+            "user_id": "u1",
+        }
+        result = CliRunner().invoke(cli, ["--base-url", PKCE_BASE_URL, "teams", "assign-key", "--team-id", "team-b"])
+
+    assert result.exit_code != 0
+    assert "requested team" in result.output
+    assert load_token() == before
+    assert "Successfully assigned" not in result.output
+    assert get.call_count == 1
+
+
+@pytest.mark.parametrize("from_env", [False, True])
+def test_assign_key_refuses_an_explicit_saved_session_override(isolated_home, monkeypatch, from_env):
+    monkeypatch.setenv(DISABLE_KEYRING_ENV_VAR, "1")
+    save_cli_token(CliTokenRecord(**_pkce_record(key="session-old", expires_at=time.time() + 3600)))
+    before = load_token()
+    with (
+        patch("litellm.proxy.client.cli.commands.teams.Client") as client,
+        patch("litellm.proxy.client.cli.commands.auth.run_pkce_login") as run_pkce,
+    ):
+        result = CliRunner().invoke(
+            cli,
+            ["--base-url", PKCE_BASE_URL]
+            + ([] if from_env else ["--api-key", "session-old"])
+            + ["teams", "assign-key", "--team-id", "team-b"],
+            env={"LITELLM_PROXY_API_KEY": "session-old"} if from_env else {},
+        )
+
+    assert result.exit_code != 0
+    assert "Unset --api-key and LITELLM_PROXY_API_KEY" in result.output
+    assert load_token() == before
+    run_pkce.assert_not_called()
+    client.return_value.keys.update.assert_not_called()

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

@@ -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: {