mirror of
https://github.com/openswarm-ai/openswarm.git
synced 2026-08-17 18:25:42 +02:00
305 lines
12 KiB
Python
305 lines
12 KiB
Python
"""Invariant + seeded-simulation tests for the persistent-client pool (lever A of the TTFT work).
|
|
Proves the red-teamed safety properties hold by construction: fingerprint-gated reuse, pop-first
|
|
disposal, never-raising teardown, idle/LRU reclaim, and (seeded sim) that random op sequences never
|
|
reuse a stale client, never double-boot needlessly, and always recover a dead one. What the
|
|
fingerprint itself hashes lives in test_boot_fingerprint.py."""
|
|
|
|
import asyncio
|
|
import random
|
|
from typing import Dict, List
|
|
|
|
from backend.apps.agents.manager.run.client_pool import (
|
|
ClientHandle,
|
|
acquire_client,
|
|
dispose_all_clients,
|
|
dispose_client,
|
|
dispose_client_soon,
|
|
start_pool_sweeper,
|
|
stop_pool_sweeper,
|
|
trim_pool_to_cap,
|
|
)
|
|
|
|
|
|
class FakeClient:
|
|
"""Stands in for ClaudeSDKClient: counts connects/disconnects, can be killed, can raise on disconnect."""
|
|
|
|
def __init__(self, registry: List["FakeClient"], raise_on_disconnect: bool = False):
|
|
self.alive = True
|
|
self.disconnected = False
|
|
self.raise_on_disconnect = raise_on_disconnect
|
|
registry.append(self)
|
|
|
|
async def disconnect(self):
|
|
self.disconnected = True
|
|
self.alive = False
|
|
if self.raise_on_disconnect:
|
|
raise RuntimeError("teardown boom")
|
|
|
|
|
|
def test_reuse_respawn_force_and_teardown():
|
|
async def run():
|
|
pool: Dict[str, ClientHandle] = {}
|
|
made: List[FakeClient] = []
|
|
|
|
async def connect():
|
|
return FakeClient(made)
|
|
|
|
h1 = await acquire_client(pool, "s1", "fpA", connect)
|
|
h2 = await acquire_client(pool, "s1", "fpA", connect)
|
|
assert h1 is h2 and len(made) == 1
|
|
|
|
h3 = await acquire_client(pool, "s1", "fpB", connect)
|
|
assert h3 is not h1 and len(made) == 2 and made[0].disconnected
|
|
|
|
h4 = await acquire_client(pool, "s1", "fpB", connect, force_respawn=True)
|
|
assert h4 is not h3 and len(made) == 3 and made[1].disconnected
|
|
|
|
await dispose_client(pool, "s1")
|
|
assert "s1" not in pool and made[2].disconnected
|
|
await dispose_client(pool, "s1") # idempotent
|
|
|
|
async def connect_bad():
|
|
return FakeClient(made, raise_on_disconnect=True)
|
|
|
|
await acquire_client(pool, "s2", "fp", connect_bad)
|
|
await dispose_client(pool, "s2") # teardown error swallowed
|
|
assert "s2" not in pool
|
|
|
|
await acquire_client(pool, "s3", "fp", connect)
|
|
dispose_client_soon(pool, "s3")
|
|
assert "s3" not in pool # pop is sync-first
|
|
await asyncio.sleep(0.01)
|
|
assert made[-1].disconnected
|
|
|
|
await acquire_client(pool, "s4", "fp", connect)
|
|
await acquire_client(pool, "s5", "fp", connect)
|
|
await dispose_all_clients(pool)
|
|
assert not pool and all(c.disconnected for c in made)
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
def test_idle_eviction():
|
|
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_ttl = cp.IDLE_EVICT_SECONDS
|
|
cp.IDLE_EVICT_SECONDS = 0.05
|
|
try:
|
|
h = await acquire_client(pool, "s1", "fp", connect)
|
|
await acquire_client(pool, "s2", "fp", connect)
|
|
await asyncio.sleep(0.1)
|
|
# s1 is mid-turn (lock held): the sweep must skip it and evict only the idle s2.
|
|
async with h.lock:
|
|
await cp.evict_idle_clients(pool)
|
|
assert "s1" in pool and "s2" not in pool and made[1].disconnected
|
|
await asyncio.sleep(0.1)
|
|
await cp.evict_idle_clients(pool)
|
|
assert "s1" not in pool and made[0].disconnected
|
|
# a fresh acquire after eviction reconnects transparently
|
|
h2 = await acquire_client(pool, "s1", "fp", connect)
|
|
assert h2.client.alive
|
|
finally:
|
|
cp.IDLE_EVICT_SECONDS = old_ttl
|
|
|
|
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)."""
|
|
async def run():
|
|
rng = random.Random(1337)
|
|
pool: Dict[str, ClientHandle] = {}
|
|
made: List[FakeClient] = []
|
|
boots = 0
|
|
turns = 0
|
|
fp = "fp0"
|
|
force = False
|
|
|
|
async def connect():
|
|
nonlocal boots
|
|
boots += 1
|
|
return FakeClient(made)
|
|
|
|
for _ in range(300):
|
|
op = rng.choice(["follow_up", "activate", "branch_or_fresh", "kill", "close"])
|
|
if op == "follow_up":
|
|
turns += 1
|
|
h = await acquire_client(pool, "sim", fp, connect, force_respawn=force)
|
|
force = False
|
|
assert h.fingerprint == fp and not h.client.disconnected
|
|
if not h.client.alive: # dead client detected by the turn -> dispose + one respawn
|
|
await dispose_client(pool, "sim")
|
|
h = await acquire_client(pool, "sim", fp, connect)
|
|
assert h.client.alive
|
|
async with h.lock:
|
|
assert h.lock.locked() # single consumer while a turn drains
|
|
h.turns_served += 1
|
|
elif op == "activate":
|
|
fp = f"fp{rng.randint(0, 10**9)}" # mcp_servers grew -> fingerprint changed
|
|
elif op == "branch_or_fresh":
|
|
force = True # needs_fresh/fork read pre-build forces respawn
|
|
elif op == "kill" and "sim" in pool:
|
|
pool["sim"].client.alive = False
|
|
elif op == "close":
|
|
await dispose_client(pool, "sim")
|
|
|
|
assert boots <= turns, f"persistent booted {boots}x for {turns} turns; one-shot baseline is {turns}"
|
|
live = [c for c in made if not c.disconnected]
|
|
assert len(live) <= 1, "at most the pooled client may be alive; everything else must be torn down"
|
|
if "sim" in pool:
|
|
assert not pool["sim"].client.disconnected
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
def test_concurrent_acquires_share_one_spawn() -> None:
|
|
"""A pre-warm and a racing first turn must never double-spawn: the second spawn used to leak the first CLI."""
|
|
import asyncio
|
|
|
|
from backend.apps.agents.manager.run.client_pool import acquire_client
|
|
|
|
async def run() -> None:
|
|
pool: dict = {}
|
|
spawns = 0
|
|
|
|
class FakeClient:
|
|
async def disconnect(self) -> None:
|
|
return None
|
|
|
|
async def connect_fn():
|
|
nonlocal spawns
|
|
spawns += 1
|
|
await asyncio.sleep(0.05)
|
|
return FakeClient()
|
|
|
|
a, b = await asyncio.gather(
|
|
acquire_client(pool, "sess-race", "fp1", connect_fn),
|
|
acquire_client(pool, "sess-race", "fp1", connect_fn),
|
|
)
|
|
assert spawns == 1, f"double spawn: {spawns}"
|
|
assert a is b
|
|
assert pool["sess-race"].client is a.client
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
def test_cancelled_waiter_does_not_kill_the_shared_spawn() -> None:
|
|
import asyncio
|
|
|
|
from backend.apps.agents.manager.run.client_pool import acquire_client
|
|
|
|
async def run() -> None:
|
|
pool: dict = {}
|
|
|
|
class FakeClient:
|
|
async def disconnect(self) -> None:
|
|
return None
|
|
|
|
async def connect_fn():
|
|
await asyncio.sleep(0.08)
|
|
return FakeClient()
|
|
|
|
first = asyncio.ensure_future(acquire_client(pool, "sess-cancel", "fp1", connect_fn))
|
|
await asyncio.sleep(0.01)
|
|
second = asyncio.ensure_future(acquire_client(pool, "sess-cancel", "fp1", connect_fn))
|
|
await asyncio.sleep(0.01)
|
|
second.cancel()
|
|
handle = await first
|
|
assert pool["sess-cancel"].client is handle.client
|
|
|
|
asyncio.run(run())
|