Compare commits

...
Author SHA1 Message Date
chengyongru be3a42ebac fix(gateway): keep event loop responsive (NAN-33)
Move blocking filesystem, persistence, subprocess, media, and DNS work off the gateway event loop while preserving existing contracts. Add bounded cancellation and responsiveness regression coverage.
2026-08-25 10:18:13 +08:00
Xubin Ren 7fb0811fbb fix(agent): honor selected project workspace 2026-08-25 02:14:02 +08:00
chengyongruandchengyongru 2ac802b2d5 feat(usage): add unified provider usage backend 2026-08-25 01:22:25 +08:00
chengyongruandchengyongru 8bb3828487 fix(tui): preserve usage telemetry fields 2026-08-25 01:04:25 +08:00
chengyongruandchengyongru 9895c23cb5 refactor(providers): define typed usage contract 2026-08-25 01:04:25 +08:00
chengyongruandchengyongru 89c94d8744 test(exec): remove Windows process timing races 2026-08-25 00:53:54 +08:00
chrischen-coderandchengyongru f5e467626d fix(agent): time out no-tools model requests 2026-08-25 00:34:01 +08:00
chengyongruandGitHub 04974b7607 fix(webui): preserve causal message order (NAN-29) (#5503) 2026-08-24 15:08:22 +08:00
chengyongruandGitHub 7f288a49fc fix(tui): preserve shell after Ctrl+C (#5502) 2026-08-24 14:32:48 +08:00
chengyongruandchengyongru 09d3bd76c9 fix(exec): disable command guard in full access 2026-08-24 11:44:00 +08:00
chengyongruandchengyongru b1cadf53c5 fix(codex): reuse TLS contexts across requests 2026-08-24 11:25:02 +08:00
Xubin Ren baa0233377 fix(tui): preserve draft scope until first message 2026-08-24 10:40:16 +08:00
Xubin Ren d50a2fab32 fix(tui): avoid saving empty sessions 2026-08-24 10:40:16 +08:00
Xubin Ren 8344066696 style(tui): frame recovery decisions 2026-08-24 00:58:04 +08:00
Xubin Ren 5accc903a0 refactor(tui): simplify session rows 2026-08-24 00:58:04 +08:00
Xubin Ren 1f0771c555 feat(tui): refine session status navigation 2026-08-24 00:58:04 +08:00
Xubin Ren 2850114eab feat(tui): surface session activity states 2026-08-24 00:58:04 +08:00
Xubin Ren 7e66375f59 fix(tui): clarify interrupted task actions 2026-08-24 00:58:04 +08:00
Xubin Ren 2cdfba38b2 fix(runtime): preserve interrupted turns on gateway exit 2026-08-24 00:58:04 +08:00
Xubin Ren c7e2a474a0 fix(runtime): normalize recovery session routing 2026-08-24 00:58:04 +08:00
Xubin Ren 58a1cc48d8 fix(tui): allow switching active sessions 2026-08-24 00:58:04 +08:00
Xubin Ren 41a2104244 docs(readme): simplify terminal quick start 2026-08-24 00:58:04 +08:00
Xubin Ren 12029f8812 feat(runtime): add user-controlled turn recovery 2026-08-24 00:58:04 +08:00
219 changed files with 14046 additions and 2565 deletions
+7 -1
View File
@@ -209,7 +209,13 @@ Use `nanobot gateway --background` for the same direct entry point without keepi
nanobot agent nanobot agent
``` ```
This opens the native terminal client with the configured model and tools, using the launch directory as its workspace. Use `/sessions` to switch saved conversations, `/new-chat` to preserve this conversation and start another one, `/branch` to fork from a completed reply, `/context` to inspect the compacted summary and raw message suffix available to the agent, or `/diff` to review the latest turn's file changes. Type `@` to mention an installed app, configured MCP server, or saved session. While nanobot is working, `Enter` steers the current turn, `Tab` queues a visible follow-up for the next turn, and `Option+Up` on macOS (`Alt+Up` on Windows/Linux) returns the latest queued message for editing. Press `Shift+Enter` to add a newline; `Ctrl+J` is the universal fallback for terminals that cannot distinguish modified Enter keys. Use `PageUp` at the top to load earlier transcript pages. Each launch starts a new session; `--session` selects an existing WebSocket session, while `--workspace` overrides the launch directory. Use `--classic` to resume a session from another channel. The existing nanobot `/new` command keeps its original behavior: it resets the current chat. `nanobot agent` and `nanobot webui` share one on-demand local gateway: either command can start it, each launcher releases only its own client, and the last interactive launcher to exit stops it. Use `/detach` to close the TUI while keeping the gateway and any active agent turn running in the background; after the terminal is restored, nanobot prints the exact `nanobot gateway stop` command for that config and workspace. Use `nanobot gateway --background` to start persistently before opening a client. Type `exit` or press `Ctrl+C` when you are done; after the terminal is restored, nanobot prints a ready-to-run `nanobot agent --session ...` command that resumes the session. Use `nanobot agent --classic` for the legacy Python prompt. This opens the native terminal client with the launch directory as its workspace. It shares saved conversations and the local gateway with the WebUI.
- Type `/` to discover commands, `/sessions` to switch conversations, or `@` to mention an app, MCP server, or saved session.
- Press `Enter` to send or steer, `Tab` to queue a follow-up, and `Shift+Enter` to add a newline (`Ctrl+J` works in terminals that cannot distinguish modified Enter keys).
- Use `/detach` to leave the current task running, or start with `nanobot gateway --background` when nanobot should stay online after all local clients exit.
Each launch starts a new session by default. Use `--session` to resume one and `--workspace` to choose another workspace. See the [CLI reference](./docs/cli-reference.md#agent-cli) for session branching, diffs, history, shortcuts, gateway lifecycle, and compatibility options.
For one request and an immediate exit, use: For one request and an immediate exit, use:
+3
View File
@@ -23,6 +23,7 @@ if TYPE_CHECKING:
STREAM_EVENT_TOOL_FAILED, STREAM_EVENT_TOOL_FAILED,
STREAM_EVENT_TOOL_STARTED, STREAM_EVENT_TOOL_STARTED,
STREAM_EVENT_TYPES, STREAM_EVENT_TYPES,
LLMUsage,
Nanobot, Nanobot,
RunResult, RunResult,
RunStream, RunStream,
@@ -56,6 +57,7 @@ __logo__ = "🐈"
_LAZY_EXPORTS = { _LAZY_EXPORTS = {
"Nanobot": ".nanobot", "Nanobot": ".nanobot",
"LLMUsage": ".nanobot",
"RunStream": ".nanobot", "RunStream": ".nanobot",
"RunResult": ".nanobot", "RunResult": ".nanobot",
"RequestContext": ".agent.tools.context", "RequestContext": ".agent.tools.context",
@@ -93,6 +95,7 @@ def __getattr__(name: str) -> Any:
__all__ = [ __all__ = [
"Nanobot", "Nanobot",
"LLMUsage",
"RunResult", "RunResult",
"RequestContext", "RequestContext",
"RuntimeContextBlock", "RuntimeContextBlock",
+94 -7
View File
@@ -2,7 +2,9 @@
from __future__ import annotations from __future__ import annotations
from collections.abc import Collection import asyncio
import inspect
from collections.abc import Awaitable, Collection
from datetime import datetime from datetime import datetime
from typing import TYPE_CHECKING, Any, Callable, Coroutine from typing import TYPE_CHECKING, Any, Callable, Coroutine
@@ -47,9 +49,27 @@ class AutoCompact:
return idle_seconds >= self._ttl * 60 return idle_seconds >= self._ttl * 60
def _has_unarchived_messages(self, key: str) -> bool: def _has_unarchived_messages(self, key: str) -> bool:
session = self.sessions.get_or_create(key) return self._session_has_unarchived_messages(self.sessions.get_or_create(key))
@staticmethod
def _session_has_unarchived_messages(session: Session) -> bool:
return session.last_consolidated < len(session.messages) return session.last_consolidated < len(session.messages)
def _has_native_async_session_method(self, name: str) -> bool:
"""Check the manager's real class, not mock-generated instance attributes."""
method = inspect.getattr_static(type(self.sessions), name, None)
return inspect.iscoroutinefunction(method)
async def _list_sessions_nonblocking(self) -> list[dict[str, Any]]:
if self._has_native_async_session_method("list_sessions_async"):
return await self.sessions.list_sessions_async()
return await asyncio.to_thread(self.sessions.list_sessions)
async def _get_or_create_nonblocking(self, key: str) -> Session:
if self._has_native_async_session_method("get_or_create_async"):
return await self.sessions.get_or_create_async(key)
return await asyncio.to_thread(self.sessions.get_or_create, key)
@classmethod @classmethod
def _is_internal_session(cls, key: str) -> bool: def _is_internal_session(cls, key: str) -> bool:
return key.startswith(cls._INTERNAL_SESSION_PREFIXES) return key.startswith(cls._INTERNAL_SESSION_PREFIXES)
@@ -79,6 +99,31 @@ class AutoCompact:
self._archiving.add(key) self._archiving.add(key)
schedule_background(self._archive(key, runtime=runtime)) schedule_background(self._archive(key, runtime=runtime))
async def check_expired_async(
self,
schedule_background: Callable[[Coroutine[Any, Any, None]], None],
resolve_runtime: Callable[[Session], Awaitable[LLMRuntime]],
active_session_keys: Collection[str] = (),
) -> None:
"""Schedule idle archival without blocking the event loop."""
now = datetime.now()
active_keys = set(active_session_keys)
for info in await self._list_sessions_nonblocking():
key = info.get("key", "")
if not key or self._is_internal_session(key) or key in self._archiving:
continue
if key in active_keys or not self._is_expired(info.get("updated_at"), now):
continue
session = await self._get_or_create_nonblocking(key)
if not self._session_has_unarchived_messages(session):
continue
try:
runtime = await resolve_runtime(session)
except (KeyError, ValueError):
continue
self._archiving.add(key)
schedule_background(self._archive_async(key, runtime=runtime))
async def _archive(self, key: str, *, runtime: LLMRuntime) -> None: async def _archive(self, key: str, *, runtime: LLMRuntime) -> None:
if self._is_internal_session(key): if self._is_internal_session(key):
self._archiving.discard(key) self._archiving.discard(key)
@@ -90,17 +135,37 @@ class AutoCompact:
max_suffix=self._RECENT_SUFFIX_MESSAGES, max_suffix=self._RECENT_SUFFIX_MESSAGES,
) )
if summary and summary != "(nothing)": if summary and summary != "(nothing)":
session = self.sessions.get_or_create(key) self._record_stored_summary(key, self.sessions.get_or_create(key))
except Exception:
logger.exception("Auto-compact: failed for {}", key)
finally:
self._archiving.discard(key)
async def _archive_async(self, key: str, *, runtime: LLMRuntime) -> None:
if self._is_internal_session(key):
self._archiving.discard(key)
return
try:
summary = await self.consolidator.compact_idle_session(
key,
runtime=runtime,
max_suffix=self._RECENT_SUFFIX_MESSAGES,
)
if summary and summary != "(nothing)":
session = await self._get_or_create_nonblocking(key)
self._record_stored_summary(key, session)
except Exception:
logger.exception("Auto-compact: failed for {}", key)
finally:
self._archiving.discard(key)
def _record_stored_summary(self, key: str, session: Session) -> None:
stored = session_summary_from_metadata( stored = session_summary_from_metadata(
session.metadata, session.metadata,
fallback_last_active=session.updated_at, fallback_last_active=session.updated_at,
) )
if stored is not None: if stored is not None:
self._summaries[key] = stored self._summaries[key] = stored
except Exception:
logger.exception("Auto-compact: failed for {}", key)
finally:
self._archiving.discard(key)
def prepare_session(self, session: Session, key: str) -> tuple[Session, SessionSummary | None]: def prepare_session(self, session: Session, key: str) -> tuple[Session, SessionSummary | None]:
if self._is_internal_session(key): if self._is_internal_session(key):
@@ -110,6 +175,28 @@ class AutoCompact:
if key in self._archiving or self._is_expired(session.updated_at): if key in self._archiving or self._is_expired(session.updated_at):
logger.info("Auto-compact: reloading session {} (archiving={})", key, key in self._archiving) logger.info("Auto-compact: reloading session {} (archiving={})", key, key in self._archiving)
session = self.sessions.get_or_create(key) session = self.sessions.get_or_create(key)
return self._prepared_summary(session, key)
async def prepare_session_async(
self,
session: Session,
key: str,
) -> tuple[Session, SessionSummary | None]:
"""Prepare a session without blocking on a reload."""
if self._is_internal_session(key):
self._archiving.discard(key)
self._summaries.pop(key, None)
return session, None
if key in self._archiving or self._is_expired(session.updated_at):
logger.info("Auto-compact: reloading session {} (archiving={})", key, key in self._archiving)
session = await self._get_or_create_nonblocking(key)
return self._prepared_summary(session, key)
def _prepared_summary(
self,
session: Session,
key: str,
) -> tuple[Session, SessionSummary | None]:
# Hot path: summary from in-memory dict (process hasn't restarted). # Hot path: summary from in-memory dict (process hasn't restarted).
entry = self._summaries.pop(key, None) entry = self._summaries.pop(key, None)
if entry: if entry:
+32 -2
View File
@@ -13,6 +13,19 @@ class AutomationTurnError(RuntimeError):
"""Raised when an automation turn reaches the agent and finishes with an error.""" """Raised when an automation turn reaches the agent and finishes with an error."""
class AutomationTurnAcceptedCancellation(asyncio.CancelledError):
"""Cancellation raised after an automation turn was accepted for processing.
Callers must not replay the turn: the accepted agent work now has independent
ownership and may continue after the submitting task is cancelled.
"""
def _consume_future_exception(future: asyncio.Future[object]) -> None:
if not future.cancelled():
future.exception()
async def publish_next_deferred_turn( async def publish_next_deferred_turn(
*, *,
deferred_queues: dict[str, list[InboundMessage]], deferred_queues: dict[str, list[InboundMessage]],
@@ -70,19 +83,36 @@ class AutomationTurnCoordinator:
future: asyncio.Future[OutboundMessage | None] = loop.create_future() future: asyncio.Future[OutboundMessage | None] = loop.create_future()
self._waiters[turn_id] = future self._waiters[turn_id] = future
self._pending_messages_by_turn_id[turn_id] = msg self._pending_messages_by_turn_id[turn_id] = msg
accepted = False
try: try:
if self._is_running(): if self._is_running():
await self._publish_inbound(msg) await self._publish_inbound(msg)
accepted = True
else: else:
await self._dispatch(msg) # Direct dispatch is given independent task ownership for the
# same reason as publishing to the inbound queue: once admitted,
# cancelling this submitter must not cancel and then replay the
# already-running agent turn.
dispatch_future: asyncio.Future[object] = asyncio.ensure_future(
self._dispatch(msg)
)
dispatch_future.add_done_callback(_consume_future_exception)
accepted = True
await asyncio.shield(dispatch_future)
try: try:
return await future return await future
except asyncio.CancelledError: except asyncio.CancelledError as exc:
if accepted:
raise AutomationTurnAcceptedCancellation(*exc.args) from None
raise raise
except AutomationTurnError: except AutomationTurnError:
raise raise
except Exception as exc: except Exception as exc:
raise AutomationTurnError(str(exc) or exc.__class__.__name__) from exc raise AutomationTurnError(str(exc) or exc.__class__.__name__) from exc
except asyncio.CancelledError as exc:
if accepted and not isinstance(exc, AutomationTurnAcceptedCancellation):
raise AutomationTurnAcceptedCancellation(*exc.args) from None
raise
finally: finally:
self._waiters.pop(turn_id, None) self._waiters.pop(turn_id, None)
self._pending_messages_by_turn_id.pop(turn_id, None) self._pending_messages_by_turn_id.pop(turn_id, None)
+8
View File
@@ -112,6 +112,14 @@ class ContextBuilder:
parts.append(render_template("agent/tool_contract.md")) parts.append(render_template("agent/tool_contract.md"))
project_path = root.expanduser().resolve()
if project_path != self.workspace.expanduser().resolve():
parts.append(
"# Current Project\n\n"
f"Working directory: {project_path}\n"
"Use it as the default root for project files and relative tool paths."
)
if include_memory: if include_memory:
memory = self.memory.read_memory() memory = self.memory.read_memory()
if memory and not self._is_template_content(memory, "memory/MEMORY.md"): if memory and not self._is_template_content(memory, "memory/MEMORY.md"):
+6 -6
View File
@@ -9,7 +9,7 @@ from typing import Any
from loguru import logger from loguru import logger
from nanobot.providers.base import LLMResponse, ToolCallRequest from nanobot.providers.base import LLMResponse, LLMUsage, ToolCallRequest
@dataclass(slots=True) @dataclass(slots=True)
@@ -19,7 +19,7 @@ class AgentHookContext:
iteration: int iteration: int
messages: list[dict[str, Any]] messages: list[dict[str, Any]]
response: LLMResponse | None = None response: LLMResponse | None = None
usage: dict[str, int] = field(default_factory=dict) usage: LLMUsage | None = None
tool_calls: list[ToolCallRequest] = field(default_factory=list) tool_calls: list[ToolCallRequest] = field(default_factory=list)
tool_results: list[Any] = field(default_factory=list) tool_results: list[Any] = field(default_factory=list)
tool_events: list[dict[str, str]] = field(default_factory=list) tool_events: list[dict[str, str]] = field(default_factory=list)
@@ -39,7 +39,7 @@ class AgentRunHookContext:
messages: list[dict[str, Any]] messages: list[dict[str, Any]]
final_content: str | None = None final_content: str | None = None
tools_used: list[str] = field(default_factory=list) tools_used: list[str] = field(default_factory=list)
usage: dict[str, int] = field(default_factory=dict) usage: LLMUsage | None = None
stop_reason: str | None = None stop_reason: str | None = None
error: str | None = None error: str | None = None
tool_events: list[dict[str, str]] = field(default_factory=list) tool_events: list[dict[str, str]] = field(default_factory=list)
@@ -284,7 +284,7 @@ class SDKCaptureHook(AgentHook):
super().__init__() super().__init__()
self.tools_used: list[str] = [] self.tools_used: list[str] = []
self.messages: list[dict[str, Any]] = [] self.messages: list[dict[str, Any]] = []
self.usage: dict[str, int] = {} self.usage: LLMUsage | None = None
self.stop_reason: str | None = None self.stop_reason: str | None = None
self.error: str | None = None self.error: str | None = None
self.tool_events: list[dict[str, str]] = [] self.tool_events: list[dict[str, str]] = []
@@ -294,7 +294,7 @@ class SDKCaptureHook(AgentHook):
for call in context.tool_calls: for call in context.tool_calls:
self.tools_used.append(call.name) self.tools_used.append(call.name)
self.messages = list(context.messages) self.messages = list(context.messages)
self.usage = dict(context.usage) self.usage = context.usage
self.stop_reason = context.stop_reason self.stop_reason = context.stop_reason
self.error = context.error self.error = context.error
self.tool_events = list(context.tool_events) self.tool_events = list(context.tool_events)
@@ -302,7 +302,7 @@ class SDKCaptureHook(AgentHook):
async def after_run(self, context: AgentRunHookContext) -> None: async def after_run(self, context: AgentRunHookContext) -> None:
self.tools_used = list(context.tools_used) self.tools_used = list(context.tools_used)
self.messages = list(context.messages) self.messages = list(context.messages)
self.usage = dict(context.usage) self.usage = context.usage
self.stop_reason = context.stop_reason self.stop_reason = context.stop_reason
self.error = context.error self.error = context.error
self.tool_events = list(context.tool_events) self.tool_events = list(context.tool_events)
+278 -180
View File
@@ -49,7 +49,8 @@ from nanobot.bus.queue import MessageBus
from nanobot.bus.runtime_events import RuntimeEventBus from nanobot.bus.runtime_events import RuntimeEventBus
from nanobot.command import CommandContext, CommandRouter, register_builtin_commands from nanobot.command import CommandContext, CommandRouter, register_builtin_commands
from nanobot.config.schema import AgentDefaults, ModelPresetConfig from nanobot.config.schema import AgentDefaults, ModelPresetConfig
from nanobot.providers.base import LLMProvider, ProviderConversationState from nanobot.llm_usage.context import source_from_request
from nanobot.providers.base import LLMProvider, LLMUsage, ProviderConversationState
from nanobot.providers.factory import ProviderSnapshot from nanobot.providers.factory import ProviderSnapshot
from nanobot.runtime_context import ( from nanobot.runtime_context import (
RUNTIME_CONTEXT_HISTORY_META, RUNTIME_CONTEXT_HISTORY_META,
@@ -66,6 +67,7 @@ from nanobot.security.workspace_access import (
reset_workspace_scope, reset_workspace_scope,
) )
from nanobot.session import turn_continuation from nanobot.session import turn_continuation
from nanobot.session.async_compat import call_session_manager
from nanobot.session.automation_turns import automation_history_overrides from nanobot.session.automation_turns import automation_history_overrides
from nanobot.session.goal_state import ( from nanobot.session.goal_state import (
goal_state_runtime_lines, goal_state_runtime_lines,
@@ -79,6 +81,15 @@ from nanobot.session.model_selection import (
SESSION_MODEL_PRESET_METADATA_KEY, SESSION_MODEL_PRESET_METADATA_KEY,
model_preset_from_metadata, model_preset_from_metadata,
) )
from nanobot.session.recovery import (
PENDING_FOLLOWUP_ID_KEY,
RECOVERY_INBOUND_METADATA_KEY,
RecoveryAdmission,
acknowledge_pending_followups,
record_pending_followup,
restore_pending_interruption,
restore_runtime_checkpoint,
)
from nanobot.session.summary import SessionSummary from nanobot.session.summary import SessionSummary
from nanobot.triggers.local_turns import LocalTriggerTurnCoordinator from nanobot.triggers.local_turns import LocalTriggerTurnCoordinator
from nanobot.utils.cancellation import task_is_cancelling from nanobot.utils.cancellation import task_is_cancelling
@@ -158,7 +169,7 @@ class TurnContext:
turn_wall_started_at: float = field(default_factory=time.time) turn_wall_started_at: float = field(default_factory=time.time)
visible_run_started_at: float | None = None visible_run_started_at: float | None = None
turn_latency_ms: int | None = None turn_latency_ms: int | None = None
usage: dict[str, int] = field(default_factory=dict) usage: LLMUsage | None = None
def require_runtime(self) -> LLMRuntime: def require_runtime(self) -> LLMRuntime:
"""Return the runtime established by the BUILD stage.""" """Return the runtime established by the BUILD stage."""
@@ -194,7 +205,7 @@ class AgentLoop:
return self.tools.tool_names return self.tools.tool_names
@property @property
def last_usage(self) -> Mapping[str, int]: def last_usage(self) -> LLMUsage | None:
"""Latest aggregate usage exposed through the runtime-control snapshot.""" """Latest aggregate usage exposed through the runtime-control snapshot."""
return self._last_usage return self._last_usage
@@ -291,12 +302,14 @@ class AgentLoop:
restart_mode: str = "auto", restart_mode: str = "auto",
local_trigger_store: LocalTriggerStore | None = None, local_trigger_store: LocalTriggerStore | None = None,
idle_compact_check_interval_seconds: int = 0, idle_compact_check_interval_seconds: int = 0,
recovery_admission: RecoveryAdmission | None = None,
): ):
from nanobot.config.schema import ToolsConfig from nanobot.config.schema import ToolsConfig
_tc = tools_config or ToolsConfig() _tc = tools_config or ToolsConfig()
defaults = AgentDefaults() defaults = AgentDefaults()
self.bus = bus self.bus = bus
self._recovery_admission = recovery_admission
if turn_delivery_factory is not None: if turn_delivery_factory is not None:
if turn_delivery_factory.bus is not bus: if turn_delivery_factory.bus is not bus:
raise ValueError("turn delivery factory must use the agent message bus") raise ValueError("turn delivery factory must use the agent message bus")
@@ -367,7 +380,7 @@ class AgentLoop:
default_restrict_to_workspace=restrict_to_workspace, default_restrict_to_workspace=restrict_to_workspace,
) )
self._start_time = time.time() self._start_time = time.time()
self._last_usage: dict[str, int] = {} self._last_usage: LLMUsage | None = None
self._extra_hooks: list[AgentHook] = hooks or [] self._extra_hooks: list[AgentHook] = hooks or []
self._hook_factories: list[AgentTurnHookFactory] = hook_factories or [] self._hook_factories: list[AgentTurnHookFactory] = hook_factories or []
@@ -409,6 +422,7 @@ class AgentLoop:
# When a session has an active task, new messages for that session # When a session has an active task, new messages for that session
# are routed here instead of creating a new task. # are routed here instead of creating a new task.
self._pending_queues: dict[str, asyncio.Queue[InboundMessage]] = {} self._pending_queues: dict[str, asyncio.Queue[InboundMessage]] = {}
self._preserve_inflight_turns_on_shutdown = False
self._deferred_automation_turns: dict[str, list[InboundMessage]] = {} self._deferred_automation_turns: dict[str, list[InboundMessage]] = {}
self._cron_turns = CronTurnCoordinator( self._cron_turns = CronTurnCoordinator(
publish_inbound=self.bus.publish_inbound, publish_inbound=self.bus.publish_inbound,
@@ -527,6 +541,33 @@ class AgentLoop:
**extra, **extra,
) )
async def _get_or_create_session(self, key: str) -> Session:
"""Use native async session loading, with a compatibility fallback."""
return await call_session_manager(
self.sessions,
"get_or_create_async",
self.sessions.get_or_create,
key,
)
async def _save_session(self, session: Session) -> None:
"""Use native async session saving, with a compatibility fallback."""
await call_session_manager(
self.sessions,
"save_async",
self.sessions.save,
session,
)
async def _save_runtime_checkpoint(self, session: Session) -> None:
"""Use native async checkpoint saving, with a compatibility fallback."""
await call_session_manager(
self.sessions,
"save_runtime_checkpoint_async",
self.sessions.save_runtime_checkpoint,
session,
)
def _sync_subagent_runtime_limits(self) -> None: def _sync_subagent_runtime_limits(self) -> None:
"""Keep subagent runtime limits aligned with mutable loop settings.""" """Keep subagent runtime limits aligned with mutable loop settings."""
self.subagents.max_iterations = self.max_iterations self.subagents.max_iterations = self.max_iterations
@@ -566,6 +607,30 @@ class AgentLoop:
self.sessions.save(session) self.sessions.save(session)
return self.llm_runtime() return self.llm_runtime()
async def runtime_for_session_async(
self,
session: Session,
*,
recover_removed: bool = True,
) -> LLMRuntime:
"""Resolve a session runtime without blocking on recovery persistence."""
name = model_preset_from_metadata(session.metadata)
if name is None:
return self.llm_runtime()
try:
return self.runtime_resolver.resolve_preset(name)
except KeyError:
if not recover_removed or name in self.runtime_resolver.model_presets:
raise
logger.warning(
"Session '{}' references removed model preset '{}'; falling back to default",
session.key,
name,
)
session.metadata.pop(SESSION_MODEL_PRESET_METADATA_KEY, None)
await self._save_session(session)
return self.llm_runtime()
def set_session_model_preset( def set_session_model_preset(
self, self,
session_key: str, session_key: str,
@@ -578,6 +643,18 @@ class AgentLoop:
self.sessions.save(session) self.sessions.save(session)
return runtime return runtime
async def set_session_model_preset_async(
self,
session_key: str,
name: str,
) -> LLMRuntime:
"""Validate and persist one session's preset selection without blocking."""
runtime = self.runtime_resolver.resolve_preset(name)
session = await self._get_or_create_session(session_key)
session.metadata[SESSION_MODEL_PRESET_METADATA_KEY] = runtime.model_preset
await self._save_session(session)
return runtime
def _publish_runtime_selection( def _publish_runtime_selection(
self, self,
runtime: LLMRuntime, runtime: LLMRuntime,
@@ -690,17 +767,14 @@ class AgentLoop:
session_key=session_key, session_key=session_key,
) )
def _persist_user_message_early( def _stage_user_message_early(
self, self,
msg: InboundMessage, msg: InboundMessage,
session: Session, session: Session,
runtime_context_blocks: list[RuntimeContextBlock] | None = None, runtime_context_blocks: list[RuntimeContextBlock] | None = None,
**kwargs: Any, **kwargs: Any,
) -> bool: ) -> bool:
"""Persist the triggering user message before the turn starts. """Add the triggering user message and recovery markers in memory."""
Returns True if the message was persisted.
"""
if not turn_continuation.should_persist_user_message(msg.metadata): if not turn_continuation.should_persist_user_message(msg.metadata):
return False return False
media_paths = [ media_paths = [
@@ -726,10 +800,48 @@ class AgentLoop:
extra[RUNTIME_CONTEXT_HISTORY_META] = runtime_context_meta extra[RUNTIME_CONTEXT_HISTORY_META] = runtime_context_meta
session.add_message("user", text, **extra) session.add_message("user", text, **extra)
self._mark_pending_user_turn(session) self._mark_pending_user_turn(session)
self.sessions.save(session) followup_id = msg.metadata.get(PENDING_FOLLOWUP_ID_KEY)
if isinstance(followup_id, str) and followup_id:
acknowledge_pending_followups(session, [followup_id])
return True return True
return False return False
def _persist_user_message_early(
self,
msg: InboundMessage,
session: Session,
runtime_context_blocks: list[RuntimeContextBlock] | None = None,
**kwargs: Any,
) -> bool:
"""Synchronously persist the user message for compatibility callers."""
persisted = self._stage_user_message_early(
msg,
session,
runtime_context_blocks,
**kwargs,
)
if persisted:
self.sessions.save(session)
return persisted
async def _persist_user_message_early_async(
self,
msg: InboundMessage,
session: Session,
runtime_context_blocks: list[RuntimeContextBlock] | None = None,
**kwargs: Any,
) -> bool:
"""Persist the user message without blocking the event loop."""
persisted = self._stage_user_message_early(
msg,
session,
runtime_context_blocks,
**kwargs,
)
if persisted:
await self._save_session(session)
return persisted
def _build_initial_messages(self, ctx: TurnContext) -> list[dict[str, Any]]: def _build_initial_messages(self, ctx: TurnContext) -> list[dict[str, Any]]:
"""Build the initial message list for the LLM turn.""" """Build the initial message list for the LLM turn."""
assert ctx.session is not None assert ctx.session is not None
@@ -827,7 +939,7 @@ class AgentLoop:
if tool is None: if tool is None:
content = "Shell execution is disabled in this nanobot configuration." content = "Shell execution is disabled in this nanobot configuration."
else: else:
session = ctx.session or self.sessions.get_or_create(ctx.key) session = ctx.session or await AgentLoop._get_or_create_session(self, ctx.key)
scope = self.workspace_scopes.for_turn( scope = self.workspace_scopes.for_turn(
channel=ctx.msg.channel, channel=ctx.msg.channel,
message_metadata=metadata, message_metadata=metadata,
@@ -985,7 +1097,7 @@ class AgentLoop:
public_payload[self._PROVIDER_STATE_CHECKPOINT_VERSION_KEY] = ( public_payload[self._PROVIDER_STATE_CHECKPOINT_VERSION_KEY] = (
self._PROVIDER_STATE_CHECKPOINT_VERSION self._PROVIDER_STATE_CHECKPOINT_VERSION
) )
self._set_runtime_checkpoint(session, public_payload) await self._set_runtime_checkpoint_async(session, public_payload)
async def _drain_pending(*, limit: int = _MAX_INJECTIONS_PER_TURN) -> list[dict[str, Any]]: async def _drain_pending(*, limit: int = _MAX_INJECTIONS_PER_TURN) -> list[dict[str, Any]]:
"""Drain follow-up messages from the pending queue. """Drain follow-up messages from the pending queue.
@@ -1061,6 +1173,9 @@ class AgentLoop:
row["subagent_task_id"] = task_id row["subagent_task_id"] = task_id
row[HIDDEN_HISTORY_META] = subagent_marker row[HIDDEN_HISTORY_META] = subagent_marker
row["injected_event"] = "subagent_result" row["injected_event"] = "subagent_result"
followup_id = metadata.get(PENDING_FOLLOWUP_ID_KEY)
if isinstance(followup_id, str) and followup_id:
row[PENDING_FOLLOWUP_ID_KEY] = followup_id
return row return row
items: list[dict[str, Any]] = [] items: list[dict[str, Any]] = []
@@ -1184,6 +1299,11 @@ class AgentLoop:
message_metadata=metadata, message_metadata=metadata,
), ),
provider_state=provider_state, provider_state=provider_state,
llm_usage_source=source_from_request(
active_session_key,
channel=channel,
metadata=metadata,
),
)) ))
finally: finally:
turn_scope_stack.close() turn_scope_stack.close()
@@ -1215,18 +1335,33 @@ class AgentLoop:
logger.error("LLM returned error: {}", (result.final_content or "")[:200]) 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 return result.final_content, result.tools_used, result.messages, result.stop_reason, result.had_injections
def _check_expired_sessions_if_due(self) -> None: def _idle_compact_scan_due(self) -> bool:
"""Scan idle sessions no more often than the configured interval."""
now = time.monotonic() now = time.monotonic()
if now < self._next_idle_compact_check_at: if now < self._next_idle_compact_check_at:
return return False
self._next_idle_compact_check_at = now + self._idle_compact_check_interval_s self._next_idle_compact_check_at = now + self._idle_compact_check_interval_s
return True
def _check_expired_sessions_if_due(self) -> None:
"""Synchronously scan idle sessions for compatibility with direct callers."""
if not self._idle_compact_scan_due():
return
self.auto_compact.check_expired( self.auto_compact.check_expired(
self.schedule_background, self.schedule_background,
self.runtime_for_session, self.runtime_for_session,
active_session_keys=self._pending_queues.keys(), active_session_keys=self._pending_queues.keys(),
) )
async def _check_expired_sessions_if_due_async(self) -> None:
"""Scan idle sessions without blocking the event loop."""
if not self._idle_compact_scan_due():
return
await self.auto_compact.check_expired_async(
self.schedule_background,
self.runtime_for_session_async,
active_session_keys=self._pending_queues.keys(),
)
async def run(self) -> None: async def run(self) -> None:
"""Run the agent loop, dispatching messages as tasks to stay responsive to /stop.""" """Run the agent loop, dispatching messages as tasks to stay responsive to /stop."""
self._running = True self._running = True
@@ -1237,7 +1372,7 @@ class AgentLoop:
try: try:
msg = await asyncio.wait_for(self.bus.consume_inbound(), timeout=1.0) msg = await asyncio.wait_for(self.bus.consume_inbound(), timeout=1.0)
except asyncio.TimeoutError: except asyncio.TimeoutError:
self._check_expired_sessions_if_due() await self._check_expired_sessions_if_due_async()
continue continue
except asyncio.CancelledError: except asyncio.CancelledError:
# Preserve real task cancellation so shutdown can complete cleanly. # Preserve real task cancellation so shutdown can complete cleanly.
@@ -1285,6 +1420,23 @@ class AgentLoop:
break break
if deferred: if deferred:
continue continue
routed_msg = msg
if effective_key != msg.session_key:
routed_msg = dataclasses.replace(
msg,
session_key_override=effective_key,
)
# A newer WebUI message must supersede an explicit recovery
# before it is injected into that recovery's pending queue.
# Without this admission point, a recovered turn could finish
# first and only then observe the user's newer request.
if (
effective_key in self._pending_queues
and msg.channel == "websocket"
and self._recovery_admission is not None
and not await self._recovery_admission.admit(routed_msg)
):
continue
# If this session already has an active pending queue (i.e. a task # If this session already has an active pending queue (i.e. a task
# is processing this session), route the message there for mid-turn # is processing this session), route the message there for mid-turn
# injection instead of creating a competing task. # injection instead of creating a competing task.
@@ -1297,12 +1449,18 @@ class AgentLoop:
self.commands.dispatch, self.commands.dispatch,
) )
continue continue
pending_msg = msg pending_msg = routed_msg
if effective_key != msg.session_key: session = await self._get_or_create_session(effective_key)
followup_id = record_pending_followup(session, pending_msg)
if followup_id is not None:
pending_msg = dataclasses.replace( pending_msg = dataclasses.replace(
msg, pending_msg,
session_key_override=effective_key, metadata={
**pending_msg.metadata,
PENDING_FOLLOWUP_ID_KEY: followup_id,
},
) )
await self._save_session(session)
try: try:
self._pending_queues[effective_key].put_nowait(pending_msg) self._pending_queues[effective_key].put_nowait(pending_msg)
except asyncio.QueueFull: except asyncio.QueueFull:
@@ -1310,6 +1468,7 @@ class AgentLoop:
"Pending queue full for session {}, falling back to queued task", "Pending queue full for session {}, falling back to queued task",
effective_key, effective_key,
) )
msg = pending_msg
else: else:
logger.info( logger.info(
"Routed follow-up message to pending queue for session {}", "Routed follow-up message to pending queue for session {}",
@@ -1319,17 +1478,45 @@ class AgentLoop:
# Compute the effective session key before dispatching # Compute the effective session key before dispatching
# This ensures /stop command can find tasks correctly when unified session is enabled # This ensures /stop command can find tasks correctly when unified session is enabled
task = asyncio.create_task(self._dispatch(msg)) task = asyncio.create_task(self._dispatch(msg))
active_tasks = self._active_tasks.setdefault(effective_key, set()) active_tasks: set[asyncio.Task[Any]] = self._active_tasks.setdefault(
effective_key,
set(),
)
active_tasks.add(task) active_tasks.add(task)
task.add_done_callback(active_tasks.discard) task.add_done_callback(active_tasks.discard)
finally: finally:
await self.aclose() await self.aclose()
def preserve_inflight_turns_on_shutdown(self) -> None:
"""Keep durable checkpoints when the owning gateway exits.
Normal cancellation intentionally materializes partial output so a
user-stopped turn leaves a readable conversation. Gateway lifecycle
shutdown is different: RecoveryCoordinator needs the checkpoint intact
to safely offer the unfinished turn for explicit continuation later.
"""
self._preserve_inflight_turns_on_shutdown = True
async def _dispatch(self, msg: InboundMessage) -> None: async def _dispatch(self, msg: InboundMessage) -> None:
"""Process a message: per-session serial, cross-session concurrent.""" """Process a message: per-session serial, cross-session concurrent."""
session_key = self._effective_session_key(msg) session_key = self._effective_session_key(msg)
if session_key != msg.session_key: if session_key != msg.session_key:
msg = dataclasses.replace(msg, session_key_override=session_key) msg = dataclasses.replace(msg, session_key_override=session_key)
recovery_task_registered = False
recovery_admission = self._recovery_admission
current_task: asyncio.Task[Any] | None = None
if recovery_admission is not None:
recovery_id = msg.metadata.get(RECOVERY_INBOUND_METADATA_KEY)
if isinstance(recovery_id, str) and recovery_id:
current_task = asyncio.current_task()
if current_task is not None:
recovery_admission.register_recovery_task(session_key, current_task)
recovery_task_registered = True
if not await recovery_admission.admit(msg):
logger.info("Skipped stale recovery for session {}", session_key)
if recovery_task_registered and current_task is not None:
recovery_admission.unregister_recovery_task(session_key, current_task)
return
lock = self._get_session_lock(session_key) lock = self._get_session_lock(session_key)
gate = self._concurrency_gate or nullcontext() gate = self._concurrency_gate or nullcontext()
@@ -1373,21 +1560,21 @@ class AgentLoop:
session_key, session_key,
exc_info=True, exc_info=True,
) )
# Preserve partial context from the interrupted turn so # An explicit turn stop materializes partial context so
# the user does not lose tool results and assistant # the next prompt can see completed tool results. Gateway
# messages accumulated before /stop. The checkpoint was # shutdown keeps the durable checkpoint untouched instead,
# already persisted to session metadata by # allowing RecoveryCoordinator to offer Continue safely.
# _emit_checkpoint during tool execution; materializing if (
# it into session history now makes it visible in the session_key in self._discarding_sessions
# next conversation turn. or self._preserve_inflight_turns_on_shutdown
if session_key in self._discarding_sessions: ):
raise raise
try: try:
key = self._effective_session_key(msg) key = self._effective_session_key(msg)
session = self.sessions.get_or_create(key) session = await self._get_or_create_session(key)
if self._restore_runtime_checkpoint(session): if self._restore_runtime_checkpoint(session):
self._clear_pending_user_turn(session) self._clear_pending_user_turn(session)
self.sessions.save(session) await self._save_session(session)
logger.info( logger.info(
"Restored partial context for cancelled session {}", "Restored partial context for cancelled session {}",
key, key,
@@ -1437,6 +1624,12 @@ class AgentLoop:
await delivery.idle() await delivery.idle()
await self._publish_next_deferred_automation_turn(session_key) await self._publish_next_deferred_automation_turn(session_key)
finally: finally:
if (
recovery_task_registered
and current_task is not None
and recovery_admission is not None
):
recovery_admission.unregister_recovery_task(session_key, current_task)
if pending is None: if pending is None:
await delivery.idle() await delivery.idle()
await self._publish_next_deferred_automation_turn(session_key) await self._publish_next_deferred_automation_turn(session_key)
@@ -1706,7 +1899,7 @@ class AgentLoop:
if ctx.session is None: if ctx.session is None:
raise RuntimeError("required session is not active") raise RuntimeError("required session is not active")
else: else:
ctx.session = self.sessions.get_or_create(ctx.session_key) ctx.session = await self._get_or_create_session(ctx.session_key)
session = ctx.session session = ctx.session
ctx.ephemeral = ctx.ephemeral or not session.policy.persist ctx.ephemeral = ctx.ephemeral or not session.policy.persist
tools = ctx.tools or self.tools tools = ctx.tools or self.tools
@@ -1737,13 +1930,16 @@ class AgentLoop:
self.workspace_scopes.persist_message_scope(session, msg) self.workspace_scopes.persist_message_scope(session, msg)
if self._restore_runtime_checkpoint(session): if self._restore_runtime_checkpoint(session):
self.sessions.save(session) await self._save_session(session)
if self._restore_pending_user_turn(session): if (
self.sessions.save(session) RECOVERY_INBOUND_METADATA_KEY not in msg.metadata
and restore_pending_interruption(session)
):
await self._save_session(session)
async def _compact_session(self, ctx: TurnContext) -> None: async def _compact_session(self, ctx: TurnContext) -> None:
session = ctx.require_session() session = ctx.require_session()
ctx.session, pending = self.auto_compact.prepare_session( ctx.session, pending = await self.auto_compact.prepare_session_async(
session, session,
ctx.session_key, ctx.session_key,
) )
@@ -1780,14 +1976,14 @@ class AgentLoop:
# them out of LLM context. /new is excluded because it # them out of LLM context. /new is excluded because it
# intentionally clears the session. # intentionally clears the session.
if cmd_ctx.raw.lower() != "/new": if cmd_ctx.raw.lower() != "/new":
ctx.input_persisted_early = self._persist_user_message_early( ctx.input_persisted_early = await self._persist_user_message_early_async(
ctx.msg, session, _command=True ctx.msg, session, _command=True
) )
session.add_message( session.add_message(
"assistant", result.content, _command=True "assistant", result.content, _command=True
) )
self._clear_pending_user_turn(session) self._clear_pending_user_turn(session)
self.sessions.save(session) await self._save_session(session)
if not ctx.ephemeral: if not ctx.ephemeral:
await self.runtime_event_publisher.session_turn_persisted( await self.runtime_event_publisher.session_turn_persisted(
ctx.msg, ctx.msg,
@@ -1802,7 +1998,7 @@ class AgentLoop:
session = ctx.require_session() session = ctx.require_session()
runtime = ctx.runtime runtime = ctx.runtime
if runtime is None: if runtime is None:
runtime = self.runtime_for_session(session) runtime = await self.runtime_for_session_async(session)
ctx.runtime = runtime ctx.runtime = runtime
if ctx.session_key.startswith("dream:"): if ctx.session_key.startswith("dream:"):
logger.info( logger.info(
@@ -1846,7 +2042,7 @@ class AgentLoop:
# provider compatibility or prompt assembly work. A compatible # provider compatibility or prompt assembly work. A compatible
# staged state replaces this in a second atomic save below. # staged state replaces this in a second atomic save below.
session.provider_state = None session.provider_state = None
self.sessions.save(session) await self._save_session(session)
ctx.input_persisted_early = True ctx.input_persisted_early = True
await ctx.delivery.runtime_admitted(runtime) await ctx.delivery.runtime_admitted(runtime)
@@ -1900,7 +2096,7 @@ class AgentLoop:
elif stored_state is not None: elif stored_state is not None:
session.provider_state = None session.provider_state = None
if ctx.kind is TurnKind.USER: if ctx.kind is TurnKind.USER:
ctx.input_persisted_early = self._persist_user_message_early( ctx.input_persisted_early = await self._persist_user_message_early_async(
ctx.msg, ctx.msg,
session, session,
runtime_context_blocks=ctx.runtime_context_blocks, runtime_context_blocks=ctx.runtime_context_blocks,
@@ -1910,7 +2106,7 @@ class AgentLoop:
elif subagent_followup_persisted and staged_provider_state: elif subagent_followup_persisted and staged_provider_state:
# Upgrade the replay-safe baseline to the resumable state before # Upgrade the replay-safe baseline to the resumable state before
# prompt assembly and the first model checkpoint. # prompt assembly and the first model checkpoint.
self.sessions.save(session) await self._save_session(session)
ctx.initial_messages = self._build_initial_messages(ctx) ctx.initial_messages = self._build_initial_messages(ctx)
if ctx.on_progress is None: if ctx.on_progress is None:
@@ -1952,7 +2148,7 @@ class AgentLoop:
ctx.all_messages = all_msgs ctx.all_messages = all_msgs
ctx.stop_reason = stop_reason ctx.stop_reason = stop_reason
ctx.had_injections = had_injections ctx.had_injections = had_injections
ctx.usage = dict(self._last_usage) ctx.usage = self._last_usage
ctx.delivery.record_usage(ctx.usage) ctx.delivery.record_usage(ctx.usage)
if ctx.kind is TurnKind.USER: if ctx.kind is TurnKind.USER:
await turn_continuation.maybe_continue_turn(ctx) await turn_continuation.maybe_continue_turn(ctx)
@@ -1979,13 +2175,16 @@ class AgentLoop:
else ctx.turn_wall_started_at else ctx.turn_wall_started_at
) )
ctx.turn_latency_ms = max(0, int((time.time() - latency_started_at) * 1000)) ctx.turn_latency_ms = max(0, int((time.time() - latency_started_at) * 1000))
if ctx.usage and not ctx.ephemeral: if ctx.usage is not None and not ctx.ephemeral:
session.metadata["_last_usage"] = dict(ctx.usage) session.metadata["_last_usage"] = ctx.usage.to_dict()
self._save_turn( self._save_turn(
session, ctx.all_messages, ctx.save_skip, session, ctx.all_messages, ctx.save_skip,
turn_latency_ms=ctx.turn_latency_ms, turn_latency_ms=ctx.turn_latency_ms,
) )
ctx.delivery.record_latency(ctx.turn_latency_ms) ctx.delivery.record_latency(ctx.turn_latency_ms)
self._clear_pending_user_turn(session)
self._clear_runtime_checkpoint(session)
await self._save_session(session)
if not ctx.ephemeral: if not ctx.ephemeral:
self.schedule_background( self.schedule_background(
self.consolidator.maybe_consolidate_by_tokens( self.consolidator.maybe_consolidate_by_tokens(
@@ -1993,10 +2192,6 @@ class AgentLoop:
runtime=runtime, runtime=runtime,
) )
) )
self._clear_pending_user_turn(session)
self._clear_runtime_checkpoint(session)
self.sessions.save(session)
if not ctx.ephemeral:
await self.runtime_event_publisher.session_turn_persisted( await self.runtime_event_publisher.session_turn_persisted(
ctx.msg, ctx.msg,
ctx.session_key, ctx.session_key,
@@ -2093,8 +2288,21 @@ class AgentLoop:
if m.get("role") == "tool" and m.get("tool_call_id") if m.get("role") == "tool" and m.get("tool_call_id")
} }
last_assistant_idx: int | None = None last_assistant_idx: int | None = None
saved_followup_ids: set[str] = set()
for m in messages[skip:]: for m in messages[skip:]:
entry = dict(m) entry = dict(m)
followup_id_value = cast(object, entry.pop(PENDING_FOLLOWUP_ID_KEY, None))
followup_ids = (
[followup_id_value]
if isinstance(followup_id_value, str)
else [
followup_id
for followup_id in cast(list[object], followup_id_value)
if isinstance(followup_id, str)
]
if isinstance(followup_id_value, list)
else []
)
internal_meta = cast(object, entry.pop("_meta", None)) internal_meta = cast(object, entry.pop("_meta", None))
runtime_context_meta = ( runtime_context_meta = (
cast(dict[str, Any], internal_meta).get( cast(dict[str, Any], internal_meta).get(
@@ -2147,6 +2355,8 @@ class AgentLoop:
entry[RUNTIME_CONTEXT_HISTORY_META] = runtime_context_meta entry[RUNTIME_CONTEXT_HISTORY_META] = runtime_context_meta
entry.setdefault("timestamp", datetime.now().isoformat()) entry.setdefault("timestamp", datetime.now().isoformat())
session.messages.append(entry) session.messages.append(entry)
if role == "user":
saved_followup_ids.update(followup_id for followup_id in followup_ids if followup_id)
if role == "assistant": if role == "assistant":
last_assistant_idx = len(session.messages) - 1 last_assistant_idx = len(session.messages) - 1
declared_tool_call_ids.update( declared_tool_call_ids.update(
@@ -2161,6 +2371,8 @@ class AgentLoop:
) )
if turn_latency_ms is not None and last_assistant_idx is not None: 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.messages[last_assistant_idx]["latency_ms"] = int(turn_latency_ms)
if saved_followup_ids:
acknowledge_pending_followups(session, saved_followup_ids)
session.updated_at = datetime.now() session.updated_at = datetime.now()
def _persist_subagent_followup(self, session: Session, msg: InboundMessage) -> bool: def _persist_subagent_followup(self, session: Session, msg: InboundMessage) -> bool:
@@ -2192,10 +2404,23 @@ class AgentLoop:
) )
return True return True
def _set_runtime_checkpoint(self, session: Session, payload: dict[str, Any]) -> None: def _set_runtime_checkpoint(
"""Persist the latest in-flight turn state into session metadata.""" self,
session: Session,
payload: dict[str, Any],
) -> None:
"""Synchronously persist a checkpoint for compatibility callers."""
session.metadata[self._RUNTIME_CHECKPOINT_KEY] = payload session.metadata[self._RUNTIME_CHECKPOINT_KEY] = payload
self.sessions.save(session) self.sessions.save_runtime_checkpoint(session)
async def _set_runtime_checkpoint_async(
self,
session: Session,
payload: dict[str, Any],
) -> None:
"""Persist the latest in-flight turn state without blocking the event loop."""
session.metadata[self._RUNTIME_CHECKPOINT_KEY] = payload
await self._save_runtime_checkpoint(session)
def _mark_pending_user_turn(self, session: Session) -> None: def _mark_pending_user_turn(self, session: Session) -> None:
session.metadata[self._PENDING_USER_TURN_KEY] = True session.metadata[self._PENDING_USER_TURN_KEY] = True
@@ -2207,136 +2432,9 @@ class AgentLoop:
if self._RUNTIME_CHECKPOINT_KEY in session.metadata: if self._RUNTIME_CHECKPOINT_KEY in session.metadata:
session.metadata.pop(self._RUNTIME_CHECKPOINT_KEY, None) 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: def _restore_runtime_checkpoint(self, session: Session) -> bool:
"""Materialize an unfinished turn into session history before a new request.""" """Materialize an unfinished turn into session history before a new request."""
from datetime import datetime return restore_runtime_checkpoint(session)
checkpoint = cast(
object,
session.metadata.get(self._RUNTIME_CHECKPOINT_KEY),
)
if not isinstance(checkpoint, dict):
return False
checkpoint_data = cast(dict[str, Any], checkpoint)
assistant_message = cast(object, checkpoint_data.get("assistant_message"))
completed_tool_results = cast(
Iterable[object],
checkpoint_data.get("completed_tool_results") or [],
)
pending_tool_calls = cast(
Iterable[object],
checkpoint_data.get("pending_tool_calls") or [],
)
restored_messages: list[dict[str, Any]] = []
if isinstance(assistant_message, dict):
restored = dict(cast(dict[str, Any], assistant_message))
restored.setdefault("timestamp", datetime.now().isoformat())
restored_messages.append(restored)
for message in completed_tool_results:
if isinstance(message, dict):
restored = dict(cast(dict[str, Any], 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_call_data = cast(dict[str, Any], tool_call)
tool_id = tool_call_data.get("id")
function_data = cast(
dict[str, Any],
tool_call_data.get("function") or {},
)
name = function_data.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
appended_messages = restored_messages[overlap:]
session.messages.extend(appended_messages)
assistant_message_data = (
cast(dict[str, Any], assistant_message)
if isinstance(assistant_message, dict)
else None
)
provider_state_is_synchronized = (
checkpoint_data.get(self._PROVIDER_STATE_CHECKPOINT_VERSION_KEY)
== self._PROVIDER_STATE_CHECKPOINT_VERSION
)
phase = checkpoint_data.get("phase")
exact_final_response = (
phase == "final_response"
and assistant_message_data is not None
and assistant_message_data.get("role") == "assistant"
and not bool(checkpoint_data.get("completed_tool_results"))
and not bool(checkpoint_data.get("pending_tool_calls"))
)
exact_completed_tools = (
phase == "tools_completed"
and assistant_message_data is not None
and assistant_message_data.get("role") == "assistant"
and not bool(checkpoint_data.get("pending_tool_calls"))
)
if not (
provider_state_is_synchronized
and (exact_final_response or exact_completed_tools)
):
session.provider_state = None
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.provider_state = None
session.updated_at = datetime.now()
self._clear_pending_user_turn(session)
return True
async def process_direct( async def process_direct(
self, self,
+28 -9
View File
@@ -20,7 +20,9 @@ from typing import TYPE_CHECKING, Any, Callable, Iterator, cast
from loguru import logger from loguru import logger
from nanobot.llm_usage.context import llm_usage_source
from nanobot.runtime_context import public_history_messages from nanobot.runtime_context import public_history_messages
from nanobot.session.async_compat import call_session_manager
from nanobot.session.manager import ( from nanobot.session.manager import (
MIN_COMPACTED_REPLAY_MESSAGES, MIN_COMPACTED_REPLAY_MESSAGES,
Session, Session,
@@ -823,6 +825,22 @@ class Consolidator:
weakref.WeakValueDictionary() weakref.WeakValueDictionary()
) )
async def _get_or_create_session(self, key: str) -> Session:
return await call_session_manager(
self.sessions,
"get_or_create_async",
self.sessions.get_or_create,
key,
)
async def _save_session(self, session: Session) -> None:
await call_session_manager(
self.sessions,
"save_async",
self.sessions.save,
session,
)
def get_lock(self, session_key: str) -> asyncio.Lock: def get_lock(self, session_key: str) -> asyncio.Lock:
"""Return the shared consolidation lock for one session.""" """Return the shared consolidation lock for one session."""
return self._locks.setdefault(session_key, asyncio.Lock()) return self._locks.setdefault(session_key, asyncio.Lock())
@@ -858,13 +876,13 @@ class Consolidator:
return [] return []
return session.get_history() return session.get_history()
def _persist_last_summary(self, session: Session, summary: str | None) -> None: async def _persist_last_summary(self, session: Session, summary: str | None) -> None:
if summary and summary != "(nothing)": if summary and summary != "(nothing)":
session.metadata["_last_summary"] = { session.metadata["_last_summary"] = {
"text": summary, "text": summary,
"last_active": session.updated_at.isoformat(), "last_active": session.updated_at.isoformat(),
} }
self.sessions.save(session) await self._save_session(session)
def estimate_session_prompt_tokens( def estimate_session_prompt_tokens(
self, self,
@@ -915,6 +933,7 @@ class Consolidator:
if not messages: if not messages:
return None return None
try: try:
with llm_usage_source("dream"):
response = await runtime.provider.chat_with_retry( response = await runtime.provider.chat_with_retry(
model=runtime.model, model=runtime.model,
messages=request_messages, messages=request_messages,
@@ -1055,7 +1074,7 @@ class Consolidator:
lock = self.get_lock(session.key) lock = self.get_lock(session.key)
async with lock: async with lock:
# Refresh session reference: AutoCompact may have replaced it. # Refresh session reference: AutoCompact may have replaced it.
fresh = self.sessions.get_or_create(session.key) fresh = await self._get_or_create_session(session.key)
if fresh is not session: if fresh is not session:
session = fresh session = fresh
if not session.messages: if not session.messages:
@@ -1069,7 +1088,7 @@ class Consolidator:
runtime=runtime, runtime=runtime,
) )
if estimated <= 0: if estimated <= 0:
self._persist_last_summary(session, last_summary) await self._persist_last_summary(session, last_summary)
return return
if estimated < budget: if estimated < budget:
unconsolidated_count = len(session.messages) - session.last_consolidated unconsolidated_count = len(session.messages) - session.last_consolidated
@@ -1081,7 +1100,7 @@ class Consolidator:
source, source,
unconsolidated_count, unconsolidated_count,
) )
self._persist_last_summary(session, last_summary) await self._persist_last_summary(session, last_summary)
return return
for round_num in range(self._MAX_CONSOLIDATION_ROUNDS): for round_num in range(self._MAX_CONSOLIDATION_ROUNDS):
@@ -1125,7 +1144,7 @@ class Consolidator:
last_summary = summary last_summary = summary
session.last_consolidated = end_idx session.last_consolidated = end_idx
session.provider_state = None session.provider_state = None
self.sessions.save(session) await self._save_session(session)
if not summary: if not summary:
# LLM is degraded — stop hammering it this call; # LLM is degraded — stop hammering it this call;
# the next invocation can retry a fresh chunk. # the next invocation can retry a fresh chunk.
@@ -1141,7 +1160,7 @@ class Consolidator:
# Persist the last summary to session metadata so it can be injected # Persist the last summary to session metadata so it can be injected
# into the runtime context on the next prepare_session() call, aligning # into the runtime context on the next prepare_session() call, aligning
# the summary injection strategy with AutoCompact._archive(). # the summary injection strategy with AutoCompact._archive().
self._persist_last_summary(session, last_summary) await self._persist_last_summary(session, last_summary)
async def compact_idle_session( async def compact_idle_session(
self, self,
@@ -1166,7 +1185,7 @@ class Consolidator:
lock = self.get_lock(session_key) lock = self.get_lock(session_key)
async with lock: async with lock:
self.sessions.invalidate(session_key) self.sessions.invalidate(session_key)
session = self.sessions.get_or_create(session_key) session = await self._get_or_create_session(session_key)
archive_start = session.last_consolidated archive_start = session.last_consolidated
messages_to_archive = list(session.messages[archive_start:]) messages_to_archive = list(session.messages[archive_start:])
@@ -1191,7 +1210,7 @@ class Consolidator:
# through the captured batch so new messages remain eligible next time. # through the captured batch so new messages remain eligible next time.
session.last_consolidated = archive_end session.last_consolidated = archive_end
session.provider_state = None session.provider_state = None
self.sessions.save(session) await self._save_session(session)
visible = session.get_history( visible = session.get_history(
max_messages=MIN_COMPACTED_REPLAY_MESSAGES, max_messages=MIN_COMPACTED_REPLAY_MESSAGES,
+7 -5
View File
@@ -210,12 +210,14 @@ class AgentProgressHook(AgentHook):
tool_hint=False, tool_hint=False,
tool_events=tool_events, tool_events=tool_events,
) )
u = context.usage or {} u = context.usage
logger.debug( logger.debug(
"LLM usage: prompt={} completion={} cached={}", "LLM usage: input={} output={} cache_read={} cache_write={} source={}",
u.get("prompt_tokens", 0), u.input_tokens if u else 0,
u.get("completion_tokens", 0), u.output_tokens if u else 0,
u.get("cached_tokens", 0), u.cache_read_tokens if u else None,
u.cache_write_tokens if u else None,
u.source if u else "missing",
) )
def finalize_content(self, context: AgentHookContext, content: str | None) -> str | None: def finalize_content(self, context: AgentHookContext, content: str | None) -> str | None:
+131 -92
View File
@@ -6,7 +6,7 @@ import asyncio
import inspect import inspect
import os import os
import time import time
from collections.abc import Awaitable, Callable, Iterable from collections.abc import Awaitable, Callable, Iterable, Sized
from copy import deepcopy from copy import deepcopy
from dataclasses import dataclass, field from dataclasses import dataclass, field
from pathlib import Path from pathlib import Path
@@ -20,9 +20,16 @@ from nanobot.agent.context_governance import (
) )
from nanobot.agent.hook import AgentHook, AgentHookContext, AgentRunHookContext from nanobot.agent.hook import AgentHook, AgentHookContext, AgentRunHookContext
from nanobot.agent.tools.registry import ToolRegistry, is_tool_error_result from nanobot.agent.tools.registry import ToolRegistry, is_tool_error_result
from nanobot.llm_usage.context import (
LLMUsageSource,
bind_llm_usage_source,
reset_llm_usage_source,
source_from_session_key,
)
from nanobot.providers.base import ( from nanobot.providers.base import (
LLMProvider, LLMProvider,
LLMResponse, LLMResponse,
LLMUsage,
ProviderCallContext, ProviderCallContext,
ProviderConversationState, ProviderConversationState,
ToolCallRequest, ToolCallRequest,
@@ -37,6 +44,7 @@ from nanobot.runtime_context import (
reattach_runtime_context, reattach_runtime_context,
) )
from nanobot.session.history_visibility import is_hidden_history_message from nanobot.session.history_visibility import is_hidden_history_message
from nanobot.session.recovery import PENDING_FOLLOWUP_ID_KEY
from nanobot.utils.helpers import ( from nanobot.utils.helpers import (
IncrementalThinkExtractor, IncrementalThinkExtractor,
build_assistant_message, build_assistant_message,
@@ -75,6 +83,22 @@ _MAX_EMPTY_RETRIES = 2
_MAX_LENGTH_RECOVERIES = 3 _MAX_LENGTH_RECOVERIES = 3
_MAX_INJECTIONS_PER_TURN = 3 _MAX_INJECTIONS_PER_TURN = 3
_MAX_INJECTION_CYCLES = 5 _MAX_INJECTION_CYCLES = 5
_SLOW_TOOL_LOG_MS = 1_000
def _tool_input_scale(params: object) -> tuple[int, int]:
"""Return bounded structural counts without logging argument content."""
if not isinstance(params, dict):
return 0, len(params) if isinstance(params, str | bytes) else 0
params_dict = cast(dict[object, object], params)
items = len(params_dict)
chars = 0
for value in params_dict.values():
if isinstance(value, str | bytes):
chars += len(value)
elif isinstance(value, list | tuple | set | dict):
items += len(cast(Sized, value))
return items, chars
def _restore_outer_whitespace(content: str, original: str | None) -> str: def _restore_outer_whitespace(content: str, original: str | None) -> str:
@@ -116,6 +140,7 @@ class AgentRunSpec:
goal_continue_message: GoalContinueMessage | None = None goal_continue_message: GoalContinueMessage | None = None
finalize_on_max_iterations: bool = True finalize_on_max_iterations: bool = True
provider_state: ProviderConversationState | None = None provider_state: ProviderConversationState | None = None
llm_usage_source: LLMUsageSource | None = None
@dataclass(slots=True) @dataclass(slots=True)
@@ -125,7 +150,7 @@ class AgentRunResult:
final_content: str | None final_content: str | None
messages: list[dict[str, Any]] messages: list[dict[str, Any]]
tools_used: list[str] = field(default_factory=list) tools_used: list[str] = field(default_factory=list)
usage: dict[str, int] = field(default_factory=dict) usage: LLMUsage | None = None
stop_reason: str = "completed" stop_reason: str = "completed"
error: str | None = None error: str | None = None
tool_events: list[dict[str, str]] = field(default_factory=list) tool_events: list[dict[str, str]] = field(default_factory=list)
@@ -234,6 +259,23 @@ class AgentRunner:
merged.get("content"), merged.get("content"),
injection.get("content"), injection.get("content"),
) )
followup_id = injection.get(PENDING_FOLLOWUP_ID_KEY)
if isinstance(followup_id, str) and followup_id:
existing = cast(object, merged.get(PENDING_FOLLOWUP_ID_KEY))
followup_ids = (
[existing]
if isinstance(existing, str)
else [
item
for item in cast(list[object], existing)
if isinstance(item, str)
]
if isinstance(existing, list)
else []
)
if followup_id not in followup_ids:
followup_ids.append(followup_id)
merged[PENDING_FOLLOWUP_ID_KEY] = followup_ids
messages[-1] = merged messages[-1] = merged
continue continue
messages.append(injection) messages.append(injection)
@@ -373,6 +415,9 @@ class AgentRunner:
hook = spec.hook or AgentHook() hook = spec.hook or AgentHook()
messages = list(spec.initial_messages) messages = list(spec.initial_messages)
context = AgentRunHookContext(messages=deepcopy(messages)) context = AgentRunHookContext(messages=deepcopy(messages))
llm_usage_source_token = bind_llm_usage_source(
spec.llm_usage_source or source_from_session_key(spec.session_key)
)
try: try:
await hook.before_run(context) await hook.before_run(context)
@@ -394,7 +439,7 @@ class AgentRunner:
context.messages = deepcopy(result.messages) context.messages = deepcopy(result.messages)
context.final_content = result.final_content context.final_content = result.final_content
context.tools_used = list(result.tools_used) context.tools_used = list(result.tools_used)
context.usage = dict(result.usage) context.usage = result.usage
context.stop_reason = result.stop_reason context.stop_reason = result.stop_reason
context.error = result.error context.error = result.error
context.tool_events = deepcopy(result.tool_events) context.tool_events = deepcopy(result.tool_events)
@@ -405,6 +450,7 @@ class AgentRunner:
await hook.after_run(context) await hook.after_run(context)
return result return result
finally: finally:
try:
context.messages = deepcopy(messages) context.messages = deepcopy(messages)
if context.exception is None: if context.exception is None:
await hook.on_finally(context) await hook.on_finally(context)
@@ -416,6 +462,8 @@ class AgentRunner:
"AgentHook.on_finally error after {}", "AgentHook.on_finally error after {}",
context.stop_reason or "run exception", context.stop_reason or "run exception",
) )
finally:
reset_llm_usage_source(llm_usage_source_token)
async def _run_core( async def _run_core(
self, self,
@@ -425,7 +473,7 @@ class AgentRunner:
) -> AgentRunResult: ) -> AgentRunResult:
final_content: str | None = None final_content: str | None = None
tools_used: list[str] = [] tools_used: list[str] = []
usage = {"prompt_tokens": 0, "completion_tokens": 0} usage: LLMUsage | None = None
error: str | None = None error: str | None = None
stop_reason = "completed" stop_reason = "completed"
tool_events: list[dict[str, str]] = [] tool_events: list[dict[str, str]] = []
@@ -501,8 +549,8 @@ class AgentRunner:
) )
response.content = cleaned_content response.content = cleaned_content
raw_usage = self._usage_or_estimate(spec, messages_for_model, response) raw_usage = self._usage_or_estimate(spec, messages_for_model, response)
context.usage = dict(raw_usage) context.usage = raw_usage
self._accumulate_usage(usage, raw_usage) usage = self._merge_usage(usage, raw_usage)
if reasoning_text and not context.streamed_reasoning: if reasoning_text and not context.streamed_reasoning:
await hook.emit_reasoning(reasoning_text) await hook.emit_reasoning(reasoning_text)
await hook.emit_reasoning_end() await hook.emit_reasoning_end()
@@ -665,10 +713,10 @@ class AgentRunner:
conversation_state=conversation_state, conversation_state=conversation_state,
) )
retry_usage = self._usage_or_estimate(spec, retry_messages, response) retry_usage = self._usage_or_estimate(spec, retry_messages, response)
self._accumulate_usage(usage, retry_usage) usage = self._merge_usage(usage, retry_usage)
raw_usage = self._merge_usage(raw_usage, retry_usage) raw_usage = self._merge_usage(raw_usage, retry_usage)
context.response = response context.response = response
context.usage = dict(raw_usage) context.usage = raw_usage
context.tool_calls = list(response.tool_calls) context.tool_calls = list(response.tool_calls)
original_content = response.content original_content = response.content
clean = hook.finalize_content(context, response.content) clean = hook.finalize_content(context, response.content)
@@ -841,7 +889,7 @@ class AgentRunner:
had_injections = True had_injections = True
terminal_content = None terminal_content = None
if spec.finalize_on_max_iterations: if spec.finalize_on_max_iterations:
terminal_content = await self._try_finalize_after_max_iterations( terminal_content, usage = await self._try_finalize_after_max_iterations(
spec, spec,
hook, hook,
messages, messages,
@@ -904,18 +952,7 @@ class AgentRunner:
conversation_state: ProviderConversationStateController, conversation_state: ProviderConversationStateController,
provider_context: ProviderCallContext | None = None, provider_context: ProviderCallContext | None = None,
) -> LLMResponse: ) -> LLMResponse:
timeout_s: float | None = spec.llm_timeout_s timeout_s = self._resolve_llm_timeout_s(spec)
if timeout_s is None:
# Default to a finite timeout to avoid per-session lock starvation when an LLM
# request hangs indefinitely (e.g. gateway/network stall).
# Set NANOBOT_LLM_TIMEOUT_S=0 to disable.
raw = os.environ.get("NANOBOT_LLM_TIMEOUT_S", "300").strip()
try:
timeout_s = float(raw)
except (TypeError, ValueError):
timeout_s = 300.0
if timeout_s <= 0:
timeout_s = None
kwargs = self._build_request_kwargs( kwargs = self._build_request_kwargs(
spec, spec,
@@ -1229,9 +1266,9 @@ class AgentRunner:
spec: AgentRunSpec, spec: AgentRunSpec,
hook: AgentHook, hook: AgentHook,
messages: list[dict[str, Any]], messages: list[dict[str, Any]],
usage: dict[str, int], usage: LLMUsage | None,
conversation_state: ProviderConversationStateController, conversation_state: ProviderConversationStateController,
) -> str | None: ) -> tuple[str | None, LLMUsage | None]:
retry_messages = self._budget_exhausted_finalization_messages(messages) retry_messages = self._budget_exhausted_finalization_messages(messages)
try: try:
response = await self._request_no_tools( response = await self._request_no_tools(
@@ -1246,10 +1283,10 @@ class AgentRunner:
"Budget-exhausted finalization failed for {}; using fallback", "Budget-exhausted finalization failed for {}; using fallback",
spec.session_key or "default", spec.session_key or "default",
) )
return None return None, usage
raw_usage = self._usage_or_estimate(spec, retry_messages, response) raw_usage = self._usage_or_estimate(spec, retry_messages, response)
self._accumulate_usage(usage, raw_usage) usage = self._merge_usage(usage, raw_usage)
if response.finish_reason == "error" or response.has_tool_calls: if response.finish_reason == "error" or response.has_tool_calls:
logger.warning( logger.warning(
"Budget-exhausted finalization returned finish_reason='{}' " "Budget-exhausted finalization returned finish_reason='{}' "
@@ -1258,19 +1295,19 @@ class AgentRunner:
len(response.tool_calls), len(response.tool_calls),
spec.session_key or "default", spec.session_key or "default",
) )
return None return None, usage
context = AgentHookContext( context = AgentHookContext(
iteration=spec.max_iterations, iteration=spec.max_iterations,
messages=messages, messages=messages,
response=response, response=response,
usage=dict(raw_usage), usage=raw_usage,
session_key=spec.session_key, session_key=spec.session_key,
) )
clean = hook.finalize_content(context, response.content) clean = hook.finalize_content(context, response.content)
if is_blank_text(clean): if is_blank_text(clean):
return None return None, usage
return clean return clean, usage
async def _request_no_tools( async def _request_no_tools(
self, self,
@@ -1284,10 +1321,38 @@ class AgentRunner:
messages, messages,
tools=None, tools=None,
) )
return await spec.runtime.provider.chat_with_retry( coro = spec.runtime.provider.chat_with_retry(
**kwargs, **kwargs,
provider_context=provider_context, provider_context=provider_context,
) )
timeout_s = self._resolve_llm_timeout_s(spec)
try:
return (
await coro
if timeout_s is None
else await asyncio.wait_for(coro, timeout=timeout_s)
)
except asyncio.TimeoutError:
return LLMResponse(
content=f"Error calling LLM: timed out after {timeout_s:g}s",
finish_reason="error",
error_kind="timeout",
)
@staticmethod
def _resolve_llm_timeout_s(spec: AgentRunSpec) -> float | None:
"""Resolve the wall-clock limit shared by every model request path."""
timeout_s = spec.llm_timeout_s
if timeout_s is None:
# Default to a finite timeout to avoid per-session lock starvation when an LLM
# request hangs indefinitely (e.g. gateway/network stall).
# Set NANOBOT_LLM_TIMEOUT_S=0 to disable.
raw = os.environ.get("NANOBOT_LLM_TIMEOUT_S", "300").strip()
try:
timeout_s = float(raw)
except (TypeError, ValueError):
timeout_s = 300.0
return timeout_s if timeout_s > 0 else None
@staticmethod @staticmethod
def _budget_exhausted_finalization_messages( def _budget_exhausted_finalization_messages(
@@ -1314,31 +1379,24 @@ class AgentRunner:
spec: AgentRunSpec, spec: AgentRunSpec,
messages: list[dict[str, Any]], messages: list[dict[str, Any]],
response: LLMResponse, response: LLMResponse,
) -> dict[str, int]: ) -> LLMUsage | None:
usage = self._usage_dict(response.usage) usage = response.usage
total = self._usage_total(usage) if response.finish_reason == "error":
if total > 0: if usage is None or usage.total_tokens == 0:
usage["total_tokens"] = total usage = LLMUsage.empty_request()
usage.setdefault("provider_tokens", total) elif usage is None or usage.total_tokens == 0:
elif response.finish_reason == "error":
return {}
else:
usage = self._estimate_response_usage(spec, messages, response) usage = self._estimate_response_usage(spec, messages, response)
completion = usage.get("completion_tokens", 0) return usage.with_timing(
if response.generation_ms is not None and completion > 0: generation_ms=response.generation_ms,
usage["generation_ms"] = response.generation_ms ttft_ms=response.ttft_ms,
usage["measured_completion_tokens"] = completion )
if response.ttft_ms is not None:
usage["ttft_ms"] = response.ttft_ms
usage["timed_requests"] = 1
return usage
def _estimate_response_usage( def _estimate_response_usage(
self, self,
spec: AgentRunSpec, spec: AgentRunSpec,
messages: list[dict[str, Any]], messages: list[dict[str, Any]],
response: LLMResponse, response: LLMResponse,
) -> dict[str, int]: ) -> LLMUsage:
try: try:
tools = spec.tools.get_definitions() tools = spec.tools.get_definitions()
except Exception: except Exception:
@@ -1356,52 +1414,21 @@ class AgentRunner:
thinking_blocks=response.thinking_blocks, thinking_blocks=response.thinking_blocks,
) )
completion_tokens = estimate_message_tokens(assistant_message) completion_tokens = estimate_message_tokens(assistant_message)
total_tokens = max(0, prompt_tokens) + max(0, completion_tokens) return LLMUsage.estimated(
if total_tokens <= 0: input_tokens=max(0, prompt_tokens),
return {} output_tokens=max(0, completion_tokens),
return { )
"prompt_tokens": max(0, prompt_tokens),
"completion_tokens": max(0, completion_tokens),
"total_tokens": total_tokens,
"estimated_tokens": total_tokens,
}
@staticmethod @staticmethod
def _usage_dict(usage: dict[str, Any] | None) -> dict[str, int]: def _merge_usage(
if not usage: left: LLMUsage | None,
return {} right: LLMUsage | None,
result: dict[str, int] = {} ) -> LLMUsage | None:
for key, value in usage.items(): if left is None:
try: return right
result[key] = int(value or 0) if right is None:
except (TypeError, ValueError): return left
continue return left + right
return result
@staticmethod
def _usage_total(usage: dict[str, int]) -> int:
return max(0, usage.get("total_tokens", 0) or (
usage.get("prompt_tokens", 0) + usage.get("completion_tokens", 0)
))
@staticmethod
def _merge_usage(left: dict[str, int], right: dict[str, int]) -> dict[str, int]:
merged = dict(left)
for key, value in right.items():
merged[key] = merged.get(key, 0) + value
return merged
@staticmethod
def _accumulate_usage(total: dict[str, int], request: dict[str, int]) -> None:
"""Fold one model request into the current turn's usage."""
total["request_count"] = total.get("request_count", 0) + 1
prompt_tokens = request.get("prompt_tokens")
if prompt_tokens is not None and prompt_tokens >= 0:
total["context_tokens"] = prompt_tokens
for key, value in request.items():
if key in {"context_tokens", "request_count"} or value < 0:
continue
total[key] = total.get(key, 0) + value
async def _execute_tools( async def _execute_tools(
self, self,
@@ -1510,6 +1537,7 @@ class AgentRunner:
RuntimeError(prep_error) if spec.fail_on_tool_error else None RuntimeError(prep_error) if spec.fail_on_tool_error else None
) )
await hook.before_execute_tool(context, tool_call, tool, params) await hook.before_execute_tool(context, tool_call, tool, params)
tool_started_at = time.perf_counter()
try: try:
if tool is not None: if tool is not None:
result = await tool.execute(**params) result = await tool.execute(**params)
@@ -1538,6 +1566,17 @@ class AgentRunner:
if spec.fail_on_tool_error: if spec.fail_on_tool_error:
return payload, event, exc return payload, event, exc
return payload, event, None return payload, event, None
finally:
duration_ms = int((time.perf_counter() - tool_started_at) * 1000)
if duration_ms >= _SLOW_TOOL_LOG_MS:
input_items, input_chars = _tool_input_scale(params)
logger.warning(
"slow tool operation={} input_items={} input_chars={} duration_ms={}",
tool_call.name,
input_items,
input_chars,
duration_ms,
)
if is_tool_error_result(result): if is_tool_error_result(result):
await hook.on_execute_tool_error(context, tool_call, tool, params, result) await hook.on_execute_tool_error(context, tool_call, tool, params, result)
+12 -4
View File
@@ -8,7 +8,7 @@ import warnings
from collections.abc import Mapping from collections.abc import Mapping
from dataclasses import dataclass, field from dataclasses import dataclass, field
from pathlib import Path from pathlib import Path
from typing import Any, Callable, TypedDict from typing import Any, Callable, NotRequired, TypedDict
from loguru import logger from loguru import logger
@@ -28,7 +28,8 @@ from nanobot.agent.tools.registry import ToolRegistry
from nanobot.bus.events import InboundMessage from nanobot.bus.events import InboundMessage
from nanobot.bus.queue import MessageBus from nanobot.bus.queue import MessageBus
from nanobot.config.schema import AgentDefaults, ToolsConfig from nanobot.config.schema import AgentDefaults, ToolsConfig
from nanobot.providers.base import LLMProvider from nanobot.llm_usage.context import LLMUsageSource, current_llm_usage_source
from nanobot.providers.base import LLMProvider, LLMUsage
from nanobot.security.workspace_access import ( from nanobot.security.workspace_access import (
WorkspaceScope, WorkspaceScope,
bind_workspace_scope, bind_workspace_scope,
@@ -43,6 +44,7 @@ class _SubagentOrigin(TypedDict):
channel: str channel: str
chat_id: str chat_id: str
session_key: str | None session_key: str | None
llm_usage_source: NotRequired[LLMUsageSource]
@dataclass(slots=True) @dataclass(slots=True)
@@ -56,7 +58,7 @@ class SubagentStatus:
phase: str = "initializing" # initializing | awaiting_tools | tools_completed | final_response | done | error phase: str = "initializing" # initializing | awaiting_tools | tools_completed | final_response | done | error
iteration: int = 0 iteration: int = 0
tool_events: list[dict[str, str]] = field(default_factory=list) tool_events: list[dict[str, str]] = field(default_factory=list)
usage: dict[str, int] = field(default_factory=dict) usage: LLMUsage | None = None
stop_reason: str | None = None stop_reason: str | None = None
error: str | None = None error: str | None = None
@@ -82,7 +84,7 @@ class _SubagentHook(AgentHook):
return return
self._status.iteration = context.iteration self._status.iteration = context.iteration
self._status.tool_events = list(context.tool_events) self._status.tool_events = list(context.tool_events)
self._status.usage = dict(context.usage) self._status.usage = context.usage
if context.error: if context.error:
self._status.error = str(context.error) self._status.error = str(context.error)
@@ -252,6 +254,7 @@ class SubagentManager:
"channel": origin_channel, "channel": origin_channel,
"chat_id": origin_chat_id, "chat_id": origin_chat_id,
"session_key": session_key, "session_key": session_key,
"llm_usage_source": current_llm_usage_source(),
} }
status = SubagentStatus( status = SubagentStatus(
@@ -315,6 +318,7 @@ class SubagentManager:
"channel": origin_channel, "channel": origin_channel,
"chat_id": origin_chat_id, "chat_id": origin_chat_id,
"session_key": session_key, "session_key": session_key,
"llm_usage_source": current_llm_usage_source(),
} }
status = SubagentStatus( status = SubagentStatus(
task_id=task_id, task_id=task_id,
@@ -417,6 +421,10 @@ class SubagentManager:
session_key=sess_key, session_key=sess_key,
workspace=root, workspace=root,
llm_timeout_s=llm_timeout, llm_timeout_s=llm_timeout,
llm_usage_source=origin.get(
"llm_usage_source",
current_llm_usage_source(),
),
)) ))
finally: finally:
if token is not None: if token is not None:
+22 -8
View File
@@ -4,7 +4,10 @@
from __future__ import annotations from __future__ import annotations
import asyncio
import inspect
from pathlib import Path from pathlib import Path
from typing import TypedDict
from pydantic import Field from pydantic import Field
@@ -24,6 +27,14 @@ from nanobot.runtime_context import RuntimeContextBlock, wrap_runtime_context_li
from nanobot.security.workspace_access import current_tool_workspace from nanobot.security.workspace_access import current_tool_workspace
class _CliAppRunKwargs(TypedDict):
args: list[str]
json_output: bool
working_dir: str | None
timeout: int | None
restrict_to_workspace: bool
class CliAppsToolConfig(Base): class CliAppsToolConfig(Base):
"""CLI Apps tool configuration.""" """CLI Apps tool configuration."""
@@ -147,14 +158,17 @@ class CliAppsTool(Tool):
) )
workspace = access.project_path or self.workspace workspace = access.project_path or self.workspace
manager = CliAppManager(workspace=workspace, runtime=self.runtime) manager = CliAppManager(workspace=workspace, runtime=self.runtime)
run_kwargs: _CliAppRunKwargs = {
"args": args or [],
"json_output": bool(json),
"working_dir": working_dir,
"timeout": timeout,
"restrict_to_workspace": access.restrict_to_workspace,
}
try: try:
return manager.run( run_async = inspect.getattr_static(type(manager), "run_async", None)
name, if inspect.iscoroutinefunction(run_async):
args=args or [], return await manager.run_async(name, **run_kwargs)
json_output=bool(json), return await asyncio.to_thread(manager.run, name, **run_kwargs)
working_dir=working_dir,
timeout=timeout,
restrict_to_workspace=access.restrict_to_workspace,
)
except CliAppError as exc: except CliAppError as exc:
return ToolResult.error(f"Error: {exc.message}") return ToolResult.error(f"Error: {exc.message}")
+29 -2
View File
@@ -143,14 +143,41 @@ class CronTool(Tool):
tz: str | None = None, tz: str | None = None,
at: str | None = None, at: str | None = None,
job_id: str | None = None, job_id: str | None = None,
) -> str:
if action == "add" and self._in_cron_context.get():
return ToolResult.error(
"Error: cannot schedule new jobs from within a cron job execution"
)
return await self._cron.run_sync(
self._execute_sync,
action,
name,
message,
every_seconds,
cron_expr,
tz,
at,
job_id,
)
def _execute_sync(
self,
action: str,
name: str | None,
message: str,
every_seconds: int | None,
cron_expr: str | None,
tz: str | None,
at: str | None,
job_id: str | None,
) -> str: ) -> str:
if action == "add": if action == "add":
if self._in_cron_context.get(): if self._in_cron_context.get():
return ToolResult.error("Error: cannot schedule new jobs from within a cron job execution") return ToolResult.error("Error: cannot schedule new jobs from within a cron job execution")
return self._add_job(name, message, every_seconds, cron_expr, tz, at) return self._add_job(name, message, every_seconds, cron_expr, tz, at)
elif action == "list": if action == "list":
return self._list_jobs() return self._list_jobs()
elif action == "remove": if action == "remove":
return self._remove_job(job_id) return self._remove_job(job_id)
return f"Unknown action: {action}" return f"Unknown action: {action}"
+85 -14
View File
@@ -2,9 +2,11 @@
# pyright: reportPrivateUsage=false, reportUnusedFunction=false # pyright: reportPrivateUsage=false, reportUnusedFunction=false
import asyncio
import difflib import difflib
import mimetypes import mimetypes
import os import os
import threading
from dataclasses import dataclass from dataclasses import dataclass
from pathlib import Path from pathlib import Path
from typing import Any from typing import Any
@@ -21,6 +23,7 @@ from nanobot.agent.tools.schema import (
) )
from nanobot.config_base import Base from nanobot.config_base import Base
from nanobot.security.workspace_access import current_tool_workspace from nanobot.security.workspace_access import current_tool_workspace
from nanobot.utils.cancellation import shield_and_drain
from nanobot.utils.helpers import build_image_content_blocks, detect_image_mime from nanobot.utils.helpers import build_image_content_blocks, detect_image_mime
@@ -664,22 +667,31 @@ def _match_covers_line(match: _MatchSpan, line: int) -> bool:
return match.line <= line <= _match_end_line(match) return match.line <= line <= _match_end_line(match)
def _find_exact_matches(content: str, old_text: str) -> list[_MatchSpan]: def _find_exact_matches(
content: str,
old_text: str,
*,
max_matches: int | None = None,
) -> list[_MatchSpan]:
matches: list[_MatchSpan] = [] matches: list[_MatchSpan] = []
start = 0 search_start = 0
while True: line_start = 0
idx = content.find(old_text, start) line = 1
while max_matches is None or len(matches) < max_matches:
idx = content.find(old_text, search_start)
if idx == -1: if idx == -1:
break break
line += content.count("\n", line_start, idx)
matches.append( matches.append(
_MatchSpan( _MatchSpan(
start=idx, start=idx,
end=idx + len(old_text), end=idx + len(old_text),
text=content[idx : idx + len(old_text)], text=content[idx : idx + len(old_text)],
line=content.count("\n", 0, idx) + 1, line=line,
) )
) )
start = idx + max(1, len(old_text)) line_start = idx
search_start = idx + max(1, len(old_text))
return matches return matches
@@ -735,27 +747,36 @@ def _find_quote_matches(content: str, old_text: str) -> list[_MatchSpan]:
norm_content = _normalize_quotes(content) norm_content = _normalize_quotes(content)
norm_old = _normalize_quotes(old_text) norm_old = _normalize_quotes(old_text)
matches: list[_MatchSpan] = [] matches: list[_MatchSpan] = []
start = 0 search_start = 0
line_start = 0
line = 1
while True: while True:
idx = norm_content.find(norm_old, start) idx = norm_content.find(norm_old, search_start)
if idx == -1: if idx == -1:
break break
line += content.count("\n", line_start, idx)
matches.append( matches.append(
_MatchSpan( _MatchSpan(
start=idx, start=idx,
end=idx + len(old_text), end=idx + len(old_text),
text=content[idx : idx + len(old_text)], text=content[idx : idx + len(old_text)],
line=content.count("\n", 0, idx) + 1, line=line,
) )
) )
start = idx + max(1, len(norm_old)) line_start = idx
search_start = idx + max(1, len(norm_old))
return matches return matches
def _find_matches(content: str, old_text: str) -> list[_MatchSpan]: def _find_matches(
"""Locate all matches using progressively looser strategies.""" content: str,
old_text: str,
*,
max_exact_matches: int | None = None,
) -> list[_MatchSpan]:
"""Locate matches using progressively looser strategies."""
for matcher in ( for matcher in (
lambda: _find_exact_matches(content, old_text), lambda: _find_exact_matches(content, old_text, max_matches=max_exact_matches),
lambda: _find_trim_matches(content, old_text), lambda: _find_trim_matches(content, old_text),
lambda: _find_trim_matches(content, old_text, normalize_quotes=True), lambda: _find_trim_matches(content, old_text, normalize_quotes=True),
lambda: _find_quote_matches(content, old_text), lambda: _find_quote_matches(content, old_text),
@@ -869,6 +890,43 @@ class EditFileTool(_FsTool):
new_text: str | None = None, new_text: str | None = None,
replace_all: bool = False, occurrence: int | None = None, replace_all: bool = False, occurrence: int | None = None,
line_hint: int | None = None, expected_replacements: int | None = None, **kwargs: Any, line_hint: int | None = None, expected_replacements: int | None = None, **kwargs: Any,
) -> str:
cancelled = threading.Event()
commit_lock = threading.Lock()
try:
return await asyncio.to_thread(
self._execute_sync,
path=path,
old_text=old_text,
new_text=new_text,
replace_all=replace_all,
occurrence=occurrence,
line_hint=line_hint,
expected_replacements=expected_replacements,
cancelled=cancelled,
commit_lock=commit_lock,
)
except asyncio.CancelledError:
cancelled.set()
# If a commit already started, do not report cancellation until the
# file bytes and FileStates record are settled. Otherwise, taking
# the lock first guarantees the worker observes ``cancelled`` before
# it can mutate the target.
await shield_and_drain(
asyncio.to_thread(self._wait_for_commit, commit_lock)
)
raise
@staticmethod
def _wait_for_commit(commit_lock: threading.Lock) -> None:
with commit_lock:
pass
def _execute_sync(
self, *, path: str | None, old_text: str | None,
new_text: str | None, replace_all: bool, occurrence: int | None,
line_hint: int | None, expected_replacements: int | None,
cancelled: threading.Event, commit_lock: threading.Lock,
) -> str: ) -> str:
try: try:
if not path: if not path:
@@ -892,6 +950,9 @@ class EditFileTool(_FsTool):
# Create-file semantics: old_text='' + file doesn't exist → create # Create-file semantics: old_text='' + file doesn't exist → create
if not file_exists: if not file_exists:
if old_text == "": if old_text == "":
with commit_lock:
if cancelled.is_set():
return ToolResult.error("Error: edit_file cancelled.")
fp.parent.mkdir(parents=True, exist_ok=True) fp.parent.mkdir(parents=True, exist_ok=True)
fp.write_text(new_text, encoding="utf-8") fp.write_text(new_text, encoding="utf-8")
self._file_states.record_write(fp) self._file_states.record_write(fp)
@@ -912,6 +973,9 @@ class EditFileTool(_FsTool):
content = raw.decode("utf-8") content = raw.decode("utf-8")
if content.strip(): if content.strip():
return ToolResult.error(f"Error: Cannot create file — {path} already exists and is not empty.") return ToolResult.error(f"Error: Cannot create file — {path} already exists and is not empty.")
with commit_lock:
if cancelled.is_set():
return ToolResult.error("Error: edit_file cancelled.")
fp.write_text(new_text, encoding="utf-8") fp.write_text(new_text, encoding="utf-8")
self._file_states.record_write(fp) self._file_states.record_write(fp)
return f"Successfully edited {fp}" return f"Successfully edited {fp}"
@@ -923,7 +987,11 @@ class EditFileTool(_FsTool):
uses_crlf = b"\r\n" in raw uses_crlf = b"\r\n" in raw
content = raw.decode("utf-8").replace("\r\n", "\n") content = raw.decode("utf-8").replace("\r\n", "\n")
norm_old = old_text.replace("\r\n", "\n") norm_old = old_text.replace("\r\n", "\n")
matches = _find_matches(content, norm_old) matches = _find_matches(
content,
norm_old,
max_exact_matches=occurrence,
)
if not matches: if not matches:
return self._not_found_msg(old_text, content, path) return self._not_found_msg(old_text, content, path)
@@ -1000,6 +1068,9 @@ class EditFileTool(_FsTool):
if uses_crlf: if uses_crlf:
new_content = new_content.replace("\n", "\r\n") new_content = new_content.replace("\n", "\r\n")
with commit_lock:
if cancelled.is_set():
return ToolResult.error("Error: edit_file cancelled.")
fp.write_bytes(new_content.encode("utf-8")) fp.write_bytes(new_content.encode("utf-8"))
self._file_states.record_write(fp) self._file_states.record_write(fp)
msg = f"Successfully edited {fp}" msg = f"Successfully edited {fp}"
+44 -14
View File
@@ -17,6 +17,7 @@ from nanobot.agent.tools.context import RequestContext, ToolContext, current_req
from nanobot.agent.tools.schema import StringSchema, tool_parameters_schema from nanobot.agent.tools.schema import StringSchema, tool_parameters_schema
from nanobot.bus.runtime_events import GoalStateChanged, RuntimeEventBus, RuntimeEventContext from nanobot.bus.runtime_events import GoalStateChanged, RuntimeEventBus, RuntimeEventContext
from nanobot.runtime_context import RuntimeContextBlock, wrap_runtime_context_lines from nanobot.runtime_context import RuntimeContextBlock, wrap_runtime_context_lines
from nanobot.session.async_compat import call_session_manager
from nanobot.session.goal_state import ( from nanobot.session.goal_state import (
GOAL_STATE_KEY, GOAL_STATE_KEY,
MAX_GOAL_OBJECTIVE_CHARS, MAX_GOAL_OBJECTIVE_CHARS,
@@ -28,6 +29,7 @@ from nanobot.session.goal_state import (
sustained_goal_active, sustained_goal_active,
) )
from nanobot.session.turn_continuation import reset_goal_continuation_rounds from nanobot.session.turn_continuation import reset_goal_continuation_rounds
from nanobot.utils.cancellation import shield_and_drain
from nanobot.utils.prompt_templates import render_template from nanobot.utils.prompt_templates import render_template
if TYPE_CHECKING: if TYPE_CHECKING:
@@ -60,36 +62,68 @@ class _GoalToolsMixin:
self._sessions = sessions self._sessions = sessions
self._runtime_events = runtime_events self._runtime_events = runtime_events
def _session(self): async def _get_or_create_session(self, key: str):
return await call_session_manager(
self._sessions,
"get_or_create_async",
self._sessions.get_or_create,
key,
)
async def _save_session(self, session: Any) -> None:
await call_session_manager(
self._sessions,
"save_async",
self._sessions.save,
session,
)
async def _session(self):
request_ctx = current_request_context() request_ctx = current_request_context()
if request_ctx is None: if request_ctx is None:
return None return None
key = request_ctx.session_key key = request_ctx.session_key
if not key: if not key:
return None return None
return self._sessions.get_or_create(key) return await self._get_or_create_session(key)
def _goal_mutation_allowed(self) -> bool: def _goal_mutation_allowed(self) -> bool:
return current_request_context() is not None and goal_mutation_allowed() return current_request_context() is not None and goal_mutation_allowed()
def _save_goal_state( async def _save_goal_state(
self, self,
sess: Any, sess: Any,
blob: dict[str, Any], blob: dict[str, Any],
*, *,
reset_continuation: bool = False, reset_continuation: bool = False,
revoke_permission: bool = False,
) -> None: ) -> None:
previous_metadata = deepcopy(sess.metadata) previous_metadata = deepcopy(sess.metadata)
saved = False
async def save_and_publish() -> None:
nonlocal saved
sess.metadata[GOAL_STATE_KEY] = blob sess.metadata[GOAL_STATE_KEY] = blob
discard_legacy_goal_state_key(sess.metadata) discard_legacy_goal_state_key(sess.metadata)
if reset_continuation: if reset_continuation:
reset_goal_continuation_rounds(sess.metadata) reset_goal_continuation_rounds(sess.metadata)
try: try:
self._sessions.save(sess) await self._save_session(sess)
except BaseException: except BaseException:
sess.metadata.clear() sess.metadata.clear()
sess.metadata.update(previous_metadata) sess.metadata.update(previous_metadata)
raise raise
saved = True
await self._publish_goal_state_changed(sess.metadata)
try:
await shield_and_drain(save_and_publish())
finally:
# This ContextVar belongs to the caller task, not the settlement task.
# Apply the post-save permission effect here even when cancellation was
# delayed until the durable save and runtime notification completed.
if revoke_permission and saved:
revoke_goal_mutation_permission()
async def _publish_goal_state_changed(self, metadata: dict[str, Any]) -> None: async def _publish_goal_state_changed(self, metadata: dict[str, Any]) -> None:
runtime_events = self._runtime_events runtime_events = self._runtime_events
@@ -175,7 +209,7 @@ class CreateGoalTool(Tool, _GoalToolsMixin):
) -> RuntimeContextBlock | None: ) -> RuntimeContextBlock | None:
if not request.session_key: if not request.session_key:
return None return None
session = self._sessions.get_or_create(request.session_key) session = await self._get_or_create_session(request.session_key)
goal_start_requested = explicit_goal_requested(request.metadata) goal_start_requested = explicit_goal_requested(request.metadata)
goal_active = sustained_goal_active(session.metadata) goal_active = sustained_goal_active(session.metadata)
if not goal_start_requested and not goal_active: if not goal_start_requested and not goal_active:
@@ -197,7 +231,7 @@ class CreateGoalTool(Tool, _GoalToolsMixin):
ui_summary: str | None = None, ui_summary: str | None = None,
**kwargs: Any, **kwargs: Any,
) -> str: ) -> str:
sess = self._session() sess = await self._session()
if sess is None: if sess is None:
return ToolResult.error( return ToolResult.error(
"Error: create_goal requires an active chat session (missing routing context)." "Error: create_goal requires an active chat session (missing routing context)."
@@ -225,8 +259,7 @@ class CreateGoalTool(Tool, _GoalToolsMixin):
"ui_summary": summary, "ui_summary": summary,
"started_at": _iso_now(), "started_at": _iso_now(),
} }
self._save_goal_state(sess, blob, reset_continuation=True) await self._save_goal_state(sess, blob, reset_continuation=True)
await self._publish_goal_state_changed(sess.metadata)
extra = f"\nSummary line: {summary}" if summary else "" extra = f"\nSummary line: {summary}" if summary else ""
return ( return (
"Goal recorded. Keep working toward the objective using ordinary tools. " "Goal recorded. Keep working toward the objective using ordinary tools. "
@@ -305,7 +338,7 @@ class UpdateGoalTool(Tool, _GoalToolsMixin):
ui_summary: str | None = None, ui_summary: str | None = None,
**kwargs: Any, **kwargs: Any,
) -> str: ) -> str:
sess = self._session() sess = await self._session()
if sess is None: if sess is None:
return ToolResult.error("Error: update_goal requires an active chat session.") return ToolResult.error("Error: update_goal requires an active chat session.")
prior = parse_goal_state(goal_state_raw(sess.metadata)) prior = parse_goal_state(goal_state_raw(sess.metadata))
@@ -340,8 +373,7 @@ class UpdateGoalTool(Tool, _GoalToolsMixin):
"previous_objective": str(prior.get("objective") or ""), "previous_objective": str(prior.get("objective") or ""),
"recap": (recap or "").strip(), "recap": (recap or "").strip(),
} }
self._save_goal_state(sess, blob, reset_continuation=True) await self._save_goal_state(sess, blob, reset_continuation=True)
await self._publish_goal_state_changed(sess.metadata)
extra = f"\nSummary line: {summary}" if summary else "" extra = f"\nSummary line: {summary}" if summary else ""
return "Goal replaced. Continue toward the new objective using ordinary tools." + extra return "Goal replaced. Continue toward the new objective using ordinary tools." + extra
@@ -359,9 +391,7 @@ class UpdateGoalTool(Tool, _GoalToolsMixin):
} }
if normalized == "complete": if normalized == "complete":
blob["completed_at"] = ended blob["completed_at"] = ended
self._save_goal_state(sess, blob) await self._save_goal_state(sess, blob, revoke_permission=True)
revoke_goal_mutation_permission()
await self._publish_goal_state_changed(sess.metadata)
tail = (recap or "").strip() tail = (recap or "").strip()
label = { label = {
+5 -5
View File
@@ -20,10 +20,10 @@ from nanobot.agent.tools.base import Tool, ToolResult
from nanobot.agent.tools.registry import ToolRegistry from nanobot.agent.tools.registry import ToolRegistry
from nanobot.security.network import ( from nanobot.security.network import (
PinnedDNSAsyncTransport, PinnedDNSAsyncTransport,
async_resolve_url_target,
async_validate_url_target,
env_proxy_applies_to_url, env_proxy_applies_to_url,
httpx_env_proxy_mounts, httpx_env_proxy_mounts,
resolve_url_target,
validate_url_target,
) )
from nanobot.utils.cancellation import task_is_cancelling from nanobot.utils.cancellation import task_is_cancelling
@@ -249,7 +249,7 @@ async def _probe_http_url(url: str, timeout: float = 3.0) -> bool:
port = parsed.port port = parsed.port
if not port: if not port:
port = 443 if parsed.scheme == "https" else 80 port = 443 if parsed.scheme == "https" else 80
ok, _, resolved_ips = resolve_url_target(url) ok, _, resolved_ips = await async_resolve_url_target(url)
if not ok: if not ok:
return False return False
if env_proxy_applies_to_url(url): if env_proxy_applies_to_url(url):
@@ -298,7 +298,7 @@ def _pinned_transport_kwargs() -> dict[str, Any]:
async def _validate_mcp_request_url(request: httpx.Request) -> None: async def _validate_mcp_request_url(request: httpx.Request) -> None:
"""Validate each outgoing MCP HTTP request, including redirect targets.""" """Validate each outgoing MCP HTTP request, including redirect targets."""
ok, error = validate_url_target(str(request.url)) ok, error = await async_validate_url_target(str(request.url))
if not ok: if not ok:
raise httpx.RequestError( raise httpx.RequestError(
f"Blocked unsafe MCP URL {_redact_url(str(request.url))} ({error})", f"Blocked unsafe MCP URL {_redact_url(str(request.url))} ({error})",
@@ -1031,7 +1031,7 @@ async def connect_mcp_servers(
return False return False
if transport_type in {"sse", "streamableHttp"}: if transport_type in {"sse", "streamableHttp"}:
ok, error = validate_url_target(cfg.url) ok, error = await async_validate_url_target(cfg.url)
if not ok: if not ok:
logger.warning( logger.warning(
"MCP server '{}': blocked unsafe URL {} ({})", "MCP server '{}': blocked unsafe URL {} ({})",
+28 -4
View File
@@ -12,6 +12,7 @@ if TYPE_CHECKING:
from nanobot.agent.tools.shell import ExecToolConfig from nanobot.agent.tools.shell import ExecToolConfig
from nanobot.agent.tools.web import WebToolsConfig from nanobot.agent.tools.web import WebToolsConfig
from nanobot.config.schema import ModelPresetConfig from nanobot.config.schema import ModelPresetConfig
from nanobot.providers.base import LLMUsage
from nanobot.utils.llm_runtime import LLMRuntime from nanobot.utils.llm_runtime import LLMRuntime
@@ -65,7 +66,7 @@ class RuntimeSnapshot:
web_config: dict[str, object] web_config: dict[str, object]
exec_config: dict[str, object] exec_config: dict[str, object]
subagent_statuses: dict[str, dict[str, object]] subagent_statuses: dict[str, dict[str, object]]
last_usage: dict[str, int] last_usage: Mapping[str, JsonScalar]
scratchpad: dict[str, JsonValue] scratchpad: dict[str, JsonValue]
def as_mapping(self) -> Mapping[str, object]: def as_mapping(self) -> Mapping[str, object]:
@@ -106,6 +107,13 @@ class RuntimeControl(Protocol):
session_key: str | None, session_key: str | None,
) -> LLMRuntime: ... ) -> LLMRuntime: ...
async def set_model_preset_async(
self,
name: str,
*,
session_key: str | None,
) -> LLMRuntime: ...
def set_max_iterations(self, value: int) -> None: ... def set_max_iterations(self, value: int) -> None: ...
def set_context_window_tokens(self, value: int) -> LLMRuntime: ... def set_context_window_tokens(self, value: int) -> LLMRuntime: ...
@@ -151,7 +159,7 @@ class _RuntimeControlTarget(Protocol):
def tool_names(self) -> list[str]: ... def tool_names(self) -> list[str]: ...
@property @property
def last_usage(self) -> Mapping[str, int]: ... def last_usage(self) -> LLMUsage | None: ...
def set_runtime_model(self, model: str) -> LLMRuntime: ... def set_runtime_model(self, model: str) -> LLMRuntime: ...
@@ -161,6 +169,12 @@ class _RuntimeControlTarget(Protocol):
def set_session_model_preset(self, session_key: str, name: str) -> LLMRuntime: ... def set_session_model_preset(self, session_key: str, name: str) -> LLMRuntime: ...
async def set_session_model_preset_async(
self,
session_key: str,
name: str,
) -> LLMRuntime: ...
class AgentRuntimeControl: class AgentRuntimeControl:
"""Allowlisted adapter from agent-loop state to ``RuntimeControl``.""" """Allowlisted adapter from agent-loop state to ``RuntimeControl``."""
@@ -190,7 +204,7 @@ class AgentRuntimeControl:
web_config=_snapshot_web_config(target.web_config), web_config=_snapshot_web_config(target.web_config),
exec_config=_snapshot_exec_config(target.exec_config), exec_config=_snapshot_exec_config(target.exec_config),
subagent_statuses=_snapshot_subagent_statuses(target.subagents), subagent_statuses=_snapshot_subagent_statuses(target.subagents),
last_usage=dict(target.last_usage), last_usage=target.last_usage.to_dict() if target.last_usage is not None else {},
scratchpad=_snapshot_json_mapping(self.__scratchpad), scratchpad=_snapshot_json_mapping(self.__scratchpad),
) )
@@ -207,6 +221,16 @@ class AgentRuntimeControl:
return self.__target.set_session_model_preset(session_key, name) return self.__target.set_session_model_preset(session_key, name)
return self.__target.set_model_preset(name) return self.__target.set_model_preset(name)
async def set_model_preset_async(
self,
name: str,
*,
session_key: str | None,
) -> LLMRuntime:
if session_key is not None:
return await self.__target.set_session_model_preset_async(session_key, name)
return self.__target.set_model_preset(name)
def set_max_iterations(self, value: int) -> None: def set_max_iterations(self, value: int) -> None:
self.__target.max_iterations = value self.__target.max_iterations = value
self.__target.subagents.max_iterations = value self.__target.subagents.max_iterations = value
@@ -297,7 +321,7 @@ def _snapshot_subagent_status(status: SubagentStatus) -> dict[str, object]:
"phase": status.phase, "phase": status.phase,
"iteration": status.iteration, "iteration": status.iteration,
"tool_events": [dict(event) for event in status.tool_events], "tool_events": [dict(event) for event in status.tool_events],
"usage": dict(status.usage), "usage": status.usage.to_dict() if status.usage is not None else None,
"stop_reason": status.stop_reason, "stop_reason": status.stop_reason,
"error": status.error, "error": status.error,
} }
+56 -6
View File
@@ -4,9 +4,12 @@
from __future__ import annotations from __future__ import annotations
import asyncio
import fnmatch import fnmatch
import os import os
import re import re
import threading
import time
from contextlib import suppress from contextlib import suppress
from pathlib import Path, PurePosixPath from pathlib import Path, PurePosixPath
from typing import Any, Iterable, TypeVar from typing import Any, Iterable, TypeVar
@@ -125,6 +128,8 @@ class _SearchTool(_FsTool):
class FindFilesTool(_SearchTool): class FindFilesTool(_SearchTool):
"""Find files by path fragment, glob, or type.""" """Find files by path fragment, glob, or type."""
_scopes = {"core", "subagent"} _scopes = {"core", "subagent"}
_MAX_SCAN_PATHS = 500_000
_MAX_SCAN_SECONDS = 30.0
@property @property
def name(self) -> str: def name(self) -> str:
@@ -218,7 +223,41 @@ class FindFilesTool(_SearchTool):
offset: int = 0, offset: int = 0,
**kwargs: Any, **kwargs: Any,
) -> str: ) -> str:
cancelled = threading.Event()
try: try:
return await asyncio.to_thread(
self._execute_sync,
path=path,
query=query,
glob=glob,
file_type=type,
include_dirs=include_dirs,
sort=sort,
head_limit=head_limit,
offset=offset,
cancelled=cancelled,
)
except asyncio.CancelledError:
cancelled.set()
raise
except PermissionError as e:
return ToolResult.error(f"Error: {e}")
except Exception as e:
return ToolResult.error(f"Error finding files: {e}")
def _execute_sync(
self,
*,
path: str,
query: str | None,
glob: str | None,
file_type: str | None,
include_dirs: bool,
sort: str,
head_limit: int | None,
offset: int,
cancelled: threading.Event,
) -> str:
target = self._resolve(path or ".") target = self._resolve(path or ".")
if not target.exists(): if not target.exists():
return ToolResult.error(f"Error: Path not found: {path}") return ToolResult.error(f"Error: Path not found: {path}")
@@ -235,8 +274,23 @@ class FindFilesTool(_SearchTool):
) )
root = target if target.is_dir() else target.parent root = target if target.is_dir() else target.parent
matches: list[tuple[str, float]] = [] matches: list[tuple[str, float]] = []
deadline = time.monotonic() + self._MAX_SCAN_SECONDS
scanned = 0
for candidate in self._iter_paths(target, include_dirs=include_dirs): for candidate in self._iter_paths(target, include_dirs=include_dirs):
if cancelled.is_set():
raise RuntimeError("find_files scan cancelled")
scanned += 1
if scanned > self._MAX_SCAN_PATHS:
return ToolResult.error(
f"Error: find_files scan exceeded {self._MAX_SCAN_PATHS} paths; "
"narrow path, query, glob, or type and retry."
)
if time.monotonic() > deadline:
return ToolResult.error(
f"Error: find_files scan exceeded {self._MAX_SCAN_SECONDS:g} seconds; "
"narrow path, query, glob, or type and retry."
)
if candidate.is_dir() and not include_dirs: if candidate.is_dir() and not include_dirs:
continue continue
rel_path = candidate.relative_to(root).as_posix() rel_path = candidate.relative_to(root).as_posix()
@@ -245,9 +299,9 @@ class FindFilesTool(_SearchTool):
if glob and not _match_glob(rel_path, name, glob): if glob and not _match_glob(rel_path, name, glob):
continue continue
if candidate.is_file() and not _matches_type(name, type): if candidate.is_file() and not _matches_type(name, file_type):
continue continue
if candidate.is_dir() and type: if candidate.is_dir() and file_type:
continue continue
if not _matches_query(display_path, query): if not _matches_query(display_path, query):
continue continue
@@ -273,10 +327,6 @@ class FindFilesTool(_SearchTool):
if note: if note:
result += "\n\n" + note result += "\n\n" + note
return result return result
except PermissionError as e:
return ToolResult.error(f"Error: {e}")
except Exception as e:
return ToolResult.error(f"Error finding files: {e}")
class GrepTool(_SearchTool): class GrepTool(_SearchTool):
+35 -2
View File
@@ -150,7 +150,7 @@ class MyTool(Tool):
"Actions: check, set.\n" "Actions: check, set.\n"
"- check (no key): full config overview — start here.\n" "- check (no key): full config overview — start here.\n"
"- check (key): drill into a value. Dot-paths allowed " "- check (key): drill into a value. Dot-paths allowed "
"(e.g. '_last_usage.prompt_tokens', 'web_config.enable').\n" "(e.g. '_last_usage.input_tokens', 'web_config.enable').\n"
"- set (key, value): change config or store notes in your scratchpad. " "- set (key, value): change config or store notes in your scratchpad. "
"Scratchpad keys persist across turns but not restarts.\n" "Scratchpad keys persist across turns but not restarts.\n"
"Key values: _current_iteration (current progress), " "Key values: _current_iteration (current progress), "
@@ -370,7 +370,7 @@ class MyTool(Tool):
if not self._modify_allowed: if not self._modify_allowed:
return ToolResult.error("Error: set is disabled (tools.my.allow_set is false)") return ToolResult.error("Error: set is disabled (tools.my.allow_set is false)")
if action in ("modify", "set"): if action in ("modify", "set"):
return self._modify(key, value) return await self._modify_async(key, value)
return f"Unknown action: {action}" return f"Unknown action: {action}"
# -- inspect -- # -- inspect --
@@ -492,6 +492,11 @@ class MyTool(Tool):
return ToolResult.error(f"Error: '{key}' is read-only and cannot be modified") return ToolResult.error(f"Error: '{key}' is read-only and cannot be modified")
return self._modify_scratchpad(key, value) return self._modify_scratchpad(key, value)
async def _modify_async(self, key: str | None, value: Any) -> str:
if key == "model_preset":
return await self._modify_model_preset_async(value)
return self._modify(key, value)
def _modify_model_preset(self, value: Any) -> str: def _modify_model_preset(self, value: Any) -> str:
if not isinstance(value, str) or not value.strip(): if not isinstance(value, str) or not value.strip():
return ToolResult.error("Error: 'model_preset' must be a non-empty string") return ToolResult.error("Error: 'model_preset' must be a non-empty string")
@@ -520,6 +525,34 @@ class MyTool(Tool):
f"context_window_tokens is now {runtime.context_window_tokens!r}" f"context_window_tokens is now {runtime.context_window_tokens!r}"
) )
async def _modify_model_preset_async(self, value: Any) -> str:
if not isinstance(value, str) or not value.strip():
return ToolResult.error("Error: 'model_preset' must be a non-empty string")
name = value.strip()
session_key = current_request_session_key()
old = self._runtime_control.snapshot().model_preset
try:
runtime = await self._runtime_control.set_model_preset_async(
name,
session_key=session_key,
)
except (KeyError, ValueError) as exc:
message = str(exc.args[0]) if exc.args else str(exc)
punctuation = "" if message.endswith((".", "!", "?")) else "."
return ToolResult.error(f"Error: {message}{punctuation}")
if session_key:
self._audit("modify", f"model_preset = {name!r}")
return (
f"Set model_preset = {name!r} for the next turn; "
f"model will be {runtime.model!r}; "
f"context_window_tokens will be {runtime.context_window_tokens!r}"
)
self._audit("modify", f"model_preset: {old!r} -> {name!r}")
return (
f"Set model_preset = {name!r} (was {old!r}); model is now {runtime.model!r}; "
f"context_window_tokens is now {runtime.context_window_tokens!r}"
)
def _modify_restricted(self, key: str, value: Any) -> str: def _modify_restricted(self, key: str, value: Any) -> str:
spec = self.RESTRICTED[key] spec = self.RESTRICTED[key]
expected = cast(type[Any], spec["type"]) expected = cast(type[Any], spec["type"])
+29 -2
View File
@@ -267,6 +267,7 @@ class ExecTool(Tool):
_MAX_TIMEOUT = 600 _MAX_TIMEOUT = 600
_MAX_OUTPUT = 10_000 _MAX_OUTPUT = 10_000
_PREPARE_TIMEOUT_SECONDS = 6.0
# Kernel device files safe as stdio redirect targets (#3599). # Kernel device files safe as stdio redirect targets (#3599).
_BENIGN_DEVICE_PATHS: frozenset[str] = frozenset({ _BENIGN_DEVICE_PATHS: frozenset[str] = frozenset({
@@ -324,7 +325,20 @@ class ExecTool(Tool):
if max_output_chars is None: if max_output_chars is None:
max_output_chars = max_output_tokens max_output_chars = max_output_tokens
prepared = self._prepare_command(command, working_dir, timeout, shell, login) try:
prepared = await asyncio.wait_for(
asyncio.to_thread(
self._prepare_command,
command,
working_dir,
timeout,
shell,
login,
),
timeout=self._PREPARE_TIMEOUT_SECONDS,
)
except asyncio.TimeoutError:
return ToolResult.error("Error: command validation timed out")
if isinstance(prepared, str): if isinstance(prepared, str):
return prepared return prepared
@@ -470,10 +484,14 @@ class ExecTool(Tool):
+ _WORKSPACE_BOUNDARY_NOTE + _WORKSPACE_BOUNDARY_NOTE
) )
# Full access is an explicit trust decision. Keep the application-level
# command guard aligned with the selected access mode instead of
# continuing to block commands after workspace restriction is disabled.
if access.restrict_to_workspace:
guard_error = self._guard_command( guard_error = self._guard_command(
command, command,
cwd, cwd,
restrict_to_workspace=access.restrict_to_workspace, restrict_to_workspace=True,
workspace_root=workspace_root, workspace_root=workspace_root,
) )
if guard_error: if guard_error:
@@ -912,6 +930,15 @@ class ExecTool(Tool):
if self._is_benign_device_path(expanded): if self._is_benign_device_path(expanded):
continue continue
except Exception: except Exception:
# ``Path.expanduser()`` raises when a named user's home
# cannot be resolved (notably on Windows). An extracted
# home path must fail closed rather than bypass the guard.
if raw.strip().startswith("~"):
return ToolResult.error(
"Error: Command blocked by safety guard "
"(path outside working dir)"
+ _WORKSPACE_BOUNDARY_NOTE
)
continue continue
if self._is_benign_device_path(str(p)): if self._is_benign_device_path(str(p)):
+12 -14
View File
@@ -96,7 +96,7 @@ def _normalize(text: str) -> str:
def _validate_url(url: str) -> tuple[bool, str]: def _validate_url(url: str) -> tuple[bool, str]:
"""Validate URL scheme/domain. Does NOT check resolved IPs (use _validate_url_safe for that).""" """Validate URL scheme/domain. Does not resolve IPs; use the async safe helper for that."""
try: try:
p = urlparse(url) p = urlparse(url)
if p.scheme not in ('http', 'https'): if p.scheme not in ('http', 'https'):
@@ -108,18 +108,16 @@ def _validate_url(url: str) -> tuple[bool, str]:
return False, str(e) return False, str(e)
def _validate_url_safe(url: str) -> tuple[bool, str]: async def _async_validate_url_safe(url: str) -> tuple[bool, str]:
"""Validate URL with SSRF protection: scheme, domain, and resolved IP check.""" from nanobot.security.network import async_validate_url_target
from nanobot.security.network import validate_url_target
return validate_url_target(url) return await async_validate_url_target(url)
def _resolve_url_safe(url: str) -> tuple[bool, str, tuple[str, ...]]: async def _async_resolve_url_safe(url: str) -> tuple[bool, str, tuple[str, ...]]:
"""Validate URL and return the resolved IPs to pin during the request.""" from nanobot.security.network import async_resolve_url_target
from nanobot.security.network import resolve_url_target
return resolve_url_target(url) return await async_resolve_url_target(url)
def _pinned_dns_transport() -> httpx.AsyncBaseTransport: def _pinned_dns_transport() -> httpx.AsyncBaseTransport:
@@ -209,7 +207,7 @@ async def _get_with_safe_redirects(
"""GET a URL while validating every redirect target before requesting it.""" """GET a URL while validating every redirect target before requesting it."""
current_url = url current_url = url
for _ in range(MAX_REDIRECTS + 1): for _ in range(MAX_REDIRECTS + 1):
is_valid, error_msg, _ = _resolve_url_safe(current_url) is_valid, error_msg, _ = await _async_resolve_url_safe(current_url)
if not is_valid: if not is_valid:
return None, f"Redirect blocked: {error_msg}" return None, f"Redirect blocked: {error_msg}"
@@ -229,7 +227,7 @@ async def _get_with_safe_redirects(
return response, None return response, None
next_url = urljoin(str(response.url), location) next_url = urljoin(str(response.url), location)
is_valid, error_msg = _validate_url_safe(next_url) is_valid, error_msg = await _async_validate_url_safe(next_url)
if not is_valid: if not is_valid:
await response.aclose() await response.aclose()
return None, f"Redirect blocked: {error_msg}" return None, f"Redirect blocked: {error_msg}"
@@ -249,7 +247,7 @@ async def _stream_with_safe_redirects(
current_url = url current_url = url
chain_carries_credentials = _url_carries_credentials(url) chain_carries_credentials = _url_carries_credentials(url)
for _ in range(MAX_REDIRECTS + 1): for _ in range(MAX_REDIRECTS + 1):
is_valid, error_msg, _ = _resolve_url_safe(current_url) is_valid, error_msg, _ = await _async_resolve_url_safe(current_url)
if not is_valid: if not is_valid:
return None, None, f"Redirect blocked: {error_msg}", chain_carries_credentials return None, None, f"Redirect blocked: {error_msg}", chain_carries_credentials
@@ -283,7 +281,7 @@ async def _stream_with_safe_redirects(
chain_carries_credentials = ( chain_carries_credentials = (
chain_carries_credentials or _url_carries_credentials(next_url) chain_carries_credentials or _url_carries_credentials(next_url)
) )
is_valid, error_msg = _validate_url_safe(next_url) is_valid, error_msg = await _async_validate_url_safe(next_url)
if not is_valid: if not is_valid:
await stream.__aexit__(None, None, None) await stream.__aexit__(None, None, None)
return None, None, f"Redirect blocked: {error_msg}", chain_carries_credentials return None, None, f"Redirect blocked: {error_msg}", chain_carries_credentials
@@ -1106,7 +1104,7 @@ class WebFetchTool(Tool):
url = url.strip(" \t\r\n`\"'") url = url.strip(" \t\r\n`\"'")
extract_mode = kwargs.pop("extractMode", extract_mode) extract_mode = kwargs.pop("extractMode", extract_mode)
max_chars = cast(int, kwargs.pop("maxChars", max_chars) or self.max_chars) max_chars = cast(int, kwargs.pop("maxChars", max_chars) or self.max_chars)
is_valid, error_msg = _validate_url_safe(url) is_valid, error_msg = await _async_validate_url_safe(url)
if not is_valid: if not is_valid:
return json.dumps({"error": f"URL validation failed: {error_msg}", "url": url}, ensure_ascii=False) return json.dumps({"error": f"URL validation failed: {error_msg}", "url": url}, ensure_ascii=False)
+3 -2
View File
@@ -4,7 +4,7 @@ from __future__ import annotations
import dataclasses import dataclasses
import time import time
from collections.abc import Awaitable, Callable, Mapping from collections.abc import Awaitable, Callable
from dataclasses import dataclass, field from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, cast from typing import TYPE_CHECKING, Any, cast
@@ -19,6 +19,7 @@ from nanobot.bus.outbound_events import (
from nanobot.bus.progress import build_bus_progress_callback from nanobot.bus.progress import build_bus_progress_callback
from nanobot.bus.queue import MessageBus from nanobot.bus.queue import MessageBus
from nanobot.bus.runtime_events import RuntimeEventBus, RuntimeEventPublisher from nanobot.bus.runtime_events import RuntimeEventBus, RuntimeEventPublisher
from nanobot.providers.base import LLMUsage
if TYPE_CHECKING: if TYPE_CHECKING:
from nanobot.utils.llm_runtime import LLMRuntime from nanobot.utils.llm_runtime import LLMRuntime
@@ -203,7 +204,7 @@ class TurnDelivery:
def record_latency(self, latency_ms: int | None) -> None: def record_latency(self, latency_ms: int | None) -> None:
self.runtime_event_publisher.record_turn_latency(self.session_key, latency_ms) self.runtime_event_publisher.record_turn_latency(self.session_key, latency_ms)
def record_usage(self, usage: Mapping[str, int]) -> None: def record_usage(self, usage: LLMUsage | None) -> None:
self.runtime_event_publisher.record_turn_usage(self.session_key, usage) self.runtime_event_publisher.record_turn_usage(self.session_key, usage)
def background_response( def background_response(
+5 -4
View File
@@ -18,6 +18,7 @@ from aiohttp import web
from loguru import logger from loguru import logger
from nanobot.config.paths import get_media_dir from nanobot.config.paths import get_media_dir
from nanobot.providers.base import LLMUsage
from nanobot.utils.helpers import safe_filename from nanobot.utils.helpers import safe_filename
from nanobot.utils.media_decode import ( from nanobot.utils.media_decode import (
MAX_FILE_SIZE, MAX_FILE_SIZE,
@@ -93,11 +94,11 @@ def _error_json(status: int, message: str, err_type: str = "invalid_request_erro
def _chat_completion_response( def _chat_completion_response(
content: str, content: str,
model: str, model: str,
usage: dict[str, int] | None = None, usage: LLMUsage | None = None,
) -> dict[str, Any]: ) -> dict[str, Any]:
prompt = (usage or {}).get("prompt_tokens", 0) prompt = usage.input_tokens if usage else 0
completion = (usage or {}).get("completion_tokens", 0) completion = usage.output_tokens if usage else 0
total = (usage or {}).get("total_tokens", 0) or prompt + completion total = usage.total_tokens if usage else 0
return { return {
"id": f"chatcmpl-{uuid.uuid4().hex[:12]}", "id": f"chatcmpl-{uuid.uuid4().hex[:12]}",
"object": "chat.completion", "object": "chat.completion",
+373 -40
View File
@@ -2,15 +2,20 @@
from __future__ import annotations from __future__ import annotations
import asyncio
import ctypes
import json import json
import os import os
import re import re
import shlex import shlex
import shutil import shutil
import signal
import subprocess import subprocess
import sys import sys
import time import time
from collections.abc import Iterable from collections.abc import Iterable
from contextlib import suppress
from ctypes import wintypes
from dataclasses import dataclass from dataclasses import dataclass
from importlib import metadata as importlib_metadata from importlib import metadata as importlib_metadata
from pathlib import Path from pathlib import Path
@@ -97,6 +102,141 @@ class CliAppsRuntimeConfig:
catalog_ttl_seconds: int = 3600 catalog_ttl_seconds: int = 3600
@dataclass(slots=True)
class _PreparedCliRun:
name: str
entry: str
resolved: str
args: list[str]
cwd: Path
timeout: int
env: dict[str, str]
artifact_snapshot: dict[Path, tuple[int, int]]
class _JobObjectBasicLimitInformation(ctypes.Structure):
_fields_ = [
("PerProcessUserTimeLimit", ctypes.c_int64),
("PerJobUserTimeLimit", ctypes.c_int64),
("LimitFlags", wintypes.DWORD),
("MinimumWorkingSetSize", ctypes.c_size_t),
("MaximumWorkingSetSize", ctypes.c_size_t),
("ActiveProcessLimit", wintypes.DWORD),
("Affinity", ctypes.c_size_t),
("PriorityClass", wintypes.DWORD),
("SchedulingClass", wintypes.DWORD),
]
class _IoCounters(ctypes.Structure):
_fields_ = [
("ReadOperationCount", ctypes.c_uint64),
("WriteOperationCount", ctypes.c_uint64),
("OtherOperationCount", ctypes.c_uint64),
("ReadTransferCount", ctypes.c_uint64),
("WriteTransferCount", ctypes.c_uint64),
("OtherTransferCount", ctypes.c_uint64),
]
class _JobObjectExtendedLimitInformation(ctypes.Structure):
_fields_ = [
("BasicLimitInformation", _JobObjectBasicLimitInformation),
("IoInfo", _IoCounters),
("ProcessMemoryLimit", ctypes.c_size_t),
("JobMemoryLimit", ctypes.c_size_t),
("PeakProcessMemoryUsed", ctypes.c_size_t),
("PeakJobMemoryUsed", ctypes.c_size_t),
]
class _WindowsJob:
"""Best-effort Windows process tree ownership for timeout/cancellation."""
_KILL_ON_JOB_CLOSE = 0x00002000
_EXTENDED_LIMIT_INFORMATION = 9
_PROCESS_TERMINATE = 0x0001
_PROCESS_SET_QUOTA = 0x0100
def __init__(self) -> None:
win_dll = getattr(ctypes, "WinDLL")
self._kernel32 = win_dll("kernel32", use_last_error=True)
self._kernel32.CreateJobObjectW.argtypes = [wintypes.LPVOID, wintypes.LPCWSTR]
self._kernel32.CreateJobObjectW.restype = wintypes.HANDLE
self._kernel32.OpenProcess.argtypes = [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD]
self._kernel32.OpenProcess.restype = wintypes.HANDLE
self._kernel32.AssignProcessToJobObject.argtypes = [wintypes.HANDLE, wintypes.HANDLE]
self._kernel32.AssignProcessToJobObject.restype = wintypes.BOOL
self._kernel32.SetInformationJobObject.argtypes = [
wintypes.HANDLE,
ctypes.c_int,
wintypes.LPVOID,
wintypes.DWORD,
]
self._kernel32.SetInformationJobObject.restype = wintypes.BOOL
self._kernel32.TerminateJobObject.argtypes = [wintypes.HANDLE, wintypes.UINT]
self._kernel32.TerminateJobObject.restype = wintypes.BOOL
self._kernel32.CloseHandle.argtypes = [wintypes.HANDLE]
self._kernel32.CloseHandle.restype = wintypes.BOOL
self._handle: Any = self._kernel32.CreateJobObjectW(None, None)
if not self._handle:
raise OSError(ctypes.get_last_error(), "CreateJobObjectW failed")
try:
self._set_kill_on_close(True)
except OSError:
self._kernel32.CloseHandle(self._handle)
self._handle = None
raise
@classmethod
def create(cls) -> _WindowsJob | None:
if os.name != "nt":
return None
try:
return cls()
except OSError as exc:
logger.debug("CLI Apps: Windows job object unavailable: {}", exc)
return None
def _set_kill_on_close(self, enabled: bool) -> None:
info = _JobObjectExtendedLimitInformation()
info.BasicLimitInformation.LimitFlags = self._KILL_ON_JOB_CLOSE if enabled else 0
ok = self._kernel32.SetInformationJobObject(
self._handle,
self._EXTENDED_LIMIT_INFORMATION,
ctypes.byref(info),
ctypes.sizeof(info),
)
if not ok:
raise OSError(ctypes.get_last_error(), "SetInformationJobObject failed")
def assign(self, pid: int) -> bool:
process_handle = self._kernel32.OpenProcess(
self._PROCESS_TERMINATE | self._PROCESS_SET_QUOTA,
False,
pid,
)
if not process_handle:
return False
try:
return bool(self._kernel32.AssignProcessToJobObject(self._handle, process_handle))
finally:
self._kernel32.CloseHandle(process_handle)
def terminate(self) -> None:
if self._handle and not self._kernel32.TerminateJobObject(self._handle, 1):
raise OSError(ctypes.get_last_error(), "TerminateJobObject failed")
def close(self, *, kill_descendants: bool) -> None:
if not self._handle:
return
if not kill_descendants:
with suppress(OSError):
self._set_kill_on_close(False)
self._kernel32.CloseHandle(self._handle)
self._handle = None
_BRANDS: dict[str, tuple[str, str]] = { _BRANDS: dict[str, tuple[str, str]] = {
"1password-cli": ("1password", "#3B66BC"), "1password-cli": ("1password", "#3B66BC"),
"arcgis": ("arcgis", "#2C7AC3"), "arcgis": ("arcgis", "#2C7AC3"),
@@ -1428,6 +1568,197 @@ Use the `run_cli_app` tool with `name="{name}"` for command execution. Do not in
lines.append(f"- {rel} ({kind}, {self._format_artifact_size(path)})") lines.append(f"- {rel} ({kind}, {self._format_artifact_size(path)})")
return lines return lines
def _prepare_run(
self,
name: str,
args: list[str] | None,
*,
json_output: bool,
working_dir: str | None,
timeout: int | None,
restrict_to_workspace: bool,
) -> _PreparedCliRun:
app = self.get_app(name)
installed = self._load_installed()
app_name = str(app["name"])
if app_name not in installed:
raise CliAppError(f"CLI app '{name}' is not installed")
cwd = self._resolve_cwd(working_dir, restrict_to_workspace=restrict_to_workspace)
entry = str(installed[app_name].get("entry_point") or app.get("entry_point") or "")
resolved = shutil.which(entry)
if not entry or not resolved:
raise CliAppError(f"{entry or name} is not available on PATH")
clean_args = [str(arg) for arg in (args or [])]
if json_output and "--json" not in clean_args:
clean_args = ["--json", *clean_args]
effective_timeout = max(1, min(timeout or self.runtime.run_timeout, 600))
return _PreparedCliRun(
name=name,
entry=entry,
resolved=resolved,
args=clean_args,
cwd=cwd,
timeout=effective_timeout,
env=self._subprocess_env(),
artifact_snapshot=self._artifact_snapshot(cwd),
)
def _format_run_result(
self,
prepared: _PreparedCliRun,
*,
returncode: int,
stdout: str,
stderr: str,
) -> str:
command = " ".join([prepared.entry, *(shlex.quote(arg) for arg in prepared.args)])
output = [
f"CLI app '{prepared.name}' exited {returncode}.",
f"Command: {command}",
]
if stdout:
output.append("\nSTDOUT:\n" + stdout.rstrip())
if stderr:
output.append("\nSTDERR:\n" + stderr.rstrip())
artifacts = self._changed_artifacts(prepared.cwd, prepared.artifact_snapshot)
if artifacts:
output.append(
"\nArtifacts created or updated:\n"
+ "\n".join(self._format_artifact_lines(prepared.cwd, artifacts))
)
if any(path.suffix.lower() in _INLINE_ARTIFACT_EXTENSIONS for path in artifacts):
output.append(
"\nTo show a preview in WebUI, reference a raster artifact with Markdown "
"using its workspace-relative path, for example `![diagram](diagram.png)`."
)
return _truncate("\n".join(output))
@staticmethod
def _terminate_run_process_sync(
process: subprocess.Popen[str],
job: _WindowsJob | None,
) -> None:
if job is not None:
with suppress(OSError):
job.terminate()
job.close(kill_descendants=True)
elif os.name == "nt":
with suppress(OSError, subprocess.TimeoutExpired):
subprocess.run(
["taskkill", "/PID", str(process.pid), "/T", "/F"],
check=False,
stdout=subprocess.DEVNULL,
stderr=subprocess.DEVNULL,
timeout=5,
)
else:
with suppress(ProcessLookupError, PermissionError):
os.killpg(process.pid, signal.SIGKILL)
if process.poll() is None:
with suppress(ProcessLookupError):
process.kill()
with suppress(subprocess.TimeoutExpired):
process.wait(timeout=5)
@staticmethod
async def _terminate_run_process(
process: asyncio.subprocess.Process,
job: _WindowsJob | None,
) -> None:
if job is not None:
with suppress(OSError):
await asyncio.to_thread(job.terminate)
job.close(kill_descendants=True)
elif os.name == "nt":
with suppress(OSError, asyncio.TimeoutError):
await asyncio.wait_for(
asyncio.to_thread(
subprocess.run,
["taskkill", "/PID", str(process.pid), "/T", "/F"],
check=False,
stdout=subprocess.DEVNULL,
stderr=subprocess.DEVNULL,
timeout=5,
),
timeout=6.0,
)
else:
with suppress(ProcessLookupError, PermissionError):
os.killpg(process.pid, signal.SIGKILL)
if process.returncode is None:
with suppress(ProcessLookupError):
process.kill()
with suppress(asyncio.TimeoutError, ProcessLookupError):
await asyncio.wait_for(process.wait(), timeout=5.0)
async def run_async(
self,
name: str,
args: list[str] | None = None,
*,
json_output: bool = False,
working_dir: str | None = None,
timeout: int | None = None,
restrict_to_workspace: bool = False,
) -> str:
prepared = await asyncio.to_thread(
self._prepare_run,
name,
args,
json_output=json_output,
working_dir=working_dir,
timeout=timeout,
restrict_to_workspace=restrict_to_workspace,
)
process_kwargs: dict[str, Any] = {}
if os.name == "nt":
process_kwargs["creationflags"] = subprocess.CREATE_NEW_PROCESS_GROUP
else:
process_kwargs["start_new_session"] = True
job = _WindowsJob.create()
try:
process = await asyncio.create_subprocess_exec(
prepared.resolved,
*prepared.args,
cwd=str(prepared.cwd),
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
env=prepared.env,
**process_kwargs,
)
except BaseException:
if job is not None:
job.close(kill_descendants=False)
raise
if job is not None and not job.assign(process.pid):
job.close(kill_descendants=False)
job = None
try:
stdout_raw, stderr_raw = await asyncio.wait_for(
process.communicate(),
timeout=prepared.timeout,
)
except asyncio.TimeoutError:
await self._terminate_run_process(process, job)
return f"CLI app '{prepared.name}' timed out after {prepared.timeout}s"
except asyncio.CancelledError:
await self._terminate_run_process(process, job)
raise
except BaseException:
await self._terminate_run_process(process, job)
raise
if job is not None:
job.close(kill_descendants=False)
stdout = stdout_raw.decode("utf-8", errors="replace")
stderr = stderr_raw.decode("utf-8", errors="replace")
return await asyncio.to_thread(
self._format_run_result,
prepared,
returncode=process.returncode or 0,
stdout=stdout,
stderr=stderr,
)
def run( def run(
self, self,
name: str, name: str,
@@ -1438,50 +1769,52 @@ Use the `run_cli_app` tool with `name="{name}"` for command execution. Do not in
timeout: int | None = None, timeout: int | None = None,
restrict_to_workspace: bool = False, restrict_to_workspace: bool = False,
) -> str: ) -> str:
app = self.get_app(name) prepared = self._prepare_run(
installed = self._load_installed() name,
if str(app["name"]) not in installed: args,
raise CliAppError(f"CLI app '{name}' is not installed") json_output=json_output,
cwd = self._resolve_cwd(working_dir, restrict_to_workspace=restrict_to_workspace) working_dir=working_dir,
entry = str(installed[str(app["name"])].get("entry_point") or app.get("entry_point") or "") timeout=timeout,
resolved = shutil.which(entry) restrict_to_workspace=restrict_to_workspace,
if not entry or not resolved: )
raise CliAppError(f"{entry or name} is not available on PATH") process_kwargs: dict[str, Any] = {}
clean_args = [str(arg) for arg in (args or [])] if os.name == "nt":
if json_output and "--json" not in clean_args: process_kwargs["creationflags"] = subprocess.CREATE_NEW_PROCESS_GROUP
clean_args = ["--json", *clean_args] else:
effective_timeout = max(1, min(timeout or self.runtime.run_timeout, 600)) process_kwargs["start_new_session"] = True
artifact_snapshot = self._artifact_snapshot(cwd) job = _WindowsJob.create()
try: try:
result = subprocess.run( process = subprocess.Popen(
[resolved, *clean_args], [prepared.resolved, *prepared.args],
cwd=str(cwd), cwd=str(prepared.cwd),
capture_output=True, stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
text=True, text=True,
encoding="utf-8", encoding="utf-8",
errors="replace", errors="replace",
timeout=effective_timeout, env=prepared.env,
env=self._subprocess_env(), **process_kwargs,
) )
except BaseException:
if job is not None:
job.close(kill_descendants=False)
raise
if job is not None and not job.assign(process.pid):
job.close(kill_descendants=False)
job = None
try:
stdout, stderr = process.communicate(timeout=prepared.timeout)
except subprocess.TimeoutExpired: except subprocess.TimeoutExpired:
return f"CLI app '{name}' timed out after {effective_timeout}s" self._terminate_run_process_sync(process, job)
output = [ return f"CLI app '{prepared.name}' timed out after {prepared.timeout}s"
f"CLI app '{name}' exited {result.returncode}.", except BaseException:
f"Command: {entry} {' '.join(shlex.quote(arg) for arg in clean_args)}".rstrip(), self._terminate_run_process_sync(process, job)
] raise
if result.stdout: if job is not None:
output.append("\nSTDOUT:\n" + result.stdout.rstrip()) job.close(kill_descendants=False)
if result.stderr: return self._format_run_result(
output.append("\nSTDERR:\n" + result.stderr.rstrip()) prepared,
artifacts = self._changed_artifacts(cwd, artifact_snapshot) returncode=process.returncode,
if artifacts: stdout=stdout,
output.append( stderr=stderr,
"\nArtifacts created or updated:\n"
+ "\n".join(self._format_artifact_lines(cwd, artifacts))
) )
if any(path.suffix.lower() in _INLINE_ARTIFACT_EXTENSIONS for path in artifacts):
output.append(
"\nTo show a preview in WebUI, reference a raster artifact with Markdown "
"using its workspace-relative path, for example `![diagram](diagram.png)`."
)
return _truncate("\n".join(output))
+11 -6
View File
@@ -12,6 +12,7 @@ from dataclasses import dataclass, replace
from typing import Any, cast from typing import Any, cast
from nanobot.bus.events import OutboundMessage from nanobot.bus.events import OutboundMessage
from nanobot.providers.base import LLMUsage
class OutboundEvent: class OutboundEvent:
@@ -58,10 +59,19 @@ class StreamedResponseEvent(OutboundEvent):
class TurnEndEvent(OutboundEvent): class TurnEndEvent(OutboundEvent):
latency_ms: int | None = None latency_ms: int | None = None
goal_state: dict[str, Any] | None = None goal_state: dict[str, Any] | None = None
usage: dict[str, int] | None = None usage: LLMUsage | None = None
context_window_tokens: int | None = None context_window_tokens: int | None = None
@dataclass(frozen=True)
class RecoveryStateEvent(OutboundEvent):
status: str
recovery_id: str
reason: str | None = None
attempts: int = 0
can_continue: bool | None = None
@dataclass(frozen=True) @dataclass(frozen=True)
class GoalStatusEvent(OutboundEvent): class GoalStatusEvent(OutboundEvent):
status: str status: str
@@ -188,11 +198,6 @@ def _legacy_event_from_metadata(msg: OutboundMessage) -> OutboundEvent | None:
return TurnEndEvent( return TurnEndEvent(
latency_ms=_metadata_int(meta, "latency_ms"), latency_ms=_metadata_int(meta, "latency_ms"),
goal_state=cast(dict[str, Any], goal_state) if isinstance(goal_state, dict) else None, goal_state=cast(dict[str, Any], goal_state) if isinstance(goal_state, dict) else None,
usage=(
cast(dict[str, int], meta.get("usage"))
if isinstance(meta.get("usage"), dict)
else None
),
context_window_tokens=_metadata_int(meta, "context_window_tokens"), context_window_tokens=_metadata_int(meta, "context_window_tokens"),
) )
if meta.get("_session_updated"): if meta.get("_session_updated"):
+8 -10
View File
@@ -10,13 +10,14 @@ from __future__ import annotations
import asyncio import asyncio
import contextlib import contextlib
import inspect import inspect
from collections.abc import Awaitable, Callable, Mapping from collections.abc import Awaitable, Callable
from dataclasses import dataclass, field from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any from typing import TYPE_CHECKING, Any
from loguru import logger from loguru import logger
from nanobot.bus.events import InboundMessage from nanobot.bus.events import InboundMessage
from nanobot.providers.base import LLMUsage
if TYPE_CHECKING: if TYPE_CHECKING:
from nanobot.utils.llm_runtime import LLMRuntime from nanobot.utils.llm_runtime import LLMRuntime
@@ -72,7 +73,7 @@ class TurnCompleted:
context: RuntimeEventContext context: RuntimeEventContext
latency_ms: int | None = None latency_ms: int | None = None
runtime: LLMRuntime | None = None runtime: LLMRuntime | None = None
usage: dict[str, int] = field(default_factory=dict) usage: LLMUsage | None = None
@dataclass(frozen=True) @dataclass(frozen=True)
@@ -180,7 +181,7 @@ class RuntimeEventPublisher:
self.bus = bus or RuntimeEventBus() self.bus = bus or RuntimeEventBus()
self._turn_latency_ms: dict[str, int] = {} self._turn_latency_ms: dict[str, int] = {}
self._turn_runtime: dict[str, LLMRuntime] = {} self._turn_runtime: dict[str, LLMRuntime] = {}
self._turn_usage: dict[str, dict[str, int]] = {} self._turn_usage: dict[str, LLMUsage] = {}
@staticmethod @staticmethod
def _context( def _context(
@@ -206,12 +207,9 @@ class RuntimeEventPublisher:
if latency_ms is not None: if latency_ms is not None:
self._turn_latency_ms[session_key] = int(latency_ms) self._turn_latency_ms[session_key] = int(latency_ms)
def record_turn_usage(self, session_key: str, usage: Mapping[str, int]) -> None: def record_turn_usage(self, session_key: str, usage: LLMUsage | None) -> None:
self._turn_usage[session_key] = { if usage is not None:
key: int(value) self._turn_usage[session_key] = usage
for key, value in usage.items()
if type(value) is int and value >= 0
}
def clear_turn(self, session_key: str) -> None: def clear_turn(self, session_key: str) -> None:
self._turn_latency_ms.pop(session_key, None) self._turn_latency_ms.pop(session_key, None)
@@ -332,7 +330,7 @@ class RuntimeEventPublisher:
), ),
latency_ms=self._turn_latency_ms.pop(session_key, None), latency_ms=self._turn_latency_ms.pop(session_key, None),
runtime=self._turn_runtime.pop(session_key, None), runtime=self._turn_runtime.pop(session_key, None),
usage=self._turn_usage.pop(session_key, {}), usage=self._turn_usage.pop(session_key, None),
) )
) )
+19 -10
View File
@@ -21,7 +21,10 @@ from nanobot.bus.events import OutboundMessage
from nanobot.bus.queue import MessageBus from nanobot.bus.queue import MessageBus
from nanobot.channels.base import BaseChannel from nanobot.channels.base import BaseChannel
from nanobot.config.schema import Base from nanobot.config.schema import Base
from nanobot.security.network import validate_resolved_url, validate_url_target from nanobot.security.network import (
async_validate_resolved_url,
async_validate_url_target,
)
DINGTALK_MAX_REMOTE_MEDIA_BYTES = 20 * 1024 * 1024 DINGTALK_MAX_REMOTE_MEDIA_BYTES = 20 * 1024 * 1024
DINGTALK_MAX_REMOTE_MEDIA_REDIRECTS = 3 DINGTALK_MAX_REMOTE_MEDIA_REDIRECTS = 3
@@ -417,8 +420,8 @@ class DingTalkChannel(BaseChannel):
return self._zip_bytes(filename, data) return self._zip_bytes(filename, data)
return data, filename, content_type return data, filename, content_type
def _validate_remote_media_url(self, media_ref: str) -> bool: async def _validate_remote_media_url(self, media_ref: str) -> bool:
ok, err = validate_url_target(media_ref) ok, err = await async_validate_url_target(media_ref)
if not ok: if not ok:
self.logger.warning("remote media URL blocked ref={} reason={}", media_ref, err) self.logger.warning("remote media URL blocked ref={} reason={}", media_ref, err)
return False return False
@@ -434,7 +437,11 @@ class DingTalkChannel(BaseChannel):
allowed_hosts = {host.lower() for host in self.config.remote_media_redirect_allowed_hosts} allowed_hosts = {host.lower() for host in self.config.remote_media_redirect_allowed_hosts}
return next_host in allowed_hosts return next_host in allowed_hosts
def _next_remote_media_url(self, current_url: str, location: str | None) -> str | None: async def _next_remote_media_url(
self,
current_url: str,
location: str | None,
) -> str | None:
if not self.config.allow_remote_media_redirects: if not self.config.allow_remote_media_redirects:
self.logger.warning("media download redirect refused ref={}", current_url) self.logger.warning("media download redirect refused ref={}", current_url)
return None return None
@@ -449,7 +456,7 @@ class DingTalkChannel(BaseChannel):
next_url, next_url,
) )
return None return None
if not self._validate_remote_media_url(next_url): if not await self._validate_remote_media_url(next_url):
return None return None
return next_url return next_url
@@ -461,7 +468,7 @@ class DingTalkChannel(BaseChannel):
if not self._http: if not self._http:
return None, None return None, None
if not self._validate_remote_media_url(media_ref): if not await self._validate_remote_media_url(media_ref):
return None, None return None, None
try: try:
@@ -473,7 +480,7 @@ class DingTalkChannel(BaseChannel):
current_url = media_ref current_url = media_ref
for _ in range(DINGTALK_MAX_REMOTE_MEDIA_REDIRECTS + 1): for _ in range(DINGTALK_MAX_REMOTE_MEDIA_REDIRECTS + 1):
async with stream("GET", current_url, follow_redirects=False) as resp: async with stream("GET", current_url, follow_redirects=False) as resp:
final_ok, final_err = validate_resolved_url(str(resp.url)) final_ok, final_err = await async_validate_resolved_url(str(resp.url))
if not final_ok: if not final_ok:
self.logger.warning( self.logger.warning(
"remote media redirect blocked ref={} final={} reason={}", "remote media redirect blocked ref={} final={} reason={}",
@@ -483,7 +490,7 @@ class DingTalkChannel(BaseChannel):
) )
return None, None return None, None
if 300 <= resp.status_code < 400: if 300 <= resp.status_code < 400:
next_url = self._next_remote_media_url( next_url = await self._next_remote_media_url(
str(resp.url), resp.headers.get("location") str(resp.url), resp.headers.get("location")
) )
if not next_url: if not next_url:
@@ -516,7 +523,9 @@ class DingTalkChannel(BaseChannel):
current_url = media_ref current_url = media_ref
for _ in range(DINGTALK_MAX_REMOTE_MEDIA_REDIRECTS + 1): for _ in range(DINGTALK_MAX_REMOTE_MEDIA_REDIRECTS + 1):
resp = await self._http.get(current_url, follow_redirects=False) resp = await self._http.get(current_url, follow_redirects=False)
final_ok, final_err = validate_resolved_url(str(getattr(resp, "url", current_url))) final_ok, final_err = await async_validate_resolved_url(
str(getattr(resp, "url", current_url))
)
if not final_ok: if not final_ok:
self.logger.warning( self.logger.warning(
"remote media redirect blocked ref={} final={} reason={}", "remote media redirect blocked ref={} final={} reason={}",
@@ -526,7 +535,7 @@ class DingTalkChannel(BaseChannel):
) )
return None, None return None, None
if 300 <= resp.status_code < 400: if 300 <= resp.status_code < 400:
next_url = self._next_remote_media_url( next_url = await self._next_remote_media_url(
str(getattr(resp, "url", current_url)), resp.headers.get("location") str(getattr(resp, "url", current_url)), resp.headers.get("location")
) )
if not next_url: if not next_url:
+11
View File
@@ -104,6 +104,9 @@ class ChannelManager:
webui_mcp_runtime_status: Callable[[], Mapping[str, str]] | None = None, webui_mcp_runtime_status: Callable[[], Mapping[str, str]] | None = None,
webui_mcp_reload: Callable[[], Awaitable[dict[str, Any]]] | None = None, webui_mcp_reload: Callable[[], Awaitable[dict[str, Any]]] | None = None,
webui_skill_state_action: Callable[[set[str]], None] | None = None, webui_skill_state_action: Callable[[set[str]], None] | None = None,
webui_recovery_action: (
Callable[[str, dict[str, Any]], Awaitable[dict[str, Any]]] | None
) = None,
config_path: Path | None = None, config_path: Path | None = None,
): ):
if config_path is None: if config_path is None:
@@ -126,6 +129,7 @@ class ChannelManager:
self._webui_mcp_runtime_status = webui_mcp_runtime_status self._webui_mcp_runtime_status = webui_mcp_runtime_status
self._webui_mcp_reload = webui_mcp_reload self._webui_mcp_reload = webui_mcp_reload
self._webui_skill_state_action = webui_skill_state_action self._webui_skill_state_action = webui_skill_state_action
self._webui_recovery_action = webui_recovery_action
self.channels: dict[str, BaseChannel] = {} self.channels: dict[str, BaseChannel] = {}
self._channel_owners: dict[str, str] = {} self._channel_owners: dict[str, str] = {}
self._channel_runtime_specs: dict[str, tuple[str, str]] = {} self._channel_runtime_specs: dict[str, tuple[str, str]] = {}
@@ -197,6 +201,7 @@ class ChannelManager:
mcp_runtime_status=self._webui_mcp_runtime_status, mcp_runtime_status=self._webui_mcp_runtime_status,
mcp_reload=self._webui_mcp_reload, mcp_reload=self._webui_mcp_reload,
skill_state_action=self._webui_skill_state_action, skill_state_action=self._webui_skill_state_action,
recovery_action=self._webui_recovery_action,
logger=logger, logger=logger,
) )
kwargs["gateway"] = gateway kwargs["gateway"] = gateway
@@ -615,6 +620,12 @@ class ChannelManager:
if target is None: if target is None:
logger.warning("Restart notice target channel is not enabled: {}", notice.channel) logger.warning("Restart notice target channel is not enabled: {}", notice.channel)
return return
if notice.channel == "websocket":
# Reconnect and recovery are already represented by WebSocket
# protocol state. A generic restart-complete notice must not
# masquerade as a recovery transition and overwrite a real
# awaiting-user checkpoint in connected clients.
return
while not target.is_running: while not target.is_running:
remaining = deadline - loop.time() remaining = deadline - loop.time()
+3 -3
View File
@@ -24,7 +24,7 @@ from nanobot.bus.queue import MessageBus
from nanobot.channels.base import BaseChannel from nanobot.channels.base import BaseChannel
from nanobot.config.paths import get_media_dir from nanobot.config.paths import get_media_dir
from nanobot.config.schema import Base from nanobot.config.schema import Base
from nanobot.security.network import validate_url_target from nanobot.security.network import async_validate_url_target
from nanobot.utils.helpers import safe_filename from nanobot.utils.helpers import safe_filename
_DOWNLOAD_TIMEOUT = aiohttp.ClientTimeout(total=60) _DOWNLOAD_TIMEOUT = aiohttp.ClientTimeout(total=60)
@@ -473,7 +473,7 @@ class NapcatChannel(BaseChannel):
if not ref: if not ref:
return None return None
if ref.startswith(("http://", "https://")): if ref.startswith(("http://", "https://")):
ok, err = validate_url_target(ref) ok, err = await async_validate_url_target(ref)
if not ok: if not ok:
logger.warning("napcat: rejected remote image '{}': {}", ref, err) logger.warning("napcat: rejected remote image '{}': {}", ref, err)
return None return None
@@ -525,7 +525,7 @@ class NapcatChannel(BaseChannel):
# logger.debug("napcat: downloading image from {}", url) # logger.debug("napcat: downloading image from {}", url)
if self._http is None: if self._http is None:
return None return None
ok, err = validate_url_target(url) ok, err = await async_validate_url_target(url)
if not ok: if not ok:
logger.warning("napcat: skip image '{}': {}", url, err) logger.warning("napcat: skip image '{}': {}", url, err)
return None return None
@@ -149,9 +149,13 @@ async def test_download_image_rejects_redirects(tmp_path, monkeypatch) -> None:
channel = _channel() channel = _channel()
channel._media_root = tmp_path channel._media_root = tmp_path
channel._http = _FakeHttp(_FakeResponse(status=302)) channel._http = _FakeHttp(_FakeResponse(status=302))
async def allow_url(_url: str) -> tuple[bool, str]:
return True, ""
monkeypatch.setattr( monkeypatch.setattr(
"nanobot.channels.napcat.runtime.validate_url_target", "nanobot.channels.napcat.runtime.async_validate_url_target",
lambda _url: (True, ""), allow_url,
) )
result = await channel._download_image({"url": "https://example.com/a.png", "file": "a.png"}) result = await channel._download_image({"url": "https://example.com/a.png", "file": "a.png"})
+2 -2
View File
@@ -40,7 +40,7 @@ from nanobot.bus.events import OutboundMessage
from nanobot.bus.queue import MessageBus from nanobot.bus.queue import MessageBus
from nanobot.channels.base import BaseChannel from nanobot.channels.base import BaseChannel
from nanobot.config.schema import Base from nanobot.config.schema import Base
from nanobot.security.network import validate_url_target from nanobot.security.network import async_validate_url_target
from nanobot.utils.logging_bridge import redirect_lib_logging from nanobot.utils.logging_bridge import redirect_lib_logging
try: try:
@@ -458,7 +458,7 @@ class QQChannel(BaseChannel):
return None, None return None, None
# Remote URL # Remote URL
ok, err = validate_url_target(media_ref) ok, err = await async_validate_url_target(media_ref)
if not ok: if not ok:
self.logger.warning("outbound media URL validation failed url={} err={}", media_ref, err) self.logger.warning("outbound media URL validation failed url={} err={}", media_ref, err)
return None, None return None, None
+2 -2
View File
@@ -23,8 +23,8 @@ from nanobot.config.schema import Base
from nanobot.pairing import is_approved from nanobot.pairing import is_approved
from nanobot.security.network import ( from nanobot.security.network import (
PinnedDNSAsyncTransport, PinnedDNSAsyncTransport,
async_validate_url_target,
httpx_env_proxy_mounts, httpx_env_proxy_mounts,
validate_url_target,
) )
from nanobot.utils.helpers import safe_filename, split_message from nanobot.utils.helpers import safe_filename, split_message
@@ -95,7 +95,7 @@ _HTML_DOWNLOAD_PREFIXES = (b"<!doctype html", b"<html")
async def _validate_slack_download_request(request: httpx.Request) -> None: async def _validate_slack_download_request(request: httpx.Request) -> None:
"""Validate every Slack file request, including redirects, before transport.""" """Validate every Slack file request, including redirects, before transport."""
ok, error = validate_url_target(str(request.url)) ok, error = await async_validate_url_target(str(request.url))
if not ok: if not ok:
raise httpx.RequestError(f"unsafe Slack file URL: {error}", request=request) raise httpx.RequestError(f"unsafe Slack file URL: {error}", request=request)
@@ -859,13 +859,13 @@ def _patch_download_validation(
monkeypatch: pytest.MonkeyPatch, monkeypatch: pytest.MonkeyPatch,
validated: list[str], validated: list[str],
) -> None: ) -> None:
def validate(url: str) -> tuple[bool, str]: async def validate(url: str) -> tuple[bool, str]:
validated.append(url) validated.append(url)
if "169.254.169.254" in url: if "169.254.169.254" in url:
return False, "blocked metadata address" return False, "blocked metadata address"
return True, "" return True, ""
monkeypatch.setattr("nanobot.channels.slack.runtime.validate_url_target", validate) monkeypatch.setattr("nanobot.channels.slack.runtime.async_validate_url_target", validate)
@pytest.mark.asyncio @pytest.mark.asyncio
+2 -2
View File
@@ -36,7 +36,7 @@ from nanobot.channels.base import BaseChannel
from nanobot.command.builtin import build_help_text from nanobot.command.builtin import build_help_text
from nanobot.config.paths import get_media_dir from nanobot.config.paths import get_media_dir
from nanobot.config.schema import Base from nanobot.config.schema import Base
from nanobot.security.network import validate_url_target from nanobot.security.network import async_validate_url_target
from nanobot.utils.helpers import split_message from nanobot.utils.helpers import split_message
from nanobot.utils.logging_bridge import redirect_lib_logging from nanobot.utils.logging_bridge import redirect_lib_logging
@@ -956,7 +956,7 @@ class TelegramChannel(BaseChannel):
# Telegram Bot API accepts HTTP(S) URLs directly for media params. # Telegram Bot API accepts HTTP(S) URLs directly for media params.
if self._is_remote_media_url(media_path): if self._is_remote_media_url(media_path):
ok, error = validate_url_target(media_path) ok, error = await async_validate_url_target(media_path)
if not ok: if not ok:
raise ValueError(f"unsafe media URL: {error}") raise ValueError(f"unsafe media URL: {error}")
await self._call_with_retry( await self._call_with_retry(
@@ -1488,7 +1488,14 @@ async def test_send_remote_media_url_after_security_validation(monkeypatch) -> N
MessageBus(), MessageBus(),
) )
_install_ready_app(channel) _install_ready_app(channel)
monkeypatch.setattr("nanobot.channels.telegram.runtime.validate_url_target", lambda url: (True, ""))
async def allow_url(_url: str) -> tuple[bool, str]:
return True, ""
monkeypatch.setattr(
"nanobot.channels.telegram.runtime.async_validate_url_target",
allow_url,
)
await channel.send( await channel.send(
OutboundMessage( OutboundMessage(
@@ -1546,9 +1553,13 @@ async def test_send_blocks_unsafe_remote_media_url(monkeypatch) -> None:
MessageBus(), MessageBus(),
) )
_install_ready_app(channel) _install_ready_app(channel)
async def deny_url(_url: str) -> tuple[bool, str]:
return False, "Blocked: example.com resolves to private/internal address 127.0.0.1"
monkeypatch.setattr( monkeypatch.setattr(
"nanobot.channels.telegram.runtime.validate_url_target", "nanobot.channels.telegram.runtime.async_validate_url_target",
lambda url: (False, "Blocked: example.com resolves to private/internal address 127.0.0.1"), deny_url,
) )
await channel.send( await channel.send(
+69 -23
View File
@@ -32,6 +32,7 @@ from nanobot.bus.outbound_events import (
GoalStateSyncEvent, GoalStateSyncEvent,
GoalStatusEvent, GoalStatusEvent,
ProgressEvent, ProgressEvent,
RecoveryStateEvent,
RuntimeModelUpdatedEvent, RuntimeModelUpdatedEvent,
SessionUpdatedEvent, SessionUpdatedEvent,
TurnEndEvent, TurnEndEvent,
@@ -43,6 +44,7 @@ from nanobot.bus.queue import MessageBus
from nanobot.channels.base import BaseChannel from nanobot.channels.base import BaseChannel
from nanobot.command.builtin import USER_SHELL_COMMAND, builtin_command_starts_agent_turn from nanobot.command.builtin import USER_SHELL_COMMAND, builtin_command_starts_agent_turn
from nanobot.config.schema import Base from nanobot.config.schema import Base
from nanobot.providers.base import LLMUsage
from nanobot.runtime_context import ( from nanobot.runtime_context import (
RUNTIME_CONTEXT_INPUT_META, RUNTIME_CONTEXT_INPUT_META,
WEBUI_QUOTE_METADATA, WEBUI_QUOTE_METADATA,
@@ -53,8 +55,10 @@ from nanobot.security.workspace_access import (
WORKSPACE_SCOPE_METADATA_KEY, WORKSPACE_SCOPE_METADATA_KEY,
WorkspaceScopeError, WorkspaceScopeError,
) )
from nanobot.session.async_compat import call_session_manager
from nanobot.session.goal_state import goal_state_ws_blob from nanobot.session.goal_state import goal_state_ws_blob
from nanobot.session.model_selection import model_preset_from_metadata from nanobot.session.model_selection import model_preset_from_metadata
from nanobot.session.recovery import recovery_state_from_metadata
from nanobot.session.webui_turns import ( from nanobot.session.webui_turns import (
clear_websocket_turn_if_current, clear_websocket_turn_if_current,
clear_websocket_turns, clear_websocket_turns,
@@ -444,6 +448,26 @@ class WebSocketChannel(BaseChannel):
if sessions is None: if sessions is None:
return {} return {}
snapshot = sessions.read_session_metadata(f"websocket:{chat_id}") snapshot = sessions.read_session_metadata(f"websocket:{chat_id}")
return self._attached_model_fields_from_snapshot(chat_id, snapshot)
async def _attached_model_fields_async(self, chat_id: str) -> dict[str, Any]:
"""Build attach fields without blocking the gateway event loop."""
sessions = self.gateway.session_manager
if sessions is None:
return {}
snapshot = await call_session_manager(
sessions,
"read_session_metadata_async",
sessions.read_session_metadata,
f"websocket:{chat_id}",
)
return self._attached_model_fields_from_snapshot(chat_id, snapshot)
def _attached_model_fields_from_snapshot(
self,
chat_id: str,
snapshot: dict[str, Any] | None,
) -> dict[str, Any]:
raw_metadata = snapshot.get("metadata") if snapshot is not None else None raw_metadata = snapshot.get("metadata") if snapshot is not None else None
metadata = cast(dict[str, object], raw_metadata) if isinstance(raw_metadata, dict) else None metadata = cast(dict[str, object], raw_metadata) if isinstance(raw_metadata, dict) else None
fields: dict[str, Any] = {} fields: dict[str, Any] = {}
@@ -453,18 +477,12 @@ class WebSocketChannel(BaseChannel):
self.logger.warning("ignoring invalid model preset metadata for chat_id={}", chat_id) self.logger.warning("ignoring invalid model preset metadata for chat_id={}", chat_id)
fields["model_preset"] = None fields["model_preset"] = None
if isinstance(metadata, dict): if isinstance(metadata, dict):
usage = metadata.get("_last_usage") recovery_state = recovery_state_from_metadata(metadata)
if isinstance(usage, dict): if recovery_state is not None:
sanitized_usage: dict[str, int | float] = {} fields["recovery_state"] = recovery_state
for key, value in cast(dict[object, object], usage).items(): usage = LLMUsage.from_dict(metadata.get("_last_usage"))
if ( if usage is not None:
isinstance(key, str) fields["usage"] = usage.to_turn_dict()
and isinstance(value, (int, float))
and not isinstance(value, bool)
and value >= 0
):
sanitized_usage[key] = value
fields["usage"] = sanitized_usage
return fields return fields
def _detach(self, connection: ServerConnection, chat_id: str) -> None: def _detach(self, connection: ServerConnection, chat_id: str) -> None:
@@ -513,13 +531,16 @@ class WebSocketChannel(BaseChannel):
fork_key: str, fork_key: str,
) -> None: ) -> None:
"""Attach and hydrate a newly created WebUI chat fork.""" """Attach and hydrate a newly created WebUI chat fork."""
scope = self._workspaces.scope_for_session_key(fork_key) scope = await asyncio.to_thread(
self._workspaces.scope_for_session_key,
fork_key,
)
self._attach(connection, fork_id) self._attach(connection, fork_id)
await self._send_event( await self._send_event(
connection, connection,
"attached", "attached",
chat_id=fork_id, chat_id=fork_id,
**self._attached_model_fields(fork_id), **await self._attached_model_fields_async(fork_id),
) )
await self._send_event( await self._send_event(
connection, connection,
@@ -901,13 +922,13 @@ class WebSocketChannel(BaseChannel):
) )
if scope is None: if scope is None:
return return
self._workspaces.persist_scope(new_id, scope) self._workspaces.stage_scope(new_id, scope)
self._attach(connection, new_id) self._attach(connection, new_id)
await self._send_event( await self._send_event(
connection, connection,
"attached", "attached",
chat_id=new_id, chat_id=new_id,
**self._attached_model_fields(new_id), **await self._attached_model_fields_async(new_id),
) )
await self._send_event( await self._send_event(
connection, connection,
@@ -963,7 +984,7 @@ class WebSocketChannel(BaseChannel):
connection, connection,
"attached", "attached",
chat_id=cid, chat_id=cid,
**self._attached_model_fields(cid), **await self._attached_model_fields_async(cid),
) )
await self._hydrate_after_subscribe(cid) await self._hydrate_after_subscribe(cid)
return return
@@ -1018,7 +1039,7 @@ class WebSocketChannel(BaseChannel):
) )
if scope is None: if scope is None:
return return
self._workspaces.persist_scope(cid, scope) self._workspaces.stage_scope(cid, scope)
# Other clients on the same gateway only need an invalidation; they # Other clients on the same gateway only need an invalidation; they
# can reload the authoritative session row without receiving a # can reload the authoritative session row without receiving a
# local project path that belongs to another connection. # local project path that belongs to another connection.
@@ -1212,7 +1233,6 @@ class WebSocketChannel(BaseChannel):
if session_mentions: if session_mentions:
metadata["session_mentions"] = session_mentions metadata["session_mentions"] = session_mentions
metadata[WORKSPACE_SCOPE_METADATA_KEY] = scope.metadata() metadata[WORKSPACE_SCOPE_METADATA_KEY] = scope.metadata()
self._workspaces.persist_scope(cid, scope)
is_webui = metadata.get("webui") is True is_webui = metadata.get("webui") is True
queued_owner = None queued_owner = None
if is_webui and not is_user_shell and builtin_command_starts_agent_turn(content): if is_webui and not is_user_shell and builtin_command_starts_agent_turn(content):
@@ -1267,6 +1287,7 @@ class WebSocketChannel(BaseChannel):
else False else False
), ),
) )
await asyncio.to_thread(self._workspaces.persist_scope, cid, scope)
accepted = True accepted = True
finally: finally:
if not accepted and queued_owner is not None: if not accepted and queued_owner is not None:
@@ -1556,7 +1577,7 @@ class WebSocketChannel(BaseChannel):
turn_id: str | None = None, turn_id: str | None = None,
) -> Any | None: ) -> Any | None:
try: try:
return resolver() return await asyncio.to_thread(resolver)
except WorkspaceScopeError as exc: except WorkspaceScopeError as exc:
await self._send_event( await self._send_event(
connection, connection,
@@ -1740,6 +1761,10 @@ class WebSocketChannel(BaseChannel):
provenance=event.provenance, provenance=event.provenance,
) )
return return
if isinstance(event, RecoveryStateEvent):
if conns:
await self.send_recovery_state(msg.chat_id, event)
return
if isinstance(event, GoalStateSyncEvent): if isinstance(event, GoalStateSyncEvent):
if conns: if conns:
await self.send_goal_state(msg.chat_id, event.goal_state or {"active": False}) await self.send_goal_state(msg.chat_id, event.goal_state or {"active": False})
@@ -2010,7 +2035,7 @@ class WebSocketChannel(BaseChannel):
latency_ms: int | None = None, latency_ms: int | None = None,
*, *,
goal_state: dict[str, Any] | None = None, goal_state: dict[str, Any] | None = None,
usage: dict[str, int] | None = None, usage: LLMUsage | None = None,
context_window_tokens: int | None = None, context_window_tokens: int | None = None,
metadata: dict[str, Any] | None = None, metadata: dict[str, Any] | None = None,
turn_owner: str | None = None, turn_owner: str | None = None,
@@ -2025,8 +2050,8 @@ class WebSocketChannel(BaseChannel):
body["latency_ms"] = int(latency_ms) body["latency_ms"] = int(latency_ms)
if goal_state is not None: if goal_state is not None:
body["goal_state"] = goal_state body["goal_state"] = goal_state
if usage: if usage is not None:
body["usage"] = usage body["usage"] = usage.to_turn_dict()
if context_window_tokens is not None: if context_window_tokens is not None:
body["context_window_tokens"] = int(context_window_tokens) body["context_window_tokens"] = int(context_window_tokens)
canonical_webui_turn = (metadata or {}).get("webui") is True canonical_webui_turn = (metadata or {}).get("webui") is True
@@ -2057,6 +2082,27 @@ class WebSocketChannel(BaseChannel):
for connection in conns: for connection in conns:
await self._safe_send_to(connection, raw, label=" turn_end ") await self._safe_send_to(connection, raw, label=" turn_end ")
async def send_recovery_state(
self,
chat_id: str,
event: RecoveryStateEvent,
) -> None:
"""Publish one structured recovery transition without chat pollution."""
body: dict[str, Any] = {
"event": "recovery_state",
"chat_id": chat_id,
"status": event.status,
"recovery_id": event.recovery_id,
"attempts": event.attempts,
}
if event.reason:
body["reason"] = event.reason
if event.can_continue is not None:
body["can_continue"] = event.can_continue
raw = json.dumps(body, ensure_ascii=False)
for connection in list(self._subs.get(chat_id, ())):
await self._safe_send_to(connection, raw, label=" recovery_state ")
async def send_goal_state(self, chat_id: str, blob: dict[str, Any]) -> None: async def send_goal_state(self, chat_id: str, blob: dict[str, Any]) -> None:
"""Push persisted goal-state snapshot for *chat_id* (multi-chat isolation).""" """Push persisted goal-state snapshot for *chat_id* (multi-chat isolation)."""
conns = list(self._subs.get(chat_id, ())) conns = list(self._subs.get(chat_id, ()))
@@ -27,6 +27,7 @@ from nanobot.bus.outbound_events import (
GoalStateSyncEvent, GoalStateSyncEvent,
GoalStatusEvent, GoalStatusEvent,
ProgressEvent, ProgressEvent,
RecoveryStateEvent,
RuntimeModelUpdatedEvent, RuntimeModelUpdatedEvent,
SessionUpdatedEvent, SessionUpdatedEvent,
TurnEndEvent, TurnEndEvent,
@@ -43,6 +44,7 @@ from nanobot.channels.websocket.runtime import (
) )
from nanobot.config.loader import load_config, save_config from nanobot.config.loader import load_config, save_config
from nanobot.config.schema import Config, ModelPresetConfig from nanobot.config.schema import Config, ModelPresetConfig
from nanobot.providers.base import LLMUsage
from nanobot.runtime_context import RUNTIME_CONTEXT_INPUT_META, WEBUI_QUOTE_SOURCE from nanobot.runtime_context import RUNTIME_CONTEXT_INPUT_META, WEBUI_QUOTE_SOURCE
from nanobot.security.workspace_access import WORKSPACE_SCOPE_METADATA_KEY from nanobot.security.workspace_access import WORKSPACE_SCOPE_METADATA_KEY
from nanobot.session import webui_turns as wth from nanobot.session import webui_turns as wth
@@ -1510,6 +1512,7 @@ async def test_webui_message_scope_inherits_persisted_session_scope(
}, },
}, },
) )
assert sessions.list_sessions() == []
await channel._dispatch_envelope( await channel._dispatch_envelope(
conn, conn,
"webui-client", "webui-client",
@@ -1523,6 +1526,87 @@ async def test_webui_message_scope_inherits_persisted_session_scope(
} }
@pytest.mark.asyncio
async def test_new_chat_without_message_does_not_create_session(
bus: MagicMock,
tmp_path,
) -> None:
sessions = SessionManager(tmp_path / "sessions")
channel = WebSocketChannel(
{"enabled": True, "allowFrom": ["*"], "host": "127.0.0.1"},
bus,
gateway=_basic_handler(bus, session_manager=sessions, workspace_path=tmp_path),
)
conn = AsyncMock()
conn.remote_address = ("127.0.0.1", 50123)
await channel._dispatch_envelope(
conn,
"tui-client",
{
"type": "new_chat",
"workspace_scope": {
"project_path": str(tmp_path),
"access_mode": "full",
},
},
)
attached = json.loads(conn.send.await_args_list[0].args[0])
assert attached["event"] == "attached"
assert sessions.list_sessions() == []
assert channel._workspaces.scope_for_session_key(
f"websocket:{attached['chat_id']}"
).access_mode == "full"
await channel._cleanup_connection(conn)
assert sessions.list_sessions() == []
@pytest.mark.asyncio
async def test_failed_first_message_does_not_persist_draft_session(
bus: MagicMock,
tmp_path,
) -> None:
sessions = SessionManager(tmp_path / "sessions")
channel = WebSocketChannel(
{"enabled": True, "allowFrom": ["*"], "host": "127.0.0.1"},
bus,
gateway=_basic_handler(bus, session_manager=sessions, workspace_path=tmp_path),
)
conn = AsyncMock()
conn.remote_address = ("127.0.0.1", 50123)
await channel._dispatch_envelope(
conn,
"tui-client",
{
"type": "new_chat",
"workspace_scope": {
"project_path": str(tmp_path),
"access_mode": "full",
},
},
)
chat_id = json.loads(conn.send.await_args_list[0].args[0])["chat_id"]
bus.publish_inbound.side_effect = RuntimeError("queue unavailable")
with pytest.raises(RuntimeError, match="queue unavailable"):
await channel._dispatch_envelope(
conn,
"tui-client",
{
"type": "message",
"chat_id": chat_id,
"content": "hello",
"webui": True,
},
)
assert sessions.list_sessions() == []
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_workspace_scope_change_invalidates_other_attached_clients( async def test_workspace_scope_change_invalidates_other_attached_clients(
bus: MagicMock, bus: MagicMock,
@@ -1729,6 +1813,10 @@ async def test_webui_set_workspace_scope_rejects_running_chat(bus: MagicMock, tm
}, },
}, },
) )
channel._workspaces.persist_scope(
"chat-running",
channel._workspaces.scope_for_session_key("websocket:chat-running"),
)
conn.send.reset_mock() conn.send.reset_mock()
wth._WEBSOCKET_TURN_WALL_STARTED_AT["chat-running"] = 123.0 wth._WEBSOCKET_TURN_WALL_STARTED_AT["chat-running"] = 123.0
@@ -1795,6 +1883,13 @@ async def test_remote_webui_scope_allows_access_reduction(
payload = json.loads(conn.send.await_args.args[0]) payload = json.loads(conn.send.await_args.args[0])
assert payload["event"] == "session_updated" assert payload["event"] == "session_updated"
assert payload["workspace_scope"]["access_mode"] == "restricted" assert payload["workspace_scope"]["access_mode"] == "restricted"
assert sessions.list_sessions() == []
await channel._dispatch_envelope(
conn,
"webui-client",
{"type": "message", "chat_id": "chat-remote", "content": "hello", "webui": True},
)
saved = sessions.read_session_file("websocket:chat-remote") saved = sessions.read_session_file("websocket:chat-remote")
assert saved["metadata"]["workspace_scope"] == { assert saved["metadata"]["workspace_scope"] == {
"project_path": str(default_workspace.resolve()), "project_path": str(default_workspace.resolve()),
@@ -1864,8 +1959,10 @@ async def test_remote_access_reduction_rejects_stale_in_flight_message_scope(
release_hydrate.set() release_hydrate.set()
await message_task await message_task
saved = sessions.read_session_file(f"websocket:{chat_id}") assert sessions.read_session_file(f"websocket:{chat_id}") is None
assert saved["metadata"]["workspace_scope"]["access_mode"] == "restricted" assert channel._workspaces.scope_for_session_key(
f"websocket:{chat_id}"
).access_mode == "restricted"
payload = json.loads(message_conn.send.await_args.args[0]) payload = json.loads(message_conn.send.await_args.args[0])
assert payload["event"] == "error" assert payload["event"] == "error"
assert payload["detail"] == "workspace_scope_rejected" assert payload["detail"] == "workspace_scope_rejected"
@@ -1953,8 +2050,10 @@ async def test_native_webui_scope_allows_custom_scope_without_loopback(
assert payload["workspace_scope"]["restrict_to_workspace"] is False assert payload["workspace_scope"]["restrict_to_workspace"] is False
assert payload["workspace_scope"]["sandbox_status"]["restrict_to_workspace"] is False assert payload["workspace_scope"]["sandbox_status"]["restrict_to_workspace"] is False
assert payload["workspace_scope"]["sandbox_status"]["workspace_root"] == str(project.resolve()) assert payload["workspace_scope"]["sandbox_status"]["workspace_root"] == str(project.resolve())
saved = sessions.read_session_file("websocket:chat-native") assert sessions.read_session_file("websocket:chat-native") is None
assert saved["metadata"]["workspace_scope"] == { assert channel._workspaces.scope_for_session_key(
"websocket:chat-native"
).metadata() == {
"project_path": str(project.resolve()), "project_path": str(project.resolve()),
"access_mode": "full", "access_mode": "full",
} }
@@ -2092,16 +2191,12 @@ async def test_send_scopes_turn_model_updates_to_the_subscribed_chat() -> None:
def test_attach_fields_restore_the_session_model_and_latest_usage() -> None: def test_attach_fields_restore_the_session_model_and_latest_usage() -> None:
usage = LLMUsage.reported(input_tokens=120, output_tokens=8, total_tokens=175)
manager = MagicMock() manager = MagicMock()
manager.read_session_metadata.return_value = { manager.read_session_metadata.return_value = {
"metadata": { "metadata": {
SESSION_MODEL_PRESET_METADATA_KEY: "Deep Research", SESSION_MODEL_PRESET_METADATA_KEY: "Deep Research",
"_last_usage": { "_last_usage": usage.to_dict(),
"prompt_tokens": 120,
"completion_tokens": 8,
"negative": -1,
"boolean": True,
},
} }
} }
bus = MagicMock() bus = MagicMock()
@@ -2113,7 +2208,7 @@ def test_attach_fields_restore_the_session_model_and_latest_usage() -> None:
assert channel._attached_model_fields("chat-1") == { assert channel._attached_model_fields("chat-1") == {
"model_preset": "Deep Research", "model_preset": "Deep Research",
"usage": {"prompt_tokens": 120, "completion_tokens": 8}, "usage": usage.to_turn_dict(),
} }
@@ -2720,6 +2815,39 @@ async def test_send_turn_end_emits_turn_end_event() -> None:
] ]
@pytest.mark.asyncio
async def test_recovery_state_is_a_structured_event_not_assistant_text() -> None:
bus = MagicMock()
channel = WebSocketChannel(
{"enabled": True, "allowFrom": ["*"]},
bus,
gateway=_basic_handler(bus),
)
mock_ws = AsyncMock()
channel._attach(mock_ws, "chat-1")
await channel.send(OutboundMessage(
channel="websocket",
chat_id="chat-1",
content="",
event=RecoveryStateEvent(
status="awaiting_user",
recovery_id="recovery-1",
reason="tool_state_unknown",
attempts=1,
),
))
assert _sent_ws_payloads(mock_ws) == [{
"event": "recovery_state",
"chat_id": "chat-1",
"status": "awaiting_user",
"recovery_id": "recovery-1",
"reason": "tool_state_unknown",
"attempts": 1,
}]
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_system_command_turn_end_only_refreshes_session_metadata() -> None: async def test_system_command_turn_end_only_refreshes_session_metadata() -> None:
bus = MagicMock() bus = MagicMock()
@@ -3191,6 +3319,11 @@ async def test_send_turn_end_includes_latency_ms_when_present() -> None:
channel = WebSocketChannel({"enabled": True, "allowFrom": ["*"]}, bus, gateway=_basic_handler(bus)) channel = WebSocketChannel({"enabled": True, "allowFrom": ["*"]}, bus, gateway=_basic_handler(bus))
mock_ws = AsyncMock() mock_ws = AsyncMock()
channel._attach(mock_ws, "chat-1") channel._attach(mock_ws, "chat-1")
usage = LLMUsage.reported(
input_tokens=80,
output_tokens=20,
cache_read_tokens=40,
).with_timing(generation_ms=500, ttft_ms=125)
await channel.send(OutboundMessage( await channel.send(OutboundMessage(
channel="websocket", channel="websocket",
@@ -3198,7 +3331,7 @@ async def test_send_turn_end_includes_latency_ms_when_present() -> None:
content="", content="",
event=TurnEndEvent( event=TurnEndEvent(
latency_ms=1500, latency_ms=1500,
usage={"prompt_tokens": 80, "completion_tokens": 20, "cached_tokens": 40}, usage=usage,
context_window_tokens=128_000, context_window_tokens=128_000,
), ),
)) ))
@@ -3208,7 +3341,19 @@ async def test_send_turn_end_includes_latency_ms_when_present() -> None:
"event": "turn_end", "event": "turn_end",
"chat_id": "chat-1", "chat_id": "chat-1",
"latency_ms": 1500, "latency_ms": 1500,
"usage": {"prompt_tokens": 80, "completion_tokens": 20, "cached_tokens": 40}, "usage": {
"prompt_tokens": 80,
"completion_tokens": 20,
"total_tokens": 100,
"context_tokens": 80,
"cached_tokens": 40,
"request_count": 1,
"estimated_tokens": 0,
"generation_ms": 500,
"measured_completion_tokens": 20,
"ttft_ms": 125,
"timed_requests": 1,
},
"context_window_tokens": 128_000, "context_window_tokens": 128_000,
}, },
{"event": "session_updated", "chat_id": "chat-1", "scope": "thread"}, {"event": "session_updated", "chat_id": "chat-1", "scope": "thread"},
@@ -5175,10 +5320,16 @@ async def test_handle_session_context_get_reads_detached_session() -> None:
from nanobot.session import Session from nanobot.session import Session
usage = LLMUsage.reported(
input_tokens=12,
output_tokens=3,
total_tokens=175,
cache_read_tokens=6,
).with_timing(generation_ms=300, ttft_ms=45)
session = Session( session = Session(
key="websocket:context-route", key="websocket:context-route",
messages=[{"role": "user", "content": "hello"}], messages=[{"role": "user", "content": "hello"}],
metadata={"_last_usage": {"prompt_tokens": 12, "completion_tokens": 3}}, metadata={"_last_usage": usage.to_dict()},
) )
manager = MagicMock() manager = MagicMock()
manager.read_session_snapshot.return_value = session manager.read_session_snapshot.return_value = session
@@ -5195,7 +5346,19 @@ async def test_handle_session_context_get_reads_detached_session() -> None:
assert response.status_code == 200 assert response.status_code == 200
body = json.loads(response.body.decode()) body = json.loads(response.body.decode())
assert body["replay_messages"] == 1 assert body["replay_messages"] == 1
assert body["last_usage"] == {"prompt_tokens": 12, "completion_tokens": 3} assert body["last_usage"] == {
"prompt_tokens": 12,
"completion_tokens": 3,
"total_tokens": 175,
"context_tokens": 12,
"cached_tokens": 6,
"request_count": 1,
"estimated_tokens": 0,
"generation_ms": 300,
"measured_completion_tokens": 3,
"ttft_ms": 45,
"timed_requests": 1,
}
manager.read_session_snapshot.assert_called_once_with(session.key) manager.read_session_snapshot.assert_called_once_with(session.key)
@@ -4,6 +4,7 @@ import asyncio
import json import json
import random import random
import socket import socket
import threading
import time import time
from contextlib import suppress from contextlib import suppress
from pathlib import Path from pathlib import Path
@@ -83,6 +84,7 @@ def _make_handler(
channel_feature_action: Any | None = None, channel_feature_action: Any | None = None,
channel_runtime_status: Any | None = None, channel_runtime_status: Any | None = None,
mcp_reload: Any | None = None, mcp_reload: Any | None = None,
recovery_action: Any | None = None,
) -> GatewayServices: ) -> GatewayServices:
config = WebSocketConfig.model_validate(cfg) if isinstance(cfg, dict) else cfg config = WebSocketConfig.model_validate(cfg) if isinstance(cfg, dict) else cfg
workspace = workspace_path or Path.cwd() workspace = workspace_path or Path.cwd()
@@ -103,6 +105,7 @@ def _make_handler(
channel_feature_action=channel_feature_action, channel_feature_action=channel_feature_action,
channel_runtime_status=channel_runtime_status, channel_runtime_status=channel_runtime_status,
mcp_reload=mcp_reload, mcp_reload=mcp_reload,
recovery_action=recovery_action,
) )
@@ -121,6 +124,7 @@ def _ch(
channel_feature_action: Any | None = None, channel_feature_action: Any | None = None,
channel_runtime_status: Any | None = None, channel_runtime_status: Any | None = None,
mcp_reload: Any | None = None, mcp_reload: Any | None = None,
recovery_action: Any | None = None,
**extra: Any, **extra: Any,
) -> WebSocketChannel: ) -> WebSocketChannel:
cfg: dict[str, Any] = { cfg: dict[str, Any] = {
@@ -145,6 +149,7 @@ def _ch(
channel_feature_action=channel_feature_action, channel_feature_action=channel_feature_action,
channel_runtime_status=channel_runtime_status, channel_runtime_status=channel_runtime_status,
mcp_reload=mcp_reload, mcp_reload=mcp_reload,
recovery_action=recovery_action,
) )
return InProcessHttpChannel(cfg, bus, gateway=gateway) return InProcessHttpChannel(cfg, bus, gateway=gateway)
@@ -2572,6 +2577,69 @@ async def test_webui_automations_route_lists_all_jobs_and_allows_user_actions(
await server_task await server_task
@pytest.mark.asyncio
async def test_webui_cron_update_rearms_started_service_on_owner_loop(
bus: MagicMock,
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
store_path = tmp_path / "cron" / "jobs.json"
cron = CronService(store_path, max_sleep_ms=60_000)
job = cron.add_job(
name="Before update",
schedule=CronSchedule(kind="every", every_ms=86_400_000),
message="Check the repo status",
session_key="websocket:abc",
origin_channel="websocket",
origin_chat_id="abc",
)
await cron.start()
owner_thread_id = threading.get_ident()
initial_timer = cron._timer_task
request_thread_ids: list[int] = []
arm_thread_ids: list[int] = []
timer_rearmed = asyncio.Event()
original_request_timer_rearm = cron._request_timer_rearm
original_arm_timer = cron._arm_timer
def tracked_request_timer_rearm() -> None:
request_thread_ids.append(threading.get_ident())
original_request_timer_rearm()
def tracked_arm_timer() -> None:
arm_thread_ids.append(threading.get_ident())
original_arm_timer()
timer_rearmed.set()
monkeypatch.setattr(cron, "_request_timer_rearm", tracked_request_timer_rearm)
monkeypatch.setattr(cron, "_arm_timer", tracked_arm_timer)
channel = _ch(bus, cron_service=cron, port=_free_port())
try:
response = await _webui_mutate(
channel,
"automation.update",
{"id": job.id, "values": {"name": "After update"}},
)
await asyncio.wait_for(timer_rearmed.wait(), timeout=1)
assert response.status_code == 200
assert request_thread_ids
assert all(thread_id != owner_thread_id for thread_id in request_thread_ids)
assert arm_thread_ids and set(arm_thread_ids) == {owner_thread_id}
assert cron._timer_task is not None
assert cron._timer_task is not initial_timer
assert not cron._timer_task.done()
stored = json.loads(store_path.read_text(encoding="utf-8"))
assert len(stored["jobs"]) == 1
assert stored["jobs"][0]["name"] == "After update"
finally:
cron.stop()
await asyncio.sleep(0)
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_webui_automations_route_manages_local_triggers( async def test_webui_automations_route_manages_local_triggers(
bus: MagicMock, tmp_path: Path bus: MagicMock, tmp_path: Path
@@ -3242,6 +3310,28 @@ async def _webui_mutate(
) )
@pytest.mark.asyncio
async def test_recovery_mutation_uses_authenticated_websocket_action(bus: MagicMock) -> None:
recovery_action = AsyncMock(return_value={
"status": "resuming",
"recovery_id": "recovery-1",
})
channel = _ch(bus, recovery_action=recovery_action)
response = await _webui_mutate(
channel,
"recovery.continue",
{"chat_id": "chat-1", "recovery_id": "recovery-1"},
)
assert response.status_code == 200
assert response.json()["status"] == "resuming"
recovery_action.assert_awaited_once_with(
"continue",
{"chat_id": "chat-1", "recovery_id": "recovery-1"},
)
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_workspace_folder_picker_is_local_authenticated_mutation( async def test_workspace_folder_picker_is_local_authenticated_mutation(
bus: MagicMock, bus: MagicMock,
@@ -3743,3 +3833,77 @@ def test_bootstrap_secret_also_enforced_on_localhost(bus: MagicMock) -> None:
channel = _ch(bus, host="0.0.0.0", tokenIssueSecret="s3cret") channel = _ch(bus, host="0.0.0.0", tokenIssueSecret="s3cret")
resp = channel.gateway.http._handle_bootstrap(_LOCAL, _NO_HEADERS) resp = channel.gateway.http._handle_bootstrap(_LOCAL, _NO_HEADERS)
assert resp.status_code == 401 assert resp.status_code == 401
@pytest.mark.asyncio
async def test_webui_skill_update_cancellation_waits_for_config_and_runtime_state(
bus: MagicMock,
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
from nanobot.webui import ws_http
skill_dir = tmp_path / "skills" / "cancel-safe-skill"
skill_dir.mkdir(parents=True)
(skill_dir / "SKILL.md").write_text(
"---\nname: cancel-safe-skill\ndescription: Cancellation test skill.\n---\n",
encoding="utf-8",
)
mutation_started = threading.Event()
release_mutation = threading.Event()
original_update = ws_http.set_webui_skill_enabled
update_calls = 0
def blocked_update(*args: Any, **kwargs: Any) -> dict[str, Any]:
nonlocal update_calls
update_calls += 1
mutation_started.set()
assert release_mutation.wait(timeout=1)
return original_update(*args, **kwargs)
monkeypatch.setattr(ws_http, "set_webui_skill_enabled", blocked_update)
channel = _ch(
bus,
session_manager=_seed_session(tmp_path),
workspace_path=tmp_path,
port=_free_port(),
)
runtime_states: list[set[str]] = []
channel.gateway.http.skill_state_action = runtime_states.append
task = asyncio.create_task(
_webui_mutate(
channel,
"skill.update",
{"name": "cancel-safe-skill", "enabled": False},
)
)
assert await asyncio.to_thread(mutation_started.wait, 1)
try:
task.cancel()
await asyncio.sleep(0)
assert not task.done()
assert runtime_states == []
finally:
release_mutation.set()
with pytest.raises(asyncio.CancelledError):
await asyncio.wait_for(task, timeout=1)
assert update_calls == 1
assert "cancel-safe-skill" in channel.gateway.http.disabled_skills
assert runtime_states == [{"cancel-safe-skill"}]
saved = load_config(channel.gateway.settings.config.path)
assert "cancel-safe-skill" in saved.agents.defaults.disabled_skills
settled_state = (
update_calls,
set(channel.gateway.http.disabled_skills),
list(runtime_states),
)
await asyncio.sleep(0.05)
assert (
update_calls,
set(channel.gateway.http.disabled_skills),
runtime_states,
) == settled_state
+109 -28
View File
@@ -34,6 +34,7 @@ from nanobot.config.paths import is_default_workspace
from nanobot.config.schema import Config from nanobot.config.schema import Config
from nanobot.gateway.runtime import GatewayInstance from nanobot.gateway.runtime import GatewayInstance
from nanobot.security.network import is_loopback_host from nanobot.security.network import is_loopback_host
from nanobot.session.async_compat import call_session_manager as _call_session_manager
from nanobot.session.keys import UNIFIED_SESSION_KEY, last_channel_from_metadata from nanobot.session.keys import UNIFIED_SESSION_KEY, last_channel_from_metadata
from nanobot.utils.evaluator import evaluate_response, resolve_evaluator_prompt from nanobot.utils.evaluator import evaluate_response, resolve_evaluator_prompt
from nanobot.utils.helpers import sync_workspace_templates from nanobot.utils.helpers import sync_workspace_templates
@@ -45,6 +46,30 @@ __all__ = ["_run_gateway"]
console = Console() console = Console()
_EVENT_LOOP_LAG_INTERVAL_S = 0.5
_EVENT_LOOP_LAG_WARNING_S = 0.25
async def _monitor_event_loop_lag(
*,
interval_s: float = _EVENT_LOOP_LAG_INTERVAL_S,
warning_threshold_s: float = _EVENT_LOOP_LAG_WARNING_S,
log: Any | None = None,
) -> None:
"""Log scheduler drift so gateway-wide stalls have direct evidence."""
loop = asyncio.get_running_loop()
lag_log = log or logger
while True:
expected = loop.time() + interval_s
await asyncio.sleep(interval_s)
lag_s = max(0.0, loop.time() - expected)
if lag_s >= warning_threshold_s:
lag_log.warning(
"event loop lag operation=gateway duration_ms={} interval_ms={}",
int(lag_s * 1000),
int(interval_s * 1000),
)
def _http_endpoint_responding(url: str, *, timeout_s: float = 0.25) -> bool: def _http_endpoint_responding(url: str, *, timeout_s: float = 0.25) -> bool:
"""Return whether an HTTP endpoint responds, including with an auth error.""" """Return whether an HTTP endpoint responds, including with an auth error."""
@@ -313,6 +338,8 @@ def _run_gateway(
from nanobot.cron.service import CronJobSkippedError, CronService from nanobot.cron.service import CronJobSkippedError, CronService
from nanobot.cron.session_turns import is_bound_cron_job from nanobot.cron.session_turns import is_bound_cron_job
from nanobot.cron.types import CronJob from nanobot.cron.types import CronJob
from nanobot.llm_usage import record_llm_call
from nanobot.llm_usage.context import llm_usage_source
from nanobot.providers.factory import ( from nanobot.providers.factory import (
ProviderSnapshot, ProviderSnapshot,
build_provider_snapshot, build_provider_snapshot,
@@ -322,6 +349,7 @@ def _run_gateway(
from nanobot.providers.fallback_provider import FallbackProvider from nanobot.providers.fallback_provider import FallbackProvider
from nanobot.providers.image_generation import image_gen_provider_configs from nanobot.providers.image_generation import image_gen_provider_configs
from nanobot.session.manager import SessionManager from nanobot.session.manager import SessionManager
from nanobot.session.recovery import RecoveryCoordinator
from nanobot.session.webui_turns import ( from nanobot.session.webui_turns import (
WebuiTurnCoordinator, WebuiTurnCoordinator,
WebuiTurnRoutePolicy, WebuiTurnRoutePolicy,
@@ -329,7 +357,6 @@ def _run_gateway(
) )
from nanobot.triggers.local_runner import run_local_trigger_queue from nanobot.triggers.local_runner import run_local_trigger_queue
from nanobot.triggers.local_store import LocalTriggerStore from nanobot.triggers.local_store import LocalTriggerStore
from nanobot.webui.token_usage import TokenUsageHook
port = port if port is not None else config.gateway.port port = port if port is not None else config.gateway.port
webui_url = _webui_browser_url(config) webui_url = _webui_browser_url(config)
@@ -360,7 +387,8 @@ def _run_gateway(
runtime_events = RuntimeEventBus() runtime_events = RuntimeEventBus()
fallback_model_observer = build_webui_fallback_model_observer(bus) fallback_model_observer = build_webui_fallback_model_observer(bus)
def _observe_fallback_models(snapshot: ProviderSnapshot) -> ProviderSnapshot: def _observe_provider(snapshot: ProviderSnapshot) -> ProviderSnapshot:
snapshot.provider.set_llm_call_observer(record_llm_call)
if isinstance(snapshot.provider, FallbackProvider): if isinstance(snapshot.provider, FallbackProvider):
snapshot.provider.set_fallback_model_observer(fallback_model_observer) snapshot.provider.set_fallback_model_observer(fallback_model_observer)
return snapshot return snapshot
@@ -370,20 +398,19 @@ def _run_gateway(
**kwargs: Any, **kwargs: Any,
) -> ProviderSnapshot: ) -> ProviderSnapshot:
try: try:
return _observe_fallback_models(load_provider_snapshot(*args, **kwargs)) return _observe_provider(load_provider_snapshot(*args, **kwargs))
except ValueError as exc: except ValueError as exc:
if unconfigured_provider_error is None: if unconfigured_provider_error is None:
raise raise
return build_unconfigured_provider_snapshot(config, str(exc)) return _observe_provider(build_unconfigured_provider_snapshot(config, str(exc)))
if unconfigured_provider_error is not None: if unconfigured_provider_error is not None:
provider_snapshot = build_unconfigured_provider_snapshot( provider_snapshot = _observe_provider(
config, build_unconfigured_provider_snapshot(config, unconfigured_provider_error)
unconfigured_provider_error,
) )
else: else:
try: try:
provider_snapshot = _observe_fallback_models(build_provider_snapshot(config)) provider_snapshot = _observe_provider(build_provider_snapshot(config))
except ValueError as exc: except ValueError as exc:
console.print(f"[red]Error: {exc}[/red]") console.print(f"[red]Error: {exc}[/red]")
raise typer.Exit(1) from exc raise typer.Exit(1) from exc
@@ -422,6 +449,12 @@ def _run_gateway(
tools = ToolRegistry() tools = ToolRegistry()
mcp_provider = MCPProvider.from_config(config, tools) mcp_provider = MCPProvider.from_config(config, tools)
recovery = RecoveryCoordinator(
sessions=session_manager,
bus=bus,
unified_session=config.agents.defaults.unified_session,
)
# Create agent with cron service # Create agent with cron service
agent = AgentLoop.from_config( agent = AgentLoop.from_config(
config, bus, config, bus,
@@ -436,10 +469,10 @@ def _run_gateway(
runtime_events=runtime_events, runtime_events=runtime_events,
turn_delivery_factory=turn_delivery_factory, turn_delivery_factory=turn_delivery_factory,
provider_signature=provider_snapshot.signature, provider_signature=provider_snapshot.signature,
hooks=[TokenUsageHook(timezone_name=config.agents.defaults.timezone)],
local_trigger_store=trigger_store, local_trigger_store=trigger_store,
hook_factories=[create_file_edit_activity_hook], hook_factories=[create_file_edit_activity_hook],
tool_registry=tools, tool_registry=tools,
recovery_admission=recovery,
) )
def _schedule_webui_background(awaitable: Awaitable[None]) -> None: def _schedule_webui_background(awaitable: Awaitable[None]) -> None:
agent.schedule_background(cast(Coroutine[Any, Any, None], awaitable)) agent.schedule_background(cast(Coroutine[Any, Any, None], awaitable))
@@ -448,6 +481,7 @@ def _run_gateway(
bus=bus, bus=bus,
sessions=session_manager, sessions=session_manager,
schedule_background=_schedule_webui_background, schedule_background=_schedule_webui_background,
recovery=recovery,
) )
webui_turn_coordinator.subscribe(runtime_events) webui_turn_coordinator.subscribe(runtime_events)
from nanobot.bus.events import OutboundMessage from nanobot.bus.events import OutboundMessage
@@ -484,12 +518,22 @@ def _run_gateway(
and hasattr(session_manager, "save") and hasattr(session_manager, "save")
): ):
key = session_key or _channel_session_key(msg.channel, msg.chat_id) key = session_key or _channel_session_key(msg.channel, msg.chat_id)
session = session_manager.get_or_create(key) session = await _call_session_manager(
session_manager,
"get_or_create_async",
session_manager.get_or_create,
key,
)
extra: dict[str, Any] = {"_channel_delivery": True} extra: dict[str, Any] = {"_channel_delivery": True}
if msg.media: if msg.media:
extra["media"] = list(msg.media) extra["media"] = list(msg.media)
session.add_message("assistant", msg.content, **extra) session.add_message("assistant", msg.content, **extra)
session_manager.save(session) await _call_session_manager(
session_manager,
"save_async",
session_manager.save,
session,
)
await bus.publish_outbound(msg) await bus.publish_outbound(msg)
message_tool = agent.tools.get("message") message_tool = agent.tools.get("message")
@@ -555,25 +599,18 @@ def _run_gateway(
except Exception: except Exception:
logger.exception("Dream cron job failed") logger.exception("Dream cron job failed")
finally: finally:
from nanobot.webui.token_usage import record_response_token_usage
record_response_token_usage(
resp,
source="dream",
timezone_name=config.agents.defaults.timezone,
)
sha = _commit_dream_changes(store) sha = _commit_dream_changes(store)
if sha: if sha:
logger.info("Dream commit: {}", sha) logger.info("Dream commit: {}", sha)
store.compact_history() store.compact_history()
prune_dream_sessions(agent.sessions) await asyncio.to_thread(prune_dream_sessions, agent.sessions)
return None return None
# Heartbeat is a system job that checks HEARTBEAT.md for active tasks. # Heartbeat is a system job that checks HEARTBEAT.md for active tasks.
if job.name == "heartbeat": if job.name == "heartbeat":
heartbeat_file = config.workspace_path / "HEARTBEAT.md" heartbeat_file = config.workspace_path / "HEARTBEAT.md"
try: try:
content = heartbeat_file.read_text(encoding="utf-8") content = await asyncio.to_thread(heartbeat_file.read_text, encoding="utf-8")
except OSError: except OSError:
logger.debug("Heartbeat: HEARTBEAT.md missing") logger.debug("Heartbeat: HEARTBEAT.md missing")
return None return None
@@ -581,7 +618,7 @@ def _run_gateway(
logger.debug("Heartbeat: HEARTBEAT.md has no active tasks") logger.debug("Heartbeat: HEARTBEAT.md has no active tasks")
return None return None
channel, chat_id = _pick_heartbeat_target() channel, chat_id = await _pick_heartbeat_target()
if channel == "cli": if channel == "cli":
return None return None
@@ -609,9 +646,19 @@ def _run_gateway(
message_tool.reset_suppress_delivery(suppress_token) message_tool.reset_suppress_delivery(suppress_token)
# Keep a small tail of heartbeat history so the loop stays bounded. # Keep a small tail of heartbeat history so the loop stays bounded.
session = agent.sessions.get_or_create("heartbeat") session = await _call_session_manager(
agent.sessions,
"get_or_create_async",
agent.sessions.get_or_create,
"heartbeat",
)
session.retain_recent_legal_suffix(hb_cfg.keep_recent_messages) session.retain_recent_legal_suffix(hb_cfg.keep_recent_messages)
agent.sessions.save(session) await _call_session_manager(
agent.sessions,
"save_async",
agent.sessions.save,
session,
)
if not resp or not resp.content: if not resp or not resp.content:
return return
@@ -621,6 +668,7 @@ def _run_gateway(
evaluator_prompt = resolve_evaluator_prompt(config.workspace_path) evaluator_prompt = resolve_evaluator_prompt(config.workspace_path)
# Fail closed: stay silent on evaluator failure instead of notifying. # Fail closed: stay silent on evaluator failure instead of notifying.
with llm_usage_source("cron"):
should_notify = await evaluate_response( should_notify = await evaluate_response(
response=response, response=response,
task_context=prompt, task_context=prompt,
@@ -683,20 +731,31 @@ def _run_gateway(
webui_mcp_runtime_status=mcp_provider.runtime_status, webui_mcp_runtime_status=mcp_provider.runtime_status,
webui_mcp_reload=mcp_provider.reload, webui_mcp_reload=mcp_provider.reload,
webui_skill_state_action=_webui_skill_state_action, webui_skill_state_action=_webui_skill_state_action,
webui_recovery_action=recovery.handle_action,
config_path=Path(config_path), config_path=Path(config_path),
) )
def _pick_heartbeat_target() -> tuple[str, str]: async def _pick_heartbeat_target() -> tuple[str, str]:
"""Pick a routable channel/chat target for heartbeat-triggered messages.""" """Pick a routable channel/chat target for heartbeat-triggered messages."""
sidebar_state = read_webui_sidebar_state() sidebar_state = await asyncio.to_thread(read_webui_sidebar_state)
unified_metadata = None unified_metadata = None
if config.agents.defaults.unified_session: if config.agents.defaults.unified_session:
record = session_manager.read_session_metadata(UNIFIED_SESSION_KEY) record = await _call_session_manager(
session_manager,
"read_session_metadata_async",
session_manager.read_session_metadata,
UNIFIED_SESSION_KEY,
)
if isinstance(record, dict) and isinstance(record.get("metadata"), dict): if isinstance(record, dict) and isinstance(record.get("metadata"), dict):
unified_metadata = record["metadata"] unified_metadata = record["metadata"]
sessions = await _call_session_manager(
session_manager,
"list_sessions_async",
session_manager.list_sessions,
)
return _pick_heartbeat_target_from_sessions( return _pick_heartbeat_target_from_sessions(
enabled_channels=channels.enabled_channels, enabled_channels=channels.enabled_channels,
sessions=session_manager.list_sessions(), sessions=sessions,
archived_keys=sidebar_state.get("archived_keys", []), archived_keys=sidebar_state.get("archived_keys", []),
unified_session_metadata=unified_metadata, unified_session_metadata=unified_metadata,
) )
@@ -849,6 +908,7 @@ def _run_gateway(
tasks: list[asyncio.Task[Any]] = [] tasks: list[asyncio.Task[Any]] = []
shutdown_task: asyncio.Task[Any] | None = None shutdown_task: asyncio.Task[Any] | None = None
runtime_tasks: asyncio.Future[list[Any]] | None = None runtime_tasks: asyncio.Future[list[Any]] | None = None
startup_complete = False
shutdown_event = asyncio.Event() shutdown_event = asyncio.Event()
cli_terminal._ensure_interactive_tty_mode() cli_terminal._ensure_interactive_tty_mode()
restore_shutdown_handlers = _install_gateway_shutdown_handlers( restore_shutdown_handlers = _install_gateway_shutdown_handlers(
@@ -861,6 +921,10 @@ def _run_gateway(
await cron.start() await cron.start()
# Re-read once on first admission to close the watcher subscription window. # Re-read once on first admission to close the watcher subscription window.
agent.runtime_resolver.invalidate() agent.runtime_resolver.invalidate()
# Recovery must finish before WebSocket and other channels begin
# accepting new input. That makes a new user message reliably
# supersede an old recoverable turn instead of racing its queue.
await recovery.scan()
async def _run_agent() -> None: async def _run_agent() -> None:
try: try:
await mcp_provider.connect() await mcp_provider.connect()
@@ -898,6 +962,10 @@ def _run_gateway(
_monitor_local_clients(), _monitor_local_clients(),
name="nanobot-gateway-client-monitor", name="nanobot-gateway-client-monitor",
), ),
asyncio.create_task(
_monitor_event_loop_lag(),
name="nanobot-event-loop-lag-monitor",
),
] ]
if health_server_enabled: if health_server_enabled:
tasks.append(asyncio.create_task( tasks.append(asyncio.create_task(
@@ -915,6 +983,7 @@ def _run_gateway(
name="nanobot-webui-dev-server", name="nanobot-webui-dev-server",
)) ))
runtime_tasks = asyncio.gather(*tasks) runtime_tasks = asyncio.gather(*tasks)
startup_complete = True
shutdown_task = asyncio.create_task( shutdown_task = asyncio.create_task(
shutdown_event.wait(), shutdown_event.wait(),
name="nanobot-gateway-shutdown", name="nanobot-gateway-shutdown",
@@ -936,6 +1005,10 @@ def _run_gateway(
console.print("\n[red]Error: Gateway crashed unexpectedly[/red]") console.print("\n[red]Error: Gateway crashed unexpectedly[/red]")
console.print(traceback.format_exc()) console.print(traceback.format_exc())
if not startup_complete:
# Do not report a successful gateway command when startup
# failed before any runtime task or listener was created.
raise typer.Exit(1)
finally: finally:
try: try:
if shutdown_task and not shutdown_task.done(): if shutdown_task and not shutdown_task.done():
@@ -943,6 +1016,10 @@ def _run_gateway(
with suppress(asyncio.CancelledError): with suppress(asyncio.CancelledError):
await shutdown_task await shutdown_task
cron.stop() cron.stop()
# A gateway exit interrupts ownership of active turns; it is
# not the same as the user stopping a turn. Keep checkpoints
# so the next gateway can offer an explicit Continue action.
agent.preserve_inflight_turns_on_shutdown()
agent.stop() agent.stop()
# Cancel runtime tasks first, then deterministically close # Cancel runtime tasks first, then deterministically close
# exec/MCP resources while the event loop is still alive. # exec/MCP resources while the event loop is still alive.
@@ -956,7 +1033,11 @@ def _run_gateway(
# Flush all cached sessions to durable storage before exit. # Flush all cached sessions to durable storage before exit.
# This prevents data loss on filesystems with write-back # This prevents data loss on filesystems with write-back
# caching (rclone VFS, NFS, FUSE mounts, etc.). # caching (rclone VFS, NFS, FUSE mounts, etc.).
flushed = agent.sessions.flush_all() flushed = await _call_session_manager(
agent.sessions,
"flush_all_async",
agent.sessions.flush_all,
)
if flushed: if flushed:
logger.info("Shutdown: flushed {} session(s) to disk", flushed) logger.info("Shutdown: flushed {} session(s) to disk", flushed)
finally: finally:
+63 -16
View File
@@ -3,6 +3,7 @@
from __future__ import annotations from __future__ import annotations
import asyncio import asyncio
import inspect
import os import os
import subprocess import subprocess
import sys import sys
@@ -14,6 +15,8 @@ from typing import TYPE_CHECKING, Any, Literal, cast
from nanobot import __version__ from nanobot import __version__
from nanobot.bus.events import INBOUND_META_USER_SHELL, OutboundMessage from nanobot.bus.events import INBOUND_META_USER_SHELL, OutboundMessage
from nanobot.command.router import CommandContext, CommandRouter, normalize_command_text from nanobot.command.router import CommandContext, CommandRouter, normalize_command_text
from nanobot.session.async_compat import call_session_manager
from nanobot.utils.cancellation import shield_and_drain
from nanobot.utils.helpers import build_status_content from nanobot.utils.helpers import build_status_content
from nanobot.utils.restart import set_restart_notice_to_env from nanobot.utils.restart import set_restart_notice_to_env
from nanobot.utils.workspace_prompts import initialize_workspace_prompt from nanobot.utils.workspace_prompts import initialize_workspace_prompt
@@ -22,6 +25,7 @@ if TYPE_CHECKING:
from nanobot.agent.loop import AgentLoop from nanobot.agent.loop import AgentLoop
from nanobot.session.manager import Session from nanobot.session.manager import Session
from nanobot.utils.gitstore import CommitInfo from nanobot.utils.gitstore import CommitInfo
from nanobot.utils.llm_runtime import LLMRuntime
# WebUI protocol contract for how a slash command participates in turn state: # WebUI protocol contract for how a slash command participates in turn state:
# - side_channel: returns control text without starting or ending an agent turn. # - side_channel: returns control text without starting or ending an agent turn.
@@ -201,6 +205,52 @@ def builtin_command_starts_agent_turn(text: str) -> bool:
return spec.lifecycle == "agent_turn_with_args" and bool(args.strip()) return spec.lifecycle == "agent_turn_with_args" and bool(args.strip())
def _has_native_coroutine_method(target: object, name: str) -> bool:
"""Check the real target class without trusting dynamic mock attributes."""
method = inspect.getattr_static(type(target), name, None)
return inspect.iscoroutinefunction(method)
async def _get_or_create_session(loop: AgentLoop, key: str) -> Session:
sessions = loop.sessions
return await call_session_manager(
sessions,
"get_or_create_async",
sessions.get_or_create,
key,
)
async def _save_session(loop: AgentLoop, session: Session) -> None:
sessions = loop.sessions
await call_session_manager(
sessions,
"save_async",
sessions.save,
session,
)
async def _runtime_for_session(loop: AgentLoop, session: Session) -> LLMRuntime:
if _has_native_coroutine_method(loop, "runtime_for_session_async"):
return await loop.runtime_for_session_async(session)
return await shield_and_drain(
asyncio.to_thread(loop.runtime_for_session, session)
)
async def _set_session_model_preset(
loop: AgentLoop,
session_key: str,
name: str,
) -> LLMRuntime:
if _has_native_coroutine_method(loop, "set_session_model_preset_async"):
return await loop.set_session_model_preset_async(session_key, name)
return await shield_and_drain(
asyncio.to_thread(loop.set_session_model_preset, session_key, name)
)
async def cmd_stop(ctx: CommandContext) -> OutboundMessage: async def cmd_stop(ctx: CommandContext) -> OutboundMessage:
"""Cancel all active tasks and subagents for the session.""" """Cancel all active tasks and subagents for the session."""
loop = ctx.loop loop = ctx.loop
@@ -257,8 +307,8 @@ async def cmd_restart(ctx: CommandContext) -> OutboundMessage:
async def cmd_status(ctx: CommandContext) -> OutboundMessage: async def cmd_status(ctx: CommandContext) -> OutboundMessage:
"""Build an outbound status message for a session.""" """Build an outbound status message for a session."""
loop = ctx.loop loop = ctx.loop
session = ctx.session or loop.sessions.get_or_create(ctx.key) session = ctx.session or await _get_or_create_session(loop, ctx.key)
runtime = ctx.runtime or loop.runtime_for_session(session) runtime = ctx.runtime or await _runtime_for_session(loop, session)
ctx_est = 0 ctx_est = 0
with suppress(Exception): with suppress(Exception):
ctx_est, _ = loop.consolidator.estimate_session_prompt_tokens( ctx_est, _ = loop.consolidator.estimate_session_prompt_tokens(
@@ -266,7 +316,8 @@ async def cmd_status(ctx: CommandContext) -> OutboundMessage:
runtime=runtime, runtime=runtime,
) )
if ctx_est <= 0: if ctx_est <= 0:
ctx_est = loop._last_usage.get("prompt_tokens", 0) # pyright: ignore[reportPrivateUsage] last_usage = loop._last_usage # pyright: ignore[reportPrivateUsage]
ctx_est = last_usage.input_tokens if last_usage is not None else 0
# Fetch web search provider usage (best-effort, never blocks the response) # Fetch web search provider usage (best-effort, never blocks the response)
search_usage_text: str | None = None search_usage_text: str | None = None
@@ -305,20 +356,21 @@ async def cmd_new(ctx: CommandContext) -> OutboundMessage:
loop = ctx.loop loop = ctx.loop
await loop._cancel_active_tasks(ctx.key) # pyright: ignore[reportPrivateUsage] await loop._cancel_active_tasks(ctx.key) # pyright: ignore[reportPrivateUsage]
loop.discard_session_file_state(ctx.key) loop.discard_session_file_state(ctx.key)
session = ctx.session or loop.sessions.get_or_create(ctx.key) session = ctx.session or await _get_or_create_session(loop, ctx.key)
snapshot = list(session.messages) snapshot = list(session.messages)
archive_snapshot = None archive_snapshot = None
runtime = None runtime = None
if session.last_consolidated < len(snapshot): if session.last_consolidated < len(snapshot):
runtime = ctx.runtime or loop.runtime_for_session(session) runtime = ctx.runtime or await _runtime_for_session(loop, session)
archive_snapshot = replace( archive_snapshot = replace(
session, session,
messages=snapshot, messages=snapshot,
metadata=dict(session.metadata), metadata=dict(session.metadata),
provider_state=None, provider_state=None,
) )
async def reset_and_schedule_archive() -> None:
session.clear() session.clear()
loop.sessions.save(session) await _save_session(loop, session)
loop.sessions.invalidate(session.key) loop.sessions.invalidate(session.key)
if archive_snapshot is not None and runtime is not None: if archive_snapshot is not None and runtime is not None:
loop.schedule_background( loop.schedule_background(
@@ -328,6 +380,8 @@ async def cmd_new(ctx: CommandContext) -> OutboundMessage:
runtime=runtime, runtime=runtime,
) )
) )
await shield_and_drain(reset_and_schedule_archive())
return OutboundMessage( return OutboundMessage(
channel=ctx.msg.channel, chat_id=ctx.msg.chat_id, channel=ctx.msg.channel, chat_id=ctx.msg.chat_id,
content="New session started.", content="New session started.",
@@ -376,7 +430,7 @@ async def cmd_model(ctx: CommandContext) -> OutboundMessage:
metadata = {**dict(ctx.msg.metadata or {}), "render_as": "text"} metadata = {**dict(ctx.msg.metadata or {}), "render_as": "text"}
if not args: if not args:
session = ctx.session or loop.sessions.get_or_create(ctx.key) session = ctx.session or await _get_or_create_session(loop, ctx.key)
return OutboundMessage( return OutboundMessage(
channel=ctx.msg.channel, channel=ctx.msg.channel,
chat_id=ctx.msg.chat_id, chat_id=ctx.msg.chat_id,
@@ -386,7 +440,7 @@ async def cmd_model(ctx: CommandContext) -> OutboundMessage:
name = args name = args
try: try:
runtime = loop.set_session_model_preset(ctx.key, name) runtime = await _set_session_model_preset(loop, ctx.key, name)
except (KeyError, ValueError) as exc: except (KeyError, ValueError) as exc:
names = _model_preset_names(loop) names = _model_preset_names(loop)
return OutboundMessage( return OutboundMessage(
@@ -478,13 +532,6 @@ async def cmd_dream(ctx: CommandContext) -> OutboundMessage:
elapsed = time.monotonic() - t0 elapsed = time.monotonic() - t0
content = f"Dream failed after {elapsed:.1f}s: {e}" content = f"Dream failed after {elapsed:.1f}s: {e}"
finally: finally:
from nanobot.webui.token_usage import record_response_token_usage
record_response_token_usage(
resp,
source="dream",
timezone_name=getattr(loop.context, "timezone", None),
)
if store.git.is_initialized(): if store.git.is_initialized():
commit_msg = build_dream_commit_message("dream: manual run", diff_body) commit_msg = build_dream_commit_message("dream: manual run", diff_body)
sha = store.git.auto_commit(commit_msg) sha = store.git.auto_commit(commit_msg)
@@ -854,7 +901,7 @@ async def cmd_history(ctx: CommandContext) -> OutboundMessage:
metadata=dict(ctx.msg.metadata or {}), metadata=dict(ctx.msg.metadata or {}),
) )
session = ctx.session or ctx.loop.sessions.get_or_create(ctx.key) session = ctx.session or await _get_or_create_session(ctx.loop, ctx.key)
history = session.get_history(max_messages=0, include_runtime_context=False) history = session.get_history(max_messages=0, include_runtime_context=False)
visible = [_format_history_message(m) for m in history] visible = [_format_history_message(m) for m in history]
visible = [m for m in visible if m is not None] visible = [m for m in visible if m is not None]
+10
View File
@@ -8,6 +8,7 @@ import time
import uuid import uuid
from typing import TYPE_CHECKING, Any, Protocol from typing import TYPE_CHECKING, Any, Protocol
from nanobot.agent.automation_turns import AutomationTurnAcceptedCancellation
from nanobot.agent.tools.cron import CronTool from nanobot.agent.tools.cron import CronTool
from nanobot.bus.events import InboundMessage, OutboundMessage from nanobot.bus.events import InboundMessage, OutboundMessage
from nanobot.cron.session_delivery import origin_delivery_context from nanobot.cron.session_delivery import origin_delivery_context
@@ -127,6 +128,15 @@ async def run_bound_cron_job(
session_key_override=session_key, session_key_override=session_key,
) )
) )
except AutomationTurnAcceptedCancellation:
cron.write_run_record(
run_id,
{
**run_record_base,
"status": "accepted",
},
)
raise
except (Exception, asyncio.CancelledError) as exc: except (Exception, asyncio.CancelledError) as exc:
error_text = str(exc) or exc.__class__.__name__ error_text = str(exc) or exc.__class__.__name__
cron.write_run_record( cron.write_run_record(
+199 -56
View File
@@ -11,11 +11,12 @@ from dataclasses import asdict
from datetime import datetime from datetime import datetime
from pathlib import Path from pathlib import Path
from types import EllipsisType from types import EllipsisType
from typing import Any, Callable, Coroutine, Literal from typing import Any, Callable, Coroutine, Literal, TypeVar
from filelock import FileLock from filelock import FileLock
from loguru import logger from loguru import logger
from nanobot.agent.automation_turns import AutomationTurnAcceptedCancellation
from nanobot.cron.session_turns import is_bound_cron_job from nanobot.cron.session_turns import is_bound_cron_job
from nanobot.cron.types import ( from nanobot.cron.types import (
CronJob, CronJob,
@@ -25,10 +26,14 @@ from nanobot.cron.types import (
CronSchedule, CronSchedule,
CronStore, CronStore,
) )
from nanobot.utils.cancellation import shield_and_drain
from nanobot.utils.run_records import ( from nanobot.utils.run_records import (
write_run_record as write_automation_run_record, write_run_record as write_automation_run_record,
) )
_FILE_LOCK_TIMEOUT_SECONDS = 5
_T = TypeVar("_T")
class CronJobSkippedError(Exception): class CronJobSkippedError(Exception):
"""Raised by cron callbacks when a job was intentionally skipped.""" """Raised by cron callbacks when a job was intentionally skipped."""
@@ -164,10 +169,16 @@ class CronService:
self.store_path = store_path self.store_path = store_path
self._action_path = store_path.parent / "action.jsonl" self._action_path = store_path.parent / "action.jsonl"
self._run_records_dir = store_path.parent / "runs" self._run_records_dir = store_path.parent / "runs"
self._lock = FileLock(str(self._action_path.parent) + ".lock") self._lock = FileLock(
str(self._action_path.parent) + ".lock",
timeout=_FILE_LOCK_TIMEOUT_SECONDS,
)
self.on_job = on_job self.on_job = on_job
self._store: CronStore | None = None self._store: CronStore | None = None
self._timer_task: asyncio.Task[None] | None = None self._timer_task: asyncio.Task[None] | None = None
self._operation_lock = asyncio.Lock()
self._claimed_job_ids: set[str] = set()
self._event_loop: asyncio.AbstractEventLoop | None = None
self._running = False self._running = False
self._active_executions = 0 self._active_executions = 0
self._store_dirty = False self._store_dirty = False
@@ -451,25 +462,58 @@ class CronService:
"""Write an internal audit record for one cron execution.""" """Write an internal audit record for one cron execution."""
write_automation_run_record(self._run_records_dir, run_id, record) write_automation_run_record(self._run_records_dir, run_id, record)
async def run_sync(
self,
operation: Callable[..., _T],
/,
*args: Any,
**kwargs: Any,
) -> _T:
"""Serialize a complete cron transaction in a worker thread.
A running thread cannot be cancelled safely. Keep the transaction lock
until it exits so cancellation is never reported while that worker can
still mutate cron state behind a later operation.
"""
async with self._operation_lock:
return await shield_and_drain(
asyncio.to_thread(operation, *args, **kwargs)
)
async def start(self) -> None: async def start(self) -> None:
"""Start the cron service.""" """Start the cron service and settle accepted work before cancellation."""
async def settle_start() -> None:
self._event_loop = asyncio.get_running_loop()
self._running = True self._running = True
loaded = self._load_store() try:
async with self._operation_lock:
loaded = await asyncio.to_thread(self._load_store)
if loaded is None: if loaded is None:
# Store file existed but was corrupt and has been preserved with # Store file existed but was corrupt and has been preserved with
# a ``.corrupt-<ts>`` suffix. Bail out instead of starting with # a ``.corrupt-<ts>`` suffix. Bail out instead of starting with
# an empty store; that would call ``_save_store`` and overwrite # an empty store; that would call ``_save_store`` and overwrite
# the now-renamed (but still recoverable) data with []. # the now-renamed (but still recoverable) data with [].
self._running = False
raise RuntimeError( raise RuntimeError(
f"cron store at {self.store_path} is corrupt and was preserved; " f"cron store at {self.store_path} is corrupt and was preserved; "
"refusing to start with an empty job list. " "refusing to start with an empty job list. "
"Inspect the .corrupt-<ts> backup and restore manually." "Inspect the .corrupt-<ts> backup and restore manually."
) )
self._recompute_next_runs() self._recompute_next_runs()
self._save_store() await asyncio.to_thread(self._save_store)
self._arm_timer() self._arm_timer()
logger.info("Cron service started with {} jobs", len(self._store.jobs if self._store else [])) logger.info(
"Cron service started with {} jobs",
len(self._store.jobs if self._store else []),
)
except BaseException:
# A failed start must not retain ownership without a timer. Caller
# cancellation is shielded until this composite either reaches the
# fully started state above or rolls back here.
self.stop()
raise
await shield_and_drain(settle_start())
def stop(self) -> None: def stop(self) -> None:
"""Stop the cron service.""" """Stop the cron service."""
@@ -497,8 +541,22 @@ class CronService:
if j.enabled and j.state.next_run_at_ms] if j.enabled and j.state.next_run_at_ms]
return min(times) if times else None return min(times) if times else None
def _request_timer_rearm(self) -> None:
"""Re-arm on the owning event loop, including from persistence workers."""
if not self._running:
return
loop = self._event_loop
try:
current_loop = asyncio.get_running_loop()
except RuntimeError:
current_loop = None
if current_loop is loop:
self._arm_timer()
elif loop is not None and loop.is_running():
loop.call_soon_threadsafe(self._arm_timer)
def _arm_timer(self) -> None: def _arm_timer(self) -> None:
"""Schedule the next timer tick.""" """Schedule the next timer tick on the owning event loop."""
if self._timer_task: if self._timer_task:
self._timer_task.cancel() self._timer_task.cancel()
@@ -520,7 +578,7 @@ class CronService:
self._timer_task = asyncio.create_task(tick()) self._timer_task = asyncio.create_task(tick())
async def _on_timer(self) -> None: async def _on_timer(self) -> None:
"""Handle timer tick - run due jobs.""" """Run due jobs while keeping persistence transactions serialized."""
reload_store = self._active_executions == 0 reload_store = self._active_executions == 0
self._active_executions += 1 self._active_executions += 1
try: try:
@@ -528,11 +586,17 @@ class CronService:
# to persist their advanced schedule. Persist that exact snapshot # to persist their advanced schedule. Persist that exact snapshot
# before reloading or executing anything else; otherwise the older # before reloading or executing anything else; otherwise the older
# disk state can replay the same job. # disk state can replay the same job.
async with self._operation_lock:
if self._store_dirty: if self._store_dirty:
self._save_store() await shield_and_drain(asyncio.to_thread(self._save_store))
return return
store = self._load_store(reload_during_execution=reload_store) store = await shield_and_drain(
asyncio.to_thread(
self._load_store,
reload_during_execution=reload_store,
)
)
# If a hot reload found a corrupt store on disk, ``self._store`` # If a hot reload found a corrupt store on disk, ``self._store``
# may still hold the previous, known-good in-memory snapshot. # may still hold the previous, known-good in-memory snapshot.
if store is None: if store is None:
@@ -547,7 +611,8 @@ class CronService:
for job in due_jobs: for job in due_jobs:
await self._execute_job(job) await self._execute_job(job)
self._save_store() async with self._operation_lock:
await shield_and_drain(asyncio.to_thread(self._save_store))
except Exception: except Exception:
# A load/persist failure must not kill the scheduler: keep the # A load/persist failure must not kill the scheduler: keep the
# in-memory store and retry on the next tick. This mirrors the # in-memory store and retry on the next tick. This mirrors the
@@ -564,59 +629,125 @@ class CronService:
# single bad tick cannot silently stop all future jobs. # single bad tick cannot silently stop all future jobs.
self._arm_timer() self._arm_timer()
async def _execute_job(self, job: CronJob) -> None: async def _claim_job(self, job_id: str) -> bool:
"""Execute a single job.""" """Claim one job without serializing callbacks for different jobs."""
start_ms = _now_ms() async with self._operation_lock:
logger.info("Cron: executing job '{}' ({})", job.name, job.id) if job_id in self._claimed_job_ids:
return False
self._claimed_job_ids.add(job_id)
return True
try: async def _release_job_claim(self, job_id: str) -> None:
if self.on_job: async with self._operation_lock:
await self.on_job(job) self._claimed_job_ids.discard(job_id)
job.state.last_status = "ok"
job.state.last_error = None
logger.info("Cron: job '{}' completed", job.name)
except CronJobSkippedError as e:
job.state.last_status = "skipped"
job.state.last_error = str(e) or None
logger.warning("Cron: job '{}' skipped: {}", job.name, job.state.last_error or "")
except asyncio.CancelledError as e:
current = asyncio.current_task()
if current is not None and current.cancelling():
raise
job.state.last_status = "error"
job.state.last_error = str(e) or e.__class__.__name__
logger.exception("Cron: job '{}' was cancelled", job.name)
except Exception as e:
job.state.last_status = "error"
job.state.last_error = str(e)
logger.exception("Cron: job '{}' failed", job.name)
async def _settle_job_execution(
self,
job: CronJob,
*,
start_ms: int,
status: Literal["ok", "error", "skipped"],
error: str | None,
persist: bool = False,
) -> None:
end_ms = _now_ms() end_ms = _now_ms()
async with self._operation_lock:
job.state.last_status = status
job.state.last_error = error
job.state.last_run_at_ms = start_ms job.state.last_run_at_ms = start_ms
job.updated_at_ms = end_ms job.updated_at_ms = end_ms
job.state.run_history.append(CronRunRecord( job.state.run_history.append(CronRunRecord(
run_at_ms=start_ms, run_at_ms=start_ms,
status=job.state.last_status, status=status,
duration_ms=end_ms - start_ms, duration_ms=end_ms - start_ms,
error=job.state.last_error, error=error,
)) ))
job.state.run_history = job.state.run_history[-self._MAX_RUN_HISTORY:] job.state.run_history = job.state.run_history[-self._MAX_RUN_HISTORY:]
# Handle one-shot jobs
if job.schedule.kind == "at": if job.schedule.kind == "at":
if job.delete_after_run: if job.delete_after_run:
store = self._require_store() store = await shield_and_drain(
asyncio.to_thread(self._require_store)
)
store.jobs = [item for item in store.jobs if item.id != job.id] store.jobs = [item for item in store.jobs if item.id != job.id]
else: else:
job.enabled = False job.enabled = False
job.state.next_run_at_ms = None job.state.next_run_at_ms = None
else: else:
# Compute next run
job.state.next_run_at_ms = _compute_next_run(job.schedule, _now_ms()) job.state.next_run_at_ms = _compute_next_run(job.schedule, _now_ms())
if persist:
await shield_and_drain(asyncio.to_thread(self._save_store))
@staticmethod
async def _drain_settlement_on_cancellation(settlement: asyncio.Task[None]) -> None:
"""Finish a short durable settlement despite repeated cancellation."""
while not settlement.done():
try:
await asyncio.shield(settlement)
except asyncio.CancelledError:
continue
settlement.result()
async def _execute_job(self, job: CronJob) -> bool:
"""Execute a claimed job and serialize its in-memory settlement."""
if not await self._claim_job(job.id):
logger.info("Cron: job '{}' ({}) is already running", job.name, job.id)
return False
start_ms = _now_ms()
logger.info("Cron: executing job '{}' ({})", job.name, job.id)
status: Literal["ok", "error", "skipped"]
error: str | None
accepted_cancellation: AutomationTurnAcceptedCancellation | None = None
try:
try:
if self.on_job:
await self.on_job(job)
status = "ok"
error = None
logger.info("Cron: job '{}' completed", job.name)
except AutomationTurnAcceptedCancellation as exc:
# The agent owns this turn now. Advance and persist the schedule
# before allowing shutdown cancellation to unwind the timer.
status = "ok"
error = None
accepted_cancellation = exc
logger.info("Cron: job '{}' was accepted before cancellation", job.name)
except CronJobSkippedError as exc:
status = "skipped"
error = str(exc) or None
logger.warning("Cron: job '{}' skipped: {}", job.name, error or "")
except asyncio.CancelledError as exc:
current = asyncio.current_task()
if current is not None and current.cancelling():
raise
status = "error"
error = str(exc) or exc.__class__.__name__
logger.exception("Cron: job '{}' was cancelled", job.name)
except Exception as exc:
status = "error"
error = str(exc)
logger.exception("Cron: job '{}' failed", job.name)
settlement = asyncio.create_task(
self._settle_job_execution(
job,
start_ms=start_ms,
status=status,
error=error,
persist=accepted_cancellation is not None,
)
)
if accepted_cancellation is not None:
await self._drain_settlement_on_cancellation(settlement)
raise accepted_cancellation
await settlement
return True
finally:
await self._release_job_claim(job.id)
def _append_action( def _append_action(
self, self,
action: Literal["add", "del", "update"], action: Literal["add", "del", "update"],
@@ -697,7 +828,7 @@ class CronService:
store = self._require_store() store = self._require_store()
store.jobs.append(job) store.jobs.append(job)
self._save_store() self._save_store()
self._arm_timer() self._request_timer_rearm()
else: else:
self._append_action("add", asdict(job)) self._append_action("add", asdict(job))
@@ -714,7 +845,7 @@ class CronService:
store.jobs = [j for j in store.jobs if j.id != job.id] store.jobs = [j for j in store.jobs if j.id != job.id]
store.jobs.append(job) store.jobs.append(job)
self._save_store() self._save_store()
self._arm_timer() self._request_timer_rearm()
logger.info("Cron: registered system job '{}' ({})", job.name, job.id) logger.info("Cron: registered system job '{}' ({})", job.name, job.id)
return job return job
@@ -726,7 +857,7 @@ class CronService:
removed = len(store.jobs) < before removed = len(store.jobs) < before
if removed: if removed:
self._save_store() self._save_store()
self._arm_timer() self._request_timer_rearm()
logger.info("Cron: removed system job {}", job_id) logger.info("Cron: removed system job {}", job_id)
return removed return removed
@@ -747,7 +878,7 @@ class CronService:
if removed: if removed:
if self._should_persist_store(): if self._should_persist_store():
self._save_store() self._save_store()
self._arm_timer() self._request_timer_rearm()
else: else:
self._append_action("del", {"job_id": job_id}) self._append_action("del", {"job_id": job_id})
logger.info("Cron: removed job {}", job_id) logger.info("Cron: removed job {}", job_id)
@@ -769,7 +900,7 @@ class CronService:
job.state.next_run_at_ms = None job.state.next_run_at_ms = None
if self._should_persist_store(): if self._should_persist_store():
self._save_store() self._save_store()
self._arm_timer() self._request_timer_rearm()
else: else:
self._append_action("update", asdict(job)) self._append_action("update", asdict(job))
return job return job
@@ -825,7 +956,7 @@ class CronService:
if self._should_persist_store(): if self._should_persist_store():
self._save_store() self._save_store()
self._arm_timer() self._request_timer_rearm()
else: else:
self._append_action("update", asdict(job)) self._append_action("update", asdict(job))
@@ -840,19 +971,31 @@ class CronService:
# A manual run is another side-effecting entrypoint. Do not start # A manual run is another side-effecting entrypoint. Do not start
# it while the result of a previous timer execution is still only # it while the result of a previous timer execution is still only
# in memory. # in memory.
async with self._operation_lock:
if self._store_dirty: if self._store_dirty:
self._save_store() await shield_and_drain(asyncio.to_thread(self._save_store))
store = self._require_store(reload_during_execution=reload_store) store = await shield_and_drain(
asyncio.to_thread(
self._require_store,
reload_during_execution=reload_store,
)
)
for job in store.jobs: for job in store.jobs:
if job.id == job_id: if job.id == job_id:
if self._is_unbound_agent_job(job): if self._is_unbound_agent_job(job):
async with self._operation_lock:
self._enforce_agent_binding(job) self._enforce_agent_binding(job)
self._save_store() await shield_and_drain(
asyncio.to_thread(self._save_store)
)
return False return False
if not force and not job.enabled: if not force and not job.enabled:
return False return False
await self._execute_job(job) executed = await self._execute_job(job)
self._save_store() if not executed:
return False
async with self._operation_lock:
await shield_and_drain(asyncio.to_thread(self._save_store))
return True return True
return False return False
finally: finally:
+86
View File
@@ -0,0 +1,86 @@
"""Unified, content-free LLM usage backend."""
from __future__ import annotations
import threading
from pathlib import Path
from typing import Any
from loguru import logger
from nanobot.config.paths import get_data_dir
from nanobot.llm_usage.models import LLMCallRecord
from nanobot.llm_usage.store import LLMUsageStore
_STORES_LOCK = threading.Lock()
_STORES: dict[Path, LLMUsageStore] = {}
def empty_usage_payload() -> dict[str, Any]:
return {
"days": [],
"total_tokens": 0,
"total_tokens_30d": 0,
"total_tokens_365d": 0,
"reported_tokens_30d": 0,
"estimated_tokens_30d": 0,
"cache_read_tokens_30d": 0,
"cache_read_observed_input_tokens_30d": 0,
"cache_read_rate_30d": None,
"peak_day_tokens": 0,
"current_streak_days": 0,
"longest_streak_days": 0,
"active_days_30d": 0,
"requests_30d": 0,
"failed_requests_30d": 0,
"providers_30d": [],
"updated_at": None,
}
def llm_usage_store_path() -> Path:
return get_data_dir() / "llm_usage.sqlite3"
def get_llm_usage_store(path: Path | None = None) -> LLMUsageStore:
resolved = (path or llm_usage_store_path()).resolve(strict=False)
with _STORES_LOCK:
store = _STORES.get(resolved)
if store is None:
store = LLMUsageStore(resolved)
_STORES[resolved] = store
return store
def record_llm_call(call: LLMCallRecord) -> None:
"""Default fail-open callback attached to gateway provider snapshots."""
try:
get_llm_usage_store().record(call)
except Exception:
logger.exception("failed to record LLM usage")
def llm_usage_payload(
*,
days: int = 371,
timezone_name: str | None = None,
) -> dict[str, Any]:
try:
return get_llm_usage_store().usage_payload(
days=days,
timezone_name=timezone_name,
)
except Exception:
logger.exception("failed to query LLM usage")
return empty_usage_payload()
__all__ = [
"LLMCallRecord",
"LLMUsageStore",
"empty_usage_payload",
"get_llm_usage_store",
"record_llm_call",
"llm_usage_store_path",
"llm_usage_payload",
]
+70
View File
@@ -0,0 +1,70 @@
"""Request-local metadata for LLM usage records."""
from __future__ import annotations
from collections.abc import Generator, Mapping
from contextlib import contextmanager
from contextvars import ContextVar, Token
from typing import Literal
LLMUsageSource = Literal["user", "api", "cron", "dream", "system"]
_CURRENT_SOURCE: ContextVar[LLMUsageSource] = ContextVar(
"nanobot_llm_usage_source",
default="system",
)
def source_from_session_key(session_key: str | None) -> LLMUsageSource:
"""Classify a private session key without persisting that key."""
key = session_key or ""
if key.startswith("dream:"):
return "dream"
if key == "heartbeat" or key.startswith("cron:"):
return "cron"
if key.startswith("api:"):
return "api"
if key.startswith("system:"):
return "system"
return "user"
def source_from_request(
session_key: str | None,
*,
channel: str | None,
metadata: Mapping[str, object] | None,
) -> LLMUsageSource:
"""Classify a turn from trusted ingress metadata without retaining identifiers."""
values = metadata or {}
if isinstance(values.get("_cron_trigger"), Mapping):
return "cron"
if isinstance(values.get("_local_trigger"), Mapping):
return "cron"
if channel == "api":
return "api"
if channel == "system":
return "system"
return source_from_session_key(session_key)
def current_llm_usage_source() -> LLMUsageSource:
return _CURRENT_SOURCE.get()
def bind_llm_usage_source(source: LLMUsageSource) -> Token[LLMUsageSource]:
return _CURRENT_SOURCE.set(source)
def reset_llm_usage_source(token: Token[LLMUsageSource]) -> None:
_CURRENT_SOURCE.reset(token)
@contextmanager
def llm_usage_source(source: LLMUsageSource) -> Generator[None]:
"""Bind a coarse usage source for nested provider calls."""
token = bind_llm_usage_source(source)
try:
yield
finally:
reset_llm_usage_source(token)
+38
View File
@@ -0,0 +1,38 @@
"""Content-free records emitted for physical LLM provider calls."""
from __future__ import annotations
from dataclasses import dataclass
from nanobot.llm_usage.context import LLMUsageSource
from nanobot.providers.base import LLMUsage
@dataclass(frozen=True, slots=True)
class LLMCallRecord:
"""The small, chart-oriented result of one provider call attempt.
Request messages, response text, reasoning, and tool payloads deliberately do
not belong to this contract. Sessions already own that content.
"""
started_at_ms: int
duration_ms: int
provider: str
model: str
source: LLMUsageSource
stream: bool
finish_reason: str
usage: LLMUsage | None = None
error_status_code: int | None = None
error_kind: str | None = None
def __post_init__(self) -> None:
if self.started_at_ms < 0 or self.duration_ms < 0:
raise ValueError("LLM usage timestamps must be non-negative")
if not self.provider.strip() or not self.model.strip():
raise ValueError("LLM usage provider and model must be non-empty")
if self.source not in {"user", "api", "cron", "dream", "system"}:
raise ValueError("invalid LLM usage source")
if not self.finish_reason.strip():
raise ValueError("LLM usage finish_reason must be non-empty")
+560
View File
@@ -0,0 +1,560 @@
"""SQLite persistence and chart queries for LLM usage records."""
from __future__ import annotations
import os
import sqlite3
import threading
import time
from collections.abc import Iterable
from copy import deepcopy
from datetime import date, datetime, timedelta, timezone
from pathlib import Path
from typing import Any, cast
from zoneinfo import ZoneInfo, ZoneInfoNotFoundError
from nanobot.llm_usage.models import LLMCallRecord
SCHEMA_VERSION = 1
MAX_DAYS_RETAINED = 400
MAX_CALLS_RETAINED = 100_000
_ERROR_KINDS = frozenset({
"authentication",
"cancelled",
"configuration",
"connection",
"content_filter",
"context_length",
"empty",
"http",
"invalid_request",
"overloaded",
"permission",
"rate_limit",
"refusal",
"server_error",
"timeout",
})
_FINISH_REASONS = frozenset({
"cancelled",
"content_filter",
"error",
"function_call",
"length",
"refusal",
"stop",
"tool_calls",
})
_USAGE_COLUMNS = (
"input_tokens",
"output_tokens",
"cache_read_tokens",
"cache_write_tokens",
"cache_read_observed_input_tokens",
"cache_write_observed_input_tokens",
"total_tokens",
"reported_tokens",
"estimated_tokens",
"generation_ms",
"measured_output_tokens",
"ttft_ms",
"timed_requests",
)
_REQUEST_COLUMNS = (
"requests",
"successful_requests",
"failed_requests",
"reported_requests",
"estimated_requests",
)
_AGGREGATE_SQL = """
COALESCE(SUM(input_tokens), 0) AS input_tokens,
COALESCE(SUM(output_tokens), 0) AS output_tokens,
COALESCE(SUM(cache_read_tokens), 0) AS cache_read_tokens,
COALESCE(SUM(cache_write_tokens), 0) AS cache_write_tokens,
COALESCE(SUM(
CASE WHEN cache_read_tokens IS NOT NULL THEN input_tokens ELSE 0 END
), 0) AS cache_read_observed_input_tokens,
COALESCE(SUM(
CASE WHEN cache_write_tokens IS NOT NULL THEN input_tokens ELSE 0 END
), 0) AS cache_write_observed_input_tokens,
COALESCE(SUM(total_tokens), 0) AS total_tokens,
COALESCE(SUM(reported_tokens), 0) AS reported_tokens,
COALESCE(SUM(estimated_tokens), 0) AS estimated_tokens,
COALESCE(SUM(generation_ms), 0) AS generation_ms,
COALESCE(SUM(measured_output_tokens), 0) AS measured_output_tokens,
COALESCE(SUM(ttft_ms), 0) AS ttft_ms,
COALESCE(SUM(timed_requests), 0) AS timed_requests,
COUNT(*) AS requests,
COALESCE(SUM(CASE WHEN finish_reason IN ('error', 'cancelled') THEN 0 ELSE 1 END), 0)
AS successful_requests,
COALESCE(SUM(CASE WHEN finish_reason IN ('error', 'cancelled') THEN 1 ELSE 0 END), 0)
AS failed_requests,
COALESCE(SUM(
CASE WHEN total_tokens IS NOT NULL AND NOT (
estimated_tokens > 0 AND reported_tokens = 0
) THEN 1 ELSE 0 END
), 0) AS reported_requests,
COALESCE(SUM(
CASE WHEN estimated_tokens > 0 AND reported_tokens = 0 THEN 1 ELSE 0 END
), 0) AS estimated_requests,
COALESCE(SUM(duration_ms), 0) AS duration_ms
"""
def _zone(timezone_name: str | None) -> timezone | ZoneInfo:
if not timezone_name:
return timezone.utc
try:
return ZoneInfo(timezone_name)
except ZoneInfoNotFoundError:
return timezone.utc
def _clean_error_kind(value: str | None) -> str | None:
if value is None:
return None
cleaned = value.strip().lower()
if not cleaned:
return None
return cleaned if cleaned in _ERROR_KINDS else "other"
def _clean_finish_reason(value: str) -> str:
cleaned = value.strip().lower()
return cleaned if cleaned in _FINISH_REASONS else "other"
def _clean_status_code(value: int | None) -> int | None:
if value is None:
return None
try:
status = int(value)
except (TypeError, ValueError):
return None
return status if 100 <= status <= 599 else None
def _as_int_row(row: sqlite3.Row) -> dict[str, int]:
return {
key: max(0, int(row[key] or 0))
for key in (*_USAGE_COLUMNS, *_REQUEST_COLUMNS, "duration_ms")
}
def _empty_totals() -> dict[str, int]:
return {key: 0 for key in (*_USAGE_COLUMNS, *_REQUEST_COLUMNS, "duration_ms")}
def _sum_rows(rows: Iterable[dict[str, Any]]) -> dict[str, int]:
totals = _empty_totals()
for row in rows:
for key in totals:
totals[key] += max(0, int(row.get(key) or 0))
return totals
class LLMUsageStore:
"""A small synchronous WAL database shared by gateway threads/processes."""
def __init__(self, path: Path) -> None:
self.path = path
self._lock = threading.RLock()
self._connection: sqlite3.Connection | None = None
self._connection_pid: int | None = None
self._last_prune_utc_day: int | None = None
self._writes_since_size_prune = 0
self._write_version = 0
self._cached_payload_key: tuple[int, str, str, int, int] | None = None
self._cached_payload: dict[str, Any] | None = None
def _connect(self) -> sqlite3.Connection:
pid = os.getpid()
if self._connection is not None and self._connection_pid == pid:
return self._connection
if self._connection is not None:
self._connection.close()
self._cached_payload_key = None
self._cached_payload = None
self.path.parent.mkdir(parents=True, exist_ok=True)
connection = sqlite3.connect(
self.path,
timeout=0.25,
isolation_level=None,
check_same_thread=False,
)
connection.row_factory = sqlite3.Row
connection.execute("PRAGMA busy_timeout = 250")
connection.execute("PRAGMA journal_mode = WAL")
connection.execute("PRAGMA synchronous = NORMAL")
connection.execute("PRAGMA temp_store = MEMORY")
connection.create_function("llm_usage_local_day", 2, self._local_day, deterministic=True)
connection.executescript(
"""
CREATE TABLE IF NOT EXISTS llm_calls (
id INTEGER PRIMARY KEY,
started_at_ms INTEGER NOT NULL,
duration_ms INTEGER NOT NULL,
provider TEXT NOT NULL,
model TEXT NOT NULL,
source TEXT NOT NULL,
stream INTEGER NOT NULL,
finish_reason TEXT NOT NULL,
input_tokens INTEGER,
output_tokens INTEGER,
total_tokens INTEGER,
cache_read_tokens INTEGER,
cache_write_tokens INTEGER,
reported_tokens INTEGER,
estimated_tokens INTEGER,
generation_ms INTEGER,
measured_output_tokens INTEGER,
ttft_ms INTEGER,
timed_requests INTEGER,
error_status_code INTEGER,
error_kind TEXT
);
CREATE INDEX IF NOT EXISTS llm_calls_started_at_idx
ON llm_calls(started_at_ms);
CREATE INDEX IF NOT EXISTS llm_calls_provider_model_time_idx
ON llm_calls(provider, model, started_at_ms);
"""
)
connection.execute(f"PRAGMA user_version = {SCHEMA_VERSION}")
self._connection = connection
self._connection_pid = pid
return connection
def _read_connection(self) -> sqlite3.Connection:
connection = sqlite3.connect(
self.path,
timeout=0.25,
isolation_level=None,
)
connection.row_factory = sqlite3.Row
connection.execute("PRAGMA busy_timeout = 250")
connection.execute("PRAGMA query_only = ON")
connection.execute("PRAGMA temp_store = MEMORY")
connection.create_function("llm_usage_local_day", 2, self._local_day, deterministic=True)
return connection
@staticmethod
def _local_day(started_at_ms: object, timezone_name: object) -> str | None:
if not isinstance(started_at_ms, int) or not isinstance(timezone_name, str):
return None
dt = datetime.fromtimestamp(started_at_ms / 1000, timezone.utc)
return dt.astimezone(_zone(timezone_name)).date().isoformat()
def close(self) -> None:
with self._lock:
if self._connection is not None:
self._connection.close()
self._connection = None
self._connection_pid = None
self._cached_payload_key = None
self._cached_payload = None
def record(self, call: LLMCallRecord) -> None:
usage = call.usage
usage_data = usage.to_dict() if usage is not None else {}
values: tuple[object, ...] = (
call.started_at_ms,
call.duration_ms,
call.provider[:120],
call.model[:240],
call.source,
int(call.stream),
_clean_finish_reason(call.finish_reason),
*(
usage_data.get(key)
for key in (
"input_tokens",
"output_tokens",
"total_tokens",
"cache_read_tokens",
"cache_write_tokens",
"reported_tokens",
"estimated_tokens",
"generation_ms",
"measured_output_tokens",
"ttft_ms",
"timed_requests",
)
),
_clean_status_code(call.error_status_code),
_clean_error_kind(call.error_kind),
)
with self._lock:
connection = self._connect()
connection.execute(
"""
INSERT INTO llm_calls (
started_at_ms, duration_ms, provider, model, source, stream,
finish_reason, input_tokens, output_tokens, total_tokens,
cache_read_tokens, cache_write_tokens, reported_tokens,
estimated_tokens, generation_ms, measured_output_tokens,
ttft_ms, timed_requests, error_status_code, error_kind
) VALUES (
?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?
)
""",
values,
)
self._write_version += 1
self._cached_payload_key = None
self._cached_payload = None
self._prune_if_due(connection)
def _prune_if_due(self, connection: sqlite3.Connection) -> None:
utc_day = int(time.time() // 86_400)
self._writes_since_size_prune += 1
prune_age = self._last_prune_utc_day != utc_day
prune_size = self._writes_since_size_prune >= 1_024
if not prune_age and not prune_size:
return
if prune_age:
cutoff_ms = int(
(datetime.now(timezone.utc) - timedelta(days=MAX_DAYS_RETAINED)).timestamp()
* 1000
)
connection.execute("DELETE FROM llm_calls WHERE started_at_ms < ?", (cutoff_ms,))
connection.execute(
"""
DELETE FROM llm_calls
WHERE id <= COALESCE((
SELECT id FROM llm_calls ORDER BY id DESC LIMIT 1 OFFSET ?
), -1)
""",
(MAX_CALLS_RETAINED,),
)
self._last_prune_utc_day = utc_day
self._writes_since_size_prune = 0
def count(self) -> int:
with self._lock:
row = self._connect().execute("SELECT COUNT(*) AS count FROM llm_calls").fetchone()
return int(row["count"] if row is not None else 0)
def _aggregate(
self,
*,
connection: sqlite3.Connection,
start_ms: int | None,
end_ms: int,
group_by: tuple[str, ...] = (),
limit: int | None = None,
) -> list[sqlite3.Row]:
selected = f"{', '.join(group_by)}, " if group_by else ""
where = "started_at_ms < ?"
params: list[object] = [end_ms]
if start_ms is not None:
where = "started_at_ms >= ? AND started_at_ms < ?"
params = [start_ms, end_ms]
query = f"SELECT {selected}{_AGGREGATE_SQL} FROM llm_calls WHERE {where}"
if group_by:
query += f" GROUP BY {', '.join(group_by)} ORDER BY total_tokens DESC"
if limit is not None:
query += " LIMIT ?"
params.append(limit)
return list(connection.execute(query, params).fetchall())
def _daily_rows(
self,
*,
connection: sqlite3.Connection,
start_ms: int,
end_ms: int,
timezone_name: str,
) -> list[dict[str, Any]]:
query = f"""
SELECT llm_usage_local_day(started_at_ms, ?) AS date, source,
{_AGGREGATE_SQL}
FROM llm_calls
WHERE started_at_ms >= ? AND started_at_ms < ?
GROUP BY date, source
ORDER BY date, source
"""
rows = connection.execute(
query,
(timezone_name, start_ms, end_ms),
).fetchall()
by_date: dict[str, dict[str, Any]] = {}
for row in rows:
day = cast(str | None, row["date"])
if day is None:
continue
values = _as_int_row(row)
aggregate = by_date.setdefault(
day,
{"date": day, **_empty_totals(), "sources": {}},
)
for key, value in values.items():
aggregate[key] += value
aggregate["sources"][str(row["source"])] = values
return list(by_date.values())
@staticmethod
def _midnight_ms(value: date, zone: timezone | ZoneInfo) -> int:
return int(datetime.combine(value, datetime.min.time(), tzinfo=zone).timestamp() * 1000)
def usage_payload(
self,
*,
days: int = 371,
timezone_name: str | None = None,
now: datetime | None = None,
) -> dict[str, Any]:
zone = _zone(timezone_name)
current = now or datetime.now(timezone.utc)
if current.tzinfo is None:
current = current.replace(tzinfo=timezone.utc)
today = current.astimezone(zone).date()
safe_days = max(1, days)
zone_name = getattr(zone, "key", "UTC")
with self._lock:
data_version_row = self._connect().execute("PRAGMA data_version").fetchone()
data_version = int(data_version_row[0]) if data_version_row is not None else 0
write_version = self._write_version
cache_key = (
safe_days,
zone_name,
today.isoformat(),
write_version,
data_version,
)
if self._cached_payload_key == cache_key and self._cached_payload is not None:
return deepcopy(self._cached_payload)
connection = self._read_connection()
try:
connection.execute("BEGIN")
end_ms = self._midnight_ms(today + timedelta(days=1), zone)
retained_start = today - timedelta(days=MAX_DAYS_RETAINED - 1)
retained_start_ms = self._midnight_ms(retained_start, zone)
daily = self._daily_rows(
connection=connection,
start_ms=retained_start_ms,
end_ms=end_ms,
timezone_name=zone_name,
)
requested_start = today - timedelta(days=safe_days - 1)
visible_days = [row for row in daily if row["date"] >= requested_start.isoformat()]
last_30_start_ms = self._midnight_ms(today - timedelta(days=29), zone)
last_30_date = (today - timedelta(days=29)).isoformat()
last_365_date = (today - timedelta(days=364)).isoformat()
all_totals = _sum_rows(daily)
totals_30 = _sum_rows(row for row in daily if row["date"] >= last_30_date)
totals_365 = _sum_rows(row for row in daily if row["date"] >= last_365_date)
provider_rows = self._aggregate(
connection=connection,
start_ms=last_30_start_ms,
end_ms=end_ms,
group_by=("provider", "model"),
limit=50,
)
providers_30d = [
{
"provider": str(row["provider"]),
"model": str(row["model"]),
**_as_int_row(row),
}
for row in provider_rows
]
active_dates = {
date.fromisoformat(row["date"]) for row in daily if row["total_tokens"] > 0
}
current_streak = 0
cursor = today
while cursor in active_dates:
current_streak += 1
cursor -= timedelta(days=1)
longest_streak = 0
running_streak = 0
previous: date | None = None
for cursor in sorted(active_dates):
running_streak = running_streak + 1 if previous == cursor - timedelta(days=1) else 1
longest_streak = max(longest_streak, running_streak)
previous = cursor
latest = (
connection
.execute("SELECT MAX(started_at_ms) AS updated_at_ms FROM llm_calls")
.fetchone()
)
updated_at_ms = int(latest["updated_at_ms"] or 0) if latest is not None else 0
denominator = totals_30["cache_read_observed_input_tokens"]
payload = {
"days": visible_days,
"total_tokens": all_totals["total_tokens"],
"total_tokens_30d": totals_30["total_tokens"],
"total_tokens_365d": totals_365["total_tokens"],
"reported_tokens_30d": totals_30["reported_tokens"],
"estimated_tokens_30d": totals_30["estimated_tokens"],
"cache_read_tokens_30d": totals_30["cache_read_tokens"],
"cache_read_observed_input_tokens_30d": denominator,
"cache_read_rate_30d": (
totals_30["cache_read_tokens"] / denominator if denominator else None
),
"peak_day_tokens": max(
(int(row["total_tokens"]) for row in daily),
default=0,
),
"current_streak_days": current_streak,
"longest_streak_days": longest_streak,
"active_days_30d": sum(
1
for row in daily
if row["date"] >= last_30_date and row["total_tokens"] > 0
),
"requests_30d": totals_30["requests"],
"failed_requests_30d": totals_30["failed_requests"],
"providers_30d": providers_30d,
"updated_at": (
datetime.fromtimestamp(updated_at_ms / 1000, timezone.utc)
.isoformat()
.replace("+00:00", "Z")
if updated_at_ms
else None
),
}
finally:
connection.close()
with self._lock:
latest_data_version_row = self._connect().execute("PRAGMA data_version").fetchone()
latest_data_version = (
int(latest_data_version_row[0])
if latest_data_version_row is not None
else 0
)
if self._write_version == write_version and latest_data_version == data_version:
self._cached_payload_key = cache_key
self._cached_payload = payload
return deepcopy(payload)
def recent_calls(self, *, limit: int = 100) -> list[dict[str, Any]]:
"""Return bounded metadata rows for diagnostics; never returns content."""
safe_limit = min(max(1, limit), 1_000)
with self._lock:
rows = (
self._connect()
.execute(
"""
SELECT * FROM llm_calls ORDER BY started_at_ms DESC, id DESC LIMIT ?
""",
(safe_limit,),
)
.fetchall()
)
return [dict(row) for row in rows]
def record_many(self, calls: Iterable[LLMCallRecord]) -> None:
for call in calls:
self.record(call)
+3 -1
View File
@@ -13,6 +13,7 @@ from nanobot.agent.loop import AgentLoop
from nanobot.agent.tools.mcp import MCPProvider from nanobot.agent.tools.mcp import MCPProvider
from nanobot.agent.tools.registry import ToolRegistry from nanobot.agent.tools.registry import ToolRegistry
from nanobot.config.schema import Config from nanobot.config.schema import Config
from nanobot.providers.base import LLMUsage
from nanobot.providers.image_generation import image_gen_provider_configs from nanobot.providers.image_generation import image_gen_provider_configs
from nanobot.sdk.clients import MemoryClient, RuntimeClient, SessionClient from nanobot.sdk.clients import MemoryClient, RuntimeClient, SessionClient
from nanobot.sdk.runtime import ( from nanobot.sdk.runtime import (
@@ -43,6 +44,7 @@ from nanobot.utils.llm_runtime import LLMRuntime
__all__ = [ __all__ = [
"Nanobot", "Nanobot",
"LLMUsage",
"RunResult", "RunResult",
"RunStream", "RunStream",
"SessionInfo", "SessionInfo",
@@ -287,7 +289,7 @@ class Nanobot:
type=STREAM_EVENT_RUN_COMPLETED, type=STREAM_EVENT_RUN_COMPLETED,
content=result.content, content=result.content,
result=result, result=result,
usage=dict(result.usage), usage=result.usage,
metadata=dict(result.metadata), metadata=dict(result.metadata),
)) ))
return result return result
+2 -1
View File
@@ -5,11 +5,12 @@ from __future__ import annotations
from importlib import import_module from importlib import import_module
from typing import TYPE_CHECKING from typing import TYPE_CHECKING
from nanobot.providers.base import LLMProvider, LLMResponse from nanobot.providers.base import LLMProvider, LLMResponse, LLMUsage
__all__ = [ __all__ = [
"LLMProvider", "LLMProvider",
"LLMResponse", "LLMResponse",
"LLMUsage",
"AnthropicProvider", "AnthropicProvider",
"OpenAICompatProvider", "OpenAICompatProvider",
"OpenAICodexProvider", "OpenAICodexProvider",
+22 -18
View File
@@ -17,6 +17,7 @@ from loguru import logger
from nanobot.providers.base import ( from nanobot.providers.base import (
LLMProvider, LLMProvider,
LLMResponse, LLMResponse,
LLMUsage,
ToolCallRequest, ToolCallRequest,
resolve_stream_idle_timeout_s, resolve_stream_idle_timeout_s,
tool_arguments_object_for_replay, tool_arguments_object_for_replay,
@@ -90,8 +91,10 @@ class AnthropicProvider(LLMProvider):
api_base: str | None = None, api_base: str | None = None,
default_model: str = "claude-sonnet-4-6", default_model: str = "claude-sonnet-4-6",
extra_headers: dict[str, str] | None = None, extra_headers: dict[str, str] | None = None,
*,
provider_name: str = "anthropic",
): ):
super().__init__(api_key, api_base) super().__init__(api_key, api_base, provider_name=provider_name)
self.default_model = default_model self.default_model = default_model
self.extra_headers = extra_headers or {} self.extra_headers = extra_headers or {}
@@ -689,24 +692,25 @@ class AnthropicProvider(LLMProvider):
stop_map = {"tool_use": "tool_calls", "end_turn": "stop", "max_tokens": "length"} stop_map = {"tool_use": "tool_calls", "end_turn": "stop", "max_tokens": "length"}
finish_reason = stop_map.get(response.stop_reason or "", response.stop_reason or "stop") finish_reason = stop_map.get(response.stop_reason or "", response.stop_reason or "stop")
usage: dict[str, int] = {} usage: LLMUsage | None = None
if response.usage: if response.usage:
input_tokens = response.usage.input_tokens cache_write_raw = getattr(
cache_creation = getattr(response.usage, "cache_creation_input_tokens", 0) or 0 response.usage,
cache_read = getattr(response.usage, "cache_read_input_tokens", 0) or 0 "cache_creation_input_tokens",
total_prompt_tokens = input_tokens + cache_creation + cache_read None,
usage = { )
"prompt_tokens": total_prompt_tokens, cache_read_raw = getattr(response.usage, "cache_read_input_tokens", None)
"completion_tokens": response.usage.output_tokens, cache_write = int(cache_write_raw) if cache_write_raw is not None else None
"total_tokens": total_prompt_tokens + response.usage.output_tokens, cache_read = int(cache_read_raw) if cache_read_raw is not None else None
} logical_input = int(response.usage.input_tokens) + (cache_write or 0) + (
for attr in ("cache_creation_input_tokens", "cache_read_input_tokens"): cache_read or 0
val = getattr(response.usage, attr, 0) )
if val: usage = LLMUsage.reported(
usage[attr] = val input_tokens=logical_input,
# Normalize to cached_tokens for downstream consistency. output_tokens=int(response.usage.output_tokens),
if cache_read: cache_read_tokens=cache_read,
usage["cached_tokens"] = cache_read cache_write_tokens=cache_write,
)
return LLMResponse( return LLMResponse(
content="".join(content_parts) or None, content="".join(content_parts) or None,
+3 -1
View File
@@ -106,8 +106,10 @@ class AzureOpenAIProvider(LLMProvider):
api_key: str = "", api_key: str = "",
api_base: str = "", api_base: str = "",
default_model: str = "gpt-5.2-chat", default_model: str = "gpt-5.2-chat",
*,
provider_name: str = "azure_openai",
): ):
super().__init__(api_key, api_base) super().__init__(api_key, api_base, provider_name=provider_name)
self.default_model = default_model self.default_model = default_model
self._native_compaction_available = True self._native_compaction_available = True
+440 -9
View File
@@ -6,6 +6,7 @@ import asyncio
import json import json
import os import os
import re import re
import time
from abc import ABC, abstractmethod from abc import ABC, abstractmethod
from collections.abc import Awaitable, Callable from collections.abc import Awaitable, Callable
from contextlib import suppress from contextlib import suppress
@@ -13,19 +14,23 @@ from copy import deepcopy
from dataclasses import dataclass, field from dataclasses import dataclass, field
from datetime import datetime, timezone from datetime import datetime, timezone
from email.utils import parsedate_to_datetime from email.utils import parsedate_to_datetime
from typing import Any, cast from typing import TYPE_CHECKING, Any, Literal, cast
import json_repair import json_repair
from loguru import logger from loguru import logger
from nanobot.utils.helpers import sanitize_surrogates_deep from nanobot.utils.helpers import sanitize_surrogates_deep
if TYPE_CHECKING:
from nanobot.llm_usage.models import LLMCallRecord
STREAM_IDLE_TIMEOUT_ENV = "NANOBOT_STREAM_IDLE_TIMEOUT_S" STREAM_IDLE_TIMEOUT_ENV = "NANOBOT_STREAM_IDLE_TIMEOUT_S"
DEFAULT_STREAM_IDLE_TIMEOUT_S = 90.0 DEFAULT_STREAM_IDLE_TIMEOUT_S = 90.0
MAX_STREAM_IDLE_TIMEOUT_S = 3600.0 MAX_STREAM_IDLE_TIMEOUT_S = 3600.0
RETRY_AFTER_BUFFER = 1 RETRY_AFTER_BUFFER = 1
RetryEventCallback = Callable[[str], Awaitable[None]] RetryEventCallback = Callable[[str], Awaitable[None]]
LLMCallObserver = Callable[["LLMCallRecord"], None]
def resolve_stream_idle_timeout_s( def resolve_stream_idle_timeout_s(
@@ -253,13 +258,298 @@ class ProviderCallContext:
context_window_tokens: int | None = None context_window_tokens: int | None = None
@dataclass(frozen=True, slots=True)
class LLMUsage:
"""Canonical token usage reported by, or estimated for, one or more LLM calls.
``input_tokens`` is the logical input total and therefore includes cache reads
and writes. ``None`` cache counts mean the wire protocol did not report that
metric, while zero means it explicitly reported no cache activity.
``total_tokens`` preserves a provider-reported total when it exceeds the
visible input plus output (for example, hidden reasoning or tool usage). It
must be at least ``input_tokens + output_tokens``. The reported and estimated
totals partition it exactly, including after multi-call aggregation.
"""
input_tokens: int
output_tokens: int
total_tokens: int
cache_read_tokens: int | None = None
cache_write_tokens: int | None = None
reported_tokens: int = 0
estimated_tokens: int = 0
generation_ms: int = 0
measured_output_tokens: int = 0
ttft_ms: int = 0
timed_requests: int = 0
context_tokens: int | None = None
request_count: int = 0
def __post_init__(self) -> None:
token_fields = {
"input_tokens": self.input_tokens,
"output_tokens": self.output_tokens,
"total_tokens": self.total_tokens,
"reported_tokens": self.reported_tokens,
"estimated_tokens": self.estimated_tokens,
"generation_ms": self.generation_ms,
"measured_output_tokens": self.measured_output_tokens,
"ttft_ms": self.ttft_ms,
"timed_requests": self.timed_requests,
"request_count": self.request_count,
}
for name, value in token_fields.items():
runtime_value = cast(object, value)
if (
not isinstance(runtime_value, int)
or isinstance(runtime_value, bool)
or runtime_value < 0
):
raise ValueError(f"{name} must be a non-negative integer")
for name, value in (
("cache_read_tokens", self.cache_read_tokens),
("cache_write_tokens", self.cache_write_tokens),
("context_tokens", self.context_tokens),
):
runtime_value = cast(object, value)
if runtime_value is not None and (
not isinstance(runtime_value, int)
or isinstance(runtime_value, bool)
or runtime_value < 0
):
raise ValueError(f"{name} must be None or a non-negative integer")
visible_total = self.input_tokens + self.output_tokens
if self.total_tokens < visible_total:
raise ValueError("total_tokens must be at least input_tokens + output_tokens")
if self.reported_tokens + self.estimated_tokens != self.total_tokens:
raise ValueError("reported_tokens + estimated_tokens must equal total_tokens")
cache_total = (self.cache_read_tokens or 0) + (self.cache_write_tokens or 0)
if cache_total > self.input_tokens:
raise ValueError("cache token counts cannot exceed logical input_tokens")
@classmethod
def reported(
cls,
*,
input_tokens: int,
output_tokens: int,
total_tokens: int | None = None,
cache_read_tokens: int | None = None,
cache_write_tokens: int | None = None,
) -> LLMUsage:
"""Build usage normalized from a provider response."""
visible_total = input_tokens + output_tokens
normalized_total = (
visible_total if total_tokens is None else max(visible_total, total_tokens)
)
return cls(
input_tokens=input_tokens,
output_tokens=output_tokens,
total_tokens=normalized_total,
cache_read_tokens=cache_read_tokens,
cache_write_tokens=cache_write_tokens,
reported_tokens=normalized_total,
context_tokens=input_tokens,
request_count=1,
)
@classmethod
def estimated(cls, *, input_tokens: int, output_tokens: int) -> LLMUsage:
"""Build usage estimated locally because the provider omitted it."""
return cls(
input_tokens=input_tokens,
output_tokens=output_tokens,
total_tokens=input_tokens + output_tokens,
estimated_tokens=input_tokens + output_tokens,
context_tokens=input_tokens,
request_count=1,
)
@classmethod
def empty_request(cls) -> LLMUsage:
"""Represent a completed model request with no measurable token usage."""
return cls(
input_tokens=0,
output_tokens=0,
total_tokens=0,
request_count=1,
)
@property
def source(self) -> Literal["reported", "estimated", "mixed"]:
if self.estimated_tokens == 0:
return "reported"
if self.reported_tokens == 0:
return "estimated"
return "mixed"
def with_timing(
self,
*,
generation_ms: int | None,
ttft_ms: int | None,
) -> LLMUsage:
"""Attach locally measured streaming telemetry to this usage value."""
return LLMUsage(
input_tokens=self.input_tokens,
output_tokens=self.output_tokens,
total_tokens=self.total_tokens,
cache_read_tokens=self.cache_read_tokens,
cache_write_tokens=self.cache_write_tokens,
reported_tokens=self.reported_tokens,
estimated_tokens=self.estimated_tokens,
generation_ms=max(0, generation_ms or 0),
measured_output_tokens=self.output_tokens if generation_ms is not None else 0,
ttft_ms=max(0, ttft_ms or 0),
timed_requests=1 if ttft_ms is not None else 0,
context_tokens=self.context_tokens,
request_count=self.request_count,
)
def __add__(self, other: LLMUsage) -> LLMUsage:
"""Aggregate calls without turning partially reported cache data into a count."""
def _sum_cache(left: int | None, right: int | None) -> int | None:
return left + right if left is not None and right is not None else None
return LLMUsage(
input_tokens=self.input_tokens + other.input_tokens,
output_tokens=self.output_tokens + other.output_tokens,
total_tokens=self.total_tokens + other.total_tokens,
cache_read_tokens=_sum_cache(self.cache_read_tokens, other.cache_read_tokens),
cache_write_tokens=_sum_cache(self.cache_write_tokens, other.cache_write_tokens),
reported_tokens=self.reported_tokens + other.reported_tokens,
estimated_tokens=self.estimated_tokens + other.estimated_tokens,
generation_ms=self.generation_ms + other.generation_ms,
measured_output_tokens=(
self.measured_output_tokens + other.measured_output_tokens
),
ttft_ms=self.ttft_ms + other.ttft_ms,
timed_requests=self.timed_requests + other.timed_requests,
context_tokens=(
other.context_tokens
if other.context_tokens is not None
else self.context_tokens
),
request_count=self.request_count + other.request_count,
)
def to_dict(self) -> dict[str, int | str | None]:
"""Serialize the canonical contract at JSON/persistence boundaries."""
return {
"input_tokens": self.input_tokens,
"output_tokens": self.output_tokens,
"total_tokens": self.total_tokens,
"cache_read_tokens": self.cache_read_tokens,
"cache_write_tokens": self.cache_write_tokens,
"reported_tokens": self.reported_tokens,
"estimated_tokens": self.estimated_tokens,
"source": self.source,
"generation_ms": self.generation_ms,
"measured_output_tokens": self.measured_output_tokens,
"ttft_ms": self.ttft_ms,
"timed_requests": self.timed_requests,
"context_tokens": self.context_tokens,
"request_count": self.request_count,
}
def to_turn_dict(self) -> dict[str, int]:
"""Project canonical usage into the compact WebUI/TUI per-turn shape."""
result: dict[str, int] = {
"prompt_tokens": self.input_tokens,
"completion_tokens": self.output_tokens,
"total_tokens": self.total_tokens,
"request_count": self.request_count,
"estimated_tokens": self.estimated_tokens,
}
if self.context_tokens is not None:
result["context_tokens"] = self.context_tokens
if self.cache_read_tokens is not None:
result["cached_tokens"] = self.cache_read_tokens
if self.cache_write_tokens is not None:
result["cache_write_tokens"] = self.cache_write_tokens
if self.generation_ms > 0 and self.measured_output_tokens > 0:
result["generation_ms"] = self.generation_ms
result["measured_completion_tokens"] = self.measured_output_tokens
if self.timed_requests > 0:
result["ttft_ms"] = self.ttft_ms
result["timed_requests"] = self.timed_requests
return result
@classmethod
def from_dict(cls, value: object) -> LLMUsage | None:
"""Validate the exact first-party serialized contract."""
if not isinstance(value, dict):
return None
data = cast(dict[object, object], value)
integer_fields = (
"input_tokens",
"output_tokens",
"reported_tokens",
"estimated_tokens",
"generation_ms",
"measured_output_tokens",
"ttft_ms",
"timed_requests",
"request_count",
)
serialized_fields = {
*integer_fields,
"total_tokens",
"cache_read_tokens",
"cache_write_tokens",
"context_tokens",
"source",
}
if set(data) != serialized_fields:
return None
if any(
not isinstance(item := data.get(name), int) or isinstance(item, bool)
for name in integer_fields
):
return None
cache_read = data.get("cache_read_tokens")
cache_write = data.get("cache_write_tokens")
context_tokens = data.get("context_tokens")
total = data.get("total_tokens")
source = data.get("source")
if any(
item is not None and (not isinstance(item, int) or isinstance(item, bool))
for item in (cache_read, cache_write, context_tokens)
) or not isinstance(total, int) or isinstance(total, bool):
return None
try:
usage = cls(
input_tokens=cast(int, data["input_tokens"]),
output_tokens=cast(int, data["output_tokens"]),
total_tokens=total,
cache_read_tokens=cast(int | None, cache_read),
cache_write_tokens=cast(int | None, cache_write),
reported_tokens=cast(int, data["reported_tokens"]),
estimated_tokens=cast(int, data["estimated_tokens"]),
generation_ms=cast(int, data["generation_ms"]),
measured_output_tokens=cast(int, data["measured_output_tokens"]),
ttft_ms=cast(int, data["ttft_ms"]),
timed_requests=cast(int, data["timed_requests"]),
context_tokens=cast(int | None, context_tokens),
request_count=cast(int, data["request_count"]),
)
except (KeyError, ValueError):
return None
if source != usage.source:
return None
return usage
@dataclass @dataclass
class LLMResponse: class LLMResponse:
"""Response from an LLM provider.""" """Response from an LLM provider."""
content: str | None content: str | None
tool_calls: list[ToolCallRequest] = field(default_factory=list) tool_calls: list[ToolCallRequest] = field(default_factory=list)
finish_reason: str = "stop" finish_reason: str = "stop"
usage: dict[str, int] = field(default_factory=dict) usage: LLMUsage | None = None
# Locally measured streaming telemetry. ``generation_ms`` excludes time to # Locally measured streaming telemetry. ``generation_ms`` excludes time to
# first token and provider retry gaps; ``ttft_ms`` measures the first # first token and provider retry gaps; ``ttft_ms`` measures the first
# streamed reasoning/content delta from request start. They stay separate # streamed reasoning/content delta from request start. They stay separate
@@ -383,10 +673,109 @@ class LLMProvider(ABC):
_SENTINEL = object() _SENTINEL = object()
def __init__(self, api_key: str | None = None, api_base: str | None = None): def __init__(
self,
api_key: str | None = None,
api_base: str | None = None,
*,
provider_name: str,
):
runtime_provider_name = cast(object, provider_name)
if not isinstance(runtime_provider_name, str) or not runtime_provider_name.strip():
raise ValueError("provider_name must be a non-empty configured identity")
self.api_key = api_key self.api_key = api_key
self.api_base = api_base self.api_base = api_base
self.provider_name = provider_name
self.generation: GenerationSettings = GenerationSettings() self.generation: GenerationSettings = GenerationSettings()
self._llm_call_observer: LLMCallObserver | None = None
def set_llm_call_observer(self, observer: LLMCallObserver | None) -> None:
"""Attach a fail-open observer for each physical retry-managed call."""
self._llm_call_observer = observer
def _usage_for_call(
self,
response: LLMResponse,
kwargs: dict[str, Any],
) -> LLMUsage | None:
usage = response.usage
if usage is None or usage.total_tokens == 0:
if response.finish_reason in {"error", "cancelled"}:
return None
messages = kwargs.get("messages")
if not isinstance(messages, list):
return usage
tools_value = kwargs.get("tools")
tools = cast(list[dict[str, Any]], tools_value) if isinstance(tools_value, list) else None
model_value = kwargs.get("model")
model = model_value if isinstance(model_value, str) else self.get_default_model()
try:
from nanobot.utils.helpers import (
build_assistant_message,
estimate_message_tokens,
estimate_prompt_tokens_chain,
)
input_tokens, _ = estimate_prompt_tokens_chain(
self,
model,
cast(list[dict[str, Any]], messages),
tools,
)
assistant_message = build_assistant_message(
response.content or "",
tool_calls=[call.to_openai_tool_call() for call in response.tool_calls],
reasoning_content=response.reasoning_content,
thinking_blocks=response.thinking_blocks,
)
usage = LLMUsage.estimated(
input_tokens=max(0, input_tokens),
output_tokens=max(0, estimate_message_tokens(assistant_message)),
)
except Exception:
logger.exception("failed to estimate usage for {}", self.provider_name)
return usage
return usage.with_timing(
generation_ms=response.generation_ms,
ttft_ms=response.ttft_ms,
)
def _observe_llm_call(
self,
response: LLMResponse,
kwargs: dict[str, Any],
*,
started_at_ms: int,
started_at_ns: int,
stream: bool,
) -> LLMResponse:
observer = self._llm_call_observer
if observer is None:
return response
usage = self._usage_for_call(response, kwargs)
if usage is not None:
response.usage = usage
model_value = kwargs.get("model")
model = model_value if isinstance(model_value, str) and model_value else self.get_default_model()
try:
from nanobot.llm_usage.context import current_llm_usage_source
from nanobot.llm_usage.models import LLMCallRecord
observer(LLMCallRecord(
started_at_ms=started_at_ms,
duration_ms=max(0, (time.monotonic_ns() - started_at_ns) // 1_000_000),
provider=self.provider_name,
model=model,
source=current_llm_usage_source(),
stream=stream,
finish_reason=response.finish_reason,
usage=usage,
error_status_code=response.error_status_code,
error_kind=response.error_kind,
))
except Exception:
logger.exception("LLM call observer failed for {}", self.provider_name)
return response
def can_resume_conversation_state( def can_resume_conversation_state(
self, self,
@@ -773,18 +1162,39 @@ class LLMProvider(ABC):
async def _safe_chat(self, **kwargs: Any) -> LLMResponse: async def _safe_chat(self, **kwargs: Any) -> LLMResponse:
"""Call chat() and convert unexpected exceptions to error responses.""" """Call chat() and convert unexpected exceptions to error responses."""
started_at_ms = time.time_ns() // 1_000_000
started_at_ns = time.monotonic_ns()
try: try:
provider_context = kwargs.pop("provider_context", None) provider_context = kwargs.pop("provider_context", None)
if isinstance(provider_context, ProviderCallContext): if isinstance(provider_context, ProviderCallContext):
return await self.chat_with_context( response = await self.chat_with_context(
provider_context=provider_context, provider_context=provider_context,
**kwargs, **kwargs,
) )
return await self.chat(**kwargs) else:
response = await self.chat(**kwargs)
except asyncio.CancelledError: except asyncio.CancelledError:
self._observe_llm_call(
LLMResponse(
content=None,
finish_reason="cancelled",
error_kind="cancelled",
),
kwargs,
started_at_ms=started_at_ms,
started_at_ns=started_at_ns,
stream=False,
)
raise raise
except Exception as exc: except Exception as exc:
return LLMResponse(content=f"Error calling LLM: {exc}", finish_reason="error") response = LLMResponse(content=f"Error calling LLM: {exc}", finish_reason="error")
return self._observe_llm_call(
response,
kwargs,
started_at_ms=started_at_ms,
started_at_ns=started_at_ns,
stream=False,
)
async def chat_stream( async def chat_stream(
self, self,
@@ -847,18 +1257,39 @@ class LLMProvider(ABC):
async def _safe_chat_stream(self, **kwargs: Any) -> LLMResponse: async def _safe_chat_stream(self, **kwargs: Any) -> LLMResponse:
"""Call chat_stream() and convert unexpected exceptions to error responses.""" """Call chat_stream() and convert unexpected exceptions to error responses."""
started_at_ms = time.time_ns() // 1_000_000
started_at_ns = time.monotonic_ns()
try: try:
provider_context = kwargs.pop("provider_context", None) provider_context = kwargs.pop("provider_context", None)
if isinstance(provider_context, ProviderCallContext): if isinstance(provider_context, ProviderCallContext):
return await self.chat_stream_with_context( response = await self.chat_stream_with_context(
provider_context=provider_context, provider_context=provider_context,
**kwargs, **kwargs,
) )
return await self.chat_stream(**kwargs) else:
response = await self.chat_stream(**kwargs)
except asyncio.CancelledError: except asyncio.CancelledError:
self._observe_llm_call(
LLMResponse(
content=None,
finish_reason="cancelled",
error_kind="cancelled",
),
kwargs,
started_at_ms=started_at_ms,
started_at_ns=started_at_ns,
stream=True,
)
raise raise
except Exception as exc: except Exception as exc:
return LLMResponse(content=f"Error calling LLM: {exc}", finish_reason="error") response = LLMResponse(content=f"Error calling LLM: {exc}", finish_reason="error")
return self._observe_llm_call(
response,
kwargs,
started_at_ms=started_at_ms,
started_at_ns=started_at_ns,
stream=True,
)
async def chat_stream_with_retry( async def chat_stream_with_retry(
self, self,
+21 -19
View File
@@ -14,6 +14,7 @@ from typing import Any, cast
from nanobot.providers.base import ( from nanobot.providers.base import (
LLMProvider, LLMProvider,
LLMResponse, LLMResponse,
LLMUsage,
ToolCallRequest, ToolCallRequest,
parse_tool_arguments, parse_tool_arguments,
resolve_stream_idle_timeout_s, resolve_stream_idle_timeout_s,
@@ -60,8 +61,9 @@ class BedrockProvider(LLMProvider):
profile: str | None = None, profile: str | None = None,
extra_body: dict[str, Any] | None = None, extra_body: dict[str, Any] | None = None,
client: Any | None = None, client: Any | None = None,
provider_name: str = "bedrock",
): ):
super().__init__(api_key, api_base) super().__init__(api_key, api_base, provider_name=provider_name)
self.default_model = default_model self.default_model = default_model
self.region = region or os.environ.get("AWS_REGION") or os.environ.get("AWS_DEFAULT_REGION") self.region = region or os.environ.get("AWS_REGION") or os.environ.get("AWS_DEFAULT_REGION")
self.profile = profile self.profile = profile
@@ -453,25 +455,25 @@ class BedrockProvider(LLMProvider):
}.get(stop_reason or "", stop_reason or "stop") }.get(stop_reason or "", stop_reason or "stop")
@staticmethod @staticmethod
def _usage(usage: dict[str, Any] | None) -> dict[str, int]: def _usage(usage: dict[str, Any] | None) -> LLMUsage | None:
if not usage: if not usage:
return {} return None
prompt = int(usage.get("inputTokens") or 0)
completion = int(usage.get("outputTokens") or 0) def _optional_count(key: str) -> int | None:
total = int(usage.get("totalTokens") or prompt + completion) raw = usage.get(key)
result = { return int(raw) if raw is not None else None
"prompt_tokens": prompt,
"completion_tokens": completion, cache_read = _optional_count("cacheReadInputTokens")
"total_tokens": total, cache_write = _optional_count("cacheWriteInputTokens")
} logical_input = int(usage.get("inputTokens") or 0) + (cache_read or 0) + (
cache_read = int(usage.get("cacheReadInputTokens") or 0) cache_write or 0
cache_write = int(usage.get("cacheWriteInputTokens") or 0) )
if cache_read: return LLMUsage.reported(
result["cached_tokens"] = cache_read input_tokens=logical_input,
result["cache_read_input_tokens"] = cache_read output_tokens=int(usage.get("outputTokens") or 0),
if cache_write: cache_read_tokens=cache_read,
result["cache_creation_input_tokens"] = cache_write cache_write_tokens=cache_write,
return result )
@staticmethod @staticmethod
def _parse_reasoning(block: dict[str, Any]) -> tuple[str | None, dict[str, Any] | None]: def _parse_reasoning(block: dict[str, Any]) -> tuple[str | None, dict[str, Any] | None]:
+7 -1
View File
@@ -172,6 +172,7 @@ def _make_provider_core(
default_model=model, default_model=model,
proxy=getattr(p, "proxy", None) if p else None, proxy=getattr(p, "proxy", None) if p else None,
extra_body=p.extra_body if p else None, extra_body=p.extra_body if p else None,
provider_name=provider_name,
) )
elif backend == "xai_grok": elif backend == "xai_grok":
from nanobot.providers.xai_grok_provider import XAIGrokProvider from nanobot.providers.xai_grok_provider import XAIGrokProvider
@@ -180,6 +181,7 @@ def _make_provider_core(
default_model=model, default_model=model,
proxy=getattr(p, "proxy", None) if p else None, proxy=getattr(p, "proxy", None) if p else None,
extra_body=p.extra_body if p else None, extra_body=p.extra_body if p else None,
provider_name=provider_name,
) )
elif backend == "azure_openai": elif backend == "azure_openai":
from nanobot.providers.azure_openai_provider import AzureOpenAIProvider from nanobot.providers.azure_openai_provider import AzureOpenAIProvider
@@ -190,11 +192,12 @@ def _make_provider_core(
api_key=p.api_key or "", api_key=p.api_key or "",
api_base=p.api_base, api_base=p.api_base,
default_model=model, default_model=model,
provider_name=provider_name,
) )
elif backend == "github_copilot": elif backend == "github_copilot":
from nanobot.providers.github_copilot_provider import GitHubCopilotProvider from nanobot.providers.github_copilot_provider import GitHubCopilotProvider
provider = GitHubCopilotProvider(default_model=model) provider = GitHubCopilotProvider(default_model=model, provider_name=provider_name)
elif backend == "anthropic": elif backend == "anthropic":
from nanobot.providers.anthropic_provider import AnthropicProvider from nanobot.providers.anthropic_provider import AnthropicProvider
@@ -203,6 +206,7 @@ def _make_provider_core(
api_base=config.get_api_base(model, preset=preset), api_base=config.get_api_base(model, preset=preset),
default_model=model, default_model=model,
extra_headers=_provider_extra_headers(spec, p), extra_headers=_provider_extra_headers(spec, p),
provider_name=provider_name,
) )
elif backend == "bedrock": elif backend == "bedrock":
from nanobot.providers.bedrock_provider import BedrockProvider from nanobot.providers.bedrock_provider import BedrockProvider
@@ -214,6 +218,7 @@ def _make_provider_core(
region=getattr(p, "region", None) if p else None, region=getattr(p, "region", None) if p else None,
profile=getattr(p, "profile", None) if p else None, profile=getattr(p, "profile", None) if p else None,
extra_body=p.extra_body if p else None, extra_body=p.extra_body if p else None,
provider_name=provider_name,
) )
else: else:
from nanobot.providers.openai_compat_provider import OpenAICompatProvider from nanobot.providers.openai_compat_provider import OpenAICompatProvider
@@ -228,6 +233,7 @@ def _make_provider_core(
api_type=p.api_type if p and provider_name == "openai" else "auto", api_type=p.api_type if p and provider_name == "openai" else "auto",
extra_query=p.extra_query if p else None, extra_query=p.extra_query if p else None,
proxy=p.proxy if p else None, proxy=p.proxy if p else None,
provider_name=provider_name,
) )
provider.generation = preset.to_generation_settings() provider.generation = preset.to_generation_settings()
+10
View File
@@ -13,6 +13,7 @@ from loguru import logger
from nanobot.providers.base import ( from nanobot.providers.base import (
GenerationSettings, GenerationSettings,
LLMCallObserver,
LLMProvider, LLMProvider,
LLMResponse, LLMResponse,
ProviderCallContext, ProviderCallContext,
@@ -124,7 +125,10 @@ class FallbackProvider(LLMProvider):
fallback_model_observer: FallbackModelObserver | None = None, fallback_model_observer: FallbackModelObserver | None = None,
primary_context_window_tokens: int | None = None, primary_context_window_tokens: int | None = None,
): ):
primary_generation = primary.generation
self._primary = primary self._primary = primary
super().__init__(provider_name=primary.provider_name)
self._primary.generation = primary_generation
self._fallback_presets = list(fallback_presets) self._fallback_presets = list(fallback_presets)
self._provider_factory = provider_factory self._provider_factory = provider_factory
self._fallback_model_observer = fallback_model_observer self._fallback_model_observer = fallback_model_observer
@@ -148,6 +152,11 @@ class FallbackProvider(LLMProvider):
"""Attach a process-level observer without changing request call signatures.""" """Attach a process-level observer without changing request call signatures."""
self._fallback_model_observer = observer self._fallback_model_observer = observer
def set_llm_call_observer(self, observer: LLMCallObserver | None) -> None:
"""Attach usage recording to the primary and future fallback leaves."""
super().set_llm_call_observer(observer)
self._primary.set_llm_call_observer(observer)
@property @property
def supports_progress_deltas(self) -> bool: def supports_progress_deltas(self) -> bool:
return bool(getattr(self._primary, "supports_progress_deltas", False)) return bool(getattr(self._primary, "supports_progress_deltas", False))
@@ -503,6 +512,7 @@ class FallbackProvider(LLMProvider):
) )
try: try:
fallback_provider = self._provider_factory(fallback) fallback_provider = self._provider_factory(fallback)
fallback_provider.set_llm_call_observer(self._llm_call_observer)
except Exception as exc: except Exception as exc:
logger.warning( logger.warning(
"Failed to create provider for fallback '{}': {}", fallback_model, exc "Failed to create provider for fallback '{}': {}", fallback_model, exc
+7 -1
View File
@@ -174,7 +174,12 @@ def login_github_copilot(
class GitHubCopilotProvider(OpenAICompatProvider): class GitHubCopilotProvider(OpenAICompatProvider):
"""Provider that exchanges a stored GitHub OAuth token for Copilot access tokens.""" """Provider that exchanges a stored GitHub OAuth token for Copilot access tokens."""
def __init__(self, default_model: str = "github-copilot/gpt-4.1"): def __init__(
self,
default_model: str = "github-copilot/gpt-4.1",
*,
provider_name: str = "github_copilot",
):
from nanobot.providers.registry import find_by_name from nanobot.providers.registry import find_by_name
self._copilot_access_token: str | None = None self._copilot_access_token: str | None = None
@@ -190,6 +195,7 @@ class GitHubCopilotProvider(OpenAICompatProvider):
"User-Agent": USER_AGENT, "User-Agent": USER_AGENT,
}, },
spec=find_by_name("github_copilot"), spec=find_by_name("github_copilot"),
provider_name=provider_name,
) )
async def _get_copilot_access_token(self) -> str: async def _get_copilot_access_token(self) -> str:
+2 -2
View File
@@ -20,7 +20,7 @@ from nanobot.providers.registry import find_by_name
from nanobot.security.network import ( from nanobot.security.network import (
PinnedDNSAsyncTransport, PinnedDNSAsyncTransport,
UnsafeURLRequestError, UnsafeURLRequestError,
resolve_url_target, async_resolve_url_target,
) )
from nanobot.utils.helpers import detect_image_mime from nanobot.utils.helpers import detect_image_mime
@@ -174,7 +174,7 @@ async def _download_image_data_url(
current_url = url current_url = url
for _ in range(_IMAGE_DOWNLOAD_MAX_REDIRECTS + 1): for _ in range(_IMAGE_DOWNLOAD_MAX_REDIRECTS + 1):
if proxy: if proxy:
ok, error, _ = resolve_url_target( ok, error, _ = await async_resolve_url_target(
current_url, current_url,
trust_remote_dns=True, trust_remote_dns=True,
) )
+19 -4
View File
@@ -7,6 +7,7 @@ from __future__ import annotations
import asyncio import asyncio
import hashlib import hashlib
import json import json
import ssl
from collections.abc import Awaitable, Callable from collections.abc import Awaitable, Callable
from typing import Any, cast from typing import Any, cast
@@ -50,12 +51,26 @@ class OpenAICodexProvider(LLMProvider):
default_model: str = "openai-codex/gpt-5.6-sol", default_model: str = "openai-codex/gpt-5.6-sol",
proxy: str | None = None, proxy: str | None = None,
extra_body: dict[str, Any] | None = None, extra_body: dict[str, Any] | None = None,
*,
provider_name: str = "openai_codex",
): ):
super().__init__(api_key=None, api_base=None) super().__init__(api_key=None, api_base=None, provider_name=provider_name)
self.default_model = default_model self.default_model = default_model
self.proxy = proxy or None self.proxy = proxy or None
self._extra_body = dict(extra_body or {}) self._extra_body = dict(extra_body or {})
self._native_compaction_available = True self._native_compaction_available = True
self._ssl_contexts: dict[bool, ssl.SSLContext] = {}
def _ssl_context(self, *, verify: bool) -> ssl.SSLContext:
"""Reuse synchronous TLS setup across requests on the shared event loop."""
context = self._ssl_contexts.get(verify)
if context is None:
context = httpx.create_ssl_context(
verify=verify,
trust_env=self.proxy is None,
)
self._ssl_contexts[verify] = context
return context
async def _call_codex( async def _call_codex(
self, self,
@@ -129,7 +144,7 @@ class OpenAICodexProvider(LLMProvider):
DEFAULT_CODEX_URL, DEFAULT_CODEX_URL,
headers, headers,
wire_body, wire_body,
verify=True, verify=self._ssl_context(verify=True),
proxy=self.proxy, proxy=self.proxy,
on_content_delta=on_content_delta if emit_deltas else None, on_content_delta=on_content_delta if emit_deltas else None,
on_thinking_delta=on_thinking_delta if emit_deltas else None, on_thinking_delta=on_thinking_delta if emit_deltas else None,
@@ -145,7 +160,7 @@ class OpenAICodexProvider(LLMProvider):
DEFAULT_CODEX_URL, DEFAULT_CODEX_URL,
headers, headers,
wire_body, wire_body,
verify=False, verify=self._ssl_context(verify=False),
proxy=self.proxy, proxy=self.proxy,
on_content_delta=on_content_delta if emit_deltas else None, on_content_delta=on_content_delta if emit_deltas else None,
on_thinking_delta=on_thinking_delta if emit_deltas else None, on_thinking_delta=on_thinking_delta if emit_deltas else None,
@@ -411,7 +426,7 @@ async def _request_codex(
url: str, url: str,
headers: dict[str, str], headers: dict[str, str],
body: dict[str, Any], body: dict[str, Any],
verify: bool, verify: ssl.SSLContext | bool,
proxy: str | None = None, proxy: str | None = None,
on_content_delta: Callable[[str], Awaitable[None]] | None = None, on_content_delta: Callable[[str], Awaitable[None]] | None = None,
on_thinking_delta: Callable[[str], Awaitable[None]] | None = None, on_thinking_delta: Callable[[str], Awaitable[None]] | None = None,
+40 -25
View File
@@ -26,6 +26,7 @@ from pydantic.alias_generators import to_snake
from nanobot.providers.base import ( from nanobot.providers.base import (
LLMProvider, LLMProvider,
LLMResponse, LLMResponse,
LLMUsage,
ProviderCallContext, ProviderCallContext,
ProviderConversationState, ProviderConversationState,
ToolCallRequest, ToolCallRequest,
@@ -517,8 +518,9 @@ class OpenAICompatProvider(LLMProvider):
api_type: str = "auto", api_type: str = "auto",
extra_query: dict[str, str] | None = None, extra_query: dict[str, str] | None = None,
proxy: str | None = None, proxy: str | None = None,
provider_name: str = "openai",
): ):
super().__init__(api_key, api_base) super().__init__(api_key, api_base, provider_name=provider_name)
self.default_model = default_model self.default_model = default_model
self.extra_headers = extra_headers or {} self.extra_headers = extra_headers or {}
self._spec = spec self._spec = spec
@@ -1428,12 +1430,12 @@ class OpenAICompatProvider(LLMProvider):
return "".join(parts) or None return "".join(parts) or None
@classmethod @classmethod
def _extract_usage(cls, response: Any) -> dict[str, int]: def _extract_usage(cls, response: Any) -> LLMUsage | None:
"""Extract token usage from an OpenAI-compatible response. """Extract token usage from an OpenAI-compatible response.
Handles both dict-based (raw JSON) and object-based (SDK Pydantic) Handles both dict-based (raw JSON) and object-based (SDK Pydantic)
responses. Provider-specific ``cached_tokens`` fields are normalised responses. Provider-specific cache fields are normalized once at
under a single key; see the priority chain inside for details. this Chat Completions wire boundary.
""" """
# --- resolve usage object --- # --- resolve usage object ---
usage_obj = None usage_obj = None
@@ -1445,21 +1447,18 @@ class OpenAICompatProvider(LLMProvider):
usage_map = cls._maybe_mapping(usage_obj) usage_map = cls._maybe_mapping(usage_obj)
if usage_map is not None: if usage_map is not None:
result = { input_tokens = int(usage_map.get("prompt_tokens") or 0)
"prompt_tokens": int(usage_map.get("prompt_tokens") or 0), output_tokens = int(usage_map.get("completion_tokens") or 0)
"completion_tokens": int(usage_map.get("completion_tokens") or 0),
"total_tokens": int(usage_map.get("total_tokens") or 0),
}
elif usage_obj: elif usage_obj:
result = { input_tokens = int(getattr(usage_obj, "prompt_tokens", 0) or 0)
"prompt_tokens": getattr(usage_obj, "prompt_tokens", 0) or 0, output_tokens = int(getattr(usage_obj, "completion_tokens", 0) or 0)
"completion_tokens": getattr(usage_obj, "completion_tokens", 0) or 0,
"total_tokens": getattr(usage_obj, "total_tokens", 0) or 0,
}
else: else:
return {} return None
# --- cached_tokens (normalised across providers) --- wire_total = cls._get_nested_int(usage_obj, ("total_tokens",))
cache_read: int | None = None
# --- cached_tokens (normalised across Chat-compatible providers) ---
# Try nested paths first (dict), fall back to attribute (SDK object). # Try nested paths first (dict), fall back to attribute (SDK object).
# Priority order ensures the most specific field wins. # Priority order ensures the most specific field wins.
for path in ( for path in (
@@ -1468,17 +1467,28 @@ class OpenAICompatProvider(LLMProvider):
("prompt_cache_hit_tokens",), # DeepSeek/SiliconFlow ("prompt_cache_hit_tokens",), # DeepSeek/SiliconFlow
): ):
cached = cls._get_nested_int(usage_map, path) cached = cls._get_nested_int(usage_map, path)
if not cached and usage_obj: if cached is None and usage_obj:
cached = cls._get_nested_int(usage_obj, path) cached = cls._get_nested_int(usage_obj, path)
if cached: if cached is not None:
result["cached_tokens"] = cached cache_read = cached
break break
return result cache_write = cls._get_nested_int(
usage_obj,
("prompt_tokens_details", "cache_write_tokens"),
)
return LLMUsage.reported(
input_tokens=input_tokens,
output_tokens=output_tokens,
total_tokens=wire_total,
cache_read_tokens=cache_read,
cache_write_tokens=cache_write,
)
@staticmethod @staticmethod
def _get_nested_int(obj: object, path: tuple[str, ...]) -> int: def _get_nested_int(obj: object, path: tuple[str, ...]) -> int | None:
"""Drill into *obj* by *path* segments and return an ``int`` value. """Return a present usage count while preserving explicit zero.
Supports both dict-key access and attribute access so it works Supports both dict-key access and attribute access so it works
uniformly with raw JSON dicts **and** SDK Pydantic models. uniformly with raw JSON dicts **and** SDK Pydantic models.
@@ -1486,12 +1496,17 @@ class OpenAICompatProvider(LLMProvider):
current: object = obj current: object = obj
for segment in path: for segment in path:
if current is None: if current is None:
return 0 return None
if isinstance(current, dict): if isinstance(current, dict):
current = cast(dict[str, Any], current).get(segment) current = cast(dict[str, Any], current).get(segment)
else: else:
current = getattr(current, segment, None) current = getattr(current, segment, None)
return int(cast(Any, current) or 0) if current is not None else 0 if current is None or isinstance(current, bool):
return None
try:
return int(cast(Any, current))
except (TypeError, ValueError):
return None
def _parse(self, response: Any) -> LLMResponse: def _parse(self, response: Any) -> LLMResponse:
if isinstance(response, str): if isinstance(response, str):
@@ -1645,7 +1660,7 @@ class OpenAICompatProvider(LLMProvider):
reasoning_parts: list[str] = [] reasoning_parts: list[str] = []
tc_bufs: dict[int, dict[str, Any]] = {} tc_bufs: dict[int, dict[str, Any]] = {}
finish_reason = "stop" finish_reason = "stop"
usage: dict[str, int] = {} usage: LLMUsage | None = None
def _accum_tc(tc: Any, idx_hint: int) -> None: def _accum_tc(tc: Any, idx_hint: int) -> None:
"""Accumulate one streaming tool-call delta into *tc_bufs*.""" """Accumulate one streaming tool-call delta into *tc_bufs*."""
+31 -37
View File
@@ -10,7 +10,7 @@ from typing import Any, AsyncGenerator, cast
import httpx import httpx
from loguru import logger from loguru import logger
from nanobot.providers.base import LLMResponse, ToolCallRequest, parse_tool_arguments from nanobot.providers.base import LLMResponse, LLMUsage, ToolCallRequest, parse_tool_arguments
from nanobot.providers.openai_responses.state import build_responses_state from nanobot.providers.openai_responses.state import build_responses_state
FINISH_REASON_MAP = { FINISH_REASON_MAP = {
@@ -186,33 +186,40 @@ def _response_finish_reason(
return map_finish_reason(terminal_status) return map_finish_reason(terminal_status)
def _usage_from_response_obj(response: object) -> dict[str, int]: def _usage_from_response_obj(response: object) -> LLMUsage | None:
response_object = _response_object(response) response_object = _response_object(response)
usage_raw: object = ( usage_raw: object = (
response_object.get("usage") response_object.get("usage")
if response_object is not None if response_object is not None
else getattr(response, "usage", None) else getattr(response, "usage", None)
) )
if not usage_raw: if usage_raw is None:
return {} return None
usage = _response_object(usage_raw) usage = _response_object(usage_raw)
if usage is None: if usage is None:
return {} return None
prompt_tokens = int(usage.get("input_tokens") or usage.get("prompt_tokens") or 0)
completion_tokens = int( def _usage_int(container: dict[str, Any] | None, key: str) -> int | None:
usage.get("output_tokens") or usage.get("completion_tokens") or 0 if container is None:
) return None
total_tokens = int(usage.get("total_tokens") or prompt_tokens + completion_tokens) raw = container.get(key)
result = { if raw is None or isinstance(raw, bool):
"prompt_tokens": prompt_tokens, return None
"completion_tokens": completion_tokens, try:
"total_tokens": total_tokens, return int(raw)
} except (TypeError, ValueError):
return None
input_tokens = _usage_int(usage, "input_tokens") or 0
output_tokens = _usage_int(usage, "output_tokens") or 0
input_details = _response_object(usage.get("input_tokens_details")) input_details = _response_object(usage.get("input_tokens_details"))
cached_tokens = int(input_details.get("cached_tokens") or 0) if input_details else 0 return LLMUsage.reported(
if cached_tokens > 0: input_tokens=input_tokens,
result["cached_tokens"] = cached_tokens output_tokens=output_tokens,
return result total_tokens=_usage_int(usage, "total_tokens"),
cache_read_tokens=_usage_int(input_details, "cached_tokens"),
cache_write_tokens=_usage_int(input_details, "cache_write_tokens"),
)
def _parse_tool_call_arguments(args_raw: Any, name: str | None) -> Any: def _parse_tool_call_arguments(args_raw: Any, name: str | None) -> Any:
@@ -352,14 +359,14 @@ async def consume_sse_with_reasoning(
on_reasoning_delta: Callable[[str], Awaitable[None]] | None = None, on_reasoning_delta: Callable[[str], Awaitable[None]] | None = None,
on_response_event: Callable[[dict[str, Any]], Awaitable[None]] | None = None, on_response_event: Callable[[dict[str, Any]], Awaitable[None]] | None = None,
capture: ResponsesStreamCapture | None = None, capture: ResponsesStreamCapture | None = None,
) -> tuple[str, list[ToolCallRequest], str, dict[str, int], str | None]: ) -> tuple[str, list[ToolCallRequest], str, LLMUsage | None, str | None]:
"""Consume a Responses API SSE stream, including visible reasoning summaries.""" """Consume a Responses API SSE stream, including visible reasoning summaries."""
content = "" content = ""
tool_calls: list[ToolCallRequest] = [] tool_calls: list[ToolCallRequest] = []
tool_call_buffers: dict[str, dict[str, Any]] = {} tool_call_buffers: dict[str, dict[str, Any]] = {}
tool_call_args_emitted: set[str] = set() tool_call_args_emitted: set[str] = set()
finish_reason = "stop" finish_reason = "stop"
usage: dict[str, int] = {} usage: LLMUsage | None = None
reasoning_content: str | None = None reasoning_content: str | None = None
streamed_reasoning = False streamed_reasoning = False
reasoning_summary_key: tuple[str | None, int] | None = None reasoning_summary_key: tuple[str | None, int] | None = None
@@ -657,14 +664,14 @@ async def consume_sdk_stream(
on_tool_call_delta: Callable[[dict[str, Any]], Awaitable[None]] | None = None, on_tool_call_delta: Callable[[dict[str, Any]], Awaitable[None]] | None = None,
on_reasoning_delta: Callable[[str], Awaitable[None]] | None = None, on_reasoning_delta: Callable[[str], Awaitable[None]] | None = None,
capture: ResponsesStreamCapture | None = None, capture: ResponsesStreamCapture | None = None,
) -> tuple[str, list[ToolCallRequest], str, dict[str, int], str | None]: ) -> tuple[str, list[ToolCallRequest], str, LLMUsage | None, str | None]:
"""Consume an SDK async stream from ``client.responses.create(stream=True)``.""" """Consume an SDK async stream from ``client.responses.create(stream=True)``."""
content = "" content = ""
tool_calls: list[ToolCallRequest] = [] tool_calls: list[ToolCallRequest] = []
tool_call_buffers: dict[str, dict[str, Any]] = {} tool_call_buffers: dict[str, dict[str, Any]] = {}
tool_call_args_emitted: set[str] = set() tool_call_args_emitted: set[str] = set()
finish_reason = "stop" finish_reason = "stop"
usage: dict[str, int] = {} usage: LLMUsage | None = None
reasoning_content: str | None = None reasoning_content: str | None = None
streamed_reasoning = False streamed_reasoning = False
refusal_seen = False refusal_seen = False
@@ -823,20 +830,7 @@ async def consume_sdk_stream(
if on_content_delta and remaining_text: if on_content_delta and remaining_text:
await on_content_delta(remaining_text) await on_content_delta(remaining_text)
if resp: if resp:
usage_obj = getattr(resp, "usage", None) usage = _usage_from_response_obj(resp) or usage
if usage_obj:
usage = {
"prompt_tokens": int(getattr(usage_obj, "input_tokens", 0) or 0),
"completion_tokens": int(getattr(usage_obj, "output_tokens", 0) or 0),
"total_tokens": int(getattr(usage_obj, "total_tokens", 0) or 0),
}
usage_data = _response_object(usage_obj) or {}
input_details = _response_object(usage_data.get("input_tokens_details"))
cached_tokens = (
int(input_details.get("cached_tokens") or 0) if input_details else 0
)
if cached_tokens > 0:
usage["cached_tokens"] = cached_tokens
if not reasoning_content: if not reasoning_content:
reasoning_content = _extract_reasoning_summary_from_output( reasoning_content = _extract_reasoning_summary_from_output(
getattr(resp, "output", None) getattr(resp, "output", None)
+4 -12
View File
@@ -7,7 +7,7 @@ from typing import Any, cast
from loguru import logger from loguru import logger
from nanobot.providers.base import ProviderConversationState from nanobot.providers.base import LLMUsage, ProviderConversationState
from nanobot.providers.openai_responses.converters import convert_messages from nanobot.providers.openai_responses.converters import convert_messages
RESPONSES_STATE_KIND = "openai_responses" RESPONSES_STATE_KIND = "openai_responses"
@@ -84,7 +84,7 @@ def build_responses_state(
model: str, model: str,
input_items: list[dict[str, Any]], input_items: list[dict[str, Any]],
output_items: list[dict[str, Any]], output_items: list[dict[str, Any]],
usage: dict[str, int] | None = None, usage: LLMUsage | None = None,
) -> ProviderConversationState: ) -> ProviderConversationState:
"""Create the canonical next state from request input and every output item.""" """Create the canonical next state from request input and every output item."""
unpruned_items = [*input_items, *output_items] unpruned_items = [*input_items, *output_items]
@@ -178,16 +178,8 @@ def _prune_before_latest_output_compaction(
return output_items[latest:] return output_items[latest:]
def _context_tokens_from_usage(usage: dict[str, int] | None) -> int: def _context_tokens_from_usage(usage: LLMUsage | None) -> int:
if not usage: return usage.total_tokens if usage is not None else 0
return 0
prompt_tokens = usage.get("prompt_tokens", 0)
completion_tokens = usage.get("completion_tokens", 0)
total_tokens = usage.get("total_tokens", 0)
values = (prompt_tokens, completion_tokens, total_tokens)
if any(isinstance(value, bool) for value in values):
return 0
return max(0, total_tokens or prompt_tokens + completion_tokens)
def _state_items( def _state_items(
+1 -1
View File
@@ -11,7 +11,7 @@ class UnconfiguredProvider(LLMProvider):
"""Keep the gateway available for settings before a model is configured.""" """Keep the gateway available for settings before a model is configured."""
def __init__(self, default_model: str) -> None: def __init__(self, default_model: str) -> None:
super().__init__() super().__init__(provider_name="unconfigured")
self._default_model = default_model self._default_model = default_model
async def chat( async def chat(
+5 -2
View File
@@ -18,6 +18,7 @@ from nanobot import __version__
from nanobot.providers.base import ( from nanobot.providers.base import (
LLMProvider, LLMProvider,
LLMResponse, LLMResponse,
LLMUsage,
ToolCallRequest, ToolCallRequest,
resolve_stream_idle_timeout_s, resolve_stream_idle_timeout_s,
) )
@@ -69,8 +70,10 @@ class XAIGrokProvider(LLMProvider):
default_model: str = DEFAULT_XAI_GROK_MODEL, default_model: str = DEFAULT_XAI_GROK_MODEL,
proxy: str | None = None, proxy: str | None = None,
extra_body: dict[str, Any] | None = None, extra_body: dict[str, Any] | None = None,
*,
provider_name: str = "xai_grok",
): ):
super().__init__(api_key=None, api_base=None) super().__init__(api_key=None, api_base=None, provider_name=provider_name)
self.default_model = default_model self.default_model = default_model
self.proxy = proxy or None self.proxy = proxy or None
self._extra_body = dict(extra_body or {}) self._extra_body = dict(extra_body or {})
@@ -436,7 +439,7 @@ async def _request_xai(
on_content_delta: Callable[[str], Awaitable[None]] | None = None, on_content_delta: Callable[[str], Awaitable[None]] | None = None,
on_thinking_delta: Callable[[str], Awaitable[None]] | None = None, on_thinking_delta: Callable[[str], Awaitable[None]] | None = None,
on_tool_call_delta: Callable[[dict[str, Any]], Awaitable[None]] | None = None, on_tool_call_delta: Callable[[dict[str, Any]], Awaitable[None]] | None = None,
) -> tuple[str, list[ToolCallRequest], str, dict[str, int], str | None]: ) -> tuple[str, list[ToolCallRequest], str, LLMUsage | None, str | None]:
async def _on_response_event(event: dict[str, Any]) -> None: async def _on_response_event(event: dict[str, Any]) -> None:
hosted_event = _xai_hosted_tool_event(event) hosted_event = _xai_hosted_tool_event(event)
if hosted_event is not None and on_tool_call_delta is not None: if hosted_event is not None and on_tool_call_delta is not None:
+7 -5
View File
@@ -210,18 +210,20 @@ class RuntimeClient:
async def compact_session(self, session_key: str) -> SessionSnapshot: async def compact_session(self, session_key: str) -> SessionSnapshot:
"""Run token consolidation for one session.""" """Run token consolidation for one session."""
session = self._loop.sessions.get_or_create(session_key) session = await self._loop.sessions.get_or_create_async(session_key)
runtime = self._loop.runtime_for_session(session) runtime = await self._loop.runtime_for_session_async(session)
await self._loop.consolidator.maybe_consolidate_by_tokens( await self._loop.consolidator.maybe_consolidate_by_tokens(
session, session,
runtime=runtime, runtime=runtime,
) )
return snapshot_from_session(self._loop.sessions.get_or_create(session_key)) return snapshot_from_session(
await self._loop.sessions.get_or_create_async(session_key)
)
async def compact_idle_session(self, session_key: str, *, max_suffix: int = 8) -> str | None: async def compact_idle_session(self, session_key: str, *, max_suffix: int = 8) -> str | None:
"""Run idle-session compaction for one session and return the summary.""" """Run idle-session compaction for one session and return the summary."""
session = self._loop.sessions.get_or_create(session_key) session = await self._loop.sessions.get_or_create_async(session_key)
runtime = self._loop.runtime_for_session(session) runtime = await self._loop.runtime_for_session_async(session)
return await self._loop.consolidator.compact_idle_session( return await self._loop.consolidator.compact_idle_session(
session_key, session_key,
runtime=runtime, runtime=runtime,
+3 -2
View File
@@ -6,6 +6,7 @@ from copy import deepcopy
from dataclasses import dataclass, field from dataclasses import dataclass, field
from typing import Any, Literal, Mapping, TypeAlias, cast from typing import Any, Literal, Mapping, TypeAlias, cast
from nanobot.providers.base import LLMUsage
from nanobot.runtime_context import public_history_messages from nanobot.runtime_context import public_history_messages
StreamEventType: TypeAlias = Literal[ StreamEventType: TypeAlias = Literal[
@@ -53,7 +54,7 @@ class RunResult:
content: str content: str
tools_used: list[str] = field(default_factory=list) tools_used: list[str] = field(default_factory=list)
messages: list[dict[str, Any]] = field(default_factory=list) messages: list[dict[str, Any]] = field(default_factory=list)
usage: dict[str, int] = field(default_factory=dict) usage: LLMUsage | None = None
stop_reason: str | None = None stop_reason: str | None = None
error: str | None = None error: str | None = None
metadata: dict[str, Any] = field(default_factory=dict) metadata: dict[str, Any] = field(default_factory=dict)
@@ -72,7 +73,7 @@ class StreamEvent:
arguments: dict[str, Any] | None = None arguments: dict[str, Any] | None = None
iteration: int | None = None iteration: int | None = None
resuming: bool | None = None resuming: bool | None = None
usage: dict[str, int] = field(default_factory=dict) usage: LLMUsage | None = None
error: str | None = None error: str | None = None
metadata: dict[str, Any] = field(default_factory=dict) metadata: dict[str, Any] = field(default_factory=dict)
+143 -41
View File
@@ -29,6 +29,7 @@ _BLOCKED_NETWORKS = [
_URL_RE = re.compile(r"https?://[^\s\"'`;|<>]+", re.IGNORECASE) _URL_RE = re.compile(r"https?://[^\s\"'`;|<>]+", re.IGNORECASE)
_allowed_networks: list[ipaddress.IPv4Network | ipaddress.IPv6Network] = [] _allowed_networks: list[ipaddress.IPv4Network | ipaddress.IPv6Network] = []
_DNS_RESOLUTION_TIMEOUT_SECONDS = 5.0
def is_loopback_host(host: str) -> bool: def is_loopback_host(host: str) -> bool:
@@ -75,6 +76,63 @@ def _is_private(addr: ipaddress.IPv4Address | ipaddress.IPv6Address) -> bool:
return any(normalized in net for net in _BLOCKED_NETWORKS) return any(normalized in net for net in _BLOCKED_NETWORKS)
def _parse_url_hostname(url: str) -> tuple[str | None, str | None]:
try:
parsed = urlparse(url)
except Exception as exc:
return None, str(exc)
if parsed.scheme not in ("http", "https"):
return None, f"Only http/https allowed, got '{parsed.scheme or 'none'}'"
if not parsed.netloc:
return None, "Missing domain"
if not parsed.hostname:
return None, "Missing hostname"
return parsed.hostname, None
def _unresolved_target_result(
hostname: str,
*,
trust_remote_dns: bool,
) -> tuple[bool, str, tuple[str, ...]]:
if not trust_remote_dns:
return False, f"Cannot resolve hostname: {hostname}", ()
normalized_hostname = hostname.rstrip(".").lower()
if normalized_hostname == "localhost" or normalized_hostname.endswith(".localhost"):
return False, f"Blocked local/internal hostname: {hostname}", ()
try:
literal_addr = ipaddress.ip_address(normalized_hostname)
except ValueError:
return True, "", ()
if _is_private(literal_addr):
return False, f"Blocked private/internal address: {literal_addr}", ()
return True, "", (str(_normalize_addr(literal_addr)),)
def _resolved_target_result(
hostname: str,
infos: list[Any],
*,
allow_loopback: bool,
) -> tuple[bool, str, tuple[str, ...]]:
addrs: list[ipaddress.IPv4Address | ipaddress.IPv6Address] = []
for info in infos:
try:
addr = ipaddress.ip_address(info[4][0])
except (IndexError, TypeError, ValueError):
continue
addrs.append(addr)
if allow_loopback and _is_allowed_loopback_target(hostname, addrs):
return True, "", tuple(dict.fromkeys(str(_normalize_addr(addr)) for addr in addrs))
for addr in addrs:
if _is_private(addr):
return False, f"Blocked: {hostname} resolves to private/internal address {addr}", ()
return True, "", tuple(dict.fromkeys(str(_normalize_addr(addr)) for addr in addrs))
def resolve_url_target( def resolve_url_target(
url: str, url: str,
*, *,
@@ -97,52 +155,43 @@ def resolve_url_target(
resolved_ips contains the public IPs that were validated for this URL, or resolved_ips contains the public IPs that were validated for this URL, or
is empty when an unresolved hostname is delegated to a trusted proxy. is empty when an unresolved hostname is delegated to a trusted proxy.
""" """
try: hostname, error = _parse_url_hostname(url)
p = urlparse(url) if hostname is None:
except Exception as e: return False, error or "Missing hostname", ()
return False, str(e), ()
if p.scheme not in ("http", "https"):
return False, f"Only http/https allowed, got '{p.scheme or 'none'}'", ()
if not p.netloc:
return False, "Missing domain", ()
hostname = p.hostname
if not hostname:
return False, "Missing hostname", ()
try: try:
infos = socket.getaddrinfo(hostname, None, socket.AF_UNSPEC, socket.SOCK_STREAM) infos = socket.getaddrinfo(hostname, None, socket.AF_UNSPEC, socket.SOCK_STREAM)
except socket.gaierror: except socket.gaierror:
if not trust_remote_dns: return _unresolved_target_result(hostname, trust_remote_dns=trust_remote_dns)
return False, f"Cannot resolve hostname: {hostname}", () return _resolved_target_result(hostname, infos, allow_loopback=allow_loopback)
normalized_hostname = hostname.rstrip(".").lower()
if normalized_hostname == "localhost" or normalized_hostname.endswith(".localhost"):
return False, f"Blocked local/internal hostname: {hostname}", ()
async def async_resolve_url_target(
url: str,
*,
allow_loopback: bool = False,
trust_remote_dns: bool = False,
timeout_s: float = _DNS_RESOLUTION_TIMEOUT_SECONDS,
) -> tuple[bool, str, tuple[str, ...]]:
"""Resolve and validate an HTTP target without blocking the event loop."""
hostname, error = _parse_url_hostname(url)
if hostname is None:
return False, error or "Missing hostname", ()
loop = asyncio.get_running_loop()
try: try:
literal_addr = ipaddress.ip_address(normalized_hostname) infos = await asyncio.wait_for(
except ValueError: loop.getaddrinfo(
return True, "", () hostname,
if _is_private(literal_addr): None,
return False, f"Blocked private/internal address: {literal_addr}", () family=socket.AF_UNSPEC,
return True, "", (str(_normalize_addr(literal_addr)),) type=socket.SOCK_STREAM,
),
addrs: list[ipaddress.IPv4Address | ipaddress.IPv6Address] = [] timeout=timeout_s,
for info in infos: )
try: except asyncio.TimeoutError:
addr = ipaddress.ip_address(info[4][0]) return False, f"Timed out resolving hostname: {hostname}", ()
except ValueError: except socket.gaierror:
continue return _unresolved_target_result(hostname, trust_remote_dns=trust_remote_dns)
addrs.append(addr) return _resolved_target_result(hostname, infos, allow_loopback=allow_loopback)
if allow_loopback and _is_allowed_loopback_target(hostname, addrs):
return True, "", tuple(dict.fromkeys(str(_normalize_addr(addr)) for addr in addrs))
for addr in addrs:
if _is_private(addr):
return False, f"Blocked: {hostname} resolves to private/internal address {addr}", ()
return True, "", tuple(dict.fromkeys(str(_normalize_addr(addr)) for addr in addrs))
def validate_url_target(url: str, *, allow_loopback: bool = False) -> tuple[bool, str]: def validate_url_target(url: str, *, allow_loopback: bool = False) -> tuple[bool, str]:
@@ -151,6 +200,16 @@ def validate_url_target(url: str, *, allow_loopback: bool = False) -> tuple[bool
return ok, error return ok, error
async def async_validate_url_target(
url: str,
*,
allow_loopback: bool = False,
) -> tuple[bool, str]:
"""Validate a URL using the event loop's asynchronous resolver."""
ok, error, _ = await async_resolve_url_target(url, allow_loopback=allow_loopback)
return ok, error
def env_proxy_applies_to_url(url: str) -> bool: def env_proxy_applies_to_url(url: str) -> bool:
"""Return True when process proxy settings would proxy this URL.""" """Return True when process proxy settings would proxy this URL."""
try: try:
@@ -277,7 +336,10 @@ class PinnedDNSAsyncTransport(httpx.AsyncBaseTransport):
async def handle_async_request(self, request: httpx.Request) -> httpx.Response: async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
url = str(request.url) url = str(request.url)
ok, error, resolved_ips = resolve_url_target(url, allow_loopback=self._allow_loopback) ok, error, resolved_ips = await async_resolve_url_target(
url,
allow_loopback=self._allow_loopback,
)
if not ok: if not ok:
raise UnsafeURLRequestError(error, request=request) raise UnsafeURLRequestError(error, request=request)
async with self._resolver_lock: async with self._resolver_lock:
@@ -320,6 +382,46 @@ def validate_resolved_url(url: str) -> tuple[bool, str]:
return True, "" return True, ""
async def async_validate_resolved_url(url: str) -> tuple[bool, str]:
"""Validate a redirect target without blocking on domain resolution."""
try:
parsed = urlparse(url)
except Exception:
return True, ""
hostname = parsed.hostname
if not hostname:
return True, ""
try:
addr = ipaddress.ip_address(hostname)
except ValueError:
loop = asyncio.get_running_loop()
try:
infos = await asyncio.wait_for(
loop.getaddrinfo(
hostname,
None,
family=socket.AF_UNSPEC,
type=socket.SOCK_STREAM,
),
timeout=_DNS_RESOLUTION_TIMEOUT_SECONDS,
)
except asyncio.TimeoutError:
return False, f"Timed out resolving redirect hostname: {hostname}"
except socket.gaierror:
return True, ""
for info in infos:
try:
addr = ipaddress.ip_address(info[4][0])
except (IndexError, TypeError, ValueError):
continue
if _is_private(addr):
return False, f"Redirect target {hostname} resolves to private address {addr}"
return True, ""
if _is_private(addr):
return False, f"Redirect target is a private address: {addr}"
return True, ""
def contains_internal_url(command: str, *, allow_loopback: bool = False) -> bool: def contains_internal_url(command: str, *, allow_loopback: bool = False) -> bool:
"""Return True if the command string contains a URL targeting an internal/private address.""" """Return True if the command string contains a URL targeting an internal/private address."""
for m in _URL_RE.finditer(command): for m in _URL_RE.finditer(command):
+29
View File
@@ -0,0 +1,29 @@
"""Compatibility bridge for asynchronous SessionManager operations."""
import asyncio
import inspect
from collections.abc import Awaitable, Callable
from typing import Any, TypeVar, cast
from nanobot.utils.cancellation import shield_and_drain
_SessionResult = TypeVar("_SessionResult")
async def call_session_manager(
manager: object,
async_method_name: str,
sync_method: Callable[..., _SessionResult],
/,
*args: Any,
**kwargs: Any,
) -> _SessionResult:
"""Prefer a class-declared coroutine, or offload the established sync contract."""
class_async_method = inspect.getattr_static(type(manager), async_method_name, None)
if inspect.iscoroutinefunction(class_async_method):
async_method = cast(
Callable[..., Awaitable[_SessionResult]],
getattr(manager, async_method_name),
)
return await async_method(*args, **kwargs)
return await shield_and_drain(asyncio.to_thread(sync_method, *args, **kwargs))
+248 -3
View File
@@ -1,5 +1,6 @@
"""Session management for conversation history.""" """Session management for conversation history."""
import asyncio
import base64 import base64
import errno import errno
import hashlib import hashlib
@@ -28,6 +29,7 @@ from nanobot.runtime_context import (
public_history_message, public_history_message,
) )
from nanobot.session.model_selection import SESSION_MODEL_PRESET_METADATA_KEY from nanobot.session.model_selection import SESSION_MODEL_PRESET_METADATA_KEY
from nanobot.utils.cancellation import shield_and_drain
from nanobot.utils.helpers import ( from nanobot.utils.helpers import (
content_with_media_breadcrumbs, content_with_media_breadcrumbs,
ensure_dir, ensure_dir,
@@ -48,15 +50,21 @@ _SESSION_PREVIEW_MAX_CHARS = 120
_SESSION_LIST_PREVIEW_MAX_RECORDS = 200 _SESSION_LIST_PREVIEW_MAX_RECORDS = 200
_SESSION_LIST_PREVIEW_MAX_CHARS = 1_000_000 _SESSION_LIST_PREVIEW_MAX_CHARS = 1_000_000
_SESSION_DATA_ERRORS = (ValueError, TypeError, AttributeError, KeyError) _SESSION_DATA_ERRORS = (ValueError, TypeError, AttributeError, KeyError)
_RUNTIME_CHECKPOINT_DATA_ERRORS = (OSError, *_SESSION_DATA_ERRORS)
_PROVIDER_STATE_RECORD_TYPE = "provider_state" _PROVIDER_STATE_RECORD_TYPE = "provider_state"
_PROVIDER_STATE_RECORD_PREFIX_RE = re.compile( _PROVIDER_STATE_RECORD_PREFIX_RE = re.compile(
r'^\s*\{\s*"_type"\s*:\s*"provider_state"\s*(?:,|\})' r'^\s*\{\s*"_type"\s*:\s*"provider_state"\s*(?:,|\})'
) )
_RUNTIME_CHECKPOINT_KEY = "runtime_checkpoint"
_RUNTIME_CHECKPOINT_VERSION = 1
_RUNTIME_CHECKPOINT_SUFFIX = ".checkpoint.json"
_FORK_VOLATILE_METADATA_KEYS = { _FORK_VOLATILE_METADATA_KEYS = {
"goal_state", "goal_state",
"pending_user_turn", "pending_user_turn",
"pending_user_followups",
"runtime_checkpoint", "runtime_checkpoint",
"session_handle", "session_handle",
"webui_recovery",
"thread_goal", "thread_goal",
"title", "title",
"title_user_edited", "title_user_edited",
@@ -65,6 +73,7 @@ _WORKSPACE_STATE_DIR = ".nanobot"
_WORKSPACE_ID_FILE = "workspace-id" _WORKSPACE_ID_FILE = "workspace-id"
_WORKSPACE_ID_RE = re.compile(r"^[0-9a-f]{32}$") _WORKSPACE_ID_RE = re.compile(r"^[0-9a-f]{32}$")
_SESSION_MIGRATION_LOCK_TIMEOUT_SECONDS = 30 _SESSION_MIGRATION_LOCK_TIMEOUT_SECONDS = 30
_SESSION_FILES_LOCK_TIMEOUT_SECONDS = 5
_SESSION_FILES_LOCK_FILENAME = ".session-files.lock" _SESSION_FILES_LOCK_FILENAME = ".session-files.lock"
_COPY_CHUNK_SIZE = 1024 * 1024 _COPY_CHUNK_SIZE = 1024 * 1024
@@ -554,7 +563,8 @@ class JsonlSessionStore:
self.sessions_dir = ensure_dir(root / workspace_id) self.sessions_dir = ensure_dir(root / workspace_id)
self.legacy_sessions_dir = get_legacy_sessions_dir() self.legacy_sessions_dir = get_legacy_sessions_dir()
self._session_files_lock = FileLock( self._session_files_lock = FileLock(
str(self.sessions_dir / _SESSION_FILES_LOCK_FILENAME) str(self.sessions_dir / _SESSION_FILES_LOCK_FILENAME),
timeout=_SESSION_FILES_LOCK_TIMEOUT_SECONDS,
) )
with self._session_files_lock: with self._session_files_lock:
self._migrate_from_workspace(canonical_workspace) self._migrate_from_workspace(canonical_workspace)
@@ -1001,6 +1011,9 @@ class JsonlSessionStore:
def get_session_path(self, key: str) -> Path: def get_session_path(self, key: str) -> Path:
return self.sessions_dir / f"{self.storage_key(key)}.jsonl" return self.sessions_dir / f"{self.storage_key(key)}.jsonl"
def get_runtime_checkpoint_path(self, key: str) -> Path:
return self.sessions_dir / f"{self.storage_key(key)}{_RUNTIME_CHECKPOINT_SUFFIX}"
def get_legacy_lossy_path(self, key: str) -> Path: def get_legacy_lossy_path(self, key: str) -> Path:
return self.sessions_dir / f"{safe_filename(key.replace(':', '_'))}.jsonl" return self.sessions_dir / f"{safe_filename(key.replace(':', '_'))}.jsonl"
@@ -1066,7 +1079,7 @@ class JsonlSessionStore:
else: else:
messages.append(data) messages.append(data)
return Session( session = Session(
key=key, key=key,
messages=messages, messages=messages,
created_at=created_at or datetime.now(), created_at=created_at or datetime.now(),
@@ -1075,6 +1088,8 @@ class JsonlSessionStore:
last_consolidated=last_consolidated, last_consolidated=last_consolidated,
provider_state=provider_state, provider_state=provider_state,
) )
self._overlay_runtime_checkpoint_unlocked(session, path)
return session
except _SESSION_DATA_ERRORS as e: except _SESSION_DATA_ERRORS as e:
logger.warning("Failed to load session {}: {}", key, e) logger.warning("Failed to load session {}: {}", key, e)
repaired = self._repair_unlocked(key) repaired = self._repair_unlocked(key)
@@ -1159,7 +1174,7 @@ class JsonlSessionStore:
if not messages and not metadata and provider_state is None: if not messages and not metadata and provider_state is None:
return None return None
return Session( session = Session(
key=key, key=key,
messages=messages, messages=messages,
created_at=created_at or datetime.now(), created_at=created_at or datetime.now(),
@@ -1168,6 +1183,8 @@ class JsonlSessionStore:
last_consolidated=last_consolidated, last_consolidated=last_consolidated,
provider_state=provider_state, provider_state=provider_state,
) )
self._overlay_runtime_checkpoint_unlocked(session, path)
return session
except _SESSION_DATA_ERRORS as e: except _SESSION_DATA_ERRORS as e:
logger.warning("Repair failed for session {}: {}", key, e) logger.warning("Repair failed for session {}: {}", key, e)
return None return None
@@ -1186,6 +1203,105 @@ class JsonlSessionStore:
with self._session_files_lock: with self._session_files_lock:
self._save_unlocked(session, fsync=fsync) self._save_unlocked(session, fsync=fsync)
def save_runtime_checkpoint(self, session: Session) -> None:
"""Atomically persist only the volatile in-flight turn state.
A checkpoint is written several times during a tool-heavy turn. Keeping it
beside the append history avoids copying the full transcript at each safe
recovery boundary.
"""
with self._session_files_lock:
path = self.get_session_path(session.key)
if not path.exists():
# A user turn normally creates the session first. Internal callers
# may checkpoint a fresh session, so establish the durable base once.
self._save_unlocked(session)
return
checkpoint = session.metadata.get(_RUNTIME_CHECKPOINT_KEY)
if not isinstance(checkpoint, dict):
self.get_runtime_checkpoint_path(session.key).unlink(missing_ok=True)
return
payload: dict[str, Any] = {
"version": _RUNTIME_CHECKPOINT_VERSION,
"session_key": session.key,
"base_updated_at": session.updated_at.isoformat(),
"base_message_count": len(session.messages),
"checkpoint": checkpoint,
"provider_state": (
session.provider_state.to_private_record()
if session.provider_state is not None
else None
),
}
target = self.get_runtime_checkpoint_path(session.key)
tmp = target.with_name(f".{target.name}.{secrets.token_hex(8)}.tmp")
try:
with open(tmp, "x", encoding="utf-8") as handle:
os.chmod(tmp, 0o600)
json.dump(
payload,
handle,
ensure_ascii=False,
separators=(",", ":"),
)
os.replace(tmp, target)
finally:
tmp.unlink(missing_ok=True)
def _overlay_runtime_checkpoint_unlocked(self, session: Session, main_path: Path) -> None:
checkpoint_path = self.get_runtime_checkpoint_path(session.key)
try:
checkpoint_stat = checkpoint_path.lstat()
if not stat.S_ISREG(checkpoint_stat.st_mode):
logger.warning(
"Ignoring non-regular runtime checkpoint for session {}",
session.key,
)
return
# A complete session save supersedes an older sidecar. This comparison
# closes the small crash window between replacing the JSONL and unlinking
# its previous checkpoint.
if main_path.stat().st_mtime_ns > checkpoint_stat.st_mtime_ns:
checkpoint_path.unlink(missing_ok=True)
return
raw = _json_object(json.loads(checkpoint_path.read_text(encoding="utf-8")))
if (
raw.get("version") != _RUNTIME_CHECKPOINT_VERSION
or raw.get("session_key") != session.key
or raw.get("base_updated_at") != session.updated_at.isoformat()
or raw.get("base_message_count") != len(session.messages)
or not isinstance(raw.get("checkpoint"), dict)
):
checkpoint_path.unlink(missing_ok=True)
return
provider_record = raw.get("provider_state")
provider_state = (
None
if provider_record is None
else ProviderConversationState.from_private_record(provider_record)
)
if provider_record is not None and provider_state is None:
raise ValueError("invalid checkpoint provider state")
session.metadata[_RUNTIME_CHECKPOINT_KEY] = cast(
dict[str, Any], raw["checkpoint"]
)
session.provider_state = provider_state
except FileNotFoundError:
return
except _RUNTIME_CHECKPOINT_DATA_ERRORS as exc:
logger.warning(
"Ignoring invalid runtime checkpoint for session {}: {}",
session.key,
exc,
)
# Atomic writes mean a malformed target cannot become valid later.
# Remove it once so future loads do not repeatedly parse and log it.
with suppress(OSError):
if checkpoint_path.is_file() and not checkpoint_path.is_symlink():
checkpoint_path.unlink()
def _save_unlocked(self, session: Session, *, fsync: bool = False) -> None: def _save_unlocked(self, session: Session, *, fsync: bool = False) -> None:
path = self.get_session_path(session.key) path = self.get_session_path(session.key)
tmp_path = path.with_name(f".{path.name}.{secrets.token_hex(8)}.tmp") tmp_path = path.with_name(f".{path.name}.{secrets.token_hex(8)}.tmp")
@@ -1215,6 +1331,10 @@ class JsonlSessionStore:
os.replace(tmp_path, path) os.replace(tmp_path, path)
# The full record now contains the authoritative checkpoint state (or
# its removal), so an older volatile overlay is no longer needed.
self.get_runtime_checkpoint_path(session.key).unlink(missing_ok=True)
if fsync: if fsync:
with suppress(PermissionError): with suppress(PermissionError):
fd = os.open(str(path.parent), os.O_RDONLY) fd = os.open(str(path.parent), os.O_RDONLY)
@@ -1278,6 +1398,7 @@ class JsonlSessionStore:
def _delete_unlocked(self, key: str) -> bool: def _delete_unlocked(self, key: str) -> bool:
paths = [ paths = [
self.get_session_path(key), self.get_session_path(key),
self.get_runtime_checkpoint_path(key),
self.get_legacy_lossy_path(key), self.get_legacy_lossy_path(key),
self.get_legacy_session_path(key), self.get_legacy_session_path(key),
] ]
@@ -1525,6 +1646,7 @@ class SessionManager:
self._cache: OrderedDict[str, Session] = OrderedDict() self._cache: OrderedDict[str, Session] = OrderedDict()
# Preserve identity for sessions held by active callers without retaining idle ones. # Preserve identity for sessions held by active callers without retaining idle ones.
self._overflow_cache: WeakValueDictionary[str, Session] = WeakValueDictionary() self._overflow_cache: WeakValueDictionary[str, Session] = WeakValueDictionary()
self._async_session_locks: WeakValueDictionary[str, asyncio.Lock] = WeakValueDictionary()
self._max_cached_sessions = SESSION_CACHE_MAX_SIZE self._max_cached_sessions = SESSION_CACHE_MAX_SIZE
self._delete_observer: Callable[[str], None] | None = None self._delete_observer: Callable[[str], None] | None = None
@@ -1585,6 +1707,10 @@ class SessionManager:
"""Get the collision-resistant workspace path for a session.""" """Get the collision-resistant workspace path for a session."""
return self._jsonl_store.get_session_path(key) return self._jsonl_store.get_session_path(key)
def _get_runtime_checkpoint_path(self, key: str) -> Path:
"""Get the private in-flight checkpoint path for a session."""
return self._jsonl_store.get_runtime_checkpoint_path(key)
def _get_legacy_lossy_path(self, key: str) -> Path: def _get_legacy_lossy_path(self, key: str) -> Path:
"""Previous workspace session path using lossy ':' to '_' replacement.""" """Previous workspace session path using lossy ':' to '_' replacement."""
return self._jsonl_store.get_legacy_lossy_path(key) return self._jsonl_store.get_legacy_lossy_path(key)
@@ -1620,6 +1746,28 @@ class SessionManager:
self._remember(session) self._remember(session)
return session return session
def _async_session_lock(self, key: str) -> asyncio.Lock:
lock = self._async_session_locks.get(key)
if lock is None:
lock = asyncio.Lock()
self._async_session_locks[key] = lock
return lock
async def get_or_create_async(self, key: str) -> Session:
"""Load a session without running file I/O or lock waits on the event loop."""
cached = self.get_cached(key)
if cached is not None:
return cached
async with self._async_session_lock(key):
cached = self.get_cached(key)
if cached is not None:
return cached
session = await asyncio.to_thread(self._load, key)
if session is None:
session = Session(key=key)
self._remember(session)
return session
def get_or_create_transient( def get_or_create_transient(
self, self,
key: str, key: str,
@@ -1653,6 +1801,46 @@ class SessionManager:
self._store.save(session, fsync=fsync) self._store.save(session, fsync=fsync)
self._remember(session) self._remember(session)
async def save_async(self, session: Session, *, fsync: bool = False) -> None:
"""Persist a session without blocking the caller's event loop."""
if not session.policy.persist:
return
async def save_and_remember() -> None:
await asyncio.to_thread(self._store.save, session, fsync=fsync)
self._remember(session)
await shield_and_drain(save_and_remember())
def save_runtime_checkpoint(self, session: Session) -> None:
"""Persist volatile recovery state without rewriting long history."""
if not session.policy.persist:
return
if self._store is self._jsonl_store:
self._jsonl_store.save_runtime_checkpoint(session)
self._remember(session)
return
# Third-party stores keep their existing all-or-nothing semantics until
# they opt into a dedicated checkpoint primitive.
self.save(session)
async def save_runtime_checkpoint_async(self, session: Session) -> None:
"""Persist an in-flight checkpoint without blocking the event loop."""
if not session.policy.persist:
return
async def save_and_remember() -> None:
if self._store is self._jsonl_store:
await asyncio.to_thread(
self._jsonl_store.save_runtime_checkpoint,
session,
)
else:
await asyncio.to_thread(self._store.save, session)
self._remember(session)
await shield_and_drain(save_and_remember())
def rename_model_preset(self, old_name: str, new_name: str) -> int: def rename_model_preset(self, old_name: str, new_name: str) -> int:
"""Rename a session-scoped model preset across durable and live sessions.""" """Rename a session-scoped model preset across durable and live sessions."""
if old_name == new_name: if old_name == new_name:
@@ -1694,6 +1882,21 @@ class SessionManager:
raise raise
return len(changed) return len(changed)
async def flush_all_async(self) -> int:
"""Re-save every cached session without blocking the event loop."""
cached = dict(self._overflow_cache.items())
cached.update(self._cache)
flushed = 0
for key, session in cached.items():
try:
await shield_and_drain(
asyncio.to_thread(self._store.save, session, fsync=True)
)
flushed += 1
except Exception:
logger.warning("Failed to flush session {}", key, exc_info=True)
return flushed
def flush_all(self) -> int: def flush_all(self) -> int:
"""Re-save every cached session with fsync for durable shutdown. """Re-save every cached session with fsync for durable shutdown.
@@ -1725,6 +1928,18 @@ class SessionManager:
self._delete_observer(key) self._delete_observer(key)
return deleted return deleted
async def delete_session_async(self, key: str) -> bool:
"""Delete a session without blocking the event loop."""
async def delete_and_notify() -> bool:
self.invalidate(key)
deleted = await asyncio.to_thread(self._store.delete, key)
if self._delete_observer is not None:
self._delete_observer(key)
return deleted
return await shield_and_drain(delete_and_notify())
def restore_sessions_to_workspace(self) -> SessionRestoreResult: def restore_sessions_to_workspace(self) -> SessionRestoreResult:
"""Restore session files to the pre-relocation path for an explicit rollback.""" """Restore session files to the pre-relocation path for an explicit rollback."""
return self._jsonl_store.restore_to_workspace() return self._jsonl_store.restore_to_workspace()
@@ -1797,6 +2012,10 @@ class SessionManager:
"""Read session metadata without loading the transcript.""" """Read session metadata without loading the transcript."""
return cast(dict[str, Any] | None, self._store.read_metadata(key)) return cast(dict[str, Any] | None, self._store.read_metadata(key))
async def read_session_metadata_async(self, key: str) -> dict[str, Any] | None:
"""Read session metadata without blocking the event loop."""
return await asyncio.to_thread(self.read_session_metadata, key)
def update_session_metadata( def update_session_metadata(
self, self,
key: str, key: str,
@@ -1810,5 +2029,31 @@ class SessionManager:
session.metadata.update(deepcopy(updates)) session.metadata.update(deepcopy(updates))
return updated return updated
async def update_session_metadata_async(
self,
key: str,
updates: dict[str, Any],
*,
fsync: bool = False,
) -> bool:
"""Update metadata without blocking the event loop."""
async def update_and_refresh_cache() -> bool:
updated = await asyncio.to_thread(
self._store.update_metadata,
key,
updates,
fsync=fsync,
)
if updated and (session := self.get_cached(key)) is not None:
session.metadata.update(deepcopy(updates))
return updated
return await shield_and_drain(update_and_refresh_cache())
def list_sessions(self) -> list[dict[str, Any]]: def list_sessions(self) -> list[dict[str, Any]]:
return cast(list[dict[str, Any]], self._store.list_sessions()) return cast(list[dict[str, Any]], self._store.list_sessions())
async def list_sessions_async(self) -> list[dict[str, Any]]:
"""List persisted sessions without blocking the event loop."""
return await asyncio.to_thread(self.list_sessions)
+971
View File
@@ -0,0 +1,971 @@
"""Durable, side-effect-safe recovery for interrupted WebUI turns.
The coordinator owns restart policy. AgentLoop only exposes checkpoint
materialization and an admission hook, so transport code never has to guess
whether an interrupted tool call is safe to replay.
"""
from __future__ import annotations
import asyncio
import dataclasses
import json
from collections.abc import Iterable, Mapping
from datetime import datetime
from typing import Any, Protocol, cast
from uuid import uuid4
from loguru import logger
from nanobot.bus.events import InboundMessage
from nanobot.bus.outbound_events import (
RecoveryStateEvent,
SessionUpdatedEvent,
outbound_message_for_event,
)
from nanobot.bus.queue import MessageBus
from nanobot.session import turn_continuation
from nanobot.session.async_compat import call_session_manager
from nanobot.session.keys import UNIFIED_SESSION_KEY, last_channel_from_metadata
from nanobot.session.manager import Session, SessionManager
from nanobot.webui.metadata import WEBUI_TURN_METADATA_KEY
RUNTIME_CHECKPOINT_KEY = "runtime_checkpoint"
PENDING_USER_TURN_KEY = "pending_user_turn"
RECOVERY_METADATA_KEY = "webui_recovery"
RECOVERY_INBOUND_METADATA_KEY = "_webui_recovery_id"
PENDING_FOLLOWUPS_KEY = "pending_user_followups"
PENDING_FOLLOWUP_ID_KEY = "_recovery_followup_id"
PROVIDER_STATE_CHECKPOINT_VERSION_KEY = "provider_state_checkpoint_version"
PROVIDER_STATE_CHECKPOINT_VERSION = "v1"
_RECOVERY_STATUSES = frozenset({"resuming", "awaiting_user", "recovered", "failed"})
_UNCERTAIN_TOOL_PHASES = frozenset({"awaiting_tools"})
_KNOWN_CHECKPOINT_PHASES = frozenset(
{"final_response", "tools_completed", "awaiting_tools", "error"}
)
class RecoveryActionError(ValueError):
"""A stale or malformed recovery action from an authenticated WebUI."""
def __init__(self, message: str, *, status: int = 400) -> None:
super().__init__(message)
self.status = status
class RecoveryAdmission(Protocol):
"""Narrow AgentLoop boundary for explicit recovery validation."""
async def admit(self, message: InboundMessage) -> bool: ...
def register_recovery_task(self, session_key: str, task: asyncio.Task[Any]) -> None: ...
def unregister_recovery_task(self, session_key: str, task: asyncio.Task[Any]) -> None: ...
def record_pending_followup(session: Session, message: InboundMessage) -> str | None:
"""Durably journal a WebUI follow-up before injecting it into a live turn."""
if message.channel != "websocket":
return None
try:
metadata_value: object = json.loads(json.dumps(message.metadata))
except (TypeError, ValueError):
logger.warning("Skipping non-serializable WebUI follow-up for recovery")
return None
if not isinstance(metadata_value, dict):
return None
metadata = cast(dict[str, Any], metadata_value)
existing_id = metadata.pop(PENDING_FOLLOWUP_ID_KEY, None)
followup_id = (
existing_id
if isinstance(existing_id, str) and existing_id
else uuid4().hex
)
records = _pending_followup_records(session)
if any(record.get("id") == followup_id for record in records):
return followup_id
records.append(
{
"id": followup_id,
"sender_id": message.sender_id,
"chat_id": message.chat_id,
"content": message.content,
"media": list(message.media or []),
"metadata": metadata,
}
)
# This journal is the recovery source of truth, not a mirror of the
# bounded in-memory injection queue. A queued turn can receive more
# follow-ups than the live queue accepts; dropping older journal entries
# would make those acknowledged user messages unrecoverable after a
# gateway restart. Entries are removed only once their user rows are
# committed by ``acknowledge_pending_followups``.
session.metadata[PENDING_FOLLOWUPS_KEY] = records
session.updated_at = datetime.now()
return followup_id
def pending_followups(session: Session) -> list[InboundMessage]:
"""Decode still-unacknowledged follow-ups from durable session metadata."""
messages: list[InboundMessage] = []
for record in _pending_followup_records(session):
followup_id = cast(object, record.get("id"))
sender_id = cast(object, record.get("sender_id"))
chat_id = cast(object, record.get("chat_id"))
content = cast(object, record.get("content"))
metadata = cast(object, record.get("metadata"))
if (
not isinstance(followup_id, str)
or not followup_id
or not isinstance(sender_id, str)
or not sender_id
or not isinstance(chat_id, str)
or not chat_id
):
continue
if not isinstance(content, str) or not isinstance(metadata, dict):
continue
media_value = cast(object, record.get("media"))
media = (
[item for item in cast(list[object], media_value) if isinstance(item, str)]
if isinstance(media_value, list)
else []
)
messages.append(
InboundMessage(
channel="websocket",
sender_id=sender_id,
chat_id=chat_id,
content=content,
media=media,
metadata={**cast(dict[str, Any], metadata), PENDING_FOLLOWUP_ID_KEY: followup_id},
session_key_override=session.key,
require_existing_session=True,
)
)
return messages
def acknowledge_pending_followups(session: Session, followup_ids: Iterable[str]) -> None:
"""Remove journal entries whose user rows were committed to history."""
acknowledged = set(followup_ids)
if not acknowledged:
return
records = [record for record in _pending_followup_records(session) if record.get("id") not in acknowledged]
if records:
session.metadata[PENDING_FOLLOWUPS_KEY] = records
else:
session.metadata.pop(PENDING_FOLLOWUPS_KEY, None)
def _pending_followup_records(session: Session) -> list[dict[str, Any]]:
raw = cast(object, session.metadata.get(PENDING_FOLLOWUPS_KEY))
if not isinstance(raw, list):
return []
values = cast(list[object], raw)
return [cast(dict[str, Any], value) for value in values if isinstance(value, dict)]
def _checkpoint_message_key(message: Mapping[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 _checkpoint_tool_call_ids(
value: object,
*,
result_rows: bool = False,
) -> list[str] | None:
"""Validate checkpoint tool rows and return their stable IDs."""
if not isinstance(value, list):
return None
ids: list[str] = []
for raw in cast(list[object], value):
if not isinstance(raw, dict):
return None
row = cast(dict[str, Any], raw)
id_key = "tool_call_id" if result_rows else "id"
call_id = cast(object, row.get(id_key))
if not isinstance(call_id, str) or not call_id:
return None
if result_rows:
if row.get("role") != "tool":
return None
else:
function_value = cast(object, row.get("function"))
if not isinstance(function_value, dict):
return None
function = cast(dict[str, Any], function_value)
name = cast(object, function.get("name"))
if not isinstance(name, str) or not name:
return None
ids.append(call_id)
return ids if len(ids) == len(set(ids)) else None
def _runtime_checkpoint_is_well_formed(checkpoint: Mapping[str, Any]) -> bool:
"""Return whether a checkpoint is safe to offer for continuation.
Restoration stays tolerant so Dismiss can always clear corrupt state.
Continue is stricter: silently dropping a malformed tool result could make
the model repeat an external side effect.
"""
assistant_value = cast(object, checkpoint.get("assistant_message"))
if not isinstance(assistant_value, dict):
return False
assistant = cast(dict[str, Any], assistant_value)
if assistant.get("role") != "assistant":
return False
completed_ids = _checkpoint_tool_call_ids(
cast(object, checkpoint.get("completed_tool_results")),
result_rows=True,
)
pending_ids = _checkpoint_tool_call_ids(
cast(object, checkpoint.get("pending_tool_calls")),
)
if completed_ids is None or pending_ids is None:
return False
assistant_calls_value = cast(object, assistant.get("tool_calls"))
assistant_call_ids = (
[]
if assistant_calls_value is None
else _checkpoint_tool_call_ids(assistant_calls_value)
)
if assistant_call_ids is None:
return False
phase = checkpoint.get("phase")
if phase == "final_response":
content = cast(object, assistant.get("content"))
return (
isinstance(content, str)
and bool(content.strip())
and not assistant_call_ids
and not completed_ids
and not pending_ids
)
if phase == "awaiting_tools":
return (
bool(assistant_call_ids)
and not completed_ids
and len(assistant_call_ids) == len(pending_ids)
and set(assistant_call_ids) == set(pending_ids)
)
if phase == "tools_completed":
return (
bool(assistant_call_ids)
and not pending_ids
and len(assistant_call_ids) == len(completed_ids)
and set(assistant_call_ids) == set(completed_ids)
)
# Error checkpoints have no current producer contract. Treat legacy or
# future instances as review-only until their exact persisted shape is
# specified; guessing here could make a partial side effect repeat.
return False
def restore_runtime_checkpoint(session: Session) -> bool:
"""Materialize the durable checkpoint exactly once and clear it.
Pending tool calls become explicit interrupted tool results. They are
never executed here. Provider-native state is retained only for the two
checkpoint shapes known to be synchronized with persisted history.
"""
checkpoint = cast(object, session.metadata.get(RUNTIME_CHECKPOINT_KEY))
if not isinstance(checkpoint, dict):
return False
data = cast(dict[str, Any], checkpoint)
assistant = cast(object, data.get("assistant_message"))
completed_value = cast(object, data.get("completed_tool_results"))
pending_value = cast(object, data.get("pending_tool_calls"))
completed = cast(list[object], completed_value) if isinstance(completed_value, list) else []
pending = cast(list[object], pending_value) if isinstance(pending_value, list) else []
restored: list[dict[str, Any]] = []
if isinstance(assistant, dict):
assistant_row = cast(dict[str, Any], assistant)
else:
assistant_row = {}
if assistant_row.get("role") == "assistant":
row = dict(assistant_row)
row.setdefault("timestamp", datetime.now().isoformat())
restored.append(row)
for value in completed:
if not isinstance(value, dict):
continue
tool_result = cast(dict[str, Any], value)
if tool_result.get("role") != "tool":
continue
row = dict(tool_result)
row.setdefault("timestamp", datetime.now().isoformat())
restored.append(row)
for value in pending:
if not isinstance(value, dict):
continue
tool_call = cast(dict[str, Any], value)
tool_call_id = tool_call.get("id")
function_value = cast(object, tool_call.get("function"))
if not isinstance(tool_call_id, str) or not tool_call_id:
continue
function = (
cast(dict[str, Any], function_value)
if isinstance(function_value, dict)
else {}
)
name = function.get("name")
restored.append(
{
"role": "tool",
"tool_call_id": tool_call_id,
"name": name if isinstance(name, str) and name else "tool",
"content": "Error: Task interrupted before this tool finished.",
"timestamp": datetime.now().isoformat(),
"_recovery_interrupted": True,
}
)
overlap = 0
for size in range(min(len(session.messages), len(restored)), 0, -1):
if all(
_checkpoint_message_key(left) == _checkpoint_message_key(right)
for left, right in zip(session.messages[-size:], restored[:size])
):
overlap = size
break
session.messages.extend(restored[overlap:])
assistant_data = cast(dict[str, Any], assistant) if isinstance(assistant, dict) else None
synchronized = (
data.get(PROVIDER_STATE_CHECKPOINT_VERSION_KEY)
== PROVIDER_STATE_CHECKPOINT_VERSION
)
phase = data.get("phase")
exact_final = (
phase == "final_response"
and assistant_data is not None
and assistant_data.get("role") == "assistant"
and not data.get("completed_tool_results")
and not data.get("pending_tool_calls")
)
exact_tools = (
phase == "tools_completed"
and assistant_data is not None
and assistant_data.get("role") == "assistant"
and not data.get("pending_tool_calls")
)
if not (synchronized and (exact_final or exact_tools)):
session.provider_state = None
session.metadata.pop(PENDING_USER_TURN_KEY, None)
session.metadata.pop(RUNTIME_CHECKPOINT_KEY, None)
session.updated_at = datetime.now()
return True
def _discard_runtime_checkpoint(session: Session) -> bool:
"""Drop checkpoint state that cannot be projected into valid history."""
if RUNTIME_CHECKPOINT_KEY not in session.metadata:
return False
session.metadata.pop(RUNTIME_CHECKPOINT_KEY, None)
session.provider_state = None
session.updated_at = datetime.now()
return True
def restore_pending_interruption(session: Session, *, superseded: bool = False) -> bool:
"""Close a persisted user-only turn without pretending it was answered."""
if not session.metadata.get(PENDING_USER_TURN_KEY):
return False
if session.messages and session.messages[-1].get("role") == "user":
content = (
"Task recovery was superseded by a newer message."
if superseded
else "Error: Task interrupted before a response was generated."
)
session.messages.append(
{
"role": "assistant",
"content": content,
"timestamp": datetime.now().isoformat(),
"_recovery_interrupted": True,
}
)
session.provider_state = None
session.updated_at = datetime.now()
session.metadata.pop(PENDING_USER_TURN_KEY, None)
return True
def append_recovery_interruption(session: Session, *, superseded: bool = False) -> None:
"""Close a restored partial turn whose last durable row is not the user message."""
if session.messages and session.messages[-1].get("_recovery_interrupted") is True:
return
session.messages.append(
{
"role": "assistant",
"content": (
"Task recovery was superseded by a newer message."
if superseded
else "Error: Task recovery was interrupted before completion."
),
"timestamp": datetime.now().isoformat(),
"_recovery_interrupted": True,
}
)
session.provider_state = None
session.updated_at = datetime.now()
def recovery_state_from_metadata(metadata: Mapping[str, Any] | None) -> dict[str, Any] | None:
"""Return a sanitized recovery state suitable for the WebSocket wire."""
value = metadata.get(RECOVERY_METADATA_KEY) if metadata else None
if not isinstance(value, dict):
return None
state = cast(dict[str, Any], value)
status = state.get("status")
recovery_id = state.get("recovery_id")
if status not in _RECOVERY_STATUSES or not isinstance(recovery_id, str):
return None
payload: dict[str, Any] = {"status": status, "recovery_id": recovery_id}
reason = state.get("reason")
if isinstance(reason, str) and reason:
payload["reason"] = reason
attempts = state.get("attempts")
if isinstance(attempts, int) and attempts >= 0:
payload["attempts"] = attempts
can_continue = state.get("can_continue")
if isinstance(can_continue, bool):
payload["can_continue"] = can_continue
return payload
@dataclasses.dataclass(slots=True)
class RecoveryCoordinator:
"""Classify, announce, and gate durable WebUI turn recovery."""
sessions: SessionManager
bus: MessageBus
unified_session: bool = False
_active_recovery_tasks: dict[str, asyncio.Task[Any]] = dataclasses.field(
default_factory=dict,
init=False,
repr=False,
)
async def _get_or_create_session(self, key: str) -> Session:
return await call_session_manager(
self.sessions,
"get_or_create_async",
self.sessions.get_or_create,
key,
)
async def _save_session(self, session: Session) -> None:
await call_session_manager(
self.sessions,
"save_async",
self.sessions.save,
session,
)
async def _read_session_metadata(self, key: str) -> dict[str, Any] | None:
return await call_session_manager(
self.sessions,
"read_session_metadata_async",
self.sessions.read_session_metadata,
key,
)
async def _list_sessions(self) -> list[dict[str, Any]]:
return await call_session_manager(
self.sessions,
"list_sessions_async",
self.sessions.list_sessions,
)
def register_recovery_task(self, session_key: str, task: asyncio.Task[Any]) -> None:
"""Track the task that owns an explicit recovery continuation."""
self._active_recovery_tasks[session_key] = task
def unregister_recovery_task(self, session_key: str, task: asyncio.Task[Any]) -> None:
"""Drop a recovery task without removing a newer task for the same session."""
if self._active_recovery_tasks.get(session_key) is task:
self._active_recovery_tasks.pop(session_key, None)
async def _cancel_active_recovery(self, session_key: str) -> None:
"""Stop an explicit continuation before accepting newer user input."""
task = self._active_recovery_tasks.get(session_key)
if task is None or task is asyncio.current_task() or task.done():
return
task.cancel()
# AgentLoop's cancellation path materializes any partial checkpoint and
# releases its pending queue. Wait for that ownership to be released
# before the newer message is routed.
await asyncio.gather(task, return_exceptions=True)
async def scan(self) -> None:
"""Recover every interrupted WebUI session once at gateway startup."""
for key in await self._recovery_candidates():
metadata_payload = await self._read_session_metadata(key)
raw_metadata = metadata_payload.get("metadata") if metadata_payload else None
metadata = cast(dict[str, Any], raw_metadata) if isinstance(raw_metadata, dict) else {}
route = self._websocket_route_for(key, metadata)
if route is None:
continue
unfinished = self._has_unfinished_webui_transcript(key)
if not self._needs_recovery(metadata) and not unfinished:
continue
session = await self._get_or_create_session(key)
try:
await self._recover_session(session, route[1])
await self._requeue_pending_followups(session)
except Exception:
logger.exception("failed to recover interrupted WebUI session {}", session.key)
state = recovery_state_from_metadata(session.metadata)
failed = self._set_state(
session,
status="failed",
recovery_id=cast(str, state["recovery_id"]) if state else uuid4().hex,
attempts=cast(int, state.get("attempts", 0)) if state else 0,
reason="recovery_failed",
can_continue=False,
)
await self._save_session(session)
await self._publish(route[1], failed)
async def _recovery_candidates(self) -> list[str]:
"""Discover canonical and transcript-only WebUI sessions cheaply."""
candidates = dict.fromkeys(
key
for item in await self._list_sessions()
if isinstance((key := item.get("key")), str)
)
try:
# Imported lazily because the sidebar index also projects recovery
# metadata. The index is the owner of transcript-only discovery;
# duplicating its filename and migration rules here would drift.
from nanobot.webui.session_list_index import list_webui_sessions
for item in await asyncio.to_thread(list_webui_sessions, self.sessions):
key = item.get("key")
if isinstance(key, str):
candidates.setdefault(key, None)
except Exception:
# Canonical checkpoint recovery remains available even if the
# optional display-history index is corrupt or unavailable.
logger.exception("failed to discover transcript-only WebUI sessions")
return list(candidates)
@staticmethod
def _needs_recovery(metadata: Mapping[str, Any]) -> bool:
if metadata.get(PENDING_USER_TURN_KEY) is True:
return True
if isinstance(metadata.get(RUNTIME_CHECKPOINT_KEY), dict):
return True
followups = metadata.get(PENDING_FOLLOWUPS_KEY)
if isinstance(followups, list) and len(cast(list[object], followups)) > 0:
return True
state = recovery_state_from_metadata(metadata)
return bool(state and state["status"] in {"resuming", "awaiting_user", "failed"})
async def admit(self, message: InboundMessage) -> bool:
"""Reject stale queued recoveries and let new user input supersede them."""
recovery_id = message.metadata.get(RECOVERY_INBOUND_METADATA_KEY)
if isinstance(recovery_id, str):
session = await self._get_or_create_session(message.session_key)
state = recovery_state_from_metadata(session.metadata)
return bool(
state
and state["status"] == "resuming"
and state["recovery_id"] == recovery_id
)
if message.channel != "websocket":
return True
session = await self._get_or_create_session(message.session_key)
state = recovery_state_from_metadata(session.metadata)
if state and state["status"] in {"resuming", "awaiting_user", "failed"}:
await self._cancel_active_recovery(message.session_key)
restore_runtime_checkpoint(session)
if not restore_pending_interruption(session, superseded=True):
append_recovery_interruption(session, superseded=True)
recovered = self._set_state(
session,
status="recovered",
recovery_id=cast(str, state["recovery_id"]),
attempts=cast(int, state.get("attempts", 0)),
reason="superseded",
)
await self._save_session(session)
await self._publish(message.chat_id, recovered)
return True
async def turn_completed(self, session_key: str) -> None:
"""Resolve a resuming state after the recovered turn commits."""
session = await self._get_or_create_session(session_key)
state = recovery_state_from_metadata(session.metadata)
if not state or state["status"] != "resuming":
return
route = self._websocket_route(session)
if route is None:
return
recovered = self._set_state(
session,
status="recovered",
recovery_id=cast(str, state["recovery_id"]),
attempts=cast(int, state.get("attempts", 0)),
reason="continued",
)
await self._save_session(session)
await self._publish(route[1], recovered)
async def handle_action(self, action: str, payload: dict[str, Any]) -> dict[str, Any]:
"""Apply an authenticated continue/dismiss operation."""
chat_id = payload.get("chat_id")
recovery_id = payload.get("recovery_id")
if not isinstance(chat_id, str) or not chat_id:
raise RecoveryActionError("missing chat_id")
if not isinstance(recovery_id, str) or not recovery_id:
raise RecoveryActionError("missing recovery_id")
session = await self._get_or_create_session(self._session_key(chat_id))
state = recovery_state_from_metadata(session.metadata)
if not state or state["recovery_id"] != recovery_id:
raise RecoveryActionError("recovery state is stale", status=409)
if action == "dismiss":
restore_runtime_checkpoint(session)
restore_pending_interruption(session)
next_state = self._set_state(
session,
status="recovered",
recovery_id=recovery_id,
attempts=cast(int, state.get("attempts", 0)),
reason="dismissed",
)
await self._save_session(session)
await self._publish(chat_id, next_state)
return next_state
if action != "continue":
raise RecoveryActionError("unknown recovery action")
if state["status"] not in {"awaiting_user", "failed"}:
raise RecoveryActionError("recovery is not waiting for confirmation", status=409)
if state.get("can_continue") is False:
raise RecoveryActionError("recovery context is unavailable", status=409)
next_state = self._set_state(
session,
status="resuming",
recovery_id=recovery_id,
attempts=cast(int, state.get("attempts", 0)) + 1,
reason="user_confirmed",
resume_message_count=len(session.messages),
)
await self._save_session(session)
await self._publish(chat_id, next_state)
await self._queue_continuation(session, chat_id, next_state)
return next_state
async def _recover_session(self, session: Session, chat_id: str) -> None:
checkpoint_value = cast(object, session.metadata.get(RUNTIME_CHECKPOINT_KEY))
checkpoint = (
cast(dict[str, Any], checkpoint_value)
if isinstance(checkpoint_value, dict)
else None
)
pending = session.metadata.get(PENDING_USER_TURN_KEY) is True
state = recovery_state_from_metadata(session.metadata)
if not pending and checkpoint is None:
if state and state["status"] == "resuming":
resume_count = self._resume_message_count(session)
if resume_count is not None and len(session.messages) > resume_count:
next_state = self._set_state(
session,
status="recovered",
recovery_id=cast(str, state["recovery_id"]),
attempts=cast(int, state.get("attempts", 0)),
reason="committed",
)
else:
next_state = self._set_state(
session,
status="awaiting_user",
recovery_id=cast(str, state["recovery_id"]),
attempts=cast(int, state.get("attempts", 1)),
reason="loop_guard",
)
await self._save_session(session)
await self._publish(chat_id, next_state)
elif self._has_unfinished_webui_transcript(session.key):
# A normal last-client shutdown can materialize the checkpoint
# before the process exits. In that path there is no pending
# marker left to classify, but the append-only transcript still
# contains an activity row without a turn_end. Treat it as an
# interrupted turn instead of letting the UI resurrect it as a
# forever-running spinner.
can_continue = self._has_saved_continuation_context(session)
waiting = self._set_state(
session,
status="awaiting_user",
recovery_id=uuid4().hex,
attempts=0,
reason=(
"interrupted_with_saved_context"
if can_continue
else "interrupted_without_checkpoint"
),
can_continue=can_continue,
)
await self._save_session(session)
await self._publish(chat_id, waiting)
return
if state and state["status"] in {"awaiting_user", "failed"}:
await self._publish(chat_id, state)
return
if state and state["status"] == "resuming":
restore_runtime_checkpoint(session)
restore_pending_interruption(session)
waiting = self._set_state(
session,
status="awaiting_user",
recovery_id=cast(str, state["recovery_id"]),
attempts=cast(int, state.get("attempts", 1)),
reason="loop_guard",
)
await self._save_session(session)
await self._publish(chat_id, waiting)
return
recovery_id = uuid4().hex
phase = checkpoint.get("phase") if checkpoint is not None else None
pending_calls = checkpoint.get("pending_tool_calls") if checkpoint is not None else None
if checkpoint is not None and phase not in _KNOWN_CHECKPOINT_PHASES:
_discard_runtime_checkpoint(session)
restore_pending_interruption(session)
waiting = self._set_state(
session,
status="awaiting_user",
recovery_id=recovery_id,
attempts=0,
reason="checkpoint_unknown",
can_continue=False,
)
await self._save_session(session)
await self._publish(chat_id, waiting)
return
if checkpoint is not None and not _runtime_checkpoint_is_well_formed(checkpoint):
_discard_runtime_checkpoint(session)
restore_pending_interruption(session)
waiting = self._set_state(
session,
status="awaiting_user",
recovery_id=recovery_id,
attempts=0,
reason="checkpoint_invalid",
can_continue=False,
)
await self._save_session(session)
await self._publish(chat_id, waiting)
return
if phase == "final_response":
restore_runtime_checkpoint(session)
recovered = self._set_state(
session,
status="recovered",
recovery_id=recovery_id,
attempts=0,
reason="answer_restored",
)
await self._save_session(session)
await self._publish(chat_id, recovered)
return
if phase in _UNCERTAIN_TOOL_PHASES or pending_calls:
restore_runtime_checkpoint(session)
waiting = self._set_state(
session,
status="awaiting_user",
recovery_id=recovery_id,
attempts=0,
reason="tool_state_unknown",
)
await self._save_session(session)
await self._publish(chat_id, waiting)
return
# A gateway restart is a lifecycle boundary. Never enqueue model work
# implicitly: even a synchronized checkpoint may sit next to an
# external side effect that the user should review first. The final
# answer path above only restores persisted output; it never executes.
restore_runtime_checkpoint(session)
waiting = self._set_state(
session,
status="awaiting_user",
recovery_id=recovery_id,
attempts=0,
reason="restart_requires_confirmation",
)
await self._save_session(session)
await self._publish(chat_id, waiting)
async def _queue_continuation(
self,
session: Session,
chat_id: str,
state: Mapping[str, Any],
) -> None:
recovery_id = cast(str, state["recovery_id"])
await self.bus.publish_inbound(
InboundMessage(
channel="websocket",
sender_id="system:recovery",
chat_id=chat_id,
content=(
"Continue the interrupted request from the saved conversation context. "
"Do not repeat completed work or mention the restart unless it affects the answer."
),
metadata={
"webui": True,
"_wants_stream": True,
WEBUI_TURN_METADATA_KEY: f"recovery:{recovery_id}",
RECOVERY_INBOUND_METADATA_KEY: recovery_id,
turn_continuation.INTERNAL_CONTINUATION_META: True,
turn_continuation.SKIP_USER_PERSIST_META: True,
},
session_key_override=session.key,
require_existing_session=True,
)
)
async def _requeue_pending_followups(self, session: Session) -> None:
"""Return durable live-turn follow-ups to the bus after a restart."""
for message in pending_followups(session):
await self.bus.publish_inbound(message)
@staticmethod
def _resume_message_count(session: Session) -> int | None:
raw_value = cast(object, session.metadata.get(RECOVERY_METADATA_KEY))
value = cast(dict[str, Any], raw_value) if isinstance(raw_value, dict) else None
if value is None:
return None
count = value.get("resume_message_count")
return count if isinstance(count, int) and count >= 0 else None
async def _publish(
self,
chat_id: str,
state: Mapping[str, Any],
) -> None:
"""Publish the recovery state and invalidate its sidebar projection."""
await self.bus.publish_outbound(
outbound_message_for_event(
channel="websocket",
chat_id=chat_id,
event=RecoveryStateEvent(
status=cast(str, state["status"]),
recovery_id=cast(str, state["recovery_id"]),
reason=cast(str | None, state.get("reason")),
attempts=cast(int, state.get("attempts", 0)),
can_continue=cast(bool | None, state.get("can_continue")),
),
)
)
await self.bus.publish_outbound(
outbound_message_for_event(
channel="websocket",
chat_id=chat_id,
event=SessionUpdatedEvent(scope="thread"),
)
)
@staticmethod
def _set_state(
session: Session,
*,
status: str,
recovery_id: str,
attempts: int,
reason: str,
resume_message_count: int | None = None,
can_continue: bool = True,
) -> dict[str, Any]:
state = {
"status": status,
"recovery_id": recovery_id,
"attempts": max(0, attempts),
"reason": reason,
"updated_at": datetime.now().isoformat(),
}
if not can_continue:
state["can_continue"] = False
if resume_message_count is not None:
state["resume_message_count"] = max(0, resume_message_count)
session.metadata[RECOVERY_METADATA_KEY] = state
session.updated_at = datetime.now()
return state
def _session_key(self, chat_id: str) -> str:
return UNIFIED_SESSION_KEY if self.unified_session else f"websocket:{chat_id}"
@staticmethod
def _has_unfinished_webui_transcript(session_key: str) -> bool:
"""Detect a stale WebUI activity tail after an unclean gateway stop.
The transcript is intentionally consulted only as a last-resort signal:
a durable pending turn or runtime checkpoint always takes precedence.
This keeps browser disconnects harmless while preventing a materialized
partial turn from being presented as active forever after a restart.
"""
try:
from nanobot.webui.transcript import has_unfinished_transcript_tail
return has_unfinished_transcript_tail(session_key)
except (OSError, ValueError, TypeError):
# Recovery must fail closed if the optional display transcript is
# corrupt or unavailable; the normal checkpoint path still applies.
return False
@staticmethod
def _has_saved_continuation_context(session: Session) -> bool:
"""Whether an interrupted turn left model-visible context to continue from."""
last_user = next(
(
index
for index in range(len(session.messages) - 1, -1, -1)
if session.messages[index].get("role") == "user"
),
None,
)
if last_user is None:
return False
tail = session.messages[last_user + 1 :]
return bool(tail) and (
tail[-1].get("role") == "tool"
or any(message.get("_recovery_interrupted") is True for message in tail)
or any(
message.get("role") == "assistant" and bool(message.get("tool_calls"))
for message in tail
)
)
@staticmethod
def _websocket_route(session: Session) -> tuple[str, str] | None:
return RecoveryCoordinator._websocket_route_for(session.key, session.metadata)
@staticmethod
def _websocket_route_for(
session_key: str,
metadata: Mapping[str, Any],
) -> tuple[str, str] | None:
if session_key.startswith("websocket:"):
chat_id = session_key.split(":", 1)[1]
return ("websocket", chat_id) if chat_id else None
if session_key == UNIFIED_SESSION_KEY:
route = last_channel_from_metadata(metadata)
if route and route[0] == "websocket":
return route
return None
+21 -12
View File
@@ -2,6 +2,7 @@
from __future__ import annotations from __future__ import annotations
import asyncio
import re import re
import time import time
from collections.abc import Awaitable, Callable from collections.abc import Awaitable, Callable
@@ -37,12 +38,14 @@ from nanobot.bus.runtime_events import (
TurnRuntimeAdmitted, TurnRuntimeAdmitted,
UserInputAccepted, UserInputAccepted,
) )
from nanobot.providers.base import LLMProvider from nanobot.llm_usage.context import llm_usage_source
from nanobot.providers.base import LLMProvider, LLMUsage
from nanobot.providers.fallback_provider import FallbackModelObserver from nanobot.providers.fallback_provider import FallbackModelObserver
from nanobot.runtime_context import public_history_message from nanobot.runtime_context import public_history_message
from nanobot.session.goal_state import goal_state_ws_blob from nanobot.session.goal_state import goal_state_ws_blob
from nanobot.session.history_visibility import is_hidden_history_message from nanobot.session.history_visibility import is_hidden_history_message
from nanobot.session.manager import Session, SessionManager from nanobot.session.manager import Session, SessionManager
from nanobot.session.recovery import RecoveryCoordinator
from nanobot.session.session_handles import session_handle_for_name from nanobot.session.session_handles import session_handle_for_name
from nanobot.session.session_messages import ( from nanobot.session.session_messages import (
SessionMessageEnvelope, SessionMessageEnvelope,
@@ -174,7 +177,7 @@ async def maybe_generate_webui_title(
model: str, model: str,
) -> bool: ) -> bool:
"""Generate and persist a short title for WebUI-owned sessions only.""" """Generate and persist a short title for WebUI-owned sessions only."""
session = sessions.get_or_create(session_key) session = await sessions.get_or_create_async(session_key)
if session.metadata.get(WEBUI_SESSION_METADATA_KEY) is not True: if session.metadata.get(WEBUI_SESSION_METADATA_KEY) is not True:
return False return False
if session.metadata.get(WEBUI_TITLE_USER_EDITED_METADATA_KEY) is True: if session.metadata.get(WEBUI_TITLE_USER_EDITED_METADATA_KEY) is True:
@@ -185,7 +188,7 @@ async def maybe_generate_webui_title(
if cleaned_current_title: if cleaned_current_title:
if cleaned_current_title != current_title: if cleaned_current_title != current_title:
session.metadata[WEBUI_TITLE_METADATA_KEY] = cleaned_current_title session.metadata[WEBUI_TITLE_METADATA_KEY] = cleaned_current_title
sessions.save(session) await sessions.save_async(session)
return False return False
session.metadata.pop(WEBUI_TITLE_METADATA_KEY, None) session.metadata.pop(WEBUI_TITLE_METADATA_KEY, None)
@@ -207,6 +210,7 @@ async def maybe_generate_webui_title(
prompt += f"\nAssistant: {truncate_text(assistant_text, 1_000)}" prompt += f"\nAssistant: {truncate_text(assistant_text, 1_000)}"
try: try:
with llm_usage_source("system"):
response = await provider.chat_with_retry( response = await provider.chat_with_retry(
[ [
{ {
@@ -238,7 +242,7 @@ async def maybe_generate_webui_title(
) )
return False return False
session.metadata[WEBUI_TITLE_METADATA_KEY] = title session.metadata[WEBUI_TITLE_METADATA_KEY] = title
sessions.save(session) await sessions.save_async(session)
return True return True
@@ -434,8 +438,8 @@ class WebuiTurnRoutePolicy:
) )
and route.channel == "websocket" and route.channel == "websocket"
): ):
session = self.sessions.get_or_create(session_key) session = self.sessions.get_cached(session_key)
if session.metadata.get(WEBUI_SESSION_METADATA_KEY) is True: if session is not None and session.metadata.get(WEBUI_SESSION_METADATA_KEY) is True:
metadata = dict(route.metadata) metadata = dict(route.metadata)
turn_prefix = "session-input" if internal_user_input else "subagent" turn_prefix = "session-input" if internal_user_input else "subagent"
metadata.update({ metadata.update({
@@ -511,6 +515,7 @@ class WebuiTurnCoordinator:
bus: MessageBus bus: MessageBus
sessions: SessionManager sessions: SessionManager
schedule_background: Callable[[Awaitable[None]], None] schedule_background: Callable[[Awaitable[None]], None]
recovery: RecoveryCoordinator | None = None
def subscribe(self, runtime_events: RuntimeEventBus) -> Callable[[], None]: def subscribe(self, runtime_events: RuntimeEventBus) -> Callable[[], None]:
"""Subscribe this coordinator to runtime events.""" """Subscribe this coordinator to runtime events."""
@@ -576,7 +581,7 @@ class WebuiTurnCoordinator:
or not session_key.startswith("websocket:") or not session_key.startswith("websocket:")
): ):
return return
persisted = self.sessions.read_session_metadata(session_key) persisted = await self.sessions.read_session_metadata_async(session_key)
metadata_value: object = persisted.get("metadata") if persisted is not None else None metadata_value: object = persisted.get("metadata") if persisted is not None else None
metadata = ( metadata = (
cast(dict[str, Any], metadata_value) cast(dict[str, Any], metadata_value)
@@ -587,7 +592,8 @@ class WebuiTurnCoordinator:
return return
public_metadata = _session_message_public_metadata(envelope) public_metadata = _session_message_public_metadata(envelope)
try: try:
append_session_message_input( await asyncio.to_thread(
append_session_message_input,
session_key, session_key,
content=event.content, content=event.content,
created_at_ms=envelope["created_at_ms"], created_at_ms=envelope["created_at_ms"],
@@ -612,7 +618,8 @@ class WebuiTurnCoordinator:
def _handle_session_turn_started(self, event: SessionTurnStarted) -> None: def _handle_session_turn_started(self, event: SessionTurnStarted) -> None:
if not self._is_websocket_event(event.context): if not self._is_websocket_event(event.context):
return return
session = self.sessions.get_or_create(event.context.session_key) session = self.sessions.get_cached(event.context.session_key)
if session is not None:
mark_webui_session(session, event.context.metadata) mark_webui_session(session, event.context.metadata)
async def _handle_run_status_changed(self, event: TurnRunStatusChanged) -> None: async def _handle_run_status_changed(self, event: TurnRunStatusChanged) -> None:
@@ -654,6 +661,8 @@ class WebuiTurnCoordinator:
event.runtime.context_window_tokens if event.runtime is not None else None event.runtime.context_window_tokens if event.runtime is not None else None
), ),
) )
if self.recovery is not None:
await self.recovery.turn_completed(event.context.session_key)
self._schedule_title_update_from_event(event) self._schedule_title_update_from_event(event)
async def _handle_goal_state_changed(self, event: GoalStateChanged) -> None: async def _handle_goal_state_changed(self, event: GoalStateChanged) -> None:
@@ -691,13 +700,13 @@ class WebuiTurnCoordinator:
*, *,
session_key: str, session_key: str,
latency_ms: int | None, latency_ms: int | None,
usage: dict[str, int] | None = None, usage: LLMUsage | None = None,
context_window_tokens: int | None = None, context_window_tokens: int | None = None,
) -> None: ) -> None:
if msg.channel != "websocket": if msg.channel != "websocket":
return return
session = self.sessions.get_or_create(session_key) session = await self.sessions.get_or_create_async(session_key)
await self.bus.publish_outbound( await self.bus.publish_outbound(
outbound_message_for_event( outbound_message_for_event(
channel=msg.channel, channel=msg.channel,
@@ -705,7 +714,7 @@ class WebuiTurnCoordinator:
event=TurnEndEvent( event=TurnEndEvent(
latency_ms=latency_ms, latency_ms=latency_ms,
goal_state=goal_state_ws_blob(session.metadata), goal_state=goal_state_ws_blob(session.metadata),
usage=usage or None, usage=usage,
context_window_tokens=context_window_tokens, context_window_tokens=context_window_tokens,
), ),
metadata=msg.metadata, metadata=msg.metadata,
+2 -2
View File
@@ -16,7 +16,7 @@ Concrete scenarios showing when and how to use the my tool effectively.
→ my(action="check", key="max_iterations") → my(action="check", key="max_iterations")
→ 40 → 40
→ my(action="check", key="_last_usage") → my(action="check", key="_last_usage")
→ {"prompt_tokens": 62000, "completion_tokens": 3000} → {"input_tokens": 62000, "output_tokens": 3000}
→ "I hit the iteration limit (40). The task was complex. I can ask the user if they want to increase it." → "I hit the iteration limit (40). The task was complex. I can ask the user if they want to increase it."
``` ```
@@ -72,6 +72,6 @@ Concrete scenarios showing when and how to use the my tool effectively.
### Token-conscious behavior ### Token-conscious behavior
``` ```
→ my(action="check", key="_last_usage") → my(action="check", key="_last_usage")
→ {"prompt_tokens": 58000, "completion_tokens": 12000} → {"input_tokens": 58000, "output_tokens": 12000}
→ "I've consumed ~70k tokens. I'll keep my remaining responses focused." → "I've consumed ~70k tokens. I'll keep my remaining responses focused."
``` ```
+206 -50
View File
@@ -5,17 +5,24 @@ from __future__ import annotations
import asyncio import asyncio
import uuid import uuid
from collections.abc import Awaitable, Callable from collections.abc import Awaitable, Callable
from typing import Any from contextlib import suppress
from typing import Any, TypeVar
from loguru import logger from loguru import logger
from nanobot.agent.automation_turns import AutomationTurnError from nanobot.agent.automation_turns import (
AutomationTurnAcceptedCancellation,
AutomationTurnError,
)
from nanobot.bus.events import InboundMessage, OutboundMessage from nanobot.bus.events import InboundMessage, OutboundMessage
from nanobot.triggers.local_session_turns import LOCAL_TRIGGER_META from nanobot.triggers.local_session_turns import LOCAL_TRIGGER_META
from nanobot.triggers.local_store import LocalTriggerStore from nanobot.triggers.local_store import LocalTriggerStore
from nanobot.triggers.local_types import LocalTrigger, TriggerDelivery from nanobot.triggers.local_types import LocalTrigger, TriggerDelivery
from nanobot.utils.cancellation import shield_and_drain
from nanobot.webui.metadata import WEBUI_MESSAGE_SOURCE_METADATA_KEY, WEBUI_TURN_METADATA_KEY from nanobot.webui.metadata import WEBUI_MESSAGE_SOURCE_METADATA_KEY, WEBUI_TURN_METADATA_KEY
_T = TypeVar("_T")
async def run_local_trigger_queue( async def run_local_trigger_queue(
*, *,
@@ -29,14 +36,16 @@ async def run_local_trigger_queue(
if submit_turn is None: if submit_turn is None:
raise ValueError("run_local_trigger_queue requires submit_turn") raise ValueError("run_local_trigger_queue requires submit_turn")
logger.info("Local trigger queue started") logger.info("Local trigger queue started")
recovered = store.recover_processing_deliveries() recovered = await shield_and_drain(asyncio.to_thread(store.recover_processing_deliveries))
if recovered: if recovered:
logger.warning( logger.warning(
"Trigger: recovered {} interrupted delivery file(s) from processing", "Trigger: recovered {} interrupted delivery file(s) from processing",
recovered, recovered,
) )
while True: while True:
deliveries = store.claim_deliveries(limit=batch_size) deliveries = await shield_and_drain(
asyncio.to_thread(store.claim_deliveries, limit=batch_size)
)
if not deliveries: if not deliveries:
await asyncio.sleep(poll_interval_s) await asyncio.sleep(poll_interval_s)
continue continue
@@ -49,30 +58,24 @@ async def run_local_trigger_queue(
submit_turn=submit_turn, submit_turn=submit_turn,
is_channel_enabled=is_channel_enabled, is_channel_enabled=is_channel_enabled,
) )
store.complete_delivery(delivery) except _DeliverySettledOnCancellation:
raise
except asyncio.CancelledError as exc: except asyncio.CancelledError as exc:
store.retry_delivery(delivery, str(exc) or exc.__class__.__name__) error = str(exc) or exc.__class__.__name__
_write_delivery_run_record( await shield_and_drain(asyncio.to_thread(store.retry_delivery, delivery, error))
await _write_delivery_run_record(
store, store,
delivery, delivery,
status="interrupted", status="interrupted",
error=str(exc) or exc.__class__.__name__, error=error,
) )
raise raise
except _TerminalDeliveryError as exc: except _TerminalDeliveryError as exc:
store.record_delivery( await _await_delivery_settlement(
delivery.trigger_id, _settle_failed_delivery(store, delivery, error=str(exc)),
status="error", store=store,
error=str(exc), delivery=delivery,
run_at_ms=delivery.created_at_ms,
) )
_write_delivery_run_record(
store,
delivery,
status="error",
error=str(exc),
)
store.complete_delivery(delivery)
logger.warning( logger.warning(
"Trigger: dropped delivery {} for {}: {}", "Trigger: dropped delivery {} for {}: {}",
delivery.id, delivery.id,
@@ -81,19 +84,11 @@ async def run_local_trigger_queue(
) )
except AutomationTurnError as exc: except AutomationTurnError as exc:
error = str(exc) or exc.__class__.__name__ error = str(exc) or exc.__class__.__name__
store.record_delivery( await _await_delivery_settlement(
delivery.trigger_id, _settle_failed_delivery(store, delivery, error=error),
status="error", store=store,
error=error, delivery=delivery,
run_at_ms=delivery.created_at_ms,
) )
_write_delivery_run_record(
store,
delivery,
status="error",
error=error,
)
store.complete_delivery(delivery)
logger.warning( logger.warning(
"Trigger: delivery {} for {} reached the agent but failed: {}", "Trigger: delivery {} for {} reached the agent but failed: {}",
delivery.id, delivery.id,
@@ -102,18 +97,10 @@ async def run_local_trigger_queue(
) )
except Exception as exc: except Exception as exc:
error = str(exc) or exc.__class__.__name__ error = str(exc) or exc.__class__.__name__
retried = store.retry_delivery(delivery, error) retried = await _await_delivery_settlement(
_write_delivery_run_record( _settle_retryable_delivery(store, delivery, error=error),
store, store=store,
delivery, delivery=delivery,
status="retrying" if retried else "error",
error=error,
)
store.record_delivery(
delivery.trigger_id,
status="error",
error=error,
run_at_ms=delivery.created_at_ms,
) )
logger.exception( logger.exception(
"Trigger: failed delivery {} for {}{}", "Trigger: failed delivery {} for {}{}",
@@ -127,6 +114,10 @@ class _TerminalDeliveryError(RuntimeError):
pass pass
class _DeliverySettledOnCancellation(asyncio.CancelledError):
"""Cancellation reported only after an already-submitted delivery is settled."""
async def _deliver_delivery( async def _deliver_delivery(
store: LocalTriggerStore, store: LocalTriggerStore,
delivery: TriggerDelivery, delivery: TriggerDelivery,
@@ -134,7 +125,7 @@ async def _deliver_delivery(
submit_turn: Callable[[InboundMessage], Awaitable[OutboundMessage | None]], submit_turn: Callable[[InboundMessage], Awaitable[OutboundMessage | None]],
is_channel_enabled: Callable[[str], bool], is_channel_enabled: Callable[[str], bool],
) -> None: ) -> None:
trigger = store.get(delivery.trigger_id) trigger = await asyncio.to_thread(store.get, delivery.trigger_id)
if trigger is None: if trigger is None:
raise _TerminalDeliveryError("trigger not found") raise _TerminalDeliveryError("trigger not found")
if not trigger.enabled: if not trigger.enabled:
@@ -142,7 +133,14 @@ async def _deliver_delivery(
if not is_channel_enabled(trigger.channel): if not is_channel_enabled(trigger.channel):
raise _TerminalDeliveryError(f"target channel is not enabled: {trigger.channel}") raise _TerminalDeliveryError(f"target channel is not enabled: {trigger.channel}")
store.write_delivery_run_record(delivery, trigger=trigger, status="processing") await shield_and_drain(
asyncio.to_thread(
store.write_delivery_run_record,
delivery,
trigger=trigger,
status="processing",
)
)
msg = InboundMessage( msg = InboundMessage(
channel=trigger.channel, channel=trigger.channel,
sender_id=trigger.sender_id, sender_id=trigger.sender_id,
@@ -151,22 +149,177 @@ async def _deliver_delivery(
metadata=_delivery_metadata(trigger, delivery), metadata=_delivery_metadata(trigger, delivery),
session_key_override=trigger.session_key, session_key_override=trigger.session_key,
) )
try:
response = await submit_turn(msg) response = await submit_turn(msg)
store.record_delivery( except AutomationTurnAcceptedCancellation:
try:
await _await_delivery_settlement(
_settle_accepted_delivery(store, delivery, trigger=trigger),
store=store,
delivery=delivery,
)
except _DeliverySettledOnCancellation:
raise
except Exception:
logger.exception(
"Trigger: failed to persist accepted delivery {}; dropping retry",
delivery.id,
)
with suppress(Exception):
await shield_and_drain(asyncio.to_thread(store.complete_delivery, delivery))
raise _DeliverySettledOnCancellation from None
try:
await _await_delivery_settlement(
_settle_submitted_delivery(store, delivery, trigger=trigger, response=response),
store=store,
delivery=delivery,
)
except Exception:
logger.exception(
"Trigger: failed to persist status for submitted delivery {}; dropping retry",
delivery.id,
)
with suppress(Exception):
await shield_and_drain(asyncio.to_thread(store.complete_delivery, delivery))
async def _await_delivery_settlement(
operation: Awaitable[_T],
*,
store: LocalTriggerStore,
delivery: TriggerDelivery,
) -> _T:
settlement = asyncio.ensure_future(operation)
try:
return await asyncio.shield(settlement)
except asyncio.CancelledError:
while not settlement.done():
try:
await asyncio.shield(settlement)
except asyncio.CancelledError:
continue
try:
settlement.result()
except Exception:
logger.exception(
"Trigger: failed to settle delivery {} during cancellation",
delivery.id,
)
completion = asyncio.create_task(
shield_and_drain(asyncio.to_thread(store.complete_delivery, delivery))
)
while not completion.done():
try:
await asyncio.shield(completion)
except asyncio.CancelledError:
continue
with suppress(Exception):
completion.result()
raise _DeliverySettledOnCancellation from None
async def _settle_failed_delivery(
store: LocalTriggerStore,
delivery: TriggerDelivery,
*,
error: str,
) -> None:
await _write_delivery_run_record(
store,
delivery,
status="error",
error=error,
)
await shield_and_drain(asyncio.to_thread(store.complete_delivery, delivery))
# Publish the terminal status only after the durable delivery state is settled.
await shield_and_drain(
asyncio.to_thread(
store.record_delivery,
delivery.trigger_id,
status="error",
error=error,
run_at_ms=delivery.created_at_ms,
)
)
async def _settle_retryable_delivery(
store: LocalTriggerStore,
delivery: TriggerDelivery,
*,
error: str,
) -> bool:
retried = await shield_and_drain(asyncio.to_thread(store.retry_delivery, delivery, error))
await _write_delivery_run_record(
store,
delivery,
status="retrying" if retried else "error",
error=error,
)
await shield_and_drain(
asyncio.to_thread(
store.record_delivery,
delivery.trigger_id,
status="error",
error=error,
run_at_ms=delivery.created_at_ms,
)
)
return retried
async def _settle_accepted_delivery(
store: LocalTriggerStore,
delivery: TriggerDelivery,
*,
trigger: LocalTrigger,
) -> None:
"""Commit an accepted delivery without claiming the agent turn completed."""
await _write_delivery_run_record(
store,
delivery,
trigger=trigger,
status="accepted",
)
await shield_and_drain(asyncio.to_thread(store.complete_delivery, delivery))
await shield_and_drain(
asyncio.to_thread(
store.record_delivery,
trigger.id, trigger.id,
status="ok", status="ok",
run_at_ms=delivery.created_at_ms, run_at_ms=delivery.created_at_ms,
) )
_write_delivery_run_record( )
async def _settle_submitted_delivery(
store: LocalTriggerStore,
delivery: TriggerDelivery,
*,
trigger: LocalTrigger,
response: OutboundMessage | None,
) -> None:
await _write_delivery_run_record(
store, store,
delivery, delivery,
trigger=trigger, trigger=trigger,
status="ok", status="ok",
response=response.content if response else "", response=response.content if response else "",
) )
await shield_and_drain(asyncio.to_thread(store.complete_delivery, delivery))
# last_status is the externally visible commit marker for a settled delivery.
await shield_and_drain(
asyncio.to_thread(
store.record_delivery,
trigger.id,
status="ok",
run_at_ms=delivery.created_at_ms,
)
)
def _write_delivery_run_record( async def _write_delivery_run_record(
store: LocalTriggerStore, store: LocalTriggerStore,
delivery: TriggerDelivery, delivery: TriggerDelivery,
*, *,
@@ -176,13 +329,16 @@ def _write_delivery_run_record(
response: str | None = None, response: str | None = None,
) -> None: ) -> None:
try: try:
store.write_delivery_run_record( await shield_and_drain(
asyncio.to_thread(
store.write_delivery_run_record,
delivery, delivery,
trigger=trigger, trigger=trigger,
status=status, status=status,
error=error, error=error,
response=response, response=response,
) )
)
except Exception: except Exception:
logger.exception( logger.exception(
"Trigger: failed to write run record for delivery {}", "Trigger: failed to write run record for delivery {}",
+5 -1
View File
@@ -24,6 +24,7 @@ _MAX_RUN_HISTORY = 20
_MAX_DELIVERY_ATTEMPTS = 10 _MAX_DELIVERY_ATTEMPTS = 10
_RUN_RECORD_TEXT_MAX_CHARS = 4000 _RUN_RECORD_TEXT_MAX_CHARS = 4000
_PROCESSING_RECOVERY_ERROR = "delivery was recovered from interrupted processing" _PROCESSING_RECOVERY_ERROR = "delivery was recovered from interrupted processing"
_FILE_LOCK_TIMEOUT_SECONDS = 5
class TriggerStoreError(RuntimeError): class TriggerStoreError(RuntimeError):
@@ -49,7 +50,10 @@ class LocalTriggerStore:
self.processing_dir = self.root / "processing" self.processing_dir = self.root / "processing"
self.failed_dir = self.root / "failed" self.failed_dir = self.root / "failed"
self.runs_dir = self.root / "runs" self.runs_dir = self.root / "runs"
self._lock = FileLock(str(self.root / ".lock")) self._lock = FileLock(
str(self.root / ".lock"),
timeout=_FILE_LOCK_TIMEOUT_SECONDS,
)
def create( def create(
self, self,
+40
View File
@@ -3,8 +3,48 @@
from __future__ import annotations from __future__ import annotations
import asyncio import asyncio
from collections.abc import Awaitable
from typing import TypeVar
_T = TypeVar("_T")
def task_is_cancelling() -> bool: def task_is_cancelling() -> bool:
task = asyncio.current_task() task = asyncio.current_task()
return task is not None and task.cancelling() > 0 return task is not None and task.cancelling() > 0
async def shield_and_drain(awaitable: Awaitable[_T]) -> _T:
"""Delay caller cancellation until an accepted operation has fully settled.
``asyncio.to_thread`` cannot stop a worker that has already started. Shielding
keeps cancellation from detaching that worker, and draining also lets any
post-write in-memory settlement in ``awaitable`` finish. Cancellation is still
re-raised as soon as the accepted operation is done.
"""
settlement = asyncio.ensure_future(awaitable)
cancellation: asyncio.CancelledError | None = None
while not settlement.done():
try:
result = await asyncio.shield(settlement)
except asyncio.CancelledError as exc:
if cancellation is None:
cancellation = exc
except BaseException:
if cancellation is None:
raise
break
else:
if cancellation is not None:
raise cancellation
return result
if cancellation is not None:
try:
settlement.result()
except BaseException:
# The caller's cancellation wins once settlement has been observed.
pass
raise cancellation
return settlement.result()
+10 -5
View File
@@ -1,5 +1,7 @@
"""Utility functions for nanobot.""" """Utility functions for nanobot."""
from __future__ import annotations
import base64 import base64
import json import json
import os import os
@@ -12,11 +14,14 @@ from contextlib import suppress
from datetime import datetime from datetime import datetime
from functools import lru_cache from functools import lru_cache
from pathlib import Path from pathlib import Path
from typing import Any, TypeVar, cast, overload from typing import TYPE_CHECKING, Any, TypeVar, cast, overload
import tiktoken import tiktoken
from loguru import logger from loguru import logger
if TYPE_CHECKING:
from nanobot.providers.base import LLMUsage
_TOOLS_TOKEN_CACHE_MAX_ENTRIES = 64 _TOOLS_TOKEN_CACHE_MAX_ENTRIES = 64
_TOOLS_TOKEN_CACHE: dict[int, tuple[tuple[int, ...], dict[bool, int]]] = {} _TOOLS_TOKEN_CACHE: dict[int, tuple[tuple[int, ...], dict[bool, int]]] = {}
_T = TypeVar("_T") _T = TypeVar("_T")
@@ -793,7 +798,7 @@ def build_status_content(
version: str, version: str,
model: str, model: str,
start_time: float, start_time: float,
last_usage: dict[str, int], last_usage: LLMUsage | None,
context_window_tokens: int, context_window_tokens: int,
session_msg_count: int, session_msg_count: int,
context_tokens_estimate: int, context_tokens_estimate: int,
@@ -814,9 +819,9 @@ def build_status_content(
if uptime_s >= 3600 if uptime_s >= 3600
else f"{uptime_s // 60}m {uptime_s % 60}s" else f"{uptime_s // 60}m {uptime_s % 60}s"
) )
last_in = last_usage.get("prompt_tokens", 0) last_in = last_usage.input_tokens if last_usage else 0
last_out = last_usage.get("completion_tokens", 0) last_out = last_usage.output_tokens if last_usage else 0
cached = last_usage.get("cached_tokens", 0) cached = last_usage.cache_read_tokens if last_usage else None
ctx_total = max(context_window_tokens, 0) ctx_total = max(context_window_tokens, 0)
# Budget mirrors Consolidator formula: ctx_window - max_completion - _SAFETY_BUFFER # Budget mirrors Consolidator formula: ctx_window - max_completion - _SAFETY_BUFFER
ctx_budget = max(ctx_total - int(max_completion_tokens) - 1024, 1) ctx_budget = max(ctx_total - int(max_completion_tokens) - 1024, 1)
+16 -3
View File
@@ -2,6 +2,7 @@
from __future__ import annotations from __future__ import annotations
import asyncio
import re import re
import uuid import uuid
from collections.abc import Mapping from collections.abc import Mapping
@@ -9,6 +10,7 @@ from typing import TYPE_CHECKING, Any, TypeGuard
from nanobot.session.manager import SessionManager from nanobot.session.manager import SessionManager
from nanobot.session.webui_turns import WEBUI_TITLE_METADATA_KEY, clean_generated_title from nanobot.session.webui_turns import WEBUI_TITLE_METADATA_KEY, clean_generated_title
from nanobot.utils.cancellation import shield_and_drain
from nanobot.webui.transcript import ( from nanobot.webui.transcript import (
append_fork_marker, append_fork_marker,
delete_webui_transcript, delete_webui_transcript,
@@ -93,15 +95,24 @@ async def handle_webui_fork_chat(
await channel.send_webui_protocol_error(connection, "session_manager_unavailable") await channel.send_webui_protocol_error(connection, "session_manager_unavailable")
return return
async def create_and_attach() -> None:
try: try:
forked = create_webui_chat_fork( forked = await asyncio.to_thread(
create_webui_chat_fork,
session_manager, session_manager,
source_chat_id=source_chat_id, source_chat_id=source_chat_id,
before_user_index=raw_index, before_user_index=raw_index,
title=envelope.get("title") if isinstance(envelope.get("title"), str) else None, title=(
envelope.get("title")
if isinstance(envelope.get("title"), str)
else None
),
) )
if forked is None: if forked is None:
await channel.send_webui_protocol_error(connection, "invalid fork source or index") await channel.send_webui_protocol_error(
connection,
"invalid fork source or index",
)
return return
fork_id, fork_key = forked fork_id, fork_key = forked
except Exception as exc: except Exception as exc:
@@ -114,3 +125,5 @@ async def handle_webui_fork_chat(
fork_id=fork_id, fork_id=fork_id,
fork_key=fork_key, fork_key=fork_key,
) )
await shield_and_drain(create_and_attach())
+2
View File
@@ -69,6 +69,7 @@ def build_gateway_services(
mcp_runtime_status: Callable[[], Mapping[str, str]] | None = None, mcp_runtime_status: Callable[[], Mapping[str, str]] | None = None,
mcp_reload: Callable[[], Awaitable[dict[str, Any]]] | None = None, mcp_reload: Callable[[], Awaitable[dict[str, Any]]] | None = None,
skill_state_action: Callable[[set[str]], None] | None = None, skill_state_action: Callable[[set[str]], None] | None = None,
recovery_action: Callable[[str, dict[str, Any]], Awaitable[dict[str, Any]]] | None = None,
logger: Any = default_logger, logger: Any = default_logger,
) -> GatewayServices: ) -> GatewayServices:
settings = WebUISettingsServices.create( settings = WebUISettingsServices.create(
@@ -131,6 +132,7 @@ def build_gateway_services(
mcp_runtime_status=mcp_runtime_status, mcp_runtime_status=mcp_runtime_status,
mcp_reload=mcp_reload, mcp_reload=mcp_reload,
skill_state_action=skill_state_action, skill_state_action=skill_state_action,
recovery_action=recovery_action,
log=logger, log=logger,
) )
return GatewayServices( return GatewayServices(
+2 -2
View File
@@ -16,7 +16,7 @@ from nanobot.agent.tools.mcp import MCPConnection, connect_mcp_servers
from nanobot.agent.tools.mcp_oauth import MCP_OAUTH_CALLBACK_PATH, MCPOAuthHandlers from nanobot.agent.tools.mcp_oauth import MCP_OAUTH_CALLBACK_PATH, MCPOAuthHandlers
from nanobot.agent.tools.registry import ToolRegistry from nanobot.agent.tools.registry import ToolRegistry
from nanobot.config.schema import MCPServerConfig from nanobot.config.schema import MCPServerConfig
from nanobot.security.network import validate_url_target from nanobot.security.network import async_validate_url_target
from nanobot.webui.http_utils import is_loopback_host from nanobot.webui.http_utils import is_loopback_host
McpReload = Callable[[], Awaitable[dict[str, Any]]] McpReload = Callable[[], Awaitable[dict[str, Any]]]
@@ -259,7 +259,7 @@ class McpOAuthManager:
): ):
flow.error = "The MCP server returned an unsafe authorization URL." flow.error = "The MCP server returned an unsafe authorization URL."
raise McpOAuthError(flow.error) raise McpOAuthError(flow.error)
ok, _error = validate_url_target(authorization_url) ok, _error = await async_validate_url_target(authorization_url)
if not ok: if not ok:
flow.error = "The MCP server returned an unsafe authorization URL." flow.error = "The MCP server returned an unsafe authorization URL."
raise McpOAuthError(flow.error) raise McpOAuthError(flow.error)
+3 -12
View File
@@ -4,6 +4,7 @@ from __future__ import annotations
from typing import Any, cast from typing import Any, cast
from nanobot.providers.base import LLMUsage
from nanobot.session.manager import Session from nanobot.session.manager import Session
from nanobot.utils.helpers import estimate_message_tokens, truncate_text from nanobot.utils.helpers import estimate_message_tokens, truncate_text
@@ -36,18 +37,8 @@ def session_context_payload(session: Session) -> dict[str, Any]:
summary_tokens = ( summary_tokens = (
estimate_message_tokens({"role": "system", "content": summary}) if summary else 0 estimate_message_tokens({"role": "system", "content": summary}) if summary else 0
) )
raw_usage = session.metadata.get("_last_usage") stored_usage = LLMUsage.from_dict(session.metadata.get("_last_usage"))
last_usage = ( last_usage = stored_usage.to_turn_dict() if stored_usage is not None else None
{
key: value
for key, value in cast(dict[object, object], raw_usage).items()
if isinstance(key, str)
and type(value) is int
and value >= 0
}
if isinstance(raw_usage, dict)
else None
)
return { return {
"schema_version": 1, "schema_version": 1,
+6 -1
View File
@@ -31,8 +31,9 @@ from nanobot.session.manager import (
_metadata_title, # pyright: ignore[reportPrivateUsage] _metadata_title, # pyright: ignore[reportPrivateUsage]
) )
from nanobot.session.model_selection import model_preset_from_metadata from nanobot.session.model_selection import model_preset_from_metadata
from nanobot.session.recovery import recovery_state_from_metadata
_INDEX_VERSION = 7 _INDEX_VERSION = 8
_INDEX_FILENAME = ".webui_session_index.json" _INDEX_FILENAME = ".webui_session_index.json"
_MODEL_PRESET_FIELD = "model_preset" _MODEL_PRESET_FIELD = "model_preset"
_ROW_SOURCE_FIELD = "_source" _ROW_SOURCE_FIELD = "_source"
@@ -245,6 +246,7 @@ def _public_row(sessions_dir: Path, webui_dir: Path, row: dict[str, Any]) -> dic
"title": row.get("title", ""), "title": row.get("title", ""),
"preview": row.get("preview", ""), "preview": row.get("preview", ""),
_MODEL_PRESET_FIELD: row.get(_MODEL_PRESET_FIELD), _MODEL_PRESET_FIELD: row.get(_MODEL_PRESET_FIELD),
"recovery_state": row.get("recovery_state"),
_WORKSPACE_SCOPE_PRESENT_FIELD: row.get(_WORKSPACE_SCOPE_PRESENT_FIELD, False), _WORKSPACE_SCOPE_PRESENT_FIELD: row.get(_WORKSPACE_SCOPE_PRESENT_FIELD, False),
_WORKSPACE_SCOPE_VALUE_FIELD: row.get(_WORKSPACE_SCOPE_VALUE_FIELD), _WORKSPACE_SCOPE_VALUE_FIELD: row.get(_WORKSPACE_SCOPE_VALUE_FIELD),
"path": str(path), "path": str(path),
@@ -485,6 +487,7 @@ def _indexed_row_for_session(session: Session, path: Path, webui_dir: Path) -> d
"title": _metadata_title(session.metadata), "title": _metadata_title(session.metadata),
"preview": _preview_from_messages(session.messages), "preview": _preview_from_messages(session.messages),
_MODEL_PRESET_FIELD: model_preset_from_metadata(session.metadata), _MODEL_PRESET_FIELD: model_preset_from_metadata(session.metadata),
"recovery_state": recovery_state_from_metadata(session.metadata),
**_indexed_workspace_scope_fields(session.metadata), **_indexed_workspace_scope_fields(session.metadata),
_ROW_SOURCE_FIELD: _SESSION_SOURCE, _ROW_SOURCE_FIELD: _SESSION_SOURCE,
"file": path.name, "file": path.name,
@@ -601,6 +604,7 @@ def _scan_transcript_row(
"title": "", "title": "",
"preview": preview or fallback_preview, "preview": preview or fallback_preview,
_MODEL_PRESET_FIELD: None, _MODEL_PRESET_FIELD: None,
"recovery_state": None,
**_indexed_workspace_scope_fields({}), **_indexed_workspace_scope_fields({}),
_ROW_SOURCE_FIELD: _TRANSCRIPT_SOURCE, _ROW_SOURCE_FIELD: _TRANSCRIPT_SOURCE,
"file": stem, "file": stem,
@@ -687,6 +691,7 @@ def _scan_session_row(
"title": _metadata_title(metadata), "title": _metadata_title(metadata),
"preview": preview or fallback_preview, "preview": preview or fallback_preview,
_MODEL_PRESET_FIELD: model_preset_from_metadata(metadata), _MODEL_PRESET_FIELD: model_preset_from_metadata(metadata),
"recovery_state": recovery_state_from_metadata(metadata),
**_indexed_workspace_scope_fields(metadata), **_indexed_workspace_scope_fields(metadata),
_ROW_SOURCE_FIELD: _SESSION_SOURCE, _ROW_SOURCE_FIELD: _SESSION_SOURCE,
"file": path.name, "file": path.name,
+38 -23
View File
@@ -27,6 +27,7 @@ from nanobot.providers.image_generation import (
) )
from nanobot.providers.registry import find_by_name from nanobot.providers.registry import find_by_name
from nanobot.security.network import is_loopback_host from nanobot.security.network import is_loopback_host
from nanobot.utils.cancellation import shield_and_drain
from nanobot.webui.settings_contracts import ( from nanobot.webui.settings_contracts import (
QueryParams, QueryParams,
SettingsRequest, SettingsRequest,
@@ -640,7 +641,11 @@ class CapabilitySettingsHandler:
) -> SettingsRouteResult: ) -> SettingsRouteResult:
if action == "api-status": if action == "api-status":
return SettingsRouteResult.success( return SettingsRouteResult.success(
api_service_payload(self.settings, operations.api_runtime()) await asyncio.to_thread(
api_service_payload,
self.settings,
operations.api_runtime(),
)
) )
if action == "api-start": if action == "api-start":
return await self._start_api(request, operations) return await self._start_api(request, operations)
@@ -673,17 +678,22 @@ class CapabilitySettingsHandler:
return SettingsRouteResult.failure(404, "unknown settings action") return SettingsRouteResult.failure(404, "unknown settings action")
operation, section, apply_image_reload = mutation operation, section, apply_image_reload = mutation
try:
payload = self.settings.mutate(operation, request.query) async def mutate_and_apply() -> tuple[dict[str, Any], bool]:
except WebUISettingsError as exc: payload = await self.settings.mutate_async(operation, request.query)
return SettingsRouteResult.failure(exc.status, exc.message) if not apply_image_reload:
if apply_image_reload: return payload, False
payload, image_restart_cleared = await self.apply_image_runtime_change( return await self.apply_image_runtime_change(
payload, payload,
operations.reload_image, operations.reload_image,
) )
else:
image_restart_cleared = False try:
payload, image_restart_cleared = await shield_and_drain(
mutate_and_apply()
)
except WebUISettingsError as exc:
return SettingsRouteResult.failure(exc.status, exc.message)
return SettingsRouteResult.success( return SettingsRouteResult.success(
payload, payload,
decorate_restart=True, decorate_restart=True,
@@ -726,16 +736,17 @@ class CapabilitySettingsHandler:
400, 400,
"API service API key must be a string", "API service API key must be a string",
) )
try: allow_install = await self._allow_feature_package_install(request)
await asyncio.to_thread(
self.settings.mutate, async def mutate_and_start() -> Any:
await self.settings.mutate_async(
operations.nanobot_features_action, operations.nanobot_features_action,
"enable", "enable",
{"name": ["api"]}, {"name": ["api"]},
allow_install=self._allow_feature_package_install(request), allow_install=allow_install,
) )
self.settings.mutate(operations.update_api, request.query) await self.settings.mutate_async(operations.update_api, request.query)
config = self.settings.config.load() config = await self.settings.config.load_async()
runtime = operations.api_runtime() runtime = operations.api_runtime()
options = ApiStartOptions( options = ApiStartOptions(
host=config.api.host, host=config.api.host,
@@ -744,10 +755,13 @@ class CapabilitySettingsHandler:
config_path=str(self.settings.config.path), config_path=str(self.settings.config.path),
) )
current = runtime.status() current = runtime.status()
result = await asyncio.to_thread( return await asyncio.to_thread(
runtime.restart if current.running else runtime.start_background, runtime.restart if current.running else runtime.start_background,
options, options,
) )
try:
result = await shield_and_drain(mutate_and_start())
if not result.ok: if not result.ok:
return SettingsRouteResult.failure( return SettingsRouteResult.failure(
500, 500,
@@ -762,7 +776,8 @@ class CapabilitySettingsHandler:
self.logger.exception("failed to start managed API service") self.logger.exception("failed to start managed API service")
return SettingsRouteResult.failure(500, str(exc)) return SettingsRouteResult.failure(500, str(exc))
return SettingsRouteResult.success( return SettingsRouteResult.success(
api_service_payload( await asyncio.to_thread(
api_service_payload,
self.settings, self.settings,
operations.api_runtime(), operations.api_runtime(),
last_action="started", last_action="started",
@@ -775,7 +790,7 @@ class CapabilitySettingsHandler:
) -> SettingsRouteResult: ) -> SettingsRouteResult:
runtime = operations.api_runtime() runtime = operations.api_runtime()
try: try:
result = await asyncio.to_thread(runtime.stop) result = await shield_and_drain(asyncio.to_thread(runtime.stop))
except Exception as exc: except Exception as exc:
self.logger.exception("failed to stop managed API service") self.logger.exception("failed to stop managed API service")
return SettingsRouteResult.failure(500, str(exc)) return SettingsRouteResult.failure(500, str(exc))
@@ -785,20 +800,20 @@ class CapabilitySettingsHandler:
api_runtime_message(result.message), api_runtime_message(result.message),
) )
return SettingsRouteResult.success( return SettingsRouteResult.success(
api_service_payload( await asyncio.to_thread(
api_service_payload,
self.settings, self.settings,
operations.api_runtime(), operations.api_runtime(),
last_action="stopped", last_action="stopped",
) )
) )
def _allow_feature_package_install(self, request: SettingsRequest) -> bool: async def _allow_feature_package_install(self, request: SettingsRequest) -> bool:
if request.local_browser: if request.local_browser:
return True return True
try: try:
return bool( config = await self.settings.config.load_async()
self.settings.config.load().tools.webui_allow_remote_package_install return bool(config.tools.webui_allow_remote_package_install)
)
except Exception: except Exception:
self.logger.exception("failed to load remote package install policy") self.logger.exception("failed to load remote package install policy")
return False return False
+49 -13
View File
@@ -29,6 +29,7 @@ from nanobot.config.schema import Config, FallbackCandidate, ModelPresetConfig,
from nanobot.providers.image_generation import get_image_gen_provider from nanobot.providers.image_generation import get_image_gen_provider
from nanobot.providers.oauth_guidance import OAUTH_CLI_KIT_MISSING_MESSAGE from nanobot.providers.oauth_guidance import OAUTH_CLI_KIT_MISSING_MESSAGE
from nanobot.providers.registry import PROVIDERS, create_dynamic_spec, find_by_name from nanobot.providers.registry import PROVIDERS, create_dynamic_spec, find_by_name
from nanobot.utils.cancellation import shield_and_drain
from nanobot.webui.settings_contracts import ( from nanobot.webui.settings_contracts import (
QueryParams, QueryParams,
SettingsRequest, SettingsRequest,
@@ -1651,6 +1652,30 @@ class ModelSettingsHandler:
if self.settings.refresh_runtime_config is not None: if self.settings.refresh_runtime_config is not None:
self.settings.refresh_runtime_config() self.settings.refresh_runtime_config()
async def _mutate_and_refresh(
self,
operation: SettingsOperation,
query: QueryParams,
**kwargs: Any,
) -> dict[str, Any]:
payload = await self.settings.mutate_async(operation, query, **kwargs)
self._refresh_runtime_config()
return payload
async def _update_provider_and_runtime(
self,
operation: SettingsOperation,
query: QueryParams,
apply_image_runtime_change: Callable[
[dict[str, Any]],
Awaitable[tuple[dict[str, Any], bool]],
],
) -> tuple[dict[str, Any], bool]:
payload = await self.settings.mutate_async(operation, query)
payload, image_restart_cleared = await apply_image_runtime_change(payload)
self._refresh_runtime_config()
return payload, image_restart_cleared
async def handle( async def handle(
self, self,
action: str, action: str,
@@ -1659,8 +1684,12 @@ class ModelSettingsHandler:
) -> SettingsRouteResult: ) -> SettingsRouteResult:
try: try:
if action == "agent-update": if action == "agent-update":
payload = self.settings.mutate(operations.update_agent, request.query) payload = await shield_and_drain(
self._refresh_runtime_config() self._mutate_and_refresh(
operations.update_agent,
request.query,
)
)
return SettingsRouteResult.success( return SettingsRouteResult.success(
payload, payload,
decorate_restart=True, decorate_restart=True,
@@ -1668,12 +1697,13 @@ class ModelSettingsHandler:
) )
if action == "model-update": if action == "model-update":
payload = self.settings.mutate( payload = await shield_and_drain(
self._mutate_and_refresh(
operations.update_model, operations.update_model,
request.query, request.query,
rename_model_preset=self.settings.rename_model_preset, rename_model_preset=self.settings.rename_model_preset,
) )
self._refresh_runtime_config() )
return SettingsRouteResult.success(payload, decorate_restart=True) return SettingsRouteResult.success(payload, decorate_restart=True)
mutation = { mutation = {
@@ -1684,19 +1714,19 @@ class ModelSettingsHandler:
"provider-create": operations.create_provider, "provider-create": operations.create_provider,
}.get(action) }.get(action)
if mutation is not None: if mutation is not None:
payload = self.settings.mutate(mutation, request.query) payload = await shield_and_drain(
self._refresh_runtime_config() self._mutate_and_refresh(mutation, request.query)
)
return SettingsRouteResult.success(payload, decorate_restart=True) return SettingsRouteResult.success(payload, decorate_restart=True)
if action == "provider-update": if action == "provider-update":
payload = self.settings.mutate( payload, image_restart_cleared = await shield_and_drain(
self._update_provider_and_runtime(
operations.update_provider, operations.update_provider,
request.query, request.query,
operations.apply_image_runtime_change,
) )
payload, image_restart_cleared = await operations.apply_image_runtime_change(
payload
) )
self._refresh_runtime_config()
return SettingsRouteResult.success( return SettingsRouteResult.success(
payload, payload,
decorate_restart=True, decorate_restart=True,
@@ -1724,32 +1754,38 @@ class ModelSettingsHandler:
return SettingsRouteResult.success(payload) return SettingsRouteResult.success(payload)
if action == "oauth-login": if action == "oauth-login":
payload = await asyncio.to_thread( payload = await shield_and_drain(
asyncio.to_thread(
self.settings.read, self.settings.read,
operations.oauth_login, operations.oauth_login,
request.query, request.query,
oauth_flows=self.settings.oauth_flows, oauth_flows=self.settings.oauth_flows,
) )
)
elif action == "oauth-complete": elif action == "oauth-complete":
raw_response = (request.payload or {}).get("authorization_response") raw_response = (request.payload or {}).get("authorization_response")
if raw_response is not None and not isinstance(raw_response, str): if raw_response is not None and not isinstance(raw_response, str):
raise WebUISettingsError( raise WebUISettingsError(
"OAuth authorization response must be a string" "OAuth authorization response must be a string"
) )
payload = await asyncio.to_thread( payload = await shield_and_drain(
asyncio.to_thread(
self.settings.read, self.settings.read,
operations.oauth_complete, operations.oauth_complete,
request.query, request.query,
raw_response or None, raw_response or None,
oauth_flows=self.settings.oauth_flows, oauth_flows=self.settings.oauth_flows,
) )
)
elif action == "oauth-logout": elif action == "oauth-logout":
payload = await asyncio.to_thread( payload = await shield_and_drain(
asyncio.to_thread(
self.settings.read, self.settings.read,
operations.oauth_logout, operations.oauth_logout,
request.query, request.query,
oauth_flows=self.settings.oauth_flows, oauth_flows=self.settings.oauth_flows,
) )
)
else: else:
return SettingsRouteResult.failure(404, "unknown settings action") return SettingsRouteResult.failure(404, "unknown settings action")
except WebUISettingsError as exc: except WebUISettingsError as exc:
+33 -18
View File
@@ -4,6 +4,7 @@ from __future__ import annotations
import asyncio import asyncio
import html import html
import inspect
import json import json
from collections.abc import Awaitable, Callable, Mapping from collections.abc import Awaitable, Callable, Mapping
from typing import Any, cast from typing import Any, cast
@@ -18,6 +19,7 @@ from nanobot.bus.queue import MessageBus
from nanobot.channels.registry import load_channel_plugin from nanobot.channels.registry import load_channel_plugin
from nanobot.channels.validation import validate_channel_config from nanobot.channels.validation import validate_channel_config
from nanobot.pairing import approve_code, deny_code, list_pending from nanobot.pairing import approve_code, deny_code, list_pending
from nanobot.utils.cancellation import shield_and_drain
from nanobot.webui import settings_capabilities as capability_domain from nanobot.webui import settings_capabilities as capability_domain
from nanobot.webui import settings_contracts as contracts from nanobot.webui import settings_contracts as contracts
from nanobot.webui import settings_models as model_domain from nanobot.webui import settings_models as model_domain
@@ -208,6 +210,16 @@ def _payload_query(payload: dict[str, Any]) -> QueryParams:
} }
async def _call_settings_handler(
handler: Callable[[], Response | Awaitable[Response]],
) -> Response:
"""Keep synchronous handlers off-loop while supporting native async handlers."""
result = await asyncio.to_thread(handler)
if inspect.isawaitable(result):
return await result
return result
class WebUISettingsRouter: class WebUISettingsRouter:
"""Authenticate and dispatch settings requests to transport-neutral domains.""" """Authenticate and dispatch settings requests to transport-neutral domains."""
@@ -284,9 +296,9 @@ class WebUISettingsRouter:
if not self._authorized(request): if not self._authorized(request):
return self._unauthorized() return self._unauthorized()
if route == ("root", "settings"): if route == ("root", "settings"):
return self._handle_settings() return await _call_settings_handler(self._handle_settings)
if route == ("root", "usage"): if route == ("root", "usage"):
return self._handle_settings_usage() return await _call_settings_handler(self._handle_settings_usage)
domain, action = route domain, action = route
domain_request = self._domain_request( domain_request = self._domain_request(
@@ -415,19 +427,18 @@ class WebUISettingsRouter:
) )
return self._json_response(payload) return self._json_response(payload)
def _handle_settings(self) -> Response: async def _handle_settings(self) -> Response:
return self._json_response( payload = await self.settings.read_async(
self._with_restart_state(
self.settings.read(
settings_payload, settings_payload,
surface=self._runtime_surface, surface=self._runtime_surface,
runtime_capability_overrides=self._runtime_capabilities, runtime_capability_overrides=self._runtime_capabilities,
) )
) return self._json_response(self._with_restart_state(payload))
)
def _handle_settings_usage(self) -> Response: async def _handle_settings_usage(self) -> Response:
return self._json_response(self.settings.read(settings_usage_payload)) return self._json_response(
await self.settings.read_async(settings_usage_payload)
)
def _model_operations(self) -> model_domain.ModelSettingsOperations: def _model_operations(self) -> model_domain.ModelSettingsOperations:
return model_domain.ModelSettingsOperations( return model_domain.ModelSettingsOperations(
@@ -560,7 +571,7 @@ class WebUISettingsRouter:
allow_install=allow_install, allow_install=allow_install,
) )
def _allow_feature_package_install( async def _allow_feature_package_install(
self, self,
connection: Any, connection: Any,
request: WsRequest, request: WsRequest,
@@ -570,29 +581,33 @@ class WebUISettingsRouter:
request, request,
needs_local_browser=True, needs_local_browser=True,
) )
return self._system.allow_feature_package_install(domain_request) return await self._system.allow_feature_package_install(domain_request)
async def _handle_mcp_oauth_start(self, request: WsRequest) -> Response: async def _handle_mcp_oauth_start(self, request: WsRequest) -> Response:
if not self._authorized(request): if not self._authorized(request):
return self._unauthorized() return self._unauthorized()
if self._mcp_oauth_redirect_uri is None: redirect_uri_for_request = self._mcp_oauth_redirect_uri
if redirect_uri_for_request is None:
return self._error_response(500, "MCP OAuth callback is not configured") return self._error_response(500, "MCP OAuth callback is not configured")
query = self._parse_mcp_settings_query(request) query = self._parse_mcp_settings_query(request)
try:
name, cfg = await asyncio.to_thread( async def mutate_and_start() -> dict[str, Any]:
self.settings.mutate, name, cfg = await self.settings.mutate_async(
ensure_mcp_oauth_server, ensure_mcp_oauth_server,
query, query,
) )
redirect_uri = self._mcp_oauth_redirect_uri(request) redirect_uri = redirect_uri_for_request(request)
reset = (_query_first(query, "reset") or "").lower() in {"1", "true", "yes"} reset = (_query_first(query, "reset") or "").lower() in {"1", "true", "yes"}
payload = await self._mcp_oauth.start( return await self._mcp_oauth.start(
name, name,
cfg, cfg,
redirect_uri, redirect_uri,
reload_mcp=self._reload_mcp_runtime, reload_mcp=self._reload_mcp_runtime,
reset_credentials=reset, reset_credentials=reset,
) )
try:
payload = await shield_and_drain(mutate_and_start())
except Exception as exc: except Exception as exc:
return self._mcp_oauth_error_response(exc, action="start") return self._mcp_oauth_error_response(exc, action="start")
return self._json_response(payload) return self._json_response(payload)
+43 -1
View File
@@ -2,6 +2,7 @@
from __future__ import annotations from __future__ import annotations
import asyncio
import threading import threading
from collections.abc import Callable from collections.abc import Callable
from dataclasses import dataclass from dataclasses import dataclass
@@ -12,9 +13,11 @@ from filelock import FileLock
from nanobot.config.loader import load_config, save_config from nanobot.config.loader import load_config, save_config
from nanobot.config.schema import Config from nanobot.config.schema import Config
from nanobot.utils.cancellation import shield_and_drain
_T = TypeVar("_T") _T = TypeVar("_T")
_WEBUI_OAUTH_MAX_FLOWS = 8 _WEBUI_OAUTH_MAX_FLOWS = 8
_SETTINGS_FILE_LOCK_TIMEOUT_SECONDS = 5
class WebUISettingsConfig: class WebUISettingsConfig:
@@ -25,13 +28,20 @@ class WebUISettingsConfig:
self.path.parent.mkdir(parents=True, exist_ok=True) self.path.parent.mkdir(parents=True, exist_ok=True)
self._lock = threading.RLock() self._lock = threading.RLock()
lock_path = self.path.with_suffix(f"{self.path.suffix}.lock") lock_path = self.path.with_suffix(f"{self.path.suffix}.lock")
self._file_lock = FileLock(str(lock_path)) self._file_lock = FileLock(
str(lock_path),
timeout=_SETTINGS_FILE_LOCK_TIMEOUT_SECONDS,
)
def load(self) -> Config: def load(self) -> Config:
"""Load this gateway's config without consulting the process-global path.""" """Load this gateway's config without consulting the process-global path."""
with self._lock: with self._lock:
return load_config(self.path) return load_config(self.path)
async def load_async(self) -> Config:
"""Load config without running file I/O or lock waits on the event loop."""
return await asyncio.to_thread(self.load)
def update(self, mutation: Callable[[Config], _T]) -> _T: def update(self, mutation: Callable[[Config], _T]) -> _T:
"""Apply and atomically persist one path-scoped read-modify-write operation.""" """Apply and atomically persist one path-scoped read-modify-write operation."""
with self._lock, self._file_lock: with self._lock, self._file_lock:
@@ -40,11 +50,21 @@ class WebUISettingsConfig:
save_config(config, self.path) save_config(config, self.path)
return result return result
async def update_async(self, mutation: Callable[[Config], _T]) -> _T:
"""Update config without blocking the event loop."""
return await shield_and_drain(asyncio.to_thread(self.update, mutation))
def run_serialized(self, operation: Callable[[Path], _T]) -> _T: def run_serialized(self, operation: Callable[[Path], _T]) -> _T:
"""Run a path-aware read-modify-write operation under the config-file lock.""" """Run a path-aware read-modify-write operation under the config-file lock."""
with self._lock, self._file_lock: with self._lock, self._file_lock:
return operation(self.path) return operation(self.path)
async def run_serialized_async(self, operation: Callable[[Path], _T]) -> _T:
"""Run a serialized config operation without blocking the event loop."""
return await shield_and_drain(
asyncio.to_thread(self.run_serialized, operation)
)
class WebUIOAuthFlowRegistry: class WebUIOAuthFlowRegistry:
"""Bounded, thread-safe OAuth flows owned by one gateway instance.""" """Bounded, thread-safe OAuth flows owned by one gateway instance."""
@@ -146,6 +166,16 @@ class WebUISettingsServices:
"""Run a settings read against this gateway's explicit config path.""" """Run a settings read against this gateway's explicit config path."""
return operation(*args, config_path=self.config.path, **kwargs) return operation(*args, config_path=self.config.path, **kwargs)
async def read_async(
self,
operation: Callable[..., _T],
/,
*args: Any,
**kwargs: Any,
) -> _T:
"""Run a settings read without blocking the event loop."""
return await asyncio.to_thread(self.read, operation, *args, **kwargs)
def mutate( def mutate(
self, self,
operation: Callable[..., _T], operation: Callable[..., _T],
@@ -161,3 +191,15 @@ class WebUISettingsServices:
**kwargs, **kwargs,
) )
) )
async def mutate_async(
self,
operation: Callable[..., _T],
/,
*args: Any,
**kwargs: Any,
) -> _T:
"""Mutate settings without blocking the event loop."""
return await shield_and_drain(
asyncio.to_thread(self.mutate, operation, *args, **kwargs)
)
+64 -24
View File
@@ -20,8 +20,10 @@ from nanobot.channels.contracts import (
channel_update_instance_config, channel_update_instance_config,
) )
from nanobot.config.schema import Config from nanobot.config.schema import Config
from nanobot.llm_usage import llm_usage_payload
from nanobot.optional_features import OptionalFeatureError, with_channel_runtime_status from nanobot.optional_features import OptionalFeatureError, with_channel_runtime_status
from nanobot.security.workspace_access import workspace_sandbox_status from nanobot.security.workspace_access import workspace_sandbox_status
from nanobot.utils.cancellation import shield_and_drain
from nanobot.webui.settings_capabilities import network_safety_payload from nanobot.webui.settings_capabilities import network_safety_payload
from nanobot.webui.settings_contracts import ( from nanobot.webui.settings_contracts import (
QueryParams, QueryParams,
@@ -31,7 +33,6 @@ from nanobot.webui.settings_contracts import (
query_first, query_first,
query_first_alias, query_first_alias,
) )
from nanobot.webui.token_usage import token_usage_payload
if TYPE_CHECKING: if TYPE_CHECKING:
from nanobot.webui.settings_services import WebUISettingsServices from nanobot.webui.settings_services import WebUISettingsServices
@@ -121,7 +122,7 @@ def system_settings_payload(
}, },
"unified_session": defaults.unified_session, "unified_session": defaults.unified_session,
}, },
"usage": token_usage_payload(timezone_name=defaults.timezone), "usage": llm_usage_payload(timezone_name=defaults.timezone),
"advanced": { "advanced": {
"restrict_to_workspace": config.tools.restrict_to_workspace, "restrict_to_workspace": config.tools.restrict_to_workspace,
"workspace_sandbox": sandbox_status.as_dict(), "workspace_sandbox": sandbox_status.as_dict(),
@@ -139,7 +140,7 @@ def system_settings_payload(
def settings_usage_payload(config: Config) -> dict[str, Any]: def settings_usage_payload(config: Config) -> dict[str, Any]:
"""Return the lightweight token usage slice for Overview refreshes.""" """Return the lightweight token usage slice for Overview refreshes."""
return token_usage_payload(timezone_name=config.agents.defaults.timezone) return llm_usage_payload(timezone_name=config.agents.defaults.timezone)
def update_agent_system_settings(config: Config, query: QueryParams) -> tuple[bool, bool]: def update_agent_system_settings(config: Config, query: QueryParams) -> tuple[bool, bool]:
@@ -446,12 +447,17 @@ class SystemSettingsHandler:
operations: SystemSettingsOperations, operations: SystemSettingsOperations,
) -> SettingsRouteResult: ) -> SettingsRouteResult:
try: try:
payload = await asyncio.to_thread( pending = asyncio.to_thread(
operations.cli_apps_action, operations.cli_apps_action,
action, action,
request.query, request.query,
config_path=self.settings.config.path, config_path=self.settings.config.path,
) )
payload = (
await shield_and_drain(pending)
if action in {"install", "update", "uninstall"}
else await pending
)
except WebUISettingsError as exc: except WebUISettingsError as exc:
return SettingsRouteResult.failure(exc.status, exc.message) return SettingsRouteResult.failure(exc.status, exc.message)
except Exception as exc: except Exception as exc:
@@ -505,17 +511,29 @@ class SystemSettingsHandler:
action: str, action: str,
operations: SystemSettingsOperations, operations: SystemSettingsOperations,
) -> SettingsRouteResult: ) -> SettingsRouteResult:
try: allow_install = (
action != "enable"
or await self.allow_feature_package_install(request)
)
async def mutate_and_apply() -> dict[str, Any]:
payload = await asyncio.to_thread( payload = await asyncio.to_thread(
self._nanobot_features_action, self._nanobot_features_action,
action, action,
request.query, request.query,
operations, operations,
allow_install=( allow_install=allow_install,
action != "enable"
or self.allow_feature_package_install(request)
),
) )
payload = await self._apply_feature_runtime_change(
action,
request.query,
payload,
operations,
)
return self._with_channel_runtime_status(payload, operations)
try:
payload = await shield_and_drain(mutate_and_apply())
except OptionalFeatureError as exc: except OptionalFeatureError as exc:
return SettingsRouteResult.failure(exc.status, exc.message) return SettingsRouteResult.failure(exc.status, exc.message)
except Exception as exc: except Exception as exc:
@@ -527,13 +545,6 @@ class SystemSettingsHandler:
action, action,
) )
return SettingsRouteResult.failure(status, message) return SettingsRouteResult.failure(status, message)
payload = await self._apply_feature_runtime_change(
action,
request.query,
payload,
operations,
)
payload = self._with_channel_runtime_status(payload, operations)
return SettingsRouteResult.success( return SettingsRouteResult.success(
payload, payload,
decorate_restart=True, decorate_restart=True,
@@ -628,6 +639,15 @@ class SystemSettingsHandler:
self, self,
request: SettingsRequest, request: SettingsRequest,
operations: SystemSettingsOperations, operations: SystemSettingsOperations,
) -> SettingsRouteResult:
return await shield_and_drain(
self._channel_configure_settled(request, operations)
)
async def _channel_configure_settled(
self,
request: SettingsRequest,
operations: SystemSettingsOperations,
) -> SettingsRouteResult: ) -> SettingsRouteResult:
name = (query_first(request.query, "name") or "").strip() name = (query_first(request.query, "name") or "").strip()
instance_id = ( instance_id = (
@@ -682,7 +702,7 @@ class SystemSettingsHandler:
"enable", "enable",
feature_query, feature_query,
operations, operations,
allow_install=self.allow_feature_package_install(request), allow_install=await self.allow_feature_package_install(request),
) )
except OptionalFeatureError as exc: except OptionalFeatureError as exc:
return SettingsRouteResult.failure( return SettingsRouteResult.failure(
@@ -825,6 +845,22 @@ class SystemSettingsHandler:
channel_name: str, channel_name: str,
payload: dict[str, Any], payload: dict[str, Any],
operations: SystemSettingsOperations, operations: SystemSettingsOperations,
) -> dict[str, Any]:
return await shield_and_drain(
self._settle_channel_connect_success(
request,
channel_name,
payload,
operations,
)
)
async def _settle_channel_connect_success(
self,
request: SettingsRequest,
channel_name: str,
payload: dict[str, Any],
operations: SystemSettingsOperations,
) -> dict[str, Any]: ) -> dict[str, Any]:
target = {"name": [channel_name]} target = {"name": [channel_name]}
if payload.get("instance_id"): if payload.get("instance_id"):
@@ -835,11 +871,11 @@ class SystemSettingsHandler:
"enable", "enable",
target, target,
operations, operations,
allow_install=self.allow_feature_package_install(request), allow_install=await self.allow_feature_package_install(request),
) )
except OptionalFeatureError as exc: except OptionalFeatureError as exc:
features = self.feature_runtime_fallback( features = self.feature_runtime_fallback(
self._nanobot_features_payload(operations), await asyncio.to_thread(self._nanobot_features_payload, operations),
message=( message=(
f"{channel_name} connected, but enabling channel support failed: " f"{channel_name} connected, but enabling channel support failed: "
f"{exc.message}" f"{exc.message}"
@@ -859,13 +895,12 @@ class SystemSettingsHandler:
) )
return updated return updated
def allow_feature_package_install(self, request: SettingsRequest) -> bool: async def allow_feature_package_install(self, request: SettingsRequest) -> bool:
if request.local_browser: if request.local_browser:
return True return True
try: try:
return bool( config = await self.settings.config.load_async()
self.settings.config.load().tools.webui_allow_remote_package_install return bool(config.tools.webui_allow_remote_package_install)
)
except Exception: except Exception:
self.logger.exception("failed to load remote package install policy") self.logger.exception("failed to load remote package install policy")
return False return False
@@ -925,13 +960,18 @@ class SystemSettingsHandler:
operations: SystemSettingsOperations, operations: SystemSettingsOperations,
) -> SettingsRouteResult: ) -> SettingsRouteResult:
try: try:
payload = await operations.mcp_presets_action( pending = operations.mcp_presets_action(
action, action,
request.query, request.query,
reload_mcp=operations.reload_mcp, reload_mcp=operations.reload_mcp,
mcp_runtime_status=operations.mcp_runtime_status, mcp_runtime_status=operations.mcp_runtime_status,
config=self.settings.config, config=self.settings.config,
) )
payload = (
await pending
if action is None
else await shield_and_drain(pending)
)
except Exception as exc: except Exception as exc:
status = getattr(exc, "status", 500) status = getattr(exc, "status", 500)
message = getattr(exc, "message", str(exc)) message = getattr(exc, "message", str(exc))
-370
View File
@@ -1,370 +0,0 @@
"""Workspace-scoped token usage telemetry for WebUI overview surfaces."""
from __future__ import annotations
import json
import os
import threading
import time
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Any, Mapping, cast
from zoneinfo import ZoneInfo, ZoneInfoNotFoundError
from loguru import logger
from nanobot.agent.hook import AgentHook, AgentHookContext
from nanobot.config.paths import get_webui_dir
TOKEN_USAGE_SCHEMA_VERSION = 1
_MAX_STATE_FILE_BYTES = 512 * 1024
_MAX_DAYS_RETAINED = 400
_USAGE_KEYS = (
"prompt_tokens",
"completion_tokens",
"cached_tokens",
"total_tokens",
"provider_tokens",
"estimated_tokens",
)
_REQUEST_KEYS = ("requests", "provider_requests", "estimated_requests")
_SOURCE_KEYS = ("user", "api", "cron", "dream", "system")
_WRITE_LOCK = threading.Lock()
def token_usage_state_path() -> Path:
return get_webui_dir() / "token-usage.json"
def default_token_usage_state() -> dict[str, Any]:
return {
"schema_version": TOKEN_USAGE_SCHEMA_VERSION,
"days": {},
"updated_at": None,
}
def _utc_now_iso() -> str:
return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
def _zone(timezone_name: str | None) -> timezone | ZoneInfo:
if not timezone_name:
return timezone.utc
try:
return ZoneInfo(timezone_name)
except ZoneInfoNotFoundError:
return timezone.utc
def _local_day(now: datetime | None = None, *, timezone_name: str | None = None) -> str:
dt = now or datetime.now(timezone.utc)
if dt.tzinfo is None:
dt = dt.replace(tzinfo=timezone.utc)
return dt.astimezone(_zone(timezone_name)).date().isoformat()
def _clean_int(value: Any) -> int:
try:
return max(0, int(value or 0))
except (TypeError, ValueError):
return 0
def _clean_source(value: str | None) -> str:
return value if value in _SOURCE_KEYS else "system"
def _source_from_session_key(session_key: str | None) -> str:
key = session_key or ""
if key.startswith("dream:"):
return "dream"
if key == "heartbeat" or key.startswith("cron:"):
return "cron"
if key.startswith("api:"):
return "api"
if key.startswith("system:"):
return "system"
return "user"
def _normalize_usage(raw: dict[str, Any] | None) -> dict[str, int]:
if not isinstance(raw, dict):
return {}
usage = {key: _clean_int(raw.get(key)) for key in _USAGE_KEYS}
fallback_total = usage["prompt_tokens"] + usage["completion_tokens"]
if usage["total_tokens"] <= 0:
usage["total_tokens"] = fallback_total
if usage["estimated_tokens"] <= 0 and usage["provider_tokens"] <= 0:
usage["provider_tokens"] = usage["total_tokens"]
elif usage["estimated_tokens"] > 0 and usage["provider_tokens"] <= 0:
usage["estimated_tokens"] = min(usage["estimated_tokens"], usage["total_tokens"])
elif usage["provider_tokens"] > 0 and usage["estimated_tokens"] <= 0:
usage["provider_tokens"] = min(usage["provider_tokens"], usage["total_tokens"])
return usage if usage["total_tokens"] > 0 else {}
def _normalize_usage_row(row: dict[str, Any]) -> dict[str, int]:
cleaned = {key: _clean_int(row.get(key)) for key in _USAGE_KEYS}
if cleaned["total_tokens"] <= 0:
cleaned["total_tokens"] = cleaned["prompt_tokens"] + cleaned["completion_tokens"]
if cleaned["provider_tokens"] <= 0 and cleaned["estimated_tokens"] <= 0:
cleaned["provider_tokens"] = cleaned["total_tokens"]
requests = {key: _clean_int(row.get(key)) for key in _REQUEST_KEYS}
if (
requests["requests"] > 0
and requests["provider_requests"] <= 0
and requests["estimated_requests"] <= 0
):
if cleaned["estimated_tokens"] > 0 and cleaned["provider_tokens"] <= 0:
requests["estimated_requests"] = requests["requests"]
else:
requests["provider_requests"] = requests["requests"]
return {**cleaned, **requests}
def _normalize_sources(raw: Any, fallback: dict[str, int]) -> dict[str, dict[str, int]]:
sources: dict[str, dict[str, int]] = {}
if isinstance(raw, dict):
for source, row_value in cast(dict[Any, Any], raw).items():
if not isinstance(row_value, dict):
continue
row = cast(dict[str, Any], row_value)
normalized = _normalize_usage_row(row)
if normalized["total_tokens"] <= 0 and normalized["requests"] <= 0:
continue
source_key = _clean_source(str(source))
current = sources.get(source_key)
if current is None:
sources[source_key] = normalized
else:
for key in (*_USAGE_KEYS, *_REQUEST_KEYS):
current[key] = _clean_int(current.get(key)) + normalized[key]
if not sources and (fallback["total_tokens"] > 0 or fallback["requests"] > 0):
sources["user"] = {key: fallback[key] for key in (*_USAGE_KEYS, *_REQUEST_KEYS)}
return sources
def normalize_token_usage_state(raw: Any) -> dict[str, Any]:
state = default_token_usage_state()
if not isinstance(raw, dict):
return state
raw = cast(dict[str, Any], raw)
days_raw = raw.get("days")
if not isinstance(days_raw, dict):
return state
days: dict[str, dict[str, Any]] = {}
for date, row_value in sorted(cast(dict[Any, Any], days_raw).items())[-_MAX_DAYS_RETAINED:]:
if not isinstance(date, str) or len(date) != 10 or not isinstance(row_value, dict):
continue
row = cast(dict[str, Any], row_value)
try:
datetime.fromisoformat(date)
except ValueError:
# A hand-edited or foreign day key that is not a real date would
# otherwise reach token_usage_payload's date parsing and fail every
# settings request; drop it like any other malformed row.
continue
normalized = _normalize_usage_row(row)
if normalized["total_tokens"] <= 0 and normalized["requests"] <= 0:
continue
days[date] = {
"date": date,
**normalized,
"sources": _normalize_sources(row.get("sources"), normalized),
}
state["days"] = days
updated_at = raw.get("updated_at")
state["updated_at"] = updated_at if isinstance(updated_at, str) else None
return state
def read_token_usage_state() -> dict[str, Any]:
path = token_usage_state_path()
if not path.is_file():
return default_token_usage_state()
try:
if path.stat().st_size > _MAX_STATE_FILE_BYTES:
logger.warning("token usage state too large, ignoring: {}", path)
return default_token_usage_state()
with open(path, encoding="utf-8") as f:
raw = json.load(f)
except (OSError, json.JSONDecodeError) as e:
logger.warning("read token usage state failed {}: {}", path, e)
return default_token_usage_state()
return normalize_token_usage_state(raw)
def write_token_usage_state(raw: dict[str, Any]) -> dict[str, Any]:
state = normalize_token_usage_state(raw)
state["updated_at"] = _utc_now_iso()
encoded = json.dumps(
state,
ensure_ascii=False,
indent=2,
sort_keys=True,
).encode("utf-8")
if len(encoded) > _MAX_STATE_FILE_BYTES:
raise ValueError("token usage state is too large")
path = token_usage_state_path()
path.parent.mkdir(parents=True, exist_ok=True)
tmp = path.with_suffix(".json.tmp")
with open(tmp, "wb") as f:
f.write(encoded)
f.write(b"\n")
f.flush()
os.fsync(f.fileno())
os.replace(tmp, path)
try:
dir_fd = os.open(path.parent, os.O_RDONLY)
except OSError:
return state
try:
os.fsync(dir_fd)
finally:
os.close(dir_fd)
return state
def record_token_usage(
usage: dict[str, Any] | None,
*,
source: str = "user",
timezone_name: str | None = None,
now: datetime | None = None,
) -> dict[str, Any]:
normalized = _normalize_usage(usage)
if not normalized:
return read_token_usage_state()
with _WRITE_LOCK:
state = read_token_usage_state()
days_by_date = cast(dict[str, dict[str, Any]], state["days"])
day = _local_day(now, timezone_name=timezone_name)
row: dict[str, Any] = dict(days_by_date.get(day) or {"date": day, "requests": 0})
for key in _USAGE_KEYS:
row[key] = _clean_int(row.get(key)) + normalized.get(key, 0)
row["requests"] = _clean_int(row.get("requests")) + 1
if normalized.get("estimated_tokens", 0) > 0 and normalized.get("provider_tokens", 0) <= 0:
row["estimated_requests"] = _clean_int(row.get("estimated_requests")) + 1
else:
row["provider_requests"] = _clean_int(row.get("provider_requests")) + 1
source_key = _clean_source(source)
sources: dict[str, dict[str, Any]] = dict(
cast(Mapping[str, dict[str, Any]], row.get("sources") or {})
)
source_row: dict[str, Any] = dict(sources.get(source_key) or {"requests": 0})
for key in _USAGE_KEYS:
source_row[key] = _clean_int(source_row.get(key)) + normalized.get(key, 0)
source_row["requests"] = _clean_int(source_row.get("requests")) + 1
if normalized.get("estimated_tokens", 0) > 0 and normalized.get("provider_tokens", 0) <= 0:
source_row["estimated_requests"] = _clean_int(source_row.get("estimated_requests")) + 1
else:
source_row["provider_requests"] = _clean_int(source_row.get("provider_requests")) + 1
sources[source_key] = source_row
row["sources"] = sources
days_by_date[day] = row
if len(days_by_date) > _MAX_DAYS_RETAINED:
state["days"] = dict(sorted(days_by_date.items())[-_MAX_DAYS_RETAINED:])
return write_token_usage_state(state)
def record_response_token_usage(
response: Any,
*,
source: str,
timezone_name: str | None = None,
) -> None:
try:
record_token_usage(
getattr(response, "usage", None),
source=source,
timezone_name=timezone_name,
)
except Exception:
logger.exception("failed to record {} token usage", source)
def token_usage_payload(
*,
days: int = 371,
timezone_name: str | None = None,
now: datetime | None = None,
) -> dict[str, Any]:
state = read_token_usage_state()
days_by_date = cast(dict[str, dict[str, Any]], state["days"])
today = datetime.fromisoformat(_local_day(now, timezone_name=timezone_name)).date()
start = today - timedelta(days=max(1, days) - 1)
day_rows = [
row
for date, row in sorted(days_by_date.items())
if start.isoformat() <= date <= today.isoformat()
]
last_30_start = today - timedelta(days=29)
last_30 = [
row
for date, row in days_by_date.items()
if last_30_start.isoformat() <= date <= today.isoformat()
]
last_365_start = today - timedelta(days=364)
last_365 = [
row
for date, row in days_by_date.items()
if last_365_start.isoformat() <= date <= today.isoformat()
]
active_dates = {
datetime.fromisoformat(date).date()
for date, row in days_by_date.items()
if _clean_int(row.get("total_tokens")) > 0
}
current_streak = 0
cursor = today
while cursor in active_dates:
current_streak += 1
cursor -= timedelta(days=1)
longest_streak = 0
running_streak = 0
for cursor in sorted(active_dates):
if cursor - timedelta(days=1) in active_dates:
running_streak += 1
else:
running_streak = 1
longest_streak = max(longest_streak, running_streak)
all_rows = list(days_by_date.values())
return {
"days": day_rows,
"total_tokens": sum(_clean_int(row.get("total_tokens")) for row in all_rows),
"total_tokens_30d": sum(_clean_int(row.get("total_tokens")) for row in last_30),
"total_tokens_365d": sum(_clean_int(row.get("total_tokens")) for row in last_365),
"peak_day_tokens": max([_clean_int(row.get("total_tokens")) for row in all_rows] or [0]),
"current_streak_days": current_streak,
"longest_streak_days": longest_streak,
"active_days_30d": sum(1 for row in last_30 if _clean_int(row.get("total_tokens")) > 0),
"requests_30d": sum(_clean_int(row.get("requests")) for row in last_30),
"updated_at": state.get("updated_at"),
}
class TokenUsageHook(AgentHook):
"""Persist provider-reported token usage without coupling it to chat messages."""
def __init__(self, *, timezone_name: str | None = None) -> None:
super().__init__()
self._timezone_name = timezone_name
async def after_iteration(self, context: AgentHookContext) -> None:
try:
record_token_usage(
context.usage,
source=_source_from_session_key(context.session_key),
timezone_name=self._timezone_name,
)
except Exception:
logger.exception("failed to record token usage")
+21 -26
View File
@@ -1813,11 +1813,16 @@ def replay_transcript_to_ui_messages(
break break
content = str(candidate.get("content") or "") content = str(candidate.get("content") or "")
has_answer = len(content) > 0 has_answer = len(content) > 0
if has_answer:
break
# A completed reasoning field is closed even while its assistant
# placeholder remains streaming for the rest of the turn.
if ( if (
candidate.get("reasoningStreaming") candidate.get("reasoningStreaming")
or candidate.get("reasoning") is not None or (
or has_answer candidate.get("isStreaming")
or candidate.get("isStreaming") and candidate.get("reasoning") is None
)
): ):
prev[i] = { prev[i] = {
**candidate, **candidate,
@@ -1827,15 +1832,6 @@ def replay_transcript_to_ui_messages(
**turn_fields, **turn_fields,
} }
return return
if not has_answer and candidate.get("isStreaming"):
prev[i] = {
**candidate,
"reasoning": chunk,
"reasoningStreaming": True,
"activitySegmentId": candidate.get("activitySegmentId") or _ensure_activity_segment(),
**turn_fields,
}
return
break break
segment = _ensure_activity_segment() segment = _ensure_activity_segment()
prev.append( prev.append(
@@ -1915,19 +1911,6 @@ def replay_transcript_to_ui_messages(
and not m.get("media") and not m.get("media")
) )
def is_tool_trace_at(index: int) -> bool:
m = messages[index] if 0 <= index < len(messages) else None
return bool(m and m.get("kind") == "trace")
def prune_reasoning_only() -> None:
nonlocal messages
kept: list[dict[str, Any]] = []
for i, m in enumerate(messages):
if is_reasoning_only_placeholder(m) and not is_tool_trace_at(i + 1):
continue
kept.append(m)
messages = kept
def stamp_completion( def stamp_completion(
*, *,
latency_ms: int | None = None, latency_ms: int | None = None,
@@ -2442,7 +2425,6 @@ def replay_transcript_to_ui_messages(
for i, m in enumerate(messages): for i, m in enumerate(messages):
if m.get("isStreaming"): if m.get("isStreaming"):
messages[i] = {**m, "isStreaming": False} messages[i] = {**m, "isStreaming": False}
prune_reasoning_only()
lat = rec.get("latency_ms") lat = rec.get("latency_ms")
usage = rec.get("usage") usage = rec.get("usage")
sanitized_usage = ( sanitized_usage = (
@@ -2541,6 +2523,19 @@ def has_pending_tool_calls(
return False return False
def has_unfinished_transcript_tail(session_key: str) -> bool:
"""Return whether the active transcript ends in an unfinished turn.
Recovery runs at gateway startup and only needs the newest, still-active
turn. Completed turns are rotated into immutable segment files, so reading
every historical segment here would make restart cost grow with the full
conversation history.
"""
return has_pending_tool_calls(
_read_transcript_file(webui_transcript_path(session_key))
)
def completed_turn_ids(lines: list[dict[str, Any]]) -> list[str]: def completed_turn_ids(lines: list[dict[str, Any]]) -> list[str]:
"""Return stable identities for turns with an explicitly persisted completion.""" """Return stable identities for turns with an explicitly persisted completion."""
completed: list[str] = [] completed: list[str] = []
+24 -1
View File
@@ -5,6 +5,7 @@ from __future__ import annotations
import json import json
import os import os
import time import time
from collections import OrderedDict
from pathlib import Path from pathlib import Path
from typing import TYPE_CHECKING, Any, cast from typing import TYPE_CHECKING, Any, cast
@@ -28,6 +29,7 @@ _MAX_STATE_FILE_BYTES = 128 * 1024
_DEFAULT_ACCESS_MODES = {"default", "full"} _DEFAULT_ACCESS_MODES = {"default", "full"}
_LEGACY_RESTRICTED_DEFAULT_ACCESS_MODE = "restricted" _LEGACY_RESTRICTED_DEFAULT_ACCESS_MODE = "restricted"
_WEBUI_SCOPE_CHANNEL = "websocket" _WEBUI_SCOPE_CHANNEL = "websocket"
_MAX_DRAFT_SCOPES = 128
def _scope_change_is_non_escalating(current: WorkspaceScope, requested: WorkspaceScope) -> bool: def _scope_change_is_non_escalating(current: WorkspaceScope, requested: WorkspaceScope) -> bool:
@@ -186,6 +188,7 @@ class WebUIWorkspaceController:
self._sessions = session_manager self._sessions = session_manager
self._default_workspace = default_workspace self._default_workspace = default_workspace
self._default_restrict_to_workspace = default_restrict_to_workspace self._default_restrict_to_workspace = default_restrict_to_workspace
self._draft_scopes: OrderedDict[str, WorkspaceScope] = OrderedDict()
def default_scope(self) -> WorkspaceScope: def default_scope(self) -> WorkspaceScope:
return default_scope_for_webui( return default_scope_for_webui(
@@ -230,6 +233,10 @@ class WebUIWorkspaceController:
return self._scope_from_metadata_value(raw_scope, default_scope=default_scope) return self._scope_from_metadata_value(raw_scope, default_scope=default_scope)
def scope_for_session_key(self, session_key: str) -> WorkspaceScope: def scope_for_session_key(self, session_key: str) -> WorkspaceScope:
draft = self._draft_scopes.get(session_key)
if draft is not None:
self._draft_scopes.move_to_end(session_key)
return draft
if self._sessions is None: if self._sessions is None:
return self.default_scope() return self.default_scope()
data = self._sessions.read_session_metadata(session_key) data = self._sessions.read_session_metadata(session_key)
@@ -328,8 +335,24 @@ class WebUIWorkspaceController:
return scope return scope
def persist_scope(self, chat_id: str, scope: WorkspaceScope) -> None: def persist_scope(self, chat_id: str, scope: WorkspaceScope) -> None:
session_key = f"websocket:{chat_id}"
if self._sessions is not None: if self._sessions is not None:
session = self._sessions.get_or_create(f"websocket:{chat_id}") session = self._sessions.get_or_create(session_key)
session.metadata["webui"] = True session.metadata["webui"] = True
session.metadata[WORKSPACE_SCOPE_METADATA_KEY] = scope.metadata() session.metadata[WORKSPACE_SCOPE_METADATA_KEY] = scope.metadata()
self._sessions.save(session) self._sessions.save(session)
self._draft_scopes.pop(session_key, None)
def stage_scope(self, chat_id: str, scope: WorkspaceScope) -> None:
"""Keep a new chat's scope transient until its first accepted message."""
session_key = f"websocket:{chat_id}"
if (
self._sessions is not None
and self._sessions.read_session_metadata(session_key) is not None
):
self.persist_scope(chat_id, scope)
return
self._draft_scopes[session_key] = scope
self._draft_scopes.move_to_end(session_key)
while len(self._draft_scopes) > _MAX_DRAFT_SCOPES:
self._draft_scopes.popitem(last=False)
+152 -24
View File
@@ -29,10 +29,12 @@ from nanobot.cron.session_turns import is_bound_cron_job
from nanobot.cron.types import CronJob, CronSchedule from nanobot.cron.types import CronJob, CronSchedule
from nanobot.security.workspace_access import WorkspaceScope from nanobot.security.workspace_access import WorkspaceScope
from nanobot.session.manager import SessionManager from nanobot.session.manager import SessionManager
from nanobot.session.recovery import RecoveryActionError
from nanobot.session.session_handles import ( from nanobot.session.session_handles import (
SessionHandleResolver, SessionHandleResolver,
) )
from nanobot.triggers.local_types import LocalTrigger from nanobot.triggers.local_types import LocalTrigger
from nanobot.utils.cancellation import shield_and_drain
from nanobot.webui.file_preview import ( from nanobot.webui.file_preview import (
WebUIFilePreviewError, WebUIFilePreviewError,
file_preview_availability_payload, file_preview_availability_payload,
@@ -134,6 +136,18 @@ _WEBUI_MUTATION_PAYLOAD_ATTR = "_nanobot_webui_mutation_payload"
_WEBUI_MUTATION_REQUEST_ATTR = "_nanobot_webui_mutation_request" _WEBUI_MUTATION_REQUEST_ATTR = "_nanobot_webui_mutation_request"
_NO_STORE_HEADERS = [("Cache-Control", "no-store")] _NO_STORE_HEADERS = [("Cache-Control", "no-store")]
def _slow_http_operation(path: str) -> str:
"""Return a route family without logging user-controlled path/query values."""
clean_path = path.split("?", 1)[0]
if clean_path == "/webui/bootstrap":
return clean_path
parts = [part for part in clean_path.split("/") if part]
if len(parts) >= 2 and parts[0] == "api":
return f"/api/{parts[1]}"
return "/webui"
_WEBUI_MUTATION_PATHS = { _WEBUI_MUTATION_PATHS = {
"automation.enable": "/api/webui/automations/enable", "automation.enable": "/api/webui/automations/enable",
"automation.disable": "/api/webui/automations/disable", "automation.disable": "/api/webui/automations/disable",
@@ -145,6 +159,8 @@ _WEBUI_MUTATION_PATHS = {
"skill.delete": "/api/webui/skills/delete", "skill.delete": "/api/webui/skills/delete",
"sidebar.update": "/api/webui/sidebar-state/update", "sidebar.update": "/api/webui/sidebar-state/update",
"workspace.pick_folder": "/api/workspaces/pick-folder", "workspace.pick_folder": "/api/workspaces/pick-folder",
"recovery.continue": "/api/webui/recovery/continue",
"recovery.dismiss": "/api/webui/recovery/dismiss",
"settings.agent.update": "/api/settings/update", "settings.agent.update": "/api/settings/update",
"settings.model_configuration.create": "/api/settings/model-configurations/create", "settings.model_configuration.create": "/api/settings/model-configurations/create",
"settings.model_configuration.update": "/api/settings/model-configurations/update", "settings.model_configuration.update": "/api/settings/model-configurations/update",
@@ -323,6 +339,9 @@ class GatewayHTTPHandler:
mcp_runtime_status: Callable[[], Mapping[str, str]] | None = None, mcp_runtime_status: Callable[[], Mapping[str, str]] | None = None,
mcp_reload: Callable[[], Awaitable[dict[str, Any]]] | None = None, mcp_reload: Callable[[], Awaitable[dict[str, Any]]] | None = None,
skill_state_action: Callable[[set[str]], None] | None = None, skill_state_action: Callable[[set[str]], None] | None = None,
recovery_action: (
Callable[[str, dict[str, Any]], Awaitable[dict[str, Any]]] | None
) = None,
log: Any = logger, log: Any = logger,
) -> None: ) -> None:
self.config = config self.config = config
@@ -340,6 +359,7 @@ class GatewayHTTPHandler:
disabled_skills if disabled_skills is not None else set() disabled_skills if disabled_skills is not None else set()
) )
self.skill_state_action = skill_state_action self.skill_state_action = skill_state_action
self.recovery_action = recovery_action
self._skill_install_lock = asyncio.Lock() self._skill_install_lock = asyncio.Lock()
self._folder_picker_lock = asyncio.Lock() self._folder_picker_lock = asyncio.Lock()
self.cron_service = cron_service self.cron_service = cron_service
@@ -413,7 +433,12 @@ class GatewayHTTPHandler:
response = await self._dispatch_resolved(connection, request, got) response = await self._dispatch_resolved(connection, request, got)
return response return response
finally: finally:
self._log_slow_http(got, response, started) self._log_slow_http(
got,
response,
started,
input_chars=len(request.path),
)
async def dispatch_webui_mutation( async def dispatch_webui_mutation(
self, self,
@@ -454,6 +479,8 @@ class GatewayHTTPHandler:
return True return True
if re.match(r"^/api/webui/automations/(enable|disable|delete|run|update)$", path): if re.match(r"^/api/webui/automations/(enable|disable|delete|run|update)$", path):
return True return True
if path in {"/api/webui/recovery/continue", "/api/webui/recovery/dismiss"}:
return True
return path in { return path in {
"/api/webui/skills/install", "/api/webui/skills/install",
"/api/webui/skills/update", "/api/webui/skills/update",
@@ -507,6 +534,11 @@ class GatewayHTTPHandler:
if response is not None: if response is not None:
return response return response
# Recovery routes
response = await self._dispatch_recovery_route(request, got)
if response is not None:
return response
# Session routes # Session routes
response = await self._dispatch_session_routes(request, got) response = await self._dispatch_session_routes(request, got)
if response is not None: if response is not None:
@@ -542,7 +574,14 @@ class GatewayHTTPHandler:
return connection.respond(404, "Not Found") return connection.respond(404, "Not Found")
def _log_slow_http(self, path: str, response: Any | None, started: float) -> None: def _log_slow_http(
self,
path: str,
response: Any | None,
started: float,
*,
input_chars: int,
) -> None:
elapsed_ms = int((time.perf_counter() - started) * 1000) elapsed_ms = int((time.perf_counter() - started) * 1000)
if elapsed_ms < _SLOW_WEBUI_HTTP_LOG_MS: if elapsed_ms < _SLOW_WEBUI_HTTP_LOG_MS:
return return
@@ -550,9 +589,10 @@ class GatewayHTTPHandler:
return return
status = getattr(response, "status_code", None) status = getattr(response, "status_code", None)
self._log.warning( self._log.warning(
"slow webui http route path={} status={} duration_ms={}", "slow webui http operation={} status={} input_chars={} duration_ms={}",
path, _slow_http_operation(path),
status if status is not None else "none", status if status is not None else "none",
input_chars,
elapsed_ms, elapsed_ms,
) )
@@ -680,7 +720,11 @@ class GatewayHTTPHandler:
async def _dispatch_session_routes(self, request: WsRequest, got: str) -> Response | None: async def _dispatch_session_routes(self, request: WsRequest, got: str) -> Response | None:
m = re.match(r"^/api/sessions/([^/]+)/webui-thread$", got) m = re.match(r"^/api/sessions/([^/]+)/webui-thread$", got)
if m: if m:
return self._handle_webui_thread_get(request, m.group(1)) return await asyncio.to_thread(
self._handle_webui_thread_get,
request,
m.group(1),
)
m = re.match(r"^/api/sessions/([^/]+)/context$", got) m = re.match(r"^/api/sessions/([^/]+)/context$", got)
if m: if m:
@@ -688,18 +732,51 @@ class GatewayHTTPHandler:
m = re.match(r"^/api/sessions/([^/]+)/file-preview$", got) m = re.match(r"^/api/sessions/([^/]+)/file-preview$", got)
if m: if m:
return self._handle_file_preview(request, m.group(1)) return await asyncio.to_thread(
self._handle_file_preview,
request,
m.group(1),
)
m = re.match(r"^/api/sessions/([^/]+)/automations$", got) m = re.match(r"^/api/sessions/([^/]+)/automations$", got)
if m: if m:
return self._handle_session_automations(request, m.group(1)) return await self._run_cron_transaction(
self._handle_session_automations,
request,
m.group(1),
)
m = re.match(r"^/api/sessions/([^/]+)/delete$", got) m = re.match(r"^/api/sessions/([^/]+)/delete$", got)
if m: if m:
return self._handle_session_delete(request, m.group(1)) return await self._run_cron_transaction(
self._handle_session_delete,
request,
m.group(1),
)
return None return None
async def _dispatch_recovery_route(
self,
request: WsRequest,
path: str,
) -> Response | None:
match = re.fullmatch(r"/api/webui/recovery/(continue|dismiss)", path)
if match is None:
return None
if not getattr(request, _WEBUI_MUTATION_REQUEST_ATTR, False):
return _http_error(405, "WebUI recovery actions require an authenticated WebSocket")
if self.recovery_action is None:
return _http_error(503, "WebUI recovery is unavailable")
payload = _mutation_payload(request)
if payload is None:
return _http_error(400, "invalid recovery payload")
try:
result = await self.recovery_action(match.group(1), payload)
except RecoveryActionError as exc:
return _http_error(exc.status, str(exc))
return _http_json_response(result)
async def _handle_session_context_get(self, request: WsRequest, key: str) -> Response: async def _handle_session_context_get(self, request: WsRequest, key: str) -> Response:
if not self.check_api_token(request): if not self.check_api_token(request):
return _http_error(401, "Unauthorized") return _http_error(401, "Unauthorized")
@@ -746,6 +823,10 @@ class GatewayHTTPHandler:
for k, v in s.items() for k, v in s.items()
if k != "path" and k not in WEBUI_SESSION_INDEX_INTERNAL_FIELDS if k != "path" and k not in WEBUI_SESSION_INDEX_INTERNAL_FIELDS
} }
# Keep the additive recovery field absent for ordinary sessions so
# older clients and compact list responses stay unchanged.
if row.get("recovery_state") is None:
row.pop("recovery_state", None)
chat_id = key.split(":", 1)[1] chat_id = key.split(":", 1)[1]
started_at = websocket_turn_wall_started_at(chat_id) started_at = websocket_turn_wall_started_at(chat_id)
if started_at is not None: if started_at is not None:
@@ -918,13 +999,24 @@ class GatewayHTTPHandler:
# -- Automation routes -------------------------------------------------- # -- Automation routes --------------------------------------------------
async def _run_cron_transaction(
self,
operation: Callable[..., Any],
/,
*args: Any,
**kwargs: Any,
) -> Any:
if self.cron_service is not None:
return await self.cron_service.run_sync(operation, *args, **kwargs)
return await shield_and_drain(asyncio.to_thread(operation, *args, **kwargs))
async def _dispatch_automation_routes( async def _dispatch_automation_routes(
self, self,
request: WsRequest, request: WsRequest,
got: str, got: str,
) -> Response | None: ) -> Response | None:
if got == "/api/webui/automations": if got == "/api/webui/automations":
return self._handle_webui_automations(request) return await self._run_cron_transaction(self._handle_webui_automations, request)
m = re.match(r"^/api/webui/automations/(enable|disable|delete|run|update)$", got) m = re.match(r"^/api/webui/automations/(enable|disable|delete|run|update)$", got)
if m: if m:
return await self._handle_webui_automation_action(request, m.group(1)) return await self._handle_webui_automation_action(request, m.group(1))
@@ -990,13 +1082,24 @@ class GatewayHTTPHandler:
job_id = (_query_first(query, "id") or _query_first(query, "job_id") or "").strip() job_id = (_query_first(query, "id") or _query_first(query, "job_id") or "").strip()
if not job_id: if not job_id:
return _http_error(400, "missing automation id") return _http_error(400, "missing automation id")
trigger = self.local_trigger_store.get(job_id) if self.local_trigger_store else None trigger = (
await asyncio.to_thread(self.local_trigger_store.get, job_id)
if self.local_trigger_store
else None
)
if trigger is not None: if trigger is not None:
return self._handle_local_trigger_action(request, action, trigger) return await shield_and_drain(
asyncio.to_thread(
self._handle_local_trigger_action,
request,
action,
trigger,
)
)
if self.cron_service is None: if self.cron_service is None:
return _http_error(404, "automation not found") return _http_error(404, "automation not found")
job = self.cron_service.get_job(job_id) job = await self.cron_service.run_sync(self.cron_service.get_job, job_id)
if job is None: if job is None:
return _http_error(404, "automation not found") return _http_error(404, "automation not found")
if job.payload.kind == "system_event": if job.payload.kind == "system_event":
@@ -1005,13 +1108,23 @@ class GatewayHTTPHandler:
return _http_error(409, "automation has no linked chat") return _http_error(409, "automation has no linked chat")
if action == "enable": if action == "enable":
if self.cron_service.enable_job(job_id, enabled=True) is None: result = await self.cron_service.run_sync(
self.cron_service.enable_job,
job_id,
enabled=True,
)
if result is None:
return _http_error(404, "automation not found") return _http_error(404, "automation not found")
elif action == "disable": elif action == "disable":
if self.cron_service.enable_job(job_id, enabled=False) is None: result = await self.cron_service.run_sync(
self.cron_service.enable_job,
job_id,
enabled=False,
)
if result is None:
return _http_error(404, "automation not found") return _http_error(404, "automation not found")
elif action == "delete": elif action == "delete":
result = self.cron_service.remove_job(job_id) result = await self.cron_service.run_sync(self.cron_service.remove_job, job_id)
if result == "not_found": if result == "not_found":
return _http_error(404, "automation not found") return _http_error(404, "automation not found")
if result == "protected": if result == "protected":
@@ -1029,7 +1142,11 @@ class GatewayHTTPHandler:
if isinstance(parsed, str): if isinstance(parsed, str):
return _http_error(400, parsed) return _http_error(400, parsed)
try: try:
result = self.cron_service.update_job(job_id, **parsed) result = await self.cron_service.run_sync(
self.cron_service.update_job,
job_id,
**parsed,
)
except ValueError as exc: except ValueError as exc:
return _http_error(400, str(exc)) return _http_error(400, str(exc))
if result == "not_found": if result == "not_found":
@@ -1039,7 +1156,7 @@ class GatewayHTTPHandler:
else: else:
return _http_error(404, "unknown automation action") return _http_error(404, "unknown automation action")
return self._handle_webui_automations(request) return await self._run_cron_transaction(self._handle_webui_automations, request)
def _handle_local_trigger_action( def _handle_local_trigger_action(
self, self,
@@ -1124,9 +1241,17 @@ class GatewayHTTPHandler:
if got == "/api/webui/skills/install": if got == "/api/webui/skills/install":
return await self._handle_webui_skill_install(connection, request) return await self._handle_webui_skill_install(connection, request)
if got == "/api/webui/skills/update": if got == "/api/webui/skills/update":
return self._handle_webui_skill_update(request) return await shield_and_drain(
asyncio.to_thread(self._handle_webui_skill_update, request)
)
if got == "/api/webui/skills/delete": if got == "/api/webui/skills/delete":
return self._handle_webui_skill_delete(connection, request) return await shield_and_drain(
asyncio.to_thread(
self._handle_webui_skill_delete,
connection,
request,
)
)
if got == "/api/webui/skills": if got == "/api/webui/skills":
return self._handle_webui_skills(request) return self._handle_webui_skills(request)
m = re.match(r"^/api/webui/skills/([^/]+)$", got) m = re.match(r"^/api/webui/skills/([^/]+)$", got)
@@ -1237,7 +1362,7 @@ class GatewayHTTPHandler:
) -> Response: ) -> Response:
if not self.check_api_token(request): if not self.check_api_token(request):
return _http_error(401, "Unauthorized") return _http_error(401, "Unauthorized")
if not self._allow_webui_package_install(connection, request): if not await self._allow_webui_package_install(connection, request):
return _http_error(403, "remote skill installation is disabled") return _http_error(403, "remote skill installation is disabled")
if self._skill_install_lock.locked(): if self._skill_install_lock.locked():
return _http_error(409, "another skill installation is already in progress") return _http_error(409, "another skill installation is already in progress")
@@ -1269,13 +1394,16 @@ class GatewayHTTPHandler:
"last_action": action, "last_action": action,
}) })
def _allow_webui_package_install(self, connection: Any, request: WsRequest) -> bool: async def _allow_webui_package_install(
self,
connection: Any,
request: WsRequest,
) -> bool:
if _is_local_browser_request(connection, request.headers): if _is_local_browser_request(connection, request.headers):
return True return True
try: try:
return bool( config = await self.settings.config.load_async()
self.settings.config.load().tools.webui_allow_remote_package_install return bool(config.tools.webui_allow_remote_package_install)
)
except Exception: except Exception:
self._log.exception("failed to load remote package install policy") self._log.exception("failed to load remote package install policy")
return False return False
+1 -1
View File
@@ -128,7 +128,7 @@ async def test_pending_document_attachment_keeps_body_out_of_prompt(
nonlocal call_count nonlocal call_count
call_count += 1 call_count += 1
captured_messages.append([dict(message) for message in messages]) captured_messages.append([dict(message) for message in messages])
return LLMResponse(content=f"answer-{call_count}", tool_calls=[], usage={}) return LLMResponse(content=f"answer-{call_count}", tool_calls=[], usage=None)
loop = _make_loop(workspace) loop = _make_loop(workspace)
loop.provider.chat_with_retry = chat_with_retry loop.provider.chat_with_retry = chat_with_retry
+23
View File
@@ -49,6 +49,29 @@ def test_system_prompt_stays_stable_when_clock_changes(tmp_path, monkeypatch) ->
assert prompt1 == prompt2 assert prompt1 == prompt2
def test_selected_project_path_follows_shared_cache_prefix(tmp_path) -> None:
"""Project paths must not invalidate the stable identity and tool contract prefix."""
agent_home = tmp_path / "agent-home"
project_a = tmp_path / "project-a"
project_b = tmp_path / "project-b"
agent_home.mkdir()
project_a.mkdir()
project_b.mkdir()
builder = ContextBuilder(agent_home)
prompt_a = builder.build_system_prompt(workspace=project_a)
prompt_b = builder.build_system_prompt(workspace=project_b)
marker = "# Current Project"
prefix_a = prompt_a[: prompt_a.index(marker)]
prefix_b = prompt_b[: prompt_b.index(marker)]
assert prefix_a == prefix_b
assert "# Tool Usage Notes" in prefix_a
assert str(project_a.resolve()) not in prefix_a
assert str(project_b.resolve()) not in prefix_b
assert prompt_a == builder.build_system_prompt(workspace=project_a)
def test_system_prompt_reflects_current_dream_memory_contract(tmp_path) -> None: def test_system_prompt_reflects_current_dream_memory_contract(tmp_path) -> None:
workspace = _make_workspace(tmp_path) workspace = _make_workspace(tmp_path)
builder = ContextBuilder(workspace) builder = ContextBuilder(workspace)
+4 -4
View File
@@ -412,7 +412,7 @@ class TestEphemeralDirect:
provider.supports_tools = True provider.supports_tools = True
provider.generation = MagicMock(max_tokens=4096) provider.generation = MagicMock(max_tokens=4096)
provider.chat_with_retry = AsyncMock( provider.chat_with_retry = AsyncMock(
return_value=LLMResponse(content="done", tool_calls=[], finish_reason="stop", usage={}) return_value=LLMResponse(content="done", tool_calls=[], finish_reason="stop", usage=None)
) )
with ( with (
@@ -556,9 +556,9 @@ class TestEphemeralDirect:
"new_text": "replacement", "new_text": "replacement",
}, },
)], )],
usage={}, usage=None,
), ),
LLMResponse(content="done", finish_reason="stop", tool_calls=[], usage={}), LLMResponse(content="done", finish_reason="stop", tool_calls=[], usage=None),
]) ])
resp = await loop.process_direct( resp = await loop.process_direct(
@@ -646,7 +646,7 @@ class TestEphemeralHooks:
provider.generation = MagicMock(max_tokens=4096) provider.generation = MagicMock(max_tokens=4096)
provider.chat_with_retry = AsyncMock( provider.chat_with_retry = AsyncMock(
return_value=LLMResponse( return_value=LLMResponse(
content="done", finish_reason="stop", tool_calls=[], usage={}, content="done", finish_reason="stop", tool_calls=[], usage=None,
) )
) )
+1 -1
View File
@@ -12,7 +12,7 @@ from nanobot.utils.evaluator import (
class DummyProvider(LLMProvider): class DummyProvider(LLMProvider):
def __init__(self, responses: list[LLMResponse]): def __init__(self, responses: list[LLMResponse]):
super().__init__() super().__init__(provider_name="dummy")
self._responses = list(responses) self._responses = list(responses)
async def chat(self, *args, **kwargs) -> LLMResponse: async def chat(self, *args, **kwargs) -> LLMResponse:
+2 -2
View File
@@ -69,7 +69,7 @@ def test_explicit_message_limit_still_starts_at_user_turn() -> None:
async def test_process_message_replays_with_token_budget_only(tmp_path: Path) -> None: async def test_process_message_replays_with_token_budget_only(tmp_path: Path) -> None:
loop = _make_loop(tmp_path, context_window_tokens=32_768) loop = _make_loop(tmp_path, context_window_tokens=32_768)
loop.provider.chat_with_retry = AsyncMock( loop.provider.chat_with_retry = AsyncMock(
return_value=LLMResponse(content="ok", tool_calls=[], usage={}) return_value=LLMResponse(content="ok", tool_calls=[], usage=None)
) )
loop.tools.get_definitions = MagicMock(return_value=[]) loop.tools.get_definitions = MagicMock(return_value=[])
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign] loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
@@ -91,7 +91,7 @@ async def test_process_message_replays_with_token_budget_only(tmp_path: Path) ->
async def test_token_budget_keeps_current_user_as_replay_boundary(tmp_path: Path) -> None: async def test_token_budget_keeps_current_user_as_replay_boundary(tmp_path: Path) -> None:
loop = _make_loop(tmp_path, context_window_tokens=8_000) loop = _make_loop(tmp_path, context_window_tokens=8_000)
loop.provider.chat_with_retry = AsyncMock( loop.provider.chat_with_retry = AsyncMock(
return_value=LLMResponse(content="ok", tool_calls=[], usage={}) return_value=LLMResponse(content="ok", tool_calls=[], usage=None)
) )
loop.tools.get_definitions = MagicMock(return_value=[]) loop.tools.get_definitions = MagicMock(return_value=[])
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign] loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
+4 -4
View File
@@ -453,7 +453,7 @@ async def test_agent_loop_extra_hook_receives_calls(tmp_path):
loop = _make_loop(tmp_path, hooks=[TrackingHook()]) loop = _make_loop(tmp_path, hooks=[TrackingHook()])
loop.provider.chat_with_retry = AsyncMock( loop.provider.chat_with_retry = AsyncMock(
return_value=LLMResponse(content="done", tool_calls=[], usage={}) return_value=LLMResponse(content="done", tool_calls=[], usage=None)
) )
loop.tools.get_definitions = MagicMock(return_value=[]) loop.tools.get_definitions = MagicMock(return_value=[])
@@ -494,7 +494,7 @@ async def test_agent_loop_turn_hook_factories_receive_context(tmp_path):
loop = _make_loop(tmp_path, hook_factories=[factory("registered")]) loop = _make_loop(tmp_path, hook_factories=[factory("registered")])
loop.provider.chat_with_retry = AsyncMock( loop.provider.chat_with_retry = AsyncMock(
return_value=LLMResponse(content="done", tool_calls=[], usage={}) return_value=LLMResponse(content="done", tool_calls=[], usage=None)
) )
loop.tools.get_definitions = MagicMock(return_value=[]) loop.tools.get_definitions = MagicMock(return_value=[])
@@ -541,7 +541,7 @@ async def test_agent_loop_extra_hook_error_isolation(tmp_path):
loop = _make_loop(tmp_path, hooks=[BadHook()]) loop = _make_loop(tmp_path, hooks=[BadHook()])
loop.provider.chat_with_retry = AsyncMock( loop.provider.chat_with_retry = AsyncMock(
return_value=LLMResponse(content="still works", tool_calls=[], usage={}) return_value=LLMResponse(content="still works", tool_calls=[], usage=None)
) )
loop.tools.get_definitions = MagicMock(return_value=[]) loop.tools.get_definitions = MagicMock(return_value=[])
@@ -562,7 +562,7 @@ async def test_agent_loop_extra_hooks_do_not_swallow_loop_hook_errors(tmp_path):
loop.provider.chat_with_retry = AsyncMock(return_value=LLMResponse( loop.provider.chat_with_retry = AsyncMock(return_value=LLMResponse(
content="working", content="working",
tool_calls=[ToolCallRequest(id="c1", name="list_dir", arguments={"path": "."})], tool_calls=[ToolCallRequest(id="c1", name="list_dir", arguments={"path": "."})],
usage={}, usage=None,
)) ))
loop.tools.get_definitions = MagicMock(return_value=[]) loop.tools.get_definitions = MagicMock(return_value=[])
loop.tools.execute = AsyncMock(return_value="ok") loop.tools.execute = AsyncMock(return_value="ok")
+2 -2
View File
@@ -393,9 +393,9 @@ class TestToolEventProgress:
}, },
) )
], ],
usage={}, usage=None,
) )
return LLMResponse(content="Done", tool_calls=[], usage={}) return LLMResponse(content="Done", tool_calls=[], usage=None)
provider.chat_stream_with_retry = chat_stream_with_retry provider.chat_stream_with_retry = chat_stream_with_retry
provider.chat_with_retry = AsyncMock() provider.chat_with_retry = AsyncMock()
+22 -22
View File
@@ -48,7 +48,7 @@ async def test_ephemeral_runner_enters_and_restores_turn_scopes(tmp_path):
async def chat_with_retry(**_kwargs): async def chat_with_retry(**_kwargs):
assert goal_mutation_allowed() is True assert goal_mutation_allowed() is True
return LLMResponse(content="done", tool_calls=[], usage={}) return LLMResponse(content="done", tool_calls=[], usage=None)
loop.provider.chat_with_retry = AsyncMock(side_effect=chat_with_retry) loop.provider.chat_with_retry = AsyncMock(side_effect=chat_with_retry)
loop.tools.get_definitions = MagicMock(return_value=[]) loop.tools.get_definitions = MagicMock(return_value=[])
@@ -83,7 +83,7 @@ async def test_goal_command_can_implement_plan_from_prior_discussion(tmp_path):
}, },
) )
], ],
usage={}, usage=None,
), ),
LLMResponse( LLMResponse(
content="closing goal", content="closing goal",
@@ -94,7 +94,7 @@ async def test_goal_command_can_implement_plan_from_prior_discussion(tmp_path):
arguments={"action": "complete", "recap": "Implemented and tested."}, arguments={"action": "complete", "recap": "Implemented and tested."},
) )
], ],
usage={}, usage=None,
), ),
LLMResponse( LLMResponse(
content="trying to start another goal", content="trying to start another goal",
@@ -105,9 +105,9 @@ async def test_goal_command_can_implement_plan_from_prior_discussion(tmp_path):
arguments={"objective": "Start an unrelated follow-up."}, arguments={"objective": "Start an unrelated follow-up."},
) )
], ],
usage={}, usage=None,
), ),
LLMResponse(content="done", tool_calls=[], usage={}), LLMResponse(content="done", tool_calls=[], usage=None),
]) ])
loop = AgentLoop(bus=MessageBus(), provider=provider, workspace=tmp_path, model="test-model") loop = AgentLoop(bus=MessageBus(), provider=provider, workspace=tmp_path, model="test-model")
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=None) loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=None)
@@ -160,8 +160,8 @@ async def test_runtime_context_is_persisted_as_next_turn_prompt_prefix(tmp_path)
provider.get_default_model.return_value = "test-model" provider.get_default_model.return_value = "test-model"
provider.generation = GenerationSettings() provider.generation = GenerationSettings()
provider.chat_with_retry = AsyncMock(side_effect=[ provider.chat_with_retry = AsyncMock(side_effect=[
LLMResponse(content="first answer", usage={}), LLMResponse(content="first answer", usage=None),
LLMResponse(content="second answer", usage={}), LLMResponse(content="second answer", usage=None),
]) ])
loop = AgentLoop(bus=MessageBus(), provider=provider, workspace=tmp_path, model="test-model") loop = AgentLoop(bus=MessageBus(), provider=provider, workspace=tmp_path, model="test-model")
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=None) loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=None)
@@ -216,7 +216,7 @@ async def test_webui_quote_reaches_model_without_leaking_into_public_history(tmp
provider = MagicMock() provider = MagicMock()
provider.get_default_model.return_value = "test-model" provider.get_default_model.return_value = "test-model"
provider.generation = GenerationSettings() provider.generation = GenerationSettings()
provider.chat_with_retry = AsyncMock(return_value=LLMResponse(content="answer", usage={})) provider.chat_with_retry = AsyncMock(return_value=LLMResponse(content="answer", usage=None))
loop = AgentLoop(bus=MessageBus(), provider=provider, workspace=tmp_path, model="test-model") loop = AgentLoop(bus=MessageBus(), provider=provider, workspace=tmp_path, model="test-model")
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=None) loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=None)
session = loop.sessions.get_or_create("websocket:chat") session = loop.sessions.get_or_create("websocket:chat")
@@ -258,9 +258,9 @@ async def test_runtime_context_provider_runs_once_across_tool_iterations(tmp_pat
name="read_file", name="read_file",
arguments={"path": "note.txt"}, arguments={"path": "note.txt"},
)], )],
usage={}, usage=None,
), ),
LLMResponse(content="done", usage={}), LLMResponse(content="done", usage=None),
]) ])
loop = AgentLoop(bus=MessageBus(), provider=provider, workspace=tmp_path, model="test-model") loop = AgentLoop(bus=MessageBus(), provider=provider, workspace=tmp_path, model="test-model")
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=None) loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=None)
@@ -303,9 +303,9 @@ async def test_non_goal_direct_turn_cannot_reuse_prior_goal_command(tmp_path):
arguments={"objective": "Unauthorized persistent objective."}, arguments={"objective": "Unauthorized persistent objective."},
) )
], ],
usage={}, usage=None,
), ),
LLMResponse(content="handled as a one-time task", tool_calls=[], usage={}), LLMResponse(content="handled as a one-time task", tool_calls=[], usage=None),
]) ])
loop = AgentLoop(bus=MessageBus(), provider=provider, workspace=tmp_path, model="test-model") loop = AgentLoop(bus=MessageBus(), provider=provider, workspace=tmp_path, model="test-model")
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=None) loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=None)
@@ -383,7 +383,7 @@ async def test_loop_stream_filter_handles_think_only_prefix_without_crashing(tmp
async def chat_stream_with_retry(*, on_content_delta, **kwargs): async def chat_stream_with_retry(*, on_content_delta, **kwargs):
await on_content_delta("<think>hidden") await on_content_delta("<think>hidden")
await on_content_delta("</think>Hello") await on_content_delta("</think>Hello")
return LLMResponse(content="<think>hidden</think>Hello", tool_calls=[], usage={}) return LLMResponse(content="<think>hidden</think>Hello", tool_calls=[], usage=None)
loop.provider.chat_stream_with_retry = chat_stream_with_retry loop.provider.chat_stream_with_retry = chat_stream_with_retry
@@ -413,7 +413,7 @@ async def test_loop_stream_filter_hides_partial_trailing_think_prefix(tmp_path):
async def chat_stream_with_retry(*, on_content_delta, **kwargs): async def chat_stream_with_retry(*, on_content_delta, **kwargs):
await on_content_delta("Hello <thin") await on_content_delta("Hello <thin")
await on_content_delta("k>hidden</think>World") await on_content_delta("k>hidden</think>World")
return LLMResponse(content="Hello <think>hidden</think>World", tool_calls=[], usage={}) return LLMResponse(content="Hello <think>hidden</think>World", tool_calls=[], usage=None)
loop.provider.chat_stream_with_retry = chat_stream_with_retry loop.provider.chat_stream_with_retry = chat_stream_with_retry
@@ -436,7 +436,7 @@ async def test_loop_stream_filter_hides_complete_trailing_think_tag(tmp_path):
async def chat_stream_with_retry(*, on_content_delta, **kwargs): async def chat_stream_with_retry(*, on_content_delta, **kwargs):
await on_content_delta("Hello <think>") await on_content_delta("Hello <think>")
await on_content_delta("hidden</think>World") await on_content_delta("hidden</think>World")
return LLMResponse(content="Hello <think>hidden</think>World", tool_calls=[], usage={}) return LLMResponse(content="Hello <think>hidden</think>World", tool_calls=[], usage=None)
loop.provider.chat_stream_with_retry = chat_stream_with_retry loop.provider.chat_stream_with_retry = chat_stream_with_retry
@@ -459,8 +459,8 @@ async def test_loop_retries_think_only_final_response(tmp_path):
async def chat_with_retry(**kwargs): async def chat_with_retry(**kwargs):
call_count["n"] += 1 call_count["n"] += 1
if call_count["n"] == 1: if call_count["n"] == 1:
return LLMResponse(content="<think>hidden</think>", tool_calls=[], usage={}) return LLMResponse(content="<think>hidden</think>", tool_calls=[], usage=None)
return LLMResponse(content="Recovered answer", tool_calls=[], usage={}) return LLMResponse(content="Recovered answer", tool_calls=[], usage=None)
loop.provider.chat_with_retry = chat_with_retry loop.provider.chat_with_retry = chat_with_retry
@@ -485,7 +485,7 @@ async def test_streamed_flag_not_set_on_llm_error(tmp_path):
provider.get_default_model.return_value = "test-model" provider.get_default_model.return_value = "test-model"
loop = AgentLoop(bus=bus, provider=provider, workspace=tmp_path, model="test-model") loop = AgentLoop(bus=bus, provider=provider, workspace=tmp_path, model="test-model")
error_resp = LLMResponse( error_resp = LLMResponse(
content="503 service unavailable", finish_reason="error", tool_calls=[], usage={}, content="503 service unavailable", finish_reason="error", tool_calls=[], usage=None,
) )
loop.provider.chat_with_retry = AsyncMock(return_value=error_resp) loop.provider.chat_with_retry = AsyncMock(return_value=error_resp)
loop.provider.chat_stream_with_retry = AsyncMock(return_value=error_resp) loop.provider.chat_stream_with_retry = AsyncMock(return_value=error_resp)
@@ -523,14 +523,14 @@ async def test_ssrf_soft_block_can_finalize_after_streamed_tool_call(tmp_path):
name="exec", name="exec",
arguments={"command": "curl http://169.254.169.254/latest/meta-data/"}, arguments={"command": "curl http://169.254.169.254/latest/meta-data/"},
)], )],
usage={}, usage=None,
) )
responses = iter([ responses = iter([
tool_call_resp, tool_call_resp,
LLMResponse( LLMResponse(
content="I cannot access private URLs. Please share the local file.", content="I cannot access private URLs. Please share the local file.",
tool_calls=[], tool_calls=[],
usage={}, usage=None,
), ),
]) ])
@@ -569,8 +569,8 @@ async def test_next_turn_after_llm_error_keeps_turn_boundary(tmp_path):
provider = MagicMock() provider = MagicMock()
provider.get_default_model.return_value = "test-model" provider.get_default_model.return_value = "test-model"
provider.chat_with_retry = AsyncMock(side_effect=[ provider.chat_with_retry = AsyncMock(side_effect=[
LLMResponse(content="429 rate limit exceeded", finish_reason="error", tool_calls=[], usage={}), LLMResponse(content="429 rate limit exceeded", finish_reason="error", tool_calls=[], usage=None),
LLMResponse(content="Recovered answer", tool_calls=[], usage={}), LLMResponse(content="Recovered answer", tool_calls=[], usage=None),
]) ])
loop = AgentLoop(bus=MessageBus(), provider=provider, workspace=tmp_path, model="test-model") loop = AgentLoop(bus=MessageBus(), provider=provider, workspace=tmp_path, model="test-model")
+63 -6
View File
@@ -20,7 +20,7 @@ from nanobot.bus.outbound_events import (
) )
from nanobot.bus.queue import MessageBus from nanobot.bus.queue import MessageBus
from nanobot.cron.session_turns import CRON_HISTORY_META, CRON_TRIGGER_META from nanobot.cron.session_turns import CRON_HISTORY_META, CRON_TRIGGER_META
from nanobot.providers.base import LLMProvider, LLMResponse, ProviderConversationState from nanobot.providers.base import LLMProvider, LLMResponse, LLMUsage, ProviderConversationState
from nanobot.providers.factory import ProviderSnapshot from nanobot.providers.factory import ProviderSnapshot
from nanobot.runtime_context import ( from nanobot.runtime_context import (
RUNTIME_CONTEXT_HISTORY_META, RUNTIME_CONTEXT_HISTORY_META,
@@ -36,6 +36,7 @@ from nanobot.session.keys import (
UNIFIED_SESSION_KEY, UNIFIED_SESSION_KEY,
) )
from nanobot.session.manager import Session from nanobot.session.manager import Session
from nanobot.session.recovery import PENDING_FOLLOWUP_ID_KEY, PENDING_FOLLOWUPS_KEY
from nanobot.session.turn_continuation import ( from nanobot.session.turn_continuation import (
INTERNAL_CONTINUATION_META, INTERNAL_CONTINUATION_META,
INTERNAL_CONTINUATION_RUN_STARTED_AT_META, INTERNAL_CONTINUATION_RUN_STARTED_AT_META,
@@ -161,6 +162,35 @@ def test_persist_cron_turn_uses_distinct_history_marker(tmp_path: Path) -> None:
assert message["cron_prompt_ref"] == prompt_ref assert message["cron_prompt_ref"] == prompt_ref
def test_persist_user_message_acknowledges_durable_followup(tmp_path: Path) -> None:
loop = _make_full_loop(tmp_path)
session = loop.sessions.get_or_create("websocket:chat")
session.metadata[PENDING_FOLLOWUPS_KEY] = [
{
"id": "followup-1",
"sender_id": "user",
"chat_id": "chat",
"content": "queued while busy",
"media": [],
"metadata": {},
}
]
persisted = loop._persist_user_message_early(
InboundMessage(
channel="websocket",
sender_id="user",
chat_id="chat",
content="queued while busy",
metadata={PENDING_FOLLOWUP_ID_KEY: "followup-1"},
),
session,
)
assert persisted is True
assert PENDING_FOLLOWUPS_KEY not in session.metadata
def test_persist_local_trigger_turn_uses_hidden_automation_marker(tmp_path: Path) -> None: def test_persist_local_trigger_turn_uses_hidden_automation_marker(tmp_path: Path) -> None:
loop = _make_full_loop(tmp_path) loop = _make_full_loop(tmp_path)
session = loop.sessions.get_or_create("websocket:auto") session = loop.sessions.get_or_create("websocket:auto")
@@ -381,6 +411,34 @@ def test_save_turn_keeps_multimodal_runtime_context_for_model_replay() -> None:
assert public_history_message(session.messages[0])["content"] == [] assert public_history_message(session.messages[0])["content"] == []
def test_save_turn_acknowledges_every_merged_recovery_followup() -> None:
"""Persisting a merged injected row retires every durable follow-up ID."""
loop = _mk_loop()
session = Session(
key="test:recovery-followups",
metadata={
PENDING_FOLLOWUPS_KEY: [
{"id": "first"},
{"id": "second"},
]
},
)
loop._save_turn(
session,
[
{
"role": "user",
"content": "first\n\nsecond",
PENDING_FOLLOWUP_ID_KEY: ["first", "second"],
}
],
skip=0,
)
assert PENDING_FOLLOWUPS_KEY not in session.metadata
def test_save_turn_keeps_image_placeholder_and_runtime_context() -> None: def test_save_turn_keeps_image_placeholder_and_runtime_context() -> None:
loop = _mk_loop() loop = _mk_loop()
session = Session(key="test:image") session = Session(key="test:image")
@@ -1825,7 +1883,7 @@ async def test_turn_usage_is_persisted_with_the_saved_session(tmp_path: Path) ->
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign] loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
async def fake_run_agent_loop(initial_messages, **_kwargs): async def fake_run_agent_loop(initial_messages, **_kwargs):
loop._last_usage = {"prompt_tokens": 64, "completion_tokens": 9} loop._last_usage = LLMUsage.reported(input_tokens=64, output_tokens=9)
return ( return (
"done", "done",
[], [],
@@ -1840,10 +1898,9 @@ async def test_turn_usage_is_persisted_with_the_saved_session(tmp_path: Path) ->
) )
loop.sessions.invalidate("cli:usage") loop.sessions.invalidate("cli:usage")
assert loop.sessions.get_or_create("cli:usage").metadata["_last_usage"] == { assert loop.sessions.get_or_create("cli:usage").metadata["_last_usage"] == (
"prompt_tokens": 64, LLMUsage.reported(input_tokens=64, output_tokens=9).to_dict()
"completion_tokens": 9, )
}
@pytest.mark.asyncio @pytest.mark.asyncio
+1 -1
View File
@@ -30,7 +30,7 @@ def _loop(tmp_path, responses: list[str], **kwargs) -> AgentLoop:
provider.get_default_model.return_value = "test-model" provider.get_default_model.return_value = "test-model"
provider.generation = GenerationSettings() provider.generation = GenerationSettings()
provider.chat_with_retry = AsyncMock( provider.chat_with_retry = AsyncMock(
side_effect=[LLMResponse(content=response, usage={}) for response in responses] side_effect=[LLMResponse(content=response, usage=None) for response in responses]
) )
return AgentLoop( return AgentLoop(
bus=MessageBus(), bus=MessageBus(),
+15 -15
View File
@@ -123,21 +123,21 @@ def allow_loopback_mcp_urls(monkeypatch: pytest.MonkeyPatch):
_resolver_lock = asyncio.Lock() _resolver_lock = asyncio.Lock()
monkeypatch.setattr(mcp_module, "PinnedDNSAsyncTransport", TestPinnedDNSAsyncTransport) monkeypatch.setattr(mcp_module, "PinnedDNSAsyncTransport", TestPinnedDNSAsyncTransport)
monkeypatch.setattr( async def allow_url(url: str, *, allow_loopback: bool = False) -> tuple[bool, str]:
mcp_module, return True, ""
"validate_url_target",
lambda url, *, allow_loopback=False: (True, ""), async def resolve_url(
) url: str,
monkeypatch.setattr( *,
mcp_module, allow_loopback: bool = False,
"resolve_url_target", trust_remote_dns: bool = False,
lambda url, *, allow_loopback=False: (True, "", ("127.0.0.1",)), timeout_s: float = 3.0,
) ) -> tuple[bool, str, tuple[str, ...]]:
monkeypatch.setattr( return True, "", ("127.0.0.1",)
security_network,
"resolve_url_target", monkeypatch.setattr(mcp_module, "async_validate_url_target", allow_url)
lambda url, *, allow_loopback=False: (True, "", ("127.0.0.1",)), monkeypatch.setattr(mcp_module, "async_resolve_url_target", resolve_url)
) monkeypatch.setattr(security_network, "async_resolve_url_target", resolve_url)
monkeypatch.setattr( monkeypatch.setattr(
mcp_module, mcp_module,
"env_proxy_applies_to_url", "env_proxy_applies_to_url",
+239 -31
View File
@@ -14,6 +14,7 @@ from nanobot.config.schema import AgentDefaults
from nanobot.providers.base import ( from nanobot.providers.base import (
LLMProvider, LLMProvider,
LLMResponse, LLMResponse,
LLMUsage,
ProviderCallContext, ProviderCallContext,
ProviderConversationState, ProviderConversationState,
ToolCallRequest, ToolCallRequest,
@@ -22,6 +23,163 @@ from nanobot.providers.base import (
_MAX_TOOL_RESULT_CHARS = AgentDefaults().max_tool_result_chars _MAX_TOOL_RESULT_CHARS = AgentDefaults().max_tool_result_chars
def _make_usage_spec(provider, tools):
return make_run_spec(
provider,
initial_messages=[{"role": "user", "content": "hello"}],
tools=tools,
model="test-model",
max_iterations=1,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
)
def test_usage_or_estimate_replaces_reported_zero_for_content(monkeypatch) -> None:
from nanobot.agent.runner import AgentRunner
provider = MagicMock(spec=LLMProvider)
tools = MagicMock()
tools.get_definitions.return_value = []
monkeypatch.setattr(
"nanobot.agent.runner.estimate_prompt_tokens_chain",
lambda provider, model, messages, definitions: (12, "test"),
)
monkeypatch.setattr("nanobot.agent.runner.estimate_message_tokens", lambda message: 7)
response = LLMResponse(
content="answer",
usage=LLMUsage.reported(input_tokens=0, output_tokens=0),
generation_ms=25,
ttft_ms=5,
)
usage = AgentRunner()._usage_or_estimate(
_make_usage_spec(provider, tools),
[{"role": "user", "content": "hello"}],
response,
)
assert usage == LLMUsage.estimated(input_tokens=12, output_tokens=7).with_timing(
generation_ms=25,
ttft_ms=5,
)
assert usage.source == "estimated"
assert usage.total_tokens == 19
def test_usage_or_estimate_counts_tool_call_output_for_reported_zero(monkeypatch) -> None:
from nanobot.agent.runner import AgentRunner
provider = MagicMock(spec=LLMProvider)
tools = MagicMock()
tools.get_definitions.return_value = []
captured_message: dict = {}
monkeypatch.setattr(
"nanobot.agent.runner.estimate_prompt_tokens_chain",
lambda provider, model, messages, definitions: (13, "test"),
)
def estimate_output(message):
captured_message.update(message)
return 9
monkeypatch.setattr("nanobot.agent.runner.estimate_message_tokens", estimate_output)
response = LLMResponse(
content=None,
tool_calls=[
ToolCallRequest(
id="call_1",
name="lookup",
arguments={"query": "nanobot"},
)
],
finish_reason="tool_calls",
usage=LLMUsage.reported(input_tokens=0, output_tokens=0),
)
usage = AgentRunner()._usage_or_estimate(
_make_usage_spec(provider, tools),
[{"role": "user", "content": "hello"}],
response,
)
assert usage == LLMUsage.estimated(input_tokens=13, output_tokens=9)
assert usage.total_tokens == 22
assert captured_message["tool_calls"][0]["function"]["name"] == "lookup"
@pytest.mark.parametrize(
"provider_usage",
[None, LLMUsage.reported(input_tokens=0, output_tokens=0)],
)
def test_usage_or_estimate_counts_error_without_estimating_tokens(
monkeypatch,
provider_usage: LLMUsage | None,
) -> None:
from nanobot.agent.runner import AgentRunner
provider = MagicMock(spec=LLMProvider)
tools = MagicMock()
estimate = MagicMock()
runner = AgentRunner()
monkeypatch.setattr(runner, "_estimate_response_usage", estimate)
response = LLMResponse(
content="upstream failed",
finish_reason="error",
usage=provider_usage,
)
usage = runner._usage_or_estimate(
_make_usage_spec(provider, tools),
[{"role": "user", "content": "hello"}],
response,
)
assert usage is not None
assert usage.total_tokens == 0
assert usage.request_count == 1
assert usage.context_tokens is None
aggregate = LLMUsage.reported(input_tokens=12, output_tokens=3) + usage
assert aggregate.context_tokens == 12
assert aggregate.request_count == 2
estimate.assert_not_called()
def test_usage_or_estimate_trusts_positive_reported_total(monkeypatch) -> None:
from nanobot.agent.runner import AgentRunner
provider = MagicMock(spec=LLMProvider)
tools = MagicMock()
estimate = MagicMock()
runner = AgentRunner()
monkeypatch.setattr(runner, "_estimate_response_usage", estimate)
response = LLMResponse(
content="answer",
usage=LLMUsage.reported(
input_tokens=15,
output_tokens=18,
total_tokens=175,
),
generation_ms=30,
ttft_ms=6,
)
usage = runner._usage_or_estimate(
_make_usage_spec(provider, tools),
[{"role": "user", "content": "hello"}],
response,
)
assert usage is not None
assert usage.source == "reported"
assert usage.input_tokens == 15
assert usage.output_tokens == 18
assert usage.total_tokens == 175
assert usage.reported_tokens == 175
assert usage.generation_ms == 30
assert usage.ttft_ms == 6
estimate.assert_not_called()
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_runner_preserves_reasoning_fields_and_tool_results(): async def test_runner_preserves_reasoning_fields_and_tool_results():
from nanobot.agent.runner import AgentRunner from nanobot.agent.runner import AgentRunner
@@ -38,10 +196,10 @@ async def test_runner_preserves_reasoning_fields_and_tool_results():
tool_calls=[ToolCallRequest(id="call_1", name="list_dir", arguments={"path": "."})], tool_calls=[ToolCallRequest(id="call_1", name="list_dir", arguments={"path": "."})],
reasoning_content="hidden reasoning", reasoning_content="hidden reasoning",
thinking_blocks=[{"type": "thinking", "thinking": "step"}], thinking_blocks=[{"type": "thinking", "thinking": "step"}],
usage={"prompt_tokens": 5, "completion_tokens": 3}, usage=LLMUsage.reported(input_tokens=5, output_tokens=3),
) )
captured_second_call[:] = messages captured_second_call[:] = messages
return LLMResponse(content="done", tool_calls=[], usage={}) return LLMResponse(content="done", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry provider.chat_with_retry = chat_with_retry
tools = MagicMock() tools = MagicMock()
@@ -441,7 +599,7 @@ async def test_runner_uses_no_tools_finalization_after_max_iterations():
return LLMResponse( return LLMResponse(
content="Read the directory twice. More investigation remains.", content="Read the directory twice. More investigation remains.",
tool_calls=[], tool_calls=[],
usage={"prompt_tokens": 10, "completion_tokens": 7}, usage=LLMUsage.reported(input_tokens=10, output_tokens=7),
) )
provider.chat_with_retry = chat_with_retry provider.chat_with_retry = chat_with_retry
@@ -499,6 +657,55 @@ async def test_runner_times_out_hung_llm_request():
assert "timed out" in (result.final_content or "").lower() assert "timed out" in (result.final_content or "").lower()
@pytest.mark.asyncio
async def test_runner_times_out_hung_max_iteration_finalization():
from nanobot.agent.runner import AgentRunner
provider = MagicMock()
calls = 0
async def chat_with_retry(**kwargs):
nonlocal calls
calls += 1
if calls == 1:
return LLMResponse(
content="",
tool_calls=[
ToolCallRequest(
id="call_1",
name="probe",
arguments={},
)
],
finish_reason="tool_calls",
)
await asyncio.Event().wait()
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
tools.get_definitions.return_value = []
tools.execute = AsyncMock(return_value="ok")
result = await asyncio.wait_for(
AgentRunner().run(make_run_spec(
provider,
initial_messages=[{"role": "user", "content": "run the probe"}],
tools=tools,
model="test-model",
max_iterations=1,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
max_iterations_message="fallback after {max_iterations} iteration",
llm_timeout_s=0.01,
)),
timeout=1.0,
)
assert calls == 2
assert result.stop_reason == "max_iterations"
assert result.error is None
assert result.final_content == "fallback after 1 iteration"
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_runner_applies_outer_wall_timeout_to_streaming_requests(): async def test_runner_applies_outer_wall_timeout_to_streaming_requests():
from nanobot.agent.hook import AgentHook, AgentHookContext from nanobot.agent.hook import AgentHook, AgentHookContext
@@ -664,10 +871,10 @@ async def test_runner_replaces_empty_tool_result_with_marker():
return LLMResponse( return LLMResponse(
content="working", content="working",
tool_calls=[ToolCallRequest(id="call_1", name="noop", arguments={})], tool_calls=[ToolCallRequest(id="call_1", name="noop", arguments={})],
usage={}, usage=None,
) )
captured_second_call[:] = messages captured_second_call[:] = messages
return LLMResponse(content="done", tool_calls=[], usage={}) return LLMResponse(content="done", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry provider.chat_with_retry = chat_with_retry
tools = MagicMock() tools = MagicMock()
@@ -702,12 +909,12 @@ async def test_runner_retries_empty_final_response_with_summary_prompt():
return LLMResponse( return LLMResponse(
content=None, content=None,
tool_calls=[], tool_calls=[],
usage={"prompt_tokens": 5, "completion_tokens": 1}, usage=LLMUsage.reported(input_tokens=5, output_tokens=1),
) )
return LLMResponse( return LLMResponse(
content="final answer", content="final answer",
tool_calls=[], tool_calls=[],
usage={"prompt_tokens": 3, "completion_tokens": 7}, usage=LLMUsage.reported(input_tokens=3, output_tokens=7),
) )
provider.chat_with_retry = chat_with_retry provider.chat_with_retry = chat_with_retry
@@ -729,8 +936,9 @@ async def test_runner_retries_empty_final_response_with_summary_prompt():
assert calls[0]["tools"] is not None assert calls[0]["tools"] is not None
assert calls[1]["tools"] is not None assert calls[1]["tools"] is not None
assert calls[2]["tools"] is None assert calls[2]["tools"] is None
assert result.usage["prompt_tokens"] == 13 assert result.usage is not None
assert result.usage["completion_tokens"] == 9 assert result.usage.input_tokens == 13
assert result.usage.output_tokens == 9
@pytest.mark.asyncio @pytest.mark.asyncio
@@ -802,7 +1010,7 @@ async def test_runner_uses_specific_message_after_empty_finalization_retry():
provider = MagicMock(spec=LLMProvider) provider = MagicMock(spec=LLMProvider)
async def chat_with_retry(*, messages, **kwargs): async def chat_with_retry(*, messages, **kwargs):
return LLMResponse(content=None, tool_calls=[], usage={}) return LLMResponse(content=None, tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry provider.chat_with_retry = chat_with_retry
tools = MagicMock() tools = MagicMock()
@@ -842,14 +1050,14 @@ async def test_empty_finalization_retry_discards_candidate_provider_state():
provider = MagicMock(spec=LLMProvider) provider = MagicMock(spec=LLMProvider)
provider.can_resume_conversation_state.return_value = True provider.can_resume_conversation_state.return_value = True
provider.chat_with_retry = AsyncMock(side_effect=[ provider.chat_with_retry = AsyncMock(side_effect=[
LLMResponse(content=None, tool_calls=[], usage={}), LLMResponse(content=None, tool_calls=[], usage=None),
LLMResponse(content=None, tool_calls=[], usage={}), LLMResponse(content=None, tool_calls=[], usage=None),
LLMResponse( LLMResponse(
content="finalized without tools", content="finalized without tools",
tool_calls=[ToolCallRequest(id="call_1", name="exec", arguments={})], tool_calls=[ToolCallRequest(id="call_1", name="exec", arguments={})],
finish_reason="stop", finish_reason="stop",
provider_state=candidate, provider_state=candidate,
usage={}, usage=None,
), ),
]) ])
tools = MagicMock() tools = MagicMock()
@@ -988,20 +1196,20 @@ async def test_runner_empty_response_does_not_break_tool_chain():
return LLMResponse( return LLMResponse(
content=None, content=None,
tool_calls=[ToolCallRequest(id="tc1", name="read_file", arguments={"path": "a.txt"})], tool_calls=[ToolCallRequest(id="tc1", name="read_file", arguments={"path": "a.txt"})],
usage={"prompt_tokens": 10, "completion_tokens": 5}, usage=LLMUsage.reported(input_tokens=10, output_tokens=5),
) )
if call_count == 2: if call_count == 2:
return LLMResponse(content=None, tool_calls=[], usage={"prompt_tokens": 10, "completion_tokens": 1}) return LLMResponse(content=None, tool_calls=[], usage=LLMUsage.reported(input_tokens=10, output_tokens=1))
if call_count == 3: if call_count == 3:
return LLMResponse( return LLMResponse(
content=None, content=None,
tool_calls=[ToolCallRequest(id="tc2", name="read_file", arguments={"path": "b.txt"})], tool_calls=[ToolCallRequest(id="tc2", name="read_file", arguments={"path": "b.txt"})],
usage={"prompt_tokens": 10, "completion_tokens": 5}, usage=LLMUsage.reported(input_tokens=10, output_tokens=5),
) )
return LLMResponse( return LLMResponse(
content="Here are the results.", content="Here are the results.",
tool_calls=[], tool_calls=[],
usage={"prompt_tokens": 10, "completion_tokens": 10}, usage=LLMUsage.reported(input_tokens=10, output_tokens=10),
) )
provider.chat_with_retry = chat_with_retry provider.chat_with_retry = chat_with_retry
@@ -1030,9 +1238,8 @@ async def test_runner_empty_response_does_not_break_tool_chain():
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_runner_accumulates_usage_and_preserves_cached_tokens(): async def test_runner_accumulates_usage_and_preserves_cache_reads():
"""Runner should accumulate prompt/completion tokens across iterations """Runner accumulates usage across iterations, including cache reads."""
and preserve cached_tokens from provider responses."""
from nanobot.agent.runner import AgentRunner from nanobot.agent.runner import AgentRunner
provider = MagicMock(spec=LLMProvider) provider = MagicMock(spec=LLMProvider)
@@ -1044,12 +1251,12 @@ async def test_runner_accumulates_usage_and_preserves_cached_tokens():
return LLMResponse( return LLMResponse(
content="thinking", content="thinking",
tool_calls=[ToolCallRequest(id="call_1", name="read_file", arguments={"path": "x"})], tool_calls=[ToolCallRequest(id="call_1", name="read_file", arguments={"path": "x"})],
usage={"prompt_tokens": 100, "completion_tokens": 10, "cached_tokens": 80}, usage=LLMUsage.reported(input_tokens=100, output_tokens=10, cache_read_tokens=80),
) )
return LLMResponse( return LLMResponse(
content="done", content="done",
tool_calls=[], tool_calls=[],
usage={"prompt_tokens": 200, "completion_tokens": 20, "cached_tokens": 150}, usage=LLMUsage.reported(input_tokens=200, output_tokens=20, cache_read_tokens=150),
) )
provider.chat_with_retry = chat_with_retry provider.chat_with_retry = chat_with_retry
@@ -1067,11 +1274,12 @@ async def test_runner_accumulates_usage_and_preserves_cached_tokens():
)) ))
# Usage should be accumulated across iterations # Usage should be accumulated across iterations
assert result.usage["prompt_tokens"] == 300 # 100 + 200 assert result.usage is not None
assert result.usage["completion_tokens"] == 30 # 10 + 20 assert result.usage.input_tokens == 300 # 100 + 200
assert result.usage["cached_tokens"] == 230 # 80 + 150 assert result.usage.output_tokens == 30 # 10 + 20
assert result.usage["context_tokens"] == 200 assert result.usage.cache_read_tokens == 230 # 80 + 150
assert result.usage["request_count"] == 2 assert result.usage.context_tokens == 200
assert result.usage.request_count == 2
@pytest.mark.asyncio @pytest.mark.asyncio
@@ -1088,7 +1296,7 @@ async def test_runner_binds_on_retry_wait_to_retry_callback_not_progress():
async def chat_with_retry(**kwargs): async def chat_with_retry(**kwargs):
captured.update(kwargs) captured.update(kwargs)
return LLMResponse(content="done", tool_calls=[], usage={}) return LLMResponse(content="done", tool_calls=[], usage=None)
provider = MagicMock(spec=LLMProvider) provider = MagicMock(spec=LLMProvider)
provider.chat_with_retry = chat_with_retry provider.chat_with_retry = chat_with_retry
@@ -1130,7 +1338,7 @@ async def test_runner_passes_temperature_to_provider():
async def chat_with_retry(**kwargs): async def chat_with_retry(**kwargs):
captured.update(kwargs) captured.update(kwargs)
return LLMResponse(content="done", tool_calls=[], usage={}) return LLMResponse(content="done", tool_calls=[], usage=None)
provider = MagicMock(spec=LLMProvider) provider = MagicMock(spec=LLMProvider)
provider.chat_with_retry = chat_with_retry provider.chat_with_retry = chat_with_retry
@@ -1159,7 +1367,7 @@ async def test_runner_passes_max_tokens_to_provider():
async def chat_with_retry(**kwargs): async def chat_with_retry(**kwargs):
captured.update(kwargs) captured.update(kwargs)
return LLMResponse(content="done", tool_calls=[], usage={}) return LLMResponse(content="done", tool_calls=[], usage=None)
provider = MagicMock(spec=LLMProvider) provider = MagicMock(spec=LLMProvider)
provider.chat_with_retry = chat_with_retry provider.chat_with_retry = chat_with_retry
@@ -1188,7 +1396,7 @@ async def test_runner_passes_reasoning_effort_to_provider():
async def chat_with_retry(**kwargs): async def chat_with_retry(**kwargs):
captured.update(kwargs) captured.update(kwargs)
return LLMResponse(content="done", tool_calls=[], usage={}) return LLMResponse(content="done", tool_calls=[], usage=None)
provider = MagicMock(spec=LLMProvider) provider = MagicMock(spec=LLMProvider)
provider.chat_with_retry = chat_with_retry provider.chat_with_retry = chat_with_retry

Some files were not shown because too many files have changed in this diff Show More