From a758d6b2d3693c14da9fcddca1c659bfab6e014f Mon Sep 17 00:00:00 2001 From: jesus Date: Wed, 9 Sep 2026 22:44:50 +0000 Subject: [PATCH 1/9] feat(cli): lite login --team option and env-var defaults for pkce and team Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp_server/discoverable_endpoints.py | 2 + .../mcp_server/gateway_dcr_flow.py | 7 ++ litellm/proxy/client/cli/README.md | 4 + litellm/proxy/client/cli/commands/auth.py | 55 ++++++++++-- .../proxy/client/cli/commands/pkce_login.py | 13 ++- .../html_forms/native_client_consent.py | 9 +- .../mcp_server/test_gateway_dcr_flow.py | 20 +++++ .../proxy/client/cli/test_auth_commands.py | 85 +++++++++++++++++++ .../proxy/client/cli/test_pkce_login.py | 10 +++ .../html_forms/test_native_client_consent.py | 12 +++ 10 files changed, 204 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index cab4b6c161a..4d56f9b7c10 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -1780,6 +1780,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): @@ -1794,6 +1795,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 3d94fa345d0..ee61d05d42c 100644 --- a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py +++ b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py @@ -509,6 +509,7 @@ async def native_client_authorize( response_type: str | None, session_user_id: str | None, lookup_consent_teams: LookupConsentTeams, + requested_team: str | None = None, ) -> Response: """The authorize verb for a native client that named the proxy API itself as its RFC 8707 ``resource``: the same client, redirect, PKCE, and sign-in checks as the @@ -530,6 +531,11 @@ async def native_client_authorize( teams: Final = await lookup_consent_teams(session_user_id) if not isinstance(teams, tuple): return _consent_lookup_failure_response(teams) + selected_team_id: Final = ( + next((t.team_id for t in teams if requested_team in (t.team_id, t.team_alias)), None) + if requested_team + else None + ) handle: Final = secrets.token_urlsafe(24) flow: Final = _new_connect_flow( session_user_id=session_user_id, @@ -544,6 +550,7 @@ async def native_client_authorize( client_origin=_origin_only(redirect_uri), user_id=session_user_id, teams=tuple((team.team_id, team.team_alias or team.team_id) for team in teams), + selected_team_id=selected_team_id, flow_handle=handle, complete_url=f"{base_url}/authorize/complete", ) diff --git a/litellm/proxy/client/cli/README.md b/litellm/proxy/client/cli/README.md index f7d9eb7da9a..5683a09cfe0 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, and `apiKeyHelper` to a `lite auth print-token` invocation, drops any stray static `ANTHROPIC_API_KEY` so the helper-issued token wins, and leaves every other setting in the file untouched. It backs up the original file before patching it. @@ -629,6 +631,8 @@ The CLI respects the following environment variables: - `LITELLM_PROXY_URL`: Base URL of the proxy server - `LITELLM_PROXY_API_KEY`: API key for authentication +- `LITELLM_PROXY_TEAM`: Default team (id or alias) for `lite login`, equivalent to `--team` +- `LITELLM_PROXY_LOGIN_PKCE`: Set to `true` to make `lite login` use the PKCE browser flow, equivalent to `--pkce` `LITELLM_PROXY_URL` takes precedence over a `base_url` stored via `lite config set`, and the `--base-url` option overrides both. See the Configuration section for the full precedence order. diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index 12a288202b6..26214f7ce5e 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -613,7 +613,24 @@ def _get_cli_sso_poll_headers(poll_secret: str) -> dict[str, str]: return {"x-litellm-cli-poll-secret": poll_secret} -def _poll_for_authentication(base_url: str, key_id: str, poll_secret: str) -> CliAuthResult | None: +def match_requested_team(teams: Sequence[CliTeam], requested_team: str | None) -> str | None: + """The team_id of the first team whose id or alias equals ``requested_team``.""" + if requested_team is None: + return None + return next( + ( + team_id + for team in teams + if (team_id := team.get("team_id")) is not None + and requested_team in (team_id, team.get("team_alias")) + ), + None, + ) + + +def _poll_for_authentication( + base_url: str, key_id: str, poll_secret: str, team: str | None = None +) -> CliAuthResult | None: """ Poll the server for authentication completion and handle team selection. @@ -643,6 +660,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 @@ -679,7 +697,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. @@ -697,7 +719,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.") @@ -816,9 +841,9 @@ def _replace_stored_token(record: CliTokenData, http: Http, vault: SecretVault) return stored -def _pkce_login(base_url: str, config_claude: bool, vault: SecretVault) -> None: +def _pkce_login(base_url: str, config_claude: bool, vault: SecretVault, team: str | None) -> None: http: Final = requests.Session() - credential: Final = run_pkce_login(base_url, http, echo=click.echo) + credential: Final = run_pkce_login(base_url, http, echo=click.echo, team=team) if isinstance(credential, PkceFailure): click.echo(f"Authentication failed: {credential.reason}") return @@ -840,14 +865,26 @@ def _pkce_login(base_url: str, config_claude: bool, vault: SecretVault) -> None: "--pkce", is_flag=True, default=False, + envvar="LITELLM_PROXY_LOGIN_PKCE", + show_envvar=True, help=( "Sign in with OAuth authorization code + PKCE through your system browser (loopback redirect), " "with a refresh token that renews the key automatically. Requires a proxy that serves " "/.well-known/litellm-cli-auth." ), ) +@click.option( + "--team", + envvar="LITELLM_PROXY_TEAM", + show_envvar=True, + default=None, + help=( + "Team id or alias to attribute this login to. Skips the team pick when it matches one of your " + "teams; otherwise you pick as usual." + ), +) @click.pass_context -def login(ctx: click.Context, config_claude: bool, pkce: bool) -> None: +def login(ctx: click.Context, config_claude: bool, pkce: bool, team: str | None) -> None: """Login to LiteLLM proxy using SSO authentication""" from litellm.constants import LITELLM_CLI_SOURCE_IDENTIFIER @@ -856,7 +893,7 @@ def login(ctx: click.Context, config_claude: bool, pkce: bool) -> None: try: if pkce: - _pkce_login(base_url, config_claude, context_secret_vault(ctx)) + _pkce_login(base_url, config_claude, context_secret_vault(ctx), team) return cli_sso_flow: Final = _start_cli_sso_flow(base_url=base_url) key_id: Final = cli_sso_flow["login_id"] @@ -887,7 +924,9 @@ def login(ctx: click.Context, config_claude: bool, pkce: bool) -> None: # Poll for authentication completion click.echo("Waiting for authentication...") - auth_result: Final = _poll_for_authentication(base_url=base_url, key_id=key_id, poll_secret=poll_secret) + auth_result: Final = _poll_for_authentication( + base_url=base_url, key_id=key_id, poll_secret=poll_secret, team=team + ) if auth_result: api_key: Final = auth_result["api_key"] diff --git a/litellm/proxy/client/cli/commands/pkce_login.py b/litellm/proxy/client/cli/commands/pkce_login.py index 93f5cfea21b..c53d45dc0c8 100644 --- a/litellm/proxy/client/cli/commands/pkce_login.py +++ b/litellm/proxy/client/cli/commands/pkce_login.py @@ -304,7 +304,14 @@ def pkce_pair() -> tuple[str, str]: return verifier, urlsafe_b64encode(digest).rstrip(b"=").decode("ascii") -def authorize_url(contract: CliAuthContract, client_id: str, redirect_uri: str, state: str, code_challenge: str) -> str: +def authorize_url( + contract: CliAuthContract, + client_id: str, + redirect_uri: str, + state: str, + code_challenge: str, + team: str | None = None, +) -> str: query: Final = urlencode( _form( response_type="code", @@ -314,6 +321,7 @@ def authorize_url(contract: CliAuthContract, client_id: str, redirect_uri: str, code_challenge=code_challenge, code_challenge_method="S256", resource=contract.resource, + **({"team": team} if team is not None else {}), ) ) return f"{contract.authorization_endpoint}?{query}" @@ -440,6 +448,7 @@ def run_pkce_login( open_browser: Callable[[str], object] = webbrowser.open, echo: Callable[[str], None] = print, timeout_seconds: float = LOGIN_TIMEOUT_SECONDS, + team: str | None = None, ) -> PkceCredential | PkceFailure: contract: Final = discover_cli_auth(base_url, http) if isinstance(contract, PkceFailure): @@ -450,7 +459,7 @@ def run_pkce_login( client_id: Final = register_client(contract, server.redirect_uri, http) if isinstance(client_id, PkceFailure): return client_id - url: Final = authorize_url(contract, client_id, server.redirect_uri, state, challenge) + url: Final = authorize_url(contract, client_id, server.redirect_uri, state, challenge, team=team) echo(f"Opening browser to: {url}") echo("Approve the sign-in in your browser. Waiting...") threading.Thread(target=open_browser, args=(url,), name="lite-login-browser", daemon=True).start() diff --git a/litellm/proxy/common_utils/html_forms/native_client_consent.py b/litellm/proxy/common_utils/html_forms/native_client_consent.py index dac92c4e787..bee56970f1b 100644 --- a/litellm/proxy/common_utils/html_forms/native_client_consent.py +++ b/litellm/proxy/common_utils/html_forms/native_client_consent.py @@ -12,6 +12,7 @@ def render_native_client_consent_page( teams: Sequence[tuple[str, str]], flow_handle: str, complete_url: str, + selected_team_id: str | None = None, ) -> str: """The consent page a native client's sign-in lands on: who is signed in, which loopback client asked, which team the credential is attributed to, and an explicit @@ -62,7 +63,7 @@ button {{ flex: 1; padding: 10px; border-radius: 6px; font-size: 15px; cursor: p

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

-{_team_field(teams)} +{_team_field(teams, selected_team_id)}
@@ -74,7 +75,7 @@ button {{ flex: 1; padding: 10px; border-radius: 6px; font-size: 15px; cursor: p """ -def _team_field(teams: Sequence[tuple[str, str]]) -> str: +def _team_field(teams: Sequence[tuple[str, str]], selected_team_id: str | None = None) -> str: if not teams: return "" if len(teams) == 1: @@ -84,7 +85,9 @@ def _team_field(teams: Sequence[tuple[str, str]]) -> str: f"

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

" ) options: Final = "".join( - f'' for team_id, team_label in teams + f'" + for team_id, team_label in teams ) return ( f'' diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py index 73a52a8d2e8..9c9b14d9aa7 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 @@ -2044,3 +2044,23 @@ async def test_introspect_fails_closed_on_dead_user_and_503s_on_outage(): status, body = await _introspect(minted.token.get_secret_value(), master_key=None) 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 From e32b3f2d76e6eaeaa83a3a86fd2a26c5ee99a9fb Mon Sep 17 00:00:00 2001 From: jesus Date: Wed, 9 Sep 2026 22:56:59 +0000 Subject: [PATCH 2/9] fix(cli): format auth.py, update polling call assertions, regenerate openapi snapshot and schema.d.ts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_lazy_openapi_snapshot.json | 32 +++++++++++++++++++ litellm/proxy/client/cli/commands/auth.py | 3 +- .../test_litellm/proxy/auth/test_cli_auth.py | 2 ++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 ++ 4 files changed, 37 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 4d5cdcc8003..b6922719924 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -19946,6 +19946,22 @@ ], "title": "Resource" } + }, + { + "in": "query", + "name": "team", + "required": false, + "schema": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Team" + } } ], "responses": { @@ -20636,6 +20652,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/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index 26214f7ce5e..4d2df54dbc6 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -621,8 +621,7 @@ def match_requested_team(teams: Sequence[CliTeam], requested_team: str | None) - ( 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")) + if (team_id := team.get("team_id")) is not None and requested_team in (team_id, team.get("team_alias")) ), None, ) diff --git a/tests/test_litellm/proxy/auth/test_cli_auth.py b/tests/test_litellm/proxy/auth/test_cli_auth.py index 5cde5522376..ade4c48f554 100644 --- a/tests/test_litellm/proxy/auth/test_cli_auth.py +++ b/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] diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index adbdb2f46d9..bd18e399fcf 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -41572,6 +41572,7 @@ export interface operations { response_type?: string | null; scope?: string | null; resource?: string | null; + team?: string | null; }; header?: never; path?: never; @@ -68967,6 +68968,7 @@ export interface operations { response_type?: string | null; scope?: string | null; resource?: string | null; + team?: string | null; }; header?: never; path: { From eef5541e60f993c80ffbe1dfad6e96761eda2e06 Mon Sep 17 00:00:00 2001 From: yassin Date: Tue, 29 Sep 2026 16:57:33 +0000 Subject: [PATCH 3/9] fix(cli): renew saved sessions when assigning teams Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/client/cli/commands/_cli_context.py | 3 +- litellm/proxy/client/cli/commands/auth.py | 57 +++-- litellm/proxy/client/cli/commands/teams.py | 23 ++ .../proxy/client/cli/test_auth_commands.py | 222 ++++++++++++++++++ 4 files changed, 290 insertions(+), 15 deletions(-) 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 5c2ce9eb53c..2520c09a535 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -634,7 +634,7 @@ def match_requested_team(teams: Sequence[CliTeam], requested_team: str | None) - def _poll_for_authentication( - base_url: str, key_id: str, poll_secret: str, team: str | None = None + 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. @@ -655,6 +655,11 @@ def _poll_for_authentication( 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 @@ -674,7 +679,7 @@ def _poll_for_authentication( "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.") @@ -852,7 +857,17 @@ 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") previous: Final = load_token(vault=vault) stored: Final = save_token(record, vault=vault) if previous is None or isinstance(stored, CredentialNotSaved): @@ -866,14 +881,17 @@ def _replace_stored_token(record: CliTokenData, http: Http, vault: SecretVault) return stored -def _pkce_login(base_url: str, config_claude: bool, vault: SecretVault, team: str | None) -> 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, 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") @@ -911,6 +929,12 @@ def _pkce_login(base_url: str, config_claude: bool, vault: SecretVault, team: st @click.pass_context 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 @@ -924,8 +948,7 @@ def login(ctx: click.Context, config_claude: bool, pkce: bool, team: str | None) try: if pkce: - _pkce_login(base_url, config_claude, context_secret_vault(ctx), team) - 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"] @@ -955,8 +978,12 @@ def login(ctx: click.Context, config_claude: bool, pkce: bool, team: str | 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, team=team + 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: @@ -975,31 +1002,33 @@ def login(ctx: click.Context, config_claude: bool, pkce: bool, team: str | 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") 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/tests/test_litellm/proxy/client/cli/test_auth_commands.py b/tests/test_litellm/proxy/client/cli/test_auth_commands.py index 543df16e23e..3e922e6c8e7 100644 --- a/tests/test_litellm/proxy/client/cli/test_auth_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_auth_commands.py @@ -2,6 +2,7 @@ import json import os import stat import time +from dataclasses import replace from pathlib import Path from unittest.mock import Mock, patch @@ -28,11 +29,13 @@ from litellm.proxy.client.cli.commands import claude_settings as claude_settings from litellm.proxy.client.cli.commands.claude_settings import SettingsFileOwner from litellm.proxy.client.cli.commands.auth import ( get_stored_api_key, + load_token, login, logout, print_token, whoami, ) +from litellm.proxy.client.cli.commands.pkce_login import PkceFailure, RevocationUnavailable @pytest.fixture @@ -2129,3 +2132,222 @@ class TestRequestedTeamLoginOption: 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 + + +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() From 08dbab791425eca74396420b04afd870dbc4479b Mon Sep 17 00:00:00 2001 From: yassin Date: Tue, 29 Sep 2026 17:20:06 +0000 Subject: [PATCH 4/9] fix(cli): recover partial credential replacement Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/client/cli/commands/auth.py | 17 +++++- .../proxy/client/cli/test_auth_commands.py | 54 +++++++++++++++++++ 2 files changed, 70 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index 2520c09a535..b0bbd28f372 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -869,8 +869,23 @@ def _replace_stored_token( ) raise click.ClickException("The login did not select the requested team; your saved login has not changed") 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: diff --git a/tests/test_litellm/proxy/client/cli/test_auth_commands.py b/tests/test_litellm/proxy/client/cli/test_auth_commands.py index 3e922e6c8e7..a2fe8e25603 100644 --- a/tests/test_litellm/proxy/client/cli/test_auth_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_auth_commands.py @@ -2274,6 +2274,60 @@ def test_assign_key_keeps_explicit_virtual_key_update_with_saved_session(isolate 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"}, + ) + ] + ) + + 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))) From 501baea61cbbef6c9c902ab535516412e0ee5d8d Mon Sep 17 00:00:00 2001 From: yassin Date: Tue, 29 Sep 2026 17:50:41 +0000 Subject: [PATCH 5/9] fix(cli): serialize session replacement with token renewal Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/client/cli/commands/auth.py | 52 +++++++++-- .../proxy/client/cli/test_auth_commands.py | 88 +++++++++++++++++++ 2 files changed, 134 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index b0bbd28f372..13266e2f72d 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -3,11 +3,13 @@ import sys import time import webbrowser from collections.abc import Callable, Mapping, Sequence +from pathlib import Path from typing import Any, Final, TypeVar from urllib.parse import urlencode import click import requests +from filelock import BaseFileLock, FileLock, Timeout from rich.console import Console from rich.table import Table from typing_extensions import NotRequired, ReadOnly, TypedDict, assert_never @@ -239,6 +241,10 @@ def _renewal_reader(vault: SecretVault) -> Callable[[], Mapping[str, object] | N return reload +def _credential_lock() -> BaseFileLock: + return FileLock(str(Path.home() / ".litellm-token.lock"), timeout=30, mode=0o600) + + def get_stored_api_key( expected_base_url: str | None = None, *, @@ -256,6 +262,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, 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), @@ -868,6 +894,14 @@ def _replace_stored_token( "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) @@ -1051,6 +1085,14 @@ def login_to_proxy( 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: @@ -1124,15 +1166,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/tests/test_litellm/proxy/client/cli/test_auth_commands.py b/tests/test_litellm/proxy/client/cli/test_auth_commands.py index a2fe8e25603..3a23c49689e 100644 --- a/tests/test_litellm/proxy/client/cli/test_auth_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_auth_commands.py @@ -2,8 +2,10 @@ import json import os import stat 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 @@ -28,6 +30,7 @@ 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, @@ -2328,6 +2331,91 @@ def test_assign_key_handles_keychain_update_without_metadata(isolated_home, secr ) +@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", + ) + + +def test_reading_a_fresh_login_does_not_need_a_writable_lock(isolated_home, secret_vault_factory, capsys): + vault = secret_vault_factory() + (isolated_home / ".litellm-token.lock").mkdir() + 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) + + 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) + (isolated_home / ".litellm-token.lock").mkdir() + http = _FakeSession() + + 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 == [] + + 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))) From 347807a3c1b30b2d5a821ddfa263adb91077860a Mon Sep 17 00:00:00 2001 From: yassin Date: Tue, 29 Sep 2026 17:54:39 +0000 Subject: [PATCH 6/9] fix(cli): keep cached credential reads from initiating renewal Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/client/cli/commands/auth.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index 13266e2f72d..95af9088c38 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -263,7 +263,7 @@ def get_stored_api_key( 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, vault) + return _key_from_record({**token_data, "refresh_token": None}, vault) try: with _credential_lock(): return _get_stored_api_key(expected_base_url, vault) From cf4088a62766f7772d0b665bca021d7e182b32b0 Mon Sep 17 00:00:00 2001 From: yassin Date: Tue, 29 Sep 2026 19:15:21 +0000 Subject: [PATCH 7/9] fix(cli): lock saved credentials without writable directories Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_core_utils/cli_credential_lock.py | 74 +++++++++ litellm/proxy/client/cli/commands/auth.py | 8 +- .../proxy/client/cli/test_auth_commands.py | 140 +++++++++++++++++- 3 files changed, 212 insertions(+), 10 deletions(-) create mode 100644 litellm/litellm_core_utils/cli_credential_lock.py 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..f38be474868 --- /dev/null +++ b/litellm/litellm_core_utils/cli_credential_lock.py @@ -0,0 +1,74 @@ +import ctypes +import errno +import os +import sys +import time +from collections.abc import Generator +from contextlib import contextmanager +from ctypes import wintypes +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 without creating or modifying a lock file.""" + if sys.platform == "win32": + with _windows_mutex(home, timeout): + yield + return + fd: Final = os.open(home, os.O_RDONLY | os.O_DIRECTORY) + 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: + yield + finally: + fcntl.flock(fd, fcntl.LOCK_UN) + finally: + os.close(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( + wintypes.HANDLE, ctypes.c_void_p, wintypes.BOOL, wintypes.LPCWSTR, use_last_error=True + )(("CreateMutexW", kernel)) + wait: Final = ctypes.WINFUNCTYPE(wintypes.DWORD, wintypes.HANDLE, wintypes.DWORD, use_last_error=True)( + ("WaitForSingleObject", kernel) + ) + release: Final = ctypes.WINFUNCTYPE(wintypes.BOOL, wintypes.HANDLE, use_last_error=True)(("ReleaseMutex", kernel)) + close: Final = ctypes.WINFUNCTYPE(wintypes.BOOL, 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/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index 95af9088c38..0bba43cd75a 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -3,18 +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 BaseFileLock, FileLock, Timeout +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, @@ -241,8 +243,8 @@ def _renewal_reader(vault: SecretVault) -> Callable[[], Mapping[str, object] | N return reload -def _credential_lock() -> BaseFileLock: - return FileLock(str(Path.home() / ".litellm-token.lock"), timeout=30, mode=0o600) +def _credential_lock() -> AbstractContextManager[None]: + return credential_lock(Path.home()) def get_stored_api_key( diff --git a/tests/test_litellm/proxy/client/cli/test_auth_commands.py b/tests/test_litellm/proxy/client/cli/test_auth_commands.py index 3a23c49689e..1c11c5fc673 100644 --- a/tests/test_litellm/proxy/client/cli/test_auth_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_auth_commands.py @@ -1,6 +1,7 @@ import json import os import stat +import sys import time from concurrent.futures import ThreadPoolExecutor from dataclasses import replace @@ -11,8 +12,10 @@ 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, @@ -39,6 +42,7 @@ from litellm.proxy.client.cli.commands.auth import ( whoami, ) from litellm.proxy.client.cli.commands.pkce_login import PkceFailure, RevocationUnavailable +from tests.test_litellm_rust.support.child_interpreter import run_child_interpreter @pytest.fixture @@ -2387,13 +2391,39 @@ def test_partial_replacement_preserves_a_concurrently_refreshed_login(isolated_h ) +@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() - (isolated_home / ".litellm-token.lock").mkdir() - assert get_stored_api_key(PKCE_BASE_URL, vault=vault) is None + 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) - assert get_stored_api_key(PKCE_BASE_URL, vault=vault) == "sk-cli-old" + 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 == "" @@ -2401,12 +2431,13 @@ def test_lock_failure_leaves_saved_login_and_refresh_token_untouched(isolated_ho vault = secret_vault_factory() save_cli_token(CliTokenRecord(**_pkce_record(team_id="team-a")), vault=vault) before = load_token(vault=vault) - (isolated_home / ".litellm-token.lock").mkdir() http = _FakeSession() - 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}) + 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 @@ -2416,6 +2447,101 @@ def test_lock_failure_leaves_saved_login_and_refresh_token_untouched(isolated_ho 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.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.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))) From 0d14946003b6da4a34a49c0fd48a2cc02f0110bd Mon Sep 17 00:00:00 2001 From: yassin Date: Tue, 29 Sep 2026 19:44:18 +0000 Subject: [PATCH 8/9] fix(cli): keep credential locks off network homes Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_core_utils/cli_credential_lock.py | 51 +++++++++++++++---- .../proxy/client/cli/test_auth_commands.py | 49 ++++++++++++++++++ 2 files changed, 91 insertions(+), 9 deletions(-) diff --git a/litellm/litellm_core_utils/cli_credential_lock.py b/litellm/litellm_core_utils/cli_credential_lock.py index f38be474868..d310936a2ee 100644 --- a/litellm/litellm_core_utils/cli_credential_lock.py +++ b/litellm/litellm_core_utils/cli_credential_lock.py @@ -1,11 +1,11 @@ -import ctypes +import ctypes.wintypes import errno import os +import stat import sys import time from collections.abc import Generator from contextlib import contextmanager -from ctypes import wintypes from hashlib import sha256 from pathlib import Path from typing import TYPE_CHECKING, Final, cast @@ -18,12 +18,12 @@ if TYPE_CHECKING or sys.platform != "win32": @contextmanager def credential_lock(home: Path, timeout: float = 30) -> Generator[None, None, None]: - """Serialize credential changes without creating or modifying a lock file.""" + """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 = os.open(home, os.O_RDONLY | os.O_DIRECTORY) + fd: Final = _open_posix_lock(home) try: deadline: Final = time.monotonic() + timeout while True: @@ -37,6 +37,9 @@ def credential_lock(home: Path, timeout: float = 30) -> Generator[None, None, No 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) @@ -44,17 +47,47 @@ def credential_lock(home: Path, timeout: float = 30) -> Generator[None, None, No 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( - wintypes.HANDLE, ctypes.c_void_p, wintypes.BOOL, wintypes.LPCWSTR, use_last_error=True + ctypes.wintypes.HANDLE, ctypes.c_void_p, ctypes.wintypes.BOOL, ctypes.wintypes.LPCWSTR, use_last_error=True )(("CreateMutexW", kernel)) - wait: Final = ctypes.WINFUNCTYPE(wintypes.DWORD, wintypes.HANDLE, wintypes.DWORD, use_last_error=True)( - ("WaitForSingleObject", 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) ) - release: Final = ctypes.WINFUNCTYPE(wintypes.BOOL, wintypes.HANDLE, use_last_error=True)(("ReleaseMutex", kernel)) - close: Final = ctypes.WINFUNCTYPE(wintypes.BOOL, 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: diff --git a/tests/test_litellm/proxy/client/cli/test_auth_commands.py b/tests/test_litellm/proxy/client/cli/test_auth_commands.py index 1c11c5fc673..01a21844877 100644 --- a/tests/test_litellm/proxy/client/cli/test_auth_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_auth_commands.py @@ -44,6 +44,9 @@ from litellm.proxy.client.cli.commands.auth import ( 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 def isolated_home(monkeypatch, tmp_path): @@ -2499,6 +2502,52 @@ def windows_mutex_api(): 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( From f7e47bd108cba03e0e5601dfa51edfbbf2ccd293 Mon Sep 17 00:00:00 2001 From: yassin Date: Tue, 29 Sep 2026 19:46:44 +0000 Subject: [PATCH 9/9] test(cli): cover a removed credential lock Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/client/cli/test_auth_commands.py | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/tests/test_litellm/proxy/client/cli/test_auth_commands.py b/tests/test_litellm/proxy/client/cli/test_auth_commands.py index 01a21844877..a80063623c4 100644 --- a/tests/test_litellm/proxy/client/cli/test_auth_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_auth_commands.py @@ -2563,6 +2563,26 @@ def test_windows_mutex_preserves_body_errors_and_closes_handle( 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)],