scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

fix(cli): keep lock contention from forcing another login

Browse source 
Wait up to sixty seconds for credential operations, distinguish locked-load failures from missing credentials, and report retryable CLI errors on renewal contention. Preserve the SDK loader contract and valid cached-key fallback.

Verify 314 affected and 101 consumer cases. Reproduce the slow-renewal and misleading-auth failures on the previous head. Targeted Ruff, format and test-quality checks pass with no new targeted type diagnostics. Live interactive and cross-platform verification remain pending.
This commit is contained in:
Jon Tang 2026-11-03 10:50:07 -08:00
parent 128930dcda
commit 711886e2c7
No known key found for this signature in database
6 changed files with 162 additions and 31 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

26
litellm/litellm_core_utils/cli_token_utils.py
View file

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

26
litellm/proxy/client/cli/commands/auth.py
View file

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

10
litellm/proxy/client/cli/commands/pkce_login.py
View file

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

80
tests/integration/sdk/test_cli_refresh_lock_sdk.py
View file

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

43
tests/unit/proxy/client/cli/test_auth_commands.py
View file

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

8
tests/unit/proxy/client/cli/test_pkce_login.py
View file

@ -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 == []