From 63bfb7cfaea8354c9fc51a2408da8e87b32854e1 Mon Sep 17 00:00:00 2001 From: ciregenz Date: Tue, 7 Jul 2026 18:09:57 -0700 Subject: [PATCH] [eric] agents: admission gate caps concurrent root turns (children bypass, queued pill) --- backend/apps/agents/agent_manager.py | 57 ++++++- backend/tests/test_turn_admission.py | 145 ++++++++++++++++++ .../app/pages/Dashboard/cards/AgentCard.tsx | 2 +- frontend/src/shared/state/agentsSlice.ts | 11 ++ frontend/src/shared/ws/WebSocketManager.ts | 11 ++ 5 files changed, 219 insertions(+), 7 deletions(-) create mode 100644 backend/tests/test_turn_admission.py diff --git a/backend/apps/agents/agent_manager.py b/backend/apps/agents/agent_manager.py index e69b44ad..37875ba6 100644 --- a/backend/apps/agents/agent_manager.py +++ b/backend/apps/agents/agent_manager.py @@ -1,7 +1,8 @@ import asyncio import logging import os -from typing import Dict, List, Optional +from contextlib import asynccontextmanager +from typing import AsyncIterator, Dict, List, Optional from typeguard import typechecked from backend.apps.agents.core.models import ( @@ -33,6 +34,9 @@ logger = logging.getLogger(__name__) os.environ.setdefault("CLAUDE_CODE_STREAM_CLOSE_TIMEOUT", "3600000") +# Cap concurrent ROOT agent turns so firing 30 agents at once doesn't spawn 30 CLIs in the same instant; the overflow queues (agents are model/IO-bound, so they're waiting anyway). Env-tunable, 0/blank disables the gate. +MAX_CONCURRENT_TURNS = int(os.environ.get("OSW_MAX_CONCURRENT_TURNS", "8") or "0") + class AgentManager(SessionLifecycle, SessionPersistence, Messaging, SessionControl, AgentLaunch, MockAgent, TurnRunner, RunOptions, RunSupport): @typechecked @@ -48,8 +52,47 @@ class AgentManager(SessionLifecycle, SessionPersistence, Messaging, SessionContr # Per-SESSION hook context + stderr buffer, updated in place each turn: a persistent client's hooks/stderr callback were bound at connect, so they must read stable objects, not per-turn rebuilds. self.hook_ctxs: Dict[str, object] = {} self.stderr_buffers: Dict[str, List[str]] = {} + # Admission gate: one shared semaphore caps concurrent ROOT turns (children bypass). (Re)created per running loop by get_turn_admission so it never binds to a dead loop across a uvicorn reload or a test's asyncio.run. + self.p_turn_admission_sema: Optional[asyncio.Semaphore] = None + self.p_turn_admission_loop: Optional[asyncio.AbstractEventLoop] = None + @typechecked + def get_turn_admission(self) -> asyncio.Semaphore: + """The shared admission semaphore for the CURRENT loop; rebuilt if the loop changed so a + reload/test-run can never await a semaphore bound to a dead loop.""" + loop = asyncio.get_running_loop() + if self.p_turn_admission_sema is None or self.p_turn_admission_loop is not loop: + self.p_turn_admission_sema = asyncio.Semaphore(MAX_CONCURRENT_TURNS) + self.p_turn_admission_loop = loop + return self.p_turn_admission_sema + + @asynccontextmanager + async def turn_admission_slot(self, session: AgentSession, session_id: str) -> AsyncIterator[None]: + """Hold one concurrency slot for the duration of a ROOT turn. Overflow turns queue on the + semaphore (emitting agent:queued, then agent:admitted when they start). Two bypasses, both + load-bearing: (1) MAX_CONCURRENT_TURNS<=0 disables the gate entirely (kill switch); (2) a + CHILD turn (parent_session_id set) is NEVER gated, because a parent holds its own slot while + awaiting a delegated child, so gating children would deadlock the pool. `async with` release + is cancellation-safe: a stop while queued never acquired, so it can't over-release.""" + if MAX_CONCURRENT_TURNS <= 0 or session.parent_session_id is not None: + yield + return + sema = self.get_turn_admission() + was_queued = sema.locked() + if was_queued: + try: + await ws_manager.send_to_session(session_id, "agent:queued", {"session_id": session_id}) + except Exception: + pass + async with sema: + if was_queued: + try: + await ws_manager.send_to_session(session_id, "agent:admitted", {"session_id": session_id}) + except Exception: + pass + yield + @typechecked async def run_agent_loop(self, session_id: str, prompt: str, images: Optional[List] = None, context_paths: Optional[List] = None, forced_tools: Optional[List[str]] = None, attached_skills: Optional[List] = None, fork_session: bool = False, selected_browser_ids: Optional[List[str]] = None, selected_app_output_ids: Optional[List[str]] = None, selected_setting_ids: Optional[List[str]] = None, context_valve_retry: bool = False): """Run the Claude Agent SDK query loop for a session.""" @@ -99,11 +142,13 @@ class AgentManager(SessionLifecycle, SessionPersistence, Messaging, SessionContr api_type = p_api_type_for_session thinking = ThinkingState() - await self.run_turn_with_retry( - session, session_id, prompt_content, options, options_kwargs, - turn, thinking, p_stderr_buffer, resolved_model, api_type, global_settings, - force_respawn=p_force_respawn, - ) + # Gate the CLI turn (spawn + stream) behind the admission slot so a burst can't run every turn at once; the slot is held ONLY for run_turn_with_retry, so the context-valve retry below re-acquires cleanly instead of nesting. + async with self.turn_admission_slot(session, session_id): + await self.run_turn_with_retry( + session, session_id, prompt_content, options, options_kwargs, + turn, thinking, p_stderr_buffer, resolved_model, api_type, global_settings, + force_respawn=p_force_respawn, + ) session.status = "completed" # Auto-continuation hook (Phase 3). If MCPActivate (or any analogous flow) flagged pending_continuation during this turn, kick off a follow-up turn immediately with the captured prompt. We dispatch as a fire-and-forget task so the current run_agent_loop frame can unwind cleanly before the next turn's options + history rebuild kicks in. The follow-up is `hidden=True` so it doesn't add a user bubble to the visible chat; the model sees it as a synthetic prompt to keep working. diff --git a/backend/tests/test_turn_admission.py b/backend/tests/test_turn_admission.py new file mode 100644 index 00000000..5186ce2c --- /dev/null +++ b/backend/tests/test_turn_admission.py @@ -0,0 +1,145 @@ +"""Invariant tests for the root-turn admission gate (lever 3: cap concurrent agent turns so a burst +doesn't spawn every CLI at once). Proves the load-bearing properties: roots are bounded to the cap, +CHILD turns bypass it (else a parent awaiting a delegated child would deadlock the pool), the slot is +cancellation-safe (a stop while queued can't over-release), and the kill switch disables the gate.""" + +import asyncio +from typing import List, Optional, Tuple + +from backend.apps.agents.core.models import AgentSession +import backend.apps.agents.agent_manager as am + + +def make_session(parent: Optional[str] = None) -> AgentSession: + s = AgentSession(name="t", model="haiku", mode="agent") + s.parent_session_id = parent + return s + + +class Harness: + """Fresh manager + a fake ws recorder + a scoped MAX_CONCURRENT_TURNS override.""" + + def __init__(self, cap: int): + self.mgr = am.AgentManager() + self.events: List[Tuple[str, str]] = [] + self.cap = cap + self.p_old_max = am.MAX_CONCURRENT_TURNS + self.p_old_send = am.ws_manager.send_to_session + + async def p_fake_send(self, session_id: str, event: str, payload: dict) -> None: + self.events.append((session_id, event)) + + def __enter__(self) -> "Harness": + am.MAX_CONCURRENT_TURNS = self.cap + am.ws_manager.send_to_session = self.p_fake_send + return self + + def __exit__(self, *exc) -> None: + am.MAX_CONCURRENT_TURNS = self.p_old_max + am.ws_manager.send_to_session = self.p_old_send + + +def test_root_turns_bounded_to_cap(): + async def run(): + with Harness(cap=3) as h: + live = 0 + peak = 0 + + async def one_root(i: int) -> None: + nonlocal live, peak + async with h.mgr.turn_admission_slot(make_session(), f"s{i}"): + live += 1 + peak = max(peak, live) + await asyncio.sleep(0.02) + live -= 1 + + await asyncio.gather(*[one_root(i) for i in range(9)]) + assert peak == 3 # never more than the cap ran concurrently + queued = [e for e in h.events if e[1] == "agent:queued"] + admitted = [e for e in h.events if e[1] == "agent:admitted"] + assert len(queued) == 6 # 9 roots, cap 3 -> 6 waited + assert len(admitted) == 6 # every queued turn was later admitted + + asyncio.run(run()) + + +def test_children_bypass_gate(): + """The deadlock guard: with the only slot held by a root, a child must still run immediately.""" + async def run(): + with Harness(cap=1) as h: + async with h.mgr.turn_admission_slot(make_session(), "root"): + ran = False + async with h.mgr.turn_admission_slot(make_session(parent="root"), "child"): + ran = True + assert ran # child bypassed the full gate + assert not any(e[0] == "child" for e in h.events) # bypass emits no queue events + + asyncio.run(run()) + + +def test_kill_switch_disables_gate(): + async def run(): + with Harness(cap=0) as h: + live = 0 + peak = 0 + + async def one(i: int) -> None: + nonlocal live, peak + async with h.mgr.turn_admission_slot(make_session(), f"s{i}"): + live += 1 + peak = max(peak, live) + await asyncio.sleep(0.01) + live -= 1 + + await asyncio.gather(*[one(i) for i in range(5)]) + assert peak == 5 # cap<=0 -> no gating, all run at once + assert not h.events # nothing queued + + asyncio.run(run()) + + +def test_cancel_while_queued_no_overrelease(): + async def run(): + with Harness(cap=1) as h: + started = asyncio.Event() + release = asyncio.Event() + + async def holder() -> None: + async with h.mgr.turn_admission_slot(make_session(), "holder"): + started.set() + await release.wait() + + async def queued() -> None: + async with h.mgr.turn_admission_slot(make_session(), "q"): + pass + + ht = asyncio.create_task(holder()) + await started.wait() # holder owns the only slot + qt = asyncio.create_task(queued()) + await asyncio.sleep(0.01) # q is now blocked on acquire + qt.cancel() + try: + await qt + except asyncio.CancelledError: + pass + release.set() + await ht + sema = h.mgr.get_turn_admission() + assert not sema.locked() # holder's slot returned; queued-then-cancelled leaked nothing + await sema.acquire() + assert sema.locked() # capacity is exactly 1: an over-release would leave it acquirable twice + sema.release() + + asyncio.run(run()) + + +def test_semaphore_rebuilds_per_loop(): + """Never bind to a dead loop: a second event loop gets a fresh semaphore, not the first's.""" + mgr = am.AgentManager() + + async def get_sema(): + return mgr.get_turn_admission() + + s1 = asyncio.run(get_sema()) + s2 = asyncio.run(get_sema()) + assert s1 is not s2 diff --git a/frontend/src/app/pages/Dashboard/cards/AgentCard.tsx b/frontend/src/app/pages/Dashboard/cards/AgentCard.tsx index 7f585b43..693e898c 100644 --- a/frontend/src/app/pages/Dashboard/cards/AgentCard.tsx +++ b/frontend/src/app/pages/Dashboard/cards/AgentCard.tsx @@ -924,7 +924,7 @@ const AgentCard: React.FC = ({ {session.status !== 'completed' && session.status !== 'stopped' && !session.is_welcome_draft && ( - {friendlyStatusLabel(session.status)} + {session.queued && session.status === 'running' ? 'queued' : friendlyStatusLabel(session.status)} )} diff --git a/frontend/src/shared/state/agentsSlice.ts b/frontend/src/shared/state/agentsSlice.ts index 6f9d6cdd..18c72f5b 100644 --- a/frontend/src/shared/state/agentsSlice.ts +++ b/frontend/src/shared/state/agentsSlice.ts @@ -100,6 +100,8 @@ export interface AgentSession { /** Browser memory signals that drive the subtle "Remembered"/"Learned" card chip. */ memory_recalled?: boolean; memory_learned?: boolean; + /** Client-only: turn is waiting on the admission gate (agent:queued -> agent:admitted). Transient; self-clears on the next full-session status update. */ + queued?: boolean; thinking_level?: 'off' | 'low' | 'medium' | 'high' | 'auto'; active_mcps?: string[]; ctx_used_pct?: number; @@ -727,6 +729,8 @@ const agentsSlice = createSlice({ if (action.payload.status === 'running' && session.status !== 'running') { session.app_deps_changed = false; } + // Leaving the running state ends any admission-queue wait (the partial-payload path; the full-session path drops it by replacing the object). + if (action.payload.status !== 'running') session.queued = false; session.status = action.payload.status; } if (action.payload.status === 'running' && !state.trackedNotificationIds.includes(action.payload.sessionId)) { @@ -855,6 +859,12 @@ const agentsSlice = createSlice({ session.turn_label = null; }, + setQueued(state, action: PayloadAction<{ sessionId: string; queued: boolean }>) { + const session = state.sessions[action.payload.sessionId]; + if (!session) return; + session.queued = action.payload.queued; + }, + // streamStart/Delta/End live in streamingSlice; keeps sessions dict stable during streaming. addApprovalRequest( @@ -1455,6 +1465,7 @@ export const { recordCompaction, setTurnLabel, clearTurnLabel, + setQueued, addApprovalRequest, removeApprovalRequest, updateSessionCost, diff --git a/frontend/src/shared/ws/WebSocketManager.ts b/frontend/src/shared/ws/WebSocketManager.ts index 9771608f..b6d2f7c4 100644 --- a/frontend/src/shared/ws/WebSocketManager.ts +++ b/frontend/src/shared/ws/WebSocketManager.ts @@ -24,6 +24,7 @@ import { fetchSession, recordCompaction, setTurnLabel, + setQueued, clearTurnLabel, } from '../state/agentsSlice'; import { streamStart, streamDelta, streamEnd, clearStreamingForSession } from '../state/streamingSlice'; @@ -591,6 +592,16 @@ class WebSocketManager { } break; + case 'agent:queued': + // Admission gate: this turn is waiting for a concurrency slot; shows a "queued" chip so it doesn't read as a hung "working". + if (session_id) store.dispatch(setQueued({ sessionId: session_id, queued: true })); + break; + + case 'agent:admitted': + // Slot acquired; the turn is about to stream. Clears the queued chip. + if (session_id) store.dispatch(setQueued({ sessionId: session_id, queued: false })); + break; + case 'agent:auth_error': // Re-uses the context_overflow card slot, both are "this session is blocked, here's what to do" cards. Reason field disambiguates. if (session_id) {