| 1 | """Agent loop: the core processing engine.""" |
| 2 | |
| 3 | from __future__ import annotations |
| 4 | |
| 5 | import asyncio |
| 6 | import dataclasses |
| 7 | import json |
| 8 | import time |
| 9 | from contextlib import AsyncExitStack, nullcontext |
| 10 | from pathlib import Path |
| 11 | from typing import TYPE_CHECKING, Any, Awaitable, Callable |
| 12 | |
| 13 | from loguru import logger |
| 14 | |
| 15 | from nanobot.agent.autocompact import AutoCompact |
| 16 | from nanobot.agent.context import ContextBuilder |
| 17 | from nanobot.agent.hook import AgentHook, AgentHookContext, CompositeHook |
| 18 | from nanobot.agent.memory import Consolidator, Dream |
| 19 | from nanobot.agent.runner import _MAX_INJECTIONS_PER_TURN, AgentRunner, AgentRunSpec |
| 20 | from nanobot.agent.skills import BUILTIN_SKILLS_DIR |
| 21 | from nanobot.agent.subagent import SubagentManager |
| 22 | from nanobot.agent.tools.ask_user import AskUserTool |
| 23 | from nanobot.agent.tools.cron import CronTool |
| 24 | from nanobot.agent.tools.director import ( |
| 25 | DIRECTOR_CONTEXT_TOOL_NAMES, |
| 26 | DIRECTOR_MUTATING_TOOL_NAMES, |
| 27 | SHOT_COUNT_NEXT_STEP_HINT, |
| 28 | CreateShotPromptTool, |
| 29 | GenerateEchoShotTool, |
| 30 | GetGuidanceTool, |
| 31 | GetShotTool, |
| 32 | GetStoryTool, |
| 33 | GetWorkplaceStatusTool, |
| 34 | MergeShotTool, |
| 35 | ReviewShotTool, |
| 36 | SetDirectorGoalTool, |
| 37 | SetShotMemoryRecommendationsTool, |
| 38 | SetShotReferencesTool, |
| 39 | StartDirectorTool, |
| 40 | WriteStoryTool, |
| 41 | consume_shot_count_next_step_hint, |
| 42 | stepwise_shot_count_next_step_hint_eligible, |
| 43 | ) |
| 44 | from nanobot.agent.tools.filesystem import EditFileTool, ListDirTool, ReadFileTool, WriteFileTool |
| 45 | from nanobot.agent.tools.message import MessageTool |
| 46 | from nanobot.agent.tools.notebook import NotebookEditTool |
| 47 | from nanobot.agent.tools.registry import ToolRegistry |
| 48 | from nanobot.agent.tools.search import GlobTool, GrepTool |
| 49 | from nanobot.agent.tools.self import MyTool |
| 50 | from nanobot.agent.tools.shell import ExecTool |
| 51 | from nanobot.agent.tools.spawn import SpawnTool |
| 52 | from nanobot.agent.tools.web import WebFetchTool, WebSearchTool |
| 53 | from nanobot.bus.events import InboundMessage, OutboundMessage |
| 54 | from nanobot.bus.queue import MessageBus |
| 55 | from nanobot.channels.webui_sessions import WEBUI_CHANNEL, parse_webui_session_key |
| 56 | from nanobot.command import CommandContext, CommandRouter, register_builtin_commands |
| 57 | from nanobot.config.schema import AgentDefaults |
| 58 | from nanobot.providers.base import LLMProvider |
| 59 | from nanobot.session.auto_generate import get_auto_generate |
| 60 | from nanobot.session.manager import Session, SessionManager |
| 61 | from nanobot.utils.document import extract_documents |
| 62 | from nanobot.utils.helpers import image_placeholder_text |
| 63 | from nanobot.utils.helpers import truncate_text as truncate_text_fn |
| 64 | from nanobot.utils.progress_events import ( |
| 65 | build_tool_event_finish_payloads, |
| 66 | build_tool_event_start_payload, |
| 67 | invoke_on_progress, |
| 68 | on_progress_accepts_tool_events, |
| 69 | ) |
| 70 | from nanobot.utils.runtime import EMPTY_FINAL_RESPONSE_MESSAGE, is_empty_final_response_message |
| 71 | |
| 72 | if TYPE_CHECKING: |
| 73 | from nanobot.config.schema import ( |
| 74 | ChannelsConfig, |
| 75 | Config, |
| 76 | ExecToolConfig, |
| 77 | ToolsConfig, |
| 78 | WebToolsConfig, |
| 79 | ) |
| 80 | from nanobot.cron.service import CronService |
| 81 | |
| 82 | |
| 83 | UNIFIED_SESSION_KEY = "unified:default" |
| 84 | _HIDDEN_WORKPLACE_INJECTION_EVENTS = frozenset( |
| 85 | { |
| 86 | "workplace_shot_revision", |
| 87 | "workplace_story_edit", |
| 88 | "workplace_beats_edit", |
| 89 | "workplace_workflow_confirm_story", |
| 90 | "workplace_workflow_start_generation", |
| 91 | "workplace_workflow_start_merge", |
| 92 | "workplace_memory_recommendation", |
| 93 | } |
| 94 | ) |
| 95 | # These must join the active turn instead of waiting on the session lock as a |
| 96 | # competing task. Otherwise the current turn still emits chat filler |
| 97 | # ("已设定为 N 个镜头") while start_generation is blocked. |
| 98 | _SERIAL_WORKPLACE_INJECTION_EVENTS = frozenset( |
| 99 | { |
| 100 | "workplace_workflow_start_generation", |
| 101 | "workplace_workflow_start_merge", |
| 102 | } |
| 103 | ) |
| 104 | |
| 105 | |
| 106 | def _suppresses_frontend_reply( |
| 107 | msg: InboundMessage, |
| 108 | *, |
| 109 | channel: str, |
| 110 | stop_reason: str | None, |
| 111 | final_content: str | None, |
| 112 | tools_used: list[str] | None = None, |
| 113 | ) -> bool: |
| 114 | """Return True when the turn should not emit a user-visible chat reply.""" |
| 115 | if msg.metadata.get("silent"): |
| 116 | return True |
| 117 | if stop_reason != "empty_final_response": |
| 118 | return False |
| 119 | if not is_empty_final_response_message(final_content): |
| 120 | return False |
| 121 | injected = msg.metadata.get("injected_event") |
| 122 | if injected in _HIDDEN_WORKPLACE_INJECTION_EVENTS: |
| 123 | return True |
| 124 | if channel not in {WEBUI_CHANNEL, "websocket"}: |
| 125 | return False |
| 126 | return bool(tools_used) and any(name in DIRECTOR_MUTATING_TOOL_NAMES for name in tools_used) |
| 127 | |
| 128 | |
| 129 | def _replace_trailing_assistant_content( |
| 130 | messages: list[dict[str, Any]] | None, |
| 131 | content: str, |
| 132 | ) -> None: |
| 133 | if not messages: |
| 134 | return |
| 135 | for item in reversed(messages): |
| 136 | if isinstance(item, dict) and item.get("role") == "assistant": |
| 137 | item["content"] = content |
| 138 | return |
| 139 | |
| 140 | |
| 141 | class _LoopHook(AgentHook): |
| 142 | """Core hook for the main loop.""" |
| 143 | |
| 144 | def __init__( |
| 145 | self, |
| 146 | agent_loop: AgentLoop, |
| 147 | on_progress: Callable[..., Awaitable[None]] | None = None, |
| 148 | on_stream: Callable[[str], Awaitable[None]] | None = None, |
| 149 | on_stream_end: Callable[..., Awaitable[None]] | None = None, |
| 150 | *, |
| 151 | channel: str = "cli", |
| 152 | chat_id: str = "direct", |
| 153 | message_id: str | None = None, |
| 154 | session_key: str | None = None, |
| 155 | ) -> None: |
| 156 | super().__init__(reraise=True) |
| 157 | self._loop = agent_loop |
| 158 | self._on_progress = on_progress |
| 159 | self._on_stream = on_stream |
| 160 | self._on_stream_end = on_stream_end |
| 161 | self._channel = channel |
| 162 | self._chat_id = chat_id |
| 163 | self._message_id = message_id |
| 164 | self._session_key = session_key |
| 165 | self._stream_buf = "" |
| 166 | |
| 167 | def wants_streaming(self) -> bool: |
| 168 | return self._on_stream is not None |
| 169 | |
| 170 | async def on_stream(self, context: AgentHookContext, delta: str) -> None: |
| 171 | from nanobot.utils.helpers import strip_think |
| 172 | |
| 173 | ask_user_tool = self._loop.tools.get("ask_user") |
| 174 | if isinstance(ask_user_tool, AskUserTool) and ask_user_tool._sent_in_turn: |
| 175 | return |
| 176 | |
| 177 | prev_clean = strip_think(self._stream_buf) |
| 178 | self._stream_buf += delta |
| 179 | new_clean = strip_think(self._stream_buf) |
| 180 | incremental = new_clean[len(prev_clean) :] |
| 181 | if incremental and self._on_stream: |
| 182 | await self._on_stream(incremental) |
| 183 | |
| 184 | async def on_stream_end(self, context: AgentHookContext, *, resuming: bool) -> None: |
| 185 | if self._on_stream_end: |
| 186 | await self._on_stream_end(resuming=resuming) |
| 187 | self._stream_buf = "" |
| 188 | |
| 189 | async def before_iteration(self, context: AgentHookContext) -> None: |
| 190 | self._loop._current_iteration = context.iteration |
| 191 | |
| 192 | async def before_execute_tools(self, context: AgentHookContext) -> None: |
| 193 | if self._on_progress: |
| 194 | if not self._on_stream: |
| 195 | thought = self._loop._strip_think( |
| 196 | context.response.content if context.response else None |
| 197 | ) |
| 198 | if thought: |
| 199 | await self._on_progress(thought) |
| 200 | tool_hint = self._loop._strip_think(self._loop._tool_hint(context.tool_calls)) |
| 201 | tool_events = [build_tool_event_start_payload(tc) for tc in context.tool_calls] |
| 202 | await invoke_on_progress( |
| 203 | self._on_progress, |
| 204 | tool_hint, |
| 205 | tool_hint=True, |
| 206 | tool_events=tool_events, |
| 207 | ) |
| 208 | for tc in context.tool_calls: |
| 209 | args_str = json.dumps(tc.arguments, ensure_ascii=False) |
| 210 | logger.info("Tool call: {}({})", tc.name, args_str[:200]) |
| 211 | self._loop._set_tool_context( |
| 212 | self._channel, |
| 213 | self._chat_id, |
| 214 | self._message_id, |
| 215 | session_key=self._session_key, |
| 216 | ) |
| 217 | |
| 218 | async def after_iteration(self, context: AgentHookContext) -> None: |
| 219 | if ( |
| 220 | self._on_progress |
| 221 | and context.tool_calls |
| 222 | and context.tool_events |
| 223 | and on_progress_accepts_tool_events(self._on_progress) |
| 224 | ): |
| 225 | tool_events = build_tool_event_finish_payloads(context) |
| 226 | if tool_events: |
| 227 | await invoke_on_progress( |
| 228 | self._on_progress, |
| 229 | "", |
| 230 | tool_hint=False, |
| 231 | tool_events=tool_events, |
| 232 | ) |
| 233 | if ( |
| 234 | self._channel == "websocket" |
| 235 | and any( |
| 236 | event.get("status") == "ok" |
| 237 | and event.get("name") in DIRECTOR_MUTATING_TOOL_NAMES |
| 238 | for event in context.tool_events |
| 239 | ) |
| 240 | ): |
| 241 | await self._loop.bus.publish_outbound(OutboundMessage( |
| 242 | channel="websocket", |
| 243 | chat_id=self._chat_id, |
| 244 | content="", |
| 245 | metadata={"_workplace_event": "updated"}, |
| 246 | )) |
| 247 | u = context.usage or {} |
| 248 | logger.debug( |
| 249 | "LLM usage: prompt={} completion={} cached={}", |
| 250 | u.get("prompt_tokens", 0), |
| 251 | u.get("completion_tokens", 0), |
| 252 | u.get("cached_tokens", 0), |
| 253 | ) |
| 254 | |
| 255 | def finalize_content(self, context: AgentHookContext, content: str | None) -> str | None: |
| 256 | return self._loop._strip_think(content) |
| 257 | |
| 258 | |
| 259 | class AgentLoop: |
| 260 | """ |
| 261 | The agent loop is the core processing engine. |
| 262 | |
| 263 | It: |
| 264 | 1. Receives messages from the bus |
| 265 | 2. Builds context with history, memory, skills |
| 266 | 3. Calls the LLM |
| 267 | 4. Executes tool calls |
| 268 | 5. Sends responses back |
| 269 | """ |
| 270 | |
| 271 | _RUNTIME_CHECKPOINT_KEY = "runtime_checkpoint" |
| 272 | _PENDING_USER_TURN_KEY = "pending_user_turn" |
| 273 | |
| 274 | def __init__( |
| 275 | self, |
| 276 | bus: MessageBus, |
| 277 | provider: LLMProvider, |
| 278 | workspace: Path, |
| 279 | model: str | None = None, |
| 280 | max_iterations: int | None = None, |
| 281 | context_window_tokens: int | None = None, |
| 282 | context_block_limit: int | None = None, |
| 283 | max_tool_result_chars: int | None = None, |
| 284 | provider_retry_mode: str = "standard", |
| 285 | web_config: WebToolsConfig | None = None, |
| 286 | exec_config: ExecToolConfig | None = None, |
| 287 | cron_service: CronService | None = None, |
| 288 | restrict_to_workspace: bool = False, |
| 289 | session_manager: SessionManager | None = None, |
| 290 | mcp_servers: dict | None = None, |
| 291 | channels_config: ChannelsConfig | None = None, |
| 292 | timezone: str | None = None, |
| 293 | session_ttl_minutes: int = 0, |
| 294 | hooks: list[AgentHook] | None = None, |
| 295 | unified_session: bool = False, |
| 296 | disabled_skills: list[str] | None = None, |
| 297 | tools_config: ToolsConfig | None = None, |
| 298 | config: Config | None = None, |
| 299 | prompt_stacker_enabled: bool = False, |
| 300 | prompt_stacker_max_traces: int = 100, |
| 301 | ): |
| 302 | from nanobot.config.schema import ExecToolConfig, ToolsConfig, WebToolsConfig |
| 303 | |
| 304 | _tc = config.tools if config is not None else (tools_config or ToolsConfig()) |
| 305 | defaults = AgentDefaults() |
| 306 | self.bus = bus |
| 307 | self.channels_config = channels_config |
| 308 | self.provider = provider |
| 309 | self.workspace = workspace |
| 310 | self.model = model or provider.get_default_model() |
| 311 | self.max_iterations = ( |
| 312 | max_iterations if max_iterations is not None else defaults.max_tool_iterations |
| 313 | ) |
| 314 | self.context_window_tokens = ( |
| 315 | context_window_tokens |
| 316 | if context_window_tokens is not None |
| 317 | else defaults.context_window_tokens |
| 318 | ) |
| 319 | self.context_block_limit = context_block_limit |
| 320 | self.max_tool_result_chars = ( |
| 321 | max_tool_result_chars |
| 322 | if max_tool_result_chars is not None |
| 323 | else defaults.max_tool_result_chars |
| 324 | ) |
| 325 | self.provider_retry_mode = provider_retry_mode |
| 326 | self.web_config = web_config or WebToolsConfig() |
| 327 | self.exec_config = exec_config or ExecToolConfig() |
| 328 | self.tools_config = _tc |
| 329 | self.config = config |
| 330 | self.cron_service = cron_service |
| 331 | self.restrict_to_workspace = restrict_to_workspace |
| 332 | self._start_time = time.time() |
| 333 | self._last_usage: dict[str, int] = {} |
| 334 | self._extra_hooks: list[AgentHook] = hooks or [] |
| 335 | |
| 336 | self.context = ContextBuilder(workspace, timezone=timezone, disabled_skills=disabled_skills) |
| 337 | self.sessions = session_manager or SessionManager(workspace) |
| 338 | self.tools = ToolRegistry() |
| 339 | # Invalidate the tool-definition cache when the active PE set changes, so |
| 340 | # director tool descriptions/schemas hot-swap on the next agent turn. |
| 341 | from nanobot.prompts import PEManager |
| 342 | |
| 343 | PEManager.instance().on_change(lambda _name: self.tools.invalidate()) |
| 344 | self.runner = AgentRunner(provider) |
| 345 | |
| 346 | from nanobot.agent.event_stacker import EventStacker |
| 347 | from nanobot.agent.prompt_stacker import PromptStacker |
| 348 | PromptStacker.init(workspace, enabled=prompt_stacker_enabled, max_traces=prompt_stacker_max_traces) |
| 349 | EventStacker.init(workspace, enabled=prompt_stacker_enabled, max_traces=prompt_stacker_max_traces) |
| 350 | |
| 351 | self.subagents = SubagentManager( |
| 352 | provider=provider, |
| 353 | workspace=workspace, |
| 354 | bus=bus, |
| 355 | model=self.model, |
| 356 | web_config=self.web_config, |
| 357 | max_tool_result_chars=self.max_tool_result_chars, |
| 358 | exec_config=self.exec_config, |
| 359 | restrict_to_workspace=restrict_to_workspace, |
| 360 | disabled_skills=disabled_skills, |
| 361 | ) |
| 362 | self._unified_session = unified_session |
| 363 | self._running = False |
| 364 | self._mcp_servers = mcp_servers or {} |
| 365 | self._mcp_stacks: dict[str, AsyncExitStack] = {} |
| 366 | self._mcp_connected = False |
| 367 | self._mcp_connecting = False |
| 368 | self._active_tasks: dict[str, list[asyncio.Task]] = {} # session_key -> tasks |
| 369 | self._background_tasks: list[asyncio.Task] = [] |
| 370 | self._session_locks: dict[str, asyncio.Lock] = {} |
| 371 | # Per-session pending queues for mid-turn message injection. |
| 372 | # When a session has an active task, new messages for that session |
| 373 | # are routed here instead of creating a new task. |
| 374 | self._pending_queues: dict[str, asyncio.Queue] = {} |
| 375 | # A non-positive value disables the global request limit. |
| 376 | _max = ( |
| 377 | config.agents.defaults.max_concurrent_requests |
| 378 | if config is not None |
| 379 | else defaults.max_concurrent_requests |
| 380 | ) |
| 381 | self._concurrency_gate: asyncio.Semaphore | None = ( |
| 382 | asyncio.Semaphore(_max) if _max > 0 else None |
| 383 | ) |
| 384 | self.consolidator = Consolidator( |
| 385 | store=self.context.memory, |
| 386 | provider=provider, |
| 387 | model=self.model, |
| 388 | sessions=self.sessions, |
| 389 | context_window_tokens=self.context_window_tokens, |
| 390 | build_messages=self.context.build_messages, |
| 391 | get_tool_definitions=self.tools.get_definitions, |
| 392 | max_completion_tokens=provider.generation.max_tokens, |
| 393 | ) |
| 394 | self.auto_compact = AutoCompact( |
| 395 | sessions=self.sessions, |
| 396 | consolidator=self.consolidator, |
| 397 | session_ttl_minutes=session_ttl_minutes, |
| 398 | ) |
| 399 | self.dream = Dream( |
| 400 | store=self.context.memory, |
| 401 | provider=provider, |
| 402 | model=self.model, |
| 403 | ) |
| 404 | self._register_default_tools() |
| 405 | if _tc.my.enable: |
| 406 | self.tools.register(MyTool(loop=self, modify_allowed=_tc.my.allow_set)) |
| 407 | self._runtime_vars: dict[str, Any] = {} |
| 408 | self._current_iteration: int = 0 |
| 409 | self.commands = CommandRouter() |
| 410 | register_builtin_commands(self.commands) |
| 411 | |
| 412 | def _register_default_tools(self) -> None: |
| 413 | """Register the default set of tools.""" |
| 414 | allowed_dir = ( |
| 415 | self.workspace if (self.restrict_to_workspace or self.exec_config.sandbox) else None |
| 416 | ) |
| 417 | extra_read = [BUILTIN_SKILLS_DIR] if allowed_dir else None |
| 418 | self.tools.register( |
| 419 | ReadFileTool( |
| 420 | workspace=self.workspace, allowed_dir=allowed_dir, extra_allowed_dirs=extra_read |
| 421 | ) |
| 422 | ) |
| 423 | for cls in (WriteFileTool, EditFileTool, ListDirTool): |
| 424 | self.tools.register(cls(workspace=self.workspace, allowed_dir=allowed_dir)) |
| 425 | for cls in (GlobTool, GrepTool): |
| 426 | self.tools.register(cls(workspace=self.workspace, allowed_dir=allowed_dir)) |
| 427 | self.tools.register(NotebookEditTool(workspace=self.workspace, allowed_dir=allowed_dir)) |
| 428 | if self.exec_config.enable: |
| 429 | self.tools.register( |
| 430 | ExecTool( |
| 431 | working_dir=str(self.workspace), |
| 432 | timeout=self.exec_config.timeout, |
| 433 | restrict_to_workspace=self.restrict_to_workspace, |
| 434 | sandbox=self.exec_config.sandbox, |
| 435 | path_append=self.exec_config.path_append, |
| 436 | allowed_env_keys=self.exec_config.allowed_env_keys, |
| 437 | ) |
| 438 | ) |
| 439 | if self.web_config.enable: |
| 440 | self.tools.register( |
| 441 | WebSearchTool(config=self.web_config.search, proxy=self.web_config.proxy) |
| 442 | ) |
| 443 | self.tools.register(WebFetchTool(proxy=self.web_config.proxy)) |
| 444 | for cls in ( |
| 445 | StartDirectorTool, |
| 446 | SetDirectorGoalTool, |
| 447 | GetWorkplaceStatusTool, |
| 448 | GetStoryTool, |
| 449 | GetGuidanceTool, |
| 450 | WriteStoryTool, |
| 451 | GetShotTool, |
| 452 | CreateShotPromptTool, |
| 453 | ReviewShotTool, |
| 454 | SetShotMemoryRecommendationsTool, |
| 455 | SetShotReferencesTool, |
| 456 | GenerateEchoShotTool, |
| 457 | MergeShotTool, |
| 458 | ): |
| 459 | self.tools.register( |
| 460 | cls( |
| 461 | workspace=self.workspace, |
| 462 | tools_config=self.tools_config, |
| 463 | ) |
| 464 | ) |
| 465 | self.tools.register(MessageTool(send_callback=self.bus.publish_outbound)) |
| 466 | |
| 467 | async def _persist_ask_user( |
| 468 | session_key: str, |
| 469 | tool_call_id: str, |
| 470 | batch_id: str, |
| 471 | content: str, |
| 472 | questions: list, |
| 473 | channel: str, |
| 474 | chat_id: str, |
| 475 | ) -> None: |
| 476 | from datetime import datetime |
| 477 | |
| 478 | from nanobot.session.question_cards import ( |
| 479 | build_ask_user_session_messages, |
| 480 | session_has_tool_turn, |
| 481 | ) |
| 482 | |
| 483 | session = self.sessions.get_or_create(session_key) |
| 484 | if session_has_tool_turn(session.messages, tool_call_id): |
| 485 | return |
| 486 | for row in build_ask_user_session_messages( |
| 487 | tool_call_id=tool_call_id, |
| 488 | content=content, |
| 489 | questions=questions, |
| 490 | batch_id=batch_id, |
| 491 | channel=channel, |
| 492 | chat_id=chat_id, |
| 493 | ): |
| 494 | session.messages.append(row) |
| 495 | session.updated_at = datetime.now() |
| 496 | self.sessions.save(session) |
| 497 | |
| 498 | ask_user_tool = AskUserTool( |
| 499 | send_callback=self.bus.publish_outbound, |
| 500 | persist_callback=_persist_ask_user, |
| 501 | ) |
| 502 | self.tools.register(ask_user_tool) |
| 503 | self.tools.register(SpawnTool(manager=self.subagents)) |
| 504 | if self.cron_service: |
| 505 | self.tools.register( |
| 506 | CronTool(self.cron_service, default_timezone=self.context.timezone or "UTC") |
| 507 | ) |
| 508 | |
| 509 | async def _connect_mcp(self) -> None: |
| 510 | """Connect to configured MCP servers (one-time, lazy).""" |
| 511 | if self._mcp_connected or self._mcp_connecting or not self._mcp_servers: |
| 512 | return |
| 513 | self._mcp_connecting = True |
| 514 | from nanobot.agent.tools.mcp import connect_mcp_servers |
| 515 | |
| 516 | try: |
| 517 | self._mcp_stacks = await connect_mcp_servers(self._mcp_servers, self.tools) |
| 518 | if self._mcp_stacks: |
| 519 | self._mcp_connected = True |
| 520 | else: |
| 521 | logger.warning("No MCP servers connected successfully (will retry next message)") |
| 522 | except asyncio.CancelledError: |
| 523 | logger.warning("MCP connection cancelled (will retry next message)") |
| 524 | self._mcp_stacks.clear() |
| 525 | except BaseException as e: |
| 526 | logger.error("Failed to connect MCP servers (will retry next message): {}", e) |
| 527 | self._mcp_stacks.clear() |
| 528 | finally: |
| 529 | self._mcp_connecting = False |
| 530 | |
| 531 | def _set_tool_context( |
| 532 | self, |
| 533 | channel: str, |
| 534 | chat_id: str, |
| 535 | message_id: str | None = None, |
| 536 | *, |
| 537 | session_key: str | None = None, |
| 538 | injected_event: str | None | object = None, |
| 539 | update_workflow_injection: bool = False, |
| 540 | ) -> None: |
| 541 | """Update context for all tools that need routing info.""" |
| 542 | # Compute the effective session key (accounts for unified sessions). |
| 543 | # Use the inbound message's resolved session_key whenever available so |
| 544 | # tools (notably director + spawn) stay aligned with session_key_override |
| 545 | # used by the local browser channel (websocket:local:{chatId}). |
| 546 | if self._unified_session: |
| 547 | effective_key = UNIFIED_SESSION_KEY |
| 548 | else: |
| 549 | effective_key = session_key or f"{channel}:{chat_id}" |
| 550 | for name in ( |
| 551 | "message", "ask_user", "spawn", "cron", "my", "get_guidance", |
| 552 | *DIRECTOR_CONTEXT_TOOL_NAMES, |
| 553 | ): |
| 554 | if tool := self.tools.get(name): |
| 555 | if hasattr(tool, "set_context"): |
| 556 | if name == "spawn": |
| 557 | tool.set_context(channel, chat_id, effective_key=effective_key) |
| 558 | elif name == "get_guidance": |
| 559 | # Needs the effective session key so it resolves the |
| 560 | # per-session PE set; must NOT carry injected_event (that |
| 561 | # would clobber the global workflow-injection state). |
| 562 | tool.set_context(channel, chat_id, effective_key=effective_key) |
| 563 | elif name in DIRECTOR_CONTEXT_TOOL_NAMES: |
| 564 | if update_workflow_injection: |
| 565 | tool.set_context( |
| 566 | channel, |
| 567 | chat_id, |
| 568 | effective_key=effective_key, |
| 569 | injected_event=injected_event, |
| 570 | ) |
| 571 | else: |
| 572 | tool.set_context(channel, chat_id, effective_key=effective_key) |
| 573 | elif name == "ask_user": |
| 574 | tool.set_context( |
| 575 | channel, |
| 576 | chat_id, |
| 577 | session_key=effective_key, |
| 578 | ) |
| 579 | else: |
| 580 | tool.set_context(channel, chat_id, *([message_id] if name == "message" else [])) |
| 581 | @staticmethod |
| 582 | def _strip_think(text: str | None) -> str | None: |
| 583 | """Remove <think>…</think> blocks that some models embed in content.""" |
| 584 | if not text: |
| 585 | return None |
| 586 | from nanobot.utils.helpers import strip_think |
| 587 | |
| 588 | return strip_think(text) or None |
| 589 | |
| 590 | @staticmethod |
| 591 | def _tool_hint(tool_calls: list) -> str: |
| 592 | """Format tool calls as concise hints with smart abbreviation.""" |
| 593 | from nanobot.utils.tool_hints import format_tool_hints |
| 594 | |
| 595 | return format_tool_hints(tool_calls) |
| 596 | |
| 597 | async def _dispatch_command_inline( |
| 598 | self, |
| 599 | msg: InboundMessage, |
| 600 | key: str, |
| 601 | raw: str, |
| 602 | dispatch_fn: Callable[[CommandContext], Awaitable[OutboundMessage | None]], |
| 603 | ) -> None: |
| 604 | """Dispatch a command directly from the run() loop and publish the result.""" |
| 605 | ctx = CommandContext(msg=msg, session=None, key=key, raw=raw, loop=self) |
| 606 | result = await dispatch_fn(ctx) |
| 607 | if result: |
| 608 | await self.bus.publish_outbound(result) |
| 609 | else: |
| 610 | logger.warning("Command '{}' matched but dispatch returned None", raw) |
| 611 | |
| 612 | async def _cancel_active_tasks(self, key: str) -> int: |
| 613 | """Cancel and await all active tasks and subagents for *key*. |
| 614 | |
| 615 | Returns the total number of cancelled tasks + subagents. |
| 616 | """ |
| 617 | tasks = self._active_tasks.pop(key, []) |
| 618 | cancelled = sum(1 for t in tasks if not t.done() and t.cancel()) |
| 619 | for t in tasks: |
| 620 | try: |
| 621 | await t |
| 622 | except (asyncio.CancelledError, Exception): |
| 623 | pass |
| 624 | sub_cancelled = await self.subagents.cancel_by_session(key) |
| 625 | return cancelled + sub_cancelled |
| 626 | |
| 627 | def _effective_session_key(self, msg: InboundMessage) -> str: |
| 628 | """Return the session key used for task routing and mid-turn injections.""" |
| 629 | if self._unified_session and not msg.session_key_override: |
| 630 | return UNIFIED_SESSION_KEY |
| 631 | return msg.session_key |
| 632 | |
| 633 | async def _run_agent_loop( |
| 634 | self, |
| 635 | initial_messages: list[dict], |
| 636 | on_progress: Callable[..., Awaitable[None]] | None = None, |
| 637 | on_stream: Callable[[str], Awaitable[None]] | None = None, |
| 638 | on_stream_end: Callable[..., Awaitable[None]] | None = None, |
| 639 | on_retry_wait: Callable[[str], Awaitable[None]] | None = None, |
| 640 | *, |
| 641 | session: Session | None = None, |
| 642 | channel: str = "cli", |
| 643 | chat_id: str = "direct", |
| 644 | message_id: str | None = None, |
| 645 | pending_queue: asyncio.Queue | None = None, |
| 646 | ) -> tuple[str | None, list[str], list[dict], str, bool]: |
| 647 | """Run the agent iteration loop. |
| 648 | |
| 649 | *on_stream*: called with each content delta during streaming. |
| 650 | *on_stream_end(resuming)*: called when a streaming session finishes. |
| 651 | ``resuming=True`` means tool calls follow (spinner should restart); |
| 652 | ``resuming=False`` means this is the final response. |
| 653 | |
| 654 | Returns (final_content, tools_used, messages, stop_reason, had_injections). |
| 655 | """ |
| 656 | loop_hook = _LoopHook( |
| 657 | self, |
| 658 | on_progress=on_progress, |
| 659 | on_stream=on_stream, |
| 660 | on_stream_end=on_stream_end, |
| 661 | channel=channel, |
| 662 | chat_id=chat_id, |
| 663 | message_id=message_id, |
| 664 | session_key=session.key if session else None, |
| 665 | ) |
| 666 | hook: AgentHook = ( |
| 667 | CompositeHook([loop_hook] + self._extra_hooks) if self._extra_hooks else loop_hook |
| 668 | ) |
| 669 | |
| 670 | async def _checkpoint(payload: dict[str, Any]) -> None: |
| 671 | if session is None: |
| 672 | return |
| 673 | self._set_runtime_checkpoint(session, payload) |
| 674 | |
| 675 | async def _drain_pending(*, limit: int = _MAX_INJECTIONS_PER_TURN) -> list[dict[str, Any]]: |
| 676 | """Drain follow-up messages from the pending queue. |
| 677 | |
| 678 | When no messages are immediately available but sub-agents |
| 679 | spawned in this dispatch are still running, blocks until at |
| 680 | least one result arrives (or timeout). This keeps the runner |
| 681 | loop alive so subsequent sub-agent completions are consumed |
| 682 | in-order rather than dispatched separately. |
| 683 | """ |
| 684 | if pending_queue is None: |
| 685 | return [] |
| 686 | |
| 687 | def _to_injected_message(pending_msg: InboundMessage) -> dict[str, Any]: |
| 688 | injected_event = pending_msg.metadata.get("injected_event") |
| 689 | if injected_event in _HIDDEN_WORKPLACE_INJECTION_EVENTS: |
| 690 | self._set_tool_context( |
| 691 | channel, |
| 692 | chat_id, |
| 693 | message_id, |
| 694 | session_key=session.key if session else None, |
| 695 | injected_event=injected_event, |
| 696 | update_workflow_injection=True, |
| 697 | ) |
| 698 | if pending_msg.metadata.get("injected_role") == "system": |
| 699 | return {"role": "system", "content": pending_msg.content.strip()} |
| 700 | if pending_msg.metadata.get("injected_event") in _HIDDEN_WORKPLACE_INJECTION_EVENTS: |
| 701 | return { |
| 702 | "role": "user", |
| 703 | "content": pending_msg.content.strip(), |
| 704 | "_hidden": True, |
| 705 | } |
| 706 | content = pending_msg.content |
| 707 | media = pending_msg.media if pending_msg.media else None |
| 708 | if media: |
| 709 | content, media = extract_documents(content, media) |
| 710 | media = media or None |
| 711 | user_content = self.context._build_user_content(content, media) |
| 712 | runtime_ctx = self.context._build_runtime_context( |
| 713 | pending_msg.channel, |
| 714 | pending_msg.chat_id, |
| 715 | self.context.timezone, |
| 716 | ) |
| 717 | if isinstance(user_content, str): |
| 718 | merged: str | list[dict[str, Any]] = f"{runtime_ctx}\n\n{user_content}" |
| 719 | else: |
| 720 | merged = [{"type": "text", "text": runtime_ctx}] + user_content |
| 721 | return {"role": "user", "content": merged} |
| 722 | |
| 723 | items: list[dict[str, Any]] = [] |
| 724 | while len(items) < limit: |
| 725 | try: |
| 726 | items.append(_to_injected_message(pending_queue.get_nowait())) |
| 727 | except asyncio.QueueEmpty: |
| 728 | break |
| 729 | |
| 730 | # Block if nothing drained but sub-agents spawned in this dispatch |
| 731 | # are still running. Keeps the runner loop alive so subsequent |
| 732 | # completions are injected in-order rather than dispatched separately. |
| 733 | if (not items |
| 734 | and session is not None |
| 735 | and self.subagents.get_running_count_by_session(session.key) > 0): |
| 736 | try: |
| 737 | msg = await asyncio.wait_for(pending_queue.get(), timeout=300) |
| 738 | except asyncio.TimeoutError: |
| 739 | logger.warning( |
| 740 | "Timeout waiting for sub-agent completion in session {}", |
| 741 | session.key, |
| 742 | ) |
| 743 | return items |
| 744 | items.append(_to_injected_message(msg)) |
| 745 | while len(items) < limit: |
| 746 | try: |
| 747 | items.append(_to_injected_message(pending_queue.get_nowait())) |
| 748 | except asyncio.QueueEmpty: |
| 749 | break |
| 750 | |
| 751 | if items: |
| 752 | if aut := self.tools.get("ask_user"): |
| 753 | if isinstance(aut, AskUserTool): |
| 754 | aut.start_turn() |
| 755 | |
| 756 | return items |
| 757 | |
| 758 | from nanobot.agent.tools.registry import set_tool_session_metadata |
| 759 | set_tool_session_metadata(session.metadata if session else None) |
| 760 | try: |
| 761 | result = await self.runner.run(AgentRunSpec( |
| 762 | initial_messages=initial_messages, |
| 763 | tools=self.tools, |
| 764 | model=self.model, |
| 765 | max_iterations=self.max_iterations, |
| 766 | max_tool_result_chars=self.max_tool_result_chars, |
| 767 | hook=hook, |
| 768 | error_message="Sorry, I encountered an error calling the AI model.", |
| 769 | concurrent_tools=True, |
| 770 | workspace=self.workspace, |
| 771 | session_key=session.key if session else None, |
| 772 | context_window_tokens=self.context_window_tokens, |
| 773 | context_block_limit=self.context_block_limit, |
| 774 | provider_retry_mode=self.provider_retry_mode, |
| 775 | progress_callback=on_progress, |
| 776 | retry_wait_callback=on_retry_wait, |
| 777 | checkpoint_callback=_checkpoint, |
| 778 | injection_callback=_drain_pending, |
| 779 | session_metadata=session.metadata if session else None, |
| 780 | )) |
| 781 | finally: |
| 782 | set_tool_session_metadata(None) |
| 783 | self._last_usage = result.usage |
| 784 | if result.stop_reason == "max_iterations": |
| 785 | logger.warning("Max iterations ({}) reached", self.max_iterations) |
| 786 | # Push final content through stream so streaming channels (e.g. Feishu) |
| 787 | # update the card instead of leaving it empty. |
| 788 | if on_stream and on_stream_end: |
| 789 | await on_stream(result.final_content or "") |
| 790 | await on_stream_end(resuming=False) |
| 791 | elif result.stop_reason == "error": |
| 792 | logger.error("LLM returned error: {}", (result.final_content or "")[:200]) |
| 793 | return result.final_content, result.tools_used, result.messages, result.stop_reason, result.had_injections |
| 794 | |
| 795 | async def run(self) -> None: |
| 796 | """Run the agent loop, dispatching messages as tasks to stay responsive to /stop.""" |
| 797 | self._running = True |
| 798 | await self._connect_mcp() |
| 799 | logger.info("Agent loop started") |
| 800 | |
| 801 | while self._running: |
| 802 | try: |
| 803 | msg = await asyncio.wait_for(self.bus.consume_inbound(), timeout=1.0) |
| 804 | except asyncio.TimeoutError: |
| 805 | self.auto_compact.check_expired( |
| 806 | self._schedule_background, |
| 807 | active_session_keys=self._pending_queues.keys(), |
| 808 | ) |
| 809 | continue |
| 810 | except asyncio.CancelledError: |
| 811 | # Preserve real task cancellation so shutdown can complete cleanly. |
| 812 | # Only ignore non-task CancelledError signals that may leak from integrations. |
| 813 | if not self._running or asyncio.current_task().cancelling(): |
| 814 | raise |
| 815 | continue |
| 816 | except Exception as e: |
| 817 | logger.warning("Error consuming inbound message: {}, continuing...", e) |
| 818 | continue |
| 819 | |
| 820 | raw = msg.content.strip() |
| 821 | if self.commands.is_priority(raw): |
| 822 | await self._dispatch_command_inline( |
| 823 | msg, msg.session_key, raw, |
| 824 | self.commands.dispatch_priority, |
| 825 | ) |
| 826 | continue |
| 827 | effective_key = self._effective_session_key(msg) |
| 828 | injected_event = msg.metadata.get("injected_event") |
| 829 | is_workplace_workflow = ( |
| 830 | msg.channel == "system" |
| 831 | and injected_event in _HIDDEN_WORKPLACE_INJECTION_EVENTS |
| 832 | ) |
| 833 | should_queue_workplace = ( |
| 834 | is_workplace_workflow |
| 835 | and injected_event in _SERIAL_WORKPLACE_INJECTION_EVENTS |
| 836 | ) |
| 837 | # If this session already has an active pending queue (i.e. a task |
| 838 | # is processing this session), route the message there for mid-turn |
| 839 | # injection instead of creating a competing task. |
| 840 | if effective_key in self._pending_queues and ( |
| 841 | not is_workplace_workflow or should_queue_workplace |
| 842 | ): |
| 843 | # Non-priority commands must not be queued for injection; |
| 844 | # dispatch them directly (same pattern as priority commands). |
| 845 | if self.commands.is_dispatchable_command(raw): |
| 846 | await self._dispatch_command_inline( |
| 847 | msg, effective_key, raw, |
| 848 | self.commands.dispatch, |
| 849 | ) |
| 850 | continue |
| 851 | pending_msg = msg |
| 852 | if effective_key != msg.session_key: |
| 853 | pending_msg = dataclasses.replace( |
| 854 | msg, |
| 855 | session_key_override=effective_key, |
| 856 | ) |
| 857 | try: |
| 858 | self._pending_queues[effective_key].put_nowait(pending_msg) |
| 859 | except asyncio.QueueFull: |
| 860 | logger.warning( |
| 861 | "Pending queue full for session {}, falling back to queued task", |
| 862 | effective_key, |
| 863 | ) |
| 864 | else: |
| 865 | logger.info( |
| 866 | "Routed follow-up message to pending queue for session {}", |
| 867 | effective_key, |
| 868 | ) |
| 869 | continue |
| 870 | # Compute the effective session key before dispatching |
| 871 | # This ensures /stop command can find tasks correctly when unified session is enabled |
| 872 | task = asyncio.create_task(self._dispatch(msg)) |
| 873 | self._active_tasks.setdefault(effective_key, []).append(task) |
| 874 | task.add_done_callback( |
| 875 | lambda t, k=effective_key: self._active_tasks.get(k, []) |
| 876 | and self._active_tasks[k].remove(t) |
| 877 | if t in self._active_tasks.get(k, []) |
| 878 | else None |
| 879 | ) |
| 880 | |
| 881 | async def _dispatch(self, msg: InboundMessage) -> None: |
| 882 | """Process a message: per-session serial, cross-session concurrent.""" |
| 883 | session_key = self._effective_session_key(msg) |
| 884 | if session_key != msg.session_key: |
| 885 | msg = dataclasses.replace(msg, session_key_override=session_key) |
| 886 | lock = self._session_locks.setdefault(session_key, asyncio.Lock()) |
| 887 | gate = self._concurrency_gate or nullcontext() |
| 888 | |
| 889 | # Register a pending queue so follow-up messages for this session are |
| 890 | # routed here (mid-turn injection) instead of spawning a new task. |
| 891 | pending = asyncio.Queue(maxsize=20) |
| 892 | self._pending_queues[session_key] = pending |
| 893 | |
| 894 | try: |
| 895 | async with lock, gate: |
| 896 | try: |
| 897 | on_stream = on_stream_end = None |
| 898 | if msg.metadata.get("_wants_stream"): |
| 899 | # Split one answer into distinct stream segments. |
| 900 | stream_base_id = f"{msg.session_key}:{time.time_ns()}" |
| 901 | stream_segment = 0 |
| 902 | |
| 903 | def _current_stream_id() -> str: |
| 904 | return f"{stream_base_id}:{stream_segment}" |
| 905 | |
| 906 | async def on_stream(delta: str) -> None: |
| 907 | meta = dict(msg.metadata or {}) |
| 908 | meta["_stream_delta"] = True |
| 909 | meta["_stream_id"] = _current_stream_id() |
| 910 | await self.bus.publish_outbound(OutboundMessage( |
| 911 | channel=msg.channel, chat_id=msg.chat_id, |
| 912 | content=delta, |
| 913 | metadata=meta, |
| 914 | )) |
| 915 | |
| 916 | async def on_stream_end(*, resuming: bool = False) -> None: |
| 917 | nonlocal stream_segment |
| 918 | meta = dict(msg.metadata or {}) |
| 919 | meta["_stream_end"] = True |
| 920 | meta["_resuming"] = resuming |
| 921 | meta["_stream_id"] = _current_stream_id() |
| 922 | await self.bus.publish_outbound(OutboundMessage( |
| 923 | channel=msg.channel, chat_id=msg.chat_id, |
| 924 | content="", |
| 925 | metadata=meta, |
| 926 | )) |
| 927 | stream_segment += 1 |
| 928 | |
| 929 | response = await self._process_message( |
| 930 | msg, on_stream=on_stream, on_stream_end=on_stream_end, |
| 931 | pending_queue=pending, |
| 932 | ) |
| 933 | if response is not None: |
| 934 | await self.bus.publish_outbound(response) |
| 935 | elif msg.channel == "cli": |
| 936 | await self.bus.publish_outbound(OutboundMessage( |
| 937 | channel=msg.channel, chat_id=msg.chat_id, |
| 938 | content="", metadata=msg.metadata or {}, |
| 939 | )) |
| 940 | except asyncio.CancelledError: |
| 941 | logger.info("Task cancelled for session {}", session_key) |
| 942 | # Preserve partial context from the interrupted turn so |
| 943 | # the user does not lose tool results and assistant |
| 944 | # messages accumulated before /stop. The checkpoint was |
| 945 | # already persisted to session metadata by |
| 946 | # _emit_checkpoint during tool execution; materializing |
| 947 | # it into session history now makes it visible in the |
| 948 | # next conversation turn. |
| 949 | try: |
| 950 | key = self._effective_session_key(msg) |
| 951 | session = self.sessions.get_or_create(key) |
| 952 | if self._restore_runtime_checkpoint(session): |
| 953 | self._clear_pending_user_turn(session) |
| 954 | self.sessions.save(session) |
| 955 | logger.info( |
| 956 | "Restored partial context for cancelled session {}", |
| 957 | key, |
| 958 | ) |
| 959 | except Exception: |
| 960 | logger.debug( |
| 961 | "Could not restore checkpoint for cancelled session {}", |
| 962 | session_key, |
| 963 | exc_info=True, |
| 964 | ) |
| 965 | raise |
| 966 | except Exception: |
| 967 | logger.exception("Error processing message for session {}", session_key) |
| 968 | await self.bus.publish_outbound(OutboundMessage( |
| 969 | channel=msg.channel, chat_id=msg.chat_id, |
| 970 | content="Sorry, I encountered an error.", |
| 971 | )) |
| 972 | finally: |
| 973 | # Drain any messages still in the pending queue and re-publish |
| 974 | # them to the bus so they are processed as fresh inbound messages |
| 975 | # rather than silently lost. |
| 976 | queue = self._pending_queues.pop(session_key, None) |
| 977 | if queue is not None: |
| 978 | leftover = 0 |
| 979 | while True: |
| 980 | try: |
| 981 | item = queue.get_nowait() |
| 982 | except asyncio.QueueEmpty: |
| 983 | break |
| 984 | await self.bus.publish_inbound(item) |
| 985 | leftover += 1 |
| 986 | if leftover: |
| 987 | logger.info( |
| 988 | "Re-published {} leftover message(s) to bus for session {}", |
| 989 | leftover, session_key, |
| 990 | ) |
| 991 | |
| 992 | async def close_mcp(self) -> None: |
| 993 | """Drain pending background archives, then close MCP connections.""" |
| 994 | if self._background_tasks: |
| 995 | await asyncio.gather(*self._background_tasks, return_exceptions=True) |
| 996 | self._background_tasks.clear() |
| 997 | for name, stack in self._mcp_stacks.items(): |
| 998 | try: |
| 999 | await stack.aclose() |
| 1000 | except (RuntimeError, BaseExceptionGroup): |
| 1001 | logger.debug("MCP server '{}' cleanup error (can be ignored)", name) |
| 1002 | self._mcp_stacks.clear() |
| 1003 | |
| 1004 | def _schedule_background(self, coro) -> None: |
| 1005 | """Schedule a coroutine as a tracked background task (drained on shutdown).""" |
| 1006 | task = asyncio.create_task(coro) |
| 1007 | self._background_tasks.append(task) |
| 1008 | task.add_done_callback(self._background_tasks.remove) |
| 1009 | |
| 1010 | def _schedule_session_preview(self, session: Session) -> None: |
| 1011 | if self.sessions.preview_generator is None: |
| 1012 | return |
| 1013 | self._schedule_background(self.sessions.ensure_preview(session)) |
| 1014 | |
| 1015 | def stop(self) -> None: |
| 1016 | """Stop the agent loop.""" |
| 1017 | self._running = False |
| 1018 | logger.info("Agent loop stopping") |
| 1019 | |
| 1020 | def _attach_shot_count_next_step_hint( |
| 1021 | self, |
| 1022 | *, |
| 1023 | msg: InboundMessage, |
| 1024 | session: Session, |
| 1025 | final_content: str | None, |
| 1026 | stop_reason: str | None, |
| 1027 | all_msgs: list[dict[str, Any]] | None, |
| 1028 | tools_used: list[str] | None = None, |
| 1029 | ) -> tuple[str | None, str | None, bool]: |
| 1030 | """Fill the stepwise 「下一步」 hint after shot_count is first locked. |
| 1031 | |
| 1032 | The third value is True when this turn should deliver the hint as a |
| 1033 | full chat message (skip ``_streamed`` so ChannelManager actually sends it). |
| 1034 | """ |
| 1035 | metadata = msg.metadata if isinstance(msg.metadata, dict) else {} |
| 1036 | skip_visible = bool(metadata.get("silent")) or ( |
| 1037 | metadata.get("injected_event") in _HIDDEN_WORKPLACE_INJECTION_EVENTS |
| 1038 | ) |
| 1039 | auto = get_auto_generate(session.metadata) |
| 1040 | hint = consume_shot_count_next_step_hint( |
| 1041 | self.workspace, |
| 1042 | session.key, |
| 1043 | auto_generate=auto, |
| 1044 | emit=not skip_visible, |
| 1045 | ) |
| 1046 | # Turn 2 calls set_director_goal + write_story together. concurrent_tools |
| 1047 | # can let write_story overwrite state.json and drop the pending flag. |
| 1048 | if ( |
| 1049 | not hint |
| 1050 | and not skip_visible |
| 1051 | and not auto |
| 1052 | and tools_used |
| 1053 | and "set_director_goal" in tools_used |
| 1054 | and stepwise_shot_count_next_step_hint_eligible( |
| 1055 | self.workspace, |
| 1056 | session.key, |
| 1057 | auto_generate=auto, |
| 1058 | ) |
| 1059 | ): |
| 1060 | hint = SHOT_COUNT_NEXT_STEP_HINT |
| 1061 | if not hint: |
| 1062 | return final_content, stop_reason, False |
| 1063 | _replace_trailing_assistant_content(all_msgs, hint) |
| 1064 | next_reason = ( |
| 1065 | "completed" if stop_reason in {None, "empty_final_response"} else stop_reason |
| 1066 | ) |
| 1067 | # Always True: even if the model streamed other text, ChannelManager |
| 1068 | # must still send this system sentence as a full chat message. |
| 1069 | return hint, next_reason, True |
| 1070 | |
| 1071 | async def _process_message( |
| 1072 | self, |
| 1073 | msg: InboundMessage, |
| 1074 | session_key: str | None = None, |
| 1075 | on_progress: Callable[..., Awaitable[None]] | None = None, |
| 1076 | on_stream: Callable[[str], Awaitable[None]] | None = None, |
| 1077 | on_stream_end: Callable[..., Awaitable[None]] | None = None, |
| 1078 | pending_queue: asyncio.Queue | None = None, |
| 1079 | ) -> OutboundMessage | None: |
| 1080 | """Process a single inbound message and return the response.""" |
| 1081 | # System messages honor local WebUI keys so outbound routing uses the |
| 1082 | # wire chat_id rather than the persisted namespace. |
| 1083 | if msg.channel == "system": |
| 1084 | key = msg.session_key_override or msg.chat_id |
| 1085 | parsed = parse_webui_session_key(key) |
| 1086 | if parsed: |
| 1087 | channel = WEBUI_CHANNEL |
| 1088 | chat_id = parsed[1] |
| 1089 | elif msg.session_key_override: |
| 1090 | channel, chat_id = key.split(":", 1) if ":" in key else ("cli", key) |
| 1091 | else: |
| 1092 | channel, chat_id = ( |
| 1093 | msg.chat_id.split(":", 1) if ":" in msg.chat_id else ("cli", msg.chat_id) |
| 1094 | ) |
| 1095 | key = f"{channel}:{chat_id}" |
| 1096 | logger.info("Processing system message from {}", msg.sender_id) |
| 1097 | session = self.sessions.get_or_create(key) |
| 1098 | if self._restore_runtime_checkpoint(session): |
| 1099 | self.sessions.save(session) |
| 1100 | if self._restore_pending_user_turn(session): |
| 1101 | self.sessions.save(session) |
| 1102 | |
| 1103 | session, pending = self.auto_compact.prepare_session(session, key) |
| 1104 | |
| 1105 | await self.consolidator.maybe_consolidate_by_tokens( |
| 1106 | session, |
| 1107 | session_summary=pending, |
| 1108 | ) |
| 1109 | # Persist subagent follow-ups into durable history BEFORE prompt |
| 1110 | # assembly. ContextBuilder merges adjacent same-role messages for |
| 1111 | # provider compatibility, which previously caused the follow-up to |
| 1112 | # disappear from session.messages while still being visible to the |
| 1113 | # LLM via the merged prompt. See _persist_subagent_followup. |
| 1114 | is_subagent = msg.sender_id == "subagent" |
| 1115 | if is_subagent and self._persist_subagent_followup(session, msg): |
| 1116 | self.sessions.save(session) |
| 1117 | self._set_tool_context( |
| 1118 | channel, |
| 1119 | chat_id, |
| 1120 | msg.metadata.get("message_id"), |
| 1121 | session_key=key, |
| 1122 | injected_event=msg.metadata.get("injected_event"), |
| 1123 | update_workflow_injection=True, |
| 1124 | ) |
| 1125 | history = session.get_history(max_messages=0) |
| 1126 | injected_role = msg.metadata.get("injected_role") |
| 1127 | current_role = "assistant" if is_subagent else "user" |
| 1128 | if injected_role == "system": |
| 1129 | current_role = "system" |
| 1130 | |
| 1131 | from nanobot.agent.event_stacker import EventStacker |
| 1132 | from nanobot.agent.prompt_stacker import PromptStacker |
| 1133 | PromptStacker.begin_turn(session.key if session else None, self.model) |
| 1134 | EventStacker.begin_turn(session.key if session else None, self.model) |
| 1135 | |
| 1136 | # Subagent content is already in `history` above; passing it again |
| 1137 | # as current_message would double-project it into the prompt. |
| 1138 | messages = self.context.build_messages( |
| 1139 | history=history, |
| 1140 | current_message="" if is_subagent else msg.content, |
| 1141 | channel=channel, |
| 1142 | chat_id=chat_id, |
| 1143 | session_summary=pending, |
| 1144 | current_role=current_role, |
| 1145 | session_metadata=session.metadata, |
| 1146 | session_key=session.key if session else None, |
| 1147 | ) |
| 1148 | if ( |
| 1149 | msg.metadata.get("injected_event") in _HIDDEN_WORKPLACE_INJECTION_EVENTS |
| 1150 | and current_role == "user" |
| 1151 | and messages |
| 1152 | and messages[-1].get("role") == "user" |
| 1153 | ): |
| 1154 | messages[-1]["_hidden"] = True |
| 1155 | final_content, tools_used, all_msgs, stop_reason, _ = await self._run_agent_loop( |
| 1156 | messages, session=session, channel=channel, chat_id=chat_id, |
| 1157 | message_id=msg.metadata.get("message_id"), |
| 1158 | pending_queue=pending_queue, |
| 1159 | ) |
| 1160 | final_content, stop_reason, _ = self._attach_shot_count_next_step_hint( |
| 1161 | msg=msg, |
| 1162 | session=session, |
| 1163 | final_content=final_content, |
| 1164 | stop_reason=stop_reason, |
| 1165 | all_msgs=all_msgs, |
| 1166 | tools_used=tools_used, |
| 1167 | ) |
| 1168 | EventStacker.end_turn(stop_reason=stop_reason) |
| 1169 | self._save_turn(session, all_msgs, 1 + len(history)) |
| 1170 | self._clear_runtime_checkpoint(session) |
| 1171 | self.sessions.save(session) |
| 1172 | self._schedule_background(self.consolidator.maybe_consolidate_by_tokens(session)) |
| 1173 | if _suppresses_frontend_reply( |
| 1174 | msg, |
| 1175 | channel=channel, |
| 1176 | stop_reason=stop_reason, |
| 1177 | final_content=final_content, |
| 1178 | tools_used=tools_used, |
| 1179 | ): |
| 1180 | injected = msg.metadata.get("injected_event") |
| 1181 | if ( |
| 1182 | injected in _HIDDEN_WORKPLACE_INJECTION_EVENTS |
| 1183 | and channel == WEBUI_CHANNEL |
| 1184 | ): |
| 1185 | await self.bus.publish_outbound( |
| 1186 | OutboundMessage( |
| 1187 | channel=channel, |
| 1188 | chat_id=chat_id, |
| 1189 | content="", |
| 1190 | metadata={ |
| 1191 | "_workplace_event": "updated", |
| 1192 | "work_id": msg.metadata.get("work_id"), |
| 1193 | }, |
| 1194 | ) |
| 1195 | ) |
| 1196 | return None |
| 1197 | return OutboundMessage( |
| 1198 | channel=channel, |
| 1199 | chat_id=chat_id, |
| 1200 | content=final_content or "Background task completed.", |
| 1201 | ) |
| 1202 | |
| 1203 | # Extract document text from media at the processing boundary so all |
| 1204 | # channels benefit without format-specific logic in ContextBuilder. |
| 1205 | if msg.media: |
| 1206 | new_content, image_only = extract_documents(msg.content, msg.media) |
| 1207 | msg = dataclasses.replace(msg, content=new_content, media=image_only) |
| 1208 | |
| 1209 | preview = msg.content[:80] + "..." if len(msg.content) > 80 else msg.content |
| 1210 | logger.info("Processing message from {}:{}: {}", msg.channel, msg.sender_id, preview) |
| 1211 | |
| 1212 | key = session_key or msg.session_key |
| 1213 | session = self.sessions.get_or_create(key) |
| 1214 | if isinstance(msg.metadata, dict): |
| 1215 | from nanobot.session.source import apply_source, resolve_source_from_wire |
| 1216 | inbound = resolve_source_from_wire(msg.metadata) |
| 1217 | if inbound and apply_source(session.metadata, inbound): |
| 1218 | self.sessions.save(session) |
| 1219 | logger.info("Persisted session source={} session_key={}", inbound, key) |
| 1220 | if self._restore_runtime_checkpoint(session): |
| 1221 | self.sessions.save(session) |
| 1222 | if self._restore_pending_user_turn(session): |
| 1223 | self.sessions.save(session) |
| 1224 | |
| 1225 | session, pending = self.auto_compact.prepare_session(session, key) |
| 1226 | |
| 1227 | # Slash commands |
| 1228 | raw = msg.content.strip() |
| 1229 | ctx = CommandContext(msg=msg, session=session, key=key, raw=raw, loop=self) |
| 1230 | if result := await self.commands.dispatch(ctx): |
| 1231 | return result |
| 1232 | |
| 1233 | await self.consolidator.maybe_consolidate_by_tokens( |
| 1234 | session, |
| 1235 | session_summary=pending, |
| 1236 | ) |
| 1237 | |
| 1238 | self._set_tool_context( |
| 1239 | msg.channel, |
| 1240 | msg.chat_id, |
| 1241 | msg.metadata.get("message_id"), |
| 1242 | session_key=key, |
| 1243 | injected_event=None, |
| 1244 | update_workflow_injection=True, |
| 1245 | ) |
| 1246 | if message_tool := self.tools.get("message"): |
| 1247 | if isinstance(message_tool, MessageTool): |
| 1248 | message_tool.start_turn() |
| 1249 | if ask_user_tool := self.tools.get("ask_user"): |
| 1250 | if isinstance(ask_user_tool, AskUserTool): |
| 1251 | ask_user_tool.start_turn() |
| 1252 | |
| 1253 | history = session.get_history(max_messages=0) |
| 1254 | |
| 1255 | from nanobot.agent.event_stacker import EventStacker |
| 1256 | from nanobot.agent.prompt_stacker import PromptStacker |
| 1257 | PromptStacker.begin_turn(session.key if session else None, self.model) |
| 1258 | EventStacker.begin_turn(session.key if session else None, self.model) |
| 1259 | |
| 1260 | initial_messages = self.context.build_messages( |
| 1261 | history=history, |
| 1262 | current_message=msg.content, |
| 1263 | session_summary=pending, |
| 1264 | media=msg.media if msg.media else None, |
| 1265 | channel=msg.channel, |
| 1266 | chat_id=msg.chat_id, |
| 1267 | session_metadata=session.metadata, |
| 1268 | session_key=key, |
| 1269 | ) |
| 1270 | |
| 1271 | async def _bus_progress( |
| 1272 | content: str, |
| 1273 | *, |
| 1274 | tool_hint: bool = False, |
| 1275 | tool_events: list[dict[str, Any]] | None = None, |
| 1276 | ) -> None: |
| 1277 | meta = dict(msg.metadata or {}) |
| 1278 | meta["_progress"] = True |
| 1279 | meta["_tool_hint"] = tool_hint |
| 1280 | if tool_events: |
| 1281 | meta["_tool_events"] = tool_events |
| 1282 | await self.bus.publish_outbound( |
| 1283 | OutboundMessage( |
| 1284 | channel=msg.channel, |
| 1285 | chat_id=msg.chat_id, |
| 1286 | content=content, |
| 1287 | metadata=meta, |
| 1288 | ) |
| 1289 | ) |
| 1290 | |
| 1291 | async def _on_retry_wait(content: str) -> None: |
| 1292 | meta = dict(msg.metadata or {}) |
| 1293 | meta["_retry_wait"] = True |
| 1294 | await self.bus.publish_outbound( |
| 1295 | OutboundMessage( |
| 1296 | channel=msg.channel, |
| 1297 | chat_id=msg.chat_id, |
| 1298 | content=content, |
| 1299 | metadata=meta, |
| 1300 | ) |
| 1301 | ) |
| 1302 | |
| 1303 | # Persist the triggering user message up front so a mid-turn crash |
| 1304 | # doesn't silently lose the prompt on recovery. ``media`` rides along |
| 1305 | # as raw on-disk paths — sanitized image blocks are stripped from |
| 1306 | # JSONL, and webui replay needs the paths to mint signed URLs. |
| 1307 | user_persisted_early = False |
| 1308 | media_paths = [p for p in (msg.media or []) if isinstance(p, str) and p] |
| 1309 | persist_text = "" |
| 1310 | if isinstance(msg.content, str): |
| 1311 | from nanobot.session.agent_inject import visible_user_content |
| 1312 | |
| 1313 | persist_text = visible_user_content(msg.content) |
| 1314 | has_text = bool(persist_text.strip()) |
| 1315 | if has_text or media_paths: |
| 1316 | extra: dict[str, Any] = {"media": list(media_paths)} if media_paths else {} |
| 1317 | from nanobot.session.question_cards import ( |
| 1318 | apply_following_user_replies_to_question_cards, |
| 1319 | session_has_user_reply, |
| 1320 | ) |
| 1321 | |
| 1322 | if not ( |
| 1323 | has_text and session_has_user_reply(session.messages, persist_text) |
| 1324 | ): |
| 1325 | session.add_message("user", persist_text, **extra) |
| 1326 | apply_following_user_replies_to_question_cards(session.messages) |
| 1327 | self._mark_pending_user_turn(session) |
| 1328 | self.sessions.save(session) |
| 1329 | user_persisted_early = True |
| 1330 | self._schedule_session_preview(session) |
| 1331 | |
| 1332 | final_content, tools_used, all_msgs, stop_reason, had_injections = await self._run_agent_loop( |
| 1333 | initial_messages, |
| 1334 | on_progress=on_progress or _bus_progress, |
| 1335 | on_stream=on_stream, |
| 1336 | on_stream_end=on_stream_end, |
| 1337 | on_retry_wait=_on_retry_wait, |
| 1338 | session=session, |
| 1339 | channel=msg.channel, |
| 1340 | chat_id=msg.chat_id, |
| 1341 | message_id=msg.metadata.get("message_id"), |
| 1342 | pending_queue=pending_queue, |
| 1343 | ) |
| 1344 | final_content, stop_reason, hint_filled_empty = self._attach_shot_count_next_step_hint( |
| 1345 | msg=msg, |
| 1346 | session=session, |
| 1347 | final_content=final_content, |
| 1348 | stop_reason=stop_reason, |
| 1349 | all_msgs=all_msgs, |
| 1350 | tools_used=tools_used, |
| 1351 | ) |
| 1352 | EventStacker.end_turn(stop_reason=stop_reason) |
| 1353 | |
| 1354 | if final_content is None or not final_content.strip(): |
| 1355 | final_content = EMPTY_FINAL_RESPONSE_MESSAGE |
| 1356 | |
| 1357 | # Skip the already-persisted user message when saving the turn |
| 1358 | save_skip = 1 + len(history) + (1 if user_persisted_early else 0) |
| 1359 | self._save_turn(session, all_msgs, save_skip) |
| 1360 | self._clear_pending_user_turn(session) |
| 1361 | self._clear_runtime_checkpoint(session) |
| 1362 | self.sessions.save(session) |
| 1363 | self._schedule_session_preview(session) |
| 1364 | self._schedule_background(self.consolidator.maybe_consolidate_by_tokens(session)) |
| 1365 | |
| 1366 | # When follow-up messages were injected mid-turn, a later natural |
| 1367 | # language reply may address those follow-ups and should not be |
| 1368 | # suppressed just because MessageTool was used earlier in the turn. |
| 1369 | # However, if the turn falls back to the empty-final-response |
| 1370 | # placeholder, suppress it when the real user-visible output already |
| 1371 | # came from MessageTool. |
| 1372 | user_visible_sent = False |
| 1373 | if (mt := self.tools.get("message")) and isinstance(mt, MessageTool) and mt._sent_in_turn: |
| 1374 | user_visible_sent = True |
| 1375 | ask_user_sent = ( |
| 1376 | (aut := self.tools.get("ask_user")) |
| 1377 | and isinstance(aut, AskUserTool) |
| 1378 | and aut._sent_in_turn |
| 1379 | ) |
| 1380 | if ask_user_sent: |
| 1381 | user_visible_sent = True |
| 1382 | if user_visible_sent: |
| 1383 | if not had_injections or stop_reason == "empty_final_response": |
| 1384 | return None |
| 1385 | if _suppresses_frontend_reply( |
| 1386 | msg, |
| 1387 | channel=msg.channel, |
| 1388 | stop_reason=stop_reason, |
| 1389 | final_content=final_content, |
| 1390 | tools_used=tools_used, |
| 1391 | ): |
| 1392 | return None |
| 1393 | |
| 1394 | preview = final_content[:120] + "..." if len(final_content) > 120 else final_content |
| 1395 | logger.info("Response to {}:{}: {}", msg.channel, msg.sender_id, preview) |
| 1396 | |
| 1397 | meta = dict(msg.metadata or {}) |
| 1398 | if on_stream is not None and stop_reason != "error" and not hint_filled_empty: |
| 1399 | meta["_streamed"] = True |
| 1400 | return OutboundMessage( |
| 1401 | channel=msg.channel, |
| 1402 | chat_id=msg.chat_id, |
| 1403 | content=final_content, |
| 1404 | metadata=meta, |
| 1405 | ) |
| 1406 | |
| 1407 | def _sanitize_persisted_blocks( |
| 1408 | self, |
| 1409 | content: list[dict[str, Any]], |
| 1410 | *, |
| 1411 | should_truncate_text: bool = False, |
| 1412 | drop_runtime: bool = False, |
| 1413 | ) -> list[dict[str, Any]]: |
| 1414 | """Strip volatile multimodal payloads before writing session history.""" |
| 1415 | filtered: list[dict[str, Any]] = [] |
| 1416 | for block in content: |
| 1417 | if not isinstance(block, dict): |
| 1418 | filtered.append(block) |
| 1419 | continue |
| 1420 | |
| 1421 | if ( |
| 1422 | drop_runtime |
| 1423 | and block.get("type") == "text" |
| 1424 | and isinstance(block.get("text"), str) |
| 1425 | and block["text"].startswith(ContextBuilder._RUNTIME_CONTEXT_TAG) |
| 1426 | ): |
| 1427 | continue |
| 1428 | |
| 1429 | if block.get("type") == "image_url" and block.get("image_url", {}).get( |
| 1430 | "url", "" |
| 1431 | ).startswith("data:image/"): |
| 1432 | path = (block.get("_meta") or {}).get("path", "") |
| 1433 | filtered.append({"type": "text", "text": image_placeholder_text(path)}) |
| 1434 | continue |
| 1435 | |
| 1436 | if block.get("type") == "text" and isinstance(block.get("text"), str): |
| 1437 | text = block["text"] |
| 1438 | if should_truncate_text and len(text) > self.max_tool_result_chars: |
| 1439 | text = truncate_text_fn(text, self.max_tool_result_chars) |
| 1440 | filtered.append({**block, "text": text}) |
| 1441 | continue |
| 1442 | |
| 1443 | filtered.append(block) |
| 1444 | |
| 1445 | return filtered |
| 1446 | |
| 1447 | def _save_turn(self, session: Session, messages: list[dict], skip: int) -> None: |
| 1448 | """Save new-turn messages into session, truncating large tool results.""" |
| 1449 | from datetime import datetime |
| 1450 | |
| 1451 | from nanobot.session.question_cards import ( |
| 1452 | apply_following_user_replies_to_question_cards, |
| 1453 | attach_batch_id_to_assistant_tool_call, |
| 1454 | extract_batch_id_from_tool_result, |
| 1455 | normalize_persisted_message, |
| 1456 | session_has_ask_user_turn, |
| 1457 | session_has_tool_turn, |
| 1458 | session_has_user_reply, |
| 1459 | should_drop_assistant_blurb_after_cards, |
| 1460 | ) |
| 1461 | |
| 1462 | for m in messages[skip:]: |
| 1463 | entry = dict(m) |
| 1464 | if entry.pop("_hidden", False): |
| 1465 | continue |
| 1466 | role, content = entry.get("role"), entry.get("content") |
| 1467 | if role == "assistant" and is_empty_final_response_message( |
| 1468 | content if isinstance(content, str) else None |
| 1469 | ): |
| 1470 | continue |
| 1471 | if role == "assistant" and not content and not entry.get("tool_calls"): |
| 1472 | continue # skip empty assistant messages — they poison session context |
| 1473 | if role == "tool": |
| 1474 | tool_call_id = entry.get("tool_call_id") |
| 1475 | if ( |
| 1476 | isinstance(tool_call_id, str) |
| 1477 | and tool_call_id.strip() |
| 1478 | and session_has_tool_turn(session.messages, tool_call_id.strip()) |
| 1479 | ): |
| 1480 | continue |
| 1481 | if isinstance(content, str) and len(content) > self.max_tool_result_chars: |
| 1482 | entry["content"] = truncate_text_fn(content, self.max_tool_result_chars) |
| 1483 | elif isinstance(content, list): |
| 1484 | filtered = self._sanitize_persisted_blocks(content, should_truncate_text=True) |
| 1485 | if not filtered: |
| 1486 | continue |
| 1487 | entry["content"] = filtered |
| 1488 | elif role == "assistant": |
| 1489 | tool_calls = entry.get("tool_calls") |
| 1490 | if isinstance(tool_calls, list): |
| 1491 | for tc in tool_calls: |
| 1492 | if not isinstance(tc, dict): |
| 1493 | continue |
| 1494 | tool_call_id = tc.get("id") |
| 1495 | if ( |
| 1496 | isinstance(tool_call_id, str) |
| 1497 | and tool_call_id.strip() |
| 1498 | and session_has_ask_user_turn( |
| 1499 | session.messages, tool_call_id.strip() |
| 1500 | ) |
| 1501 | ): |
| 1502 | entry = None |
| 1503 | break |
| 1504 | if entry is None: |
| 1505 | continue |
| 1506 | elif ( |
| 1507 | isinstance(content, str) |
| 1508 | and content.strip() |
| 1509 | and should_drop_assistant_blurb_after_cards(session.messages) |
| 1510 | ): |
| 1511 | # Drop redundant natural-language blurb before the user taps a card. |
| 1512 | continue |
| 1513 | elif role == "user": |
| 1514 | if isinstance(content, str): |
| 1515 | from nanobot.session.agent_inject import visible_user_content |
| 1516 | |
| 1517 | visible = visible_user_content(content) |
| 1518 | if visible != content: |
| 1519 | if not visible.strip(): |
| 1520 | continue |
| 1521 | entry["content"] = visible |
| 1522 | content = visible |
| 1523 | if ( |
| 1524 | isinstance(content, str) |
| 1525 | and content.strip() |
| 1526 | and session_has_user_reply(session.messages, content) |
| 1527 | ): |
| 1528 | continue |
| 1529 | if isinstance(content, str) and content.startswith(ContextBuilder._RUNTIME_CONTEXT_TAG): |
| 1530 | # Strip the entire runtime-context block (including any session summary). |
| 1531 | # The block is bounded by _RUNTIME_CONTEXT_TAG and _RUNTIME_CONTEXT_END. |
| 1532 | end_marker = ContextBuilder._RUNTIME_CONTEXT_END |
| 1533 | end_pos = content.find(end_marker) |
| 1534 | if end_pos >= 0: |
| 1535 | after = content[end_pos + len(end_marker):].lstrip("\n") |
| 1536 | if after: |
| 1537 | entry["content"] = after |
| 1538 | else: |
| 1539 | continue |
| 1540 | else: |
| 1541 | # Fallback: no end marker found, strip the tag prefix |
| 1542 | after_tag = content[len(ContextBuilder._RUNTIME_CONTEXT_TAG):].lstrip("\n") |
| 1543 | if after_tag.strip(): |
| 1544 | entry["content"] = after_tag |
| 1545 | else: |
| 1546 | continue |
| 1547 | if isinstance(content, list): |
| 1548 | filtered = self._sanitize_persisted_blocks(content, drop_runtime=True) |
| 1549 | if not filtered: |
| 1550 | continue |
| 1551 | entry["content"] = filtered |
| 1552 | entry = normalize_persisted_message(entry) |
| 1553 | entry.setdefault("timestamp", datetime.now().isoformat()) |
| 1554 | session.messages.append(entry) |
| 1555 | if role == "tool" and entry.get("name") == "ask_user": |
| 1556 | tool_call_id = entry.get("tool_call_id") |
| 1557 | batch_id = extract_batch_id_from_tool_result( |
| 1558 | entry.get("content") if isinstance(entry.get("content"), str) else "" |
| 1559 | ) |
| 1560 | if ( |
| 1561 | isinstance(tool_call_id, str) |
| 1562 | and tool_call_id.strip() |
| 1563 | and isinstance(batch_id, str) |
| 1564 | and batch_id.strip() |
| 1565 | ): |
| 1566 | attach_batch_id_to_assistant_tool_call( |
| 1567 | session.messages, |
| 1568 | tool_call_id=tool_call_id.strip(), |
| 1569 | question_batch_id=batch_id.strip(), |
| 1570 | ) |
| 1571 | apply_following_user_replies_to_question_cards(session.messages) |
| 1572 | session.updated_at = datetime.now() |
| 1573 | |
| 1574 | def _persist_subagent_followup(self, session: Session, msg: InboundMessage) -> bool: |
| 1575 | """Persist subagent follow-ups before prompt assembly so history stays durable. |
| 1576 | |
| 1577 | Returns True if a new entry was appended; False if the follow-up was |
| 1578 | deduped (same ``subagent_task_id`` already in session) or carries no |
| 1579 | content worth persisting. |
| 1580 | """ |
| 1581 | if not msg.content: |
| 1582 | return False |
| 1583 | task_id = msg.metadata.get("subagent_task_id") if isinstance(msg.metadata, dict) else None |
| 1584 | if task_id and any( |
| 1585 | m.get("injected_event") == "subagent_result" and m.get("subagent_task_id") == task_id |
| 1586 | for m in session.messages |
| 1587 | ): |
| 1588 | return False |
| 1589 | session.add_message( |
| 1590 | "assistant", |
| 1591 | msg.content, |
| 1592 | sender_id=msg.sender_id, |
| 1593 | injected_event="subagent_result", |
| 1594 | subagent_task_id=task_id, |
| 1595 | ) |
| 1596 | return True |
| 1597 | |
| 1598 | def _set_runtime_checkpoint(self, session: Session, payload: dict[str, Any]) -> None: |
| 1599 | """Persist the latest in-flight turn state into session metadata.""" |
| 1600 | session.metadata[self._RUNTIME_CHECKPOINT_KEY] = payload |
| 1601 | self.sessions.save(session) |
| 1602 | |
| 1603 | def _mark_pending_user_turn(self, session: Session) -> None: |
| 1604 | session.metadata[self._PENDING_USER_TURN_KEY] = True |
| 1605 | |
| 1606 | def _clear_pending_user_turn(self, session: Session) -> None: |
| 1607 | session.metadata.pop(self._PENDING_USER_TURN_KEY, None) |
| 1608 | |
| 1609 | def _clear_runtime_checkpoint(self, session: Session) -> None: |
| 1610 | if self._RUNTIME_CHECKPOINT_KEY in session.metadata: |
| 1611 | session.metadata.pop(self._RUNTIME_CHECKPOINT_KEY, None) |
| 1612 | |
| 1613 | @staticmethod |
| 1614 | def _checkpoint_message_key(message: dict[str, Any]) -> tuple[Any, ...]: |
| 1615 | return ( |
| 1616 | message.get("role"), |
| 1617 | message.get("content"), |
| 1618 | message.get("tool_call_id"), |
| 1619 | message.get("name"), |
| 1620 | message.get("tool_calls"), |
| 1621 | message.get("reasoning_content"), |
| 1622 | message.get("thinking_blocks"), |
| 1623 | ) |
| 1624 | |
| 1625 | def _restore_runtime_checkpoint(self, session: Session) -> bool: |
| 1626 | """Materialize an unfinished turn into session history before a new request.""" |
| 1627 | from datetime import datetime |
| 1628 | |
| 1629 | checkpoint = session.metadata.get(self._RUNTIME_CHECKPOINT_KEY) |
| 1630 | if not isinstance(checkpoint, dict): |
| 1631 | return False |
| 1632 | |
| 1633 | assistant_message = checkpoint.get("assistant_message") |
| 1634 | completed_tool_results = checkpoint.get("completed_tool_results") or [] |
| 1635 | pending_tool_calls = checkpoint.get("pending_tool_calls") or [] |
| 1636 | |
| 1637 | restored_messages: list[dict[str, Any]] = [] |
| 1638 | if isinstance(assistant_message, dict): |
| 1639 | restored = dict(assistant_message) |
| 1640 | restored.setdefault("timestamp", datetime.now().isoformat()) |
| 1641 | restored_messages.append(restored) |
| 1642 | for message in completed_tool_results: |
| 1643 | if isinstance(message, dict): |
| 1644 | restored = dict(message) |
| 1645 | restored.setdefault("timestamp", datetime.now().isoformat()) |
| 1646 | restored_messages.append(restored) |
| 1647 | for tool_call in pending_tool_calls: |
| 1648 | if not isinstance(tool_call, dict): |
| 1649 | continue |
| 1650 | tool_id = tool_call.get("id") |
| 1651 | name = ((tool_call.get("function") or {}).get("name")) or "tool" |
| 1652 | restored_messages.append( |
| 1653 | { |
| 1654 | "role": "tool", |
| 1655 | "tool_call_id": tool_id, |
| 1656 | "name": name, |
| 1657 | "content": "Error: Task interrupted before this tool finished.", |
| 1658 | "timestamp": datetime.now().isoformat(), |
| 1659 | } |
| 1660 | ) |
| 1661 | |
| 1662 | overlap = 0 |
| 1663 | max_overlap = min(len(session.messages), len(restored_messages)) |
| 1664 | for size in range(max_overlap, 0, -1): |
| 1665 | existing = session.messages[-size:] |
| 1666 | restored = restored_messages[:size] |
| 1667 | if all( |
| 1668 | self._checkpoint_message_key(left) == self._checkpoint_message_key(right) |
| 1669 | for left, right in zip(existing, restored) |
| 1670 | ): |
| 1671 | overlap = size |
| 1672 | break |
| 1673 | session.messages.extend(restored_messages[overlap:]) |
| 1674 | |
| 1675 | self._clear_pending_user_turn(session) |
| 1676 | self._clear_runtime_checkpoint(session) |
| 1677 | return True |
| 1678 | |
| 1679 | def _restore_pending_user_turn(self, session: Session) -> bool: |
| 1680 | """Close a turn that only persisted the user message before crashing.""" |
| 1681 | from datetime import datetime |
| 1682 | |
| 1683 | if not session.metadata.get(self._PENDING_USER_TURN_KEY): |
| 1684 | return False |
| 1685 | |
| 1686 | if session.messages and session.messages[-1].get("role") == "user": |
| 1687 | session.messages.append( |
| 1688 | { |
| 1689 | "role": "assistant", |
| 1690 | "content": "Error: Task interrupted before a response was generated.", |
| 1691 | "timestamp": datetime.now().isoformat(), |
| 1692 | } |
| 1693 | ) |
| 1694 | session.updated_at = datetime.now() |
| 1695 | |
| 1696 | self._clear_pending_user_turn(session) |
| 1697 | return True |
| 1698 | |
| 1699 | async def process_direct( |
| 1700 | self, |
| 1701 | content: str, |
| 1702 | session_key: str = "cli:direct", |
| 1703 | channel: str = "cli", |
| 1704 | chat_id: str = "direct", |
| 1705 | media: list[str] | None = None, |
| 1706 | on_progress: Callable[..., Awaitable[None]] | None = None, |
| 1707 | on_stream: Callable[[str], Awaitable[None]] | None = None, |
| 1708 | on_stream_end: Callable[..., Awaitable[None]] | None = None, |
| 1709 | ) -> OutboundMessage | None: |
| 1710 | """Process a message directly and return the outbound payload.""" |
| 1711 | await self._connect_mcp() |
| 1712 | msg = InboundMessage( |
| 1713 | channel=channel, sender_id="user", chat_id=chat_id, |
| 1714 | content=content, media=media or [], |
| 1715 | ) |
| 1716 | return await self._process_message( |
| 1717 | msg, |
| 1718 | session_key=session_key, |
| 1719 | on_progress=on_progress, |
| 1720 | on_stream=on_stream, |
| 1721 | on_stream_end=on_stream_end, |
| 1722 | ) |
| 1723 |