"""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