From a1a461eb578bdae1b56f703d1bef8c62604ac3f3 Mon Sep 17 00:00:00 2001 From: ciregenz Date: Tue, 7 Jul 2026 17:33:35 -0700 Subject: [PATCH] [eric] agents: cap persistent client pool + background sweeper (bound resident memory) --- backend/apps/agents/agents.py | 5 +- backend/apps/agents/manager/run/TurnRunner.py | 2 + .../apps/agents/manager/run/client_pool.py | 66 +++++++++++++- backend/tests/test_client_pool.py | 90 +++++++++++++++++++ 4 files changed, 159 insertions(+), 4 deletions(-) diff --git a/backend/apps/agents/agents.py b/backend/apps/agents/agents.py index 8f70b0c7..6ac4deae 100644 --- a/backend/apps/agents/agents.py +++ b/backend/apps/agents/agents.py @@ -24,13 +24,16 @@ async def agents_lifespan(): logger.info("Agents sub-app starting") await agent_manager.reconcile_on_startup() await agent_manager.restore_all_sessions() + from backend.apps.agents.manager.run.client_pool import start_pool_sweeper, stop_pool_sweeper, dispose_all_clients + pool_sweeper = start_pool_sweeper(agent_manager.client_pool) yield logger.info("Agents sub-app shutting down") for session_id in list(agent_manager.tasks.keys()): await agent_manager.stop_agent(session_id) await agent_manager.persist_all_sessions() + # Cancel the sweeper before disposing so a background sweep can't race the shutdown teardown. + await stop_pool_sweeper(pool_sweeper) # Persistent CLI clients outlive turns; without this a uvicorn reload/quit orphans one subprocess per live session. - from backend.apps.agents.manager.run.client_pool import dispose_all_clients await dispose_all_clients(agent_manager.client_pool) agents = SubApp("agents", agents_lifespan) diff --git a/backend/apps/agents/manager/run/TurnRunner.py b/backend/apps/agents/manager/run/TurnRunner.py index 6acecd96..1167f547 100644 --- a/backend/apps/agents/manager/run/TurnRunner.py +++ b/backend/apps/agents/manager/run/TurnRunner.py @@ -134,6 +134,8 @@ class TurnRunner(AgentManagerProtocol): try: await handle.client.query(prompt_stream()) await p_run_streaming_turn(p_stream=handle.client.receive_response()) + # LRU by turn-END so a session mid-long-turn isn't first cap-evicted the instant it finishes. + handle.last_used = time.monotonic() except BaseException: # Fail-safe: an error or stop mid-turn poisons the live conversation; drop the client so the next attempt/turn reconnects fresh (== today's one-shot behavior, never worse). Pool pop is sync-first, so even a cancelled disconnect can't leave a reusable stale handle. await dispose_client(self.client_pool, session_id) diff --git a/backend/apps/agents/manager/run/client_pool.py b/backend/apps/agents/manager/run/client_pool.py index e895451b..368a94d5 100644 --- a/backend/apps/agents/manager/run/client_pool.py +++ b/backend/apps/agents/manager/run/client_pool.py @@ -1,6 +1,6 @@ -"""Per-session persistent SDK client pool (lever A of the TTFT work, gated by -OSW_TTFT_PERSISTENT_CLIENT=1, default OFF). One live Claude CLI per session, reused across -follow-up turns so the ~0.5s subprocess + MCP boot is paid once, not per message. +"""Per-session persistent SDK client pool (lever A of the TTFT work, default ON, kill switch +OPENSWARM_PERSISTENT_CLIENT=0). One live Claude CLI per session, reused across follow-up turns so +the ~0.5s subprocess + MCP boot is paid once, not per message. Safety model, from the red-teamed plan: reuse is gated on a BOOT FINGERPRINT (a hash of every boot-frozen input), never on session flags. Any change to the booted config (MCPActivate growing @@ -72,6 +72,13 @@ class ClientHandle(BaseModel): # A pooled CLI holds ~100MB+ per session; evict clients idle past this so parked chats don't accumulate subprocesses (respawn on the next message is the normal cold path). IDLE_EVICT_SECONDS = float(os.environ.get("OSW_CLIENT_IDLE_EVICT_SECONDS", "1800")) +# Hard ceiling on warm CLIs regardless of idle age: past this, the least-recently-used IDLE sessions are disposed (they respawn ~0.5s on their next message), bounding the "30 chats open" resident-memory case. Kept a SOFT cap: a mid-turn or just-acquired client is never evicted, so a burst of live turns may exceed it rather than kill work. +MAX_LIVE_CLIENTS = int(os.environ.get("OSW_CLIENT_MAX_LIVE", "12")) +# Never cap-evict a client used this recently; far larger than the acquire->lock window, so a just-acquired client can't be reaped before its turn takes the lock. +LRU_GUARD_SECONDS = float(os.environ.get("OSW_CLIENT_LRU_GUARD_SECONDS", "5")) +# Timer cadence for the background reclaim; the acquire-time sweep is lazy (fires only when some session takes a turn), this one catches an all-quiet pool. +SWEEP_INTERVAL_SECONDS = float(os.environ.get("OSW_CLIENT_SWEEP_INTERVAL_SECONDS", "60")) + @typechecked async def evict_idle_clients(pool: Dict[str, "ClientHandle"]) -> None: @@ -86,6 +93,24 @@ async def evict_idle_clients(pool: Dict[str, "ClientHandle"]) -> None: await dispose_client(pool, sid) +@typechecked +async def trim_pool_to_cap(pool: Dict[str, "ClientHandle"]) -> None: + """Dispose least-recently-used IDLE clients until the pool is back under MAX_LIVE_CLIENTS. Soft + cap: rechecks lock + recency immediately before each dispose (the check->pop is await-free), so a + client that went mid-turn or was just re-acquired is skipped and the pool temporarily exceeds the + cap rather than killing live work.""" + if len(pool) <= MAX_LIVE_CLIENTS: + return + for sid, _ in sorted(pool.items(), key=lambda kv: kv[1].last_used): + if len(pool) <= MAX_LIVE_CLIENTS: + break + handle = pool.get(sid) + if handle is None or handle.lock.locked() or time.monotonic() - handle.last_used <= LRU_GUARD_SECONDS: + continue + logger.info(f"[client-pool] {sid}: cap-evict (pool {len(pool)} > {MAX_LIVE_CLIENTS})") + await dispose_client(pool, sid) + + @typechecked async def acquire_client( pool: Dict[str, ClientHandle], @@ -113,6 +138,7 @@ async def acquire_client( ) pool[session_id] = handle logger.info(f"[client-pool] {session_id}: connected fresh client") + await trim_pool_to_cap(pool) return handle @@ -154,3 +180,37 @@ async def dispose_all_clients(pool: Dict[str, ClientHandle]) -> None: orphan one CLI per live session without this.""" for sid in list(pool.keys()): await dispose_client(pool, sid) + + +@typechecked +async def p_pool_sweeper_loop(pool: Dict[str, ClientHandle]) -> None: + """Timer-driven reclaim: runs the idle-TTL sweep AND the cap trim so a pool that went all-quiet + frees its subprocesses instead of holding them until the next turn's lazy acquire-time sweep.""" + while True: + try: + await asyncio.sleep(SWEEP_INTERVAL_SECONDS) + await evict_idle_clients(pool) + await trim_pool_to_cap(pool) + except asyncio.CancelledError: + break + except Exception: + logger.exception("[client-pool] sweeper iteration failed") + + +@typechecked +def start_pool_sweeper(pool: Dict[str, ClientHandle]) -> asyncio.Task: + """Launch the background reclaim loop (call from within the running loop); hold the returned task + and pass it to stop_pool_sweeper on shutdown.""" + return asyncio.get_running_loop().create_task(p_pool_sweeper_loop(pool)) + + +@typechecked +async def stop_pool_sweeper(task: Optional[asyncio.Task]) -> None: + """Cancel + await the sweeper. Call BEFORE dispose_all_clients so a sweep can't race teardown.""" + if task is None: + return + task.cancel() + try: + await task + except asyncio.CancelledError: + pass diff --git a/backend/tests/test_client_pool.py b/backend/tests/test_client_pool.py index 2722e5be..4e0f3e14 100644 --- a/backend/tests/test_client_pool.py +++ b/backend/tests/test_client_pool.py @@ -17,6 +17,9 @@ from backend.apps.agents.manager.run.client_pool import ( dispose_all_clients, dispose_client, dispose_client_soon, + start_pool_sweeper, + stop_pool_sweeper, + trim_pool_to_cap, ) @@ -160,6 +163,93 @@ def test_idle_eviction(): asyncio.run(run()) +def test_cap_lru_eviction(): + """Over MAX_LIVE_CLIENTS, acquire trims the least-recently-used IDLE sessions and keeps the newest.""" + async def run(): + import backend.apps.agents.manager.run.client_pool as cp + pool: Dict[str, ClientHandle] = {} + made: List[FakeClient] = [] + + async def connect(): + return FakeClient(made) + + old_max, old_guard = cp.MAX_LIVE_CLIENTS, cp.LRU_GUARD_SECONDS + cp.MAX_LIVE_CLIENTS, cp.LRU_GUARD_SECONDS = 3, 0.0 + try: + for i in range(5): + await acquire_client(pool, f"s{i}", "fp", connect) + await asyncio.sleep(0.001) # distinct last_used so LRU order is deterministic + assert len(pool) == 3 + assert "s0" not in pool and "s1" not in pool # two oldest reaped + assert {"s2", "s3", "s4"} <= set(pool) + assert made[0].disconnected and made[1].disconnected + assert not made[3].disconnected and not made[4].disconnected + finally: + cp.MAX_LIVE_CLIENTS, cp.LRU_GUARD_SECONDS = old_max, old_guard + + asyncio.run(run()) + + +def test_cap_soft_exceeds_when_busy(): + """A cap can't evict mid-turn clients: a new acquire over the cap exceeds it rather than kill a + live turn, then trims back once they go idle.""" + async def run(): + import backend.apps.agents.manager.run.client_pool as cp + pool: Dict[str, ClientHandle] = {} + made: List[FakeClient] = [] + + async def connect(): + return FakeClient(made) + + old_max, old_guard = cp.MAX_LIVE_CLIENTS, cp.LRU_GUARD_SECONDS + cp.MAX_LIVE_CLIENTS, cp.LRU_GUARD_SECONDS = 2, 0.4 + try: + h0 = await acquire_client(pool, "s0", "fp", connect) + h1 = await acquire_client(pool, "s1", "fp", connect) + async with h0.lock, h1.lock: + await acquire_client(pool, "s2", "fp", connect) + # s0/s1 locked, s2 just-acquired (guard-protected): nothing is eligible, so the pool exceeds the cap. + assert len(pool) == 3 + assert not made[0].disconnected and not made[1].disconnected + await asyncio.sleep(0.5) # past the guard: the now-idle sessions become eligible + await trim_pool_to_cap(pool) + assert len(pool) == 2 and made[0].disconnected # oldest idle reaped back to cap + finally: + cp.MAX_LIVE_CLIENTS, cp.LRU_GUARD_SECONDS = old_max, old_guard + + asyncio.run(run()) + + +def test_pool_sweeper_reclaims_over_cap(): + """The background sweeper trims an over-cap pool on its timer, with no new turn to trigger it.""" + async def run(): + import backend.apps.agents.manager.run.client_pool as cp + pool: Dict[str, ClientHandle] = {} + made: List[FakeClient] = [] + + async def connect(): + return FakeClient(made) + + old_max, old_guard, old_int = cp.MAX_LIVE_CLIENTS, cp.LRU_GUARD_SECONDS, cp.SWEEP_INTERVAL_SECONDS + cp.MAX_LIVE_CLIENTS, cp.LRU_GUARD_SECONDS, cp.SWEEP_INTERVAL_SECONDS = 10, 0.0, 0.02 + try: + for i in range(5): + await acquire_client(pool, f"s{i}", "fp", connect) + await asyncio.sleep(0.001) + assert len(pool) == 5 # under the temporary high cap + cp.MAX_LIVE_CLIENTS = 3 + task = start_pool_sweeper(pool) + await asyncio.sleep(0.1) # several sweep cycles + await stop_pool_sweeper(task) + assert len(pool) == 3 + assert "s0" not in pool and "s1" not in pool + await stop_pool_sweeper(None) # None is a no-op, must not raise + finally: + cp.MAX_LIVE_CLIENTS, cp.LRU_GUARD_SECONDS, cp.SWEEP_INTERVAL_SECONDS = old_max, old_guard, old_int + + asyncio.run(run()) + + def test_seeded_simulation_invariants(): """Random op sequences: reuse only on identical fingerprint, dead clients always replaced, pool never re-serves a disposed client, and boots never exceed the one-shot baseline (one per turn)."""