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
181 changed files with 9461 additions and 2378 deletions
+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",
+96 -9
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,18 +135,38 @@ 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))
stored = session_summary_from_metadata(
session.metadata,
fallback_last_active=session.updated_at,
)
if stored is not None:
self._summaries[key] = stored
except Exception: except Exception:
logger.exception("Auto-compact: failed for {}", key) logger.exception("Auto-compact: failed for {}", key)
finally: finally:
self._archiving.discard(key) 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(
session.metadata,
fallback_last_active=session.updated_at,
)
if stored is not None:
self._summaries[key] = stored
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):
self._archiving.discard(key) self._archiving.discard(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)
+168 -39
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,
@@ -167,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."""
@@ -203,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
@@ -378,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 []
@@ -539,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
@@ -578,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,
@@ -590,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,
@@ -702,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 = [
@@ -741,10 +803,45 @@ class AgentLoop:
followup_id = msg.metadata.get(PENDING_FOLLOWUP_ID_KEY) followup_id = msg.metadata.get(PENDING_FOLLOWUP_ID_KEY)
if isinstance(followup_id, str) and followup_id: if isinstance(followup_id, str) and followup_id:
acknowledge_pending_followups(session, [followup_id]) acknowledge_pending_followups(session, [followup_id])
self.sessions.save(session)
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
@@ -842,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,
@@ -1000,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.
@@ -1202,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()
@@ -1233,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
@@ -1255,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.
@@ -1333,7 +1450,7 @@ class AgentLoop:
) )
continue continue
pending_msg = routed_msg pending_msg = routed_msg
session = self.sessions.get_or_create(effective_key) session = await self._get_or_create_session(effective_key)
followup_id = record_pending_followup(session, pending_msg) followup_id = record_pending_followup(session, pending_msg)
if followup_id is not None: if followup_id is not None:
pending_msg = dataclasses.replace( pending_msg = dataclasses.replace(
@@ -1343,7 +1460,7 @@ class AgentLoop:
PENDING_FOLLOWUP_ID_KEY: followup_id, PENDING_FOLLOWUP_ID_KEY: followup_id,
}, },
) )
self.sessions.save(session) 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:
@@ -1454,10 +1571,10 @@ class AgentLoop:
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,
@@ -1782,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
@@ -1813,16 +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 ( if (
RECOVERY_INBOUND_METADATA_KEY not in msg.metadata RECOVERY_INBOUND_METADATA_KEY not in msg.metadata
and restore_pending_interruption(session) and restore_pending_interruption(session)
): ):
self.sessions.save(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,
) )
@@ -1859,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,
@@ -1881,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(
@@ -1925,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)
@@ -1979,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,
@@ -1989,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:
@@ -2031,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)
@@ -2058,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(
@@ -2072,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,
@@ -2288,11 +2404,24 @@ 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_runtime_checkpoint(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
+37 -18
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,15 +933,16 @@ class Consolidator:
if not messages: if not messages:
return None return None
try: try:
response = await runtime.provider.chat_with_retry( with llm_usage_source("dream"):
model=runtime.model, response = await runtime.provider.chat_with_retry(
messages=request_messages, model=runtime.model,
tools=request_tools, messages=request_messages,
tool_choice="none", tools=request_tools,
temperature=runtime.generation.temperature, tool_choice="none",
max_tokens=runtime.generation.max_tokens, temperature=runtime.generation.temperature,
reasoning_effort=runtime.generation.reasoning_effort, max_tokens=runtime.generation.max_tokens,
) reasoning_effort=runtime.generation.reasoning_effort,
)
except Exception: except Exception:
logger.warning("Consolidation provider call failed, raw-dumping to history") logger.warning("Consolidation provider call failed, raw-dumping to history")
self.store.raw_archive(messages, session_key=session_key) self.store.raw_archive(messages, session_key=session_key)
@@ -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:
+123 -102
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,
@@ -76,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:
@@ -117,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)
@@ -126,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)
@@ -391,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)
@@ -412,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)
@@ -423,17 +450,20 @@ class AgentRunner:
await hook.after_run(context) await hook.after_run(context)
return result return result
finally: finally:
context.messages = deepcopy(messages) try:
if context.exception is None: context.messages = deepcopy(messages)
await hook.on_finally(context) if context.exception is None:
else:
try:
await hook.on_finally(context) await hook.on_finally(context)
except Exception: else:
logger.exception( try:
"AgentHook.on_finally error after {}", await hook.on_finally(context)
context.stop_reason or "run exception", except Exception:
) logger.exception(
"AgentHook.on_finally error after {}",
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,
@@ -443,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]] = []
@@ -519,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()
@@ -683,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)
@@ -859,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,
@@ -922,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,
@@ -1247,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(
@@ -1264,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='{}' "
@@ -1276,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,
@@ -1302,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(
@@ -1332,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:
@@ -1374,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,
@@ -1528,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)
@@ -1556,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}"
+92 -21
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,9 +950,12 @@ 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 == "":
fp.parent.mkdir(parents=True, exist_ok=True) with commit_lock:
fp.write_text(new_text, encoding="utf-8") if cancelled.is_set():
self._file_states.record_write(fp) return ToolResult.error("Error: edit_file cancelled.")
fp.parent.mkdir(parents=True, exist_ok=True)
fp.write_text(new_text, encoding="utf-8")
self._file_states.record_write(fp)
return f"Successfully created {fp}" return f"Successfully created {fp}"
return self._file_not_found_msg(path, fp) return self._file_not_found_msg(path, fp)
@@ -912,8 +973,11 @@ 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.")
fp.write_text(new_text, encoding="utf-8") with commit_lock:
self._file_states.record_write(fp) if cancelled.is_set():
return ToolResult.error("Error: edit_file cancelled.")
fp.write_text(new_text, encoding="utf-8")
self._file_states.record_write(fp)
return f"Successfully edited {fp}" return f"Successfully edited {fp}"
# Read-before-edit check # Read-before-edit check
@@ -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,8 +1068,11 @@ 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")
fp.write_bytes(new_content.encode("utf-8")) with commit_lock:
self._file_states.record_write(fp) if cancelled.is_set():
return ToolResult.error("Error: edit_file cancelled.")
fp.write_bytes(new_content.encode("utf-8"))
self._file_states.record_write(fp)
msg = f"Successfully edited {fp}" msg = f"Successfully edited {fp}"
if warning: if warning:
msg = f"{warning}\n{msg}" msg = f"{warning}\n{msg}"
+52 -22
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)
sess.metadata[GOAL_STATE_KEY] = blob saved = False
discard_legacy_goal_state_key(sess.metadata)
if reset_continuation: async def save_and_publish() -> None:
reset_goal_continuation_rounds(sess.metadata) nonlocal saved
sess.metadata[GOAL_STATE_KEY] = blob
discard_legacy_goal_state_key(sess.metadata)
if reset_continuation:
reset_goal_continuation_rounds(sess.metadata)
try:
await self._save_session(sess)
except BaseException:
sess.metadata.clear()
sess.metadata.update(previous_metadata)
raise
saved = True
await self._publish_goal_state_changed(sess.metadata)
try: try:
self._sessions.save(sess) await shield_and_drain(save_and_publish())
except BaseException: finally:
sess.metadata.clear() # This ContextVar belongs to the caller task, not the settlement task.
sess.metadata.update(previous_metadata) # Apply the post-save permission effect here even when cancellation was
raise # 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,
} }
+103 -53
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,66 +223,111 @@ class FindFilesTool(_SearchTool):
offset: int = 0, offset: int = 0,
**kwargs: Any, **kwargs: Any,
) -> str: ) -> str:
cancelled = threading.Event()
try: try:
target = self._resolve(path or ".") return await asyncio.to_thread(
if not target.exists(): self._execute_sync,
return ToolResult.error(f"Error: Path not found: {path}") path=path,
if not (target.is_dir() or target.is_file()): query=query,
return ToolResult.error(f"Error: Unsupported path: {path}") glob=glob,
file_type=type,
if sort not in {"path", "modified"}: include_dirs=include_dirs,
return ToolResult.error("Error: sort must be 'path' or 'modified'") sort=sort,
head_limit=head_limit,
limit = ( offset=offset,
_DEFAULT_FILE_HEAD_LIMIT cancelled=cancelled,
if head_limit is None
else None if head_limit == 0 else head_limit
) )
root = target if target.is_dir() else target.parent except asyncio.CancelledError:
matches: list[tuple[str, float]] = [] cancelled.set()
raise
for candidate in self._iter_paths(target, include_dirs=include_dirs):
if candidate.is_dir() and not include_dirs:
continue
rel_path = candidate.relative_to(root).as_posix()
display_path = self._display_path(candidate, root)
name = candidate.name
if glob and not _match_glob(rel_path, name, glob):
continue
if candidate.is_file() and not _matches_type(name, type):
continue
if candidate.is_dir() and type:
continue
if not _matches_query(display_path, query):
continue
try:
mtime = candidate.stat().st_mtime
except OSError:
mtime = 0.0
suffix = "/" if candidate.is_dir() else ""
matches.append((display_path + suffix, mtime))
if sort == "modified":
matches.sort(key=lambda item: (-item[1], item[0]))
else:
matches.sort(key=lambda item: item[0])
paths = [item[0] for item in matches]
paged, truncated = _paginate(paths, limit, offset)
if not paged:
return "No files found"
result = "\n".join(paged)
note = _pagination_note(limit, offset, truncated)
if note:
result += "\n\n" + note
return result
except PermissionError as e: except PermissionError as e:
return ToolResult.error(f"Error: {e}") return ToolResult.error(f"Error: {e}")
except Exception as e: except Exception as e:
return ToolResult.error(f"Error finding files: {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 ".")
if not target.exists():
return ToolResult.error(f"Error: Path not found: {path}")
if not (target.is_dir() or target.is_file()):
return ToolResult.error(f"Error: Unsupported path: {path}")
if sort not in {"path", "modified"}:
return ToolResult.error("Error: sort must be 'path' or 'modified'")
limit = (
_DEFAULT_FILE_HEAD_LIMIT
if head_limit is None
else None if head_limit == 0 else head_limit
)
root = target if target.is_dir() else target.parent
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):
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:
continue
rel_path = candidate.relative_to(root).as_posix()
display_path = self._display_path(candidate, root)
name = candidate.name
if glob and not _match_glob(rel_path, name, glob):
continue
if candidate.is_file() and not _matches_type(name, file_type):
continue
if candidate.is_dir() and file_type:
continue
if not _matches_query(display_path, query):
continue
try:
mtime = candidate.stat().st_mtime
except OSError:
mtime = 0.0
suffix = "/" if candidate.is_dir() else ""
matches.append((display_path + suffix, mtime))
if sort == "modified":
matches.sort(key=lambda item: (-item[1], item[0]))
else:
matches.sort(key=lambda item: item[0])
paths = [item[0] for item in matches]
paged, truncated = _paginate(paths, limit, offset)
if not paged:
return "No files found"
result = "\n".join(paged)
note = _pagination_note(limit, offset, truncated)
if note:
result += "\n\n" + note
return result
class GrepTool(_SearchTool): class GrepTool(_SearchTool):
"""Search file contents using a regex-like pattern.""" """Search file contents using a regex-like pattern."""
+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"])
+36 -9
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,14 +484,18 @@ class ExecTool(Tool):
+ _WORKSPACE_BOUNDARY_NOTE + _WORKSPACE_BOUNDARY_NOTE
) )
guard_error = self._guard_command( # Full access is an explicit trust decision. Keep the application-level
command, # command guard aligned with the selected access mode instead of
cwd, # continuing to block commands after workspace restriction is disabled.
restrict_to_workspace=access.restrict_to_workspace, if access.restrict_to_workspace:
workspace_root=workspace_root, guard_error = self._guard_command(
) command,
if guard_error: cwd,
return guard_error restrict_to_workspace=True,
workspace_root=workspace_root,
)
if guard_error:
return guard_error
if self.sandbox: if self.sandbox:
if _IS_WINDOWS: if _IS_WINDOWS:
@@ -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",
+374 -41
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))
+2 -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,7 +59,7 @@ 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
@@ -197,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:
+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(
+39 -23
View File
@@ -44,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,
@@ -54,6 +55,7 @@ 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.recovery import recovery_state_from_metadata
@@ -446,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] = {}
@@ -458,18 +480,9 @@ class WebSocketChannel(BaseChannel):
recovery_state = recovery_state_from_metadata(metadata) recovery_state = recovery_state_from_metadata(metadata)
if recovery_state is not None: if recovery_state is not None:
fields["recovery_state"] = recovery_state fields["recovery_state"] = recovery_state
usage = metadata.get("_last_usage") usage = LLMUsage.from_dict(metadata.get("_last_usage"))
if isinstance(usage, dict): if usage is not None:
sanitized_usage: dict[str, int | float] = {} fields["usage"] = usage.to_turn_dict()
for key, value in cast(dict[object, object], usage).items():
if (
isinstance(key, str)
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:
@@ -518,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,
@@ -906,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,
@@ -968,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
@@ -1023,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.
@@ -1217,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):
@@ -1272,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:
@@ -1561,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,
@@ -2019,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,
@@ -2034,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
@@ -44,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
@@ -1511,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",
@@ -1524,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,
@@ -1730,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
@@ -1796,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()),
@@ -1865,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"
@@ -1954,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",
} }
@@ -2093,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()
@@ -2114,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(),
} }
@@ -3225,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",
@@ -3232,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,
), ),
)) ))
@@ -3242,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"},
@@ -5209,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
@@ -5229,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
@@ -2576,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
@@ -3769,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
+93 -36
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,
@@ -330,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)
@@ -361,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
@@ -371,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
@@ -443,7 +469,6 @@ 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,
@@ -493,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")
@@ -564,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
@@ -590,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
@@ -618,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
@@ -630,14 +668,15 @@ 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.
should_notify = await evaluate_response( with llm_usage_source("cron"):
response=response, should_notify = await evaluate_response(
task_context=prompt, response=response,
provider=agent.provider, task_context=prompt,
model=agent.model, provider=agent.provider,
evaluator_prompt=evaluator_prompt, model=agent.model,
default_notify=False, evaluator_prompt=evaluator_prompt,
) default_notify=False,
)
if should_notify: if should_notify:
logger.info("Heartbeat: completed, delivering response") logger.info("Heartbeat: completed, delivering response")
@@ -696,17 +735,27 @@ def _run_gateway(
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,
) )
@@ -913,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(
@@ -980,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:
+72 -25
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,29 +356,32 @@ 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,
) )
session.clear() async def reset_and_schedule_archive() -> None:
loop.sessions.save(session) session.clear()
loop.sessions.invalidate(session.key) await _save_session(loop, session)
if archive_snapshot is not None and runtime is not None: loop.sessions.invalidate(session.key)
loop.schedule_background( if archive_snapshot is not None and runtime is not None:
loop.consolidator.archive_session( # pyright: ignore[reportUnknownMemberType] loop.schedule_background(
archive_snapshot, loop.consolidator.archive_session( # pyright: ignore[reportUnknownMemberType]
archive_end=len(snapshot), archive_snapshot,
runtime=runtime, archive_end=len(snapshot),
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(
+230 -87
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 start(self) -> None: async def run_sync(
"""Start the cron service.""" self,
self._running = True operation: Callable[..., _T],
loaded = self._load_store() /,
if loaded is None: *args: Any,
# Store file existed but was corrupt and has been preserved with **kwargs: Any,
# a ``.corrupt-<ts>`` suffix. Bail out instead of starting with ) -> _T:
# an empty store; that would call ``_save_store`` and overwrite """Serialize a complete cron transaction in a worker thread.
# the now-renamed (but still recoverable) data with [].
self._running = False A running thread cannot be cancelled safely. Keep the transaction lock
raise RuntimeError( until it exits so cancellation is never reported while that worker can
f"cron store at {self.store_path} is corrupt and was preserved; " still mutate cron state behind a later operation.
"refusing to start with an empty job list. " """
"Inspect the .corrupt-<ts> backup and restore manually." async with self._operation_lock:
return await shield_and_drain(
asyncio.to_thread(operation, *args, **kwargs)
) )
self._recompute_next_runs()
self._save_store() async def start(self) -> None:
self._arm_timer() """Start the cron service and settle accepted work before cancellation."""
logger.info("Cron service started with {} jobs", len(self._store.jobs if self._store else []))
async def settle_start() -> None:
self._event_loop = asyncio.get_running_loop()
self._running = True
try:
async with self._operation_lock:
loaded = await asyncio.to_thread(self._load_store)
if loaded is None:
# Store file existed but was corrupt and has been preserved with
# a ``.corrupt-<ts>`` suffix. Bail out instead of starting with
# an empty store; that would call ``_save_store`` and overwrite
# the now-renamed (but still recoverable) data with [].
raise RuntimeError(
f"cron store at {self.store_path} is corrupt and was preserved; "
"refusing to start with an empty job list. "
"Inspect the .corrupt-<ts> backup and restore manually."
)
self._recompute_next_runs()
await asyncio.to_thread(self._save_store)
self._arm_timer()
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.
if self._store_dirty: async with self._operation_lock:
self._save_store() if self._store_dirty:
return await shield_and_drain(asyncio.to_thread(self._save_store))
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,58 +629,124 @@ 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."""
async with self._operation_lock:
if job_id in self._claimed_job_ids:
return False
self._claimed_job_ids.add(job_id)
return True
async def _release_job_claim(self, job_id: str) -> None:
async with self._operation_lock:
self._claimed_job_ids.discard(job_id)
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()
async with self._operation_lock:
job.state.last_status = status
job.state.last_error = error
job.state.last_run_at_ms = start_ms
job.updated_at_ms = end_ms
job.state.run_history.append(CronRunRecord(
run_at_ms=start_ms,
status=status,
duration_ms=end_ms - start_ms,
error=error,
))
job.state.run_history = job.state.run_history[-self._MAX_RUN_HISTORY:]
if job.schedule.kind == "at":
if job.delete_after_run:
store = await shield_and_drain(
asyncio.to_thread(self._require_store)
)
store.jobs = [item for item in store.jobs if item.id != job.id]
else:
job.enabled = False
job.state.next_run_at_ms = None
else:
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() start_ms = _now_ms()
logger.info("Cron: executing job '{}' ({})", job.name, job.id) 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: try:
await self.on_job(job) 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)
job.state.last_status = "ok" settlement = asyncio.create_task(
job.state.last_error = None self._settle_job_execution(
logger.info("Cron: job '{}' completed", job.name) job,
start_ms=start_ms,
except CronJobSkippedError as e: status=status,
job.state.last_status = "skipped" error=error,
job.state.last_error = str(e) or None persist=accepted_cancellation is not None,
logger.warning("Cron: job '{}' skipped: {}", job.name, job.state.last_error or "") )
except asyncio.CancelledError as e: )
current = asyncio.current_task() if accepted_cancellation is not None:
if current is not None and current.cancelling(): await self._drain_settlement_on_cancellation(settlement)
raise raise accepted_cancellation
job.state.last_status = "error" await settlement
job.state.last_error = str(e) or e.__class__.__name__ return True
logger.exception("Cron: job '{}' was cancelled", job.name) finally:
except Exception as e: await self._release_job_claim(job.id)
job.state.last_status = "error"
job.state.last_error = str(e)
logger.exception("Cron: job '{}' failed", job.name)
end_ms = _now_ms()
job.state.last_run_at_ms = start_ms
job.updated_at_ms = end_ms
job.state.run_history.append(CronRunRecord(
run_at_ms=start_ms,
status=job.state.last_status,
duration_ms=end_ms - start_ms,
error=job.state.last_error,
))
job.state.run_history = job.state.run_history[-self._MAX_RUN_HISTORY:]
# Handle one-shot jobs
if job.schedule.kind == "at":
if job.delete_after_run:
store = self._require_store()
store.jobs = [item for item in store.jobs if item.id != job.id]
else:
job.enabled = False
job.state.next_run_at_ms = None
else:
# Compute next run
job.state.next_run_at_ms = _compute_next_run(job.schedule, _now_ms())
def _append_action( def _append_action(
self, self,
@@ -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.
if self._store_dirty: async with self._operation_lock:
self._save_store() if self._store_dirty:
store = self._require_store(reload_during_execution=reload_store) await shield_and_drain(asyncio.to_thread(self._save_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):
self._enforce_agent_binding(job) async with self._operation_lock:
self._save_store() self._enforce_agent_binding(job)
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)
+144 -42
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}", ()
try: async def async_resolve_url_target(
literal_addr = ipaddress.ip_address(normalized_hostname) url: str,
except ValueError: *,
return True, "", () allow_loopback: bool = False,
if _is_private(literal_addr): trust_remote_dns: bool = False,
return False, f"Blocked private/internal address: {literal_addr}", () timeout_s: float = _DNS_RESOLUTION_TIMEOUT_SECONDS,
return True, "", (str(_normalize_addr(literal_addr)),) ) -> tuple[bool, str, tuple[str, ...]]:
"""Resolve and validate an HTTP target without blocking the event loop."""
addrs: list[ipaddress.IPv4Address | ipaddress.IPv6Address] = [] hostname, error = _parse_url_hostname(url)
for info in infos: if hostname is None:
try: return False, error or "Missing hostname", ()
addr = ipaddress.ip_address(info[4][0]) loop = asyncio.get_running_loop()
except ValueError: try:
continue infos = await asyncio.wait_for(
addrs.append(addr) loop.getaddrinfo(
if allow_loopback and _is_allowed_loopback_target(hostname, addrs): hostname,
return True, "", tuple(dict.fromkeys(str(_normalize_addr(addr)) for addr in addrs)) None,
for addr in addrs: family=socket.AF_UNSPEC,
if _is_private(addr): type=socket.SOCK_STREAM,
return False, f"Blocked: {hostname} resolves to private/internal address {addr}", () ),
timeout=timeout_s,
return True, "", tuple(dict.fromkeys(str(_normalize_addr(addr)) for addr in addrs)) )
except asyncio.TimeoutError:
return False, f"Timed out resolving hostname: {hostname}", ()
except socket.gaierror:
return _unresolved_target_result(hostname, trust_remote_dns=trust_remote_dns)
return _resolved_target_result(hostname, infos, allow_loopback=allow_loopback)
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))
+113 -1
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,
@@ -71,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
@@ -560,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)
@@ -1642,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
@@ -1741,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,
@@ -1774,6 +1801,17 @@ 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: def save_runtime_checkpoint(self, session: Session) -> None:
"""Persist volatile recovery state without rewriting long history.""" """Persist volatile recovery state without rewriting long history."""
if not session.policy.persist: if not session.policy.persist:
@@ -1786,6 +1824,23 @@ class SessionManager:
# they opt into a dedicated checkpoint primitive. # they opt into a dedicated checkpoint primitive.
self.save(session) 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:
@@ -1827,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.
@@ -1858,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()
@@ -1930,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,
@@ -1943,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)
+55 -23
View File
@@ -25,6 +25,7 @@ from nanobot.bus.outbound_events import (
) )
from nanobot.bus.queue import MessageBus from nanobot.bus.queue import MessageBus
from nanobot.session import turn_continuation 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.keys import UNIFIED_SESSION_KEY, last_channel_from_metadata
from nanobot.session.manager import Session, SessionManager from nanobot.session.manager import Session, SessionManager
from nanobot.webui.metadata import WEBUI_TURN_METADATA_KEY from nanobot.webui.metadata import WEBUI_TURN_METADATA_KEY
@@ -460,6 +461,37 @@ class RecoveryCoordinator:
repr=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: def register_recovery_task(self, session_key: str, task: asyncio.Task[Any]) -> None:
"""Track the task that owns an explicit recovery continuation.""" """Track the task that owns an explicit recovery continuation."""
self._active_recovery_tasks[session_key] = task self._active_recovery_tasks[session_key] = task
@@ -482,8 +514,8 @@ class RecoveryCoordinator:
async def scan(self) -> None: async def scan(self) -> None:
"""Recover every interrupted WebUI session once at gateway startup.""" """Recover every interrupted WebUI session once at gateway startup."""
for key in self._recovery_candidates(): for key in await self._recovery_candidates():
metadata_payload = self.sessions.read_session_metadata(key) metadata_payload = await self._read_session_metadata(key)
raw_metadata = metadata_payload.get("metadata") if metadata_payload else None raw_metadata = metadata_payload.get("metadata") if metadata_payload else None
metadata = cast(dict[str, Any], raw_metadata) if isinstance(raw_metadata, dict) else {} metadata = cast(dict[str, Any], raw_metadata) if isinstance(raw_metadata, dict) else {}
route = self._websocket_route_for(key, metadata) route = self._websocket_route_for(key, metadata)
@@ -492,7 +524,7 @@ class RecoveryCoordinator:
unfinished = self._has_unfinished_webui_transcript(key) unfinished = self._has_unfinished_webui_transcript(key)
if not self._needs_recovery(metadata) and not unfinished: if not self._needs_recovery(metadata) and not unfinished:
continue continue
session = self.sessions.get_or_create(key) session = await self._get_or_create_session(key)
try: try:
await self._recover_session(session, route[1]) await self._recover_session(session, route[1])
await self._requeue_pending_followups(session) await self._requeue_pending_followups(session)
@@ -507,14 +539,14 @@ class RecoveryCoordinator:
reason="recovery_failed", reason="recovery_failed",
can_continue=False, can_continue=False,
) )
self.sessions.save(session) await self._save_session(session)
await self._publish(route[1], failed) await self._publish(route[1], failed)
def _recovery_candidates(self) -> list[str]: async def _recovery_candidates(self) -> list[str]:
"""Discover canonical and transcript-only WebUI sessions cheaply.""" """Discover canonical and transcript-only WebUI sessions cheaply."""
candidates = dict.fromkeys( candidates = dict.fromkeys(
key key
for item in self.sessions.list_sessions() for item in await self._list_sessions()
if isinstance((key := item.get("key")), str) if isinstance((key := item.get("key")), str)
) )
try: try:
@@ -523,7 +555,7 @@ class RecoveryCoordinator:
# duplicating its filename and migration rules here would drift. # duplicating its filename and migration rules here would drift.
from nanobot.webui.session_list_index import list_webui_sessions from nanobot.webui.session_list_index import list_webui_sessions
for item in list_webui_sessions(self.sessions): for item in await asyncio.to_thread(list_webui_sessions, self.sessions):
key = item.get("key") key = item.get("key")
if isinstance(key, str): if isinstance(key, str):
candidates.setdefault(key, None) candidates.setdefault(key, None)
@@ -549,7 +581,7 @@ class RecoveryCoordinator:
"""Reject stale queued recoveries and let new user input supersede them.""" """Reject stale queued recoveries and let new user input supersede them."""
recovery_id = message.metadata.get(RECOVERY_INBOUND_METADATA_KEY) recovery_id = message.metadata.get(RECOVERY_INBOUND_METADATA_KEY)
if isinstance(recovery_id, str): if isinstance(recovery_id, str):
session = self.sessions.get_or_create(message.session_key) session = await self._get_or_create_session(message.session_key)
state = recovery_state_from_metadata(session.metadata) state = recovery_state_from_metadata(session.metadata)
return bool( return bool(
state state
@@ -558,7 +590,7 @@ class RecoveryCoordinator:
) )
if message.channel != "websocket": if message.channel != "websocket":
return True return True
session = self.sessions.get_or_create(message.session_key) session = await self._get_or_create_session(message.session_key)
state = recovery_state_from_metadata(session.metadata) state = recovery_state_from_metadata(session.metadata)
if state and state["status"] in {"resuming", "awaiting_user", "failed"}: if state and state["status"] in {"resuming", "awaiting_user", "failed"}:
await self._cancel_active_recovery(message.session_key) await self._cancel_active_recovery(message.session_key)
@@ -572,13 +604,13 @@ class RecoveryCoordinator:
attempts=cast(int, state.get("attempts", 0)), attempts=cast(int, state.get("attempts", 0)),
reason="superseded", reason="superseded",
) )
self.sessions.save(session) await self._save_session(session)
await self._publish(message.chat_id, recovered) await self._publish(message.chat_id, recovered)
return True return True
async def turn_completed(self, session_key: str) -> None: async def turn_completed(self, session_key: str) -> None:
"""Resolve a resuming state after the recovered turn commits.""" """Resolve a resuming state after the recovered turn commits."""
session = self.sessions.get_or_create(session_key) session = await self._get_or_create_session(session_key)
state = recovery_state_from_metadata(session.metadata) state = recovery_state_from_metadata(session.metadata)
if not state or state["status"] != "resuming": if not state or state["status"] != "resuming":
return return
@@ -592,7 +624,7 @@ class RecoveryCoordinator:
attempts=cast(int, state.get("attempts", 0)), attempts=cast(int, state.get("attempts", 0)),
reason="continued", reason="continued",
) )
self.sessions.save(session) await self._save_session(session)
await self._publish(route[1], recovered) await self._publish(route[1], recovered)
async def handle_action(self, action: str, payload: dict[str, Any]) -> dict[str, Any]: async def handle_action(self, action: str, payload: dict[str, Any]) -> dict[str, Any]:
@@ -603,7 +635,7 @@ class RecoveryCoordinator:
raise RecoveryActionError("missing chat_id") raise RecoveryActionError("missing chat_id")
if not isinstance(recovery_id, str) or not recovery_id: if not isinstance(recovery_id, str) or not recovery_id:
raise RecoveryActionError("missing recovery_id") raise RecoveryActionError("missing recovery_id")
session = self.sessions.get_or_create(self._session_key(chat_id)) session = await self._get_or_create_session(self._session_key(chat_id))
state = recovery_state_from_metadata(session.metadata) state = recovery_state_from_metadata(session.metadata)
if not state or state["recovery_id"] != recovery_id: if not state or state["recovery_id"] != recovery_id:
raise RecoveryActionError("recovery state is stale", status=409) raise RecoveryActionError("recovery state is stale", status=409)
@@ -618,7 +650,7 @@ class RecoveryCoordinator:
attempts=cast(int, state.get("attempts", 0)), attempts=cast(int, state.get("attempts", 0)),
reason="dismissed", reason="dismissed",
) )
self.sessions.save(session) await self._save_session(session)
await self._publish(chat_id, next_state) await self._publish(chat_id, next_state)
return next_state return next_state
if action != "continue": if action != "continue":
@@ -635,7 +667,7 @@ class RecoveryCoordinator:
reason="user_confirmed", reason="user_confirmed",
resume_message_count=len(session.messages), resume_message_count=len(session.messages),
) )
self.sessions.save(session) await self._save_session(session)
await self._publish(chat_id, next_state) await self._publish(chat_id, next_state)
await self._queue_continuation(session, chat_id, next_state) await self._queue_continuation(session, chat_id, next_state)
return next_state return next_state
@@ -668,7 +700,7 @@ class RecoveryCoordinator:
attempts=cast(int, state.get("attempts", 1)), attempts=cast(int, state.get("attempts", 1)),
reason="loop_guard", reason="loop_guard",
) )
self.sessions.save(session) await self._save_session(session)
await self._publish(chat_id, next_state) await self._publish(chat_id, next_state)
elif self._has_unfinished_webui_transcript(session.key): elif self._has_unfinished_webui_transcript(session.key):
# A normal last-client shutdown can materialize the checkpoint # A normal last-client shutdown can materialize the checkpoint
@@ -690,7 +722,7 @@ class RecoveryCoordinator:
), ),
can_continue=can_continue, can_continue=can_continue,
) )
self.sessions.save(session) await self._save_session(session)
await self._publish(chat_id, waiting) await self._publish(chat_id, waiting)
return return
if state and state["status"] in {"awaiting_user", "failed"}: if state and state["status"] in {"awaiting_user", "failed"}:
@@ -706,7 +738,7 @@ class RecoveryCoordinator:
attempts=cast(int, state.get("attempts", 1)), attempts=cast(int, state.get("attempts", 1)),
reason="loop_guard", reason="loop_guard",
) )
self.sessions.save(session) await self._save_session(session)
await self._publish(chat_id, waiting) await self._publish(chat_id, waiting)
return return
@@ -724,7 +756,7 @@ class RecoveryCoordinator:
reason="checkpoint_unknown", reason="checkpoint_unknown",
can_continue=False, can_continue=False,
) )
self.sessions.save(session) await self._save_session(session)
await self._publish(chat_id, waiting) await self._publish(chat_id, waiting)
return return
if checkpoint is not None and not _runtime_checkpoint_is_well_formed(checkpoint): if checkpoint is not None and not _runtime_checkpoint_is_well_formed(checkpoint):
@@ -738,7 +770,7 @@ class RecoveryCoordinator:
reason="checkpoint_invalid", reason="checkpoint_invalid",
can_continue=False, can_continue=False,
) )
self.sessions.save(session) await self._save_session(session)
await self._publish(chat_id, waiting) await self._publish(chat_id, waiting)
return return
if phase == "final_response": if phase == "final_response":
@@ -750,7 +782,7 @@ class RecoveryCoordinator:
attempts=0, attempts=0,
reason="answer_restored", reason="answer_restored",
) )
self.sessions.save(session) await self._save_session(session)
await self._publish(chat_id, recovered) await self._publish(chat_id, recovered)
return return
if phase in _UNCERTAIN_TOOL_PHASES or pending_calls: if phase in _UNCERTAIN_TOOL_PHASES or pending_calls:
@@ -762,7 +794,7 @@ class RecoveryCoordinator:
attempts=0, attempts=0,
reason="tool_state_unknown", reason="tool_state_unknown",
) )
self.sessions.save(session) await self._save_session(session)
await self._publish(chat_id, waiting) await self._publish(chat_id, waiting)
return return
# A gateway restart is a lifecycle boundary. Never enqueue model work # A gateway restart is a lifecycle boundary. Never enqueue model work
@@ -777,7 +809,7 @@ class RecoveryCoordinator:
attempts=0, attempts=0,
reason="restart_requires_confirmation", reason="restart_requires_confirmation",
) )
self.sessions.save(session) await self._save_session(session)
await self._publish(chat_id, waiting) await self._publish(chat_id, waiting)
async def _queue_continuation( async def _queue_continuation(
+36 -31
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,7 +38,8 @@ 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
@@ -175,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:
@@ -186,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)
@@ -208,24 +210,25 @@ 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:
response = await provider.chat_with_retry( with llm_usage_source("system"):
[ response = await provider.chat_with_retry(
{ [
"role": "system", {
"content": ( "role": "system",
"You write short, neutral chat titles. " "content": (
"Return only the title text." "You write short, neutral chat titles. "
), "Return only the title text."
}, ),
{"role": "user", "content": prompt}, },
], {"role": "user", "content": prompt},
tools=None, ],
model=model, tools=None,
max_tokens=TITLE_GENERATION_MAX_TOKENS, model=model,
temperature=0.2, max_tokens=TITLE_GENERATION_MAX_TOKENS,
reasoning_effort=TITLE_GENERATION_REASONING_EFFORT, temperature=0.2,
retry_mode="standard", reasoning_effort=TITLE_GENERATION_REASONING_EFFORT,
) retry_mode="standard",
)
except Exception: except Exception:
logger.debug("Failed to generate webui session title for {}", session_key, exc_info=True) logger.debug("Failed to generate webui session title for {}", session_key, exc_info=True)
return False return False
@@ -239,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
@@ -435,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({
@@ -578,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)
@@ -589,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"],
@@ -614,8 +618,9 @@ 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)
mark_webui_session(session, event.context.metadata) if session is not None:
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:
if not self._is_websocket_event(event.context): if not self._is_websocket_event(event.context):
@@ -695,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,
@@ -709,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."
``` ```
+215 -59
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,
) )
response = await submit_turn(msg) try:
store.record_delivery( response = await submit_turn(msg)
trigger.id, except AutomationTurnAcceptedCancellation:
status="ok", try:
run_at_ms=delivery.created_at_ms, 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,
) )
_write_delivery_run_record( 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,
status="ok",
run_at_ms=delivery.created_at_ms,
)
)
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,12 +329,15 @@ 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(
delivery, asyncio.to_thread(
trigger=trigger, store.write_delivery_run_record,
status=status, delivery,
error=error, trigger=trigger,
response=response, status=status,
error=error,
response=response,
)
) )
except Exception: except Exception:
logger.exception( logger.exception(
+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)
+32 -19
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,24 +95,35 @@ 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
try: async def create_and_attach() -> None:
forked = create_webui_chat_fork( try:
session_manager, forked = await asyncio.to_thread(
source_chat_id=source_chat_id, create_webui_chat_fork,
before_user_index=raw_index, session_manager,
title=envelope.get("title") if isinstance(envelope.get("title"), str) else None, source_chat_id=source_chat_id,
) before_user_index=raw_index,
if forked is None: title=(
await channel.send_webui_protocol_error(connection, "invalid fork source or index") envelope.get("title")
if isinstance(envelope.get("title"), str)
else None
),
)
if forked is None:
await channel.send_webui_protocol_error(
connection,
"invalid fork source or index",
)
return
fork_id, fork_key = forked
except Exception as exc:
channel.logger.warning("fork_chat failed: {}", exc)
await channel.send_webui_protocol_error(connection, "fork_chat_failed")
return return
fork_id, fork_key = forked
except Exception as exc:
channel.logger.warning("fork_chat failed: {}", exc)
await channel.send_webui_protocol_error(connection, "fork_chat_failed")
return
await channel.attach_webui_fork( await channel.attach_webui_fork(
connection, connection,
fork_id=fork_id, fork_id=fork_id,
fork_key=fork_key, fork_key=fork_key,
) )
await shield_and_drain(create_and_attach())
+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,
+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
+68 -32
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(
operations.update_model, self._mutate_and_refresh(
request.query, operations.update_model,
rename_model_preset=self.settings.rename_model_preset, request.query,
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(
operations.update_provider, self._update_provider_and_runtime(
request.query, operations.update_provider,
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,11 +1754,13 @@ 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(
self.settings.read, asyncio.to_thread(
operations.oauth_login, self.settings.read,
request.query, operations.oauth_login,
oauth_flows=self.settings.oauth_flows, request.query,
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")
@@ -1736,19 +1768,23 @@ class ModelSettingsHandler:
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(
self.settings.read, asyncio.to_thread(
operations.oauth_complete, self.settings.read,
request.query, operations.oauth_complete,
raw_response or None, request.query,
oauth_flows=self.settings.oauth_flows, raw_response or None,
oauth_flows=self.settings.oauth_flows,
)
) )
elif action == "oauth-logout": elif action == "oauth-logout":
payload = await asyncio.to_thread( payload = await shield_and_drain(
self.settings.read, asyncio.to_thread(
operations.oauth_logout, self.settings.read,
request.query, operations.oauth_logout,
oauth_flows=self.settings.oauth_flows, request.query,
oauth_flows=self.settings.oauth_flows,
)
) )
else: else:
return SettingsRouteResult.failure(404, "unknown settings action") return SettingsRouteResult.failure(404, "unknown settings action")
+36 -21
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( settings_payload,
self.settings.read( surface=self._runtime_surface,
settings_payload, runtime_capability_overrides=self._runtime_capabilities,
surface=self._runtime_surface,
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")
+8 -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 = (
+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)
+113 -24
View File
@@ -34,6 +34,7 @@ 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,
@@ -135,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",
@@ -420,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,
@@ -556,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
@@ -564,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,
) )
@@ -694,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:
@@ -702,15 +732,27 @@ 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
@@ -957,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))
@@ -1029,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":
@@ -1044,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":
@@ -1068,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":
@@ -1078,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,
@@ -1163,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)
@@ -1276,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")
@@ -1308,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")
+5 -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,
@@ -1883,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",
[], [],
@@ -1898,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
+55 -8
View File
@@ -90,7 +90,7 @@ async def test_llm_error_not_appended_to_session_messages():
provider = MagicMock(spec=LLMProvider) provider = MagicMock(spec=LLMProvider)
provider.chat_with_retry = AsyncMock(return_value=LLMResponse( provider.chat_with_retry = AsyncMock(return_value=LLMResponse(
content="429 rate limit exceeded", finish_reason="error", tool_calls=[], usage={}, content="429 rate limit exceeded", finish_reason="error", tool_calls=[], usage=None,
)) ))
tools = MagicMock() tools = MagicMock()
tools.get_definitions.return_value = [] tools.get_definitions.return_value = []
@@ -158,7 +158,7 @@ async def test_runner_ignores_tool_calls_when_finish_reason_blocks_execution(
content="Request blocked by provider policy.", content="Request blocked by provider policy.",
finish_reason=finish_reason, finish_reason=finish_reason,
tool_calls=[ToolCallRequest(id="call_1", name="exec", arguments={"command": "echo nope"})], tool_calls=[ToolCallRequest(id="call_1", name="exec", arguments={"command": "echo nope"})],
usage={}, usage=None,
)) ))
tools = MagicMock() tools = MagicMock()
tools.get_definitions.return_value = [] tools.get_definitions.return_value = []
@@ -189,7 +189,7 @@ async def test_runner_tool_error_sets_final_content():
return LLMResponse( return LLMResponse(
content="working", content="working",
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={}, usage=None,
) )
provider.chat_with_retry = chat_with_retry provider.chat_with_retry = chat_with_retry
@@ -224,9 +224,9 @@ async def test_runner_preserves_successful_exec_output_that_starts_with_error():
tool_calls=[ tool_calls=[
ToolCallRequest(id="call_1", name="exec", arguments={"command": "report"}) ToolCallRequest(id="call_1", name="exec", arguments={"command": "report"})
], ],
usage={}, usage=None,
) )
return LLMResponse(content="done", usage={}) return LLMResponse(content="done", usage=None)
provider.chat_with_retry = chat_with_retry provider.chat_with_retry = chat_with_retry
output = "Error: generated report successfully\n\nExit code: 0" output = "Error: generated report successfully\n\nExit code: 0"
@@ -266,7 +266,7 @@ async def test_runner_tool_error_preserves_tool_results_in_messages():
ToolCallRequest(id="tc1", name="read_file", arguments={"path": "a"}), ToolCallRequest(id="tc1", name="read_file", arguments={"path": "a"}),
ToolCallRequest(id="tc2", name="exec", arguments={"cmd": "bad"}), ToolCallRequest(id="tc2", name="exec", arguments={"cmd": "bad"}),
], ],
usage={}, usage=None,
) )
provider.chat_with_retry = chat_with_retry provider.chat_with_retry = chat_with_retry
@@ -333,9 +333,9 @@ async def test_length_finish_with_blank_content_routes_to_length_recovery():
content="", content="",
finish_reason="length", finish_reason="length",
tool_calls=[ToolCallRequest(id="call_1", name="exec", arguments={})], tool_calls=[ToolCallRequest(id="call_1", name="exec", arguments={})],
usage={}, usage=None,
), ),
LLMResponse(content="done", finish_reason="stop", tool_calls=[], usage={}), LLMResponse(content="done", finish_reason="stop", tool_calls=[], usage=None),
]) ])
tools = MagicMock() tools = MagicMock()
tools.get_definitions.return_value = [] tools.get_definitions.return_value = []
@@ -357,3 +357,50 @@ async def test_length_finish_with_blank_content_routes_to_length_recovery():
"finish_reason='length' response with blank content" "finish_reason='length' response with blank content"
) )
assert result.final_content == "done" assert result.final_content == "done"
@pytest.mark.asyncio
async def test_slow_tool_log_records_scale_without_argument_content(monkeypatch) -> None:
from nanobot.agent import runner as runner_module
from nanobot.agent.runner import AgentRunner
records: list[str] = []
class _Logger:
def warning(self, message: str, *args: object) -> None:
records.append(message.format(*args))
secret = "customer-token-do-not-log"
async def execute(_name, _args):
return "ok"
monkeypatch.setattr(runner_module, "_SLOW_TOOL_LOG_MS", 0)
monkeypatch.setattr(runner_module, "logger", _Logger())
runner = AgentRunner()
spec = make_run_spec(
MagicMock(spec=LLMProvider),
initial_messages=[],
tools=SimpleNamespace(execute=execute),
model="test-model",
max_iterations=1,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
)
await runner._run_tool(
spec,
ToolCallRequest(
id="call_1",
name="edit_file",
arguments={"old_text": secret, "paths": ["one", "two"]},
),
external_lookup_counts={},
workspace_violation_counts={},
)
assert len(records) == 1
assert "operation=edit_file" in records[0]
assert "input_items=4" in records[0]
assert f"input_chars={len(secret)}" in records[0]
assert "duration_ms=" in records[0]
assert secret not in records[0]
+36 -1
View File
@@ -78,7 +78,7 @@ class _FakeProvider(LLMProvider):
*, *,
responses: list[LLMResponse] | None = None, responses: list[LLMResponse] | None = None,
): ):
super().__init__() super().__init__(provider_name=name)
self.name = name self.name = name
self._response = response or _make_response() self._response = response or _make_response()
self._responses = iter(responses) if responses is not None else None self._responses = iter(responses) if responses is not None else None
@@ -260,6 +260,41 @@ def test_provider_snapshot_uses_smallest_fallback_context_window() -> None:
assert snapshot.provider._primary_context_window_tokens == 128000 assert snapshot.provider._primary_context_window_tokens == 128000
def test_factory_injects_configured_identity_into_primary_and_fallback_leaves() -> None:
from nanobot.config.schema import Config
from nanobot.providers.factory import build_provider_snapshot
config = Config.model_validate({
"agents": {
"defaults": {
"modelPreset": "primary",
"fallbackModels": ["backup"],
}
},
"modelPresets": {
"primary": {"model": "primary-model", "provider": "primary_edge"},
"backup": {"model": "backup-model", "provider": "backup_edge"},
},
"providers": {
"primary_edge": {
"apiKey": "primary-key",
"apiBase": "https://primary.example/v1",
},
"backup_edge": {
"apiKey": "backup-key",
"apiBase": "https://backup.example/v1",
},
},
})
snapshot = build_provider_snapshot(config)
assert isinstance(snapshot.provider, FallbackProvider)
assert snapshot.provider._primary.provider_name == "primary_edge"
fallback = snapshot.provider._provider_factory(snapshot.provider._fallback_presets[0])
assert fallback.provider_name == "backup_edge"
def test_inline_fallback_reasoning_effort_does_not_inherit_primary() -> None: def test_inline_fallback_reasoning_effort_does_not_inherit_primary() -> None:
from nanobot.config.schema import Config from nanobot.config.schema import Config
from nanobot.providers.factory import provider_signature from nanobot.providers.factory import provider_signature
+8 -8
View File
@@ -25,7 +25,7 @@ async def test_runner_exits_normally_without_predicate():
provider = MagicMock(spec=LLMProvider) provider = MagicMock(spec=LLMProvider)
provider.chat_with_retry = AsyncMock(return_value=LLMResponse( provider.chat_with_retry = AsyncMock(return_value=LLMResponse(
content="all done", tool_calls=[], usage={}, content="all done", tool_calls=[], usage=None,
)) ))
tools = MagicMock() tools = MagicMock()
tools.get_definitions.return_value = [] tools.get_definitions.return_value = []
@@ -50,7 +50,7 @@ async def test_runner_exits_normally_with_inactive_goal():
provider = MagicMock(spec=LLMProvider) provider = MagicMock(spec=LLMProvider)
provider.chat_with_retry = AsyncMock(return_value=LLMResponse( provider.chat_with_retry = AsyncMock(return_value=LLMResponse(
content="all done", tool_calls=[], usage={}, content="all done", tool_calls=[], usage=None,
)) ))
tools = MagicMock() tools = MagicMock()
tools.get_definitions.return_value = [] tools.get_definitions.return_value = []
@@ -82,7 +82,7 @@ async def test_runner_forces_continue_when_goal_active():
provider = MagicMock(spec=LLMProvider) provider = MagicMock(spec=LLMProvider)
provider.chat_with_retry = AsyncMock(return_value=LLMResponse( provider.chat_with_retry = AsyncMock(return_value=LLMResponse(
content="still working", tool_calls=[], usage={}, content="still working", tool_calls=[], usage=None,
)) ))
tools = MagicMock() tools = MagicMock()
tools.get_definitions.return_value = [] tools.get_definitions.return_value = []
@@ -112,7 +112,7 @@ async def test_runner_respects_max_iterations_even_with_active_goal():
provider = MagicMock(spec=LLMProvider) provider = MagicMock(spec=LLMProvider)
provider.chat_with_retry = AsyncMock(return_value=LLMResponse( provider.chat_with_retry = AsyncMock(return_value=LLMResponse(
content="still working", tool_calls=[], usage={}, content="still working", tool_calls=[], usage=None,
)) ))
tools = MagicMock() tools = MagicMock()
tools.get_definitions.return_value = [] tools.get_definitions.return_value = []
@@ -137,7 +137,7 @@ async def test_runner_goal_continue_not_limited_by_injection_cycle_cap():
provider = MagicMock(spec=LLMProvider) provider = MagicMock(spec=LLMProvider)
provider.chat_with_retry = AsyncMock(return_value=LLMResponse( provider.chat_with_retry = AsyncMock(return_value=LLMResponse(
content="still working", tool_calls=[], usage={}, content="still working", tool_calls=[], usage=None,
)) ))
tools = MagicMock() tools = MagicMock()
tools.get_definitions.return_value = [] tools.get_definitions.return_value = []
@@ -165,7 +165,7 @@ async def test_runner_does_not_force_continue_on_error():
provider = MagicMock(spec=LLMProvider) provider = MagicMock(spec=LLMProvider)
provider.chat_with_retry = AsyncMock(return_value=LLMResponse( provider.chat_with_retry = AsyncMock(return_value=LLMResponse(
content=None, tool_calls=[], usage={}, content=None, tool_calls=[], usage=None,
finish_reason="error", finish_reason="error",
)) ))
tools = MagicMock() tools = MagicMock()
@@ -191,7 +191,7 @@ async def test_runner_uses_custom_goal_continue_message():
provider = MagicMock(spec=LLMProvider) provider = MagicMock(spec=LLMProvider)
provider.chat_with_retry = AsyncMock(return_value=LLMResponse( provider.chat_with_retry = AsyncMock(return_value=LLMResponse(
content="still working", tool_calls=[], usage={}, content="still working", tool_calls=[], usage=None,
)) ))
tools = MagicMock() tools = MagicMock()
tools.get_definitions.return_value = [] tools.get_definitions.return_value = []
@@ -220,7 +220,7 @@ async def test_runner_resolves_goal_continue_message_lazily():
provider = MagicMock(spec=LLMProvider) provider = MagicMock(spec=LLMProvider)
provider.chat_with_retry = AsyncMock(return_value=LLMResponse( provider.chat_with_retry = AsyncMock(return_value=LLMResponse(
content="still working", tool_calls=[], usage={}, content="still working", tool_calls=[], usage=None,
)) ))
tools = MagicMock() tools = MagicMock()
tools.get_definitions.return_value = [] tools.get_definitions.return_value = []
+3 -3
View File
@@ -273,7 +273,7 @@ async def test_runner_drops_orphan_tool_results_before_model_request():
async def chat_with_retry(*, messages, **kwargs): async def chat_with_retry(*, messages, **kwargs):
captured_messages[:] = messages captured_messages[:] = 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()
@@ -312,7 +312,7 @@ async def test_backfill_repairs_model_context_without_shifting_save_turn_boundar
provider = MagicMock() provider = MagicMock()
provider.get_default_model.return_value = "test-model" provider.get_default_model.return_value = "test-model"
response = LLMResponse(content="new answer", tool_calls=[], usage={}) response = LLMResponse(content="new answer", tool_calls=[], usage=None)
provider.chat_with_retry = AsyncMock(return_value=response) provider.chat_with_retry = AsyncMock(return_value=response)
provider.chat_stream_with_retry = AsyncMock(return_value=response) provider.chat_stream_with_retry = AsyncMock(return_value=response)
@@ -397,7 +397,7 @@ async def test_runner_backfill_only_mutates_model_context_not_returned_messages(
async def chat_with_retry(*, messages, **kwargs): async def chat_with_retry(*, messages, **kwargs):
captured_messages[:] = messages captured_messages[:] = 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()

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