| 1 | """Subagent manager for background task execution.""" |
| 2 | |
| 3 | import asyncio |
| 4 | import json |
| 5 | import time |
| 6 | import uuid |
| 7 | from dataclasses import dataclass, field |
| 8 | from pathlib import Path |
| 9 | from typing import Any |
| 10 | |
| 11 | from loguru import logger |
| 12 | |
| 13 | from nanobot.agent.hook import AgentHook, AgentHookContext |
| 14 | from nanobot.utils.prompt_templates import render_template |
| 15 | from nanobot.agent.runner import AgentRunSpec, AgentRunner |
| 16 | from nanobot.agent.skills import BUILTIN_SKILLS_DIR |
| 17 | from nanobot.agent.tools.filesystem import EditFileTool, ListDirTool, ReadFileTool, WriteFileTool |
| 18 | from nanobot.agent.tools.registry import ToolRegistry |
| 19 | from nanobot.agent.tools.search import GlobTool, GrepTool |
| 20 | from nanobot.agent.tools.shell import ExecTool |
| 21 | from nanobot.agent.tools.web import WebFetchTool, WebSearchTool |
| 22 | from nanobot.bus.events import InboundMessage |
| 23 | from nanobot.bus.queue import MessageBus |
| 24 | from nanobot.config.schema import ExecToolConfig, WebToolsConfig |
| 25 | from nanobot.providers.base import LLMProvider |
| 26 | |
| 27 | |
| 28 | @dataclass(slots=True) |
| 29 | class SubagentStatus: |
| 30 | """Real-time status of a running subagent.""" |
| 31 | |
| 32 | task_id: str |
| 33 | label: str |
| 34 | task_description: str |
| 35 | started_at: float # time.monotonic() |
| 36 | phase: str = "initializing" # initializing | awaiting_tools | tools_completed | final_response | done | error |
| 37 | iteration: int = 0 |
| 38 | tool_events: list = field(default_factory=list) # [{name, status, detail}, ...] |
| 39 | usage: dict = field(default_factory=dict) # token usage |
| 40 | stop_reason: str | None = None |
| 41 | error: str | None = None |
| 42 | |
| 43 | |
| 44 | class _SubagentHook(AgentHook): |
| 45 | """Hook for subagent execution — logs tool calls and updates status.""" |
| 46 | |
| 47 | def __init__(self, task_id: str, status: SubagentStatus | None = None) -> None: |
| 48 | super().__init__() |
| 49 | self._task_id = task_id |
| 50 | self._status = status |
| 51 | |
| 52 | async def before_execute_tools(self, context: AgentHookContext) -> None: |
| 53 | for tool_call in context.tool_calls: |
| 54 | args_str = json.dumps(tool_call.arguments, ensure_ascii=False) |
| 55 | logger.debug( |
| 56 | "Subagent [{}] executing: {} with arguments: {}", |
| 57 | self._task_id, tool_call.name, args_str, |
| 58 | ) |
| 59 | |
| 60 | async def after_iteration(self, context: AgentHookContext) -> None: |
| 61 | if self._status is None: |
| 62 | return |
| 63 | self._status.iteration = context.iteration |
| 64 | self._status.tool_events = list(context.tool_events) |
| 65 | self._status.usage = dict(context.usage) |
| 66 | if context.error: |
| 67 | self._status.error = str(context.error) |
| 68 | |
| 69 | |
| 70 | class SubagentManager: |
| 71 | """Manages background subagent execution.""" |
| 72 | |
| 73 | def __init__( |
| 74 | self, |
| 75 | provider: LLMProvider, |
| 76 | workspace: Path, |
| 77 | bus: MessageBus, |
| 78 | max_tool_result_chars: int, |
| 79 | model: str | None = None, |
| 80 | web_config: "WebToolsConfig | None" = None, |
| 81 | exec_config: "ExecToolConfig | None" = None, |
| 82 | restrict_to_workspace: bool = False, |
| 83 | disabled_skills: list[str] | None = None, |
| 84 | ): |
| 85 | self.provider = provider |
| 86 | self.workspace = workspace |
| 87 | self.bus = bus |
| 88 | self.model = model or provider.get_default_model() |
| 89 | self.web_config = web_config or WebToolsConfig() |
| 90 | self.max_tool_result_chars = max_tool_result_chars |
| 91 | self.exec_config = exec_config or ExecToolConfig() |
| 92 | self.restrict_to_workspace = restrict_to_workspace |
| 93 | self.disabled_skills = set(disabled_skills or []) |
| 94 | self.runner = AgentRunner(provider) |
| 95 | self._running_tasks: dict[str, asyncio.Task[None]] = {} |
| 96 | self._task_statuses: dict[str, SubagentStatus] = {} |
| 97 | self._session_tasks: dict[str, set[str]] = {} # session_key -> {task_id, ...} |
| 98 | |
| 99 | async def spawn( |
| 100 | self, |
| 101 | task: str, |
| 102 | label: str | None = None, |
| 103 | origin_channel: str = "cli", |
| 104 | origin_chat_id: str = "direct", |
| 105 | session_key: str | None = None, |
| 106 | ) -> str: |
| 107 | """Spawn a subagent to execute a task in the background.""" |
| 108 | task_id = str(uuid.uuid4())[:8] |
| 109 | display_label = label or task[:30] + ("..." if len(task) > 30 else "") |
| 110 | origin = {"channel": origin_channel, "chat_id": origin_chat_id, "session_key": session_key} |
| 111 | |
| 112 | status = SubagentStatus( |
| 113 | task_id=task_id, |
| 114 | label=display_label, |
| 115 | task_description=task, |
| 116 | started_at=time.monotonic(), |
| 117 | ) |
| 118 | self._task_statuses[task_id] = status |
| 119 | |
| 120 | bg_task = asyncio.create_task( |
| 121 | self._run_subagent(task_id, task, display_label, origin, status) |
| 122 | ) |
| 123 | self._running_tasks[task_id] = bg_task |
| 124 | if session_key: |
| 125 | self._session_tasks.setdefault(session_key, set()).add(task_id) |
| 126 | |
| 127 | def _cleanup(_: asyncio.Task) -> None: |
| 128 | self._running_tasks.pop(task_id, None) |
| 129 | self._task_statuses.pop(task_id, None) |
| 130 | if session_key and (ids := self._session_tasks.get(session_key)): |
| 131 | ids.discard(task_id) |
| 132 | if not ids: |
| 133 | del self._session_tasks[session_key] |
| 134 | |
| 135 | bg_task.add_done_callback(_cleanup) |
| 136 | |
| 137 | logger.info("Spawned subagent [{}]: {}", task_id, display_label) |
| 138 | return f"Subagent [{display_label}] started (id: {task_id}). I'll notify you when it completes." |
| 139 | |
| 140 | async def _run_subagent( |
| 141 | self, |
| 142 | task_id: str, |
| 143 | task: str, |
| 144 | label: str, |
| 145 | origin: dict[str, str], |
| 146 | status: SubagentStatus, |
| 147 | ) -> None: |
| 148 | """Execute the subagent task and announce the result.""" |
| 149 | logger.info("Subagent [{}] starting task: {}", task_id, label) |
| 150 | |
| 151 | async def _on_checkpoint(payload: dict) -> None: |
| 152 | status.phase = payload.get("phase", status.phase) |
| 153 | status.iteration = payload.get("iteration", status.iteration) |
| 154 | |
| 155 | try: |
| 156 | # Build subagent tools (no message tool, no spawn tool) |
| 157 | tools = ToolRegistry() |
| 158 | allowed_dir = self.workspace if (self.restrict_to_workspace or self.exec_config.sandbox) else None |
| 159 | extra_read = [BUILTIN_SKILLS_DIR] if allowed_dir else None |
| 160 | tools.register(ReadFileTool(workspace=self.workspace, allowed_dir=allowed_dir, extra_allowed_dirs=extra_read)) |
| 161 | tools.register(WriteFileTool(workspace=self.workspace, allowed_dir=allowed_dir)) |
| 162 | tools.register(EditFileTool(workspace=self.workspace, allowed_dir=allowed_dir)) |
| 163 | tools.register(ListDirTool(workspace=self.workspace, allowed_dir=allowed_dir)) |
| 164 | tools.register(GlobTool(workspace=self.workspace, allowed_dir=allowed_dir)) |
| 165 | tools.register(GrepTool(workspace=self.workspace, allowed_dir=allowed_dir)) |
| 166 | if self.exec_config.enable: |
| 167 | tools.register(ExecTool( |
| 168 | working_dir=str(self.workspace), |
| 169 | timeout=self.exec_config.timeout, |
| 170 | restrict_to_workspace=self.restrict_to_workspace, |
| 171 | sandbox=self.exec_config.sandbox, |
| 172 | path_append=self.exec_config.path_append, |
| 173 | allowed_env_keys=self.exec_config.allowed_env_keys, |
| 174 | )) |
| 175 | if self.web_config.enable: |
| 176 | tools.register(WebSearchTool(config=self.web_config.search, proxy=self.web_config.proxy)) |
| 177 | tools.register(WebFetchTool(proxy=self.web_config.proxy)) |
| 178 | system_prompt = self._build_subagent_prompt() |
| 179 | messages: list[dict[str, Any]] = [ |
| 180 | {"role": "system", "content": system_prompt}, |
| 181 | {"role": "user", "content": task}, |
| 182 | ] |
| 183 | |
| 184 | result = await self.runner.run(AgentRunSpec( |
| 185 | initial_messages=messages, |
| 186 | tools=tools, |
| 187 | model=self.model, |
| 188 | max_iterations=15, |
| 189 | max_tool_result_chars=self.max_tool_result_chars, |
| 190 | hook=_SubagentHook(task_id, status), |
| 191 | max_iterations_message="Task completed but no final response was generated.", |
| 192 | error_message=None, |
| 193 | fail_on_tool_error=True, |
| 194 | checkpoint_callback=_on_checkpoint, |
| 195 | )) |
| 196 | status.phase = "done" |
| 197 | status.stop_reason = result.stop_reason |
| 198 | |
| 199 | if result.stop_reason == "tool_error": |
| 200 | status.tool_events = list(result.tool_events) |
| 201 | await self._announce_result( |
| 202 | task_id, label, task, |
| 203 | self._format_partial_progress(result), |
| 204 | origin, "error", |
| 205 | ) |
| 206 | elif result.stop_reason == "error": |
| 207 | await self._announce_result( |
| 208 | task_id, label, task, |
| 209 | result.error or "Error: subagent execution failed.", |
| 210 | origin, "error", |
| 211 | ) |
| 212 | else: |
| 213 | final_result = result.final_content or "Task completed but no final response was generated." |
| 214 | logger.info("Subagent [{}] completed successfully", task_id) |
| 215 | await self._announce_result(task_id, label, task, final_result, origin, "ok") |
| 216 | |
| 217 | except Exception as e: |
| 218 | status.phase = "error" |
| 219 | status.error = str(e) |
| 220 | logger.error("Subagent [{}] failed: {}", task_id, e) |
| 221 | await self._announce_result(task_id, label, task, f"Error: {e}", origin, "error") |
| 222 | |
| 223 | async def _announce_result( |
| 224 | self, |
| 225 | task_id: str, |
| 226 | label: str, |
| 227 | task: str, |
| 228 | result: str, |
| 229 | origin: dict[str, str], |
| 230 | status: str, |
| 231 | ) -> None: |
| 232 | """Announce the subagent result to the main agent via the message bus.""" |
| 233 | status_text = "completed successfully" if status == "ok" else "failed" |
| 234 | |
| 235 | announce_content = render_template( |
| 236 | "agent/subagent_announce.md", |
| 237 | label=label, |
| 238 | status_text=status_text, |
| 239 | task=task, |
| 240 | result=result, |
| 241 | ) |
| 242 | |
| 243 | # Inject as system message to trigger main agent. |
| 244 | # Use session_key_override to align with the main agent's effective |
| 245 | # session key (which accounts for unified sessions) so the result is |
| 246 | # routed to the correct pending queue (mid-turn injection) instead of |
| 247 | # being dispatched as a competing independent task. |
| 248 | override = origin.get("session_key") or f"{origin['channel']}:{origin['chat_id']}" |
| 249 | msg = InboundMessage( |
| 250 | channel="system", |
| 251 | sender_id="subagent", |
| 252 | chat_id=f"{origin['channel']}:{origin['chat_id']}", |
| 253 | content=announce_content, |
| 254 | session_key_override=override, |
| 255 | metadata={ |
| 256 | "injected_event": "subagent_result", |
| 257 | "subagent_task_id": task_id, |
| 258 | }, |
| 259 | ) |
| 260 | |
| 261 | await self.bus.publish_inbound(msg) |
| 262 | logger.debug("Subagent [{}] announced result to {}:{}", task_id, origin['channel'], origin['chat_id']) |
| 263 | |
| 264 | @staticmethod |
| 265 | def _format_partial_progress(result) -> str: |
| 266 | completed = [e for e in result.tool_events if e["status"] == "ok"] |
| 267 | failure = next((e for e in reversed(result.tool_events) if e["status"] == "error"), None) |
| 268 | lines: list[str] = [] |
| 269 | if completed: |
| 270 | lines.append("Completed steps:") |
| 271 | for event in completed[-3:]: |
| 272 | lines.append(f"- {event['name']}: {event['detail']}") |
| 273 | if failure: |
| 274 | if lines: |
| 275 | lines.append("") |
| 276 | lines.append("Failure:") |
| 277 | lines.append(f"- {failure['name']}: {failure['detail']}") |
| 278 | if result.error and not failure: |
| 279 | if lines: |
| 280 | lines.append("") |
| 281 | lines.append("Failure:") |
| 282 | lines.append(f"- {result.error}") |
| 283 | return "\n".join(lines) or (result.error or "Error: subagent execution failed.") |
| 284 | |
| 285 | def _build_subagent_prompt(self) -> str: |
| 286 | """Build a focused system prompt for the subagent.""" |
| 287 | from nanobot.agent.context import ContextBuilder |
| 288 | from nanobot.agent.skills import SkillsLoader |
| 289 | |
| 290 | time_ctx = ContextBuilder._build_runtime_context(None, None) |
| 291 | skills_summary = SkillsLoader( |
| 292 | self.workspace, |
| 293 | disabled_skills=self.disabled_skills, |
| 294 | ).build_skills_summary() |
| 295 | return render_template( |
| 296 | "agent/subagent_system.md", |
| 297 | time_ctx=time_ctx, |
| 298 | workspace=str(self.workspace), |
| 299 | skills_summary=skills_summary or "", |
| 300 | ) |
| 301 | |
| 302 | async def cancel_by_session(self, session_key: str) -> int: |
| 303 | """Cancel all subagents for the given session. Returns count cancelled.""" |
| 304 | tasks = [self._running_tasks[tid] for tid in self._session_tasks.get(session_key, []) |
| 305 | if tid in self._running_tasks and not self._running_tasks[tid].done()] |
| 306 | for t in tasks: |
| 307 | t.cancel() |
| 308 | if tasks: |
| 309 | await asyncio.gather(*tasks, return_exceptions=True) |
| 310 | return len(tasks) |
| 311 | |
| 312 | def get_running_count(self) -> int: |
| 313 | """Return the number of currently running subagents.""" |
| 314 | return len(self._running_tasks) |
| 315 | |
| 316 | def get_running_count_by_session(self, session_key: str) -> int: |
| 317 | """Return the number of currently running subagents for a session.""" |
| 318 | tids = self._session_tasks.get(session_key, set()) |
| 319 | return sum( |
| 320 | 1 for tid in tids |
| 321 | if tid in self._running_tasks and not self._running_tasks[tid].done() |
| 322 | ) |
| 323 |