mirror of
https://github.com/HKUDS/nanobot.git
synced 2026-08-04 16:38:49 +00:00
feat(agent): persist subagent result delivery
This commit is contained in:
parent
4a6853f0ff
commit
43d592f8e4
@ -29,7 +29,7 @@ def session_extra(metadata: Mapping[str, Any] | None) -> dict[str, Any]:
|
||||
|
||||
def runtime_lines(state: Any, msg: Any, workspace: Path, *, skip: bool = False) -> list[str]:
|
||||
"""Return model-visible runtime annotations for turn-attached capabilities."""
|
||||
return [
|
||||
lines = [
|
||||
*cli_app_utils.runtime_lines(msg, workspace, skip=skip),
|
||||
*mcp_tools.runtime_lines(
|
||||
msg,
|
||||
@ -38,6 +38,11 @@ def runtime_lines(state: Any, msg: Any, workspace: Path, *, skip: bool = False)
|
||||
skip=skip,
|
||||
),
|
||||
]
|
||||
if not skip and getattr(state, "subagents", None) is not None:
|
||||
session_key = getattr(msg, "session_key", None)
|
||||
if session_key:
|
||||
lines.extend(state.subagents.runtime_status_lines(session_key))
|
||||
return lines
|
||||
|
||||
|
||||
async def connect_mcp(state: Any, tools: ToolRegistry) -> None:
|
||||
|
||||
@ -25,6 +25,10 @@ from nanobot.agent.memory import Consolidator
|
||||
from nanobot.agent.progress_hook import AgentProgressHook
|
||||
from nanobot.agent.runner import _MAX_INJECTIONS_PER_TURN, AgentRunner, AgentRunSpec
|
||||
from nanobot.agent.subagent import SubagentManager
|
||||
from nanobot.agent.subagent_delivery import (
|
||||
build_subagent_result_continuation,
|
||||
materialize_subagent_result_continuation,
|
||||
)
|
||||
from nanobot.agent.tools.context import RequestContext, bind_request_context, reset_request_context
|
||||
from nanobot.agent.tools.file_state import FileStateStore, bind_file_states, reset_file_states
|
||||
from nanobot.agent.tools.message import MessageTool
|
||||
@ -287,6 +291,7 @@ class AgentLoop:
|
||||
max_iterations=self.max_iterations,
|
||||
max_concurrent_subagents=max_concurrent_subagents,
|
||||
llm_wall_timeout_for_session=lambda sk: runner_wall_llm_timeout_s(self.sessions, sk),
|
||||
on_result_ready=self._on_subagent_result_ready,
|
||||
)
|
||||
self._unified_session = unified_session
|
||||
self._max_messages = max_messages if max_messages > 0 else 120
|
||||
@ -548,6 +553,21 @@ class AgentLoop:
|
||||
"""Build a progress callback that publishes to the message bus."""
|
||||
return build_bus_progress_callback(self.bus, msg)
|
||||
|
||||
async def _on_subagent_result_ready(self, result: Any) -> None:
|
||||
"""Wake the owning session when a subagent result becomes ready."""
|
||||
msg = build_subagent_result_continuation(result)
|
||||
queue = self._pending_queues.get(result.session_key)
|
||||
if queue is not None:
|
||||
try:
|
||||
queue.put_nowait(msg)
|
||||
return
|
||||
except asyncio.QueueFull:
|
||||
logger.warning(
|
||||
"Pending queue full for subagent result in session {}; queueing fresh turn",
|
||||
result.session_key,
|
||||
)
|
||||
await self.bus.publish_inbound(msg)
|
||||
|
||||
async def _build_retry_wait_callback(
|
||||
self, msg: InboundMessage
|
||||
) -> Callable[[str], Awaitable[None]]:
|
||||
@ -750,9 +770,15 @@ class AgentLoop:
|
||||
items: list[dict[str, Any]] = []
|
||||
while len(items) < limit:
|
||||
try:
|
||||
items.append(_to_user_message(pending_queue.get_nowait()))
|
||||
pending_msg = pending_queue.get_nowait()
|
||||
except asyncio.QueueEmpty:
|
||||
break
|
||||
pending_msg = await materialize_subagent_result_continuation(
|
||||
pending_msg,
|
||||
session_key=active_session_key or pending_msg.session_key,
|
||||
subagents=self.subagents,
|
||||
)
|
||||
items.append(_to_user_message(pending_msg))
|
||||
|
||||
return items
|
||||
|
||||
@ -1409,6 +1435,11 @@ class AgentLoop:
|
||||
ctx.session,
|
||||
replay_max_messages=self._max_messages,
|
||||
)
|
||||
ctx.msg = await materialize_subagent_result_continuation(
|
||||
ctx.msg,
|
||||
session_key=ctx.session_key,
|
||||
subagents=self.subagents,
|
||||
)
|
||||
self._set_tool_context(
|
||||
ctx.msg.channel,
|
||||
ctx.msg.chat_id,
|
||||
|
||||
@ -1,12 +1,19 @@
|
||||
"""Mailbox primitives for manager-worker task coordination."""
|
||||
"""Durable mailbox primitives for manager-worker task coordination."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
import uuid
|
||||
from contextlib import suppress
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from nanobot.utils.helpers import ensure_dir, safe_filename
|
||||
|
||||
TaskState = str # running | completed | failed | cancelled
|
||||
MailboxReadState = str # ready | running | not_found | consumed | timeout
|
||||
|
||||
@ -75,40 +82,36 @@ class _TaskRecord:
|
||||
|
||||
|
||||
class MailboxStore:
|
||||
"""In-memory mailbox for worker task/result records.
|
||||
"""Durable task mailbox for local subagent coordination.
|
||||
|
||||
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.
|
||||
JSON files are the source of truth. The condition variable only wakes
|
||||
waiters inside this process; persisted records remain readable after a
|
||||
manager restart.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._records: dict[str, _TaskRecord] = {}
|
||||
self._session_tasks: dict[str, set[str]] = {}
|
||||
self._dedupe_keys: set[str] = set()
|
||||
def __init__(self, workspace: str | Path, *, root: str | Path | None = None) -> None:
|
||||
base = Path(root).expanduser() if root is not None else Path(workspace) / "tasks" / "subagents"
|
||||
self.root = ensure_dir(base)
|
||||
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:
|
||||
path, record = self._load_by_task_id(request.task_id, session_key=request.session_key)
|
||||
if record is not None:
|
||||
return
|
||||
self._records[request.task_id] = _TaskRecord(request=request)
|
||||
self._session_tasks.setdefault(request.session_key, set()).add(request.task_id)
|
||||
path = self._record_path(request.session_key, request.task_id)
|
||||
self._write_record(path, _TaskRecord(request=request))
|
||||
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.
|
||||
Returns ``True`` when this call writes a new terminal result and
|
||||
``False`` when 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)
|
||||
path, record = self._load_by_task_id(result.task_id, session_key=result.session_key)
|
||||
if record is None:
|
||||
request = TaskRequest(
|
||||
task_id=result.task_id,
|
||||
@ -119,17 +122,15 @@ class MailboxStore:
|
||||
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)
|
||||
path = self._record_path(result.session_key, result.task_id)
|
||||
elif record.result is not None or record.state != "running":
|
||||
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._write_record(path, record)
|
||||
self._changed.notify_all()
|
||||
return True
|
||||
|
||||
@ -142,12 +143,8 @@ class MailboxStore:
|
||||
) -> 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:
|
||||
path, record = self._load_by_task_id(task_id, session_key=session_key)
|
||||
if record is None or record.result is not None or record.state != "running":
|
||||
return False
|
||||
result = TaskResult(
|
||||
task_id=task_id,
|
||||
@ -162,7 +159,7 @@ class MailboxStore:
|
||||
record.completed_at = result.completed_at
|
||||
record.state = "cancelled"
|
||||
record.error = reason
|
||||
self._dedupe_keys.add(task_id)
|
||||
self._write_record(path, record)
|
||||
self._changed.notify_all()
|
||||
return True
|
||||
|
||||
@ -174,20 +171,25 @@ class MailboxStore:
|
||||
) -> 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)]
|
||||
return self.snapshot_sync(session_key, task_id=task_id)
|
||||
|
||||
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
|
||||
def snapshot_sync(
|
||||
self,
|
||||
session_key: str,
|
||||
*,
|
||||
task_id: str | None = None,
|
||||
) -> list[TaskSnapshot]:
|
||||
"""Synchronous snapshot used while building runtime context."""
|
||||
if task_id is not None:
|
||||
_, record = self._load_by_task_id(task_id, session_key=session_key)
|
||||
if record is None:
|
||||
return []
|
||||
return [self._snapshot(record)]
|
||||
|
||||
records = self._load_session_records(session_key)
|
||||
snapshots = [self._snapshot(record) for record in 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,
|
||||
@ -217,8 +219,8 @@ class MailboxStore:
|
||||
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:
|
||||
path, record = self._load_by_task_id(task_id, session_key=session_key)
|
||||
if record is None:
|
||||
return MailboxRead("not_found")
|
||||
snapshot = self._snapshot(record)
|
||||
if record.result is None:
|
||||
@ -226,35 +228,168 @@ class MailboxStore:
|
||||
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)
|
||||
self._write_record(path, record)
|
||||
return MailboxRead("ready", task=self._snapshot(record), result=record.result)
|
||||
|
||||
ids = self._session_tasks.get(session_key, set())
|
||||
records = [
|
||||
self._records[tid]
|
||||
for tid in ids
|
||||
if tid in self._records
|
||||
]
|
||||
records_with_paths = self._load_session_records_with_paths(session_key)
|
||||
ready = [
|
||||
record
|
||||
for record in records
|
||||
(path, record)
|
||||
for path, record in records_with_paths
|
||||
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]
|
||||
ready.sort(key=lambda item: (
|
||||
item[1].completed_at or item[1].request.created_at,
|
||||
item[1].request.task_id,
|
||||
))
|
||||
path, record = ready[0]
|
||||
record.consumed_at = time.time()
|
||||
self._write_record(path, record)
|
||||
return MailboxRead("ready", task=self._snapshot(record), result=record.result)
|
||||
|
||||
running = [record for record in records if record.result is None]
|
||||
running = [record for _, record in records_with_paths 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))
|
||||
if records_with_paths:
|
||||
records = [record for _, record in records_with_paths]
|
||||
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")
|
||||
|
||||
def _session_dir(self, session_key: str) -> Path:
|
||||
return self.root / safe_filename(session_key)
|
||||
|
||||
def _record_path(self, session_key: str, task_id: str) -> Path:
|
||||
return ensure_dir(self._session_dir(session_key)) / f"{safe_filename(task_id)}.json"
|
||||
|
||||
def _load_by_task_id(
|
||||
self,
|
||||
task_id: str,
|
||||
*,
|
||||
session_key: str | None = None,
|
||||
) -> tuple[Path, _TaskRecord | None]:
|
||||
if session_key is not None:
|
||||
path = self._record_path(session_key, task_id)
|
||||
return path, self._read_record(path)
|
||||
|
||||
filename = f"{safe_filename(task_id)}.json"
|
||||
for path in self.root.glob(f"*/{filename}"):
|
||||
record = self._read_record(path)
|
||||
if record is not None:
|
||||
return path, record
|
||||
return self.root / "_missing" / filename, None
|
||||
|
||||
def _load_session_records(self, session_key: str) -> list[_TaskRecord]:
|
||||
return [record for _, record in self._load_session_records_with_paths(session_key)]
|
||||
|
||||
def _load_session_records_with_paths(self, session_key: str) -> list[tuple[Path, _TaskRecord]]:
|
||||
directory = self._session_dir(session_key)
|
||||
if not directory.exists():
|
||||
return []
|
||||
records: list[tuple[Path, _TaskRecord]] = []
|
||||
for path in directory.glob("*.json"):
|
||||
record = self._read_record(path)
|
||||
if record is not None:
|
||||
records.append((path, record))
|
||||
return records
|
||||
|
||||
def _read_record(self, path: Path) -> _TaskRecord | None:
|
||||
if not path.exists():
|
||||
return None
|
||||
try:
|
||||
data = json.loads(path.read_text(encoding="utf-8"))
|
||||
return self._record_from_json(data)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
def _write_record(self, path: Path, record: _TaskRecord) -> None:
|
||||
ensure_dir(path.parent)
|
||||
payload = json.dumps(self._record_to_json(record), ensure_ascii=False, indent=2)
|
||||
tmp = path.with_name(f".{path.name}.{uuid.uuid4().hex}.tmp")
|
||||
try:
|
||||
with open(tmp, "w", encoding="utf-8") as f:
|
||||
f.write(payload)
|
||||
f.write("\n")
|
||||
with suppress(OSError):
|
||||
os.fsync(f.fileno())
|
||||
os.replace(tmp, path)
|
||||
with suppress(OSError):
|
||||
fd = os.open(str(path.parent), os.O_RDONLY)
|
||||
try:
|
||||
os.fsync(fd)
|
||||
finally:
|
||||
os.close(fd)
|
||||
finally:
|
||||
tmp.unlink(missing_ok=True)
|
||||
|
||||
@staticmethod
|
||||
def _record_to_json(record: _TaskRecord) -> dict[str, Any]:
|
||||
result = record.result
|
||||
return {
|
||||
"version": 1,
|
||||
"task_id": record.request.task_id,
|
||||
"session_key": record.request.session_key,
|
||||
"label": record.request.label,
|
||||
"task": record.request.task,
|
||||
"origin": record.request.origin,
|
||||
"state": record.state,
|
||||
"result": None if result is None else {
|
||||
"task_id": result.task_id,
|
||||
"session_key": result.session_key,
|
||||
"label": result.label,
|
||||
"task": result.task,
|
||||
"status": result.status,
|
||||
"content": result.content,
|
||||
"sender": result.sender,
|
||||
"completed_at": result.completed_at,
|
||||
"dedupe_key": result.dedupe_key,
|
||||
"metadata": result.metadata,
|
||||
},
|
||||
"consumed_at": record.consumed_at,
|
||||
"created_at": record.request.created_at,
|
||||
"completed_at": record.completed_at,
|
||||
"updated_at": time.time(),
|
||||
"error": record.error,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _record_from_json(data: dict[str, Any]) -> _TaskRecord:
|
||||
request = TaskRequest(
|
||||
task_id=str(data["task_id"]),
|
||||
session_key=str(data["session_key"]),
|
||||
label=str(data.get("label") or data["task_id"]),
|
||||
task=str(data.get("task") or ""),
|
||||
origin=dict(data.get("origin") or {}),
|
||||
created_at=float(data.get("created_at") or time.time()),
|
||||
)
|
||||
raw_result = data.get("result")
|
||||
result = None
|
||||
if isinstance(raw_result, dict):
|
||||
result = TaskResult(
|
||||
task_id=str(raw_result.get("task_id") or request.task_id),
|
||||
session_key=str(raw_result.get("session_key") or request.session_key),
|
||||
label=str(raw_result.get("label") or request.label),
|
||||
task=str(raw_result.get("task") or request.task),
|
||||
status=str(raw_result.get("status") or "error"),
|
||||
content=str(raw_result.get("content") or ""),
|
||||
sender=str(raw_result.get("sender") or "subagent"),
|
||||
completed_at=float(raw_result.get("completed_at") or time.time()),
|
||||
dedupe_key=raw_result.get("dedupe_key"),
|
||||
metadata=dict(raw_result.get("metadata") or {}),
|
||||
)
|
||||
return _TaskRecord(
|
||||
request=request,
|
||||
state=str(data.get("state") or "running"),
|
||||
result=result,
|
||||
consumed_at=data.get("consumed_at"),
|
||||
completed_at=data.get("completed_at"),
|
||||
error=data.get("error"),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _state_for_result(status: str) -> TaskState:
|
||||
if status == "ok":
|
||||
|
||||
@ -7,7 +7,7 @@ import uuid
|
||||
from contextlib import suppress
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable
|
||||
from typing import Any, Awaitable, Callable
|
||||
|
||||
from loguru import logger
|
||||
|
||||
@ -89,6 +89,7 @@ class SubagentManager:
|
||||
max_concurrent_subagents: int | None = None,
|
||||
llm_wall_timeout_for_session: Callable[[str | None], float | None] | None = None,
|
||||
mailbox: MailboxStore | None = None,
|
||||
on_result_ready: Callable[[TaskResult], Awaitable[None]] | None = None,
|
||||
):
|
||||
defaults = AgentDefaults()
|
||||
self.provider = provider
|
||||
@ -111,7 +112,8 @@ class SubagentManager:
|
||||
)
|
||||
self.runner = AgentRunner(provider)
|
||||
self._llm_wall_timeout_for_session = llm_wall_timeout_for_session
|
||||
self.mailbox = mailbox or MailboxStore()
|
||||
self.mailbox = mailbox or MailboxStore(workspace)
|
||||
self._on_result_ready = on_result_ready
|
||||
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, ...}
|
||||
@ -332,7 +334,7 @@ class SubagentManager:
|
||||
if origin_message_id:
|
||||
metadata["origin_message_id"] = origin_message_id
|
||||
|
||||
written = await self.mailbox.record_result(TaskResult(
|
||||
task_result = TaskResult(
|
||||
task_id=task_id,
|
||||
session_key=override,
|
||||
label=label,
|
||||
@ -341,7 +343,8 @@ class SubagentManager:
|
||||
content=result,
|
||||
dedupe_key=task_id,
|
||||
metadata=metadata,
|
||||
))
|
||||
)
|
||||
written = await self.mailbox.record_result(task_result)
|
||||
|
||||
if written:
|
||||
logger.debug(
|
||||
@ -349,6 +352,11 @@ class SubagentManager:
|
||||
task_id,
|
||||
override,
|
||||
)
|
||||
if self._on_result_ready is not None:
|
||||
try:
|
||||
await self._on_result_ready(task_result)
|
||||
except Exception:
|
||||
logger.exception("Subagent result-ready callback failed")
|
||||
else:
|
||||
logger.debug("Subagent [{}] result already recorded", task_id)
|
||||
|
||||
@ -449,6 +457,41 @@ class SubagentManager:
|
||||
timeout_seconds=timeout_seconds,
|
||||
)
|
||||
|
||||
def runtime_status_lines(self, session_key: str, *, limit: int = 8) -> list[str]:
|
||||
"""Return compact model-visible task status lines for runtime context."""
|
||||
snapshots = self.mailbox.snapshot_sync(session_key)
|
||||
if not snapshots:
|
||||
return []
|
||||
|
||||
now = time.time()
|
||||
ordered = sorted(
|
||||
snapshots,
|
||||
key=lambda item: (
|
||||
item.consumed_at is not None,
|
||||
item.completed_at is None,
|
||||
item.created_at,
|
||||
item.task_id,
|
||||
),
|
||||
)
|
||||
lines = ["Subagent tasks:"]
|
||||
for snapshot in ordered[: max(0, limit)]:
|
||||
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"
|
||||
elapsed = max(0, int((snapshot.completed_at or now) - snapshot.created_at))
|
||||
label = " ".join(snapshot.label.split())
|
||||
if len(label) > 48:
|
||||
label = label[:45] + "..."
|
||||
lines.append(
|
||||
f"- {snapshot.task_id}: {state}, label=\"{label}\", elapsed={elapsed}s"
|
||||
)
|
||||
remaining = len(ordered) - limit
|
||||
if remaining > 0:
|
||||
lines.append(f"- ... {remaining} more subagent task(s)")
|
||||
return lines
|
||||
|
||||
def get_running_count(self) -> int:
|
||||
"""Return the number of currently running subagents."""
|
||||
return len(self._running_tasks)
|
||||
|
||||
94
nanobot/agent/subagent_delivery.py
Normal file
94
nanobot/agent/subagent_delivery.py
Normal file
@ -0,0 +1,94 @@
|
||||
"""Runtime delivery helpers for completed subagent task results."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import dataclasses
|
||||
from typing import Any
|
||||
|
||||
from nanobot.bus.events import InboundMessage
|
||||
from nanobot.session import turn_continuation
|
||||
|
||||
_FORWARDED_METADATA_KEYS = frozenset({
|
||||
"message_id",
|
||||
"origin_message_id",
|
||||
"_wants_stream",
|
||||
"webui",
|
||||
"slack",
|
||||
})
|
||||
|
||||
|
||||
def build_subagent_result_continuation(result: Any) -> InboundMessage:
|
||||
"""Build an internal inbound wake-up for a ready subagent result."""
|
||||
metadata = dict(result.metadata or {})
|
||||
channel = str(metadata.get("origin_channel") or "")
|
||||
chat_id = str(metadata.get("origin_chat_id") or "")
|
||||
if not channel or not chat_id:
|
||||
channel, chat_id = _channel_chat_from_session_key(result.session_key)
|
||||
|
||||
wake_meta = turn_continuation.subagent_result_continuation_metadata(
|
||||
{key: value for key, value in metadata.items() if key in _FORWARDED_METADATA_KEYS},
|
||||
task_id=result.task_id,
|
||||
)
|
||||
return InboundMessage(
|
||||
channel=channel,
|
||||
sender_id="system:continuation",
|
||||
chat_id=chat_id,
|
||||
content=(
|
||||
"A subagent task result is ready. The runtime will attach the "
|
||||
"result to this continuation turn."
|
||||
),
|
||||
metadata=wake_meta,
|
||||
session_key_override=result.session_key,
|
||||
)
|
||||
|
||||
|
||||
async def materialize_subagent_result_continuation(
|
||||
msg: InboundMessage,
|
||||
*,
|
||||
session_key: str,
|
||||
subagents: Any,
|
||||
) -> InboundMessage:
|
||||
"""Replace a subagent-result continuation placeholder with the mailbox result."""
|
||||
task_id = turn_continuation.subagent_result_continuation_task_id(msg.metadata)
|
||||
if not task_id:
|
||||
return msg
|
||||
read = await subagents.wait_for_result(
|
||||
session_key,
|
||||
task_id=task_id,
|
||||
timeout_seconds=0,
|
||||
)
|
||||
return dataclasses.replace(msg, content=_subagent_result_continuation_content(read, task_id))
|
||||
|
||||
|
||||
def _channel_chat_from_session_key(session_key: str) -> tuple[str, str]:
|
||||
channel, _, chat_id = session_key.partition(":")
|
||||
return channel or "cli", chat_id or "direct"
|
||||
|
||||
|
||||
def _subagent_result_continuation_content(read: Any, requested_task_id: str) -> str:
|
||||
if read.state == "ready" and read.result is not None:
|
||||
status_text = {
|
||||
"ok": "completed",
|
||||
"error": "failed",
|
||||
"cancelled": "cancelled",
|
||||
}.get(read.result.status, read.result.status)
|
||||
return (
|
||||
"A subagent result was delivered by the runtime. Use this result "
|
||||
"as authoritative context for the next answer; do not mention the "
|
||||
"internal continuation boundary.\n\n"
|
||||
f"Subagent [{read.result.label}] "
|
||||
f"(id: {read.result.task_id}, status: {status_text})\n\n"
|
||||
f"Task:\n{read.result.task}\n\n"
|
||||
f"Result:\n{read.result.content}"
|
||||
)
|
||||
if read.state == "consumed":
|
||||
return (
|
||||
f"Subagent task {requested_task_id} already has a consumed result. "
|
||||
"Check poll_subagents if you need its current status."
|
||||
)
|
||||
if read.state == "running":
|
||||
return (
|
||||
f"Subagent task {requested_task_id} is still running. "
|
||||
"Use poll_subagents or wait_subagents if you need to block."
|
||||
)
|
||||
return f"Subagent task {requested_task_id} result is not available ({read.state})."
|
||||
@ -25,6 +25,8 @@ INTERNAL_CONTINUATION_RUN_STARTED_AT_META = "_internal_continuation_run_started_
|
||||
SKIP_USER_PERSIST_META = "_skip_user_persist"
|
||||
|
||||
_GOAL_CONTINUATION_KIND = "sustained_goal"
|
||||
SUBAGENT_RESULT_CONTINUATION_KIND = "subagent_result"
|
||||
SUBAGENT_RESULT_TASK_ID_META = "_subagent_result_task_id"
|
||||
_GOAL_CONTINUATION_SENDER = "system:continuation"
|
||||
_GOAL_CONTINUATION_ROUNDS_KEY = "_sustained_goal_continuation_rounds"
|
||||
_MAX_GOAL_CONTINUATION_ROUNDS = 12
|
||||
@ -58,6 +60,38 @@ def internal_continuation_run_started_at(metadata: Mapping[str, Any] | None) ->
|
||||
return started_at if started_at > 0 else None
|
||||
|
||||
|
||||
def subagent_result_continuation_inbound(metadata: Mapping[str, Any] | None) -> bool:
|
||||
"""True for an internal continuation caused by a ready subagent result."""
|
||||
return bool(
|
||||
internal_continuation_inbound(metadata)
|
||||
and metadata.get(INTERNAL_CONTINUATION_KIND_META) == SUBAGENT_RESULT_CONTINUATION_KIND
|
||||
)
|
||||
|
||||
|
||||
def subagent_result_continuation_task_id(metadata: Mapping[str, Any] | None) -> str | None:
|
||||
"""Return the ready subagent task id carried by a continuation message."""
|
||||
if not subagent_result_continuation_inbound(metadata):
|
||||
return None
|
||||
value = metadata.get(SUBAGENT_RESULT_TASK_ID_META) if metadata else None
|
||||
return value if isinstance(value, str) and value else None
|
||||
|
||||
|
||||
def subagent_result_continuation_metadata(
|
||||
message_metadata: Mapping[str, Any] | None,
|
||||
*,
|
||||
task_id: str,
|
||||
run_started_at: float | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Build sanitized metadata for a subagent-result continuation turn."""
|
||||
metadata = _internal_continuation_metadata(
|
||||
message_metadata,
|
||||
kind=SUBAGENT_RESULT_CONTINUATION_KIND,
|
||||
run_started_at=run_started_at,
|
||||
)
|
||||
metadata[SUBAGENT_RESULT_TASK_ID_META] = task_id
|
||||
return metadata
|
||||
|
||||
|
||||
def should_persist_user_message(metadata: Mapping[str, Any] | None) -> bool:
|
||||
"""Return whether this inbound message should be persisted as user input."""
|
||||
if metadata and metadata.get(SKIP_USER_PERSIST_META) is True:
|
||||
@ -223,11 +257,12 @@ def _increment_goal_continuation_round(session_metadata: MutableMapping[str, Any
|
||||
def _internal_continuation_metadata(
|
||||
message_metadata: Mapping[str, Any] | None,
|
||||
*,
|
||||
kind: str = _GOAL_CONTINUATION_KIND,
|
||||
run_started_at: float | None = None,
|
||||
) -> dict[str, Any]:
|
||||
metadata = dict(message_metadata or {})
|
||||
metadata[INTERNAL_CONTINUATION_META] = True
|
||||
metadata[INTERNAL_CONTINUATION_KIND_META] = _GOAL_CONTINUATION_KIND
|
||||
metadata[INTERNAL_CONTINUATION_KIND_META] = kind
|
||||
if run_started_at is not None:
|
||||
metadata[INTERNAL_CONTINUATION_RUN_STARTED_AT_META] = float(run_started_at)
|
||||
for key in _STRIPPED_INBOUND_META_KEYS:
|
||||
|
||||
@ -14,8 +14,10 @@ from nanobot.providers.base import LLMResponse
|
||||
from nanobot.session.goal_state import GOAL_STATE_KEY
|
||||
from nanobot.session.manager import Session, SessionManager
|
||||
from nanobot.session.turn_continuation import (
|
||||
INTERNAL_CONTINUATION_KIND_META,
|
||||
INTERNAL_CONTINUATION_META,
|
||||
INTERNAL_CONTINUATION_RUN_STARTED_AT_META,
|
||||
SUBAGENT_RESULT_CONTINUATION_KIND,
|
||||
)
|
||||
from nanobot.session.webui_turns import (
|
||||
TITLE_GENERATION_MAX_TOKENS,
|
||||
@ -864,6 +866,100 @@ async def test_websocket_internal_continuation_keeps_single_visible_run(
|
||||
assert isinstance(turn_end[0].metadata.get("latency_ms"), int)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_runtime_context_lists_ready_subagent_result(tmp_path: Path) -> None:
|
||||
loop = _make_full_loop(tmp_path)
|
||||
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
|
||||
|
||||
await loop.subagents._announce_result(
|
||||
"sub-ready",
|
||||
"research",
|
||||
"look up the answer",
|
||||
"worker answer",
|
||||
{"channel": "cli", "chat_id": "test", "session_key": "cli:test"},
|
||||
"ok",
|
||||
)
|
||||
|
||||
seen: dict[str, list[dict]] = {}
|
||||
|
||||
async def fake_run_agent_loop(initial_messages, **_kwargs):
|
||||
seen["initial_messages"] = initial_messages
|
||||
return (
|
||||
"done",
|
||||
[],
|
||||
[*initial_messages, {"role": "assistant", "content": "done"}],
|
||||
"completed",
|
||||
False,
|
||||
)
|
||||
|
||||
loop._run_agent_loop = fake_run_agent_loop # type: ignore[method-assign]
|
||||
|
||||
await loop._process_message(
|
||||
InboundMessage(channel="cli", sender_id="user", chat_id="test", content="continue")
|
||||
)
|
||||
|
||||
rendered = "\n".join(str(msg.get("content", "")) for msg in seen["initial_messages"])
|
||||
assert "Subagent tasks:" in rendered
|
||||
assert "sub-ready: completed, result ready" in rendered
|
||||
assert "worker answer" not in rendered
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_subagent_result_continuation_delivers_result_without_user_history(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
loop = _make_full_loop(tmp_path)
|
||||
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
|
||||
|
||||
await loop.subagents._announce_result(
|
||||
"sub-deliver",
|
||||
"worker",
|
||||
"calculate the answer",
|
||||
"the worker result",
|
||||
{"channel": "cli", "chat_id": "test", "session_key": "cli:test"},
|
||||
"ok",
|
||||
)
|
||||
queued = await asyncio.wait_for(loop.bus.consume_inbound(), timeout=0.5)
|
||||
|
||||
assert queued.metadata[INTERNAL_CONTINUATION_META] is True
|
||||
assert queued.metadata[INTERNAL_CONTINUATION_KIND_META] == SUBAGENT_RESULT_CONTINUATION_KIND
|
||||
assert "the worker result" not in queued.content
|
||||
|
||||
seen: dict[str, list[dict]] = {}
|
||||
|
||||
async def fake_run_agent_loop(initial_messages, **_kwargs):
|
||||
seen["initial_messages"] = initial_messages
|
||||
return (
|
||||
"reported",
|
||||
[],
|
||||
[*initial_messages, {"role": "assistant", "content": "reported"}],
|
||||
"completed",
|
||||
False,
|
||||
)
|
||||
|
||||
loop._run_agent_loop = fake_run_agent_loop # type: ignore[method-assign]
|
||||
|
||||
response = await loop._process_message(queued, pending_queue=asyncio.Queue())
|
||||
|
||||
assert response is not None
|
||||
assert response.content == "reported"
|
||||
rendered = "\n".join(str(msg.get("content", "")) for msg in seen["initial_messages"])
|
||||
assert "the worker result" in rendered
|
||||
|
||||
read = await loop.subagents.wait_for_result(
|
||||
"cli:test",
|
||||
task_id="sub-deliver",
|
||||
timeout_seconds=0,
|
||||
)
|
||||
assert read.state == "consumed"
|
||||
|
||||
session = loop.sessions.get_or_create("cli:test")
|
||||
assert [
|
||||
{k: v for k, v in m.items() if k in {"role", "content"}}
|
||||
for m in session.messages
|
||||
] == [{"role": "assistant", "content": "reported"}]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_message_uses_context_chat_id_for_runtime_prompt(tmp_path: Path) -> None:
|
||||
loop = _make_full_loop(tmp_path)
|
||||
|
||||
@ -64,6 +64,30 @@ async def test_wait_subagents_returns_result_once(tmp_path: Path) -> None:
|
||||
assert "already consumed" in second
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_wait_subagents_reads_result_after_manager_recreation(tmp_path: Path) -> None:
|
||||
mgr = _manager(tmp_path)
|
||||
mgr.runner.run = AsyncMock(
|
||||
return_value=AgentRunResult(final_content="durable worker result", messages=[], stop_reason="completed")
|
||||
)
|
||||
|
||||
await mgr.spawn("do durable work", label="worker", session_key="cli:test")
|
||||
task_id = next(iter(mgr._running_tasks))
|
||||
await _drain(mgr)
|
||||
|
||||
recreated = _manager(tmp_path)
|
||||
wait_tool = WaitSubagentsTool(recreated)
|
||||
poll_tool = PollSubagentsTool(recreated)
|
||||
_bind(wait_tool)
|
||||
_bind(poll_tool)
|
||||
|
||||
first = await wait_tool.execute(task_id=task_id, timeout_seconds=0)
|
||||
after = await poll_tool.execute(task_id=task_id)
|
||||
|
||||
assert "durable worker result" in first
|
||||
assert "result consumed" in after
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_poll_subagents_reports_running_completed_and_not_found(tmp_path: Path) -> None:
|
||||
mgr = _manager(tmp_path)
|
||||
|
||||
@ -14,12 +14,16 @@ from nanobot.session.turn_continuation import (
|
||||
INTERNAL_CONTINUATION_META,
|
||||
INTERNAL_CONTINUATION_PENDING_META,
|
||||
INTERNAL_CONTINUATION_RUN_STARTED_AT_META,
|
||||
SUBAGENT_RESULT_CONTINUATION_KIND,
|
||||
_save_skip_for_turn,
|
||||
internal_continuation_pending,
|
||||
internal_continuation_run_started_at,
|
||||
maybe_continue_turn,
|
||||
should_finalize_on_max_iterations,
|
||||
should_stream_budget_response,
|
||||
subagent_result_continuation_inbound,
|
||||
subagent_result_continuation_metadata,
|
||||
subagent_result_continuation_task_id,
|
||||
)
|
||||
|
||||
|
||||
@ -165,3 +169,23 @@ def test_save_skip_unchanged_for_standalone_current_message():
|
||||
history_count=1,
|
||||
user_persisted_early=False,
|
||||
) == 2
|
||||
|
||||
|
||||
def test_subagent_result_continuation_metadata():
|
||||
meta = subagent_result_continuation_metadata(
|
||||
{
|
||||
"message_id": "msg-1",
|
||||
"_stream_id": "old-stream",
|
||||
"_stream_delta": True,
|
||||
},
|
||||
task_id="sub-1",
|
||||
run_started_at=42.0,
|
||||
)
|
||||
|
||||
assert meta[INTERNAL_CONTINUATION_META] is True
|
||||
assert meta[INTERNAL_CONTINUATION_KIND_META] == SUBAGENT_RESULT_CONTINUATION_KIND
|
||||
assert meta[INTERNAL_CONTINUATION_RUN_STARTED_AT_META] == 42.0
|
||||
assert subagent_result_continuation_inbound(meta)
|
||||
assert subagent_result_continuation_task_id(meta) == "sub-1"
|
||||
assert "_stream_id" not in meta
|
||||
assert "_stream_delta" not in meta
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user