fix(runtime): make child agent limit reservation atomic

This commit is contained in:
Ousama Ben Younes 2026-07-27 19:45:51 +00:00
parent 077a72d560
commit 799ef24842
3 changed files with 87 additions and 14 deletions

View file

@ -168,6 +168,34 @@ class AgentCoordinator:
logger.info("agent.register %s (%s) parent=%s", agent_id, name, parent_id or "-")
await self._maybe_snapshot()
async def register_child_if_capacity(
self,
agent_id: str,
name: str,
parent_id: str,
*,
max_child_agents: int,
task: str | None = None,
skills: list[str] | None = None,
) -> tuple[bool, int]:
async with self._lock:
child_count = sum(parent is not None for parent in self.parent_of.values())
if max_child_agents > 0 and child_count >= max_child_agents:
return False, child_count
self.statuses[agent_id] = "running"
self.parent_of[agent_id] = parent_id
self.names[agent_id] = name
self.pending_counts.setdefault(agent_id, 0)
self.metadata[agent_id] = {
"task": task or "",
"skills": list(skills or []),
}
self.runtimes.setdefault(agent_id, AgentRuntime())
child_count += 1
logger.info("agent.register %s (%s) parent=%s", agent_id, name, parent_id)
await self._maybe_snapshot()
return True, child_count
async def attach_runtime(
self,
agent_id: str,

View file

@ -310,8 +310,16 @@ async def spawn_child_agent(
if not isinstance(parent_id, str):
raise TypeError("Parent agent_id missing from context")
child_count = _child_agent_count(coordinator)
if max_child_agents > _UNLIMITED_CHILD_AGENTS and child_count >= max_child_agents:
child_id = uuid.uuid4().hex[:8]
registered, child_count = await coordinator.register_child_if_capacity(
child_id,
name,
parent_id,
max_child_agents=max_child_agents,
task=task,
skills=skills,
)
if not registered:
logger.info(
"refusing to spawn child agent %r: limit %d already reached",
name,
@ -328,15 +336,7 @@ async def spawn_child_agent(
"current_child_agents": child_count,
}
child_id = uuid.uuid4().hex[:8]
child_agent = factory(name=name, skills=skills)
await coordinator.register(
child_id,
name,
parent_id,
task=task,
skills=skills,
)
await _start_child_runner(
parent_ctx=parent_ctx,
@ -371,10 +371,6 @@ async def spawn_child_agent(
}
def _child_agent_count(coordinator: AgentCoordinator) -> int:
return sum(parent_id is not None for parent_id in coordinator.parent_of.values())
async def respawn_subagents(
*,
coordinator: AgentCoordinator,

View file

@ -168,6 +168,55 @@ async def test_spawn_child_agent_respects_child_limit(
assert set(coordinator.parent_of) == {"root", "child-1"}
@pytest.mark.asyncio
async def test_concurrent_spawn_child_agent_reserves_limit_atomically(
monkeypatch: pytest.MonkeyPatch, tmp_path: Any
) -> None:
coordinator = AgentCoordinator()
await coordinator.register("root", "strix", parent_id=None)
async def _noop_start(**_kwargs: Any) -> None:
return None
monkeypatch.setattr("strix.core.execution._start_child_runner", _noop_start)
def _factory(**_kwargs: Any) -> object:
return object()
async def _spawn(name: str) -> dict[str, Any]:
return await spawn_child_agent(
coordinator=coordinator,
factory=_factory,
agents_db_path=tmp_path / "agents.db",
sessions_to_close=[],
run_config=cast("Any", object()),
max_turns=1,
interactive=False,
parent_ctx={"agent_id": "root"},
name=name,
task="do more recon",
skills=[],
parent_history=[],
max_child_agents=1,
)
async with coordinator._lock:
tasks = [
asyncio.create_task(_spawn("extra-a")),
asyncio.create_task(_spawn("extra-b")),
]
await asyncio.sleep(0)
results = await asyncio.gather(*tasks)
assert [result["success"] for result in results].count(True) == 1
assert _child_count(coordinator) == 1
def _child_count(coordinator: AgentCoordinator) -> int:
return sum(parent_id is not None for parent_id in coordinator.parent_of.values())
@pytest.mark.asyncio
async def test_concurrent_reserve_claims_yield_single_root() -> None:
coordinator = AgentCoordinator()