diff --git a/litellm/litellm_core_utils/cli_token_utils.py b/litellm/litellm_core_utils/cli_token_utils.py index 287c418a654..70ae0e292a9 100644 --- a/litellm/litellm_core_utils/cli_token_utils.py +++ b/litellm/litellm_core_utils/cli_token_utils.py @@ -48,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. @@ -148,18 +153,31 @@ def cli_token_lock() -> FileLock: 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=10, mode=0o600, is_singleton=True) + 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 (OSError, Timeout): - return None + 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: diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index cc74d0487cd..9f3dc12e96d 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -4,7 +4,7 @@ import time import webbrowser from collections.abc import Callable, Generator, Mapping, Sequence from contextlib import contextmanager -from typing import Any, Final, Protocol, TypeVar +from typing import Any, Final, NoReturn, Protocol, TypeVar from urllib.parse import urlencode import click @@ -33,6 +33,7 @@ from litellm.litellm_core_utils.cli_keyring import ( from litellm.litellm_core_utils.cli_token_utils import ( CliTokenRecord, CredentialNotCleared, + CredentialNotLoaded, CredentialNotRecorded, CredentialNotSaved, SecretSave, @@ -40,7 +41,7 @@ from litellm.litellm_core_utils.cli_token_utils import ( cli_token_lock, get_cli_token_file_path, is_cli_token_fresh, - load_cli_token, + read_cli_token, save_cli_token, ) @@ -226,13 +227,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: @@ -287,6 +289,7 @@ def get_stored_api_key( requests.Session(), reload=_renewal_reader(vault), warn=_warn, + lock_failure=_raise_credential_error, ) @@ -294,6 +297,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" @@ -1094,6 +1101,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 3ec6962613b..4bc5a2dd377 100644 --- a/litellm/proxy/client/cli/commands/pkce_login.py +++ b/litellm/proxy/client/cli/commands/pkce_login.py @@ -499,6 +499,7 @@ 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. @@ -513,15 +514,20 @@ def fresh_api_key( ): return _fresh_api_key(token_data, save, http, reload=reload, now=now, warn=warn) key: Final = token_data.get("key") - still_valid: Final = key if isinstance(key, str) and now() < expires_at else None 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 still_valid + 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)): diff --git a/tests/integration/sdk/test_cli_refresh_lock_sdk.py b/tests/integration/sdk/test_cli_refresh_lock_sdk.py index d92f994f695..d97831901fb 100644 --- a/tests/integration/sdk/test_cli_refresh_lock_sdk.py +++ b/tests/integration/sdk/test_cli_refresh_lock_sdk.py @@ -47,6 +47,24 @@ _HOLD_LOCK: Final = textwrap.dedent( """ ) _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): @@ -56,6 +74,8 @@ class _TokenEndpoint(ThreadingHTTPServer): 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: @@ -84,6 +104,8 @@ class _TokenHandler(BaseHTTPRequestHandler): 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") @@ -152,9 +174,9 @@ def _seed(home: Path, origin: str, *, fresh: bool = False) -> None: @contextmanager -def _worker(script: str, home: Path) -> Generator[subprocess.Popen[str], None, None]: +def _worker(script: str, home: Path, args: tuple[str, ...] = ()) -> Generator[subprocess.Popen[str], None, None]: process: Final = subprocess.Popen( - (sys.executable, "-P", "-c", script), + (sys.executable, "-P", "-c", script, *args), env=_environment(home), stdin=subprocess.PIPE, stdout=subprocess.PIPE, @@ -170,9 +192,9 @@ def _worker(script: str, home: Path) -> Generator[subprocess.Popen[str], None, N process.communicate(timeout=20) -def _print_token(home: Path, origin: str) -> subprocess.CompletedProcess[str]: +def _print_token(home: Path, origin: str, script: str = _CLI) -> subprocess.CompletedProcess[str]: return subprocess.run( - (sys.executable, "-P", "-c", _CLI, "--base-url", origin, "auth", "print-token"), + (sys.executable, "-P", "-c", script, "--base-url", origin, "auth", "print-token"), env=_environment(home), capture_output=True, text=True, @@ -231,8 +253,56 @@ def test_lock_contention_never_sends_an_unlocked_refresh( token_endpoint: _TokenEndpoint, ) -> None: _seed(tmp_path, token_endpoint.origin) - with _worker(_REFRESH, tmp_path) as reader, _worker(_HOLD_LOCK, tmp_path): + 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/proxy/client/cli/test_auth_commands.py b/tests/unit/proxy/client/cli/test_auth_commands.py index f8db47f47f5..0c6a198f514 100644 --- a/tests/unit/proxy/client/cli/test_auth_commands.py +++ b/tests/unit/proxy/client/cli/test_auth_commands.py @@ -680,7 +680,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 @@ -692,7 +692,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 @@ -709,7 +709,7 @@ class TestWhoamiCommand: timestamp=time.time() - (25 * 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 @@ -720,7 +720,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 @@ -790,7 +790,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), @@ -959,7 +959,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 @@ -971,7 +971,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", @@ -993,7 +993,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", @@ -1012,7 +1012,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", @@ -1033,7 +1033,7 @@ class TestPrintTokenCommand: 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", @@ -1057,7 +1057,7 @@ class TestPrintTokenCommand: 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-stale-key", @@ -2206,6 +2206,27 @@ def test_logout_reports_a_lock_failure_without_revoking_or_clearing_the_credenti 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 84f0580f64f..46a9da29f70 100644 --- a/tests/unit/proxy/client/cli/test_pkce_login.py +++ b/tests/unit/proxy/client/cli/test_pkce_login.py @@ -12,6 +12,7 @@ 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 @@ -719,6 +720,7 @@ def test_fresh_api_key_reports_a_lock_failure_without_reading_or_spending_the_re ) -> 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"), @@ -726,6 +728,12 @@ def test_fresh_api_key_reports_a_lock_failure_without_reading_or_spending_the_re 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 == []