From 12aa72dad4f8a18ea958534a9776ea5ea3aecc91 Mon Sep 17 00:00:00 2001 From: bunnysayzz Date: Thu, 1 Oct 2026 10:14:49 +0530 Subject: [PATCH] fix(cli): stage keyring writes so an interrupted write cannot burn the stored credential A PKCE token refresh spends the single-use server refresh token before the keychain save lands. The keyring write replaced the live entry in one shot, so a failure or kill between deleting the old entry and adding the new one left the user logged out with no way back. Writes now park the new secret in a staging slot and read it back before touching the live slot, so a failed write keeps the old credential and a killed write is recovered from staging by the next read. Logout clears both slots so a stranded staging copy cannot resurrect an ended login. Fixes #43374 --- litellm/litellm_core_utils/cli_keyring.py | 73 +++++++++++++-- .../test_cli_token_utils.py | 93 ++++++++++++++++++- 2 files changed, 158 insertions(+), 8 deletions(-) diff --git a/litellm/litellm_core_utils/cli_keyring.py b/litellm/litellm_core_utils/cli_keyring.py index 70b1773739d..63ef2b68a88 100644 --- a/litellm/litellm_core_utils/cli_keyring.py +++ b/litellm/litellm_core_utils/cli_keyring.py @@ -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,21 @@ 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() 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..17bb0a0908a 100644 --- a/tests/unit/litellm_core_utils/test_cli_token_utils.py +++ b/tests/unit/litellm_core_utils/test_cli_token_utils.py @@ -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, @@ -1058,6 +1059,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 +1105,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 +1113,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 @@ -1232,6 +1259,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