diff --git a/litellm/litellm_core_utils/cli_token_utils.py b/litellm/litellm_core_utils/cli_token_utils.py index a55eb4522d2..3ea369767cf 100644 --- a/litellm/litellm_core_utils/cli_token_utils.py +++ b/litellm/litellm_core_utils/cli_token_utils.py @@ -250,12 +250,19 @@ def clear_cli_token(*, vault: SecretVault = SYSTEM_KEYRING) -> SecretClear: record: Final = _read_token_file() settled: Final = _nothing_left_behind(outcome, record) left_on_disk: Final = _end_the_local_login(record) - _note_what_the_keychain_answered(outcome, record) + unnoted: Final = _note_what_the_keychain_answered(outcome, record) + if unnoted is not None: + _put_the_token_file_back(unnoted) if left_on_disk is not None: return left_on_disk return SecretErased() if settled else outcome +def _put_the_token_file_back(record: CliTokenRecord) -> None: + with suppress(OSError): + _write_token_file(record) + + def _end_the_local_login(record: CliTokenRecord | None) -> CredentialNotCleared | None: removal: Final = _remove_token_file() if removal is None or record is None or _scrub_file_secret(record): @@ -311,7 +318,7 @@ def _read_keychain_not_cleared_note() -> CliTokenRecord | None: return None -def _note_what_the_keychain_answered(outcome: SecretErase, record: CliTokenRecord | None) -> None: +def _note_what_the_keychain_answered(outcome: SecretErase, record: CliTokenRecord | None) -> CliTokenRecord | None: """Keep the note that a logout left the keychain uncleared for exactly as long as that is true. A keychain that could not be read kept this logout from revoking the login, so the note holds @@ -321,28 +328,38 @@ def _note_what_the_keychain_answered(outcome: SecretErase, record: CliTokenRecor kept its secret in the file never reached the keychain, so the note an earlier logout left still names the login the keychain holds and stays, and a logout with no file to end keeps it too. - The secret is off disk by the time this runs, so a note that will not be written or removed is - no reason to fail the logout. The warning printed now is then the only one the user gets. + Only a note that cannot take an unrevoked login's record keeps that login from being revoked + later, so that record comes back for the token file to keep, as it did before the note existed. + Any other note that will not be written or removed is no reason to fail the logout, and the + warning printed now is then the only one the user gets. """ note: Final = _keychain_not_cleared_note_path() - with suppress(OSError): - match outcome: - case SecretErased(): + match outcome: + case SecretErased(): + with suppress(OSError): note.unlink(missing_ok=True) - case SecretStranded(): + return None + case SecretStranded(): + with suppress(OSError): _write_keychain_not_cleared_note(note, None) - case KeyringDisabled() | KeyringNotInstalled() | KeyringUnreachable(): - _note_a_keychain_left_unchecked(note, record) + return None + case KeyringDisabled() | KeyringNotInstalled() | KeyringUnreachable(): + return _note_a_keychain_left_unchecked(note, record) -def _note_a_keychain_left_unchecked(note: Path, record: CliTokenRecord | None) -> None: +def _note_a_keychain_left_unchecked(note: Path, record: CliTokenRecord | None) -> CliTokenRecord | None: if record is None: - return + return None if not _holds_a_secret(record): - _write_keychain_not_cleared_note(note, record) - return - if not note.exists(): - _write_keychain_not_cleared_note(note, None) + try: + _write_keychain_not_cleared_note(note, record) + except OSError: + return record + return None + with suppress(OSError): + if not note.exists(): + _write_keychain_not_cleared_note(note, None) + return None def _write_keychain_not_cleared_note(note: Path, record: CliTokenRecord | None) -> None: diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index eb010808238..ccead6c8366 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -1017,7 +1017,21 @@ def login(ctx: click.Context, config_claude: bool, pkce: bool | None) -> None: return -def _revoke_before_logout(token_data: Mapping[str, object] | None, http: Http) -> None: +def _revoke_before_logout(left: Mapping[str, object] | None, current: Mapping[str, object] | None, http: Http) -> None: + """Revoke every login this logout ends, holding the logout back only for the one it signs out. + + An earlier login left in the keychain is the previous login once another one is current, so a + proxy that cannot revoke it only earns the warning `lite login` gives, never a logout that the + current proxy is ready to finish. + """ + if left is None or current is None or _same_refresh_token(left, current): + _revoke_the_login_or_keep_it(current if left is None else left, http) + return + _revoke_the_replaced_login(left, http, click.echo) + _revoke_the_login_or_keep_it(current, http) + + +def _revoke_the_login_or_keep_it(token_data: Mapping[str, object] | None, http: Http) -> None: revocation: Final = None if token_data is None else revoke_stored_credential(token_data, http) match revocation: case RevocationUnavailable(reason=reason): @@ -1040,9 +1054,8 @@ def logout(ctx: click.Context): vault: Final = context_secret_vault(ctx) http: Final = requests.Session() left: Final = _login_left_in_keychain(vault) - _revoke_before_logout(left, http) current: Final = load_token(vault=vault) - _revoke_before_logout(None if _same_refresh_token(left, current) else current, http) + _revoke_before_logout(left, current, http) path: Final = get_cli_token_file_path() match clear_cli_token(vault=vault): diff --git a/tests/unit/proxy/client/cli/test_auth_commands.py b/tests/unit/proxy/client/cli/test_auth_commands.py index 156999f37ce..8737bd2e69d 100644 --- a/tests/unit/proxy/client/cli/test_auth_commands.py +++ b/tests/unit/proxy/client/cli/test_auth_commands.py @@ -1971,6 +1971,28 @@ class TestPkceLogoutCommand: ) ] + def test_a_logout_that_cannot_note_the_keychain_login_leaves_it_for_the_next_logout_to_revoke( + self, isolated_home, secret_vault_factory + ): + unlocked = secret_vault_factory() + save_cli_token(CliTokenRecord(**_pkce_record()), vault=unlocked) + locked = secret_vault_factory(blob=unlocked.blob, available=False) + (isolated_home / ".litellm" / "keychain_not_cleared").mkdir() + + with patch("litellm.proxy.client.cli.commands.auth.requests.Session", _FakeSession): + locked_logout = self.runner.invoke(logout, obj={"secret_vault": locked}) + result = self.runner.invoke(logout, obj={"secret_vault": unlocked}) + + assert locked_logout.exit_code == 0, locked_logout.output + assert result.output == "Logged out successfully. Authentication token cleared.\n" + assert unlocked.blob is None + assert [post for session in _FakeSession.instances for post in session.posts] == [ + ( + f"{PKCE_BASE_URL}/revoke", + {"token": "llm_srefresh_old", "token_type_hint": "refresh_token", "client_id": "llm_dcrc_abc"}, + ) + ] + def test_a_login_after_a_logout_that_missed_the_keychain_revokes_the_login_it_replaces( self, isolated_home, secret_vault_factory ): @@ -2050,6 +2072,62 @@ class TestPkceLogoutCommand: for token in ("llm_srefresh_old", "llm_srefresh_fresh") ] + def test_a_proxy_that_cannot_revoke_the_earlier_login_never_holds_back_the_current_logout( + self, isolated_home, secret_vault_factory + ): + class _EarlierProxyUnavailableSession(_FakeSession): + def post(self, url, *, data=None, json=None, timeout, allow_redirects): + self.posts.append((url, data)) + if data["token"] == "llm_srefresh_old": + return _FakeHttpResponse(503, {"error": "temporarily_unavailable"}) + return _FakeHttpResponse(200, {}) + + unlocked = secret_vault_factory() + save_cli_token(CliTokenRecord(**_pkce_record()), vault=unlocked) + locked = secret_vault_factory(blob=unlocked.blob, available=False) + + with ( + patch("litellm.proxy.client.cli.commands.auth.requests.Session", _EarlierProxyUnavailableSession), + patch("litellm.proxy.client.cli.commands.auth.run_pkce_login", return_value=_pkce_credential()), + patch("litellm.proxy.client.cli.interface.show_commands"), + ): + self.runner.invoke(cli, ["logout"], obj={"secret_vault": locked}) + self.runner.invoke(cli, ["--base-url", PKCE_BASE_URL, "login", "--pkce"], obj={"secret_vault": locked}) + result = self.runner.invoke(cli, ["logout"], obj={"secret_vault": unlocked}) + + assert result.exit_code == 0, result.output + assert ( + "Could not revoke the previous login's refresh token on the proxy (revocation failed with 503" + in result.output + ) + assert result.output.endswith("Logged out successfully. Authentication token cleared.\n") + assert unlocked.blob is None + assert list((isolated_home / ".litellm").iterdir()) == [] + assert [post["token"] for session in _FakeSession.instances for _, post in session.posts] == [ + "llm_srefresh_old", + "llm_srefresh_fresh", + ] + + def test_a_logout_whose_only_login_is_the_one_left_in_the_keychain_waits_for_its_revocation( + self, isolated_home, secret_vault_factory + ): + class _UnavailableSession(_FakeSession): + def __init__(self): + super().__init__() + self.response = _FakeHttpResponse(503, {"error": "temporarily_unavailable"}) + + unlocked = secret_vault_factory() + save_cli_token(CliTokenRecord(**_pkce_record()), vault=unlocked) + locked = secret_vault_factory(blob=unlocked.blob, available=False) + + with patch("litellm.proxy.client.cli.commands.auth.requests.Session", _UnavailableSession): + self.runner.invoke(cli, ["logout"], obj={"secret_vault": locked}) + result = self.runner.invoke(cli, ["logout"], obj={"secret_vault": unlocked}) + + assert result.exit_code == 1 + assert "Nothing was cleared; run `lite logout` again shortly." in result.output + assert unlocked.blob is not None + def test_a_login_after_a_locked_keychain_login_revokes_the_login_the_keychain_kept( self, isolated_home, secret_vault_factory ):