diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index 2520c09a535..b0bbd28f372 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -869,8 +869,23 @@ def _replace_stored_token( ) raise click.ClickException("The login did not select the requested team; your saved login has not changed") previous: Final = load_token(vault=vault) + previous_secret: Final = vault.read() stored: Final = save_token(record, vault=vault) - if previous is None or isinstance(stored, CredentialNotSaved): + if previous is not None and isinstance(stored, CredentialNotRecorded): + restored: Final = ( + vault.write(previous_secret.blob) if isinstance(previous_secret, SecretFound) else previous_secret + ) + if isinstance(restored, SecretStored): + abandoned_revocation: Final = revoke_stored_credential(record, http) + if abandoned_revocation is not None: + click.echo( + "Could not revoke the abandoned login's refresh token " + f"on the proxy ({abandoned_revocation.reason}); " + "it expires on its own." + ) + return CredentialNotSaved("The replacement could not be recorded; your previous login was restored") + click.echo("Could not restore the previous login after the partial save; sign in again to repair it.") + if previous is None or isinstance(stored, (CredentialNotSaved, CredentialNotRecorded)): return stored revocation: Final = revoke_stored_credential(previous, http) if revocation is not None: diff --git a/tests/test_litellm/proxy/client/cli/test_auth_commands.py b/tests/test_litellm/proxy/client/cli/test_auth_commands.py index 3e922e6c8e7..a2fe8e25603 100644 --- a/tests/test_litellm/proxy/client/cli/test_auth_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_auth_commands.py @@ -2274,6 +2274,60 @@ def test_assign_key_keeps_explicit_virtual_key_update_with_saved_session(isolate assert load_token() == before +@pytest.mark.parametrize("restore_fails", [False, True]) +def test_assign_key_handles_keychain_update_without_metadata(isolated_home, secret_vault_factory, restore_fails): + vault = secret_vault_factory() + save_cli_token( + CliTokenRecord( + **_pkce_record(key="session-old", expires_at=time.time() + 3600, team_id="team-a", client_id="old-client") + ), + vault=vault, + ) + before = load_token(vault=vault) + metadata_before = (isolated_home / ".litellm" / "token.json").read_bytes() + _FakeSession.instances.clear() + + def refuse_metadata(*_): + if restore_fails: + vault.writable = False + raise OSError("read-only") + + with ( + patch("litellm.proxy.client.cli.commands.teams.Client"), + patch( + "litellm.proxy.client.cli.commands.auth.run_pkce_login", + return_value=replace(_pkce_credential(), access_token="session-new", team_id="team-b"), + ), + patch("litellm.litellm_core_utils.cli_token_utils.commit_staged_json", side_effect=refuse_metadata), + patch("litellm.proxy.client.cli.commands.auth.requests.Session", _FakeSession), + ): + result = CliRunner().invoke( + cli, + ["--base-url", PKCE_BASE_URL, "teams", "assign-key", "--team-id", "team-b"], + obj={"secret_vault": vault}, + ) + + assert result.exit_code != 0, result.output + assert (isolated_home / ".litellm" / "token.json").read_bytes() == metadata_before + if restore_fails: + assert load_token(vault=vault)["key"] == "session-new" + assert "Could not restore the previous login" in result.output + else: + assert load_token(vault=vault) == before + assert "previous login was restored" in result.output + assert "Successfully assigned" not in result.output + assert [post for session in _FakeSession.instances for post in session.posts] == ( + [] + if restore_fails + else [ + ( + f"{PKCE_BASE_URL}/revoke", + {"token": "llm_srefresh_fresh", "token_type_hint": "refresh_token", "client_id": "llm_dcrc_abc"}, + ) + ] + ) + + def test_assign_key_warns_when_a_rejected_login_cannot_be_revoked(isolated_home, monkeypatch): monkeypatch.setenv(DISABLE_KEYRING_ENV_VAR, "1") save_cli_token(CliTokenRecord(**_pkce_record(key="session-old", expires_at=time.time() + 3600)))