diff --git a/litellm/litellm_core_utils/cli_token_utils.py b/litellm/litellm_core_utils/cli_token_utils.py index 37ed4efa0ca..2a95bd3c836 100644 --- a/litellm/litellm_core_utils/cli_token_utils.py +++ b/litellm/litellm_core_utils/cli_token_utils.py @@ -20,6 +20,7 @@ from pathlib import Path from types import MappingProxyType from typing import Final, TypeAlias +from filelock import FileLock, Timeout from pydantic import ConfigDict, ValidationError from litellm.litellm_core_utils.cli_keyring import ( @@ -47,6 +48,11 @@ from litellm.litellm_core_utils.private_json import ( from litellm.types.llms.base_model import LiteLLMBaseModel +@dataclass(frozen=True, slots=True) +class CredentialNotLoaded: + detail: str + + @dataclass(frozen=True, slots=True) class CredentialNotSaved: """The credential was minted but no store would keep it, so this machine has none. @@ -139,8 +145,39 @@ def get_cli_token_file_path() -> str: return str(config_dir / "token.json") +def cli_token_lock() -> FileLock: + """Lock metadata and keychain operations beside the directory so read-only stores still allow logout.""" + directory: Final = Path(get_cli_token_file_path()).parent + directory.parent.mkdir(parents=True, exist_ok=True) + path: Final = directory.with_name(f"{directory.name}-token.lock") + return FileLock(str(path), timeout=60, mode=0o600, is_singleton=True) + + def load_cli_token(*, vault: SecretVault = SYSTEM_KEYRING) -> CliTokenRecord | None: - """Load the stored CLI credential, or None when this machine has none""" + """Load a stored CLI credential, or None when it cannot be loaded.""" + outcome: Final = read_cli_token(vault=vault) + return outcome if isinstance(outcome, CliTokenRecord) else None + + +def read_cli_token(*, vault: SecretVault = SYSTEM_KEYRING) -> CliTokenRecord | CredentialNotLoaded | None: + if not Path(get_cli_token_file_path()).exists(): + return None + try: + with cli_token_lock(): + return _load_cli_token(vault) + except Timeout: + return CredentialNotLoaded( + "CLI credential store is busy. Retry this command after the other credential operation finishes. " + "This command did not change your stored login." + ) + except OSError: + return CredentialNotLoaded( + "Could not lock the CLI credential store. Retry this command after checking the lock file permissions. " + "This command did not change your stored login." + ) + + +def _load_cli_token(vault: SecretVault) -> CliTokenRecord | None: record: Final = _read_token_file() if record is None: return None @@ -148,6 +185,14 @@ def load_cli_token(*, vault: SecretVault = SYSTEM_KEYRING) -> CliTokenRecord | N def save_cli_token(record: CliTokenRecord, *, vault: SecretVault = SYSTEM_KEYRING) -> SecretSave: + try: + with cli_token_lock(): + return _save_cli_token(record, vault) + except (OSError, Timeout) as error: + return CredentialNotSaved(str(error)) + + +def _save_cli_token(record: CliTokenRecord, vault: SecretVault) -> SecretSave: """Store a freshly minted credential. Reports where its secret material ended up, and why. The token file is what makes a keychain-backed credential findable again, and it is also the @@ -218,6 +263,14 @@ def _keep_the_secret_in_the_file(record: CliTokenRecord, outcome: SecretWrite) - def clear_cli_token(*, vault: SecretVault = SYSTEM_KEYRING) -> SecretClear: + try: + with cli_token_lock(): + return _clear_cli_token(vault) + except (OSError, Timeout) as error: + return CredentialNotCleared(str(error)) + + +def _clear_cli_token(vault: SecretVault) -> SecretClear: """Remove the credential from both stores. Reports whether the keychain is now free of it. A logout the keychain never answered keeps the token file, with its secret taken out, because @@ -339,13 +392,17 @@ def get_litellm_gateway_api_key( >>> base_url="https://your-proxy.com/v1" >>> ) """ - record: Final = _read_token_file() - if record is None: + try: + with cli_token_lock(): + record: Final = _read_token_file() + if record is None: + return None + if expected_base_url is not None and record.base_url != expected_base_url.rstrip("/"): + return None + resolved: Final = _resolve_secret(record, vault) + return None if resolved is None else resolved.key + except (OSError, Timeout): return None - if expected_base_url is not None and record.base_url != expected_base_url.rstrip("/"): - return None - resolved: Final = _resolve_secret(record, vault) - return None if resolved is None else resolved.key def is_cli_token_fresh( diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index e38f54adc4f..38a1e2d8534 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -2,12 +2,14 @@ import os import sys import time import webbrowser -from collections.abc import Callable, Mapping, Sequence -from typing import Any, Final, Protocol, TypeVar +from collections.abc import Callable, Generator, Mapping, Sequence +from contextlib import contextmanager +from typing import Any, Final, NoReturn, Protocol, TypeVar from urllib.parse import urlencode import click import requests +from filelock import Timeout from pydantic import ConfigDict, TypeAdapter from rich.console import Console from rich.table import Table @@ -30,13 +32,15 @@ from litellm.litellm_core_utils.cli_keyring import ( from litellm.litellm_core_utils.cli_token_utils import ( CliTokenRecord, CredentialNotCleared, + CredentialNotLoaded, CredentialNotRecorded, CredentialNotSaved, SecretSave, clear_cli_token, + cli_token_lock, get_cli_token_file_path, is_cli_token_fresh, - load_cli_token, + read_cli_token, save_cli_token, ) @@ -225,13 +229,14 @@ def context_secret_vault(ctx: click.Context) -> SecretVault: def load_token(*, vault: SecretVault = SYSTEM_KEYRING) -> Mapping[str, object] | None: - """The stored credential as a plain mapping, with the secret resolved out of the vault. - - The PKCE renewal and revocation helpers read records by field name, so this is the - shape they get; the keychain split lives underneath, in `load_cli_token`. - """ - record: Final = load_cli_token(vault=vault) - return None if record is None else record.model_dump(exclude_none=True) + """Read a stored credential, reporting lock failures separately from a missing login.""" + match read_cli_token(vault=vault): + case CliTokenRecord() as record: + return record.model_dump(exclude_none=True) + case CredentialNotLoaded(detail=detail): + _raise_credential_error(detail) + case None: + return None def save_token(record: CliTokenData, *, vault: SecretVault = SYSTEM_KEYRING) -> SecretSave: @@ -286,6 +291,7 @@ def get_stored_api_key( requests.Session(), reload=_renewal_reader(vault), warn=_warn, + lock_failure=_raise_credential_error, ) @@ -293,6 +299,10 @@ def _warn(message: str) -> None: click.echo(message, err=True) +def _raise_credential_error(message: str) -> NoReturn: + raise click.ClickException(message) + + def _login_command(renews: bool) -> str: return "lite login --pkce" if renews else "lite login" @@ -852,6 +862,22 @@ def _finish_login(base_url: str, api_key: str, config_claude: bool, stored: Secr show_commands() +@contextmanager +def _locked_credentials() -> Generator[None, None, None]: + try: + lock: Final = cli_token_lock() + lock.acquire() + except (OSError, Timeout) as error: + raise click.ClickException( + "Could not lock the CLI credential store. No credential operation was started" + ) from error + try: + yield + finally: + lock.release() + + +@_locked_credentials() def _replace_stored_token(record: CliTokenData, http: Http, vault: SecretVault) -> SecretSave: previous: Final = load_token(vault=vault) stored: Final = save_token(record, vault=vault) @@ -989,6 +1015,7 @@ def login(ctx: click.Context, config_claude: bool, pkce: bool | None) -> None: @click.command(name="logout") @click.pass_context +@_locked_credentials() def logout(ctx: click.Context): """Logout and clear stored authentication""" vault: Final = context_secret_vault(ctx) @@ -1067,8 +1094,10 @@ def print_token(ctx: click.Context): click.echo(keychain_unreadable_notice(vault), err=True) sys.exit(1) + stored_key: Final = token_data.get("key") + fresh_stored_key: Final = stored_key if isinstance(stored_key, str) and is_cli_token_fresh(token_data) else None api_key: Final = ( - ctx_obj.get("api_key") + ctx_obj.get("api_key") or fresh_stored_key if issued_for_this_server and ctx_obj.get("api_key_from_token_file") else fresh_api_key( token_data, @@ -1076,6 +1105,7 @@ def print_token(ctx: click.Context): requests.Session(), reload=_renewal_reader(vault), warn=_warn, + lock_failure=_raise_credential_error, ) ) if not api_key: diff --git a/litellm/proxy/client/cli/commands/pkce_login.py b/litellm/proxy/client/cli/commands/pkce_login.py index c9dd542eb36..4bc5a2dd377 100644 --- a/litellm/proxy/client/cli/commands/pkce_login.py +++ b/litellm/proxy/client/cli/commands/pkce_login.py @@ -20,10 +20,11 @@ from typing import TYPE_CHECKING, Final, Literal, Protocol from urllib.parse import parse_qs, urlencode, urlparse import requests +from filelock import Timeout from pydantic import ConfigDict, Field, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict -from litellm.litellm_core_utils.cli_token_utils import CLI_TOKEN_FRESHNESS_BUFFER_SECONDS +from litellm.litellm_core_utils.cli_token_utils import CLI_TOKEN_FRESHNESS_BUFFER_SECONDS, cli_token_lock from litellm.types.llms.base_model import LiteLLMBaseModel if TYPE_CHECKING: @@ -498,17 +499,53 @@ def fresh_api_key( reload: Callable[[], Mapping[str, object] | None], now: Callable[[], float] = time.time, warn: Callable[[str], None] = _ignore_warning, + lock_failure: Callable[[str], None] = _ignore_warning, +) -> str | None: + """Refresh under the shared credential lock, reloading before spending a single-use token. + + A replaced or cleared login must not be renewed from the caller's stale snapshot. + The lock covers persistence, but cannot recover a lost token response or interrupted save. + """ + expires_at: Final = token_data.get("expires_at") + if ( + not isinstance(expires_at, (int, float)) + or now() < expires_at - CLI_TOKEN_FRESHNESS_BUFFER_SECONDS + or _refresh_inputs(token_data) is None + ): + return _fresh_api_key(token_data, save, http, reload=reload, now=now, warn=warn) + key: Final = token_data.get("key") + try: + lock: Final = cli_token_lock() + lock.acquire() + except (OSError, Timeout): + fallback: Final = key if isinstance(key, str) and now() < expires_at else None + if fallback is None: + lock_failure( + "Could not lock the CLI credential store for renewal. Retry this command. No refresh was attempted." + ) + warn("Could not renew the key: the CLI credential store could not be locked; no refresh was attempted") + return fallback + try: + current: Final = reload() + still_valid: Final = key if isinstance(key, str) and now() < expires_at else None + if current is None or any(current.get(field) != token_data.get(field) for field in _CREDENTIAL_IDENTITY_FIELDS): + return still_valid + if not isinstance(current.get("expires_at"), (int, float)): + return still_valid + return _fresh_api_key(current, save, http, reload=reload, now=now, warn=warn) or still_valid + finally: + lock.release() + + +def _fresh_api_key( + token_data: Mapping[str, object], + save: Callable[[CliTokenData], None], + http: Http, + *, + reload: Callable[[], Mapping[str, object] | None], + now: Callable[[], float], + warn: Callable[[str], None], ) -> str | None: - """The stored key, refreshed first when it is about to expire and a refresh token is - on file. The refresh fires at the same moment ``is_cli_token_fresh`` stops calling the - key fresh, so a command that checks freshness and then asks for the key never disagrees - with itself. The rotated pair is saved before the new key is returned, so a crash after - this point never strands the CLI with a burned refresh token. A refresh that fails - reads the record again, because a sibling ``lite`` process may have rotated the pair - first, in which case the key it saved for this same proxy is the live one; when no sibling - did, the reason the proxy gave goes to ``warn`` so a revoked or refused refresh token is - never a silent failure. A record without ``expires_at`` (the classic ``lite login`` - credential) is returned as stored.""" key: Final = token_data.get("key") if not isinstance(key, str) or not key: return None diff --git a/tests/integration/sdk/test_cli_refresh_lock_sdk.py b/tests/integration/sdk/test_cli_refresh_lock_sdk.py new file mode 100644 index 00000000000..d97831901fb --- /dev/null +++ b/tests/integration/sdk/test_cli_refresh_lock_sdk.py @@ -0,0 +1,308 @@ +from __future__ import annotations + +import json +import os +import subprocess +import sys +import textwrap +import threading +from collections.abc import Generator, Iterator +from contextlib import contextmanager +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from typing import Final +from urllib.parse import parse_qs + +import pytest + +from tests.integration._support.process import kill_stray_process, stop_child + +EXPIRED_KEY: Final = "fixture-expired-key" +FRESH_KEY: Final = "fixture-fresh-key" +REFRESH_TOKEN: Final = "fixture-single-use-refresh" +ROTATED_REFRESH_TOKEN: Final = "fixture-rotated-refresh" + +_REFRESH: Final = textwrap.dedent( + """ + import sys, requests + from litellm.proxy.client.cli.commands.auth import load_token, save_token + from litellm.proxy.client.cli.commands.pkce_login import fresh_api_key + record = load_token() + assert record is not None + print("ready", flush=True) + sys.stdin.readline() + key = fresh_api_key(record, save_token, requests.Session(), reload=load_token) + if key is None: + raise SystemExit(1) + print(key, flush=True) + """ +) +_HOLD_LOCK: Final = textwrap.dedent( + """ + import sys + from litellm.litellm_core_utils.cli_token_utils import cli_token_lock + with cli_token_lock(): + print("ready", flush=True) + sys.stdin.readline() + """ +) +_CLI: Final = "from litellm_cli.main import cli; cli()" +_PRINT_TOKEN: Final = textwrap.dedent( + """ + import sys + from litellm.proxy.client.cli.main import cli + print("ready", flush=True) + sys.stdin.readline() + cli(["--base-url", sys.argv[1], "auth", "print-token"]) + """ +) +_SHORT_WAIT: Final = textwrap.dedent( + """ + from unittest.mock import patch + from filelock import FileLock + def bounded_lock(path, **options): + return FileLock(path, **{**options, "timeout": 0.1}) + with patch("litellm.litellm_core_utils.cli_token_utils.FileLock", side_effect=bounded_lock): + """ +) + + +class _TokenEndpoint(ThreadingHTTPServer): + daemon_threads = True + + def __init__(self) -> None: + super().__init__(("127.0.0.1", 0), _TokenHandler) + self.requests: list[dict[str, list[str]]] = [] + self.guard = threading.Lock() + self.response_delay = 0.0 + self.request_started = threading.Event() + + @property + def origin(self) -> str: + return f"http://127.0.0.1:{self.server_address[1]}" + + def redeem(self, form: dict[str, list[str]]) -> bool: + with self.guard: + self.requests.append(form) + return len(self.requests) == 1 and form.get("refresh_token") == [REFRESH_TOKEN] + + +class _TokenHandler(BaseHTTPRequestHandler): + server: _TokenEndpoint + + def do_POST(self) -> None: + form: Final = parse_qs(self.rfile.read(int(self.headers["Content-Length"])).decode()) + accepted: Final = self.path == "/token" and self.server.redeem(form) + payload: Final = ( + { + "access_token": FRESH_KEY, + "refresh_token": ROTATED_REFRESH_TOKEN, + "expires_in": 3600, + "user_id": "fixture-user", + "team_id": "fixture-team", + } + if accepted + else {"error": "invalid_grant", "error_description": "the refresh token was already used"} + ) + self.server.request_started.set() + threading.Event().wait(self.server.response_delay) + body: Final = json.dumps(payload).encode() + self.send_response(200 if accepted else 400) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def log_message(self, format: str, *args: object) -> None: + return + + +@pytest.fixture +def token_endpoint() -> Iterator[_TokenEndpoint]: + server: Final = _TokenEndpoint() + thread: Final = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield server + finally: + server.shutdown() + server.server_close() + thread.join(timeout=10) + + +def _environment(home: Path) -> dict[str, str]: + return { + **{ + name: os.environ[name] + for name in ("PATH", "SYSTEMROOT", "WINDIR", "TMPDIR", "TEMP", "TMP", "LANG", "LC_ALL") + if name in os.environ + }, + "HOME": str(home), + "USERPROFILE": str(home), + "PYTHONPATH": str(Path(__file__).resolve().parents[3]), + "PYTHON_DOTENV_DISABLED": "1", + "LITELLM_CLI_DISABLE_KEYRING": "1", + "PYTHON_KEYRING_BACKEND": "keyring.backends.null.Keyring", + "LITELLM_LOCAL_MODEL_COST_MAP": "True", + "LITELLM_API_KEY": "", + "LITELLM_PROXY_URL": "", + } + + +def _seed(home: Path, origin: str, *, fresh: bool = False) -> None: + directory: Final = home / ".litellm" + directory.mkdir(mode=0o700) + path: Final = directory / "token.json" + path.write_text( + json.dumps( + { + "base_url": origin, + "key": FRESH_KEY if fresh else EXPIRED_KEY, + "refresh_token": ROTATED_REFRESH_TOKEN if fresh else REFRESH_TOKEN, + "timestamp": 1.0, + "expires_at": 1e20 if fresh else 0.0, + "client_id": "fixture-client", + "token_endpoint": f"{origin}/token", + "revocation_endpoint": f"{origin}/revoke", + "resource": origin, + "user_id": "fixture-user", + "team_id": "fixture-team", + } + ) + ) + path.chmod(0o600) + + +@contextmanager +def _worker(script: str, home: Path, args: tuple[str, ...] = ()) -> Generator[subprocess.Popen[str], None, None]: + process: Final = subprocess.Popen( + (sys.executable, "-P", "-c", script, *args), + env=_environment(home), + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + try: + assert process.stdout is not None + assert process.stdout.readline() == "ready\n" + yield process + finally: + stop_child(process) + process.communicate(timeout=20) + + +def _print_token(home: Path, origin: str, script: str = _CLI) -> subprocess.CompletedProcess[str]: + return subprocess.run( + (sys.executable, "-P", "-c", script, "--base-url", origin, "auth", "print-token"), + env=_environment(home), + capture_output=True, + text=True, + timeout=30, + check=False, + ) + + +@pytest.mark.timeout(60) +def test_two_processes_with_the_same_stale_snapshot_spend_the_refresh_token_once( + tmp_path: Path, + token_endpoint: _TokenEndpoint, +) -> None: + _seed(tmp_path, token_endpoint.origin) + with _worker(_REFRESH, tmp_path) as first, _worker(_REFRESH, tmp_path) as second: + assert first.stdin is not None and second.stdin is not None + first.stdin.write("refresh\n") + first.stdin.flush() + second.stdin.write("refresh\n") + second.stdin.flush() + first_output: Final = first.communicate(timeout=20) + second_output: Final = second.communicate(timeout=20) + assert first.returncode == second.returncode == 0, (first_output, second_output) + assert first_output == second_output == (f"{FRESH_KEY}\n", "") + result: Final = _print_token(tmp_path, token_endpoint.origin) + assert result.returncode == 0, result.stderr + assert result.stdout == f"{FRESH_KEY}\n" + assert token_endpoint.requests == [ + { + "grant_type": ["refresh_token"], + "refresh_token": [REFRESH_TOKEN], + "client_id": ["fixture-client"], + "resource": [token_endpoint.origin], + } + ] + + +@pytest.mark.timeout(60) +def test_a_killed_credential_writer_releases_the_lock_for_the_next_cli( + tmp_path: Path, + token_endpoint: _TokenEndpoint, +) -> None: + _seed(tmp_path, token_endpoint.origin, fresh=True) + with _worker(_HOLD_LOCK, tmp_path) as holder: + kill_stray_process(holder.pid) + holder.communicate(timeout=20) + result: Final = _print_token(tmp_path, token_endpoint.origin) + assert result.returncode == 0, result.stderr + assert result.stdout == f"{FRESH_KEY}\n" + assert token_endpoint.requests == [] + + +@pytest.mark.timeout(60) +def test_lock_contention_never_sends_an_unlocked_refresh( + tmp_path: Path, + token_endpoint: _TokenEndpoint, +) -> None: + _seed(tmp_path, token_endpoint.origin) + bounded_refresh: Final = _SHORT_WAIT + textwrap.indent(_REFRESH, " ") + with _worker(bounded_refresh, tmp_path) as reader, _worker(_HOLD_LOCK, tmp_path): + output: Final = reader.communicate(input="refresh\n", timeout=20) + assert reader.returncode == 1, output + assert output == ("", "") + assert token_endpoint.requests == [] + + +@pytest.mark.timeout(60) +@pytest.mark.parametrize("reader_script", (_REFRESH, _PRINT_TOKEN), ids=("renewal-helper", "cli")) +def test_helpers_wait_for_a_slow_sibling_renewal_without_logging_in_again( + tmp_path: Path, + token_endpoint: _TokenEndpoint, + reader_script: str, +) -> None: + _seed(tmp_path, token_endpoint.origin) + token_endpoint.response_delay = 11.0 + with ( + _worker(_REFRESH, tmp_path) as renewer, + _worker(reader_script, tmp_path, (token_endpoint.origin,)) as reader, + ): + assert renewer.stdin is not None + renewer.stdin.write("refresh\n") + renewer.stdin.flush() + assert token_endpoint.request_started.wait(timeout=10) + reader_output: Final = reader.communicate(input="read\n", timeout=25) + renewer_output: Final = renewer.communicate(timeout=20) + assert reader.returncode == renewer.returncode == 0, (reader_output, renewer_output) + assert reader_output == renewer_output == (f"{FRESH_KEY}\n", "") + assert len(token_endpoint.requests) == 1 + + +@pytest.mark.timeout(60) +def test_print_token_reports_lock_contention_and_can_retry_without_logging_in_again( + tmp_path: Path, + token_endpoint: _TokenEndpoint, +) -> None: + _seed(tmp_path, token_endpoint.origin) + path: Final = tmp_path / ".litellm" / "token.json" + original: Final = path.read_bytes() + with _worker(_HOLD_LOCK, tmp_path): + result: Final = _print_token(tmp_path, token_endpoint.origin, _SHORT_WAIT + textwrap.indent(_CLI, " ")) + assert result.returncode == 1, result.stderr + assert result.stdout == "" + assert "Retry this command" in result.stderr + assert "lite login" not in result.stderr + assert "Not authenticated" not in result.stderr + assert path.read_bytes() == original + assert token_endpoint.requests == [] + retried: Final = _print_token(tmp_path, token_endpoint.origin) + assert retried.returncode == 0, retried.stderr + assert retried.stdout == f"{FRESH_KEY}\n" + assert len(token_endpoint.requests) == 1 diff --git a/tests/unit/litellm_core_utils/test_cli_token_utils.py b/tests/unit/litellm_core_utils/test_cli_token_utils.py index 7e7eee5373f..11700c9ed66 100644 --- a/tests/unit/litellm_core_utils/test_cli_token_utils.py +++ b/tests/unit/litellm_core_utils/test_cli_token_utils.py @@ -6,6 +6,8 @@ import sys import tempfile import threading import time +from pathlib import Path +from typing import Final import pytest @@ -32,6 +34,7 @@ from litellm.litellm_core_utils.cli_token_utils import ( CredentialNotRecorded, CredentialNotSaved, clear_cli_token, + cli_token_lock, get_cli_token_file_path, get_litellm_gateway_api_key, is_cli_token_fresh, @@ -39,6 +42,8 @@ from litellm.litellm_core_utils.cli_token_utils import ( save_cli_token, ) +from tests.unit.litellm_core_utils.fake_secret_vault import FakeSecretVault + SERVER = "https://proxy.example.com" OTHER_SERVER = "https://other-proxy.example.com" @@ -389,6 +394,58 @@ class TestLoadCliToken: assert load_cli_token(vault=secret_vault_factory(blob=_blob())) is None +def test_load_does_not_migrate_the_credential_when_the_lock_cannot_be_opened( + isolated_home: Path, + secret_vault_factory: type[FakeSecretVault], +) -> None: + path: Final = _write_legacy_file(isolated_home, timestamp=1000.0) + original: Final = path.read_bytes() + vault: Final = secret_vault_factory() + Path(cli_token_lock().lock_file).mkdir() + + assert load_cli_token(vault=vault) is None + assert path.read_bytes() == original + assert vault.reads == 0 + assert vault.writes == [] + + +def test_save_preserves_both_stores_when_the_lock_cannot_be_opened( + isolated_home: Path, + secret_vault_factory: type[FakeSecretVault], +) -> None: + path: Final = _write_legacy_file(isolated_home, timestamp=1000.0) + original: Final = path.read_bytes() + vault: Final = secret_vault_factory(blob=_blob()) + original_secret: Final = vault.blob + Path(cli_token_lock().lock_file).mkdir() + + outcome: Final = save_cli_token( + CliTokenRecord(base_url=SERVER, key="sk-replacement", timestamp=2000.0), vault=vault + ) + assert isinstance(outcome, CredentialNotSaved) + assert path.read_bytes() == original + assert vault.blob == original_secret + assert vault.reads == 0 + assert vault.writes == [] + + +def test_clear_preserves_both_stores_when_the_lock_cannot_be_opened( + isolated_home: Path, + secret_vault_factory: type[FakeSecretVault], +) -> None: + path: Final = _write_legacy_file(isolated_home, timestamp=1000.0) + original: Final = path.read_bytes() + vault: Final = secret_vault_factory(blob=_blob()) + original_secret: Final = vault.blob + Path(cli_token_lock().lock_file).mkdir() + + outcome: Final = clear_cli_token(vault=vault) + assert isinstance(outcome, CredentialNotCleared) + assert path.read_bytes() == original + assert vault.blob == original_secret + assert vault.erases == 0 + + class TestGetLitellmGatewayApiKey: def test_returns_the_vault_secret_when_the_origin_matches(self, isolated_home, secret_vault_factory): _write_metadata_only_file(isolated_home) @@ -417,6 +474,21 @@ class TestGetLitellmGatewayApiKey: assert get_litellm_gateway_api_key(vault=secret_vault_factory(blob=_blob())) is None +def test_sdk_read_does_not_migrate_credentials_when_the_store_cannot_be_locked( + isolated_home: Path, + secret_vault_factory: type[FakeSecretVault], +) -> None: + path: Final = _write_legacy_file(isolated_home, timestamp=1000.0) + original: Final = path.read_bytes() + vault: Final = secret_vault_factory() + Path(cli_token_lock().lock_file).mkdir() + + assert get_litellm_gateway_api_key(expected_base_url=SERVER, vault=vault) is None + assert path.read_bytes() == original + assert vault.reads == 0 + assert vault.writes == [] + + class TestSaveCliToken: def test_secret_goes_to_the_keychain_and_never_to_the_file(self, isolated_home, secret_vault_factory): vault = secret_vault_factory() diff --git a/tests/unit/proxy/client/cli/test_auth_commands.py b/tests/unit/proxy/client/cli/test_auth_commands.py index 5d5ebcf5dfd..bf82ab870ce 100644 --- a/tests/unit/proxy/client/cli/test_auth_commands.py +++ b/tests/unit/proxy/client/cli/test_auth_commands.py @@ -23,6 +23,7 @@ from litellm.litellm_core_utils.cli_token_utils import ( CliTokenRecord, CredentialNotRecorded, CredentialNotSaved, + cli_token_lock, save_cli_token, ) from litellm.proxy.client.cli import cli @@ -35,6 +36,7 @@ from litellm.proxy.client.cli.commands.auth import ( whoami, ) from litellm.proxy.client.cli.commands.claude_settings import SettingsFileOwner +from tests.unit.litellm_core_utils.fake_secret_vault import FakeSecretVault @pytest.fixture @@ -693,7 +695,7 @@ class TestWhoamiCommand: timestamp=time.time() - 3600, ) - with patch("litellm.proxy.client.cli.commands.auth.load_cli_token", return_value=token_data): + with patch("litellm.proxy.client.cli.commands.auth.read_cli_token", return_value=token_data): result = self.runner.invoke(whoami) assert result.exit_code == 0 @@ -705,7 +707,7 @@ class TestWhoamiCommand: def test_whoami_not_authenticated(self): """Test whoami when user is not authenticated""" - with patch("litellm.proxy.client.cli.commands.auth.load_cli_token", return_value=None): + with patch("litellm.proxy.client.cli.commands.auth.read_cli_token", return_value=None): result = self.runner.invoke(whoami) assert result.exit_code == 0 @@ -731,7 +733,7 @@ class TestWhoamiCommand: timestamp=STALE_TOKEN_TIMESTAMP, ) - with patch("litellm.proxy.client.cli.commands.auth.load_cli_token", return_value=token_data): + with patch("litellm.proxy.client.cli.commands.auth.read_cli_token", return_value=token_data): result: Final = self.runner.invoke(whoami) assert result.exit_code == 0 @@ -743,7 +745,7 @@ class TestWhoamiCommand: """Test whoami with token missing some fields""" token_data = CliTokenRecord(key="sk-live", timestamp=time.time() - 3600) - with patch("litellm.proxy.client.cli.commands.auth.load_cli_token", return_value=token_data): + with patch("litellm.proxy.client.cli.commands.auth.read_cli_token", return_value=token_data): result = self.runner.invoke(whoami) assert result.exit_code == 0 @@ -813,7 +815,7 @@ class TestWhoamiCommand: with ( patch( - "litellm.proxy.client.cli.commands.auth.load_cli_token", + "litellm.proxy.client.cli.commands.auth.read_cli_token", return_value=token_data, ), patch("time.time", return_value=1000), @@ -982,7 +984,7 @@ class TestPrintTokenCommand: self.runner = CliRunner() def test_no_stored_token_fails_cleanly(self): - with patch("litellm.proxy.client.cli.commands.auth.load_cli_token", return_value=None): + with patch("litellm.proxy.client.cli.commands.auth.read_cli_token", return_value=None): result = self.runner.invoke(print_token, obj={}) assert result.exit_code != 0 @@ -994,7 +996,7 @@ class TestPrintTokenCommand: one). Must use token.json's own base_url, not a hardcoded default.""" with ( patch( - "litellm.proxy.client.cli.commands.auth.load_cli_token", + "litellm.proxy.client.cli.commands.auth.read_cli_token", return_value=CliTokenRecord( base_url="https://litellm-proxy.corp.com", key="sk-prod-fresh", @@ -1016,7 +1018,7 @@ class TestPrintTokenCommand: token minted for proxy A must not reach a helper invocation aimed at proxy B, even though the token itself is otherwise fresh.""" with patch( - "litellm.proxy.client.cli.commands.auth.load_cli_token", + "litellm.proxy.client.cli.commands.auth.read_cli_token", return_value=CliTokenRecord( base_url="https://other-server.com", key="sk-should-not-print", @@ -1035,7 +1037,7 @@ class TestPrintTokenCommand: """`lite up`'s own bound invocation shape: --base-url matching the token's origin must succeed exactly like the bare/legacy invocation does.""" with patch( - "litellm.proxy.client.cli.commands.auth.load_cli_token", + "litellm.proxy.client.cli.commands.auth.read_cli_token", return_value=CliTokenRecord( base_url="http://localhost:4000", key="sk-matches", @@ -1056,7 +1058,7 @@ class TestPrintTokenCommand: refresh call on every single invocation (apiKeyHelper gets called frequently).""" with patch( - "litellm.proxy.client.cli.commands.auth.load_cli_token", + "litellm.proxy.client.cli.commands.auth.read_cli_token", return_value=CliTokenRecord( base_url="http://localhost:4000", key="sk-cached-fresh", @@ -1091,7 +1093,7 @@ class TestPrintTokenCommand: _register_models_answer("sk-stale-key", proxy_answer) with patch( - "litellm.proxy.client.cli.commands.auth.load_cli_token", + "litellm.proxy.client.cli.commands.auth.read_cli_token", return_value=CliTokenRecord( base_url=STALE_TOKEN_BASE_URL, key="sk-stale-key", @@ -1110,7 +1112,7 @@ class TestPrintTokenCommand: _register_models_answer("sk-stale-key", requests.ReadTimeout("proxy sent nothing")) with patch( - "litellm.proxy.client.cli.commands.auth.load_cli_token", + "litellm.proxy.client.cli.commands.auth.read_cli_token", return_value=CliTokenRecord( base_url=STALE_TOKEN_BASE_URL, key="sk-stale-key", @@ -1128,7 +1130,7 @@ class TestPrintTokenCommand: @responses.activate def test_key_without_a_recorded_proxy_still_expires_after_the_client_window(self): with patch( - "litellm.proxy.client.cli.commands.auth.load_cli_token", + "litellm.proxy.client.cli.commands.auth.read_cli_token", return_value=CliTokenRecord(key="sk-old-format", timestamp=STALE_TOKEN_TIMESTAMP), ): result: Final = self.runner.invoke(print_token, obj={}) @@ -2229,6 +2231,78 @@ def test_print_token_reads_the_connection_from_the_context_object( assert result.stderr == expected_stderr +def test_print_token_uses_the_saved_key_when_the_initializer_did_not_get_one( + isolated_home: Path, + secret_vault_factory: type[FakeSecretVault], +) -> None: + vault: Final = secret_vault_factory() + record: Final = CliTokenRecord( + base_url="https://test.example.com", + key="sk-fixture-reloaded", + timestamp=1000.0, + expires_at=1e20, + refresh_token="fixture-refresh", + ) + assert save_cli_token(record, vault=vault) == SecretStored() + result: Final = CliRunner().invoke( + print_token, + obj={ + "base_url": record.base_url, + "api_key": None, + "api_key_from_token_file": True, + "base_url_explicit": True, + "secret_vault": vault, + }, + ) + assert result.exit_code == 0, result.output + assert result.output == f"{record.key}\n" + + +@responses.activate +def test_logout_reports_a_lock_failure_without_revoking_or_clearing_the_credential( + isolated_home: Path, + secret_vault_factory: type[FakeSecretVault], +) -> None: + _write_home_json(isolated_home, "token.json", _pkce_record(timestamp=1000.0, expires_at=0.0)) + responses.add(responses.POST, f"{PKCE_BASE_URL}/revoke", json={}, status=200) + path: Final = isolated_home / ".litellm" / "token.json" + original: Final = path.read_bytes() + vault: Final = secret_vault_factory(blob=_secret_blob("https://test.example.com", "sk-fixture-vault")) + original_secret: Final = vault.blob + Path(cli_token_lock().lock_file).mkdir() + + result: Final = CliRunner().invoke(logout, obj={"secret_vault": vault}) + assert result.exit_code == 1 + assert result.stdout == "" + assert "Could not lock the CLI credential store. No credential operation was started" in result.stderr + assert len(responses.calls) == 0 + assert path.read_bytes() == original + assert vault.blob == original_secret + assert vault.reads == 0 + assert vault.erases == 0 + + +def test_print_token_reports_an_unavailable_lock_without_claiming_the_login_is_missing( + isolated_home: Path, + secret_vault_factory: type[FakeSecretVault], +) -> None: + _write_home_json(isolated_home, "token.json", _pkce_record(timestamp=1000.0, expires_at=0.0)) + path: Final = isolated_home / ".litellm" / "token.json" + original: Final = path.read_bytes() + vault: Final = secret_vault_factory() + Path(cli_token_lock().lock_file).mkdir() + + result: Final = CliRunner().invoke(print_token, obj={"secret_vault": vault}) + assert result.exit_code == 1 + assert result.stdout == "" + assert "Retry this command" in result.stderr + assert "lite login" not in result.stderr + assert "Not authenticated" not in result.stderr + assert path.read_bytes() == original + assert vault.reads == 0 + assert vault.writes == [] + + def test_print_token_without_a_context_object_reports_no_login_when_nothing_is_stored(isolated_home): result = CliRunner().invoke(print_token) diff --git a/tests/unit/proxy/client/cli/test_pkce_login.py b/tests/unit/proxy/client/cli/test_pkce_login.py index f58bd0ff412..46a9da29f70 100644 --- a/tests/unit/proxy/client/cli/test_pkce_login.py +++ b/tests/unit/proxy/client/cli/test_pkce_login.py @@ -4,17 +4,21 @@ import hashlib import json import re import socket +import sys import threading import time import urllib.error import urllib.request from base64 import urlsafe_b64encode +from pathlib import Path +from typing import Final +from unittest.mock import Mock, call from urllib.parse import parse_qs, urlparse import pytest import requests -from litellm.litellm_core_utils.cli_token_utils import is_cli_token_fresh +from litellm.litellm_core_utils.cli_token_utils import cli_token_lock, is_cli_token_fresh from litellm.proxy.client.cli.commands.pkce_login import ( CallbackCode, CallbackDenied, @@ -125,8 +129,9 @@ def _get(url): return error.code, error.read().decode() -def _fresh(token_data, save, http, reload=lambda: None, **options): - return fresh_api_key(token_data, save, http, reload=reload, **options) +def _fresh(token_data, save, http, reload=None, **options): + reader = reload if reload is not None else lambda: token_data + return fresh_api_key(token_data, save, http, reload=reader, **options) def _challenge_for(verifier: str) -> str: @@ -600,9 +605,7 @@ def test_fresh_api_key_refreshes_near_expiry_and_saves_before_returning(): @pytest.mark.parametrize(("seconds_left", "refreshes"), [(200, True), (600, False)]) -def test_fresh_api_key_refreshes_exactly_when_is_cli_token_fresh_stops_calling_the_key_fresh( - seconds_left, refreshes -): +def test_fresh_api_key_refreshes_exactly_when_is_cli_token_fresh_stops_calling_the_key_fresh(seconds_left, refreshes): """`lite up` and `lite auth print-token` ask `is_cli_token_fresh` and then `fresh_api_key`; a key the first calls stale that the second would not renew left `lite up` refusing a renewable login.""" record = {**STORED, "expires_at": time.time() + seconds_left} @@ -643,17 +646,20 @@ def test_fresh_api_key_reports_why_a_renewal_failed_unless_a_sibling_rotated_fir assert _fresh(STORED, lambda _: None, failing, now=lambda: 999_950.0, warn=warnings.append) == "sk-cli-old" assert _fresh(STORED, lambda _: None, failing, now=lambda: 1_000_001.0, warn=warnings.append) is None assert warnings == ["Could not renew the key: token request failed with 503: temporarily_unavailable"] * 2 - sibling = _fresh(STORED, lambda _: None, failing, reload=lambda: rotated, now=lambda: 1_000_001.0, warn=warnings.append) + sibling = _fresh( + STORED, lambda _: None, failing, reload=lambda: rotated, now=lambda: 1_000_001.0, warn=warnings.append + ) assert sibling == "sk-cli-sibling" assert len(warnings) == 2 -def test_fresh_api_key_uses_a_sibling_rotation_when_its_own_refresh_loses_the_race(): +def test_fresh_api_key_reloads_a_sibling_rotation_before_spending_the_old_refresh_token(): rotated = {**STORED, "key": "sk-cli-sibling", "refresh_token": "llm_srefresh_sibling", "expires_at": 1_003_600.0} http = _FakeHttp({("POST", f"{BASE}/token"): _FakeResponse(400, {"error": "invalid_grant"})}) saved = [] assert _fresh(STORED, saved.append, http, reload=lambda: rotated, now=lambda: 999_950.0) == "sk-cli-sibling" assert _fresh(STORED, saved.append, http, reload=lambda: rotated, now=lambda: 1_000_001.0) == "sk-cli-sibling" + assert http.calls == [] assert _fresh(STORED, saved.append, http, reload=lambda: STORED, now=lambda: 1_000_001.0) is None assert _fresh(STORED, saved.append, http, reload=lambda: None, now=lambda: 1_000_001.0) is None assert saved == [] @@ -696,6 +702,49 @@ def test_fresh_api_key_never_substitutes_a_sibling_record_for_another_proxy_iden assert _fresh(STORED, lambda _: None, http, reload=lambda: foreign, now=lambda: 999_950.0) == "sk-cli-old" +@pytest.mark.parametrize("current", [None, {**STORED, "user_id": "another-user"}]) +def test_fresh_api_key_does_not_renew_a_cleared_or_replaced_login(current): + http = _FakeHttp({("POST", f"{BASE}/token"): _FakeResponse(200, TOKEN_JSON)}) + assert ( + _fresh(STORED, lambda _: pytest.fail("Unexpected save"), http, reload=lambda: current, now=lambda: 1_000_001.0) + is None + ) + assert http.calls == [] + + +@pytest.mark.parametrize(("now", "expected"), [(999_950.0, "sk-cli-old"), (1_000_001.0, None)]) +def test_fresh_api_key_reports_a_lock_failure_without_reading_or_spending_the_refresh_token( + capsys: pytest.CaptureFixture[str], + now: float, + expected: str | None, +) -> None: + Path(cli_token_lock().lock_file).mkdir() + http: Final = _FakeHttp({("POST", f"{BASE}/token"): _FakeResponse(200, TOKEN_JSON)}) + lock_failure: Final = Mock() + key: Final = _fresh( + STORED, + lambda _: pytest.fail("Unexpected credential save"), + http, + reload=lambda: pytest.fail("Credential read without the lock"), + now=lambda: now, + warn=lambda message: print(message, file=sys.stderr), + lock_failure=lock_failure, + ) + assert tuple(lock_failure.call_args_list) == ( + (call("Could not lock the CLI credential store for renewal. Retry this command. No refresh was attempted."),) + if expected is None + else () + ) + assert key == expected + assert http.calls == [] + captured: Final = capsys.readouterr() + assert captured.out == "" + assert ( + captured.err + == "Could not renew the key: the CLI credential store could not be locked; no refresh was attempted\n" + ) + + def test_fresh_api_key_without_refresh_inputs_expires_like_a_classic_key(): http = _FakeHttp() no_refresh = {key: value for key, value in STORED.items() if key != "refresh_token"}