scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 8

fix: split rollback into separate DB and memory try blocks for accurate error reporting

Browse source
This commit is contained in:
Harshit28j 2026-03-16 19:52:27 +05:30
parent 69f5784875
commit 2eff336815
2 changed files with 50 additions and 24 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

66
litellm/proxy/guardrails/guardrail_endpoints.py
View file

@ -462,29 +462,38 @@ async def update_guardrail(
f"Immediate sync: Failed to update '{guardrail_name}' (ID: {guardrail_id}) in memory: {update_error}"
)
# Rollback: restore previous guardrail data in DB
rollback_ok = False
db_rollback_ok = False
mem_rollback_ok = False
try:
await GUARDRAIL_REGISTRY.update_guardrail_in_db(
guardrail_id=guardrail_id,
guardrail=cast(Guardrail, existing_guardrail),
prisma_client=prisma_client,
)
# Re-initialize the old guardrail in memory
IN_MEMORY_GUARDRAIL_HANDLER.update_in_memory_guardrail(
guardrail_id=guardrail_id,
guardrail=cast(Guardrail, existing_guardrail),
)
rollback_ok = True
db_rollback_ok = True
except Exception:
verbose_proxy_logger.error(
"Failed to rollback guardrail DB entry after update failure"
)
status = (
"rolled back"
if rollback_ok
else "rollback also failed, DB/memory may be inconsistent"
)
if db_rollback_ok:
try:
IN_MEMORY_GUARDRAIL_HANDLER.update_in_memory_guardrail(
guardrail_id=guardrail_id,
guardrail=cast(Guardrail, existing_guardrail),
)
mem_rollback_ok = True
except Exception:
verbose_proxy_logger.error(
"Failed to rollback guardrail in memory after update failure"
)
if db_rollback_ok and mem_rollback_ok:
status = "rolled back"
elif db_rollback_ok:
status = "DB rolled back but memory rollback failed"
else:
status = "rollback also failed, DB/memory may be inconsistent"
raise HTTPException(
status_code=422,
detail=f"Guardrail update failed, {status}: {update_error}",
@ -1043,7 +1052,7 @@ async def reject_guardrail_submission(
"/guardrails/{guardrail_id}",
tags=["Guardrails"],
)
async def patch_guardrail(
async def patch_guardrail( # noqa: PLR0915
guardrail_id: str,
request: PatchGuardrailRequest,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
@ -1169,28 +1178,37 @@ async def patch_guardrail(
f"Immediate sync: Failed to update '{guardrail_name}' (ID: {guardrail_id}) in memory: {update_error}"
)
# Rollback: restore previous guardrail data in DB
rollback_ok = False
db_rollback_ok = False
mem_rollback_ok = False
try:
await GUARDRAIL_REGISTRY.update_guardrail_in_db(
guardrail_id=guardrail_id,
guardrail=cast(Guardrail, existing_guardrail),
prisma_client=prisma_client,
)
# Re-sync the old guardrail in memory
IN_MEMORY_GUARDRAIL_HANDLER.sync_guardrail_from_db(
guardrail=cast(Guardrail, existing_guardrail),
)
rollback_ok = True
db_rollback_ok = True
except Exception:
verbose_proxy_logger.error(
"Failed to rollback guardrail DB entry after patch failure"
)
status = (
"rolled back"
if rollback_ok
else "rollback also failed, DB/memory may be inconsistent"
)
if db_rollback_ok:
try:
IN_MEMORY_GUARDRAIL_HANDLER.sync_guardrail_from_db(
guardrail=cast(Guardrail, existing_guardrail),
)
mem_rollback_ok = True
except Exception:
verbose_proxy_logger.error(
"Failed to rollback guardrail in memory after patch failure"
)
if db_rollback_ok and mem_rollback_ok:
status = "rolled back"
elif db_rollback_ok:
status = "DB rolled back but memory rollback failed"
else:
status = "rollback also failed, DB/memory may be inconsistent"
raise HTTPException(
status_code=422,
detail=f"Guardrail patch failed, {status}: {update_error}",

8
tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py
View file

@ -838,6 +838,10 @@ async def test_update_guardrail_endpoint(
elif scenario == "success_sync_fails":
assert exc_info.value.status_code == 422
assert "Guardrail update failed" in str(exc_info.value.detail)
# DB rollback succeeds, memory rollback fails → accurate message
assert "DB rolled back but memory rollback failed" in str(
exc_info.value.detail
)
# Verify rollback was attempted: update_in_db called twice (initial + rollback)
assert mock_guardrail_registry.update_guardrail_in_db.call_count == 2
# Verify rollback attempted in-memory re-init with old config
@ -947,6 +951,10 @@ async def test_patch_guardrail_endpoint(
if scenario == "success_sync_fails":
assert exc_info.value.status_code == 422
assert "Guardrail patch failed" in str(exc_info.value.detail)
# DB rollback succeeds, memory rollback fails → accurate message
assert "DB rolled back but memory rollback failed" in str(
exc_info.value.detail
)
# Verify rollback was attempted: update_in_db called twice (initial + rollback)
assert mock_guardrail_registry.update_guardrail_in_db.call_count == 2
# Verify rollback attempted in-memory re-sync with old config