feat(subagent): add mailbox-backed worker results

This commit is contained in:
chengyongru 2026-06-05 18:54:29 +08:00 committed by chengyongru
parent c29601d303
commit 3b03cc2079
10 changed files with 1074 additions and 126 deletions

View File

@ -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.

View File

@ -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

280
nanobot/agent/mailbox.py Normal file
View File

@ -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,
)

View File

@ -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)

View File

@ -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,

View File

@ -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}."

View File

@ -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"
# ---------------------------------------------------------------------------

View File

@ -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"

View File

@ -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

View File

@ -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()