fix: queue concurrent subagents (#5566)

* fix: queue concurrent subagents * chore: keep spawn schema concise
2026-09-03 01:31:47 +03:00 · 2026-08-27 17:53:28 +08:00
parent 39de4594d7
commit b9e7c7f6fe
6 changed files with 197 additions and 38 deletions
@@ -211,8 +211,8 @@ async def test_spawn_forwards_temperature_to_run_spec(tmp_path):


@pytest.mark.asyncio
-async def test_spawn_tool_rejects_when_at_concurrency_limit(tmp_path):
-    """SpawnTool should return an error string when the concurrency limit is reached."""
+async def test_background_spawn_waits_for_concurrency_capacity(tmp_path):
+    """Background tasks should be accepted and start when capacity becomes available."""
    from nanobot.agent.subagent import SubagentManager
    from nanobot.agent.tools.spawn import SpawnTool
    from nanobot.bus.queue import MessageBus
@@ -224,14 +224,23 @@ async def test_spawn_tool_rejects_when_at_concurrency_limit(tmp_path):
        workspace=tmp_path,
        bus=bus,
        max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
+        max_concurrent_subagents=1,
    )
    mgr._announce_result = AsyncMock()

-    # Block the first subagent so it stays "running"
-    release = asyncio.Event()
+    first_entered = asyncio.Event()
+    second_entered = asyncio.Event()
+    release_first = asyncio.Event()
+    release_second = asyncio.Event()

    async def fake_run(spec):
-        await release.wait()
+        task = spec.initial_messages[-1]["content"]
+        if task == "first task":
+            first_entered.set()
+            await release_first.wait()
+        else:
+            second_entered.set()
+            await release_second.wait()
        return SimpleNamespace(
            stop_reason="done",
            final_content="done",
@@ -250,19 +259,24 @@ async def test_spawn_tool_rejects_when_at_concurrency_limit(tmp_path):
        session_key="test:c1",
        runtime=_runtime(provider),
    )):
-        # First spawn succeeds
-        result = await tool.execute(task="first task")
-        assert "started" in result
+        first_result = await tool.execute(task="first task")
+        assert "started" in first_result
+        await asyncio.wait_for(first_entered.wait(), timeout=1.0)

-        # Second spawn should be rejected (default limit is 1)
-        result = await tool.execute(task="second task")
-        assert "Cannot spawn subagent" in result
-        assert "concurrency limit reached" in result
+        second_result = await tool.execute(task="second task")
+        assert "started" in second_result
+        tasks = list(mgr._running_tasks.values())
+        await asyncio.sleep(0)
+        assert not second_entered.is_set()
+        phases = {status.task_description: status.phase for status in mgr._task_statuses.values()}
+        assert phases == {"first task": "initializing", "second task": "queued"}

-    # Release the first subagent
-    release.set()
-    # Allow cleanup
-    await asyncio.gather(*mgr._running_tasks.values(), return_exceptions=True)
+    release_first.set()
+    await asyncio.wait_for(second_entered.wait(), timeout=1.0)
+    release_second.set()
+    await asyncio.gather(*tasks, return_exceptions=True)
+    await asyncio.sleep(0)
+    assert mgr._running_tasks == {}


@pytest.mark.asyncio
@@ -300,7 +314,7 @@ async def test_spawn_tool_waits_for_inline_result():


@pytest.mark.asyncio
-async def test_inline_spawn_counts_toward_concurrency_limit(tmp_path):
+async def test_inline_spawn_waits_for_concurrency_capacity(tmp_path):
    from nanobot.agent.subagent import SubagentManager
    from nanobot.agent.tools.context import RequestContext, request_context
    from nanobot.agent.tools.spawn import SpawnTool
@@ -312,12 +326,19 @@ async def test_inline_spawn_counts_toward_concurrency_limit(tmp_path):
        max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
        max_concurrent_subagents=1,
    )
-    release = asyncio.Event()
-    entered = asyncio.Event()
+    first_entered = asyncio.Event()
+    second_entered = asyncio.Event()
+    release_first = asyncio.Event()
+    release_second = asyncio.Event()

    async def fake_run(spec):
-        entered.set()
-        await release.wait()
+        task = spec.initial_messages[-1]["content"]
+        if task == "first":
+            first_entered.set()
+            await release_first.wait()
+        else:
+            second_entered.set()
+            await release_second.wait()
        return SimpleNamespace(
            stop_reason="done",
            final_content="done",
@@ -334,19 +355,99 @@ async def test_inline_spawn_counts_toward_concurrency_limit(tmp_path):
        runtime=_runtime(MagicMock()),
    )):
        first = asyncio.create_task(tool.execute(task="first", wait=True))
-        await asyncio.wait_for(entered.wait(), timeout=1.0)
+        await asyncio.wait_for(first_entered.wait(), timeout=1.0)

-        second = await tool.execute(task="second", wait=True)
+        second = asyncio.create_task(tool.execute(task="second", wait=True))
+        await asyncio.sleep(0)

-        assert "concurrency limit reached" in second
-        assert manager.get_running_count() == 1
-        release.set()
+        assert not second.done()
+        assert not second_entered.is_set()
+        assert manager.get_running_count() == 2
+        release_first.set()
        assert await first == "done"
+        await asyncio.wait_for(second_entered.wait(), timeout=1.0)
+        release_second.set()
+        assert await second == "done"

    assert manager.get_running_count() == 0
    assert manager._session_tasks == {}


+@pytest.mark.asyncio
+async def test_runner_executes_inline_spawn_batch_concurrently(tmp_path):
+    """Adjacent blocking consultations should share one concurrent tool batch."""
+    from nanobot.agent.runner import AgentRunner, AgentRunSpec
+    from nanobot.agent.subagent import SubagentManager
+    from nanobot.agent.tools.context import RequestContext, request_context
+    from nanobot.agent.tools.registry import ToolRegistry
+    from nanobot.agent.tools.spawn import SpawnTool
+    from nanobot.bus.queue import MessageBus
+    from nanobot.providers.base import ToolCallRequest
+
+    manager = SubagentManager(
+        workspace=tmp_path,
+        bus=MessageBus(),
+        max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
+        max_concurrent_subagents=2,
+    )
+    both_entered = asyncio.Event()
+    release = asyncio.Event()
+    entered: list[str] = []
+
+    async def fake_run(spec):
+        entered.append(spec.initial_messages[-1]["content"])
+        if len(entered) == 2:
+            both_entered.set()
+        await release.wait()
+        return SimpleNamespace(
+            stop_reason="done",
+            final_content=spec.initial_messages[-1]["content"],
+            error=None,
+            tool_events=[],
+        )
+
+    manager.runner.run = AsyncMock(side_effect=fake_run)
+    tools = ToolRegistry()
+    tools.register(SpawnTool(manager))
+    runtime = _runtime(MagicMock())
+    spec = AgentRunSpec(
+        initial_messages=[],
+        tools=tools,
+        runtime=runtime,
+        max_iterations=1,
+        max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
+        concurrent_tools=True,
+    )
+    calls = [
+        ToolCallRequest(
+            id="spawn-1",
+            name="spawn",
+            arguments={"task": "first", "wait": True},
+        ),
+        ToolCallRequest(
+            id="spawn-2",
+            name="spawn",
+            arguments={"task": "second", "wait": True},
+        ),
+    ]
+
+    with request_context(RequestContext(
+        channel="test",
+        chat_id="c1",
+        session_key="test:c1",
+        runtime=runtime,
+    )):
+        execution = asyncio.create_task(AgentRunner()._execute_tools(spec, calls, {}, {}))
+        await asyncio.wait_for(both_entered.wait(), timeout=1.0)
+        release.set()
+        results, events = await execution
+
+    assert set(entered) == {"first", "second"}
+    assert results == ["first", "second"]
+    assert [event["status"] for event in events] == ["ok", "ok"]
+    assert manager._running_tasks == {}
+
+
@pytest.mark.asyncio
 async def test_cancel_by_session_cancels_inline_subagent(tmp_path):
    from nanobot.agent.subagent import SubagentManager
@@ -391,6 +492,7 @@ def test_subagent_default_max_concurrent_matches_agent_defaults(tmp_path):
        max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
    )

+    assert AgentDefaults().max_concurrent_subagents == 4
    assert mgr.max_concurrent_subagents == AgentDefaults().max_concurrent_subagents