From 3b03cc20791ff8282ce4bc5118dd88cb126a9e26 Mon Sep 17 00:00:00 2001 From: chengyongru Date: Fri, 5 Jun 2026 18:54:29 +0800 Subject: [PATCH] feat(subagent): add mailbox-backed worker results --- .agent/subagent-mailbox-mvp-plan.md | 290 ++++++++++++++++++ nanobot/agent/loop.py | 29 +- nanobot/agent/mailbox.py | 280 +++++++++++++++++ nanobot/agent/subagent.py | 130 ++++++-- nanobot/agent/tools/spawn.py | 7 +- nanobot/agent/tools/subagent_mailbox.py | 207 +++++++++++++ tests/agent/test_subagent_lifecycle.py | 70 +++-- tests/agent/test_task_cancel.py | 30 +- .../tools/test_subagent_mailbox_tools.py | 118 +++++++ tests/agent/tools/test_subagent_tools.py | 39 +-- 10 files changed, 1074 insertions(+), 126 deletions(-) create mode 100644 .agent/subagent-mailbox-mvp-plan.md create mode 100644 nanobot/agent/mailbox.py create mode 100644 nanobot/agent/tools/subagent_mailbox.py create mode 100644 tests/agent/tools/test_subagent_mailbox_tools.py diff --git a/.agent/subagent-mailbox-mvp-plan.md b/.agent/subagent-mailbox-mvp-plan.md new file mode 100644 index 000000000..01af27e54 --- /dev/null +++ b/.agent/subagent-mailbox-mvp-plan.md @@ -0,0 +1,290 @@ +# Subagent Mailbox MVP Plan + +This document records the agreed implementation direction for replacing the +current subagent wait behavior. It is intentionally narrower than PR #3461. + +## Goal Prompt + +Use this goal when asking an agent to implement the change: + +```text +Implement a minimal, mergeable mailbox-backed manager-worker coordination MVP +that replaces the current implicit subagent-result path through +pending_queue/mid-turn injection. + +Goals: +1. Borrow only the core ideas from PR #3461: manager-worker mailbox mechanics, + task/result messages, and explicit poll/wait. Treat task-scoped session state + as an optional supporting idea, not a required first-pass feature. Do not copy the + broadcast/bid/aggregation/circuit-breaker/create-instance skill pieces. +2. Ensure subagent/task results no longer masquerade as ordinary inbound + messages routed through the main session pending_queue, and ensure + _drain_pending does not block the current turn just because a subagent is + still running. +3. Add a clear mailbox protocol layer with minimal dispatch/spawn, result, + poll/wait, cancel/finalize semantics. This is a strict manager-worker model, + not peer-to-peer agent collaboration. +4. Keep responsibilities separated: AgentLoop handles user turn scheduling, + AgentRunner handles model/tool execution, and mailbox/worker management + handles manager-worker messages and task lifecycle. +5. Choose the smallest deterministic implementation that can be tested cleanly. + In-process and filesystem-backed stores are both acceptable design options; + if filesystem storage is used, atomic writes, result deduplication, and + restart-safe reads must be covered. +6. Preserve existing chat behavior: ordinary user follow-up messages, + streaming stream_end, /stop, session history, and long_task goal state must + keep working. +7. Add focused tests covering: dispatch does not block the current turn, + explicit result wait/poll, user follow-up is not confused with worker result, + cancellation/finalization semantics if implemented, result deduplication, and + removal of the old implicit subagent wait path. +8. Run the relevant pytest and ruff checks. If any cannot run, report the reason + and residual risk. + +Non-goals: +- Do not implement a full multi-process agent network. +- Do not implement decentralized P2P agent collaboration. Prior P2P-style agent + exchange has already proven unsuitable for the current agent behavior. +- Do not implement broadcast/bid marketplace behavior. +- Do not implement create-instance skill. +- Do not do broad WebUI work. Any status-event adaptation must be minimal and + justified by correctness or test evidence. + +Done means: +The main agent can dispatch a background worker and finish the current turn +normally. The worker result lands in a mailbox/result store. The main agent can +consume it through an explicit wait/poll tool. User messages continue through +the normal user-turn/pending-message path and are not blocked by hidden subagent +waiting. +``` + +## Problem To Solve + +The current implementation couples subagent completion to the main agent's +mid-turn injection queue: + +- Subagent completion is published as a synthetic inbound message. +- The main loop routes same-session inbound messages to the active pending queue. +- The runner drains that queue as mid-turn injections. +- If the queue is empty but a subagent is still running, the drain callback can + wait for the queue for a long time. + +This makes subagent lifecycle a hidden dependency of the current turn. It also +mixes user follow-up messages with worker results in the same queue. + +The target behavior is explicit message passing: + +- User messages and worker results have separate paths. +- Worker results are stored as task results, not injected as ordinary inbound + messages by default. +- Waiting is an explicit tool action, not a hidden AgentLoop behavior. + +## Current Code Points + +These line references were taken from the origin/main-based worktree at the time +this plan was written. They are landmarks, not a requirement to preserve exact +line numbers. + +- `nanobot/agent/loop.py:712` defines `_drain_pending`. +- `nanobot/agent/loop.py:733-738` drains pending messages with `get_nowait`. +- `nanobot/agent/loop.py:743-747` blocks when no pending items exist but the + session still has running subagents. +- `nanobot/agent/loop.py:880-896` routes same-session inbound messages into the + active pending queue. +- `nanobot/agent/subagent.py:174-198` creates a background task for a subagent. +- `nanobot/agent/subagent.py:242-257` runs the subagent using `AgentRunner`. +- `nanobot/agent/subagent.py:309-330` publishes the subagent result back as an + inbound message. +- `nanobot/agent/tools/spawn.py:71-88` exposes the current spawn tool. + +## Consensus Direction + +Use a small mailbox-backed coordination layer. Do not move directly to a full +Codex-style thread tree. Do not keep fixing the hidden pending queue wait as the +long-term architecture. + +This is explicitly a master/worker model: + +- The main agent is the only orchestrator for the user-facing task. +- Workers receive bounded delegated tasks and return results. +- Workers do not negotiate with each other, bid on work, form a decentralized + network, or independently decide to report to the user. +- Any further delegation must be a deliberate future design, not an accidental + property of the mailbox protocol. + +The key property is the protocol boundary, not the storage backend. Backend +choice is an implementation decision: choose the smallest option that proves the +manager-worker behavior and keeps tests deterministic. + +Possible concepts. These names are examples, not settled API: + +- `TaskId`: stable id returned by dispatch/spawn. +- `WorkerId`: logical worker id. For the first version this can be in-process + worker ids owned by the main agent. +- `TaskMessage`: request payload with task description, origin session, created + time, deadline, and cancellation metadata. +- `TaskResult`: completion payload with task id, status, content, error, sender, + completed time, and dedupe key. +- `MailboxStore`: append/read/claim result records. +- `WorkerManager`: starts in-process workers and writes results to the mailbox. +- `wait_subagents` or `wait_agent`: explicit tool that consumes mailbox results. +- `poll_subagents` or `poll_agent`: non-blocking status/result check. +- `cancel_subagent` or `finalize_task`: explicit cancellation/finalization. + +## Borrow From PR #3461 + +The useful ideas from PR #3461 are: + +- Filesystem-backed inbox/processed layout as a possible persistence model. +- Dispatch writes a task message and returns immediately. +- Result reporting writes a separate result message to the requester. +- Polling is explicit. +- Task-scoped sessions may give delegated work isolated context. + +Borrow only mailbox mechanics, not P2P semantics. PR #3461's broader +decentralized collaboration direction is not a good fit for the current agent. +The MVP should preserve a clear main-agent-to-worker hierarchy. + +Specific PR #3461 landmarks: + +- `nanobot/p2p/shell.py:18` defines a mailbox-like shell. +- `nanobot/p2p/shell.py:21-29` gives each agent inbox and processed dirs. +- `nanobot/p2p/shell.py:66-122` dispatches a task by writing to target inbox. +- `nanobot/p2p/shell.py:124-158` polls task status/results. +- `nanobot/p2p/shell.py:259-283` writes result messages. +- `nanobot/session/manager.py:584-612` sketches task-scoped sessions. +- `nanobot/agent/context.py:77-97` adds task-session collaboration hints. + +Do not borrow these parts for the MVP: + +- Peer-to-peer/decentralized agent exchange. +- Broadcast/bid aggregation. +- Circuit breaker/failover. +- Create-instance skill. +- Heartbeat-based inbox scanning. +- Default-channel `report_user` delivery without reliable callback metadata. + +## Architecture Boundary + +Keep these responsibilities separate: + +- `AgentLoop`: owns user turn scheduling, session locks, user pending messages, + commands, streaming callbacks, and runtime events. +- `AgentRunner`: owns model/tool iteration and injection callback execution. +- `MailboxStore`: owns task/result records and deduplication. +- `WorkerManager`: owns worker lifecycle, cancellation, and result publication. +- Tools: expose explicit operations to the LLM: spawn/dispatch, wait/poll, + cancel/finalize. + +Workers should not have tools that let them directly orchestrate peer workers in +the first version. They may use ordinary task tools to complete their delegated +work, then return a result to the main agent. + +The mailbox layer should not know about WebUI-specific wire details. If UI status +is needed, emit generic runtime events or expose status through existing session +state patterns. + +## Implementation Shape + +Suggested first pass. This is a starting shape, not a fixed design: + +1. Add a small mailbox module, for example `nanobot/agent/mailbox.py` or + `nanobot/session/mailbox.py`. +2. Add dataclasses for task request/result/status. Keep them JSON-serializable. +3. Choose a simple store backend. In-memory is fine for a first implementation; + filesystem is fine only if it stays simple and is tested for atomicity and + deduplication. +4. Modify `SubagentManager` into a worker supervisor that writes completion to + the mailbox instead of publishing inbound results. +5. Remove the long subagent-running wait from `_drain_pending`. +6. Add explicit wait/poll/cancel tools. +7. Decide whether to keep the existing `spawn` tool name for compatibility or + introduce a clearer worker-specific name. +8. Add focused tests before broad refactors. + +Open transitional behaviors: + +- A worker completion notification may be useful when no active turn exists, but + this is not part of the agreed MVP unless explicitly chosen. It must not + reintroduce hidden waiting in `_drain_pending`. +- Session history persistence for worker results needs a deliberate choice: + write after explicit wait/poll, and decide whether the durable entry is + assistant, system, or metadata-only. + +## Open Decisions + +These points are not yet consensus and should not be treated as requirements: + +- Store backend: in-memory first, filesystem first, or a small interface with one + concrete implementation. +- Task-scoped sessions: useful idea from PR #3461, but optional for the MVP. +- Public tool names: keep `spawn`, add `wait_subagents`, use `dispatch_task`, or + choose clearer worker-specific names. +- Worker completion notification: explicit wait/poll only, or a minimal + notification when no active turn exists. +- Session history semantics: when and how consumed worker results become durable + conversation history. +- UI/runtime status: no broad WebUI work; any minimal status event needs a clear + correctness reason. +- Exact module/class names: `MailboxStore`, `WorkerManager`, and `WorkerId` are + placeholders for the implementation discussion. + +## Test Plan + +Minimum tests: + +- Dispatch returns before worker completion. +- Current turn reaches final response/stream_end while worker is still running. +- Worker completion is stored in mailbox. +- Explicit wait returns the result once and does not duplicate it. +- Explicit poll reports running/completed/not_found states. +- User follow-up during an active main turn still uses ordinary pending queue + behavior. +- User follow-up is not ordered behind hidden subagent waits. +- `/stop` cancels active workers for the session. +- Finalize/cancel marks task state and prevents later result injection. +- Existing long_task goal state continuation still works. + +Useful regression target: + +- A test should fail on the old implementation because `_drain_pending` waits on + `pending_queue.get()` solely due to a running subagent, then pass after the + hidden wait is removed. + +## Risks And Guardrails + +Main risks: + +- Accidentally replacing one hidden queue with another hidden queue. +- Accidentally recreating PR #3461's P2P agent network instead of a strict + manager-worker boundary. +- Duplicating results after repeated wait/poll calls; if persistent storage is + chosen, duplicating results after restart. +- Losing compatibility with existing spawn tool expectations. +- Making `AgentLoop` larger instead of reducing its subagent-specific knowledge. +- Over-scoping the first PR with marketplace or multi-process behavior. + +Guardrails: + +- Keep the MVP small. +- Keep the main agent in charge of orchestration. +- Keep waiting explicit. +- Keep user messages and worker results on separate paths. +- Avoid new WebUI behavior unless needed for correctness. +- Preserve existing tests for pending messages, streaming, stop, session history, + and long_task. + +## Review Checklist + +Before considering the implementation complete: + +- `_drain_pending` no longer blocks for running subagents. +- Subagent/worker result publication does not call `bus.publish_inbound` as the + primary result path. +- There is a clear task id in every spawn/dispatch response. +- There is a clear explicit way to wait or poll for a task result. +- Repeated wait/poll calls do not duplicate consumed results. +- Cancellation has a defined state. +- The implementation does not include PR #3461 broadcast/bid/create-instance + features. +- Tests cover the old stuck-turn behavior and the new explicit mailbox behavior. diff --git a/nanobot/agent/loop.py b/nanobot/agent/loop.py index 050f0cf7d..208e287d5 100644 --- a/nanobot/agent/loop.py +++ b/nanobot/agent/loop.py @@ -731,11 +731,9 @@ class AgentLoop: async def _drain_pending(*, limit: int = _MAX_INJECTIONS_PER_TURN) -> list[dict[str, Any]]: """Drain follow-up messages from the pending queue. - When no messages are immediately available but sub-agents - spawned in this dispatch are still running, blocks until at - least one result arrives (or timeout). This keeps the runner - loop alive so subsequent sub-agent completions are consumed - in-order rather than dispatched separately. + This path is only for real same-session user follow-up messages. + Worker results are read explicitly through the subagent mailbox + tools instead of being injected as ordinary inbound messages. """ if pending_queue is None: return [] @@ -756,27 +754,6 @@ class AgentLoop: except asyncio.QueueEmpty: break - # Block if nothing drained but sub-agents spawned in this dispatch - # are still running. Keeps the runner loop alive so subsequent - # completions are injected in-order rather than dispatched separately. - if (not items - and session is not None - and self.subagents.get_running_count_by_session(session.key) > 0): - try: - msg = await asyncio.wait_for(pending_queue.get(), timeout=300) - except asyncio.TimeoutError: - logger.warning( - "Timeout waiting for sub-agent completion in session {}", - session.key, - ) - return items - items.append(_to_user_message(msg)) - while len(items) < limit: - try: - items.append(_to_user_message(pending_queue.get_nowait())) - except asyncio.QueueEmpty: - break - return items active_session_key = session.key if session else session_key diff --git a/nanobot/agent/mailbox.py b/nanobot/agent/mailbox.py new file mode 100644 index 000000000..e3e0f0037 --- /dev/null +++ b/nanobot/agent/mailbox.py @@ -0,0 +1,280 @@ +"""Mailbox primitives for manager-worker task coordination.""" + +from __future__ import annotations + +import asyncio +import time +from dataclasses import dataclass, field +from typing import Any + +TaskState = str # running | completed | failed | cancelled +MailboxReadState = str # ready | running | not_found | consumed | timeout + + +@dataclass(slots=True) +class TaskRequest: + """Task request recorded when the manager dispatches a worker.""" + + task_id: str + session_key: str + label: str + task: str + origin: dict[str, Any] = field(default_factory=dict) + created_at: float = field(default_factory=time.time) + + +@dataclass(slots=True) +class TaskResult: + """Worker result written to the manager mailbox.""" + + task_id: str + session_key: str + label: str + task: str + status: str + content: str + sender: str = "subagent" + completed_at: float = field(default_factory=time.time) + dedupe_key: str | None = None + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(slots=True) +class TaskSnapshot: + """Read-only view of a task in the mailbox.""" + + task_id: str + session_key: str + label: str + task: str + state: TaskState + created_at: float + completed_at: float | None = None + consumed_at: float | None = None + result_status: str | None = None + error: str | None = None + + +@dataclass(slots=True) +class MailboxRead: + """Result of a mailbox wait/consume operation.""" + + state: MailboxReadState + task: TaskSnapshot | None = None + result: TaskResult | None = None + + +@dataclass(slots=True) +class _TaskRecord: + request: TaskRequest + state: TaskState = "running" + result: TaskResult | None = None + consumed_at: float | None = None + completed_at: float | None = None + error: str | None = None + + +class MailboxStore: + """In-memory mailbox for worker task/result records. + + The store owns result deduplication and one-time result consumption. It is + intentionally small; persistence can be added behind this protocol later + without putting worker results back on the user pending queue. + """ + + def __init__(self) -> None: + self._records: dict[str, _TaskRecord] = {} + self._session_tasks: dict[str, set[str]] = {} + self._dedupe_keys: set[str] = set() + self._changed = asyncio.Condition() + + async def dispatch(self, request: TaskRequest) -> None: + """Record that a task was dispatched.""" + async with self._changed: + if request.task_id in self._records: + return + self._records[request.task_id] = _TaskRecord(request=request) + self._session_tasks.setdefault(request.session_key, set()).add(request.task_id) + self._changed.notify_all() + + async def record_result(self, result: TaskResult) -> bool: + """Record a worker result. + + Returns ``True`` when this call writes a new result and ``False`` when + the result is a duplicate or the task was already finalized. + """ + async with self._changed: + dedupe_key = result.dedupe_key or result.task_id + if dedupe_key in self._dedupe_keys: + return False + + record = self._records.get(result.task_id) + if record is None: + request = TaskRequest( + task_id=result.task_id, + session_key=result.session_key, + label=result.label, + task=result.task, + origin=dict(result.metadata), + created_at=result.completed_at, + ) + record = _TaskRecord(request=request) + self._records[result.task_id] = record + self._session_tasks.setdefault(result.session_key, set()).add(result.task_id) + elif record.result is not None: + self._dedupe_keys.add(dedupe_key) + return False + + record.result = result + record.completed_at = result.completed_at + record.state = self._state_for_result(result.status) + record.error = result.content if result.status in {"error", "cancelled"} else None + self._dedupe_keys.add(dedupe_key) + self._changed.notify_all() + return True + + async def mark_cancelled( + self, + task_id: str, + *, + session_key: str | None = None, + reason: str = "Cancelled.", + ) -> bool: + """Mark a task cancelled and make the cancellation consumable once.""" + async with self._changed: + record = self._records.get(task_id) + if record is None: + return False + if session_key is not None and record.request.session_key != session_key: + return False + if record.result is not None: + return False + result = TaskResult( + task_id=task_id, + session_key=record.request.session_key, + label=record.request.label, + task=record.request.task, + status="cancelled", + content=reason, + dedupe_key=task_id, + ) + record.result = result + record.completed_at = result.completed_at + record.state = "cancelled" + record.error = reason + self._dedupe_keys.add(task_id) + self._changed.notify_all() + return True + + async def poll( + self, + session_key: str, + *, + task_id: str | None = None, + ) -> list[TaskSnapshot]: + """Return snapshots for one task or all tasks in a session.""" + async with self._changed: + if task_id is not None: + record = self._records.get(task_id) + if record is None or record.request.session_key != session_key: + return [] + return [self._snapshot(record)] + + ids = self._session_tasks.get(session_key, set()) + snapshots = [ + self._snapshot(self._records[tid]) + for tid in ids + if tid in self._records + ] + snapshots.sort(key=lambda item: (item.completed_at is None, item.created_at, item.task_id)) + return snapshots + + async def wait_for_result( + self, + session_key: str, + *, + task_id: str | None = None, + timeout_seconds: float = 30.0, + ) -> MailboxRead: + """Wait for and consume a result once.""" + deadline = time.monotonic() + max(0.0, timeout_seconds) + async with self._changed: + while True: + read = self._consume_ready_locked(session_key, task_id) + if read.state != "running": + return read + remaining = deadline - time.monotonic() + if remaining <= 0: + return MailboxRead("timeout", task=read.task) + try: + await asyncio.wait_for(self._changed.wait(), timeout=remaining) + except asyncio.TimeoutError: + return MailboxRead("timeout", task=read.task) + + def _consume_ready_locked( + self, + session_key: str, + task_id: str | None, + ) -> MailboxRead: + if task_id is not None: + record = self._records.get(task_id) + if record is None or record.request.session_key != session_key: + return MailboxRead("not_found") + snapshot = self._snapshot(record) + if record.result is None: + return MailboxRead("running", task=snapshot) + if record.consumed_at is not None: + return MailboxRead("consumed", task=snapshot, result=record.result) + record.consumed_at = time.time() + snapshot = self._snapshot(record) + return MailboxRead("ready", task=snapshot, result=record.result) + + ids = self._session_tasks.get(session_key, set()) + records = [ + self._records[tid] + for tid in ids + if tid in self._records + ] + ready = [ + record + for record in records + if record.result is not None and record.consumed_at is None + ] + if ready: + ready.sort(key=lambda record: (record.completed_at or record.request.created_at, record.request.task_id)) + record = ready[0] + record.consumed_at = time.time() + return MailboxRead("ready", task=self._snapshot(record), result=record.result) + + running = [record for record in records if record.result is None] + if running: + running.sort(key=lambda record: (record.request.created_at, record.request.task_id)) + return MailboxRead("running", task=self._snapshot(running[0])) + if records: + records.sort(key=lambda record: (record.completed_at or record.request.created_at, record.request.task_id)) + return MailboxRead("consumed", task=self._snapshot(records[-1])) + return MailboxRead("not_found") + + @staticmethod + def _state_for_result(status: str) -> TaskState: + if status == "ok": + return "completed" + if status == "cancelled": + return "cancelled" + return "failed" + + @staticmethod + def _snapshot(record: _TaskRecord) -> TaskSnapshot: + result = record.result + return TaskSnapshot( + task_id=record.request.task_id, + session_key=record.request.session_key, + label=record.request.label, + task=record.request.task, + state=record.state, + created_at=record.request.created_at, + completed_at=record.completed_at, + consumed_at=record.consumed_at, + result_status=result.status if result is not None else None, + error=record.error, + ) diff --git a/nanobot/agent/subagent.py b/nanobot/agent/subagent.py index d93f8419e..768b9fe0d 100644 --- a/nanobot/agent/subagent.py +++ b/nanobot/agent/subagent.py @@ -4,6 +4,7 @@ import asyncio import json import time import uuid +from contextlib import suppress from dataclasses import dataclass, field from pathlib import Path from typing import Any, Callable @@ -11,12 +12,12 @@ from typing import Any, Callable from loguru import logger from nanobot.agent.hook import AgentHook, AgentHookContext +from nanobot.agent.mailbox import MailboxRead, MailboxStore, TaskRequest, TaskResult, TaskSnapshot from nanobot.agent.runner import AgentRunner, AgentRunSpec from nanobot.agent.tools.context import ToolContext from nanobot.agent.tools.file_state import FileStates from nanobot.agent.tools.loader import ToolLoader from nanobot.agent.tools.registry import ToolRegistry -from nanobot.bus.events import InboundMessage from nanobot.bus.queue import MessageBus from nanobot.config.schema import AgentDefaults, ToolsConfig from nanobot.providers.base import LLMProvider @@ -87,6 +88,7 @@ class SubagentManager: max_iterations: int | None = None, max_concurrent_subagents: int | None = None, llm_wall_timeout_for_session: Callable[[str | None], float | None] | None = None, + mailbox: MailboxStore | None = None, ): defaults = AgentDefaults() self.provider = provider @@ -109,6 +111,7 @@ class SubagentManager: ) self.runner = AgentRunner(provider) self._llm_wall_timeout_for_session = llm_wall_timeout_for_session + self.mailbox = mailbox or MailboxStore() self._running_tasks: dict[str, asyncio.Task[None]] = {} self._task_statuses: dict[str, SubagentStatus] = {} self._session_tasks: dict[str, set[str]] = {} # session_key -> {task_id, ...} @@ -162,6 +165,7 @@ class SubagentManager: """Spawn a subagent to execute a task in the background.""" task_id = str(uuid.uuid4())[:8] display_label = label or task[:30] + ("..." if len(task) > 30 else "") + mailbox_session_key = session_key or f"{origin_channel}:{origin_chat_id}" origin = {"channel": origin_channel, "chat_id": origin_chat_id, "session_key": session_key} status = SubagentStatus( @@ -171,6 +175,18 @@ class SubagentManager: started_at=time.monotonic(), ) self._task_statuses[task_id] = status + await self.mailbox.dispatch(TaskRequest( + task_id=task_id, + session_key=mailbox_session_key, + label=display_label, + task=task, + origin={ + "channel": origin_channel, + "chat_id": origin_chat_id, + "session_key": session_key, + "origin_message_id": origin_message_id, + }, + )) bg_task = asyncio.create_task( self._run_subagent( @@ -199,14 +215,17 @@ class SubagentManager: bg_task.add_done_callback(_cleanup) logger.info("Spawned subagent [{}]: {}", task_id, display_label) - return f"Subagent [{display_label}] started (id: {task_id}). I'll notify you when it completes." + return ( + f"Subagent [{display_label}] started (id: {task_id}). " + f"Use poll_subagents or wait_subagents with id {task_id} to get the result." + ) async def _run_subagent( self, task_id: str, task: str, label: str, - origin: dict[str, str], + origin: dict[str, Any], status: SubagentStatus, origin_message_id: str | None = None, temperature: float | None = None, @@ -281,6 +300,12 @@ class SubagentManager: logger.info("Subagent [{}] completed successfully", task_id) await self._announce_result(task_id, label, task, final_result, origin, "ok", origin_message_id) + except asyncio.CancelledError: + status.phase = "cancelled" + status.stop_reason = "cancelled" + await self.mailbox.mark_cancelled(task_id, reason="Cancelled.") + logger.info("Subagent [{}] cancelled", task_id) + raise except Exception as e: status.phase = "error" status.error = str(e) @@ -293,44 +318,39 @@ class SubagentManager: label: str, task: str, result: str, - origin: dict[str, str], + origin: dict[str, Any], status: str, origin_message_id: str | None = None, ) -> None: - """Announce the subagent result to the main agent via the message bus.""" - status_text = "completed successfully" if status == "ok" else "failed" - - announce_content = render_template( - "agent/subagent_announce.md", - label=label, - status_text=status_text, - task=task, - result=result, - ) - - # Inject as system message to trigger main agent. - # Use session_key_override to align with the main agent's effective - # session key (which accounts for unified sessions) so the result is - # routed to the correct pending queue (mid-turn injection) instead of - # being dispatched as a competing independent task. + """Record the subagent result in the mailbox for explicit manager polling.""" override = origin.get("session_key") or f"{origin['channel']}:{origin['chat_id']}" metadata: dict[str, Any] = { - "injected_event": "subagent_result", "subagent_task_id": task_id, + "origin_channel": origin.get("channel"), + "origin_chat_id": origin.get("chat_id"), } if origin_message_id: metadata["origin_message_id"] = origin_message_id - msg = InboundMessage( - channel="system", - sender_id="subagent", - chat_id=f"{origin['channel']}:{origin['chat_id']}", - content=announce_content, - session_key_override=override, - metadata=metadata, - ) - await self.bus.publish_inbound(msg) - logger.debug("Subagent [{}] announced result to {}:{}", task_id, origin['channel'], origin['chat_id']) + written = await self.mailbox.record_result(TaskResult( + task_id=task_id, + session_key=override, + label=label, + task=task, + status=status, + content=result, + dedupe_key=task_id, + metadata=metadata, + )) + + if written: + logger.debug( + "Subagent [{}] wrote result to mailbox for session {}", + task_id, + override, + ) + else: + logger.debug("Subagent [{}] result already recorded", task_id) @staticmethod def _format_partial_progress(result) -> str: @@ -375,12 +395,60 @@ class SubagentManager: """Cancel all subagents for the given session. Returns count cancelled.""" tasks = [self._running_tasks[tid] for tid in self._session_tasks.get(session_key, []) if tid in self._running_tasks and not self._running_tasks[tid].done()] + for tid in list(self._session_tasks.get(session_key, [])): + if tid in self._running_tasks and not self._running_tasks[tid].done(): + await self.mailbox.mark_cancelled( + tid, + session_key=session_key, + reason="Cancelled by /stop.", + ) for t in tasks: t.cancel() if tasks: await asyncio.gather(*tasks, return_exceptions=True) return len(tasks) + async def cancel_task(self, task_id: str, session_key: str | None = None) -> str: + """Cancel one running subagent task and record a cancelled mailbox state.""" + snapshots = await self.mailbox.poll(session_key, task_id=task_id) if session_key else [] + if session_key and not snapshots: + return "not_found" + task = self._running_tasks.get(task_id) + if task is None or task.done(): + if snapshots: + return snapshots[0].state + return "not_found" + await self.mailbox.mark_cancelled( + task_id, + session_key=session_key, + reason="Cancelled by manager.", + ) + task.cancel() + with suppress(asyncio.CancelledError, Exception): + await task + return "cancelled" + + async def poll( + self, + session_key: str, + task_id: str | None = None, + ) -> list[TaskSnapshot]: + """Return mailbox task status snapshots for a session.""" + return await self.mailbox.poll(session_key, task_id=task_id) + + async def wait_for_result( + self, + session_key: str, + task_id: str | None = None, + timeout_seconds: float = 30.0, + ) -> MailboxRead: + """Wait for and consume a mailbox result for a session.""" + return await self.mailbox.wait_for_result( + session_key, + task_id=task_id, + timeout_seconds=timeout_seconds, + ) + def get_running_count(self) -> int: """Return the number of currently running subagents.""" return len(self._running_tasks) diff --git a/nanobot/agent/tools/spawn.py b/nanobot/agent/tools/spawn.py index 420afc048..780b32927 100644 --- a/nanobot/agent/tools/spawn.py +++ b/nanobot/agent/tools/spawn.py @@ -63,7 +63,8 @@ class SpawnTool(Tool, ContextAware): return ( "Spawn a subagent to handle a task in the background. " "Use this for complex or time-consuming tasks that can run independently. " - "The subagent will complete the task and report back when done. " + "The subagent writes its result to a mailbox; use poll_subagents " + "or wait_subagents to retrieve it explicitly. " "For deliverables or existing projects, inspect the workspace first " "and use a dedicated subdirectory when helpful." ) @@ -81,8 +82,8 @@ class SpawnTool(Tool, ContextAware): if running >= limit: return ( f"Cannot spawn subagent: concurrency limit reached " - f"({running}/{limit} running). Wait for a running subagent " - f"to complete before spawning a new one." + f"({running}/{limit} running). Use wait_subagents or cancel_subagent " + f"before spawning a new one." ) return await self._manager.spawn( task=task, diff --git a/nanobot/agent/tools/subagent_mailbox.py b/nanobot/agent/tools/subagent_mailbox.py new file mode 100644 index 000000000..fa761cbbe --- /dev/null +++ b/nanobot/agent/tools/subagent_mailbox.py @@ -0,0 +1,207 @@ +"""Explicit mailbox tools for subagent coordination.""" + +from __future__ import annotations + +from contextvars import ContextVar +from typing import TYPE_CHECKING, Any + +from nanobot.agent.mailbox import MailboxRead, TaskSnapshot +from nanobot.agent.tools.base import Tool, tool_parameters +from nanobot.agent.tools.context import ContextAware, RequestContext +from nanobot.agent.tools.schema import NumberSchema, StringSchema, tool_parameters_schema + +if TYPE_CHECKING: + from nanobot.agent.subagent import SubagentManager + + +def _normalize_task_id(task_id: str | None) -> str | None: + if task_id is None: + return None + task_id = task_id.strip() + return task_id or None + + +def _truncate(text: str, limit: int = 120) -> str: + text = " ".join(text.split()) + return text if len(text) <= limit else text[: limit - 3] + "..." + + +class _SubagentMailboxTool(Tool, ContextAware): + """Shared context plumbing for subagent mailbox tools.""" + + def __init__(self, manager: "SubagentManager"): + self._manager = manager + self._session_key: ContextVar[str] = ContextVar( + f"{self.__class__.__name__}_session_key", + default="cli:direct", + ) + + @classmethod + def enabled(cls, ctx: Any) -> bool: + return getattr(ctx, "subagent_manager", None) is not None + + @classmethod + def create(cls, ctx: Any) -> Tool: + return cls(manager=ctx.subagent_manager) + + def set_context(self, ctx: RequestContext) -> None: + self._session_key.set(ctx.session_key or f"{ctx.channel}:{ctx.chat_id}") + + +@tool_parameters( + tool_parameters_schema( + task_id=StringSchema( + "Optional subagent task id. Omit to list all subagent tasks for this session.", + nullable=True, + ), + ) +) +class PollSubagentsTool(_SubagentMailboxTool): + """Non-blocking task status check.""" + + @property + def name(self) -> str: + return "poll_subagents" + + @property + def description(self) -> str: + return ( + "Check subagent task status without blocking. Use this to see whether a " + "spawned subagent is still running or has a result ready to consume." + ) + + @property + def read_only(self) -> bool: + return True + + async def execute(self, task_id: str | None = None, **_: Any) -> str: + task_id = _normalize_task_id(task_id) + session_key = self._session_key.get() + snapshots = await self._manager.poll(session_key, task_id=task_id) + if not snapshots: + if task_id: + return f"Subagent task {task_id} not found for this session." + return "No subagent tasks found for this session." + return self._format_snapshots(snapshots) + + @staticmethod + def _format_snapshots(snapshots: list[TaskSnapshot]) -> str: + lines = ["Subagent task status:"] + for snapshot in snapshots: + state = snapshot.state + if snapshot.result_status and snapshot.consumed_at is None: + state = f"{state}, result ready" + elif snapshot.consumed_at is not None: + state = f"{state}, result consumed" + lines.append( + f"- id: {snapshot.task_id} | label: {snapshot.label} | " + f"status: {state} | task: {_truncate(snapshot.task)}" + ) + return "\n".join(lines) + + +@tool_parameters( + tool_parameters_schema( + task_id=StringSchema( + "Optional subagent task id. Omit to consume the next ready result.", + nullable=True, + ), + timeout_seconds=NumberSchema( + description="How long to wait for a result before returning. Defaults to 30 seconds.", + minimum=0.0, + maximum=300.0, + ), + ) +) +class WaitSubagentsTool(_SubagentMailboxTool): + """Wait for and consume one task result.""" + + @property + def name(self) -> str: + return "wait_subagents" + + @property + def description(self) -> str: + return ( + "Wait for a subagent result and consume it once. Use this after spawn " + "when you need the worker's result before continuing." + ) + + async def execute( + self, + task_id: str | None = None, + timeout_seconds: float = 30.0, + **_: Any, + ) -> str: + task_id = _normalize_task_id(task_id) + read = await self._manager.wait_for_result( + self._session_key.get(), + task_id=task_id, + timeout_seconds=timeout_seconds, + ) + return self._format_read(read, task_id) + + @staticmethod + def _format_read(read: MailboxRead, requested_task_id: str | None) -> str: + if read.state == "not_found": + target = f" {requested_task_id}" if requested_task_id else "" + return f"Subagent task{target} not found for this session." + if read.state == "timeout": + target = f" {read.task.task_id}" if read.task is not None else "" + return f"Timed out waiting for subagent task{target}." + if read.state == "consumed": + target = f" {read.task.task_id}" if read.task is not None else "" + return f"Subagent result for task{target} was already consumed." + if read.result is None or read.task is None: + return "No subagent result is ready." + + status_text = { + "ok": "completed", + "error": "failed", + "cancelled": "cancelled", + }.get(read.result.status, read.result.status) + return ( + f"Subagent result for [{read.result.label}] " + f"(id: {read.result.task_id}, status: {status_text}).\n\n" + f"Task: {read.result.task}\n\n" + f"Result:\n{read.result.content}" + ) + + +@tool_parameters( + tool_parameters_schema( + task_id=StringSchema("Subagent task id to cancel"), + required=["task_id"], + ) +) +class CancelSubagentTool(_SubagentMailboxTool): + """Cancel one running task.""" + + @property + def name(self) -> str: + return "cancel_subagent" + + @property + def description(self) -> str: + return ( + "Cancel a running subagent task and record a cancelled mailbox state. " + "Use this only when the delegated task is no longer needed." + ) + + async def execute(self, task_id: str, **_: Any) -> str: + task_id = _normalize_task_id(task_id) + if task_id is None: + return "Error: task_id is required." + state = await self._manager.cancel_task(task_id, session_key=self._session_key.get()) + if state == "cancelled": + return f"Cancelled subagent task {task_id}." + if state == "not_found": + return f"Subagent task {task_id} not found for this session." + if state in {"completed", "failed"}: + return ( + f"Subagent task {task_id} already {state}; " + "use wait_subagents to consume its result if needed." + ) + if state == "cancelled": + return f"Subagent task {task_id} is already cancelled." + return f"Subagent task {task_id} is {state}." diff --git a/tests/agent/test_subagent_lifecycle.py b/tests/agent/test_subagent_lifecycle.py index d68b3f7c1..7a07ff3bb 100644 --- a/tests/agent/test_subagent_lifecycle.py +++ b/tests/agent/test_subagent_lifecycle.py @@ -285,80 +285,76 @@ class TestRunSubagent: class TestAnnounceResult: @pytest.mark.asyncio - async def test_publishes_inbound_message(self, tmp_path): + async def test_records_mailbox_result_without_publishing_inbound(self, tmp_path): sm = _manager(tmp_path) - published = [] - sm.bus.publish_inbound = AsyncMock(side_effect=lambda msg: published.append(msg)) + sm.bus.publish_inbound = AsyncMock() await sm._announce_result( "t1", "label", "task", "result text", {"channel": "cli", "chat_id": "direct"}, "ok", ) - assert len(published) == 1 - msg = published[0] - assert msg.channel == "system" - assert msg.sender_id == "subagent" - assert msg.metadata["injected_event"] == "subagent_result" - assert msg.metadata["subagent_task_id"] == "t1" + sm.bus.publish_inbound.assert_not_awaited() + snapshots = await sm.mailbox.poll("cli:direct", task_id="t1") + assert snapshots[0].state == "completed" + read = await sm.mailbox.wait_for_result("cli:direct", task_id="t1", timeout_seconds=0) + assert read.state == "ready" + assert read.result is not None + assert read.result.content == "result text" + assert read.result.metadata["subagent_task_id"] == "t1" @pytest.mark.asyncio async def test_session_key_override(self, tmp_path): sm = _manager(tmp_path) - published = [] - sm.bus.publish_inbound = AsyncMock(side_effect=lambda msg: published.append(msg)) await sm._announce_result( "t1", "label", "task", "result", {"channel": "telegram", "chat_id": "123", "session_key": "s1"}, "ok", ) - assert published[0].session_key_override == "s1" + assert await sm.mailbox.poll("s1", task_id="t1") + assert await sm.mailbox.poll("telegram:123", task_id="t1") == [] @pytest.mark.asyncio async def test_session_key_override_fallback(self, tmp_path): sm = _manager(tmp_path) - published = [] - sm.bus.publish_inbound = AsyncMock(side_effect=lambda msg: published.append(msg)) await sm._announce_result( "t1", "label", "task", "result", {"channel": "telegram", "chat_id": "123"}, "ok", ) - assert published[0].session_key_override == "telegram:123" + snapshots = await sm.mailbox.poll("telegram:123", task_id="t1") + assert snapshots[0].session_key == "telegram:123" @pytest.mark.asyncio - async def test_ok_status_text(self, tmp_path): + async def test_ok_status_records_completed_state(self, tmp_path): sm = _manager(tmp_path) - published = [] - sm.bus.publish_inbound = AsyncMock(side_effect=lambda msg: published.append(msg)) await sm._announce_result( "t1", "label", "task", "result", {"channel": "cli", "chat_id": "direct"}, "ok", ) - assert "completed successfully" in published[0].content + snapshots = await sm.mailbox.poll("cli:direct", task_id="t1") + assert snapshots[0].state == "completed" @pytest.mark.asyncio - async def test_error_status_text(self, tmp_path): + async def test_error_status_records_failed_state(self, tmp_path): sm = _manager(tmp_path) - published = [] - sm.bus.publish_inbound = AsyncMock(side_effect=lambda msg: published.append(msg)) await sm._announce_result( "t1", "label", "task", "error details", {"channel": "cli", "chat_id": "direct"}, "error", ) - assert "failed" in published[0].content + snapshots = await sm.mailbox.poll("cli:direct", task_id="t1") + assert snapshots[0].state == "failed" + assert snapshots[0].error == "error details" @pytest.mark.asyncio async def test_origin_message_id_in_metadata(self, tmp_path): sm = _manager(tmp_path) - published = [] - sm.bus.publish_inbound = AsyncMock(side_effect=lambda msg: published.append(msg)) await sm._announce_result( "t1", "label", "task", "result", @@ -366,7 +362,29 @@ class TestAnnounceResult: origin_message_id="msg-123", ) - assert published[0].metadata["origin_message_id"] == "msg-123" + read = await sm.mailbox.wait_for_result("cli:direct", task_id="t1", timeout_seconds=0) + assert read.result is not None + assert read.result.metadata["origin_message_id"] == "msg-123" + + @pytest.mark.asyncio + async def test_duplicate_results_are_not_consumed_twice(self, tmp_path): + sm = _manager(tmp_path) + + await sm._announce_result( + "t1", "label", "task", "first", + {"channel": "cli", "chat_id": "direct"}, "ok", + ) + await sm._announce_result( + "t1", "label", "task", "second", + {"channel": "cli", "chat_id": "direct"}, "ok", + ) + + first = await sm.mailbox.wait_for_result("cli:direct", task_id="t1", timeout_seconds=0) + second = await sm.mailbox.wait_for_result("cli:direct", task_id="t1", timeout_seconds=0) + assert first.state == "ready" + assert first.result is not None + assert first.result.content == "first" + assert second.state == "consumed" # --------------------------------------------------------------------------- diff --git a/tests/agent/test_task_cancel.py b/tests/agent/test_task_cancel.py index e6a59d5e1..7a5366530 100644 --- a/tests/agent/test_task_cancel.py +++ b/tests/agent/test_task_cancel.py @@ -427,7 +427,7 @@ class TestSubagentCancellation: class TestSubagentAnnounceSessionKey: - """Verify _announce_result uses the effective session key for mid-turn routing.""" + """Verify _announce_result stores results under the effective session key.""" def _make_mgr(self): """Create a SubagentManager with mocked deps and its bus.""" @@ -448,27 +448,27 @@ class TestSubagentAnnounceSessionKey: @pytest.mark.asyncio async def test_announce_uses_effective_key_in_unified_mode(self): """In unified session mode, session_key_override must be 'unified:default' - so the result matches the pending queue key.""" + so the result matches the manager mailbox session key.""" mgr, bus = self._make_mgr() origin = {"channel": "telegram", "chat_id": "111", "session_key": UNIFIED_SESSION_KEY} await mgr._announce_result("sub-1", "label", "task", "result", origin, "ok") - msg = await bus.consume_inbound() - assert msg.session_key_override == UNIFIED_SESSION_KEY - assert msg.session_key == UNIFIED_SESSION_KEY + assert bus.inbound.empty() + snapshots = await mgr.mailbox.poll("unified:default", task_id="sub-1") + assert snapshots[0].session_key == "unified:default" @pytest.mark.asyncio async def test_announce_uses_raw_key_in_normal_mode(self): - """Without unified sessions, session_key_override is the raw channel:chat_id.""" + """Without unified sessions, the mailbox session is the raw channel:chat_id.""" mgr, bus = self._make_mgr() origin = {"channel": "telegram", "chat_id": "222", "session_key": "telegram:222"} await mgr._announce_result("sub-2", "label", "task", "result", origin, "ok") - msg = await bus.consume_inbound() - assert msg.session_key_override == "telegram:222" - assert msg.session_key == "telegram:222" + assert bus.inbound.empty() + snapshots = await mgr.mailbox.poll("telegram:222", task_id="sub-2") + assert snapshots[0].session_key == "telegram:222" @pytest.mark.asyncio async def test_announce_falls_back_to_origin_when_no_session_key(self): @@ -478,10 +478,9 @@ class TestSubagentAnnounceSessionKey: origin = {"channel": "discord", "chat_id": "333", "session_key": None} await mgr._announce_result("sub-3", "label", "task", "result", origin, "ok") - msg = await bus.consume_inbound() - assert msg.session_key_override == "discord:333" - assert msg.channel == "system" - assert msg.chat_id == "discord:333" + assert bus.inbound.empty() + snapshots = await mgr.mailbox.poll("discord:333", task_id="sub-3") + assert snapshots[0].session_key == "discord:333" @pytest.mark.asyncio async def test_session_key_flows_through_run_subagent(self): @@ -510,5 +509,6 @@ class TestSubagentAnnounceSessionKey: status, ) - msg = await bus.consume_inbound() - assert msg.session_key_override == UNIFIED_SESSION_KEY + assert bus.inbound.empty() + snapshots = await mgr.mailbox.poll("unified:default", task_id="sub-4") + assert snapshots[0].session_key == "unified:default" diff --git a/tests/agent/tools/test_subagent_mailbox_tools.py b/tests/agent/tools/test_subagent_mailbox_tools.py new file mode 100644 index 000000000..a0b5e2203 --- /dev/null +++ b/tests/agent/tools/test_subagent_mailbox_tools.py @@ -0,0 +1,118 @@ +"""Tests for explicit subagent mailbox tools.""" + +from __future__ import annotations + +import asyncio +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from nanobot.agent.runner import AgentRunResult +from nanobot.agent.subagent import SubagentManager +from nanobot.agent.tools.context import RequestContext +from nanobot.agent.tools.subagent_mailbox import ( + CancelSubagentTool, + PollSubagentsTool, + WaitSubagentsTool, +) +from nanobot.bus.queue import MessageBus +from nanobot.config.schema import AgentDefaults + + +def _manager(tmp_path: Path) -> SubagentManager: + provider = MagicMock() + provider.get_default_model.return_value = "test-model" + return SubagentManager( + provider=provider, + workspace=tmp_path, + bus=MessageBus(), + max_tool_result_chars=AgentDefaults().max_tool_result_chars, + ) + + +def _bind(tool, session_key: str = "cli:test") -> None: + tool.set_context(RequestContext(channel="cli", chat_id="test", session_key=session_key)) + + +async def _drain(mgr: SubagentManager) -> None: + tasks = list(mgr._running_tasks.values()) + if tasks: + await asyncio.gather(*tasks, return_exceptions=True) + await asyncio.sleep(0) + + +@pytest.mark.asyncio +async def test_wait_subagents_returns_result_once(tmp_path: Path) -> None: + mgr = _manager(tmp_path) + mgr.runner.run = AsyncMock( + return_value=AgentRunResult(final_content="worker result", messages=[], stop_reason="completed") + ) + + await mgr.spawn("do work", label="worker", session_key="cli:test") + task_id = next(iter(mgr._running_tasks)) + await _drain(mgr) + + wait_tool = WaitSubagentsTool(mgr) + _bind(wait_tool) + + first = await wait_tool.execute(task_id=task_id, timeout_seconds=0) + second = await wait_tool.execute(task_id=task_id, timeout_seconds=0) + + assert "worker result" in first + assert f"id: {task_id}" in first + assert "already consumed" in second + + +@pytest.mark.asyncio +async def test_poll_subagents_reports_running_completed_and_not_found(tmp_path: Path) -> None: + mgr = _manager(tmp_path) + release = asyncio.Event() + + async def _run(_spec): + await release.wait() + return AgentRunResult(final_content="done", messages=[], stop_reason="completed") + + mgr.runner.run = AsyncMock(side_effect=_run) + await mgr.spawn("slow work", label="slow", session_key="cli:test") + task_id = next(iter(mgr._running_tasks)) + + poll_tool = PollSubagentsTool(mgr) + _bind(poll_tool) + + running = await poll_tool.execute(task_id=task_id) + missing = await poll_tool.execute(task_id="missing") + release.set() + await _drain(mgr) + completed = await poll_tool.execute(task_id=task_id) + + assert "status: running" in running + assert "not found" in missing + assert "completed, result ready" in completed + + +@pytest.mark.asyncio +async def test_cancel_subagent_marks_cancelled_result(tmp_path: Path) -> None: + mgr = _manager(tmp_path) + started = asyncio.Event() + + async def _run(_spec): + started.set() + await asyncio.Event().wait() + + mgr.runner.run = AsyncMock(side_effect=_run) + await mgr.spawn("slow work", label="slow", session_key="cli:test") + task_id = next(iter(mgr._running_tasks)) + await asyncio.wait_for(started.wait(), timeout=1.0) + + cancel_tool = CancelSubagentTool(mgr) + wait_tool = WaitSubagentsTool(mgr) + _bind(cancel_tool) + _bind(wait_tool) + + cancelled = await cancel_tool.execute(task_id=task_id) + result = await wait_tool.execute(task_id=task_id, timeout_seconds=0) + + assert cancelled == f"Cancelled subagent task {task_id}." + assert "status: cancelled" in result + assert "Cancelled by manager." in result diff --git a/tests/agent/tools/test_subagent_tools.py b/tests/agent/tools/test_subagent_tools.py index 6a037981d..351b0ff34 100644 --- a/tests/agent/tools/test_subagent_tools.py +++ b/tests/agent/tools/test_subagent_tools.py @@ -279,8 +279,8 @@ async def test_agent_loop_syncs_updated_max_iterations_before_run(tmp_path): @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.""" +async def test_drain_pending_does_not_block_while_subagents_running(tmp_path): + """_drain_pending should ignore running workers unless user messages are queued.""" from nanobot.agent.loop import AgentLoop from nanobot.bus.events import InboundMessage from nanobot.bus.queue import MessageBus @@ -336,31 +336,24 @@ async def test_drain_pending_blocks_while_subagents_running(tmp_path): 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()) + # Running subagents alone must not keep the current turn alive. + results = await asyncio.wait_for(injection_callback(), timeout=1.0) + assert results == [] - # Let the task enter the blocking queue wait. - await asyncio.sleep(0) - - # 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) + # Real follow-up messages still use the ordinary pending queue path. await pending_queue.put(InboundMessage( - sender_id="subagent", + sender_id="user", channel="test", chat_id="c1", - content="Sub-agent result", + content="User follow-up", media=None, metadata={}, )) - # Should unblock and return results - results = await asyncio.wait_for(drain_task, timeout=2.0) + results = await asyncio.wait_for(injection_callback(), timeout=1.0) assert len(results) >= 1 assert results[0]["role"] == "user" - assert "Sub-agent result" in str(results[0]["content"]) + assert "User follow-up" in str(results[0]["content"]) # Cleanup hang_task.cancel() @@ -417,8 +410,8 @@ async def test_drain_pending_no_block_when_no_subagents(tmp_path): @pytest.mark.asyncio -async def test_drain_pending_timeout(tmp_path): - """_drain_pending should return empty after timeout when sub-agents hang.""" +async def test_drain_pending_does_not_wait_for_hung_subagents(tmp_path): + """_drain_pending should not call asyncio.wait_for for hung subagents.""" from nanobot.agent.loop import AgentLoop from nanobot.bus.queue import MessageBus from nanobot.session.manager import Session @@ -467,14 +460,10 @@ async def test_drain_pending_timeout(tmp_path): assert injection_callback is not None - # Patch the timeout path without leaking the queue.get() coroutine. - async def _timeout(awaitable, timeout): - awaitable.close() - raise asyncio.TimeoutError - - with patch("nanobot.agent.loop.asyncio.wait_for", side_effect=_timeout): + with patch("nanobot.agent.loop.asyncio.wait_for") as wait_for: results = await injection_callback() assert results == [] + wait_for.assert_not_called() # Cleanup hang_task.cancel()