mirror of
https://github.com/HKUDS/nanobot.git
synced 2026-08-04 16:38:49 +00:00
feat(subagent): add mailbox-backed worker results
This commit is contained in:
parent
c29601d303
commit
3b03cc2079
290
.agent/subagent-mailbox-mvp-plan.md
Normal file
290
.agent/subagent-mailbox-mvp-plan.md
Normal 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.
|
||||
@ -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
280
nanobot/agent/mailbox.py
Normal 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,
|
||||
)
|
||||
@ -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)
|
||||
|
||||
@ -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,
|
||||
|
||||
207
nanobot/agent/tools/subagent_mailbox.py
Normal file
207
nanobot/agent/tools/subagent_mailbox.py
Normal 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}."
|
||||
@ -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"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@ -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"
|
||||
|
||||
118
tests/agent/tools/test_subagent_mailbox_tools.py
Normal file
118
tests/agent/tools/test_subagent_mailbox_tools.py
Normal 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
|
||||
@ -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()
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user