scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-07 02:59:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

fix(proxy): reinit env-configured hashicorp vault manager after cyberark persist rollback

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yassin 2026-08-27 02:33:17 +00:00
parent 054ac7b3de
commit 4084d25ffa
2 changed files with 53 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

16
litellm/proxy/management_endpoints/config_override_endpoints.py
View file

@ -274,11 +274,17 @@ def _restore_cyberark_runtime(proxy_config: "ProxyConfig", env_values: Mapping[s
verbose_proxy_logger.exception("Failed to restore previous CyberArk configuration")
else:
return
if litellm._key_management_system == KeyManagementSystem.CYBERARK: # pyright: ignore[reportPrivateUsage] # proxy-internal helper, mirrors hashicorp endpoint usage
litellm.secret_manager_client = None
litellm._key_management_system = None # pyright: ignore[reportPrivateUsage] # proxy-internal helper, mirrors hashicorp endpoint usage
# Force the vault reload to re-init from its own row so no manager is stranded inactive
proxy_config._last_hashicorp_vault_config = None # pyright: ignore[reportPrivateUsage] # proxy-internal change-detection cache
if litellm._key_management_system != KeyManagementSystem.CYBERARK: # pyright: ignore[reportPrivateUsage] # proxy-internal helper, mirrors hashicorp endpoint usage
return
litellm.secret_manager_client = None
litellm._key_management_system = None # pyright: ignore[reportPrivateUsage] # proxy-internal helper, mirrors hashicorp endpoint usage
# Force the vault reload to re-init from its own row so no manager is stranded inactive
proxy_config._last_hashicorp_vault_config = None # pyright: ignore[reportPrivateUsage] # proxy-internal change-detection cache
if os.environ.get("HCP_VAULT_ADDR"):
try:
proxy_config.initialize_secret_manager(key_management_system="hashicorp_vault")
except Exception: # noqa: BLE001 # restore is best-effort; the vault reload loop retries from its own row
verbose_proxy_logger.exception("Failed to reinitialize Hashicorp Vault after CyberArk rollback")
def _clear_cyberark_state(proxy_config: "ProxyConfig") -> None:

42
tests/test_litellm/proxy/management_endpoints/test_config_override_endpoints.py
View file

@ -529,6 +529,48 @@ async def test_cyberark_persist_failure_rolls_back_runtime_state(client, monkeyp
_cleanup()
@pytest.mark.asyncio
async def test_cyberark_persist_failure_restores_hashicorp_manager(client, monkeypatch):
"""If CyberArk init displaced an env-configured Hashicorp manager and the DB
upsert then fails, rollback must bring the Hashicorp manager back."""
mock_prisma, mock_db = _make_mock_db()
mock_cfg = _make_mock_proxy_config()
mock_cfg._last_cyberark_config = None
mock_db.upsert = AsyncMock(side_effect=Exception("db write failed"))
def _fake_init(key_management_system):
litellm._key_management_system = ( # test-quality-ok: endpoint hot-reloads litellm globals; test must set and restore them
KeyManagementSystem.CYBERARK
if key_management_system == "cyberark"
else KeyManagementSystem.HASHICORP_VAULT
)
mock_cfg.initialize_secret_manager = MagicMock(side_effect=_fake_init)
monkeypatch.setattr(ps, "prisma_client", mock_prisma)
monkeypatch.setattr(ps, "proxy_config", mock_cfg)
old_client, old_kms = litellm.secret_manager_client, litellm._key_management_system
_set_admin()
try:
monkeypatch.setenv("HCP_VAULT_ADDR", "https://vault.example.com")
litellm._key_management_system = KeyManagementSystem.HASHICORP_VAULT # test-quality-ok: endpoint hot-reloads litellm globals; test must set and restore them
r = client.post(
CYBERARK_URL,
json={"cyberark_api_base": "https://conjur.new.com", "cyberark_api_key": "new-key"},
)
assert r.status_code == 500
assert litellm._key_management_system == KeyManagementSystem.HASHICORP_VAULT
assert (
mock_cfg.initialize_secret_manager.call_args_list[-1].kwargs["key_management_system"] == "hashicorp_vault"
)
finally:
litellm.secret_manager_client = old_client # test-quality-ok: endpoint hot-reloads litellm globals; test must set and restore them
litellm._key_management_system = old_kms # test-quality-ok: endpoint hot-reloads litellm globals; test must set and restore them
os.environ.pop("HCP_VAULT_ADDR", None)
_cleanup()
@pytest.mark.asyncio
async def test_cyberark_audit_log_redacts_values(client, monkeypatch):
monkeypatch.setattr(litellm, "store_audit_logs", True)