chengyongruandXubin Ren
fc1c8ea770
The runtime media-attachment mechanism was broken for streaming channels (e.g. WebSocket): the _streamed flag caused _send_once to skip the final OutboundMessage that carried generated media, so images were never delivered. Rather than adding complex coordination between streaming and media delivery, delegate image delivery to the LLM: after generate_image returns artifact paths, the next_step prompt now instructs the LLM to call the message tool with the paths in the media parameter. This works uniformly across all channels, streaming or not. Remove generated_media from TurnContext, _assemble_outbound, and _state_save. Update prompts in identity.md, SKILL.md, message tool description, and artifacts.py to reflect the new flow.
1614 lines
66 KiB
Python
1614 lines
66 KiB
Python
"""Agent loop: the core processing engine."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import dataclasses
|
|
import os
|
|
import time
|
|
from contextlib import AsyncExitStack, nullcontext, suppress
|
|
from dataclasses import dataclass, field
|
|
from enum import Enum, auto
|
|
from pathlib import Path
|
|
from typing import TYPE_CHECKING, Any, Awaitable, Callable
|
|
|
|
from loguru import logger
|
|
|
|
from nanobot.agent import model_presets as preset_helpers
|
|
from nanobot.agent.autocompact import AutoCompact
|
|
from nanobot.agent.context import ContextBuilder
|
|
from nanobot.agent.hook import AgentHook, CompositeHook
|
|
from nanobot.agent.memory import Consolidator, Dream
|
|
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.tools.file_state import FileStateStore, bind_file_states, reset_file_states
|
|
from nanobot.agent.tools.message import MessageTool
|
|
from nanobot.agent.tools.registry import ToolRegistry
|
|
from nanobot.agent.tools.self import MyTool
|
|
from nanobot.bus.events import InboundMessage, OutboundMessage
|
|
from nanobot.bus.queue import MessageBus
|
|
from nanobot.command import CommandContext, CommandRouter, register_builtin_commands
|
|
from nanobot.config.schema import AgentDefaults, ModelPresetConfig
|
|
from nanobot.providers.base import LLMProvider
|
|
from nanobot.providers.factory import ProviderSnapshot
|
|
from nanobot.session.goal_state import (
|
|
runner_wall_llm_timeout_s,
|
|
)
|
|
from nanobot.session.manager import Session, SessionManager
|
|
from nanobot.utils.document import extract_documents
|
|
from nanobot.utils.helpers import image_placeholder_text
|
|
from nanobot.utils.helpers import truncate_text as truncate_text_fn
|
|
from nanobot.utils.image_generation_intent import image_generation_prompt
|
|
from nanobot.utils.llm_runtime import LLMRuntime
|
|
from nanobot.utils.runtime import EMPTY_FINAL_RESPONSE_MESSAGE
|
|
from nanobot.utils.webui_turn_helpers import (
|
|
WebuiTurnCoordinator,
|
|
build_bus_progress_callback,
|
|
mark_webui_session,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
from nanobot.config.schema import (
|
|
ChannelsConfig,
|
|
ProviderConfig,
|
|
ToolsConfig,
|
|
)
|
|
from nanobot.cron.service import CronService
|
|
|
|
|
|
UNIFIED_SESSION_KEY = "unified:default"
|
|
|
|
|
|
class TurnState(Enum):
|
|
RESTORE = auto()
|
|
COMPACT = auto()
|
|
COMMAND = auto()
|
|
BUILD = auto()
|
|
RUN = auto()
|
|
SAVE = auto()
|
|
RESPOND = auto()
|
|
DONE = auto()
|
|
|
|
|
|
@dataclass
|
|
class StateTraceEntry:
|
|
state: TurnState
|
|
started_at: float
|
|
duration_ms: float
|
|
event: str
|
|
error: str | None = None
|
|
|
|
|
|
@dataclass
|
|
class TurnContext:
|
|
msg: InboundMessage
|
|
session_key: str
|
|
state: TurnState
|
|
turn_id: str
|
|
session: Session | None = None
|
|
|
|
history: list[dict[str, Any]] = field(default_factory=list)
|
|
initial_messages: list[dict[str, Any]] = field(default_factory=list)
|
|
|
|
final_content: str | None = None
|
|
tools_used: list[str] = field(default_factory=list)
|
|
all_messages: list[dict[str, Any]] = field(default_factory=list)
|
|
stop_reason: str = ""
|
|
had_injections: bool = False
|
|
|
|
user_persisted_early: bool = False
|
|
save_skip: int = 0
|
|
|
|
outbound: OutboundMessage | None = None
|
|
|
|
on_progress: Callable[..., Awaitable[None]] | None = None
|
|
on_stream: Callable[[str], Awaitable[None]] | None = None
|
|
on_stream_end: Callable[..., Awaitable[None]] | None = None
|
|
on_retry_wait: Callable[[str], Awaitable[None]] | None = None
|
|
|
|
pending_queue: asyncio.Queue | None = None
|
|
pending_summary: str | None = None
|
|
|
|
turn_wall_started_at: float = field(default_factory=time.time)
|
|
turn_latency_ms: int | None = None
|
|
|
|
trace: list[StateTraceEntry] = field(default_factory=list)
|
|
|
|
|
|
class AgentLoop:
|
|
"""
|
|
The agent loop is the core processing engine.
|
|
|
|
It:
|
|
1. Receives messages from the bus
|
|
2. Builds context with history, memory, skills
|
|
3. Calls the LLM
|
|
4. Executes tool calls
|
|
5. Sends responses back
|
|
"""
|
|
|
|
@property
|
|
def current_iteration(self) -> int:
|
|
return self._current_iteration
|
|
|
|
@property
|
|
def tool_names(self) -> list[str]:
|
|
return self.tools.tool_names
|
|
|
|
def llm_runtime(self) -> LLMRuntime:
|
|
"""Return the current provider/model pair owned by this loop."""
|
|
self._refresh_provider_snapshot()
|
|
return LLMRuntime(self.provider, self.model)
|
|
|
|
_RUNTIME_CHECKPOINT_KEY = "runtime_checkpoint"
|
|
_PENDING_USER_TURN_KEY = "pending_user_turn"
|
|
|
|
# Event-driven state transition table.
|
|
# Handlers return an event string; the driver looks up the next state here.
|
|
_TRANSITIONS: dict[tuple[TurnState, str], TurnState] = {
|
|
(TurnState.RESTORE, "ok"): TurnState.COMPACT,
|
|
(TurnState.COMPACT, "ok"): TurnState.COMMAND,
|
|
(TurnState.COMMAND, "dispatch"): TurnState.BUILD,
|
|
(TurnState.COMMAND, "shortcut"): TurnState.DONE,
|
|
(TurnState.BUILD, "ok"): TurnState.RUN,
|
|
(TurnState.RUN, "ok"): TurnState.SAVE,
|
|
(TurnState.SAVE, "ok"): TurnState.RESPOND,
|
|
(TurnState.RESPOND, "ok"): TurnState.DONE,
|
|
}
|
|
|
|
def __init__(
|
|
self,
|
|
bus: MessageBus,
|
|
provider: LLMProvider,
|
|
workspace: Path,
|
|
model: str | None = None,
|
|
max_iterations: int | None = None,
|
|
context_window_tokens: int | None = None,
|
|
context_block_limit: int | None = None,
|
|
max_tool_result_chars: int | None = None,
|
|
provider_retry_mode: str = "standard",
|
|
tool_hint_max_length: int | None = None,
|
|
cron_service: CronService | None = None,
|
|
restrict_to_workspace: bool = False,
|
|
session_manager: SessionManager | None = None,
|
|
mcp_servers: dict | None = None,
|
|
channels_config: ChannelsConfig | None = None,
|
|
timezone: str | None = None,
|
|
session_ttl_minutes: int = 0,
|
|
consolidation_ratio: float = 0.5,
|
|
max_messages: int = 120,
|
|
hooks: list[AgentHook] | None = None,
|
|
unified_session: bool = False,
|
|
disabled_skills: list[str] | None = None,
|
|
tools_config: ToolsConfig | None = None,
|
|
image_generation_provider_config: ProviderConfig | None = None,
|
|
image_generation_provider_configs: dict[str, ProviderConfig] | None = None,
|
|
provider_snapshot_loader: Callable[..., ProviderSnapshot] | None = None,
|
|
provider_signature: tuple[object, ...] | None = None,
|
|
model_presets: dict[str, ModelPresetConfig] | None = None,
|
|
model_preset: str | None = None,
|
|
preset_snapshot_loader: preset_helpers.PresetSnapshotLoader | None = None,
|
|
runtime_model_publisher: Callable[[str, str | None], None] | None = None,
|
|
):
|
|
from nanobot.config.schema import ToolsConfig
|
|
|
|
_tc = tools_config or ToolsConfig()
|
|
defaults = AgentDefaults()
|
|
self.bus = bus
|
|
self.channels_config = channels_config
|
|
self.provider = provider
|
|
self._provider_snapshot_loader = provider_snapshot_loader
|
|
self._preset_snapshot_loader = preset_snapshot_loader
|
|
self._runtime_model_publisher = runtime_model_publisher
|
|
self._provider_signature = provider_signature
|
|
self._default_selection_signature = preset_helpers.default_selection_signature(provider_signature)
|
|
self.workspace = workspace
|
|
self.model = model or provider.get_default_model()
|
|
self.max_iterations = (
|
|
max_iterations if max_iterations is not None else defaults.max_tool_iterations
|
|
)
|
|
self.context_window_tokens = (
|
|
context_window_tokens
|
|
if context_window_tokens is not None
|
|
else defaults.context_window_tokens
|
|
)
|
|
self.context_block_limit = context_block_limit
|
|
self.max_tool_result_chars = (
|
|
max_tool_result_chars
|
|
if max_tool_result_chars is not None
|
|
else defaults.max_tool_result_chars
|
|
)
|
|
self.provider_retry_mode = provider_retry_mode
|
|
self.tool_hint_max_length = (
|
|
tool_hint_max_length if tool_hint_max_length is not None
|
|
else defaults.tool_hint_max_length
|
|
)
|
|
self.tools_config = _tc
|
|
self.web_config = _tc.web
|
|
self.exec_config = _tc.exec
|
|
self._image_generation_provider_configs = dict(image_generation_provider_configs or {})
|
|
if (
|
|
image_generation_provider_config is not None
|
|
and "openrouter" not in self._image_generation_provider_configs
|
|
):
|
|
self._image_generation_provider_configs["openrouter"] = image_generation_provider_config
|
|
self.cron_service = cron_service
|
|
self.restrict_to_workspace = restrict_to_workspace
|
|
self._start_time = time.time()
|
|
self._last_usage: dict[str, int] = {}
|
|
self._pending_turn_latency_ms: dict[str, int] = {}
|
|
self._extra_hooks: list[AgentHook] = hooks or []
|
|
|
|
self.context = ContextBuilder(workspace, timezone=timezone, disabled_skills=disabled_skills)
|
|
self.sessions = session_manager or SessionManager(workspace)
|
|
self._webui_turns = WebuiTurnCoordinator(
|
|
bus=self.bus,
|
|
sessions=self.sessions,
|
|
schedule_background=lambda coro: self._schedule_background(coro),
|
|
)
|
|
self.tools = ToolRegistry()
|
|
# One file-read/write tracker per logical session. The tool registry is
|
|
# shared by this loop, so tools resolve the active state via contextvars.
|
|
self._file_state_store = FileStateStore()
|
|
self.runner = AgentRunner(provider)
|
|
self.subagents = SubagentManager(
|
|
provider=provider,
|
|
workspace=workspace,
|
|
bus=bus,
|
|
model=self.model,
|
|
tools_config=_tc,
|
|
max_tool_result_chars=self.max_tool_result_chars,
|
|
restrict_to_workspace=restrict_to_workspace,
|
|
disabled_skills=disabled_skills,
|
|
max_iterations=self.max_iterations,
|
|
llm_wall_timeout_for_session=lambda sk: runner_wall_llm_timeout_s(self.sessions, sk),
|
|
)
|
|
self._unified_session = unified_session
|
|
self._max_messages = max_messages if max_messages > 0 else 120
|
|
self._running = False
|
|
self._mcp_servers = mcp_servers or {}
|
|
self._mcp_stacks: dict[str, AsyncExitStack] = {}
|
|
self._mcp_connected = False
|
|
self._mcp_connecting = False
|
|
self._active_tasks: dict[str, list[asyncio.Task]] = {} # session_key -> tasks
|
|
self._background_tasks: list[asyncio.Task] = []
|
|
self._session_locks: dict[str, asyncio.Lock] = {}
|
|
# Per-session pending queues for mid-turn message injection.
|
|
# When a session has an active task, new messages for that session
|
|
# are routed here instead of creating a new task.
|
|
self._pending_queues: dict[str, asyncio.Queue] = {}
|
|
# NANOBOT_MAX_CONCURRENT_REQUESTS: <=0 means unlimited; default 3.
|
|
_max = int(os.environ.get("NANOBOT_MAX_CONCURRENT_REQUESTS", "3"))
|
|
self._concurrency_gate: asyncio.Semaphore | None = (
|
|
asyncio.Semaphore(_max) if _max > 0 else None
|
|
)
|
|
self.consolidator = Consolidator(
|
|
store=self.context.memory,
|
|
provider=provider,
|
|
model=self.model,
|
|
sessions=self.sessions,
|
|
context_window_tokens=self.context_window_tokens,
|
|
build_messages=self.context.build_messages,
|
|
get_tool_definitions=self.tools.get_definitions,
|
|
max_completion_tokens=provider.generation.max_tokens,
|
|
consolidation_ratio=consolidation_ratio,
|
|
)
|
|
self.auto_compact = AutoCompact(
|
|
sessions=self.sessions,
|
|
consolidator=self.consolidator,
|
|
session_ttl_minutes=session_ttl_minutes,
|
|
)
|
|
self.dream = Dream(
|
|
store=self.context.memory,
|
|
provider=provider,
|
|
model=self.model,
|
|
)
|
|
self.model_presets: dict[str, ModelPresetConfig] = model_presets or {}
|
|
self._active_preset: str | None = None
|
|
if model_preset:
|
|
self.set_model_preset(model_preset, publish_update=False)
|
|
self._register_default_tools()
|
|
self._runtime_vars: dict[str, Any] = {}
|
|
self._current_iteration: int = 0
|
|
self.commands = CommandRouter()
|
|
register_builtin_commands(self.commands)
|
|
|
|
@classmethod
|
|
def from_config(
|
|
cls,
|
|
config: Any,
|
|
bus: MessageBus | None = None,
|
|
**extra: Any,
|
|
) -> AgentLoop:
|
|
"""Create an AgentLoop from config with the common parameter set.
|
|
|
|
Extra keyword arguments are forwarded to ``AgentLoop.__init__``,
|
|
allowing callers to override or extend the standard config-derived
|
|
parameters (e.g. ``cron_service``, ``session_manager``).
|
|
"""
|
|
from nanobot.providers.factory import make_provider
|
|
|
|
if bus is None:
|
|
bus = MessageBus()
|
|
defaults = config.agents.defaults
|
|
provider = extra.pop("provider", None) or make_provider(config)
|
|
resolved = config.resolve_preset()
|
|
model = extra.pop("model", None) or resolved.model
|
|
context_window_tokens = extra.pop("context_window_tokens", None) or resolved.context_window_tokens
|
|
provider_snapshot_loader = extra.pop("provider_snapshot_loader", None)
|
|
preset_snapshot_loader = extra.pop("preset_snapshot_loader", None) or preset_helpers.make_preset_snapshot_loader(
|
|
config,
|
|
provider_snapshot_loader,
|
|
)
|
|
return cls(
|
|
bus=bus,
|
|
provider=provider,
|
|
workspace=config.workspace_path,
|
|
model=model,
|
|
max_iterations=defaults.max_tool_iterations,
|
|
context_window_tokens=context_window_tokens,
|
|
context_block_limit=defaults.context_block_limit,
|
|
max_tool_result_chars=defaults.max_tool_result_chars,
|
|
provider_retry_mode=defaults.provider_retry_mode,
|
|
tool_hint_max_length=defaults.tool_hint_max_length,
|
|
restrict_to_workspace=config.tools.restrict_to_workspace,
|
|
mcp_servers=config.tools.mcp_servers,
|
|
channels_config=config.channels,
|
|
timezone=defaults.timezone,
|
|
unified_session=defaults.unified_session,
|
|
disabled_skills=defaults.disabled_skills,
|
|
session_ttl_minutes=defaults.session_ttl_minutes,
|
|
consolidation_ratio=defaults.consolidation_ratio,
|
|
max_messages=defaults.max_messages,
|
|
tools_config=config.tools,
|
|
model_presets=preset_helpers.configured_model_presets(config),
|
|
model_preset=defaults.model_preset,
|
|
provider_snapshot_loader=provider_snapshot_loader,
|
|
preset_snapshot_loader=preset_snapshot_loader,
|
|
**extra,
|
|
)
|
|
|
|
def _sync_subagent_runtime_limits(self) -> None:
|
|
"""Keep subagent runtime limits aligned with mutable loop settings."""
|
|
self.subagents.max_iterations = self.max_iterations
|
|
|
|
def _apply_provider_snapshot(
|
|
self,
|
|
snapshot: ProviderSnapshot,
|
|
*,
|
|
publish_update: bool = True,
|
|
model_preset: str | None = None,
|
|
) -> None:
|
|
"""Swap model/provider for future turns without disturbing an active one."""
|
|
provider = snapshot.provider
|
|
model = snapshot.model
|
|
context_window_tokens = snapshot.context_window_tokens
|
|
old_model = self.model
|
|
self.provider = provider
|
|
self.model = model
|
|
self.context_window_tokens = context_window_tokens
|
|
self.runner.provider = provider
|
|
self.subagents.set_provider(provider, model)
|
|
self.consolidator.set_provider(provider, model, context_window_tokens)
|
|
self.dream.set_provider(provider, model)
|
|
self._provider_signature = snapshot.signature
|
|
if publish_update and self._runtime_model_publisher is not None:
|
|
self._runtime_model_publisher(
|
|
self.model,
|
|
model_preset if model_preset is not None else self.model_preset,
|
|
)
|
|
logger.info("Runtime model switched for next turn: {} -> {}", old_model, model)
|
|
|
|
def _refresh_provider_snapshot(self) -> None:
|
|
if self._provider_snapshot_loader is None:
|
|
return
|
|
try:
|
|
snapshot = self._provider_snapshot_loader()
|
|
except Exception:
|
|
logger.exception("Failed to refresh provider config")
|
|
return
|
|
default_selection = preset_helpers.default_selection_signature(snapshot.signature)
|
|
if self._active_preset and self._default_selection_signature in (None, default_selection):
|
|
self._default_selection_signature = default_selection
|
|
try:
|
|
snapshot = self._build_model_preset_snapshot(self._active_preset)
|
|
except Exception:
|
|
logger.exception("Failed to refresh active model preset")
|
|
return
|
|
else:
|
|
self._active_preset = None
|
|
self._default_selection_signature = default_selection
|
|
if snapshot.signature == self._provider_signature:
|
|
return
|
|
self._default_selection_signature = preset_helpers.default_selection_signature(snapshot.signature)
|
|
self._apply_provider_snapshot(snapshot)
|
|
|
|
@property
|
|
def model_preset(self) -> str | None:
|
|
return self._active_preset
|
|
|
|
@model_preset.setter
|
|
def model_preset(self, name: str | None) -> None:
|
|
self.set_model_preset(name)
|
|
|
|
def _build_model_preset_snapshot(self, name: str) -> ProviderSnapshot:
|
|
return preset_helpers.build_runtime_preset_snapshot(
|
|
name=name,
|
|
presets=self.model_presets,
|
|
provider=self.provider,
|
|
loader=self._preset_snapshot_loader,
|
|
)
|
|
|
|
def set_model_preset(self, name: str | None, *, publish_update: bool = True) -> None:
|
|
"""Resolve a preset by name and apply all runtime model dependents."""
|
|
name = preset_helpers.normalize_preset_name(name, self.model_presets)
|
|
snapshot = self._build_model_preset_snapshot(name)
|
|
self._apply_provider_snapshot(snapshot, publish_update=publish_update, model_preset=name)
|
|
self._active_preset = name
|
|
|
|
def _register_default_tools(self) -> None:
|
|
"""Register the default set of tools via plugin loader."""
|
|
from nanobot.agent.tools.context import ToolContext
|
|
from nanobot.agent.tools.loader import ToolLoader
|
|
|
|
ctx = ToolContext(
|
|
config=self.tools_config,
|
|
workspace=str(self.workspace),
|
|
bus=self.bus,
|
|
subagent_manager=self.subagents,
|
|
cron_service=self.cron_service,
|
|
sessions=self.sessions,
|
|
provider_snapshot_loader=self._provider_snapshot_loader,
|
|
image_generation_provider_configs=self._image_generation_provider_configs,
|
|
timezone=self.context.timezone or "UTC",
|
|
)
|
|
loader = ToolLoader()
|
|
registered = loader.load(ctx, self.tools)
|
|
|
|
# MyTool needs runtime state reference — manual registration
|
|
if self.tools_config.my.enable:
|
|
self.tools.register(
|
|
MyTool(runtime_state=self, modify_allowed=self.tools_config.my.allow_set)
|
|
)
|
|
registered.append("my")
|
|
|
|
logger.info("Registered {} tools: {}", len(registered), registered)
|
|
|
|
async def _connect_mcp(self) -> None:
|
|
"""Connect to configured MCP servers (one-time, lazy)."""
|
|
if self._mcp_connected or self._mcp_connecting or not self._mcp_servers:
|
|
return
|
|
self._mcp_connecting = True
|
|
from nanobot.agent.tools.mcp import connect_mcp_servers
|
|
|
|
try:
|
|
self._mcp_stacks = await connect_mcp_servers(self._mcp_servers, self.tools)
|
|
if self._mcp_stacks:
|
|
self._mcp_connected = True
|
|
else:
|
|
logger.warning("No MCP servers connected successfully (will retry next message)")
|
|
except asyncio.CancelledError:
|
|
logger.warning("MCP connection cancelled (will retry next message)")
|
|
self._mcp_stacks.clear()
|
|
except BaseException as e:
|
|
logger.warning("Failed to connect MCP servers (will retry next message): {}", e)
|
|
self._mcp_stacks.clear()
|
|
finally:
|
|
self._mcp_connecting = False
|
|
|
|
def _set_tool_context(
|
|
self, channel: str, chat_id: str,
|
|
message_id: str | None = None, metadata: dict | None = None,
|
|
session_key: str | None = None,
|
|
) -> None:
|
|
"""Update context for all tools that need routing info."""
|
|
from nanobot.agent.tools.context import ContextAware, RequestContext
|
|
|
|
if session_key is not None:
|
|
effective_key = session_key
|
|
elif self._unified_session:
|
|
effective_key = UNIFIED_SESSION_KEY
|
|
else:
|
|
effective_key = f"{channel}:{chat_id}"
|
|
|
|
request_ctx = RequestContext(
|
|
channel=channel,
|
|
chat_id=chat_id,
|
|
message_id=message_id,
|
|
session_key=effective_key,
|
|
metadata=dict(metadata or {}),
|
|
)
|
|
|
|
for name in self.tools.tool_names:
|
|
tool = self.tools.get(name)
|
|
if tool and isinstance(tool, ContextAware):
|
|
tool.set_context(request_ctx)
|
|
|
|
@staticmethod
|
|
def _runtime_chat_id(msg: InboundMessage) -> str:
|
|
"""Return the chat id shown in runtime metadata for the model."""
|
|
return str(msg.metadata.get("context_chat_id") or msg.chat_id)
|
|
|
|
async def _build_bus_progress_callback(
|
|
self, msg: InboundMessage
|
|
) -> Callable[..., Awaitable[None]]:
|
|
"""Build a progress callback that publishes to the message bus."""
|
|
return build_bus_progress_callback(self.bus, msg)
|
|
|
|
async def _build_retry_wait_callback(
|
|
self, msg: InboundMessage
|
|
) -> Callable[[str], Awaitable[None]]:
|
|
"""Build a retry-wait callback that publishes to the message bus."""
|
|
|
|
async def _on_retry_wait(content: str) -> None:
|
|
meta = dict(msg.metadata or {})
|
|
meta["_retry_wait"] = True
|
|
await self.bus.publish_outbound(
|
|
OutboundMessage(
|
|
channel=msg.channel,
|
|
chat_id=msg.chat_id,
|
|
content=content,
|
|
metadata=meta,
|
|
)
|
|
)
|
|
|
|
return _on_retry_wait
|
|
|
|
def _persist_user_message_early(
|
|
self,
|
|
msg: InboundMessage,
|
|
session: Session,
|
|
**kwargs: Any,
|
|
) -> bool:
|
|
"""Persist the triggering user message before the turn starts.
|
|
|
|
Returns True if the message was persisted.
|
|
"""
|
|
media_paths = [p for p in (msg.media or []) if isinstance(p, str) and p]
|
|
has_text = isinstance(msg.content, str) and msg.content.strip()
|
|
if has_text or media_paths:
|
|
extra: dict[str, Any] = {"media": list(media_paths)} if media_paths else {}
|
|
extra.update(kwargs)
|
|
text = msg.content if isinstance(msg.content, str) else ""
|
|
session.add_message("user", text, **extra)
|
|
self._mark_pending_user_turn(session)
|
|
self.sessions.save(session)
|
|
return True
|
|
return False
|
|
|
|
def _build_initial_messages(
|
|
self,
|
|
msg: InboundMessage,
|
|
session: Session,
|
|
history: list[dict[str, Any]],
|
|
pending_summary: str | None,
|
|
) -> list[dict[str, Any]]:
|
|
"""Build the initial message list for the LLM turn."""
|
|
return self.context.build_messages(
|
|
history=history,
|
|
current_message=image_generation_prompt(msg.content, msg.metadata),
|
|
media=msg.media if msg.media else None,
|
|
channel=msg.channel,
|
|
chat_id=self._runtime_chat_id(msg),
|
|
sender_id=msg.sender_id,
|
|
session_summary=pending_summary,
|
|
session_metadata=session.metadata,
|
|
)
|
|
|
|
async def _dispatch_command_inline(
|
|
self,
|
|
msg: InboundMessage,
|
|
key: str,
|
|
raw: str,
|
|
dispatch_fn: Callable[[CommandContext], Awaitable[OutboundMessage | None]],
|
|
) -> None:
|
|
"""Dispatch a command directly from the run() loop and publish the result."""
|
|
ctx = CommandContext(msg=msg, session=None, key=key, raw=raw, loop=self)
|
|
result = await dispatch_fn(ctx)
|
|
if result:
|
|
await self.bus.publish_outbound(result)
|
|
else:
|
|
logger.warning("Command '{}' matched but dispatch returned None", raw)
|
|
|
|
async def _cancel_active_tasks(self, key: str) -> int:
|
|
"""Cancel and await all active tasks and subagents for *key*.
|
|
|
|
Returns the total number of cancelled tasks + subagents.
|
|
"""
|
|
tasks = self._active_tasks.pop(key, [])
|
|
cancelled = sum(1 for t in tasks if not t.done() and t.cancel())
|
|
for t in tasks:
|
|
with suppress(asyncio.CancelledError, Exception):
|
|
await t
|
|
sub_cancelled = await self.subagents.cancel_by_session(key)
|
|
return cancelled + sub_cancelled
|
|
|
|
def _effective_session_key(self, msg: InboundMessage) -> str:
|
|
"""Return the session key used for task routing and mid-turn injections."""
|
|
if self._unified_session and not msg.session_key_override:
|
|
return UNIFIED_SESSION_KEY
|
|
return msg.session_key
|
|
|
|
def _replay_token_budget(self) -> int:
|
|
"""Derive a token budget for session history replay from the context window."""
|
|
if self.context_window_tokens <= 0:
|
|
return 0
|
|
max_output = getattr(getattr(self.provider, "generation", None), "max_tokens", 4096)
|
|
try:
|
|
reserved_output = int(max_output)
|
|
except (TypeError, ValueError):
|
|
reserved_output = 4096
|
|
budget = self.context_window_tokens - max(1, reserved_output) - 1024
|
|
return budget if budget > 0 else max(128, self.context_window_tokens // 2)
|
|
|
|
async def _run_agent_loop(
|
|
self,
|
|
initial_messages: list[dict],
|
|
on_progress: Callable[..., Awaitable[None]] | None = None,
|
|
on_stream: Callable[[str], Awaitable[None]] | None = None,
|
|
on_stream_end: Callable[..., Awaitable[None]] | None = None,
|
|
on_retry_wait: Callable[[str], Awaitable[None]] | None = None,
|
|
*,
|
|
session: Session | None = None,
|
|
channel: str = "cli",
|
|
chat_id: str = "direct",
|
|
message_id: str | None = None,
|
|
metadata: dict[str, Any] | None = None,
|
|
session_key: str | None = None,
|
|
pending_queue: asyncio.Queue | None = None,
|
|
) -> tuple[str | None, list[str], list[dict], str, bool]:
|
|
"""Run the agent iteration loop.
|
|
|
|
*on_stream*: called with each content delta during streaming.
|
|
*on_stream_end(resuming)*: called when a streaming session finishes.
|
|
``resuming=True`` means tool calls follow (spinner should restart);
|
|
``resuming=False`` means this is the final response.
|
|
|
|
Returns (final_content, tools_used, messages, stop_reason, had_injections).
|
|
"""
|
|
self._sync_subagent_runtime_limits()
|
|
|
|
loop_hook = AgentProgressHook(
|
|
on_progress=on_progress,
|
|
on_stream=on_stream,
|
|
on_stream_end=on_stream_end,
|
|
channel=channel,
|
|
chat_id=chat_id,
|
|
message_id=message_id,
|
|
metadata=metadata,
|
|
session_key=session_key,
|
|
tool_hint_max_length=self.tool_hint_max_length,
|
|
set_tool_context=self._set_tool_context,
|
|
on_iteration=lambda iteration: setattr(self, "_current_iteration", iteration),
|
|
)
|
|
hook: AgentHook = (
|
|
CompositeHook([loop_hook] + self._extra_hooks) if self._extra_hooks else loop_hook
|
|
)
|
|
|
|
async def _checkpoint(payload: dict[str, Any]) -> None:
|
|
if session is None:
|
|
return
|
|
self._set_runtime_checkpoint(session, payload)
|
|
|
|
async def _drain_pending(*, limit: int = _MAX_INJECTIONS_PER_TURN) -> list[dict[str, Any]]:
|
|
"""Drain follow-up messages from the pending queue.
|
|
|
|
When no messages are immediately available but sub-agents
|
|
spawned in this dispatch are still running, blocks until at
|
|
least one result arrives (or timeout). This keeps the runner
|
|
loop alive so subsequent sub-agent completions are consumed
|
|
in-order rather than dispatched separately.
|
|
"""
|
|
if pending_queue is None:
|
|
return []
|
|
|
|
def _to_user_message(pending_msg: InboundMessage) -> dict[str, Any]:
|
|
content = pending_msg.content
|
|
media = pending_msg.media if pending_msg.media else None
|
|
if media:
|
|
content, media = extract_documents(content, media)
|
|
media = media or None
|
|
user_content = self.context._build_user_content(content, media)
|
|
return {"role": "user", "content": user_content}
|
|
|
|
items: list[dict[str, Any]] = []
|
|
while len(items) < limit:
|
|
try:
|
|
items.append(_to_user_message(pending_queue.get_nowait()))
|
|
except asyncio.QueueEmpty:
|
|
break
|
|
|
|
# Block if nothing drained but sub-agents spawned in this dispatch
|
|
# are still running. Keeps the runner loop alive so subsequent
|
|
# completions are injected in-order rather than dispatched separately.
|
|
if (not items
|
|
and session is not None
|
|
and self.subagents.get_running_count_by_session(session.key) > 0):
|
|
try:
|
|
msg = await asyncio.wait_for(pending_queue.get(), timeout=300)
|
|
except asyncio.TimeoutError:
|
|
logger.warning(
|
|
"Timeout waiting for sub-agent completion in session {}",
|
|
session.key,
|
|
)
|
|
return items
|
|
items.append(_to_user_message(msg))
|
|
while len(items) < limit:
|
|
try:
|
|
items.append(_to_user_message(pending_queue.get_nowait()))
|
|
except asyncio.QueueEmpty:
|
|
break
|
|
|
|
return items
|
|
|
|
active_session_key = session.key if session else session_key
|
|
file_state_token = bind_file_states(self._file_state_store.for_session(active_session_key))
|
|
try:
|
|
result = await self.runner.run(AgentRunSpec(
|
|
initial_messages=initial_messages,
|
|
tools=self.tools,
|
|
model=self.model,
|
|
max_iterations=self.max_iterations,
|
|
max_tool_result_chars=self.max_tool_result_chars,
|
|
hook=hook,
|
|
error_message="Sorry, I encountered an error calling the AI model.",
|
|
concurrent_tools=True,
|
|
workspace=self.workspace,
|
|
session_key=session.key if session else None,
|
|
context_window_tokens=self.context_window_tokens,
|
|
context_block_limit=self.context_block_limit,
|
|
provider_retry_mode=self.provider_retry_mode,
|
|
progress_callback=on_progress,
|
|
stream_progress_deltas=on_stream is not None,
|
|
retry_wait_callback=on_retry_wait,
|
|
checkpoint_callback=_checkpoint,
|
|
injection_callback=_drain_pending,
|
|
# Sustained goals may legitimately exceed NANOBOT_LLM_TIMEOUT_S; idle stall
|
|
# is still capped by NANOBOT_STREAM_IDLE_TIMEOUT_S in streaming providers.
|
|
llm_timeout_s=runner_wall_llm_timeout_s(
|
|
self.sessions,
|
|
session.key if session is not None else session_key,
|
|
metadata=(session.metadata if session is not None else None),
|
|
),
|
|
))
|
|
finally:
|
|
reset_file_states(file_state_token)
|
|
self._last_usage = result.usage
|
|
if result.stop_reason == "max_iterations":
|
|
logger.warning("Max iterations ({}) reached", self.max_iterations)
|
|
# Push final content through stream so streaming channels (e.g. Feishu)
|
|
# update the card instead of leaving it empty.
|
|
if on_stream and on_stream_end:
|
|
await on_stream(result.final_content or "")
|
|
await on_stream_end(resuming=False)
|
|
elif result.stop_reason == "error":
|
|
logger.error("LLM returned error: {}", (result.final_content or "")[:200])
|
|
return result.final_content, result.tools_used, result.messages, result.stop_reason, result.had_injections
|
|
|
|
async def run(self) -> None:
|
|
"""Run the agent loop, dispatching messages as tasks to stay responsive to /stop."""
|
|
self._running = True
|
|
await self._connect_mcp()
|
|
logger.info("Agent loop started")
|
|
|
|
while self._running:
|
|
try:
|
|
msg = await asyncio.wait_for(self.bus.consume_inbound(), timeout=1.0)
|
|
except asyncio.TimeoutError:
|
|
self.auto_compact.check_expired(
|
|
self._schedule_background,
|
|
active_session_keys=self._pending_queues.keys(),
|
|
)
|
|
continue
|
|
except asyncio.CancelledError:
|
|
# Preserve real task cancellation so shutdown can complete cleanly.
|
|
# Only ignore non-task CancelledError signals that may leak from integrations.
|
|
if not self._running or asyncio.current_task().cancelling():
|
|
raise
|
|
continue
|
|
except Exception as e:
|
|
logger.warning("Error consuming inbound message: {}, continuing...", e)
|
|
continue
|
|
|
|
raw = msg.content.strip()
|
|
if self.commands.is_priority(raw):
|
|
await self._dispatch_command_inline(
|
|
msg, msg.session_key, raw,
|
|
self.commands.dispatch_priority,
|
|
)
|
|
continue
|
|
effective_key = self._effective_session_key(msg)
|
|
# If this session already has an active pending queue (i.e. a task
|
|
# is processing this session), route the message there for mid-turn
|
|
# injection instead of creating a competing task.
|
|
if effective_key in self._pending_queues:
|
|
# Non-priority commands must not be queued for injection;
|
|
# dispatch them directly (same pattern as priority commands).
|
|
if self.commands.is_dispatchable_command(raw):
|
|
await self._dispatch_command_inline(
|
|
msg, effective_key, raw,
|
|
self.commands.dispatch,
|
|
)
|
|
continue
|
|
pending_msg = msg
|
|
if effective_key != msg.session_key:
|
|
pending_msg = dataclasses.replace(
|
|
msg,
|
|
session_key_override=effective_key,
|
|
)
|
|
try:
|
|
self._pending_queues[effective_key].put_nowait(pending_msg)
|
|
except asyncio.QueueFull:
|
|
logger.warning(
|
|
"Pending queue full for session {}, falling back to queued task",
|
|
effective_key,
|
|
)
|
|
else:
|
|
logger.info(
|
|
"Routed follow-up message to pending queue for session {}",
|
|
effective_key,
|
|
)
|
|
continue
|
|
# Compute the effective session key before dispatching
|
|
# This ensures /stop command can find tasks correctly when unified session is enabled
|
|
task = asyncio.create_task(self._dispatch(msg))
|
|
self._active_tasks.setdefault(effective_key, []).append(task)
|
|
task.add_done_callback(
|
|
lambda t, k=effective_key: self._active_tasks.get(k, [])
|
|
and self._active_tasks[k].remove(t)
|
|
if t in self._active_tasks.get(k, [])
|
|
else None
|
|
)
|
|
|
|
async def _dispatch(self, msg: InboundMessage) -> None:
|
|
"""Process a message: per-session serial, cross-session concurrent."""
|
|
session_key = self._effective_session_key(msg)
|
|
if session_key != msg.session_key:
|
|
msg = dataclasses.replace(msg, session_key_override=session_key)
|
|
lock = self._session_locks.setdefault(session_key, asyncio.Lock())
|
|
gate = self._concurrency_gate or nullcontext()
|
|
|
|
# Register a pending queue so follow-up messages for this session are
|
|
# routed here (mid-turn injection) instead of spawning a new task.
|
|
pending = asyncio.Queue(maxsize=20)
|
|
self._pending_queues[session_key] = pending
|
|
|
|
try:
|
|
async with lock, gate:
|
|
try:
|
|
on_stream = on_stream_end = None
|
|
if msg.metadata.get("_wants_stream"):
|
|
# Split one answer into distinct stream segments.
|
|
stream_base_id = f"{msg.session_key}:{time.time_ns()}"
|
|
stream_segment = 0
|
|
|
|
def _current_stream_id() -> str:
|
|
return f"{stream_base_id}:{stream_segment}"
|
|
|
|
async def on_stream(delta: str) -> None:
|
|
meta = dict(msg.metadata or {})
|
|
meta["_stream_delta"] = True
|
|
meta["_stream_id"] = _current_stream_id()
|
|
await self.bus.publish_outbound(OutboundMessage(
|
|
channel=msg.channel, chat_id=msg.chat_id,
|
|
content=delta,
|
|
metadata=meta,
|
|
))
|
|
|
|
async def on_stream_end(*, resuming: bool = False) -> None:
|
|
nonlocal stream_segment
|
|
meta = dict(msg.metadata or {})
|
|
meta["_stream_end"] = True
|
|
meta["_resuming"] = resuming
|
|
meta["_stream_id"] = _current_stream_id()
|
|
await self.bus.publish_outbound(OutboundMessage(
|
|
channel=msg.channel, chat_id=msg.chat_id,
|
|
content="",
|
|
metadata=meta,
|
|
))
|
|
stream_segment += 1
|
|
|
|
response = await self._process_message(
|
|
msg, on_stream=on_stream, on_stream_end=on_stream_end,
|
|
pending_queue=pending,
|
|
)
|
|
if response is not None:
|
|
await self.bus.publish_outbound(response)
|
|
elif msg.channel == "cli":
|
|
await self.bus.publish_outbound(OutboundMessage(
|
|
channel=msg.channel, chat_id=msg.chat_id,
|
|
content="", metadata=msg.metadata or {},
|
|
))
|
|
if msg.channel == "websocket":
|
|
turn_lat = self._pending_turn_latency_ms.pop(session_key, None)
|
|
await self._webui_turns.handle_turn_end(
|
|
msg,
|
|
session_key=session_key,
|
|
latency_ms=turn_lat,
|
|
)
|
|
except asyncio.CancelledError:
|
|
logger.info("Task cancelled for session {}", session_key)
|
|
# Preserve partial context from the interrupted turn so
|
|
# the user does not lose tool results and assistant
|
|
# messages accumulated before /stop. The checkpoint was
|
|
# already persisted to session metadata by
|
|
# _emit_checkpoint during tool execution; materializing
|
|
# it into session history now makes it visible in the
|
|
# next conversation turn.
|
|
try:
|
|
key = self._effective_session_key(msg)
|
|
session = self.sessions.get_or_create(key)
|
|
if self._restore_runtime_checkpoint(session):
|
|
self._clear_pending_user_turn(session)
|
|
self.sessions.save(session)
|
|
logger.info(
|
|
"Restored partial context for cancelled session {}",
|
|
key,
|
|
)
|
|
except Exception:
|
|
logger.debug(
|
|
"Could not restore checkpoint for cancelled session {}",
|
|
session_key,
|
|
exc_info=True,
|
|
)
|
|
raise
|
|
except Exception:
|
|
logger.exception("Error processing message for session {}", session_key)
|
|
await self.bus.publish_outbound(OutboundMessage(
|
|
channel=msg.channel, chat_id=msg.chat_id,
|
|
content="Sorry, I encountered an error.",
|
|
))
|
|
finally:
|
|
# Drain any messages still in the pending queue and re-publish
|
|
# them to the bus so they are processed as fresh inbound messages
|
|
# rather than silently lost.
|
|
queue = self._pending_queues.pop(session_key, None)
|
|
if queue is not None:
|
|
leftover = 0
|
|
while True:
|
|
try:
|
|
item = queue.get_nowait()
|
|
except asyncio.QueueEmpty:
|
|
break
|
|
await self.bus.publish_inbound(item)
|
|
leftover += 1
|
|
if leftover:
|
|
logger.info(
|
|
"Re-published {} leftover message(s) to bus for session {}",
|
|
leftover, session_key,
|
|
)
|
|
await self._webui_turns.publish_run_status(msg, "idle")
|
|
self._pending_turn_latency_ms.pop(session_key, None)
|
|
self._webui_turns.discard(session_key)
|
|
|
|
async def close_mcp(self) -> None:
|
|
"""Drain pending background archives, then close MCP connections."""
|
|
if self._background_tasks:
|
|
await asyncio.gather(*self._background_tasks, return_exceptions=True)
|
|
self._background_tasks.clear()
|
|
for name, stack in self._mcp_stacks.items():
|
|
try:
|
|
await stack.aclose()
|
|
except (RuntimeError, BaseExceptionGroup):
|
|
logger.debug("MCP server '{}' cleanup error (can be ignored)", name)
|
|
self._mcp_stacks.clear()
|
|
|
|
def _schedule_background(self, coro) -> None:
|
|
"""Schedule a coroutine as a tracked background task (drained on shutdown)."""
|
|
task = asyncio.create_task(coro)
|
|
self._background_tasks.append(task)
|
|
task.add_done_callback(self._background_tasks.remove)
|
|
|
|
def stop(self) -> None:
|
|
"""Stop the agent loop."""
|
|
self._running = False
|
|
logger.info("Agent loop stopping")
|
|
|
|
async def _process_system_message(
|
|
self,
|
|
msg: InboundMessage,
|
|
session_key: str | None = None,
|
|
on_progress: Callable[..., Awaitable[None]] | None = None,
|
|
on_stream: Callable[[str], Awaitable[None]] | None = None,
|
|
on_stream_end: Callable[..., Awaitable[None]] | None = None,
|
|
pending_queue: asyncio.Queue | None = None,
|
|
) -> OutboundMessage | None:
|
|
"""Process a system inbound message (e.g. subagent announce)."""
|
|
channel, chat_id = (
|
|
msg.chat_id.split(":", 1) if ":" in msg.chat_id else ("cli", msg.chat_id)
|
|
)
|
|
logger.info("Processing system message from {}", msg.sender_id)
|
|
key = msg.session_key_override or f"{channel}:{chat_id}"
|
|
session = self.sessions.get_or_create(key)
|
|
if self._restore_runtime_checkpoint(session):
|
|
self.sessions.save(session)
|
|
if self._restore_pending_user_turn(session):
|
|
self.sessions.save(session)
|
|
|
|
session, pending = self.auto_compact.prepare_session(session, key)
|
|
if pending:
|
|
logger.info("Memory compact triggered for session {}", key)
|
|
|
|
await self.consolidator.maybe_consolidate_by_tokens(
|
|
session,
|
|
replay_max_messages=self._max_messages,
|
|
)
|
|
is_subagent = msg.sender_id == "subagent"
|
|
if is_subagent and self._persist_subagent_followup(session, msg):
|
|
logger.debug("Subagent result persisted for session {}", key)
|
|
self.sessions.save(session)
|
|
self._set_tool_context(
|
|
channel, chat_id, msg.metadata.get("message_id"),
|
|
msg.metadata, session_key=key,
|
|
)
|
|
_hist_kwargs: dict[str, Any] = {
|
|
"max_messages": self._max_messages,
|
|
"max_tokens": self._replay_token_budget(),
|
|
"include_timestamps": True,
|
|
}
|
|
history = session.get_history(**_hist_kwargs)
|
|
current_role = "assistant" if is_subagent else "user"
|
|
|
|
messages = self.context.build_messages(
|
|
history=history,
|
|
current_message="" if is_subagent else msg.content,
|
|
channel=channel,
|
|
chat_id=chat_id,
|
|
current_role=current_role,
|
|
sender_id=msg.sender_id,
|
|
session_summary=pending,
|
|
session_metadata=session.metadata,
|
|
)
|
|
t_wall = time.time()
|
|
final_content, _, all_msgs, stop_reason, _ = await self._run_agent_loop(
|
|
messages, session=session, channel=channel, chat_id=chat_id,
|
|
message_id=msg.metadata.get("message_id"),
|
|
metadata=msg.metadata,
|
|
session_key=key,
|
|
pending_queue=pending_queue,
|
|
)
|
|
wall_done = time.time()
|
|
latency_ms = max(0, int((wall_done - t_wall) * 1000))
|
|
self._save_turn(session, all_msgs, 1 + len(history), turn_latency_ms=latency_ms)
|
|
if channel == "websocket":
|
|
self._pending_turn_latency_ms[key] = latency_ms
|
|
session.enforce_file_cap(on_archive=self.context.memory.raw_archive)
|
|
self._clear_runtime_checkpoint(session)
|
|
self.sessions.save(session)
|
|
self._schedule_background(
|
|
self.consolidator.maybe_consolidate_by_tokens(
|
|
session,
|
|
replay_max_messages=self._max_messages,
|
|
)
|
|
)
|
|
content = final_content or "Background task completed."
|
|
outbound_metadata: dict[str, Any] = {}
|
|
if channel == "slack" and key.startswith("slack:") and key.count(":") >= 2:
|
|
outbound_metadata["slack"] = {"thread_ts": key.split(":", 2)[2]}
|
|
if origin_message_id := msg.metadata.get("origin_message_id"):
|
|
outbound_metadata["origin_message_id"] = origin_message_id
|
|
return OutboundMessage(
|
|
channel=channel,
|
|
chat_id=chat_id,
|
|
content=content,
|
|
metadata=outbound_metadata,
|
|
)
|
|
|
|
async def _process_message(
|
|
self,
|
|
msg: InboundMessage,
|
|
session_key: str | None = None,
|
|
on_progress: Callable[..., Awaitable[None]] | None = None,
|
|
on_stream: Callable[[str], Awaitable[None]] | None = None,
|
|
on_stream_end: Callable[..., Awaitable[None]] | None = None,
|
|
pending_queue: asyncio.Queue | None = None,
|
|
) -> OutboundMessage | None:
|
|
"""Process a single inbound message and return the response."""
|
|
self._refresh_provider_snapshot()
|
|
|
|
if msg.channel == "system":
|
|
return await self._process_system_message(
|
|
msg,
|
|
session_key=session_key,
|
|
on_progress=on_progress,
|
|
on_stream=on_stream,
|
|
on_stream_end=on_stream_end,
|
|
pending_queue=pending_queue,
|
|
)
|
|
|
|
key = session_key or msg.session_key
|
|
ctx = TurnContext(
|
|
msg=msg,
|
|
session=None,
|
|
session_key=key,
|
|
state=TurnState.RESTORE,
|
|
turn_id=f"{key}:{time.time_ns()}",
|
|
on_progress=on_progress,
|
|
on_stream=on_stream,
|
|
on_stream_end=on_stream_end,
|
|
pending_queue=pending_queue,
|
|
)
|
|
|
|
while ctx.state is not TurnState.DONE:
|
|
handler_name = f"_state_{ctx.state.name.lower()}"
|
|
handler = getattr(self, handler_name, None)
|
|
if handler is None:
|
|
raise RuntimeError(f"Missing state handler for {ctx.state}")
|
|
|
|
t0 = time.perf_counter()
|
|
try:
|
|
event = await handler(ctx)
|
|
except Exception:
|
|
duration = (time.perf_counter() - t0) * 1000
|
|
ctx.trace.append(
|
|
StateTraceEntry(
|
|
state=ctx.state,
|
|
started_at=t0,
|
|
duration_ms=duration,
|
|
event="",
|
|
error="exception",
|
|
)
|
|
)
|
|
raise
|
|
|
|
duration = (time.perf_counter() - t0) * 1000
|
|
ctx.trace.append(
|
|
StateTraceEntry(
|
|
state=ctx.state,
|
|
started_at=t0,
|
|
duration_ms=duration,
|
|
event=event,
|
|
)
|
|
)
|
|
logger.debug(
|
|
"[turn {}] State {} took {:.1f}ms -> event {}",
|
|
ctx.turn_id,
|
|
ctx.state.name,
|
|
duration,
|
|
event,
|
|
)
|
|
|
|
next_state = self._TRANSITIONS.get((ctx.state, event))
|
|
if next_state is None:
|
|
raise RuntimeError(
|
|
f"[turn {ctx.turn_id}] No transition from {ctx.state} "
|
|
f"on event {event!r}"
|
|
)
|
|
ctx.state = next_state
|
|
|
|
logger.debug(
|
|
"[turn {}] Turn completed after {} states",
|
|
ctx.turn_id,
|
|
len(ctx.trace),
|
|
)
|
|
return ctx.outbound
|
|
|
|
def _assemble_outbound(
|
|
self,
|
|
msg: InboundMessage,
|
|
final_content: str,
|
|
all_msgs: list[dict[str, Any]],
|
|
stop_reason: str,
|
|
had_injections: bool,
|
|
on_stream: Callable[[str], Awaitable[None]] | None,
|
|
*,
|
|
turn_latency_ms: int | None = None,
|
|
) -> OutboundMessage | None:
|
|
"""Assemble the final outbound message from turn results."""
|
|
# MessageTool suppression
|
|
if (mt := self.tools.get("message")) and isinstance(mt, MessageTool) and mt._sent_in_turn:
|
|
if not had_injections or stop_reason == "empty_final_response":
|
|
return None
|
|
|
|
preview = final_content[:120] + "..." if len(final_content) > 120 else final_content
|
|
logger.info("Response to {}:{}: {}", msg.channel, msg.sender_id, preview)
|
|
|
|
meta = dict(msg.metadata or {})
|
|
if on_stream is not None and stop_reason not in {"error", "tool_error"}:
|
|
meta["_streamed"] = True
|
|
if turn_latency_ms is not None:
|
|
meta["latency_ms"] = int(turn_latency_ms)
|
|
|
|
return OutboundMessage(
|
|
channel=msg.channel,
|
|
chat_id=msg.chat_id,
|
|
content=final_content,
|
|
metadata=meta,
|
|
)
|
|
|
|
async def _state_restore(self, ctx: TurnContext) -> TurnState:
|
|
"""Restore checkpoint / pending user turn; extract documents."""
|
|
msg = ctx.msg
|
|
|
|
if msg.media:
|
|
new_content, image_only = extract_documents(msg.content, msg.media)
|
|
ctx.msg = dataclasses.replace(msg, content=new_content, media=image_only)
|
|
msg = ctx.msg
|
|
|
|
preview = msg.content[:80] + "..." if len(msg.content) > 80 else msg.content
|
|
logger.info("Processing message from {}:{}: {}", msg.channel, msg.sender_id, preview)
|
|
|
|
# Session is already fetched by the caller (_process_message) but
|
|
# ensure it exists in case this handler is invoked independently.
|
|
if ctx.session is None:
|
|
ctx.session = self.sessions.get_or_create(ctx.session_key)
|
|
mark_webui_session(ctx.session, msg.metadata)
|
|
|
|
if self._restore_runtime_checkpoint(ctx.session):
|
|
self.sessions.save(ctx.session)
|
|
if self._restore_pending_user_turn(ctx.session):
|
|
self.sessions.save(ctx.session)
|
|
|
|
return "ok"
|
|
|
|
async def _state_compact(self, ctx: TurnContext) -> str:
|
|
ctx.session, pending = self.auto_compact.prepare_session(ctx.session, ctx.session_key)
|
|
ctx.pending_summary = pending
|
|
return "ok"
|
|
|
|
async def _state_command(self, ctx: TurnContext) -> str:
|
|
raw = ctx.msg.content.strip()
|
|
cmd_ctx = CommandContext(
|
|
msg=ctx.msg, session=ctx.session, key=ctx.session_key, raw=raw, loop=self
|
|
)
|
|
result = await self.commands.dispatch(cmd_ctx)
|
|
if result is not None:
|
|
ctx.outbound = result
|
|
# Shortcut commands skip BUILD and SAVE, so we must persist the
|
|
# turn here so WebUI history hydration after _turn_end sees the
|
|
# message. Mark messages with _command so get_history can filter
|
|
# them out of LLM context. /new is excluded because it
|
|
# intentionally clears the session.
|
|
if raw.lower() != "/new":
|
|
ctx.user_persisted_early = self._persist_user_message_early(
|
|
ctx.msg, ctx.session, _command=True
|
|
)
|
|
ctx.session.add_message(
|
|
"assistant", result.content, _command=True
|
|
)
|
|
self.sessions.save(ctx.session)
|
|
self._clear_pending_user_turn(ctx.session)
|
|
return "shortcut"
|
|
return "dispatch"
|
|
|
|
async def _state_build(self, ctx: TurnContext) -> str:
|
|
await self.consolidator.maybe_consolidate_by_tokens(
|
|
ctx.session,
|
|
replay_max_messages=self._max_messages,
|
|
)
|
|
self._set_tool_context(
|
|
ctx.msg.channel,
|
|
ctx.msg.chat_id,
|
|
ctx.msg.metadata.get("message_id"),
|
|
ctx.msg.metadata,
|
|
session_key=ctx.session_key,
|
|
)
|
|
if message_tool := self.tools.get("message"):
|
|
if isinstance(message_tool, MessageTool):
|
|
message_tool.start_turn()
|
|
|
|
_hist_kwargs: dict[str, Any] = {
|
|
"max_messages": self._max_messages,
|
|
"max_tokens": self._replay_token_budget(),
|
|
"include_timestamps": True,
|
|
}
|
|
ctx.history = ctx.session.get_history(**_hist_kwargs)
|
|
self._webui_turns.capture_title_context(
|
|
ctx.session_key,
|
|
ctx.msg,
|
|
self.llm_runtime(),
|
|
)
|
|
|
|
ctx.initial_messages = self._build_initial_messages(
|
|
ctx.msg, ctx.session, ctx.history, ctx.pending_summary
|
|
)
|
|
ctx.user_persisted_early = self._persist_user_message_early(
|
|
ctx.msg, ctx.session
|
|
)
|
|
|
|
if ctx.on_progress is None:
|
|
ctx.on_progress = await self._build_bus_progress_callback(ctx.msg)
|
|
if ctx.on_retry_wait is None:
|
|
ctx.on_retry_wait = await self._build_retry_wait_callback(ctx.msg)
|
|
|
|
return "ok"
|
|
|
|
async def _state_run(self, ctx: TurnContext) -> str:
|
|
await self._webui_turns.publish_run_status(ctx.msg, "running")
|
|
result = await self._run_agent_loop(
|
|
ctx.initial_messages,
|
|
on_progress=ctx.on_progress,
|
|
on_stream=ctx.on_stream,
|
|
on_stream_end=ctx.on_stream_end,
|
|
on_retry_wait=ctx.on_retry_wait,
|
|
session=ctx.session,
|
|
channel=ctx.msg.channel,
|
|
chat_id=ctx.msg.chat_id,
|
|
message_id=ctx.msg.metadata.get("message_id"),
|
|
metadata=ctx.msg.metadata,
|
|
session_key=ctx.session_key,
|
|
pending_queue=ctx.pending_queue,
|
|
)
|
|
final_content, tools_used, all_msgs, stop_reason, had_injections = result
|
|
ctx.final_content = final_content
|
|
ctx.tools_used = tools_used
|
|
ctx.all_messages = all_msgs
|
|
ctx.stop_reason = stop_reason
|
|
ctx.had_injections = had_injections
|
|
return "ok"
|
|
|
|
async def _state_save(self, ctx: TurnContext) -> str:
|
|
if ctx.final_content is None or not ctx.final_content.strip():
|
|
ctx.final_content = EMPTY_FINAL_RESPONSE_MESSAGE
|
|
|
|
ctx.save_skip = 1 + len(ctx.history) + (1 if ctx.user_persisted_early else 0)
|
|
|
|
ctx.turn_latency_ms = max(0, int((time.time() - ctx.turn_wall_started_at) * 1000))
|
|
self._save_turn(
|
|
ctx.session, ctx.all_messages, ctx.save_skip,
|
|
turn_latency_ms=ctx.turn_latency_ms,
|
|
)
|
|
if ctx.msg.channel == "websocket":
|
|
self._pending_turn_latency_ms[ctx.session_key] = ctx.turn_latency_ms
|
|
ctx.session.enforce_file_cap(on_archive=self.context.memory.raw_archive)
|
|
self._clear_pending_user_turn(ctx.session)
|
|
self._clear_runtime_checkpoint(ctx.session)
|
|
self.sessions.save(ctx.session)
|
|
self._schedule_background(
|
|
self.consolidator.maybe_consolidate_by_tokens(
|
|
ctx.session,
|
|
replay_max_messages=self._max_messages,
|
|
)
|
|
)
|
|
return "ok"
|
|
|
|
async def _state_respond(self, ctx: TurnContext) -> str:
|
|
ctx.outbound = self._assemble_outbound(
|
|
ctx.msg,
|
|
ctx.final_content,
|
|
ctx.all_messages,
|
|
ctx.stop_reason,
|
|
ctx.had_injections,
|
|
ctx.on_stream,
|
|
turn_latency_ms=ctx.turn_latency_ms,
|
|
)
|
|
return "ok"
|
|
|
|
def _sanitize_persisted_blocks(
|
|
self,
|
|
content: list[dict[str, Any]],
|
|
*,
|
|
should_truncate_text: bool = False,
|
|
drop_runtime: bool = False,
|
|
) -> list[dict[str, Any]]:
|
|
"""Strip volatile multimodal payloads before writing session history."""
|
|
filtered: list[dict[str, Any]] = []
|
|
for block in content:
|
|
if not isinstance(block, dict):
|
|
filtered.append(block)
|
|
continue
|
|
|
|
if (
|
|
drop_runtime
|
|
and block.get("type") == "text"
|
|
and isinstance(block.get("text"), str)
|
|
and block["text"].startswith(ContextBuilder._RUNTIME_CONTEXT_TAG)
|
|
):
|
|
continue
|
|
|
|
if block.get("type") == "image_url" and block.get("image_url", {}).get(
|
|
"url", ""
|
|
).startswith("data:image/"):
|
|
path = (block.get("_meta") or {}).get("path", "")
|
|
filtered.append({"type": "text", "text": image_placeholder_text(path)})
|
|
continue
|
|
|
|
if block.get("type") == "text" and isinstance(block.get("text"), str):
|
|
text = block["text"]
|
|
if should_truncate_text and len(text) > self.max_tool_result_chars:
|
|
text = truncate_text_fn(text, self.max_tool_result_chars)
|
|
filtered.append({**block, "text": text})
|
|
continue
|
|
|
|
filtered.append(block)
|
|
|
|
return filtered
|
|
|
|
def _save_turn(
|
|
self,
|
|
session: Session,
|
|
messages: list[dict],
|
|
skip: int,
|
|
*,
|
|
turn_latency_ms: int | None = None,
|
|
) -> None:
|
|
"""Save new-turn messages into session, truncating large tool results."""
|
|
from datetime import datetime
|
|
|
|
last_assistant_idx: int | None = None
|
|
for m in messages[skip:]:
|
|
entry = dict(m)
|
|
role, content = entry.get("role"), entry.get("content")
|
|
if role == "assistant" and not content and not entry.get("tool_calls"):
|
|
continue # skip empty assistant messages — they poison session context
|
|
if role == "tool":
|
|
if isinstance(content, str) and len(content) > self.max_tool_result_chars:
|
|
entry["content"] = truncate_text_fn(content, self.max_tool_result_chars)
|
|
elif isinstance(content, list):
|
|
filtered = self._sanitize_persisted_blocks(content, should_truncate_text=True)
|
|
if not filtered:
|
|
continue
|
|
entry["content"] = filtered
|
|
elif role == "user":
|
|
if isinstance(content, str) and ContextBuilder._RUNTIME_CONTEXT_TAG in content:
|
|
# Strip the runtime-context block appended at the end.
|
|
tag_pos = content.find(ContextBuilder._RUNTIME_CONTEXT_TAG)
|
|
before = content[:tag_pos].rstrip("\n ")
|
|
if before:
|
|
entry["content"] = before
|
|
else:
|
|
continue
|
|
if isinstance(content, list):
|
|
filtered = self._sanitize_persisted_blocks(content, drop_runtime=True)
|
|
if not filtered:
|
|
continue
|
|
entry["content"] = filtered
|
|
entry.setdefault("timestamp", datetime.now().isoformat())
|
|
session.messages.append(entry)
|
|
if role == "assistant":
|
|
last_assistant_idx = len(session.messages) - 1
|
|
if turn_latency_ms is not None and last_assistant_idx is not None:
|
|
session.messages[last_assistant_idx]["latency_ms"] = int(turn_latency_ms)
|
|
session.updated_at = datetime.now()
|
|
|
|
def _persist_subagent_followup(self, session: Session, msg: InboundMessage) -> bool:
|
|
"""Persist subagent follow-ups before prompt assembly so history stays durable.
|
|
|
|
Returns True if a new entry was appended; False if the follow-up was
|
|
deduped (same ``subagent_task_id`` already in session) or carries no
|
|
content worth persisting.
|
|
"""
|
|
if not msg.content:
|
|
return False
|
|
task_id = msg.metadata.get("subagent_task_id") if isinstance(msg.metadata, dict) else None
|
|
if task_id and any(
|
|
m.get("injected_event") == "subagent_result" and m.get("subagent_task_id") == task_id
|
|
for m in session.messages
|
|
):
|
|
return False
|
|
session.add_message(
|
|
"assistant",
|
|
msg.content,
|
|
sender_id=msg.sender_id,
|
|
injected_event="subagent_result",
|
|
subagent_task_id=task_id,
|
|
)
|
|
return True
|
|
|
|
def _set_runtime_checkpoint(self, session: Session, payload: dict[str, Any]) -> None:
|
|
"""Persist the latest in-flight turn state into session metadata."""
|
|
session.metadata[self._RUNTIME_CHECKPOINT_KEY] = payload
|
|
self.sessions.save(session)
|
|
|
|
def _mark_pending_user_turn(self, session: Session) -> None:
|
|
session.metadata[self._PENDING_USER_TURN_KEY] = True
|
|
|
|
def _clear_pending_user_turn(self, session: Session) -> None:
|
|
session.metadata.pop(self._PENDING_USER_TURN_KEY, None)
|
|
|
|
def _clear_runtime_checkpoint(self, session: Session) -> None:
|
|
if self._RUNTIME_CHECKPOINT_KEY in session.metadata:
|
|
session.metadata.pop(self._RUNTIME_CHECKPOINT_KEY, None)
|
|
|
|
@staticmethod
|
|
def _checkpoint_message_key(message: dict[str, Any]) -> tuple[Any, ...]:
|
|
return (
|
|
message.get("role"),
|
|
message.get("content"),
|
|
message.get("tool_call_id"),
|
|
message.get("name"),
|
|
message.get("tool_calls"),
|
|
message.get("reasoning_content"),
|
|
message.get("thinking_blocks"),
|
|
)
|
|
|
|
def _restore_runtime_checkpoint(self, session: Session) -> bool:
|
|
"""Materialize an unfinished turn into session history before a new request."""
|
|
from datetime import datetime
|
|
|
|
checkpoint = session.metadata.get(self._RUNTIME_CHECKPOINT_KEY)
|
|
if not isinstance(checkpoint, dict):
|
|
return False
|
|
|
|
assistant_message = checkpoint.get("assistant_message")
|
|
completed_tool_results = checkpoint.get("completed_tool_results") or []
|
|
pending_tool_calls = checkpoint.get("pending_tool_calls") or []
|
|
|
|
restored_messages: list[dict[str, Any]] = []
|
|
if isinstance(assistant_message, dict):
|
|
restored = dict(assistant_message)
|
|
restored.setdefault("timestamp", datetime.now().isoformat())
|
|
restored_messages.append(restored)
|
|
for message in completed_tool_results:
|
|
if isinstance(message, dict):
|
|
restored = dict(message)
|
|
restored.setdefault("timestamp", datetime.now().isoformat())
|
|
restored_messages.append(restored)
|
|
for tool_call in pending_tool_calls:
|
|
if not isinstance(tool_call, dict):
|
|
continue
|
|
tool_id = tool_call.get("id")
|
|
name = ((tool_call.get("function") or {}).get("name")) or "tool"
|
|
restored_messages.append(
|
|
{
|
|
"role": "tool",
|
|
"tool_call_id": tool_id,
|
|
"name": name,
|
|
"content": "Error: Task interrupted before this tool finished.",
|
|
"timestamp": datetime.now().isoformat(),
|
|
}
|
|
)
|
|
|
|
overlap = 0
|
|
max_overlap = min(len(session.messages), len(restored_messages))
|
|
for size in range(max_overlap, 0, -1):
|
|
existing = session.messages[-size:]
|
|
restored = restored_messages[:size]
|
|
if all(
|
|
self._checkpoint_message_key(left) == self._checkpoint_message_key(right)
|
|
for left, right in zip(existing, restored)
|
|
):
|
|
overlap = size
|
|
break
|
|
session.messages.extend(restored_messages[overlap:])
|
|
|
|
self._clear_pending_user_turn(session)
|
|
self._clear_runtime_checkpoint(session)
|
|
return True
|
|
|
|
def _restore_pending_user_turn(self, session: Session) -> bool:
|
|
"""Close a turn that only persisted the user message before crashing."""
|
|
from datetime import datetime
|
|
|
|
if not session.metadata.get(self._PENDING_USER_TURN_KEY):
|
|
return False
|
|
|
|
if session.messages and session.messages[-1].get("role") == "user":
|
|
session.messages.append(
|
|
{
|
|
"role": "assistant",
|
|
"content": "Error: Task interrupted before a response was generated.",
|
|
"timestamp": datetime.now().isoformat(),
|
|
}
|
|
)
|
|
session.updated_at = datetime.now()
|
|
|
|
self._clear_pending_user_turn(session)
|
|
return True
|
|
|
|
async def process_direct(
|
|
self,
|
|
content: str,
|
|
session_key: str = "cli:direct",
|
|
channel: str = "cli",
|
|
chat_id: str = "direct",
|
|
media: list[str] | None = None,
|
|
on_progress: Callable[..., Awaitable[None]] | None = None,
|
|
on_stream: Callable[[str], Awaitable[None]] | None = None,
|
|
on_stream_end: Callable[..., Awaitable[None]] | None = None,
|
|
) -> OutboundMessage | None:
|
|
"""Process a message directly and return the outbound payload."""
|
|
await self._connect_mcp()
|
|
msg = InboundMessage(
|
|
channel=channel, sender_id="user", chat_id=chat_id,
|
|
content=content, media=media or [],
|
|
)
|
|
return await self._process_message(
|
|
msg,
|
|
session_key=session_key,
|
|
on_progress=on_progress,
|
|
on_stream=on_stream,
|
|
on_stream_end=on_stream_end,
|
|
)
|