mirror of
https://github.com/HKUDS/nanobot.git
synced 2026-04-30 14:56:01 +00:00
42c4af2118
When the main agent spawns multiple sub-agents, each completion independently triggered a new _dispatch, causing 3-4 user-visible responses instead of a single comprehensive report. - Extend _drain_pending to block-wait on pending_queue when sub-agents are still running, keeping the runner loop alive for in-order injection - Pass pending_queue in the system message path so subsequent sub-agent results can still be injected mid-turn via a new dispatch
261 lines
8.2 KiB
Python
261 lines
8.2 KiB
Python
"""Tests for subagent tool registration and wiring."""
|
|
|
|
import asyncio
|
|
import time
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from nanobot.config.schema import AgentDefaults
|
|
|
|
_MAX_TOOL_RESULT_CHARS = AgentDefaults().max_tool_result_chars
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_subagent_exec_tool_receives_allowed_env_keys(tmp_path):
|
|
"""allowed_env_keys from ExecToolConfig must be forwarded to the subagent's ExecTool."""
|
|
from nanobot.agent.subagent import SubagentManager, SubagentStatus
|
|
from nanobot.bus.queue import MessageBus
|
|
from nanobot.config.schema import ExecToolConfig
|
|
|
|
bus = MessageBus()
|
|
provider = MagicMock()
|
|
provider.get_default_model.return_value = "test-model"
|
|
mgr = SubagentManager(
|
|
provider=provider,
|
|
workspace=tmp_path,
|
|
bus=bus,
|
|
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
|
|
exec_config=ExecToolConfig(allowed_env_keys=["GOPATH", "JAVA_HOME"]),
|
|
)
|
|
mgr._announce_result = AsyncMock()
|
|
|
|
async def fake_run(spec):
|
|
exec_tool = spec.tools.get("exec")
|
|
assert exec_tool is not None
|
|
assert exec_tool.allowed_env_keys == ["GOPATH", "JAVA_HOME"]
|
|
return SimpleNamespace(
|
|
stop_reason="done",
|
|
final_content="done",
|
|
error=None,
|
|
tool_events=[],
|
|
)
|
|
|
|
mgr.runner.run = AsyncMock(side_effect=fake_run)
|
|
|
|
status = SubagentStatus(
|
|
task_id="sub-1", label="label", task_description="do task", started_at=time.monotonic()
|
|
)
|
|
await mgr._run_subagent(
|
|
"sub-1", "do task", "label", {"channel": "test", "chat_id": "c1"}, status
|
|
)
|
|
|
|
mgr.runner.run.assert_awaited_once()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_drain_pending_blocks_while_subagents_running(tmp_path):
|
|
"""_drain_pending should block when no messages are available but sub-agents are still running."""
|
|
from nanobot.agent.loop import AgentLoop
|
|
from nanobot.agent.subagent import SubagentManager
|
|
from nanobot.bus.events import InboundMessage
|
|
from nanobot.bus.queue import MessageBus
|
|
from nanobot.session.manager import Session
|
|
|
|
bus = MessageBus()
|
|
provider = MagicMock()
|
|
provider.get_default_model.return_value = "test-model"
|
|
|
|
loop = AgentLoop(bus=bus, provider=provider, workspace=tmp_path, model="test-model")
|
|
|
|
pending_queue: asyncio.Queue[InboundMessage] = asyncio.Queue()
|
|
session = Session(key="test:drain-block")
|
|
injection_callback = None
|
|
|
|
# Capture the injection_callback that _run_agent_loop creates
|
|
original_run = loop.runner.run
|
|
|
|
async def fake_runner_run(spec):
|
|
nonlocal injection_callback
|
|
injection_callback = spec.injection_callback
|
|
|
|
# Simulate: first call to injection_callback should block because
|
|
# sub-agents are running and no messages are in the queue yet.
|
|
# We'll resolve this from a concurrent task.
|
|
return SimpleNamespace(
|
|
stop_reason="done",
|
|
final_content="done",
|
|
error=None,
|
|
tool_events=[],
|
|
messages=[],
|
|
usage={},
|
|
had_injections=False,
|
|
tools_used=[],
|
|
)
|
|
|
|
loop.runner.run = AsyncMock(side_effect=fake_runner_run)
|
|
|
|
# Register a running sub-agent in the SubagentManager for this session
|
|
async def _hang_forever():
|
|
await asyncio.Event().wait()
|
|
|
|
hang_task = asyncio.create_task(_hang_forever())
|
|
loop.subagents._session_tasks.setdefault(session.key, set()).add("sub-drain-1")
|
|
loop.subagents._running_tasks["sub-drain-1"] = hang_task
|
|
|
|
# Run _run_agent_loop — this defines the _drain_pending closure
|
|
await loop._run_agent_loop(
|
|
[{"role": "user", "content": "test"}],
|
|
session=session,
|
|
channel="test",
|
|
chat_id="c1",
|
|
pending_queue=pending_queue,
|
|
)
|
|
|
|
assert injection_callback is not None
|
|
|
|
# Now test the callback directly
|
|
# With sub-agents running and an empty queue, it should block
|
|
drain_task = asyncio.create_task(injection_callback())
|
|
|
|
# Give it a moment to enter the blocking wait
|
|
await asyncio.sleep(0.05)
|
|
|
|
# Should still be running (blocked on pending_queue.get())
|
|
assert not drain_task.done(), "drain should block while sub-agents are running"
|
|
|
|
# Now put a message in the queue (simulating sub-agent completion)
|
|
await pending_queue.put(InboundMessage(
|
|
sender_id="subagent",
|
|
channel="test",
|
|
chat_id="c1",
|
|
content="Sub-agent result",
|
|
media=None,
|
|
metadata={},
|
|
))
|
|
|
|
# Should unblock and return results
|
|
results = await asyncio.wait_for(drain_task, timeout=2.0)
|
|
assert len(results) >= 1
|
|
assert results[0]["role"] == "user"
|
|
assert "Sub-agent result" in str(results[0]["content"])
|
|
|
|
# Cleanup
|
|
hang_task.cancel()
|
|
try:
|
|
await hang_task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_drain_pending_no_block_when_no_subagents(tmp_path):
|
|
"""_drain_pending should not block when no sub-agents are running."""
|
|
from nanobot.agent.loop import AgentLoop
|
|
from nanobot.bus.queue import MessageBus
|
|
|
|
bus = MessageBus()
|
|
provider = MagicMock()
|
|
provider.get_default_model.return_value = "test-model"
|
|
|
|
loop = AgentLoop(bus=bus, provider=provider, workspace=tmp_path, model="test-model")
|
|
|
|
pending_queue: asyncio.Queue = asyncio.Queue()
|
|
injection_callback = None
|
|
|
|
async def fake_runner_run(spec):
|
|
nonlocal injection_callback
|
|
injection_callback = spec.injection_callback
|
|
return SimpleNamespace(
|
|
stop_reason="done",
|
|
final_content="done",
|
|
error=None,
|
|
tool_events=[],
|
|
messages=[],
|
|
usage={},
|
|
had_injections=False,
|
|
tools_used=[],
|
|
)
|
|
|
|
loop.runner.run = AsyncMock(side_effect=fake_runner_run)
|
|
|
|
await loop._run_agent_loop(
|
|
[{"role": "user", "content": "test"}],
|
|
session=None,
|
|
channel="test",
|
|
chat_id="c1",
|
|
pending_queue=pending_queue,
|
|
)
|
|
|
|
assert injection_callback is not None
|
|
|
|
# With no sub-agents and empty queue, should return immediately
|
|
results = await asyncio.wait_for(injection_callback(), timeout=1.0)
|
|
assert results == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_drain_pending_timeout(tmp_path):
|
|
"""_drain_pending should return empty after timeout when sub-agents hang."""
|
|
from nanobot.agent.loop import AgentLoop
|
|
from nanobot.bus.queue import MessageBus
|
|
from nanobot.session.manager import Session
|
|
|
|
bus = MessageBus()
|
|
provider = MagicMock()
|
|
provider.get_default_model.return_value = "test-model"
|
|
|
|
loop = AgentLoop(bus=bus, provider=provider, workspace=tmp_path, model="test-model")
|
|
|
|
pending_queue: asyncio.Queue = asyncio.Queue()
|
|
session = Session(key="test:drain-timeout")
|
|
injection_callback = None
|
|
|
|
async def fake_runner_run(spec):
|
|
nonlocal injection_callback
|
|
injection_callback = spec.injection_callback
|
|
return SimpleNamespace(
|
|
stop_reason="done",
|
|
final_content="done",
|
|
error=None,
|
|
tool_events=[],
|
|
messages=[],
|
|
usage={},
|
|
had_injections=False,
|
|
tools_used=[],
|
|
)
|
|
|
|
loop.runner.run = AsyncMock(side_effect=fake_runner_run)
|
|
|
|
# Register a "running" sub-agent that will never complete
|
|
async def _hang_forever():
|
|
await asyncio.Event().wait()
|
|
|
|
hang_task = asyncio.create_task(_hang_forever())
|
|
loop.subagents._session_tasks.setdefault(session.key, set()).add("sub-timeout-1")
|
|
loop.subagents._running_tasks["sub-timeout-1"] = hang_task
|
|
|
|
await loop._run_agent_loop(
|
|
[{"role": "user", "content": "test"}],
|
|
session=session,
|
|
channel="test",
|
|
chat_id="c1",
|
|
pending_queue=pending_queue,
|
|
)
|
|
|
|
assert injection_callback is not None
|
|
|
|
# Patch the timeout to be very short for testing
|
|
with patch("nanobot.agent.loop.asyncio.wait_for") as mock_wait:
|
|
mock_wait.side_effect = asyncio.TimeoutError
|
|
results = await injection_callback()
|
|
assert results == []
|
|
|
|
# Cleanup
|
|
hang_task.cancel()
|
|
try:
|
|
await hang_task
|
|
except asyncio.CancelledError:
|
|
pass
|