mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
Merge 128930dcda into a392a39627
This commit is contained in:
commit
4465cc9fae
7 changed files with 501 additions and 21 deletions
|
|
@ -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 (
|
||||
|
|
@ -139,8 +140,29 @@ def get_cli_token_file_path() -> str:
|
|||
return str(config_dir / "token.json")
|
||||
|
||||
|
||||
def cli_token_lock() -> FileLock:
|
||||
"""One reentrant process lock for the metadata file and shared keychain entry.
|
||||
|
||||
Keep it beside the credential directory so a read-only directory still permits 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=10, 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"""
|
||||
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
|
||||
|
||||
|
||||
def _load_cli_token(vault: SecretVault) -> CliTokenRecord | None:
|
||||
record: Final = _read_token_file()
|
||||
if record is None:
|
||||
return None
|
||||
|
|
@ -148,6 +170,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 +248,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
|
||||
|
|
|
|||
|
|
@ -2,12 +2,14 @@ import os
|
|||
import sys
|
||||
import time
|
||||
import webbrowser
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from collections.abc import Callable, Generator, Mapping, Sequence
|
||||
from contextlib import contextmanager
|
||||
from typing import Any, Final, 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
|
||||
|
|
@ -34,6 +36,7 @@ from litellm.litellm_core_utils.cli_token_utils import (
|
|||
CredentialNotSaved,
|
||||
SecretSave,
|
||||
clear_cli_token,
|
||||
cli_token_lock,
|
||||
get_cli_token_file_path,
|
||||
is_cli_token_fresh,
|
||||
load_cli_token,
|
||||
|
|
@ -852,6 +855,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 +1008,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 +1087,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,
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
@ -499,16 +500,46 @@ def fresh_api_key(
|
|||
now: Callable[[], float] = time.time,
|
||||
warn: Callable[[str], None] = _ignore_warning,
|
||||
) -> 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."""
|
||||
"""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")
|
||||
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):
|
||||
warn("Could not renew the key: the CLI credential store could not be locked; no refresh was attempted")
|
||||
return still_valid
|
||||
try:
|
||||
current: Final = reload()
|
||||
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:
|
||||
key: Final = token_data.get("key")
|
||||
if not isinstance(key, str) or not key:
|
||||
return None
|
||||
|
|
|
|||
238
tests/integration/sdk/test_cli_refresh_lock_sdk.py
Normal file
238
tests/integration/sdk/test_cli_refresh_lock_sdk.py
Normal file
|
|
@ -0,0 +1,238 @@
|
|||
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()"
|
||||
|
||||
|
||||
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()
|
||||
|
||||
@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"}
|
||||
)
|
||||
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) -> Generator[subprocess.Popen[str], None, None]:
|
||||
process: Final = subprocess.Popen(
|
||||
(sys.executable, "-P", "-c", script),
|
||||
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) -> subprocess.CompletedProcess[str]:
|
||||
return subprocess.run(
|
||||
(sys.executable, "-P", "-c", _CLI, "--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)
|
||||
with _worker(_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 == []
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -2229,6 +2231,57 @@ 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_without_a_context_object_reports_no_login_when_nothing_is_stored(isolated_home):
|
||||
result = CliRunner().invoke(print_token)
|
||||
|
||||
|
|
|
|||
|
|
@ -4,17 +4,20 @@ 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 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 +128,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 +604,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 +645,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 +701,42 @@ 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)})
|
||||
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),
|
||||
)
|
||||
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"}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue