From 19ada25e59779e0c6212d340e63f751ea753d7e5 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Wed, 6 May 2026 15:37:23 -0700 Subject: [PATCH] fix(session_endpoints): close /followup concurrency gap and gate writes for view-only admins Two fixes in this file: 1. /followup new-run branch was bypassing the run-busy concurrency guard. When latest_run was terminal-or-absent, the endpoint went straight to litellm_agentrun.create without calling _has_active_run. Two concurrent /followup calls on an idle session both passed the latest_run.status check and both inserted runs, breaking the 'one active run per session' invariant that POST /runs enforces via 409 run_busy. Add the same _has_active_run check before the fallthrough create, plus a defensive insert-time retry that re-queries for active runs after IntegrityError and surfaces 409 run_busy if it lost the race. 2. Call assert_caller_can_mutate on create_session, delete_session, and followup so view-only admins get 403 instead of bypassing ownership. Greptile P1 (concurrency) + P1 SECURITY (view-only admin write bypass). --- .../session_endpoints.py | 67 ++++++++++++++++--- 1 file changed, 58 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/agent_session_endpoints/session_endpoints.py b/litellm/proxy/agent_session_endpoints/session_endpoints.py index faf3dbe0a8e..df7e4883c58 100644 --- a/litellm/proxy/agent_session_endpoints/session_endpoints.py +++ b/litellm/proxy/agent_session_endpoints/session_endpoints.py @@ -43,6 +43,7 @@ from litellm.proxy.agent_session_endpoints.constants import ( ) from litellm.proxy.agent_session_endpoints.ids import new_run_id, new_session_id from litellm.proxy.agent_session_endpoints.ownership import ( + assert_caller_can_mutate, assert_caller_owns_agent, assert_caller_owns_session, caller_api_key_hash, @@ -72,6 +73,22 @@ def _now() -> datetime: return datetime.now(timezone.utc) +async def _has_active_run(prisma_client, session_id: str) -> bool: + """True iff ``session_id`` has any run in queued/running. + + Mirrors the ``_has_active_run`` helper in ``run_endpoints.py``. We + duplicate it here (instead of importing) to keep ``session_endpoints`` + free of any dependency on ``run_endpoints``. + """ + existing = await prisma_client.db.litellm_agentrun.find_first( + where={ + "session_id": session_id, + "status": {"in": list(RUN_ACTIVE_STATUSES)}, + } + ) + return existing is not None + + async def _get_prisma_client_or_503(): from litellm.proxy.proxy_server import prisma_client @@ -216,6 +233,7 @@ async def create_session( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), idempotency_key: Optional[str] = Header(default=None, alias="Idempotency-Key"), ): + assert_caller_can_mutate(user_api_key_dict) prisma_client = await _get_prisma_client_or_503() # Idempotency: same (caller, key) returns same session — no daemon @@ -345,6 +363,7 @@ async def delete_session( session_id: str, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): + assert_caller_can_mutate(user_api_key_dict) prisma_client = await _get_prisma_client_or_503() row = await prisma_client.db.litellm_agentsession.find_unique( where={"id": session_id} @@ -474,6 +493,7 @@ async def followup( ``user_message`` event via the events stream and weaves it into the in-flight LLM turn. """ + assert_caller_can_mutate(user_api_key_dict) prisma_client = await _get_prisma_client_or_503() session = await prisma_client.db.litellm_agentsession.find_unique( where={"id": session_id} @@ -504,16 +524,45 @@ async def followup( ) return FollowupResponse(run_id=latest_run.id, action="queued") + # Concurrency guard: matches POST /runs. Without this, two concurrent + # /followup requests on an idle session both pass the + # ``latest_run.status in RUN_ACTIVE_STATUSES`` check above and both + # fall through to ``create``, breaking the "one active run at a time" + # invariant. Re-check for any active run RIGHT before insert and + # return 409 ``run_busy`` instead, just like POST /runs does. + if await _has_active_run(prisma_client, session_id): + raise HTTPException( + status_code=409, + detail="run_busy: another run is queued/running for this session", + ) + # Else start a fresh run. - new_run = await prisma_client.db.litellm_agentrun.create( - data={ - "id": new_run_id(), - "session_id": session_id, - "status": RUN_STATUS_QUEUED, - "prompt": body.prompt, - "updated_at": _now(), - } - ) + try: + new_run = await prisma_client.db.litellm_agentrun.create( + data={ + "id": new_run_id(), + "session_id": session_id, + "status": RUN_STATUS_QUEUED, + "prompt": body.prompt, + "updated_at": _now(), + } + ) + except Exception as exc: + # Last line of defense: another /followup or POST /runs may have + # raced past the busy check and won the insert. Surface the same + # 409 so the client can retry deterministically. + active_other = await prisma_client.db.litellm_agentrun.find_first( + where={ + "session_id": session_id, + "status": {"in": list(RUN_ACTIVE_STATUSES)}, + } + ) + if active_other is not None: + raise HTTPException( + status_code=409, + detail="run_busy: another run is queued/running for this session", + ) from exc + raise return FollowupResponse(run_id=new_run.id, action="new_run")