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

Merge 526a8cecc0 into aeede53112

Browse source
This commit is contained in:
Azhar 2026-11-06 15:53:26 -05:00 • committed by GitHub
commit 1763eb6d77
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 175 additions and 46 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

75
litellm/litellm_core_utils/cli_keyring.py
View file

@ -13,6 +13,10 @@ A write is only reported as stored once it has been read back, because keyring's
null backend, which `keyring --disable` and headless CI images both select,
accepts every write and keeps nothing. Writes are also pre-flighted with a
throwaway value, because a keychain can answer neither way and block forever.
The new secret is parked in a staging slot and read back before the live slot is
touched, so a write that fails or is killed halfway cannot take the previous
credential down with it; a staging copy outliving its write is promoted by the
next read.
"""
import os
@ -23,6 +27,7 @@ from typing import Final, Protocol, TypeAlias
KEYRING_SERVICE: Final = "litellm-cli"
KEYRING_ACCOUNT: Final = "credential"
KEYRING_STAGING_ACCOUNT: Final = "credential-staging"
KEYRING_PREFLIGHT_ACCOUNT: Final = "credential-preflight"
DISABLE_KEYRING_ENV_VAR: Final = "LITELLM_CLI_DISABLE_KEYRING"
@ -164,6 +169,14 @@ class KeyringVault:
stopped_answering: threading.Event = field(default_factory=threading.Event, compare=False, repr=False)
def read(self) -> SecretRead:
"""The stored secret, promoting a write that died after staging but before replacing.
A write parks the new secret in the staging slot before it touches the live one, so a
kill between the two leaves the staging copy as the newest credential this machine saved.
Finding the live slot empty with a staging copy present means exactly that happened, and
handing the staged secret back (with a best-effort copy into the live slot) recovers it
instead of reporting the login gone.
"""
if self.stopped_answering.is_set():
return KeyringUnreachable()
api: Final = _keyring_api()
@ -173,7 +186,17 @@ class KeyringVault:
blob: Final = api.get_password(KEYRING_SERVICE, KEYRING_ACCOUNT)
except Exception: # noqa: BLE001 # backends raise outside keyring.errors; never break the SDK
return KeyringUnreachable()
return SecretMissing() if blob is None else SecretFound(blob)
if blob is not None:
return SecretFound(blob)
try:
staged: Final = api.get_password(KEYRING_SERVICE, KEYRING_STAGING_ACCOUNT)
except Exception: # noqa: BLE001 # backends raise outside keyring.errors; never break the SDK
return KeyringUnreachable()
if staged is None:
return SecretMissing()
with suppress(Exception):
api.set_password(KEYRING_SERVICE, KEYRING_ACCOUNT, staged)
return SecretFound(staged)
def write(self, blob: str) -> SecretWrite:
"""Store the secret, reporting stored only once the keychain hands the same bytes back.
@ -183,6 +206,12 @@ class KeyringVault:
distinguish itself. Reading the value back is the only way to tell it apart from a keychain
that really stored the credential, and the caller is about to drop its own copy on our word.
The new secret is parked in the staging slot and read back before the live slot is touched.
Replacing the live entry is not one keychain operation, so a write that fails or is killed
halfway can no longer take the previous credential down with it: the old secret is still in
the live slot when staging fails, and the staged copy is what the next read promotes when
the live slot did not survive. The staging copy is removed once the live slot reads back.
The keychain is pre-flighted first, because one that blocks rather than answering would
otherwise hang `lite login` outright.
"""
@ -195,14 +224,39 @@ class KeyringVault:
self.stopped_answering.set()
return KeyringUnreachable()
_forget_the_preflight(api)
staged: Final = self._stage(api, blob)
if staged is not None:
return staged
try:
api.set_password(KEYRING_SERVICE, KEYRING_ACCOUNT, blob)
except Exception: # noqa: BLE001 # a keychain that refuses the write falls back to the token file
return KeyringUnreachable()
return SecretStored() if self.read() == SecretFound(blob) else KeyringDiscardsWrites()
stored: Final = self.read() == SecretFound(blob)
self._clear_staging(api)
return SecretStored() if stored else KeyringDiscardsWrites()
def _stage(self, api: KeyringApi, blob: str) -> SecretWrite | None:
"""Park the new secret where a failed live write cannot strand it. None means parked."""
try:
api.set_password(KEYRING_SERVICE, KEYRING_STAGING_ACCOUNT, blob)
except Exception: # noqa: BLE001 # backends raise outside keyring.errors; never break the SDK
return KeyringUnreachable()
try:
parked: Final = api.get_password(KEYRING_SERVICE, KEYRING_STAGING_ACCOUNT)
except Exception: # noqa: BLE001 # backends raise outside keyring.errors; never break the SDK
return KeyringUnreachable()
return None if parked == blob else KeyringDiscardsWrites()
def _clear_staging(self, api: KeyringApi) -> None:
with suppress(Exception):
api.delete_password(KEYRING_SERVICE, KEYRING_STAGING_ACCOUNT)
def erase(self) -> SecretErase:
"""Remove our entry, reporting whether the keychain is guaranteed to be free of it.
"""Remove our entries, live and staging, reporting whether the keychain is guaranteed
to be free of them.
The staging slot has to go too: a copy a killed write left behind would otherwise be
promoted by the next read and resurrect the login that was just ended.
A keychain out of reach is never an erasure: the entry belongs to the OS, not to this
install, so it outlives an uninstalled `keyring` package and a kill switch set after login.
@ -215,14 +269,23 @@ class KeyringVault:
case SecretMissing():
return SecretErased()
case SecretFound():
return self._delete()
return self._delete_all()
def _delete(self) -> SecretErase:
def _delete_all(self) -> SecretErase:
api: Final = _keyring_api()
if isinstance(api, (KeyringNotInstalled, KeyringDisabled)):
return api
live: Final = self._delete_account(api, KEYRING_ACCOUNT)
staged: Final = self._delete_account(api, KEYRING_STAGING_ACCOUNT)
return (
SecretErased() if isinstance(live, SecretErased) and isinstance(staged, SecretErased) else SecretStranded()
)
def _delete_account(self, api: KeyringApi, account: str) -> SecretErase:
try:
api.delete_password(KEYRING_SERVICE, KEYRING_ACCOUNT)
if api.get_password(KEYRING_SERVICE, account) is None:
return SecretErased()
api.delete_password(KEYRING_SERVICE, account)
except Exception: # noqa: BLE001 # report the failure as a value so `lite logout` can warn
return SecretStranded()
return SecretErased()

146
tests/unit/litellm_core_utils/test_cli_token_utils.py
View file

@ -15,6 +15,7 @@ from litellm.litellm_core_utils.cli_keyring import (
KEYRING_ACCOUNT,
KEYRING_PREFLIGHT_ACCOUNT,
KEYRING_SERVICE,
KEYRING_STAGING_ACCOUNT,
KeyringDisabled,
KeyringDiscardsWrites,
KeyringNotInstalled,
@ -217,9 +218,7 @@ class TestLoadCliToken:
assert json.loads(vault.blob)["key"] == "sk-live"
assert json.loads(vault.blob)["refresh_token"] == "rt-live"
def test_a_superseded_refresh_token_on_disk_never_outlives_the_keychain(
self, isolated_home, secret_vault_factory
):
def test_a_superseded_refresh_token_on_disk_never_outlives_the_keychain(self, isolated_home, secret_vault_factory):
"""Two stores, two sign-ins, and the newer one is in the keychain. Handing back its key with
the older one's refresh token would build a credential neither store ever held, and would
renew the login the user already replaced."""
@ -280,9 +279,7 @@ class TestLoadCliToken:
assert json.loads(vault.blob)["key"] == "sk-fresh"
assert "key" not in json.loads(path.read_text())
def test_a_login_the_file_could_not_record_is_the_one_that_gets_used(
self, isolated_home, secret_vault_factory
):
def test_a_login_the_file_could_not_record_is_the_one_that_gets_used(self, isolated_home, secret_vault_factory):
"""A login the keychain took and the file could not be pointed at afterwards leaves the
superseded secret sitting on disk in front of the fresh one. Serving the file's copy would
put a credential the user just replaced, and may well have just revoked, back into every
@ -297,9 +294,7 @@ class TestLoadCliToken:
assert json.loads(vault.blob)["key"] == "sk-fresh"
assert "key" not in json.loads(path.read_text())
def test_a_secret_written_to_disk_after_the_keychain_entry_still_wins(
self, isolated_home, secret_vault_factory
):
def test_a_secret_written_to_disk_after_the_keychain_entry_still_wins(self, isolated_home, secret_vault_factory):
"""The other direction of the same rule, which is the common one: a login that fell back to
the file because the keychain refused it is newer than whatever the keychain kept."""
path = _write_legacy_file(isolated_home, key="sk-fresh", timestamp=2000.0)
@ -311,9 +306,7 @@ class TestLoadCliToken:
assert json.loads(vault.blob)["key"] == "sk-fresh"
assert "key" not in json.loads(path.read_text())
def test_a_disk_secret_survives_when_the_stale_vault_refuses_the_rewrite(
self, isolated_home, secret_vault_factory
):
def test_a_disk_secret_survives_when_the_stale_vault_refuses_the_rewrite(self, isolated_home, secret_vault_factory):
path = _write_legacy_file(isolated_home, key="sk-fresh")
before = path.read_text()
@ -431,9 +424,7 @@ class TestSaveCliToken:
assert json.loads(vault.blob)["key"] == "sk-new"
assert load_cli_token(vault=vault).key == "sk-new"
def test_the_refresh_token_goes_to_the_keychain_and_never_to_the_file(
self, isolated_home, secret_vault_factory
):
def test_the_refresh_token_goes_to_the_keychain_and_never_to_the_file(self, isolated_home, secret_vault_factory):
vault = secret_vault_factory()
stored = save_cli_token(
@ -523,9 +514,7 @@ class TestSaveCliToken:
assert vault.blob is None
@pytest.mark.skipif(os.geteuid() == 0, reason="root ignores directory permissions")
def test_a_login_that_cannot_be_saved_leaves_the_working_one_alone(
self, isolated_home, secret_vault_factory
):
def test_a_login_that_cannot_be_saved_leaves_the_working_one_alone(self, isolated_home, secret_vault_factory):
"""Signing in again on a machine whose ~/.litellm has gone read-only must not cost the user
the credential they already had. Overwriting the keychain and then failing to record it, or
undoing that write afterwards, would take a login that still works out from under them."""
@ -558,9 +547,7 @@ class TestSaveCliToken:
assert isinstance(outcome, CredentialNotRecorded)
assert json.loads(vault.blob)["key"] == "sk-new"
def test_the_credential_the_file_cannot_name_is_left_in_the_keychain(
self, isolated_home, secret_vault_factory
):
def test_the_credential_the_file_cannot_name_is_left_in_the_keychain(self, isolated_home, secret_vault_factory):
"""The keychain holds one entry, so the secret that was there went the moment this one
landed. Taking the new one back out would turn a login this machine may still be able to
use into no login at all, and it cannot restore the old one either way."""
@ -573,7 +560,9 @@ class TestSaveCliToken:
assert vault.blob is not None
def test_a_failed_write_leaves_the_previous_credential_intact(self, isolated_home, secret_vault_factory, monkeypatch):
def test_a_failed_write_leaves_the_previous_credential_intact(
self, isolated_home, secret_vault_factory, monkeypatch
):
path = _write_legacy_file(isolated_home)
before = path.read_text()
@ -615,9 +604,7 @@ class TestSaveCliToken:
assert load_cli_token(vault=vault).key == "sk-fresh"
def test_a_login_on_a_clock_that_moved_forwards_keeps_its_own_time(
self, isolated_home, secret_vault_factory
):
def test_a_login_on_a_clock_that_moved_forwards_keeps_its_own_time(self, isolated_home, secret_vault_factory):
"""Pinning the stamp above the previous login is only ever a floor. The ordinary case has
to record when the user actually signed in, because that is what decides expiry."""
_write_legacy_file(isolated_home, key="sk-old", timestamp=1000.0)
@ -727,7 +714,6 @@ class TestScrubFailure:
assert load_cli_token(vault=vault).key == "sk-legacy"
assert json.loads(path.read_text()).get("key") is None
@pytest.mark.skipif(os.geteuid() == 0, reason="root ignores file permissions")
def test_a_rejoin_the_file_refuses_never_takes_the_key_with_it(
self, isolated_home, secret_vault_factory, monkeypatch
@ -777,9 +763,7 @@ class TestClearCliToken:
assert clear_cli_token(vault=vault) == SecretStranded()
assert load_cli_token(vault=vault) is None
@pytest.mark.parametrize(
"failure", [KeyringDisabled(), KeyringUnreachable(), KeyringNotInstalled()]
)
@pytest.mark.parametrize("failure", [KeyringDisabled(), KeyringUnreachable(), KeyringNotInstalled()])
def test_a_secret_in_the_file_is_no_evidence_about_a_keychain_that_exists(
self, isolated_home, secret_vault_factory, failure
):
@ -796,9 +780,7 @@ class TestClearCliToken:
assert clear_cli_token(vault=vault) == failure
assert json.loads(_token_file(isolated_home).read_text()).get("key") is None
def test_a_second_logout_still_reports_the_keychain_it_could_not_clear(
self, isolated_home, secret_vault_factory
):
def test_a_second_logout_still_reports_the_keychain_it_could_not_clear(self, isolated_home, secret_vault_factory):
"""The first logout deletes the file and tells the user to run it again once the keychain is
reachable. If the second run reads that missing file as proof of a clean keychain, the advice
turns into the very false all-clear it was issued to prevent."""
@ -846,9 +828,7 @@ class TestClearCliToken:
assert "sk-legacy" not in left_on_disk
assert "rt-legacy" not in left_on_disk
def test_a_repeat_logout_never_answers_its_own_warning_with_an_all_clear(
self, isolated_home, secret_vault_factory
):
def test_a_repeat_logout_never_answers_its_own_warning_with_an_all_clear(self, isolated_home, secret_vault_factory):
"""Sign in while the keychain works, sign in again once it has gone out of reach so the
second secret lands in the file, then log out twice. The first logout cannot say the first
login's entry is gone, and says so. If the second one reads the file the first one took
@ -1058,6 +1038,31 @@ class _KeychainHeldByABlockedWrite(_NeverAnsweringKeyringModule):
return self.stored
class _StagingKeyring:
"""A keychain that keeps every account separately, and can die between the delete and the add
of a write the way an interrupted macOS delete-then-add does."""
def __init__(self, *, die_after_delete=()):
self.store = {}
self.calls = []
self.die_after_delete = set(die_after_delete)
def get_password(self, service_name, username):
self.calls.append(("get", service_name, username))
return self.store.get(username)
def set_password(self, service_name, username, password):
self.calls.append(("set", service_name, username))
self.store.pop(username, None)
if username in self.die_after_delete:
raise RuntimeError("killed mid-write")
self.store[username] = password
def delete_password(self, service_name, username):
self.calls.append(("delete", service_name, username))
self.store.pop(username, None)
def _answered_within(seconds, call):
answers = []
worker = threading.Thread(target=lambda: answers.append(call()), daemon=True)
@ -1079,7 +1084,7 @@ def install_fake_keyring(monkeypatch):
class TestKeyringVault:
def test_round_trips_through_the_installed_keyring(self, install_fake_keyring):
fake = install_fake_keyring(_FakeKeyringModule())
fake = install_fake_keyring(_StagingKeyring())
vault = KeyringVault()
assert vault.write("blob-1") == SecretStored()
@ -1087,7 +1092,8 @@ class TestKeyringVault:
assert vault.erase() == SecretErased()
assert vault.read() == SecretMissing()
assert {call[1] for call in fake.calls} == {KEYRING_SERVICE}
assert {call[2] for call in fake.calls} == {KEYRING_ACCOUNT, KEYRING_PREFLIGHT_ACCOUNT}
assert {call[2] for call in fake.calls} == {KEYRING_ACCOUNT, KEYRING_STAGING_ACCOUNT, KEYRING_PREFLIGHT_ACCOUNT}
assert KEYRING_STAGING_ACCOUNT not in fake.store
def test_the_kill_switch_reports_no_keychain(self, monkeypatch):
"""`LITELLM_CLI_DISABLE_KEYRING` has to work without importing keyring, because keyring
@ -1214,9 +1220,7 @@ class TestKeyringVault:
finally:
keyring.set_keyring(previous)
def test_a_credential_survives_a_backend_that_keeps_nothing(
self, isolated_home, install_fake_keyring
):
def test_a_credential_survives_a_backend_that_keeps_nothing(self, isolated_home, install_fake_keyring):
"""The end of the same story: the credential must still be usable afterwards. Reporting the
discard is only worth anything if the token file then keeps the copy the keychain refused."""
install_fake_keyring(_FakeKeyringModule(discard=True))
@ -1232,6 +1236,68 @@ class TestKeyringVault:
assert KeyringVault().erase() == KeyringUnreachable()
def test_a_write_parks_the_secret_in_staging_before_touching_the_live_slot(self, install_fake_keyring):
"""The live slot is only replaced once the new secret has been read back from staging, so a
write the keychain takes badly can no longer strand the previous credential."""
fake = install_fake_keyring(_StagingKeyring())
vault = KeyringVault()
fake.store[KEYRING_ACCOUNT] = "blob-old"
assert vault.write("blob-new") == SecretStored()
sets = [call for call in fake.calls if call[0] == "set" and call[2] != KEYRING_PREFLIGHT_ACCOUNT]
assert [call[2] for call in sets] == [KEYRING_STAGING_ACCOUNT, KEYRING_ACCOUNT]
assert vault.read() == SecretFound("blob-new")
assert KEYRING_STAGING_ACCOUNT not in fake.store
def test_a_staging_write_that_fails_leaves_the_previous_credential_untouched(self, install_fake_keyring):
"""A keychain that refuses the new secret must cost the caller that secret, never the one it
is replacing. This is the shape of #43374: the refresh already burned the old token
server-side, so losing the stored copy too is what logged the CLI out."""
fake = install_fake_keyring(_StagingKeyring(die_after_delete={KEYRING_STAGING_ACCOUNT}))
vault = KeyringVault()
fake.store[KEYRING_ACCOUNT] = "blob-old"
assert vault.write("blob-new") == KeyringUnreachable()
assert vault.read() == SecretFound("blob-old")
def test_an_interrupted_live_write_is_recovered_from_staging_on_the_next_read(self, install_fake_keyring):
"""A kill between the live slot's delete and add used to leave no credential at all: the
server had already burned the old refresh token, and the new pair died in memory with the
process. The staged copy is the newest credential this machine saved, so the next read
hands it back instead of reporting the login gone."""
fake = install_fake_keyring(_StagingKeyring(die_after_delete={KEYRING_ACCOUNT}))
vault = KeyringVault()
fake.store[KEYRING_ACCOUNT] = "blob-old"
assert vault.write("blob-new") == KeyringUnreachable()
assert fake.store.get(KEYRING_ACCOUNT) is None
assert fake.store[KEYRING_STAGING_ACCOUNT] == "blob-new"
assert vault.read() == SecretFound("blob-new")
assert fake.store[KEYRING_ACCOUNT] == "blob-new"
def test_a_stranded_staging_copy_does_not_replace_a_live_credential(self, install_fake_keyring):
"""Promotion is only for a live slot that did not survive its write. A live slot that still
answers keeps serving, and the next save overwrites the stranded staging copy."""
fake = install_fake_keyring(_StagingKeyring())
vault = KeyringVault()
fake.store[KEYRING_ACCOUNT] = "blob-old"
fake.store[KEYRING_STAGING_ACCOUNT] = "blob-stranded"
assert vault.read() == SecretFound("blob-old")
def test_logout_removes_a_stranded_staging_copy(self, install_fake_keyring):
"""A staging copy a killed write left behind must not resurrect the login that was just
ended the next time anything reads the vault."""
fake = install_fake_keyring(_StagingKeyring())
vault = KeyringVault()
fake.store[KEYRING_STAGING_ACCOUNT] = "blob-stranded"
assert vault.erase() == SecretErased()
assert fake.store == {}
assert vault.read() == SecretMissing()
class TestIsCliTokenFreshWithExpiresAt:
"""A ``lite login --pkce`` record carries the proxy's own ``expires_at``, which wins