scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-09 03:18:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 109

address greptile review feedback (greploop iteration 2)

Browse source 
- Fix TOCTOU race in report_task_result using update_many with status guard
- Fix test fake agent_id filter to match production SQL semantics

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Krrish Dholakia 2026-05-01 12:33:16 -07:00
parent 933f38362c
commit 621d74a478
2 changed files with 27 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

31
litellm/proxy/scheduled_tasks/store.py
View file

@ -174,21 +174,36 @@ async def report_task_result(
return None
if result == "success":
return await prisma_client.db.litellm_scheduledtasktable.update(
where={"task_id": task_id},
affected = await prisma_client.db.litellm_scheduledtasktable.update_many(
where={"task_id": task_id, "owner_token": owner_token},
data={"consecutive_errors": 0, "last_error": None},
)
if affected == 0:
return None
return await prisma_client.db.litellm_scheduledtasktable.find_first(
where={"task_id": task_id, "owner_token": owner_token},
)
new_count = (existing.consecutive_errors or 0) + 1
update: Dict[str, Any] = {
update_data: Dict[str, Any] = {
"consecutive_errors": new_count,
"last_error": reason,
}
if new_count >= MAX_CONSECUTIVE_ERRORS and existing.status == "pending":
update["status"] = "failed"
return await prisma_client.db.litellm_scheduledtasktable.update(
where={"task_id": task_id},
data=update,
where_clause: Dict[str, Any] = {
"task_id": task_id,
"owner_token": owner_token,
}
if new_count >= MAX_CONSECUTIVE_ERRORS:
update_data["status"] = "failed"
where_clause["status"] = "pending"
affected = await prisma_client.db.litellm_scheduledtasktable.update_many(
where=where_clause,
data=update_data,
)
if affected == 0:
return None
return await prisma_client.db.litellm_scheduledtasktable.find_first(
where={"task_id": task_id, "owner_token": owner_token},
)

5
tests/test_litellm/proxy/scheduled_tasks/test_endpoints.py
View file

@ -205,7 +205,10 @@ class _FakeTx:
continue
if r.owner_token != owner_token:
continue
if agent_id is not None and r.agent_id != agent_id:
if agent_id is None:
if r.agent_id is not None:
continue
elif r.agent_id != agent_id:
continue
if actions is not None and r.action not in actions:
continue