Merge 5ea9261e1f into 982d3a5476

litellm/litellm_core_utils/cli_keyring.py

@@ -17,9 +17,10 @@ throwaway value, because a keychain can answer neither way and block forever.
 
 import os
 import threading
-from contextlib import suppress
+from collections.abc import Callable
 from dataclasses import dataclass, field
-from typing import Final, Protocol, TypeAlias
+from queue import Empty, Queue
+from typing import Final, Generic, Protocol, TypeAlias, TypeVar
 
 KEYRING_SERVICE: Final = "litellm-cli"
 KEYRING_ACCOUNT: Final = "credential"
@@ -29,6 +30,9 @@ DISABLE_KEYRING_ENV_VAR: Final = "LITELLM_CLI_DISABLE_KEYRING"
 _DISABLED_VALUES: Final = frozenset(("1", "true", "yes", "on"))
 _PREFLIGHT_VALUE: Final = "preflight"
 _PREFLIGHT_TIMEOUT_SECONDS: Final = 5.0
+_MAX_CREDENTIAL_BLOB_BYTES: Final = 5 * 512
+
+_T = TypeVar("_T")
 
 
 @dataclass(frozen=True, slots=True)
@@ -76,9 +80,14 @@ class KeyringDiscardsWrites:
     pass
 
 
+@dataclass(frozen=True, slots=True)
+class SecretTooLarge:
+    pass
+
+
 KeyringUnusable: TypeAlias = KeyringNotInstalled | KeyringDisabled | KeyringUnreachable
 SecretRead: TypeAlias = SecretFound | SecretMissing | KeyringUnusable
-SecretWrite: TypeAlias = SecretStored | KeyringUnusable | KeyringDiscardsWrites
+SecretWrite: TypeAlias = SecretStored | KeyringUnusable | KeyringDiscardsWrites | SecretTooLarge
 SecretErase: TypeAlias = SecretErased | SecretStranded | KeyringUnusable
 
 
@@ -100,6 +109,42 @@ class KeyringApi(Protocol):
     def delete_password(self, service_name: str, username: str) -> None: ...
 
 
+@dataclass(frozen=True, slots=True)
+class _KeyringCallAnswered(Generic[_T]):
+    value: _T
+
+
+@dataclass(frozen=True, slots=True)
+class _KeyringCallFailed:
+    error: Exception
+
+
+@dataclass(frozen=True, slots=True)
+class _KeyringCallTimedOut:
+    pass
+
+
+_KeyringCallResult: TypeAlias = _KeyringCallAnswered[_T] | _KeyringCallFailed | _KeyringCallTimedOut
+
+
+@dataclass(frozen=True, slots=True)
+class _ProbeStored:
+    pass
+
+
+@dataclass(frozen=True, slots=True)
+class _ProbeRefused:
+    pass
+
+
+@dataclass(frozen=True, slots=True)
+class _ProbeSilent:
+    pass
+
+
+_WriteProbe: TypeAlias = _ProbeStored | _ProbeRefused | _ProbeSilent
+
+
 def _keyring_disabled() -> bool:
     return os.getenv(DISABLE_KEYRING_ENV_VAR, "").strip().lower() in _DISABLED_VALUES
 
@@ -119,7 +164,23 @@ def _keyring_api() -> KeyringApi | KeyringNotInstalled | KeyringDisabled:
     return KeyringNotInstalled() if api is None else api
 
 
-def _answers_a_write(api: KeyringApi, timeout_seconds: float) -> bool:
+def _bounded_keyring_call(call: Callable[[], _T], timeout_seconds: float) -> _KeyringCallResult[_T]:
+    answers: Final = Queue[_KeyringCallResult[_T]](maxsize=1)
+
+    def run() -> None:
+        try:
+            answers.put(_KeyringCallAnswered(call()))
+        except Exception as error:  # noqa: BLE001  # keyring backends raise outside keyring.errors
+            answers.put(_KeyringCallFailed(error))
+
+    threading.Thread(target=run, daemon=True, name="litellm-cli-keyring-call").start()
+    try:
+        return answers.get(timeout=timeout_seconds)
+    except Empty:
+        return _KeyringCallTimedOut()
+
+
+def _probe_a_write(api: KeyringApi, timeout_seconds: float) -> _WriteProbe:
     """Whether the keychain answers a write at all, asked with a value worth nothing.
 
     macOS derives the login keychain from `$HOME`, and `set_password` against a HOME with no usable
@@ -129,25 +190,30 @@ def _answers_a_write(api: KeyringApi, timeout_seconds: float) -> bool:
     it, and keeps the real credential out of a store that might accept it long after we gave up.
     A keychain that refuses the probe outright still answered it, so only silence counts against it.
     """
-    answered: Final = threading.Event()
-
-    def ask() -> None:
-        with suppress(Exception):
-            api.set_password(KEYRING_SERVICE, KEYRING_PREFLIGHT_ACCOUNT, _PREFLIGHT_VALUE)
-        answered.set()
-
-    threading.Thread(target=ask, daemon=True, name="litellm-cli-keyring-preflight").start()
-    return answered.wait(timeout_seconds)
+    result: Final = _bounded_keyring_call(
+        lambda: api.set_password(KEYRING_SERVICE, KEYRING_PREFLIGHT_ACCOUNT, _PREFLIGHT_VALUE),
+        timeout_seconds,
+    )
+    match result:
+        case _KeyringCallAnswered():
+            return _ProbeStored()
+        case _KeyringCallFailed():
+            return _ProbeRefused()
+        case _KeyringCallTimedOut():
+            return _ProbeSilent()
 
 
-def _forget_the_preflight(api: KeyringApi) -> None:
+def _forget_the_preflight(api: KeyringApi, timeout_seconds: float) -> bool:
     """Take the throwaway probe back out.
 
     A backend that kept nothing has nothing to remove, and the probe is worth nothing either way,
     so a keychain that refuses to give it up costs the caller nothing.
     """
-    with suppress(Exception):
-        api.delete_password(KEYRING_SERVICE, KEYRING_PREFLIGHT_ACCOUNT)
+    result: Final = _bounded_keyring_call(
+        lambda: api.delete_password(KEYRING_SERVICE, KEYRING_PREFLIGHT_ACCOUNT),
+        timeout_seconds,
+    )
+    return not isinstance(result, _KeyringCallTimedOut)
 
 
 @dataclass(frozen=True, slots=True)
@@ -162,18 +228,26 @@ class KeyringVault:
 
     preflight_timeout_seconds: float = _PREFLIGHT_TIMEOUT_SECONDS
     stopped_answering: threading.Event = field(default_factory=threading.Event, compare=False, repr=False)
+    keyring_api: Callable[[], KeyringApi | KeyringNotInstalled | KeyringDisabled] = _keyring_api
 
     def read(self) -> SecretRead:
         if self.stopped_answering.is_set():
             return KeyringUnreachable()
-        api: Final = _keyring_api()
+        api: Final = self.keyring_api()
         if isinstance(api, (KeyringNotInstalled, KeyringDisabled)):
             return api
-        try:
-            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)
+        result: Final = _bounded_keyring_call(
+            lambda: api.get_password(KEYRING_SERVICE, KEYRING_ACCOUNT),
+            self.preflight_timeout_seconds,
+        )
+        match result:
+            case _KeyringCallAnswered(value=blob):
+                return SecretMissing() if blob is None else SecretFound(blob)
+            case _KeyringCallFailed():
+                return KeyringUnreachable()
+            case _KeyringCallTimedOut():
+                self.stopped_answering.set()
+                return KeyringUnreachable()
 
     def write(self, blob: str) -> SecretWrite:
         """Store the secret, reporting stored only once the keychain hands the same bytes back.
@@ -188,18 +262,40 @@ class KeyringVault:
         """
         if self.stopped_answering.is_set():
             return KeyringUnreachable()
-        api: Final = _keyring_api()
+        api: Final = self.keyring_api()
         if isinstance(api, (KeyringNotInstalled, KeyringDisabled)):
             return api
-        if not _answers_a_write(api, self.preflight_timeout_seconds):
+        probe: Final = _probe_a_write(api, self.preflight_timeout_seconds)
+        match probe:
+            case _ProbeStored() | _ProbeRefused():
+                pass
+            case _ProbeSilent():
+                self.stopped_answering.set()
+                return KeyringUnreachable()
+        if not _forget_the_preflight(api, self.preflight_timeout_seconds):
             self.stopped_answering.set()
             return KeyringUnreachable()
-        _forget_the_preflight(api)
-        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()
+        result: Final = _bounded_keyring_call(
+            lambda: api.set_password(KEYRING_SERVICE, KEYRING_ACCOUNT, blob),
+            self.preflight_timeout_seconds,
+        )
+        match result:
+            case _KeyringCallFailed():
+                match probe:
+                    case _ProbeStored():
+                        return SecretTooLarge() if _payload_exceeds_windows_limit(blob) else KeyringUnreachable()
+                    case _ProbeRefused():
+                        return KeyringUnreachable()
+            case _KeyringCallTimedOut():
+                self.stopped_answering.set()
+                return KeyringUnreachable()
+            case _KeyringCallAnswered():
+                read_back: Final = self.read()
+                if read_back == SecretFound(blob):
+                    return SecretStored()
+                if read_back == KeyringUnreachable():
+                    return KeyringUnreachable()
+                return KeyringDiscardsWrites()
 
     def erase(self) -> SecretErase:
         """Remove our entry, reporting whether the keychain is guaranteed to be free of it.
@@ -218,14 +314,25 @@ class KeyringVault:
                 return self._delete()
 
     def _delete(self) -> SecretErase:
-        api: Final = _keyring_api()
+        api: Final = self.keyring_api()
         if isinstance(api, (KeyringNotInstalled, KeyringDisabled)):
             return api
-        try:
-            api.delete_password(KEYRING_SERVICE, KEYRING_ACCOUNT)
-        except Exception:  # noqa: BLE001  # report the failure as a value so `lite logout` can warn
-            return SecretStranded()
-        return SecretErased()
+        result: Final = _bounded_keyring_call(
+            lambda: api.delete_password(KEYRING_SERVICE, KEYRING_ACCOUNT),
+            self.preflight_timeout_seconds,
+        )
+        match result:
+            case _KeyringCallAnswered():
+                return SecretErased()
+            case _KeyringCallFailed():
+                return SecretStranded()
+            case _KeyringCallTimedOut():
+                self.stopped_answering.set()
+                return SecretStranded()
+
+
+def _payload_exceeds_windows_limit(blob: str) -> bool:
+    return len(blob.encode("utf-16-le")) > _MAX_CREDENTIAL_BLOB_BYTES
 
 
 SYSTEM_KEYRING: Final[SecretVault] = KeyringVault()

litellm/proxy/client/README.md

@@ -378,7 +378,7 @@ The key itself, together with the refresh token that renews a `--pkce` credentia
 }
 ```
 
-Keychain storage needs the `keyring` package, which ships with `pip install 'litellm[cli]'`. Headless boxes and CI runners usually have no keychain either. In all of those cases the key and the refresh token stay in the same `0600` file alongside the metadata, exactly as they did before, and `lite login` names which one applies: the package is missing, the machine has no keychain, or you set `LITELLM_CLI_DISABLE_KEYRING=1` to force the file even where a keychain exists. A `token.json` written by an older `lite` keeps working and is moved into the keychain, and scrubbed from the file, the first time a keychain-capable `lite` reads it. That includes a refresh token left behind by the release that moved only the key.
+Keychain storage needs the `keyring` package, which ships with `pip install 'litellm[cli]'`. Headless boxes and CI runners usually have no keychain either. In all of those cases the key and the refresh token stay in the same `0600` file alongside the metadata, exactly as they did before, and `lite login` names which one applies: the package is missing, the machine has no keychain, the keychain is unreachable, the backend discards writes, the credential is too large for the backend, or you set `LITELLM_CLI_DISABLE_KEYRING=1` to force the file even where a keychain exists. A `token.json` written by an older `lite` keeps working and is moved into the keychain, and scrubbed from the file, the first time a keychain-capable `lite` reads it. That includes a refresh token left behind by the release that moved only the key.
 
 `lite logout` clears both stores. If the keychain is locked at that moment it says so, and re-running it once the keychain is unlocked finishes the job.
 

litellm/proxy/client/cli/commands/auth.py

@@ -24,6 +24,7 @@ from litellm.litellm_core_utils.cli_keyring import (
     SecretMissing,
     SecretStored,
     SecretStranded,
+    SecretTooLarge,
     SecretVault,
 )
 from litellm.litellm_core_utils.cli_token_utils import (
@@ -151,6 +152,8 @@ def storage_notice(outcome: SecretSave) -> str:
                 f"Your keyring backend keeps nothing it is given, so the credential was stored in {path} "
                 f"(owner-only) instead. For OS keychain storage, run: {KEYRING_ENABLE_HINT}"
             )
+        case SecretTooLarge():
+            return f"Your credential is too large for the OS keychain, so it was stored in {path} (owner-only) instead."
         case CredentialNotSaved(detail=detail):
             return (
                 f"Signed in, but the credential could not be saved to {path}: {detail}. "

tests/test_litellm/litellm_core_utils/test_cli_keyring.py

@@ -0,0 +1,151 @@
+import threading
+import time
+
+from litellm.litellm_core_utils.cli_keyring import (
+    KEYRING_ACCOUNT,
+    KEYRING_PREFLIGHT_ACCOUNT,
+    KeyringApi,
+    KeyringUnreachable,
+    KeyringVault,
+    SecretErased,
+    SecretFound,
+    SecretStranded,
+    SecretTooLarge,
+)
+
+
+class _BlockingReadKeyring:
+    def __init__(self) -> None:
+        self.blocked = threading.Event()
+        self.calls: list[tuple[str, str]] = []
+
+    def get_password(self, service_name: str, username: str) -> str | None:
+        self.calls.append(("get", username))
+        self.blocked.set()
+        threading.Event().wait()
+        return None
+
+    def set_password(self, service_name: str, username: str, password: str) -> None:
+        self.calls.append(("set", username))
+
+    def delete_password(self, service_name: str, username: str) -> None:
+        self.calls.append(("delete", username))
+
+
+class _BlockingDeleteKeyring:
+    def __init__(self) -> None:
+        self.blocked = threading.Event()
+        self.calls: list[tuple[str, str]] = []
+
+    def get_password(self, service_name: str, username: str) -> str | None:
+        self.calls.append(("get", username))
+        return "stored-blob"
+
+    def set_password(self, service_name: str, username: str, password: str) -> None:
+        self.calls.append(("set", username))
+
+    def delete_password(self, service_name: str, username: str) -> None:
+        self.calls.append(("delete", username))
+        self.blocked.set()
+        threading.Event().wait()
+
+
+class _OversizedWriteKeyring:
+    def __init__(self) -> None:
+        self.calls: list[tuple[str, str]] = []
+
+    def get_password(self, service_name: str, username: str) -> str | None:
+        self.calls.append(("get", username))
+        return None
+
+    def set_password(self, service_name: str, username: str, password: str) -> None:
+        self.calls.append(("set", username))
+        if username != KEYRING_PREFLIGHT_ACCOUNT:
+            raise ValueError("credential blob too large")
+
+    def delete_password(self, service_name: str, username: str) -> None:
+        self.calls.append(("delete", username))
+
+
+class _RefusingWriteKeyring:
+    def set_password(self, service_name: str, username: str, password: str) -> None:
+        raise ValueError("keychain refused write")
+
+    def get_password(self, service_name: str, username: str) -> str | None:
+        return None
+
+    def delete_password(self, service_name: str, username: str) -> None:
+        return None
+
+
+def _vault(api: KeyringApi, timeout_seconds: float = 0.02) -> KeyringVault:
+    return KeyringVault(
+        preflight_timeout_seconds=timeout_seconds,
+        keyring_api=lambda: api,
+    )
+
+
+def test_blocking_read_is_bounded_and_latches():
+    keyring = _BlockingReadKeyring()
+    vault = _vault(keyring)
+
+    started = time.monotonic()
+    outcome = vault.read()
+
+    assert outcome == KeyringUnreachable()
+    assert time.monotonic() - started < 1
+    assert keyring.blocked.is_set()
+    assert vault.read() == KeyringUnreachable()
+    assert keyring.calls == [("get", KEYRING_ACCOUNT)]
+
+
+def test_blocking_delete_is_bounded_and_stranded():
+    keyring = _BlockingDeleteKeyring()
+    vault = _vault(keyring)
+
+    started = time.monotonic()
+    outcome = vault.erase()
+
+    assert outcome == SecretStranded()
+    assert time.monotonic() - started < 1
+    assert keyring.blocked.is_set()
+    assert vault.read() == KeyringUnreachable()
+    assert keyring.calls == [("get", KEYRING_ACCOUNT), ("delete", KEYRING_ACCOUNT)]
+
+
+def test_successful_read_and_delete_use_the_injected_backend():
+    class _WorkingKeyring:
+        def get_password(self, service_name: str, username: str) -> str | None:
+            return "stored-blob"
+
+        def set_password(self, service_name: str, username: str, password: str) -> None:
+            return None
+
+        def delete_password(self, service_name: str, username: str) -> None:
+            return None
+
+    vault = _vault(_WorkingKeyring())
+
+    assert vault.read() == SecretFound("stored-blob")
+    assert vault.erase() == SecretErased()
+
+
+def test_oversized_write_is_classified_from_utf16_payload_size():
+    keyring = _OversizedWriteKeyring()
+    vault = _vault(keyring)
+
+    assert vault.write("a" * 1281) == SecretTooLarge()
+    assert ("set", KEYRING_PREFLIGHT_ACCOUNT) in keyring.calls
+    assert ("set", KEYRING_ACCOUNT) in keyring.calls
+
+
+def test_small_write_failure_remains_unreachable():
+    keyring = _OversizedWriteKeyring()
+
+    assert _vault(keyring).write("small") == KeyringUnreachable()
+
+
+def test_oversized_write_from_refusing_probe_remains_unreachable():
+    keyring = _RefusingWriteKeyring()
+
+    assert _vault(keyring).write("a" * 1281) == KeyringUnreachable()

tests/test_litellm/litellm_core_utils/test_cli_token_utils.py

@@ -25,6 +25,7 @@ from litellm.litellm_core_utils.cli_keyring import (
     SecretMissing,
     SecretStored,
     SecretStranded,
+    SecretTooLarge,
 )
 from litellm.litellm_core_utils.cli_token_utils import (
     CliTokenRecord,
@@ -473,6 +474,66 @@ class TestSaveCliToken:
         assert stat.S_IMODE(path.stat().st_mode) == 0o600
         assert list(path.parent.glob(".tmp-*")) == []
 
+    def test_an_oversized_keychain_credential_falls_back_to_a_private_file(self, isolated_home):
+        class _OversizedKeyring:
+            def get_password(self, service_name, username):
+                return None
+
+            def set_password(self, service_name, username, password):
+                if username != KEYRING_PREFLIGHT_ACCOUNT:
+                    raise ValueError("credential blob too large")
+
+            def delete_password(self, service_name, username):
+                return None
+
+        path = _token_file(isolated_home)
+        outcome = save_cli_token(
+            CliTokenRecord(base_url=SERVER, key="a" * 1281),
+            vault=KeyringVault(keyring_api=lambda: _OversizedKeyring()),
+        )
+
+        assert outcome == SecretTooLarge()
+        assert stat.S_IMODE(path.stat().st_mode) == 0o600
+        assert json.loads(path.read_text())["key"] == "a" * 1281
+
+    def test_a_blocking_keychain_read_does_not_hold_up_saving(self, isolated_home):
+        class _BlockingReadKeyring:
+            def get_password(self, service_name, username):
+                threading.Event().wait()
+                return None
+
+            def set_password(self, service_name, username, password):
+                return None
+
+            def delete_password(self, service_name, username):
+                return None
+
+        started = time.monotonic()
+        outcome = save_cli_token(
+            CliTokenRecord(base_url=SERVER, key="sk-new"),
+            vault=KeyringVault(
+                preflight_timeout_seconds=0.02,
+                keyring_api=lambda: _BlockingReadKeyring(),
+            ),
+        )
+
+        assert outcome == KeyringUnreachable()
+        assert time.monotonic() - started < 1
+        assert json.loads(_token_file(isolated_home).read_text())["key"] == "sk-new"
+
+    def test_keychain_backed_token_file_contains_no_secret_fields(self, isolated_home, secret_vault_factory):
+        path = _token_file(isolated_home)
+
+        assert save_cli_token(
+            CliTokenRecord(base_url=SERVER, key="sk-new", jwt_token="jwt-new", refresh_token="rt-new"),
+            vault=secret_vault_factory(),
+        ) == SecretStored()
+
+        token_file = json.loads(path.read_text())
+        assert token_file.get("key") is None
+        assert not token_file.get("jwt_token")
+        assert token_file.get("refresh_token") is None
+
     def test_creates_the_config_directory_owner_only(self, isolated_home, secret_vault_factory):
         """A 0755 ~/.litellm lets any local process list, and in the fallback case read, the
         credential's directory."""
@@ -749,6 +810,31 @@ class TestScrubFailure:
 
 
 class TestClearCliToken:
+    def test_a_blocking_keychain_delete_reports_a_stranded_credential(self, isolated_home):
+        _write_metadata_only_file(isolated_home)
+
+        class _BlockingDeleteKeyring:
+            def get_password(self, service_name, username):
+                return _blob()
+
+            def set_password(self, service_name, username, password):
+                return None
+
+            def delete_password(self, service_name, username):
+                threading.Event().wait()
+
+        started = time.monotonic()
+        outcome = clear_cli_token(
+            vault=KeyringVault(
+                preflight_timeout_seconds=0.02,
+                keyring_api=lambda: _BlockingDeleteKeyring(),
+            )
+        )
+
+        assert outcome == SecretStranded()
+        assert time.monotonic() - started < 1
+        assert not _token_file(isolated_home).exists()
+
     def test_removes_the_credential_from_both_stores(self, isolated_home, secret_vault_factory):
         vault = secret_vault_factory()
         save_cli_token(CliTokenRecord(base_url=SERVER, key="sk-new"), vault=vault)