mirror of
https://github.com/usestrix/strix.git
synced 2026-09-09 22:31:07 +00:00
fix(runtime): make child agent limit reservation atomic
This commit is contained in:
parent
077a72d560
commit
799ef24842
3 changed files with 87 additions and 14 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue