mirror of
https://github.com/HKUDS/nanobot.git
synced 2026-08-31 08:13:11 +03:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
be3a42ebac | ||
|
|
7fb0811fbb | ||
|
|
2ac802b2d5 | ||
|
|
8bb3828487 | ||
|
|
9895c23cb5 | ||
|
|
89c94d8744 | ||
|
|
f5e467626d | ||
|
|
04974b7607 | ||
|
|
7f288a49fc | ||
|
|
09d3bd76c9 | ||
|
|
b1cadf53c5 | ||
|
|
baa0233377 | ||
|
|
d50a2fab32 | ||
|
|
8344066696 | ||
|
|
5accc903a0 | ||
|
|
1f0771c555 | ||
|
|
2850114eab | ||
|
|
7e66375f59 | ||
|
|
2cdfba38b2 | ||
|
|
c7e2a474a0 | ||
|
|
58a1cc48d8 | ||
|
|
41a2104244 | ||
|
|
12029f8812 |
@@ -23,6 +23,7 @@ if TYPE_CHECKING:
|
||||
STREAM_EVENT_TOOL_FAILED,
|
||||
STREAM_EVENT_TOOL_STARTED,
|
||||
STREAM_EVENT_TYPES,
|
||||
LLMUsage,
|
||||
Nanobot,
|
||||
RunResult,
|
||||
RunStream,
|
||||
@@ -56,6 +57,7 @@ __logo__ = "🐈"
|
||||
|
||||
_LAZY_EXPORTS = {
|
||||
"Nanobot": ".nanobot",
|
||||
"LLMUsage": ".nanobot",
|
||||
"RunStream": ".nanobot",
|
||||
"RunResult": ".nanobot",
|
||||
"RequestContext": ".agent.tools.context",
|
||||
@@ -93,6 +95,7 @@ def __getattr__(name: str) -> Any:
|
||||
|
||||
__all__ = [
|
||||
"Nanobot",
|
||||
"LLMUsage",
|
||||
"RunResult",
|
||||
"RequestContext",
|
||||
"RuntimeContextBlock",
|
||||
|
||||
@@ -2,7 +2,9 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Collection
|
||||
import asyncio
|
||||
import inspect
|
||||
from collections.abc import Awaitable, Collection
|
||||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Any, Callable, Coroutine
|
||||
|
||||
@@ -47,9 +49,27 @@ class AutoCompact:
|
||||
return idle_seconds >= self._ttl * 60
|
||||
|
||||
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)
|
||||
|
||||
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
|
||||
def _is_internal_session(cls, key: str) -> bool:
|
||||
return key.startswith(cls._INTERNAL_SESSION_PREFIXES)
|
||||
@@ -79,6 +99,31 @@ class AutoCompact:
|
||||
self._archiving.add(key)
|
||||
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:
|
||||
if self._is_internal_session(key):
|
||||
self._archiving.discard(key)
|
||||
@@ -90,18 +135,38 @@ class AutoCompact:
|
||||
max_suffix=self._RECENT_SUFFIX_MESSAGES,
|
||||
)
|
||||
if summary and summary != "(nothing)":
|
||||
session = 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
|
||||
self._record_stored_summary(key, self.sessions.get_or_create(key))
|
||||
except Exception:
|
||||
logger.exception("Auto-compact: failed for {}", key)
|
||||
finally:
|
||||
self._archiving.discard(key)
|
||||
|
||||
async def _archive_async(self, key: str, *, runtime: LLMRuntime) -> None:
|
||||
if self._is_internal_session(key):
|
||||
self._archiving.discard(key)
|
||||
return
|
||||
try:
|
||||
summary = await self.consolidator.compact_idle_session(
|
||||
key,
|
||||
runtime=runtime,
|
||||
max_suffix=self._RECENT_SUFFIX_MESSAGES,
|
||||
)
|
||||
if summary and summary != "(nothing)":
|
||||
session = await self._get_or_create_nonblocking(key)
|
||||
self._record_stored_summary(key, session)
|
||||
except Exception:
|
||||
logger.exception("Auto-compact: failed for {}", key)
|
||||
finally:
|
||||
self._archiving.discard(key)
|
||||
|
||||
def _record_stored_summary(self, key: str, session: Session) -> None:
|
||||
stored = session_summary_from_metadata(
|
||||
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]:
|
||||
if self._is_internal_session(key):
|
||||
self._archiving.discard(key)
|
||||
@@ -110,6 +175,28 @@ class AutoCompact:
|
||||
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 = 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).
|
||||
entry = self._summaries.pop(key, None)
|
||||
if entry:
|
||||
|
||||
@@ -13,6 +13,19 @@ class AutomationTurnError(RuntimeError):
|
||||
"""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(
|
||||
*,
|
||||
deferred_queues: dict[str, list[InboundMessage]],
|
||||
@@ -70,19 +83,36 @@ class AutomationTurnCoordinator:
|
||||
future: asyncio.Future[OutboundMessage | None] = loop.create_future()
|
||||
self._waiters[turn_id] = future
|
||||
self._pending_messages_by_turn_id[turn_id] = msg
|
||||
accepted = False
|
||||
try:
|
||||
if self._is_running():
|
||||
await self._publish_inbound(msg)
|
||||
accepted = True
|
||||
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:
|
||||
return await future
|
||||
except asyncio.CancelledError:
|
||||
except asyncio.CancelledError as exc:
|
||||
if accepted:
|
||||
raise AutomationTurnAcceptedCancellation(*exc.args) from None
|
||||
raise
|
||||
except AutomationTurnError:
|
||||
raise
|
||||
except Exception as 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:
|
||||
self._waiters.pop(turn_id, None)
|
||||
self._pending_messages_by_turn_id.pop(turn_id, None)
|
||||
|
||||
@@ -112,6 +112,14 @@ class ContextBuilder:
|
||||
|
||||
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:
|
||||
memory = self.memory.read_memory()
|
||||
if memory and not self._is_template_content(memory, "memory/MEMORY.md"):
|
||||
|
||||
@@ -9,7 +9,7 @@ from typing import Any
|
||||
|
||||
from loguru import logger
|
||||
|
||||
from nanobot.providers.base import LLMResponse, ToolCallRequest
|
||||
from nanobot.providers.base import LLMResponse, LLMUsage, ToolCallRequest
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
@@ -19,7 +19,7 @@ class AgentHookContext:
|
||||
iteration: int
|
||||
messages: list[dict[str, Any]]
|
||||
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_results: list[Any] = 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]]
|
||||
final_content: str | None = None
|
||||
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
|
||||
error: str | None = None
|
||||
tool_events: list[dict[str, str]] = field(default_factory=list)
|
||||
@@ -284,7 +284,7 @@ class SDKCaptureHook(AgentHook):
|
||||
super().__init__()
|
||||
self.tools_used: list[str] = []
|
||||
self.messages: list[dict[str, Any]] = []
|
||||
self.usage: dict[str, int] = {}
|
||||
self.usage: LLMUsage | None = None
|
||||
self.stop_reason: str | None = None
|
||||
self.error: str | None = None
|
||||
self.tool_events: list[dict[str, str]] = []
|
||||
@@ -294,7 +294,7 @@ class SDKCaptureHook(AgentHook):
|
||||
for call in context.tool_calls:
|
||||
self.tools_used.append(call.name)
|
||||
self.messages = list(context.messages)
|
||||
self.usage = dict(context.usage)
|
||||
self.usage = context.usage
|
||||
self.stop_reason = context.stop_reason
|
||||
self.error = context.error
|
||||
self.tool_events = list(context.tool_events)
|
||||
@@ -302,7 +302,7 @@ class SDKCaptureHook(AgentHook):
|
||||
async def after_run(self, context: AgentRunHookContext) -> None:
|
||||
self.tools_used = list(context.tools_used)
|
||||
self.messages = list(context.messages)
|
||||
self.usage = dict(context.usage)
|
||||
self.usage = context.usage
|
||||
self.stop_reason = context.stop_reason
|
||||
self.error = context.error
|
||||
self.tool_events = list(context.tool_events)
|
||||
|
||||
+168
-39
@@ -49,7 +49,8 @@ from nanobot.bus.queue import MessageBus
|
||||
from nanobot.bus.runtime_events import RuntimeEventBus
|
||||
from nanobot.command import CommandContext, CommandRouter, register_builtin_commands
|
||||
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.runtime_context import (
|
||||
RUNTIME_CONTEXT_HISTORY_META,
|
||||
@@ -66,6 +67,7 @@ from nanobot.security.workspace_access import (
|
||||
reset_workspace_scope,
|
||||
)
|
||||
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.goal_state import (
|
||||
goal_state_runtime_lines,
|
||||
@@ -167,7 +169,7 @@ class TurnContext:
|
||||
turn_wall_started_at: float = field(default_factory=time.time)
|
||||
visible_run_started_at: float | 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:
|
||||
"""Return the runtime established by the BUILD stage."""
|
||||
@@ -203,7 +205,7 @@ class AgentLoop:
|
||||
return self.tools.tool_names
|
||||
|
||||
@property
|
||||
def last_usage(self) -> Mapping[str, int]:
|
||||
def last_usage(self) -> LLMUsage | None:
|
||||
"""Latest aggregate usage exposed through the runtime-control snapshot."""
|
||||
return self._last_usage
|
||||
|
||||
@@ -378,7 +380,7 @@ class AgentLoop:
|
||||
default_restrict_to_workspace=restrict_to_workspace,
|
||||
)
|
||||
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._hook_factories: list[AgentTurnHookFactory] = hook_factories or []
|
||||
|
||||
@@ -539,6 +541,33 @@ class AgentLoop:
|
||||
**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:
|
||||
"""Keep subagent runtime limits aligned with mutable loop settings."""
|
||||
self.subagents.max_iterations = self.max_iterations
|
||||
@@ -578,6 +607,30 @@ class AgentLoop:
|
||||
self.sessions.save(session)
|
||||
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(
|
||||
self,
|
||||
session_key: str,
|
||||
@@ -590,6 +643,18 @@ class AgentLoop:
|
||||
self.sessions.save(session)
|
||||
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(
|
||||
self,
|
||||
runtime: LLMRuntime,
|
||||
@@ -702,17 +767,14 @@ class AgentLoop:
|
||||
session_key=session_key,
|
||||
)
|
||||
|
||||
def _persist_user_message_early(
|
||||
def _stage_user_message_early(
|
||||
self,
|
||||
msg: InboundMessage,
|
||||
session: Session,
|
||||
runtime_context_blocks: list[RuntimeContextBlock] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> bool:
|
||||
"""Persist the triggering user message before the turn starts.
|
||||
|
||||
Returns True if the message was persisted.
|
||||
"""
|
||||
"""Add the triggering user message and recovery markers in memory."""
|
||||
if not turn_continuation.should_persist_user_message(msg.metadata):
|
||||
return False
|
||||
media_paths = [
|
||||
@@ -741,10 +803,45 @@ class AgentLoop:
|
||||
followup_id = msg.metadata.get(PENDING_FOLLOWUP_ID_KEY)
|
||||
if isinstance(followup_id, str) and followup_id:
|
||||
acknowledge_pending_followups(session, [followup_id])
|
||||
self.sessions.save(session)
|
||||
return True
|
||||
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]]:
|
||||
"""Build the initial message list for the LLM turn."""
|
||||
assert ctx.session is not None
|
||||
@@ -842,7 +939,7 @@ class AgentLoop:
|
||||
if tool is None:
|
||||
content = "Shell execution is disabled in this nanobot configuration."
|
||||
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(
|
||||
channel=ctx.msg.channel,
|
||||
message_metadata=metadata,
|
||||
@@ -1000,7 +1097,7 @@ class AgentLoop:
|
||||
public_payload[self._PROVIDER_STATE_CHECKPOINT_VERSION_KEY] = (
|
||||
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]]:
|
||||
"""Drain follow-up messages from the pending queue.
|
||||
@@ -1202,6 +1299,11 @@ class AgentLoop:
|
||||
message_metadata=metadata,
|
||||
),
|
||||
provider_state=provider_state,
|
||||
llm_usage_source=source_from_request(
|
||||
active_session_key,
|
||||
channel=channel,
|
||||
metadata=metadata,
|
||||
),
|
||||
))
|
||||
finally:
|
||||
turn_scope_stack.close()
|
||||
@@ -1233,18 +1335,33 @@ class AgentLoop:
|
||||
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
|
||||
|
||||
def _check_expired_sessions_if_due(self) -> None:
|
||||
"""Scan idle sessions no more often than the configured interval."""
|
||||
def _idle_compact_scan_due(self) -> bool:
|
||||
now = time.monotonic()
|
||||
if now < self._next_idle_compact_check_at:
|
||||
return
|
||||
return False
|
||||
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.schedule_background,
|
||||
self.runtime_for_session,
|
||||
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:
|
||||
"""Run the agent loop, dispatching messages as tasks to stay responsive to /stop."""
|
||||
self._running = True
|
||||
@@ -1255,7 +1372,7 @@ class AgentLoop:
|
||||
try:
|
||||
msg = await asyncio.wait_for(self.bus.consume_inbound(), timeout=1.0)
|
||||
except asyncio.TimeoutError:
|
||||
self._check_expired_sessions_if_due()
|
||||
await self._check_expired_sessions_if_due_async()
|
||||
continue
|
||||
except asyncio.CancelledError:
|
||||
# Preserve real task cancellation so shutdown can complete cleanly.
|
||||
@@ -1333,7 +1450,7 @@ class AgentLoop:
|
||||
)
|
||||
continue
|
||||
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)
|
||||
if followup_id is not None:
|
||||
pending_msg = dataclasses.replace(
|
||||
@@ -1343,7 +1460,7 @@ class AgentLoop:
|
||||
PENDING_FOLLOWUP_ID_KEY: followup_id,
|
||||
},
|
||||
)
|
||||
self.sessions.save(session)
|
||||
await self._save_session(session)
|
||||
try:
|
||||
self._pending_queues[effective_key].put_nowait(pending_msg)
|
||||
except asyncio.QueueFull:
|
||||
@@ -1454,10 +1571,10 @@ class AgentLoop:
|
||||
raise
|
||||
try:
|
||||
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):
|
||||
self._clear_pending_user_turn(session)
|
||||
self.sessions.save(session)
|
||||
await self._save_session(session)
|
||||
logger.info(
|
||||
"Restored partial context for cancelled session {}",
|
||||
key,
|
||||
@@ -1782,7 +1899,7 @@ class AgentLoop:
|
||||
if ctx.session is None:
|
||||
raise RuntimeError("required session is not active")
|
||||
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
|
||||
ctx.ephemeral = ctx.ephemeral or not session.policy.persist
|
||||
tools = ctx.tools or self.tools
|
||||
@@ -1813,16 +1930,16 @@ class AgentLoop:
|
||||
self.workspace_scopes.persist_message_scope(session, msg)
|
||||
|
||||
if self._restore_runtime_checkpoint(session):
|
||||
self.sessions.save(session)
|
||||
await self._save_session(session)
|
||||
if (
|
||||
RECOVERY_INBOUND_METADATA_KEY not in msg.metadata
|
||||
and restore_pending_interruption(session)
|
||||
):
|
||||
self.sessions.save(session)
|
||||
await self._save_session(session)
|
||||
|
||||
async def _compact_session(self, ctx: TurnContext) -> None:
|
||||
session = ctx.require_session()
|
||||
ctx.session, pending = self.auto_compact.prepare_session(
|
||||
ctx.session, pending = await self.auto_compact.prepare_session_async(
|
||||
session,
|
||||
ctx.session_key,
|
||||
)
|
||||
@@ -1859,14 +1976,14 @@ class AgentLoop:
|
||||
# them out of LLM context. /new is excluded because it
|
||||
# intentionally clears the session.
|
||||
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
|
||||
)
|
||||
session.add_message(
|
||||
"assistant", result.content, _command=True
|
||||
)
|
||||
self._clear_pending_user_turn(session)
|
||||
self.sessions.save(session)
|
||||
await self._save_session(session)
|
||||
if not ctx.ephemeral:
|
||||
await self.runtime_event_publisher.session_turn_persisted(
|
||||
ctx.msg,
|
||||
@@ -1881,7 +1998,7 @@ class AgentLoop:
|
||||
session = ctx.require_session()
|
||||
runtime = ctx.runtime
|
||||
if runtime is None:
|
||||
runtime = self.runtime_for_session(session)
|
||||
runtime = await self.runtime_for_session_async(session)
|
||||
ctx.runtime = runtime
|
||||
if ctx.session_key.startswith("dream:"):
|
||||
logger.info(
|
||||
@@ -1925,7 +2042,7 @@ class AgentLoop:
|
||||
# provider compatibility or prompt assembly work. A compatible
|
||||
# staged state replaces this in a second atomic save below.
|
||||
session.provider_state = None
|
||||
self.sessions.save(session)
|
||||
await self._save_session(session)
|
||||
ctx.input_persisted_early = True
|
||||
await ctx.delivery.runtime_admitted(runtime)
|
||||
|
||||
@@ -1979,7 +2096,7 @@ class AgentLoop:
|
||||
elif stored_state is not None:
|
||||
session.provider_state = None
|
||||
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,
|
||||
session,
|
||||
runtime_context_blocks=ctx.runtime_context_blocks,
|
||||
@@ -1989,7 +2106,7 @@ class AgentLoop:
|
||||
elif subagent_followup_persisted and staged_provider_state:
|
||||
# Upgrade the replay-safe baseline to the resumable state before
|
||||
# prompt assembly and the first model checkpoint.
|
||||
self.sessions.save(session)
|
||||
await self._save_session(session)
|
||||
ctx.initial_messages = self._build_initial_messages(ctx)
|
||||
|
||||
if ctx.on_progress is None:
|
||||
@@ -2031,7 +2148,7 @@ class AgentLoop:
|
||||
ctx.all_messages = all_msgs
|
||||
ctx.stop_reason = stop_reason
|
||||
ctx.had_injections = had_injections
|
||||
ctx.usage = dict(self._last_usage)
|
||||
ctx.usage = self._last_usage
|
||||
ctx.delivery.record_usage(ctx.usage)
|
||||
if ctx.kind is TurnKind.USER:
|
||||
await turn_continuation.maybe_continue_turn(ctx)
|
||||
@@ -2058,13 +2175,16 @@ class AgentLoop:
|
||||
else ctx.turn_wall_started_at
|
||||
)
|
||||
ctx.turn_latency_ms = max(0, int((time.time() - latency_started_at) * 1000))
|
||||
if ctx.usage and not ctx.ephemeral:
|
||||
session.metadata["_last_usage"] = dict(ctx.usage)
|
||||
if ctx.usage is not None and not ctx.ephemeral:
|
||||
session.metadata["_last_usage"] = ctx.usage.to_dict()
|
||||
self._save_turn(
|
||||
session, ctx.all_messages, ctx.save_skip,
|
||||
turn_latency_ms=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:
|
||||
self.schedule_background(
|
||||
self.consolidator.maybe_consolidate_by_tokens(
|
||||
@@ -2072,10 +2192,6 @@ class AgentLoop:
|
||||
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(
|
||||
ctx.msg,
|
||||
ctx.session_key,
|
||||
@@ -2288,11 +2404,24 @@ class AgentLoop:
|
||||
)
|
||||
return True
|
||||
|
||||
def _set_runtime_checkpoint(self, session: Session, payload: dict[str, Any]) -> None:
|
||||
"""Persist the latest in-flight turn state into session metadata."""
|
||||
def _set_runtime_checkpoint(
|
||||
self,
|
||||
session: Session,
|
||||
payload: dict[str, Any],
|
||||
) -> None:
|
||||
"""Synchronously persist a checkpoint for compatibility callers."""
|
||||
session.metadata[self._RUNTIME_CHECKPOINT_KEY] = payload
|
||||
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:
|
||||
session.metadata[self._PENDING_USER_TURN_KEY] = True
|
||||
|
||||
|
||||
+37
-18
@@ -20,7 +20,9 @@ from typing import TYPE_CHECKING, Any, Callable, Iterator, cast
|
||||
|
||||
from loguru import logger
|
||||
|
||||
from nanobot.llm_usage.context import llm_usage_source
|
||||
from nanobot.runtime_context import public_history_messages
|
||||
from nanobot.session.async_compat import call_session_manager
|
||||
from nanobot.session.manager import (
|
||||
MIN_COMPACTED_REPLAY_MESSAGES,
|
||||
Session,
|
||||
@@ -823,6 +825,22 @@ class Consolidator:
|
||||
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:
|
||||
"""Return the shared consolidation lock for one session."""
|
||||
return self._locks.setdefault(session_key, asyncio.Lock())
|
||||
@@ -858,13 +876,13 @@ class Consolidator:
|
||||
return []
|
||||
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)":
|
||||
session.metadata["_last_summary"] = {
|
||||
"text": summary,
|
||||
"last_active": session.updated_at.isoformat(),
|
||||
}
|
||||
self.sessions.save(session)
|
||||
await self._save_session(session)
|
||||
|
||||
def estimate_session_prompt_tokens(
|
||||
self,
|
||||
@@ -915,15 +933,16 @@ class Consolidator:
|
||||
if not messages:
|
||||
return None
|
||||
try:
|
||||
response = await runtime.provider.chat_with_retry(
|
||||
model=runtime.model,
|
||||
messages=request_messages,
|
||||
tools=request_tools,
|
||||
tool_choice="none",
|
||||
temperature=runtime.generation.temperature,
|
||||
max_tokens=runtime.generation.max_tokens,
|
||||
reasoning_effort=runtime.generation.reasoning_effort,
|
||||
)
|
||||
with llm_usage_source("dream"):
|
||||
response = await runtime.provider.chat_with_retry(
|
||||
model=runtime.model,
|
||||
messages=request_messages,
|
||||
tools=request_tools,
|
||||
tool_choice="none",
|
||||
temperature=runtime.generation.temperature,
|
||||
max_tokens=runtime.generation.max_tokens,
|
||||
reasoning_effort=runtime.generation.reasoning_effort,
|
||||
)
|
||||
except Exception:
|
||||
logger.warning("Consolidation provider call failed, raw-dumping to history")
|
||||
self.store.raw_archive(messages, session_key=session_key)
|
||||
@@ -1055,7 +1074,7 @@ class Consolidator:
|
||||
lock = self.get_lock(session.key)
|
||||
async with lock:
|
||||
# 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:
|
||||
session = fresh
|
||||
if not session.messages:
|
||||
@@ -1069,7 +1088,7 @@ class Consolidator:
|
||||
runtime=runtime,
|
||||
)
|
||||
if estimated <= 0:
|
||||
self._persist_last_summary(session, last_summary)
|
||||
await self._persist_last_summary(session, last_summary)
|
||||
return
|
||||
if estimated < budget:
|
||||
unconsolidated_count = len(session.messages) - session.last_consolidated
|
||||
@@ -1081,7 +1100,7 @@ class Consolidator:
|
||||
source,
|
||||
unconsolidated_count,
|
||||
)
|
||||
self._persist_last_summary(session, last_summary)
|
||||
await self._persist_last_summary(session, last_summary)
|
||||
return
|
||||
|
||||
for round_num in range(self._MAX_CONSOLIDATION_ROUNDS):
|
||||
@@ -1125,7 +1144,7 @@ class Consolidator:
|
||||
last_summary = summary
|
||||
session.last_consolidated = end_idx
|
||||
session.provider_state = None
|
||||
self.sessions.save(session)
|
||||
await self._save_session(session)
|
||||
if not summary:
|
||||
# LLM is degraded — stop hammering it this call;
|
||||
# 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
|
||||
# into the runtime context on the next prepare_session() call, aligning
|
||||
# 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(
|
||||
self,
|
||||
@@ -1166,7 +1185,7 @@ class Consolidator:
|
||||
lock = self.get_lock(session_key)
|
||||
async with lock:
|
||||
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
|
||||
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.
|
||||
session.last_consolidated = archive_end
|
||||
session.provider_state = None
|
||||
self.sessions.save(session)
|
||||
await self._save_session(session)
|
||||
|
||||
visible = session.get_history(
|
||||
max_messages=MIN_COMPACTED_REPLAY_MESSAGES,
|
||||
|
||||
@@ -210,12 +210,14 @@ class AgentProgressHook(AgentHook):
|
||||
tool_hint=False,
|
||||
tool_events=tool_events,
|
||||
)
|
||||
u = context.usage or {}
|
||||
u = context.usage
|
||||
logger.debug(
|
||||
"LLM usage: prompt={} completion={} cached={}",
|
||||
u.get("prompt_tokens", 0),
|
||||
u.get("completion_tokens", 0),
|
||||
u.get("cached_tokens", 0),
|
||||
"LLM usage: input={} output={} cache_read={} cache_write={} source={}",
|
||||
u.input_tokens if u else 0,
|
||||
u.output_tokens if u else 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:
|
||||
|
||||
+123
-102
@@ -6,7 +6,7 @@ import asyncio
|
||||
import inspect
|
||||
import os
|
||||
import time
|
||||
from collections.abc import Awaitable, Callable, Iterable
|
||||
from collections.abc import Awaitable, Callable, Iterable, Sized
|
||||
from copy import deepcopy
|
||||
from dataclasses import dataclass, field
|
||||
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.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 (
|
||||
LLMProvider,
|
||||
LLMResponse,
|
||||
LLMUsage,
|
||||
ProviderCallContext,
|
||||
ProviderConversationState,
|
||||
ToolCallRequest,
|
||||
@@ -76,6 +83,22 @@ _MAX_EMPTY_RETRIES = 2
|
||||
_MAX_LENGTH_RECOVERIES = 3
|
||||
_MAX_INJECTIONS_PER_TURN = 3
|
||||
_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:
|
||||
@@ -117,6 +140,7 @@ class AgentRunSpec:
|
||||
goal_continue_message: GoalContinueMessage | None = None
|
||||
finalize_on_max_iterations: bool = True
|
||||
provider_state: ProviderConversationState | None = None
|
||||
llm_usage_source: LLMUsageSource | None = None
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
@@ -126,7 +150,7 @@ class AgentRunResult:
|
||||
final_content: str | None
|
||||
messages: list[dict[str, Any]]
|
||||
tools_used: list[str] = field(default_factory=list)
|
||||
usage: dict[str, int] = field(default_factory=dict)
|
||||
usage: LLMUsage | None = None
|
||||
stop_reason: str = "completed"
|
||||
error: str | None = None
|
||||
tool_events: list[dict[str, str]] = field(default_factory=list)
|
||||
@@ -391,6 +415,9 @@ class AgentRunner:
|
||||
hook = spec.hook or AgentHook()
|
||||
messages = list(spec.initial_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:
|
||||
await hook.before_run(context)
|
||||
@@ -412,7 +439,7 @@ class AgentRunner:
|
||||
context.messages = deepcopy(result.messages)
|
||||
context.final_content = result.final_content
|
||||
context.tools_used = list(result.tools_used)
|
||||
context.usage = dict(result.usage)
|
||||
context.usage = result.usage
|
||||
context.stop_reason = result.stop_reason
|
||||
context.error = result.error
|
||||
context.tool_events = deepcopy(result.tool_events)
|
||||
@@ -423,17 +450,20 @@ class AgentRunner:
|
||||
await hook.after_run(context)
|
||||
return result
|
||||
finally:
|
||||
context.messages = deepcopy(messages)
|
||||
if context.exception is None:
|
||||
await hook.on_finally(context)
|
||||
else:
|
||||
try:
|
||||
try:
|
||||
context.messages = deepcopy(messages)
|
||||
if context.exception is None:
|
||||
await hook.on_finally(context)
|
||||
except Exception:
|
||||
logger.exception(
|
||||
"AgentHook.on_finally error after {}",
|
||||
context.stop_reason or "run exception",
|
||||
)
|
||||
else:
|
||||
try:
|
||||
await hook.on_finally(context)
|
||||
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(
|
||||
self,
|
||||
@@ -443,7 +473,7 @@ class AgentRunner:
|
||||
) -> AgentRunResult:
|
||||
final_content: str | None = None
|
||||
tools_used: list[str] = []
|
||||
usage = {"prompt_tokens": 0, "completion_tokens": 0}
|
||||
usage: LLMUsage | None = None
|
||||
error: str | None = None
|
||||
stop_reason = "completed"
|
||||
tool_events: list[dict[str, str]] = []
|
||||
@@ -519,8 +549,8 @@ class AgentRunner:
|
||||
)
|
||||
response.content = cleaned_content
|
||||
raw_usage = self._usage_or_estimate(spec, messages_for_model, response)
|
||||
context.usage = dict(raw_usage)
|
||||
self._accumulate_usage(usage, raw_usage)
|
||||
context.usage = raw_usage
|
||||
usage = self._merge_usage(usage, raw_usage)
|
||||
if reasoning_text and not context.streamed_reasoning:
|
||||
await hook.emit_reasoning(reasoning_text)
|
||||
await hook.emit_reasoning_end()
|
||||
@@ -683,10 +713,10 @@ class AgentRunner:
|
||||
conversation_state=conversation_state,
|
||||
)
|
||||
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)
|
||||
context.response = response
|
||||
context.usage = dict(raw_usage)
|
||||
context.usage = raw_usage
|
||||
context.tool_calls = list(response.tool_calls)
|
||||
original_content = response.content
|
||||
clean = hook.finalize_content(context, response.content)
|
||||
@@ -859,7 +889,7 @@ class AgentRunner:
|
||||
had_injections = True
|
||||
terminal_content = None
|
||||
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,
|
||||
hook,
|
||||
messages,
|
||||
@@ -922,18 +952,7 @@ class AgentRunner:
|
||||
conversation_state: ProviderConversationStateController,
|
||||
provider_context: ProviderCallContext | None = None,
|
||||
) -> LLMResponse:
|
||||
timeout_s: float | None = 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
|
||||
if timeout_s <= 0:
|
||||
timeout_s = None
|
||||
timeout_s = self._resolve_llm_timeout_s(spec)
|
||||
|
||||
kwargs = self._build_request_kwargs(
|
||||
spec,
|
||||
@@ -1247,9 +1266,9 @@ class AgentRunner:
|
||||
spec: AgentRunSpec,
|
||||
hook: AgentHook,
|
||||
messages: list[dict[str, Any]],
|
||||
usage: dict[str, int],
|
||||
usage: LLMUsage | None,
|
||||
conversation_state: ProviderConversationStateController,
|
||||
) -> str | None:
|
||||
) -> tuple[str | None, LLMUsage | None]:
|
||||
retry_messages = self._budget_exhausted_finalization_messages(messages)
|
||||
try:
|
||||
response = await self._request_no_tools(
|
||||
@@ -1264,10 +1283,10 @@ class AgentRunner:
|
||||
"Budget-exhausted finalization failed for {}; using fallback",
|
||||
spec.session_key or "default",
|
||||
)
|
||||
return None
|
||||
return None, usage
|
||||
|
||||
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:
|
||||
logger.warning(
|
||||
"Budget-exhausted finalization returned finish_reason='{}' "
|
||||
@@ -1276,19 +1295,19 @@ class AgentRunner:
|
||||
len(response.tool_calls),
|
||||
spec.session_key or "default",
|
||||
)
|
||||
return None
|
||||
return None, usage
|
||||
|
||||
context = AgentHookContext(
|
||||
iteration=spec.max_iterations,
|
||||
messages=messages,
|
||||
response=response,
|
||||
usage=dict(raw_usage),
|
||||
usage=raw_usage,
|
||||
session_key=spec.session_key,
|
||||
)
|
||||
clean = hook.finalize_content(context, response.content)
|
||||
if is_blank_text(clean):
|
||||
return None
|
||||
return clean
|
||||
return None, usage
|
||||
return clean, usage
|
||||
|
||||
async def _request_no_tools(
|
||||
self,
|
||||
@@ -1302,10 +1321,38 @@ class AgentRunner:
|
||||
messages,
|
||||
tools=None,
|
||||
)
|
||||
return await spec.runtime.provider.chat_with_retry(
|
||||
coro = spec.runtime.provider.chat_with_retry(
|
||||
**kwargs,
|
||||
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
|
||||
def _budget_exhausted_finalization_messages(
|
||||
@@ -1332,31 +1379,24 @@ class AgentRunner:
|
||||
spec: AgentRunSpec,
|
||||
messages: list[dict[str, Any]],
|
||||
response: LLMResponse,
|
||||
) -> dict[str, int]:
|
||||
usage = self._usage_dict(response.usage)
|
||||
total = self._usage_total(usage)
|
||||
if total > 0:
|
||||
usage["total_tokens"] = total
|
||||
usage.setdefault("provider_tokens", total)
|
||||
elif response.finish_reason == "error":
|
||||
return {}
|
||||
else:
|
||||
) -> LLMUsage | None:
|
||||
usage = response.usage
|
||||
if response.finish_reason == "error":
|
||||
if usage is None or usage.total_tokens == 0:
|
||||
usage = LLMUsage.empty_request()
|
||||
elif usage is None or usage.total_tokens == 0:
|
||||
usage = self._estimate_response_usage(spec, messages, response)
|
||||
completion = usage.get("completion_tokens", 0)
|
||||
if response.generation_ms is not None and completion > 0:
|
||||
usage["generation_ms"] = response.generation_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
|
||||
return usage.with_timing(
|
||||
generation_ms=response.generation_ms,
|
||||
ttft_ms=response.ttft_ms,
|
||||
)
|
||||
|
||||
def _estimate_response_usage(
|
||||
self,
|
||||
spec: AgentRunSpec,
|
||||
messages: list[dict[str, Any]],
|
||||
response: LLMResponse,
|
||||
) -> dict[str, int]:
|
||||
) -> LLMUsage:
|
||||
try:
|
||||
tools = spec.tools.get_definitions()
|
||||
except Exception:
|
||||
@@ -1374,52 +1414,21 @@ class AgentRunner:
|
||||
thinking_blocks=response.thinking_blocks,
|
||||
)
|
||||
completion_tokens = estimate_message_tokens(assistant_message)
|
||||
total_tokens = max(0, prompt_tokens) + max(0, completion_tokens)
|
||||
if total_tokens <= 0:
|
||||
return {}
|
||||
return {
|
||||
"prompt_tokens": max(0, prompt_tokens),
|
||||
"completion_tokens": max(0, completion_tokens),
|
||||
"total_tokens": total_tokens,
|
||||
"estimated_tokens": total_tokens,
|
||||
}
|
||||
return LLMUsage.estimated(
|
||||
input_tokens=max(0, prompt_tokens),
|
||||
output_tokens=max(0, completion_tokens),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _usage_dict(usage: dict[str, Any] | None) -> dict[str, int]:
|
||||
if not usage:
|
||||
return {}
|
||||
result: dict[str, int] = {}
|
||||
for key, value in usage.items():
|
||||
try:
|
||||
result[key] = int(value or 0)
|
||||
except (TypeError, ValueError):
|
||||
continue
|
||||
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
|
||||
def _merge_usage(
|
||||
left: LLMUsage | None,
|
||||
right: LLMUsage | None,
|
||||
) -> LLMUsage | None:
|
||||
if left is None:
|
||||
return right
|
||||
if right is None:
|
||||
return left
|
||||
return left + right
|
||||
|
||||
async def _execute_tools(
|
||||
self,
|
||||
@@ -1528,6 +1537,7 @@ class AgentRunner:
|
||||
RuntimeError(prep_error) if spec.fail_on_tool_error else None
|
||||
)
|
||||
await hook.before_execute_tool(context, tool_call, tool, params)
|
||||
tool_started_at = time.perf_counter()
|
||||
try:
|
||||
if tool is not None:
|
||||
result = await tool.execute(**params)
|
||||
@@ -1556,6 +1566,17 @@ class AgentRunner:
|
||||
if spec.fail_on_tool_error:
|
||||
return payload, event, exc
|
||||
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):
|
||||
await hook.on_execute_tool_error(context, tool_call, tool, params, result)
|
||||
|
||||
@@ -8,7 +8,7 @@ import warnings
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable, TypedDict
|
||||
from typing import Any, Callable, NotRequired, TypedDict
|
||||
|
||||
from loguru import logger
|
||||
|
||||
@@ -28,7 +28,8 @@ from nanobot.agent.tools.registry import ToolRegistry
|
||||
from nanobot.bus.events import InboundMessage
|
||||
from nanobot.bus.queue import MessageBus
|
||||
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 (
|
||||
WorkspaceScope,
|
||||
bind_workspace_scope,
|
||||
@@ -43,6 +44,7 @@ class _SubagentOrigin(TypedDict):
|
||||
channel: str
|
||||
chat_id: str
|
||||
session_key: str | None
|
||||
llm_usage_source: NotRequired[LLMUsageSource]
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
@@ -56,7 +58,7 @@ class SubagentStatus:
|
||||
phase: str = "initializing" # initializing | awaiting_tools | tools_completed | final_response | done | error
|
||||
iteration: int = 0
|
||||
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
|
||||
error: str | None = None
|
||||
|
||||
@@ -82,7 +84,7 @@ class _SubagentHook(AgentHook):
|
||||
return
|
||||
self._status.iteration = context.iteration
|
||||
self._status.tool_events = list(context.tool_events)
|
||||
self._status.usage = dict(context.usage)
|
||||
self._status.usage = context.usage
|
||||
if context.error:
|
||||
self._status.error = str(context.error)
|
||||
|
||||
@@ -252,6 +254,7 @@ class SubagentManager:
|
||||
"channel": origin_channel,
|
||||
"chat_id": origin_chat_id,
|
||||
"session_key": session_key,
|
||||
"llm_usage_source": current_llm_usage_source(),
|
||||
}
|
||||
|
||||
status = SubagentStatus(
|
||||
@@ -315,6 +318,7 @@ class SubagentManager:
|
||||
"channel": origin_channel,
|
||||
"chat_id": origin_chat_id,
|
||||
"session_key": session_key,
|
||||
"llm_usage_source": current_llm_usage_source(),
|
||||
}
|
||||
status = SubagentStatus(
|
||||
task_id=task_id,
|
||||
@@ -417,6 +421,10 @@ class SubagentManager:
|
||||
session_key=sess_key,
|
||||
workspace=root,
|
||||
llm_timeout_s=llm_timeout,
|
||||
llm_usage_source=origin.get(
|
||||
"llm_usage_source",
|
||||
current_llm_usage_source(),
|
||||
),
|
||||
))
|
||||
finally:
|
||||
if token is not None:
|
||||
|
||||
@@ -4,7 +4,10 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import inspect
|
||||
from pathlib import Path
|
||||
from typing import TypedDict
|
||||
|
||||
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
|
||||
|
||||
|
||||
class _CliAppRunKwargs(TypedDict):
|
||||
args: list[str]
|
||||
json_output: bool
|
||||
working_dir: str | None
|
||||
timeout: int | None
|
||||
restrict_to_workspace: bool
|
||||
|
||||
|
||||
class CliAppsToolConfig(Base):
|
||||
"""CLI Apps tool configuration."""
|
||||
|
||||
@@ -147,14 +158,17 @@ class CliAppsTool(Tool):
|
||||
)
|
||||
workspace = access.project_path or self.workspace
|
||||
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:
|
||||
return manager.run(
|
||||
name,
|
||||
args=args or [],
|
||||
json_output=bool(json),
|
||||
working_dir=working_dir,
|
||||
timeout=timeout,
|
||||
restrict_to_workspace=access.restrict_to_workspace,
|
||||
)
|
||||
run_async = inspect.getattr_static(type(manager), "run_async", None)
|
||||
if inspect.iscoroutinefunction(run_async):
|
||||
return await manager.run_async(name, **run_kwargs)
|
||||
return await asyncio.to_thread(manager.run, name, **run_kwargs)
|
||||
except CliAppError as exc:
|
||||
return ToolResult.error(f"Error: {exc.message}")
|
||||
|
||||
@@ -143,14 +143,41 @@ class CronTool(Tool):
|
||||
tz: str | None = None,
|
||||
at: 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:
|
||||
if action == "add":
|
||||
if self._in_cron_context.get():
|
||||
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)
|
||||
elif action == "list":
|
||||
if action == "list":
|
||||
return self._list_jobs()
|
||||
elif action == "remove":
|
||||
if action == "remove":
|
||||
return self._remove_job(job_id)
|
||||
return f"Unknown action: {action}"
|
||||
|
||||
|
||||
@@ -2,9 +2,11 @@
|
||||
|
||||
# pyright: reportPrivateUsage=false, reportUnusedFunction=false
|
||||
|
||||
import asyncio
|
||||
import difflib
|
||||
import mimetypes
|
||||
import os
|
||||
import threading
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
@@ -21,6 +23,7 @@ from nanobot.agent.tools.schema import (
|
||||
)
|
||||
from nanobot.config_base import Base
|
||||
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
|
||||
|
||||
|
||||
@@ -664,22 +667,31 @@ def _match_covers_line(match: _MatchSpan, line: int) -> bool:
|
||||
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] = []
|
||||
start = 0
|
||||
while True:
|
||||
idx = content.find(old_text, start)
|
||||
search_start = 0
|
||||
line_start = 0
|
||||
line = 1
|
||||
while max_matches is None or len(matches) < max_matches:
|
||||
idx = content.find(old_text, search_start)
|
||||
if idx == -1:
|
||||
break
|
||||
line += content.count("\n", line_start, idx)
|
||||
matches.append(
|
||||
_MatchSpan(
|
||||
start=idx,
|
||||
end=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
|
||||
|
||||
|
||||
@@ -735,27 +747,36 @@ def _find_quote_matches(content: str, old_text: str) -> list[_MatchSpan]:
|
||||
norm_content = _normalize_quotes(content)
|
||||
norm_old = _normalize_quotes(old_text)
|
||||
matches: list[_MatchSpan] = []
|
||||
start = 0
|
||||
search_start = 0
|
||||
line_start = 0
|
||||
line = 1
|
||||
while True:
|
||||
idx = norm_content.find(norm_old, start)
|
||||
idx = norm_content.find(norm_old, search_start)
|
||||
if idx == -1:
|
||||
break
|
||||
line += content.count("\n", line_start, idx)
|
||||
matches.append(
|
||||
_MatchSpan(
|
||||
start=idx,
|
||||
end=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
|
||||
|
||||
|
||||
def _find_matches(content: str, old_text: str) -> list[_MatchSpan]:
|
||||
"""Locate all matches using progressively looser strategies."""
|
||||
def _find_matches(
|
||||
content: str,
|
||||
old_text: str,
|
||||
*,
|
||||
max_exact_matches: int | None = None,
|
||||
) -> list[_MatchSpan]:
|
||||
"""Locate matches using progressively looser strategies."""
|
||||
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, normalize_quotes=True),
|
||||
lambda: _find_quote_matches(content, old_text),
|
||||
@@ -869,6 +890,43 @@ class EditFileTool(_FsTool):
|
||||
new_text: str | None = None,
|
||||
replace_all: bool = False, occurrence: int | None = None,
|
||||
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:
|
||||
try:
|
||||
if not path:
|
||||
@@ -892,9 +950,12 @@ class EditFileTool(_FsTool):
|
||||
# Create-file semantics: old_text='' + file doesn't exist → create
|
||||
if not file_exists:
|
||||
if old_text == "":
|
||||
fp.parent.mkdir(parents=True, exist_ok=True)
|
||||
fp.write_text(new_text, encoding="utf-8")
|
||||
self._file_states.record_write(fp)
|
||||
with commit_lock:
|
||||
if cancelled.is_set():
|
||||
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 self._file_not_found_msg(path, fp)
|
||||
|
||||
@@ -912,8 +973,11 @@ class EditFileTool(_FsTool):
|
||||
content = raw.decode("utf-8")
|
||||
if content.strip():
|
||||
return ToolResult.error(f"Error: Cannot create file — {path} already exists and is not empty.")
|
||||
fp.write_text(new_text, encoding="utf-8")
|
||||
self._file_states.record_write(fp)
|
||||
with commit_lock:
|
||||
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}"
|
||||
|
||||
# Read-before-edit check
|
||||
@@ -923,7 +987,11 @@ class EditFileTool(_FsTool):
|
||||
uses_crlf = b"\r\n" in raw
|
||||
content = raw.decode("utf-8").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:
|
||||
return self._not_found_msg(old_text, content, path)
|
||||
@@ -1000,8 +1068,11 @@ class EditFileTool(_FsTool):
|
||||
if uses_crlf:
|
||||
new_content = new_content.replace("\n", "\r\n")
|
||||
|
||||
fp.write_bytes(new_content.encode("utf-8"))
|
||||
self._file_states.record_write(fp)
|
||||
with commit_lock:
|
||||
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}"
|
||||
if warning:
|
||||
msg = f"{warning}\n{msg}"
|
||||
|
||||
@@ -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.bus.runtime_events import GoalStateChanged, RuntimeEventBus, RuntimeEventContext
|
||||
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 (
|
||||
GOAL_STATE_KEY,
|
||||
MAX_GOAL_OBJECTIVE_CHARS,
|
||||
@@ -28,6 +29,7 @@ from nanobot.session.goal_state import (
|
||||
sustained_goal_active,
|
||||
)
|
||||
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
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -60,36 +62,68 @@ class _GoalToolsMixin:
|
||||
self._sessions = sessions
|
||||
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()
|
||||
if request_ctx is None:
|
||||
return None
|
||||
key = request_ctx.session_key
|
||||
if not key:
|
||||
return None
|
||||
return self._sessions.get_or_create(key)
|
||||
return await self._get_or_create_session(key)
|
||||
|
||||
def _goal_mutation_allowed(self) -> bool:
|
||||
return current_request_context() is not None and goal_mutation_allowed()
|
||||
|
||||
def _save_goal_state(
|
||||
async def _save_goal_state(
|
||||
self,
|
||||
sess: Any,
|
||||
blob: dict[str, Any],
|
||||
*,
|
||||
reset_continuation: bool = False,
|
||||
revoke_permission: bool = False,
|
||||
) -> None:
|
||||
previous_metadata = deepcopy(sess.metadata)
|
||||
sess.metadata[GOAL_STATE_KEY] = blob
|
||||
discard_legacy_goal_state_key(sess.metadata)
|
||||
if reset_continuation:
|
||||
reset_goal_continuation_rounds(sess.metadata)
|
||||
saved = False
|
||||
|
||||
async def save_and_publish() -> None:
|
||||
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:
|
||||
self._sessions.save(sess)
|
||||
except BaseException:
|
||||
sess.metadata.clear()
|
||||
sess.metadata.update(previous_metadata)
|
||||
raise
|
||||
await shield_and_drain(save_and_publish())
|
||||
finally:
|
||||
# This ContextVar belongs to the caller task, not the settlement task.
|
||||
# Apply the post-save permission effect here even when cancellation was
|
||||
# delayed until the durable save and runtime notification completed.
|
||||
if revoke_permission and saved:
|
||||
revoke_goal_mutation_permission()
|
||||
|
||||
async def _publish_goal_state_changed(self, metadata: dict[str, Any]) -> None:
|
||||
runtime_events = self._runtime_events
|
||||
@@ -175,7 +209,7 @@ class CreateGoalTool(Tool, _GoalToolsMixin):
|
||||
) -> RuntimeContextBlock | None:
|
||||
if not request.session_key:
|
||||
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_active = sustained_goal_active(session.metadata)
|
||||
if not goal_start_requested and not goal_active:
|
||||
@@ -197,7 +231,7 @@ class CreateGoalTool(Tool, _GoalToolsMixin):
|
||||
ui_summary: str | None = None,
|
||||
**kwargs: Any,
|
||||
) -> str:
|
||||
sess = self._session()
|
||||
sess = await self._session()
|
||||
if sess is None:
|
||||
return ToolResult.error(
|
||||
"Error: create_goal requires an active chat session (missing routing context)."
|
||||
@@ -225,8 +259,7 @@ class CreateGoalTool(Tool, _GoalToolsMixin):
|
||||
"ui_summary": summary,
|
||||
"started_at": _iso_now(),
|
||||
}
|
||||
self._save_goal_state(sess, blob, reset_continuation=True)
|
||||
await self._publish_goal_state_changed(sess.metadata)
|
||||
await self._save_goal_state(sess, blob, reset_continuation=True)
|
||||
extra = f"\nSummary line: {summary}" if summary else ""
|
||||
return (
|
||||
"Goal recorded. Keep working toward the objective using ordinary tools. "
|
||||
@@ -305,7 +338,7 @@ class UpdateGoalTool(Tool, _GoalToolsMixin):
|
||||
ui_summary: str | None = None,
|
||||
**kwargs: Any,
|
||||
) -> str:
|
||||
sess = self._session()
|
||||
sess = await self._session()
|
||||
if sess is None:
|
||||
return ToolResult.error("Error: update_goal requires an active chat session.")
|
||||
prior = parse_goal_state(goal_state_raw(sess.metadata))
|
||||
@@ -340,8 +373,7 @@ class UpdateGoalTool(Tool, _GoalToolsMixin):
|
||||
"previous_objective": str(prior.get("objective") or ""),
|
||||
"recap": (recap or "").strip(),
|
||||
}
|
||||
self._save_goal_state(sess, blob, reset_continuation=True)
|
||||
await self._publish_goal_state_changed(sess.metadata)
|
||||
await self._save_goal_state(sess, blob, reset_continuation=True)
|
||||
extra = f"\nSummary line: {summary}" if summary else ""
|
||||
return "Goal replaced. Continue toward the new objective using ordinary tools." + extra
|
||||
|
||||
@@ -359,9 +391,7 @@ class UpdateGoalTool(Tool, _GoalToolsMixin):
|
||||
}
|
||||
if normalized == "complete":
|
||||
blob["completed_at"] = ended
|
||||
self._save_goal_state(sess, blob)
|
||||
revoke_goal_mutation_permission()
|
||||
await self._publish_goal_state_changed(sess.metadata)
|
||||
await self._save_goal_state(sess, blob, revoke_permission=True)
|
||||
|
||||
tail = (recap or "").strip()
|
||||
label = {
|
||||
|
||||
@@ -20,10 +20,10 @@ from nanobot.agent.tools.base import Tool, ToolResult
|
||||
from nanobot.agent.tools.registry import ToolRegistry
|
||||
from nanobot.security.network import (
|
||||
PinnedDNSAsyncTransport,
|
||||
async_resolve_url_target,
|
||||
async_validate_url_target,
|
||||
env_proxy_applies_to_url,
|
||||
httpx_env_proxy_mounts,
|
||||
resolve_url_target,
|
||||
validate_url_target,
|
||||
)
|
||||
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
|
||||
if not port:
|
||||
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:
|
||||
return False
|
||||
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:
|
||||
"""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:
|
||||
raise httpx.RequestError(
|
||||
f"Blocked unsafe MCP URL {_redact_url(str(request.url))} ({error})",
|
||||
@@ -1031,7 +1031,7 @@ async def connect_mcp_servers(
|
||||
return False
|
||||
|
||||
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:
|
||||
logger.warning(
|
||||
"MCP server '{}': blocked unsafe URL {} ({})",
|
||||
|
||||
@@ -12,6 +12,7 @@ if TYPE_CHECKING:
|
||||
from nanobot.agent.tools.shell import ExecToolConfig
|
||||
from nanobot.agent.tools.web import WebToolsConfig
|
||||
from nanobot.config.schema import ModelPresetConfig
|
||||
from nanobot.providers.base import LLMUsage
|
||||
from nanobot.utils.llm_runtime import LLMRuntime
|
||||
|
||||
|
||||
@@ -65,7 +66,7 @@ class RuntimeSnapshot:
|
||||
web_config: dict[str, object]
|
||||
exec_config: dict[str, object]
|
||||
subagent_statuses: dict[str, dict[str, object]]
|
||||
last_usage: dict[str, int]
|
||||
last_usage: Mapping[str, JsonScalar]
|
||||
scratchpad: dict[str, JsonValue]
|
||||
|
||||
def as_mapping(self) -> Mapping[str, object]:
|
||||
@@ -106,6 +107,13 @@ class RuntimeControl(Protocol):
|
||||
session_key: str | None,
|
||||
) -> 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_context_window_tokens(self, value: int) -> LLMRuntime: ...
|
||||
@@ -151,7 +159,7 @@ class _RuntimeControlTarget(Protocol):
|
||||
def tool_names(self) -> list[str]: ...
|
||||
|
||||
@property
|
||||
def last_usage(self) -> Mapping[str, int]: ...
|
||||
def last_usage(self) -> LLMUsage | None: ...
|
||||
|
||||
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: ...
|
||||
|
||||
async def set_session_model_preset_async(
|
||||
self,
|
||||
session_key: str,
|
||||
name: str,
|
||||
) -> LLMRuntime: ...
|
||||
|
||||
|
||||
class AgentRuntimeControl:
|
||||
"""Allowlisted adapter from agent-loop state to ``RuntimeControl``."""
|
||||
@@ -190,7 +204,7 @@ class AgentRuntimeControl:
|
||||
web_config=_snapshot_web_config(target.web_config),
|
||||
exec_config=_snapshot_exec_config(target.exec_config),
|
||||
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),
|
||||
)
|
||||
|
||||
@@ -207,6 +221,16 @@ class AgentRuntimeControl:
|
||||
return self.__target.set_session_model_preset(session_key, 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:
|
||||
self.__target.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,
|
||||
"iteration": status.iteration,
|
||||
"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,
|
||||
"error": status.error,
|
||||
}
|
||||
|
||||
+103
-53
@@ -4,9 +4,12 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import fnmatch
|
||||
import os
|
||||
import re
|
||||
import threading
|
||||
import time
|
||||
from contextlib import suppress
|
||||
from pathlib import Path, PurePosixPath
|
||||
from typing import Any, Iterable, TypeVar
|
||||
@@ -125,6 +128,8 @@ class _SearchTool(_FsTool):
|
||||
class FindFilesTool(_SearchTool):
|
||||
"""Find files by path fragment, glob, or type."""
|
||||
_scopes = {"core", "subagent"}
|
||||
_MAX_SCAN_PATHS = 500_000
|
||||
_MAX_SCAN_SECONDS = 30.0
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
@@ -218,66 +223,111 @@ class FindFilesTool(_SearchTool):
|
||||
offset: int = 0,
|
||||
**kwargs: Any,
|
||||
) -> str:
|
||||
cancelled = threading.Event()
|
||||
try:
|
||||
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
|
||||
return await asyncio.to_thread(
|
||||
self._execute_sync,
|
||||
path=path,
|
||||
query=query,
|
||||
glob=glob,
|
||||
file_type=type,
|
||||
include_dirs=include_dirs,
|
||||
sort=sort,
|
||||
head_limit=head_limit,
|
||||
offset=offset,
|
||||
cancelled=cancelled,
|
||||
)
|
||||
root = target if target.is_dir() else target.parent
|
||||
matches: list[tuple[str, float]] = []
|
||||
|
||||
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 asyncio.CancelledError:
|
||||
cancelled.set()
|
||||
raise
|
||||
except PermissionError as e:
|
||||
return ToolResult.error(f"Error: {e}")
|
||||
except Exception as e:
|
||||
return ToolResult.error(f"Error finding files: {e}")
|
||||
|
||||
def _execute_sync(
|
||||
self,
|
||||
*,
|
||||
path: str,
|
||||
query: str | None,
|
||||
glob: str | None,
|
||||
file_type: str | None,
|
||||
include_dirs: bool,
|
||||
sort: str,
|
||||
head_limit: int | None,
|
||||
offset: int,
|
||||
cancelled: threading.Event,
|
||||
) -> str:
|
||||
target = self._resolve(path or ".")
|
||||
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):
|
||||
"""Search file contents using a regex-like pattern."""
|
||||
|
||||
@@ -150,7 +150,7 @@ class MyTool(Tool):
|
||||
"Actions: check, set.\n"
|
||||
"- check (no key): full config overview — start here.\n"
|
||||
"- 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. "
|
||||
"Scratchpad keys persist across turns but not restarts.\n"
|
||||
"Key values: _current_iteration (current progress), "
|
||||
@@ -370,7 +370,7 @@ class MyTool(Tool):
|
||||
if not self._modify_allowed:
|
||||
return ToolResult.error("Error: set is disabled (tools.my.allow_set is false)")
|
||||
if action in ("modify", "set"):
|
||||
return self._modify(key, value)
|
||||
return await self._modify_async(key, value)
|
||||
return f"Unknown action: {action}"
|
||||
|
||||
# -- inspect --
|
||||
@@ -492,6 +492,11 @@ class MyTool(Tool):
|
||||
return ToolResult.error(f"Error: '{key}' is read-only and cannot be modified")
|
||||
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:
|
||||
if not isinstance(value, str) or not value.strip():
|
||||
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}"
|
||||
)
|
||||
|
||||
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:
|
||||
spec = self.RESTRICTED[key]
|
||||
expected = cast(type[Any], spec["type"])
|
||||
|
||||
@@ -267,6 +267,7 @@ class ExecTool(Tool):
|
||||
|
||||
_MAX_TIMEOUT = 600
|
||||
_MAX_OUTPUT = 10_000
|
||||
_PREPARE_TIMEOUT_SECONDS = 6.0
|
||||
|
||||
# Kernel device files safe as stdio redirect targets (#3599).
|
||||
_BENIGN_DEVICE_PATHS: frozenset[str] = frozenset({
|
||||
@@ -324,7 +325,20 @@ class ExecTool(Tool):
|
||||
if max_output_chars is None:
|
||||
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):
|
||||
return prepared
|
||||
|
||||
@@ -470,14 +484,18 @@ class ExecTool(Tool):
|
||||
+ _WORKSPACE_BOUNDARY_NOTE
|
||||
)
|
||||
|
||||
guard_error = self._guard_command(
|
||||
command,
|
||||
cwd,
|
||||
restrict_to_workspace=access.restrict_to_workspace,
|
||||
workspace_root=workspace_root,
|
||||
)
|
||||
if guard_error:
|
||||
return guard_error
|
||||
# Full access is an explicit trust decision. Keep the application-level
|
||||
# command guard aligned with the selected access mode instead of
|
||||
# continuing to block commands after workspace restriction is disabled.
|
||||
if access.restrict_to_workspace:
|
||||
guard_error = self._guard_command(
|
||||
command,
|
||||
cwd,
|
||||
restrict_to_workspace=True,
|
||||
workspace_root=workspace_root,
|
||||
)
|
||||
if guard_error:
|
||||
return guard_error
|
||||
|
||||
if self.sandbox:
|
||||
if _IS_WINDOWS:
|
||||
@@ -912,6 +930,15 @@ class ExecTool(Tool):
|
||||
if self._is_benign_device_path(expanded):
|
||||
continue
|
||||
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
|
||||
|
||||
if self._is_benign_device_path(str(p)):
|
||||
|
||||
+12
-14
@@ -96,7 +96,7 @@ def _normalize(text: str) -> 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:
|
||||
p = urlparse(url)
|
||||
if p.scheme not in ('http', 'https'):
|
||||
@@ -108,18 +108,16 @@ def _validate_url(url: str) -> tuple[bool, str]:
|
||||
return False, str(e)
|
||||
|
||||
|
||||
def _validate_url_safe(url: str) -> tuple[bool, str]:
|
||||
"""Validate URL with SSRF protection: scheme, domain, and resolved IP check."""
|
||||
from nanobot.security.network import validate_url_target
|
||||
async def _async_validate_url_safe(url: str) -> tuple[bool, str]:
|
||||
from nanobot.security.network import async_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, ...]]:
|
||||
"""Validate URL and return the resolved IPs to pin during the request."""
|
||||
from nanobot.security.network import resolve_url_target
|
||||
async def _async_resolve_url_safe(url: str) -> tuple[bool, str, tuple[str, ...]]:
|
||||
from nanobot.security.network import async_resolve_url_target
|
||||
|
||||
return resolve_url_target(url)
|
||||
return await async_resolve_url_target(url)
|
||||
|
||||
|
||||
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."""
|
||||
current_url = url
|
||||
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:
|
||||
return None, f"Redirect blocked: {error_msg}"
|
||||
|
||||
@@ -229,7 +227,7 @@ async def _get_with_safe_redirects(
|
||||
return response, None
|
||||
|
||||
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:
|
||||
await response.aclose()
|
||||
return None, f"Redirect blocked: {error_msg}"
|
||||
@@ -249,7 +247,7 @@ async def _stream_with_safe_redirects(
|
||||
current_url = url
|
||||
chain_carries_credentials = _url_carries_credentials(url)
|
||||
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:
|
||||
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 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:
|
||||
await stream.__aexit__(None, None, None)
|
||||
return None, None, f"Redirect blocked: {error_msg}", chain_carries_credentials
|
||||
@@ -1106,7 +1104,7 @@ class WebFetchTool(Tool):
|
||||
url = url.strip(" \t\r\n`\"'")
|
||||
extract_mode = kwargs.pop("extractMode", extract_mode)
|
||||
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:
|
||||
return json.dumps({"error": f"URL validation failed: {error_msg}", "url": url}, ensure_ascii=False)
|
||||
|
||||
|
||||
@@ -4,7 +4,7 @@ from __future__ import annotations
|
||||
|
||||
import dataclasses
|
||||
import time
|
||||
from collections.abc import Awaitable, Callable, Mapping
|
||||
from collections.abc import Awaitable, Callable
|
||||
from dataclasses import dataclass, field
|
||||
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.queue import MessageBus
|
||||
from nanobot.bus.runtime_events import RuntimeEventBus, RuntimeEventPublisher
|
||||
from nanobot.providers.base import LLMUsage
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from nanobot.utils.llm_runtime import LLMRuntime
|
||||
@@ -203,7 +204,7 @@ class TurnDelivery:
|
||||
def record_latency(self, latency_ms: int | None) -> None:
|
||||
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)
|
||||
|
||||
def background_response(
|
||||
|
||||
@@ -18,6 +18,7 @@ from aiohttp import web
|
||||
from loguru import logger
|
||||
|
||||
from nanobot.config.paths import get_media_dir
|
||||
from nanobot.providers.base import LLMUsage
|
||||
from nanobot.utils.helpers import safe_filename
|
||||
from nanobot.utils.media_decode import (
|
||||
MAX_FILE_SIZE,
|
||||
@@ -93,11 +94,11 @@ def _error_json(status: int, message: str, err_type: str = "invalid_request_erro
|
||||
def _chat_completion_response(
|
||||
content: str,
|
||||
model: str,
|
||||
usage: dict[str, int] | None = None,
|
||||
usage: LLMUsage | None = None,
|
||||
) -> dict[str, Any]:
|
||||
prompt = (usage or {}).get("prompt_tokens", 0)
|
||||
completion = (usage or {}).get("completion_tokens", 0)
|
||||
total = (usage or {}).get("total_tokens", 0) or prompt + completion
|
||||
prompt = usage.input_tokens if usage else 0
|
||||
completion = usage.output_tokens if usage else 0
|
||||
total = usage.total_tokens if usage else 0
|
||||
return {
|
||||
"id": f"chatcmpl-{uuid.uuid4().hex[:12]}",
|
||||
"object": "chat.completion",
|
||||
|
||||
+374
-41
@@ -2,15 +2,20 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import ctypes
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import shlex
|
||||
import shutil
|
||||
import signal
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from collections.abc import Iterable
|
||||
from contextlib import suppress
|
||||
from ctypes import wintypes
|
||||
from dataclasses import dataclass
|
||||
from importlib import metadata as importlib_metadata
|
||||
from pathlib import Path
|
||||
@@ -97,6 +102,141 @@ class CliAppsRuntimeConfig:
|
||||
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]] = {
|
||||
"1password-cli": ("1password", "#3B66BC"),
|
||||
"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)})")
|
||||
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 ``."
|
||||
)
|
||||
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(
|
||||
self,
|
||||
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,
|
||||
restrict_to_workspace: bool = False,
|
||||
) -> str:
|
||||
app = self.get_app(name)
|
||||
installed = self._load_installed()
|
||||
if str(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[str(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))
|
||||
artifact_snapshot = self._artifact_snapshot(cwd)
|
||||
prepared = 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:
|
||||
result = subprocess.run(
|
||||
[resolved, *clean_args],
|
||||
cwd=str(cwd),
|
||||
capture_output=True,
|
||||
process = subprocess.Popen(
|
||||
[prepared.resolved, *prepared.args],
|
||||
cwd=str(prepared.cwd),
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
text=True,
|
||||
encoding="utf-8",
|
||||
errors="replace",
|
||||
timeout=effective_timeout,
|
||||
env=self._subprocess_env(),
|
||||
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, stderr = process.communicate(timeout=prepared.timeout)
|
||||
except subprocess.TimeoutExpired:
|
||||
return f"CLI app '{name}' timed out after {effective_timeout}s"
|
||||
output = [
|
||||
f"CLI app '{name}' exited {result.returncode}.",
|
||||
f"Command: {entry} {' '.join(shlex.quote(arg) for arg in clean_args)}".rstrip(),
|
||||
]
|
||||
if result.stdout:
|
||||
output.append("\nSTDOUT:\n" + result.stdout.rstrip())
|
||||
if result.stderr:
|
||||
output.append("\nSTDERR:\n" + result.stderr.rstrip())
|
||||
artifacts = self._changed_artifacts(cwd, artifact_snapshot)
|
||||
if artifacts:
|
||||
output.append(
|
||||
"\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 ``."
|
||||
)
|
||||
return _truncate("\n".join(output))
|
||||
self._terminate_run_process_sync(process, job)
|
||||
return f"CLI app '{prepared.name}' timed out after {prepared.timeout}s"
|
||||
except BaseException:
|
||||
self._terminate_run_process_sync(process, job)
|
||||
raise
|
||||
if job is not None:
|
||||
job.close(kill_descendants=False)
|
||||
return self._format_run_result(
|
||||
prepared,
|
||||
returncode=process.returncode,
|
||||
stdout=stdout,
|
||||
stderr=stderr,
|
||||
)
|
||||
|
||||
@@ -12,6 +12,7 @@ from dataclasses import dataclass, replace
|
||||
from typing import Any, cast
|
||||
|
||||
from nanobot.bus.events import OutboundMessage
|
||||
from nanobot.providers.base import LLMUsage
|
||||
|
||||
|
||||
class OutboundEvent:
|
||||
@@ -58,7 +59,7 @@ class StreamedResponseEvent(OutboundEvent):
|
||||
class TurnEndEvent(OutboundEvent):
|
||||
latency_ms: int | 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
|
||||
|
||||
|
||||
@@ -197,11 +198,6 @@ def _legacy_event_from_metadata(msg: OutboundMessage) -> OutboundEvent | None:
|
||||
return TurnEndEvent(
|
||||
latency_ms=_metadata_int(meta, "latency_ms"),
|
||||
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"),
|
||||
)
|
||||
if meta.get("_session_updated"):
|
||||
|
||||
@@ -10,13 +10,14 @@ from __future__ import annotations
|
||||
import asyncio
|
||||
import contextlib
|
||||
import inspect
|
||||
from collections.abc import Awaitable, Callable, Mapping
|
||||
from collections.abc import Awaitable, Callable
|
||||
from dataclasses import dataclass, field
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from loguru import logger
|
||||
|
||||
from nanobot.bus.events import InboundMessage
|
||||
from nanobot.providers.base import LLMUsage
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from nanobot.utils.llm_runtime import LLMRuntime
|
||||
@@ -72,7 +73,7 @@ class TurnCompleted:
|
||||
context: RuntimeEventContext
|
||||
latency_ms: int | None = None
|
||||
runtime: LLMRuntime | None = None
|
||||
usage: dict[str, int] = field(default_factory=dict)
|
||||
usage: LLMUsage | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
@@ -180,7 +181,7 @@ class RuntimeEventPublisher:
|
||||
self.bus = bus or RuntimeEventBus()
|
||||
self._turn_latency_ms: dict[str, int] = {}
|
||||
self._turn_runtime: dict[str, LLMRuntime] = {}
|
||||
self._turn_usage: dict[str, dict[str, int]] = {}
|
||||
self._turn_usage: dict[str, LLMUsage] = {}
|
||||
|
||||
@staticmethod
|
||||
def _context(
|
||||
@@ -206,12 +207,9 @@ class RuntimeEventPublisher:
|
||||
if latency_ms is not None:
|
||||
self._turn_latency_ms[session_key] = int(latency_ms)
|
||||
|
||||
def record_turn_usage(self, session_key: str, usage: Mapping[str, int]) -> None:
|
||||
self._turn_usage[session_key] = {
|
||||
key: int(value)
|
||||
for key, value in usage.items()
|
||||
if type(value) is int and value >= 0
|
||||
}
|
||||
def record_turn_usage(self, session_key: str, usage: LLMUsage | None) -> None:
|
||||
if usage is not None:
|
||||
self._turn_usage[session_key] = usage
|
||||
|
||||
def clear_turn(self, session_key: str) -> 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),
|
||||
runtime=self._turn_runtime.pop(session_key, None),
|
||||
usage=self._turn_usage.pop(session_key, {}),
|
||||
usage=self._turn_usage.pop(session_key, None),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@@ -21,7 +21,10 @@ from nanobot.bus.events import OutboundMessage
|
||||
from nanobot.bus.queue import MessageBus
|
||||
from nanobot.channels.base import BaseChannel
|
||||
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_REDIRECTS = 3
|
||||
@@ -417,8 +420,8 @@ class DingTalkChannel(BaseChannel):
|
||||
return self._zip_bytes(filename, data)
|
||||
return data, filename, content_type
|
||||
|
||||
def _validate_remote_media_url(self, media_ref: str) -> bool:
|
||||
ok, err = validate_url_target(media_ref)
|
||||
async def _validate_remote_media_url(self, media_ref: str) -> bool:
|
||||
ok, err = await async_validate_url_target(media_ref)
|
||||
if not ok:
|
||||
self.logger.warning("remote media URL blocked ref={} reason={}", media_ref, err)
|
||||
return False
|
||||
@@ -434,7 +437,11 @@ class DingTalkChannel(BaseChannel):
|
||||
allowed_hosts = {host.lower() for host in self.config.remote_media_redirect_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:
|
||||
self.logger.warning("media download redirect refused ref={}", current_url)
|
||||
return None
|
||||
@@ -449,7 +456,7 @@ class DingTalkChannel(BaseChannel):
|
||||
next_url,
|
||||
)
|
||||
return None
|
||||
if not self._validate_remote_media_url(next_url):
|
||||
if not await self._validate_remote_media_url(next_url):
|
||||
return None
|
||||
return next_url
|
||||
|
||||
@@ -461,7 +468,7 @@ class DingTalkChannel(BaseChannel):
|
||||
if not self._http:
|
||||
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
|
||||
|
||||
try:
|
||||
@@ -473,7 +480,7 @@ class DingTalkChannel(BaseChannel):
|
||||
current_url = media_ref
|
||||
for _ in range(DINGTALK_MAX_REMOTE_MEDIA_REDIRECTS + 1):
|
||||
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:
|
||||
self.logger.warning(
|
||||
"remote media redirect blocked ref={} final={} reason={}",
|
||||
@@ -483,7 +490,7 @@ class DingTalkChannel(BaseChannel):
|
||||
)
|
||||
return None, None
|
||||
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")
|
||||
)
|
||||
if not next_url:
|
||||
@@ -516,7 +523,9 @@ class DingTalkChannel(BaseChannel):
|
||||
current_url = media_ref
|
||||
for _ in range(DINGTALK_MAX_REMOTE_MEDIA_REDIRECTS + 1):
|
||||
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:
|
||||
self.logger.warning(
|
||||
"remote media redirect blocked ref={} final={} reason={}",
|
||||
@@ -526,7 +535,7 @@ class DingTalkChannel(BaseChannel):
|
||||
)
|
||||
return None, None
|
||||
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")
|
||||
)
|
||||
if not next_url:
|
||||
|
||||
@@ -24,7 +24,7 @@ from nanobot.bus.queue import MessageBus
|
||||
from nanobot.channels.base import BaseChannel
|
||||
from nanobot.config.paths import get_media_dir
|
||||
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
|
||||
|
||||
_DOWNLOAD_TIMEOUT = aiohttp.ClientTimeout(total=60)
|
||||
@@ -473,7 +473,7 @@ class NapcatChannel(BaseChannel):
|
||||
if not ref:
|
||||
return None
|
||||
if ref.startswith(("http://", "https://")):
|
||||
ok, err = validate_url_target(ref)
|
||||
ok, err = await async_validate_url_target(ref)
|
||||
if not ok:
|
||||
logger.warning("napcat: rejected remote image '{}': {}", ref, err)
|
||||
return None
|
||||
@@ -525,7 +525,7 @@ class NapcatChannel(BaseChannel):
|
||||
# logger.debug("napcat: downloading image from {}", url)
|
||||
if self._http is None:
|
||||
return None
|
||||
ok, err = validate_url_target(url)
|
||||
ok, err = await async_validate_url_target(url)
|
||||
if not ok:
|
||||
logger.warning("napcat: skip image '{}': {}", url, err)
|
||||
return None
|
||||
|
||||
@@ -149,9 +149,13 @@ async def test_download_image_rejects_redirects(tmp_path, monkeypatch) -> None:
|
||||
channel = _channel()
|
||||
channel._media_root = tmp_path
|
||||
channel._http = _FakeHttp(_FakeResponse(status=302))
|
||||
|
||||
async def allow_url(_url: str) -> tuple[bool, str]:
|
||||
return True, ""
|
||||
|
||||
monkeypatch.setattr(
|
||||
"nanobot.channels.napcat.runtime.validate_url_target",
|
||||
lambda _url: (True, ""),
|
||||
"nanobot.channels.napcat.runtime.async_validate_url_target",
|
||||
allow_url,
|
||||
)
|
||||
|
||||
result = await channel._download_image({"url": "https://example.com/a.png", "file": "a.png"})
|
||||
|
||||
@@ -40,7 +40,7 @@ from nanobot.bus.events import OutboundMessage
|
||||
from nanobot.bus.queue import MessageBus
|
||||
from nanobot.channels.base import BaseChannel
|
||||
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
|
||||
|
||||
try:
|
||||
@@ -458,7 +458,7 @@ class QQChannel(BaseChannel):
|
||||
return None, None
|
||||
|
||||
# Remote URL
|
||||
ok, err = validate_url_target(media_ref)
|
||||
ok, err = await async_validate_url_target(media_ref)
|
||||
if not ok:
|
||||
self.logger.warning("outbound media URL validation failed url={} err={}", media_ref, err)
|
||||
return None, None
|
||||
|
||||
@@ -23,8 +23,8 @@ from nanobot.config.schema import Base
|
||||
from nanobot.pairing import is_approved
|
||||
from nanobot.security.network import (
|
||||
PinnedDNSAsyncTransport,
|
||||
async_validate_url_target,
|
||||
httpx_env_proxy_mounts,
|
||||
validate_url_target,
|
||||
)
|
||||
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:
|
||||
"""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:
|
||||
raise httpx.RequestError(f"unsafe Slack file URL: {error}", request=request)
|
||||
|
||||
|
||||
@@ -859,13 +859,13 @@ def _patch_download_validation(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
validated: list[str],
|
||||
) -> None:
|
||||
def validate(url: str) -> tuple[bool, str]:
|
||||
async def validate(url: str) -> tuple[bool, str]:
|
||||
validated.append(url)
|
||||
if "169.254.169.254" in url:
|
||||
return False, "blocked metadata address"
|
||||
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
|
||||
|
||||
@@ -36,7 +36,7 @@ from nanobot.channels.base import BaseChannel
|
||||
from nanobot.command.builtin import build_help_text
|
||||
from nanobot.config.paths import get_media_dir
|
||||
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.logging_bridge import redirect_lib_logging
|
||||
|
||||
@@ -956,7 +956,7 @@ class TelegramChannel(BaseChannel):
|
||||
|
||||
# Telegram Bot API accepts HTTP(S) URLs directly for media params.
|
||||
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:
|
||||
raise ValueError(f"unsafe media URL: {error}")
|
||||
await self._call_with_retry(
|
||||
|
||||
@@ -1488,7 +1488,14 @@ async def test_send_remote_media_url_after_security_validation(monkeypatch) -> N
|
||||
MessageBus(),
|
||||
)
|
||||
_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(
|
||||
OutboundMessage(
|
||||
@@ -1546,9 +1553,13 @@ async def test_send_blocks_unsafe_remote_media_url(monkeypatch) -> None:
|
||||
MessageBus(),
|
||||
)
|
||||
_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(
|
||||
"nanobot.channels.telegram.runtime.validate_url_target",
|
||||
lambda url: (False, "Blocked: example.com resolves to private/internal address 127.0.0.1"),
|
||||
"nanobot.channels.telegram.runtime.async_validate_url_target",
|
||||
deny_url,
|
||||
)
|
||||
|
||||
await channel.send(
|
||||
|
||||
@@ -44,6 +44,7 @@ from nanobot.bus.queue import MessageBus
|
||||
from nanobot.channels.base import BaseChannel
|
||||
from nanobot.command.builtin import USER_SHELL_COMMAND, builtin_command_starts_agent_turn
|
||||
from nanobot.config.schema import Base
|
||||
from nanobot.providers.base import LLMUsage
|
||||
from nanobot.runtime_context import (
|
||||
RUNTIME_CONTEXT_INPUT_META,
|
||||
WEBUI_QUOTE_METADATA,
|
||||
@@ -54,6 +55,7 @@ from nanobot.security.workspace_access import (
|
||||
WORKSPACE_SCOPE_METADATA_KEY,
|
||||
WorkspaceScopeError,
|
||||
)
|
||||
from nanobot.session.async_compat import call_session_manager
|
||||
from nanobot.session.goal_state import goal_state_ws_blob
|
||||
from nanobot.session.model_selection import model_preset_from_metadata
|
||||
from nanobot.session.recovery import recovery_state_from_metadata
|
||||
@@ -446,6 +448,26 @@ class WebSocketChannel(BaseChannel):
|
||||
if sessions is None:
|
||||
return {}
|
||||
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
|
||||
metadata = cast(dict[str, object], raw_metadata) if isinstance(raw_metadata, dict) else None
|
||||
fields: dict[str, Any] = {}
|
||||
@@ -458,18 +480,9 @@ class WebSocketChannel(BaseChannel):
|
||||
recovery_state = recovery_state_from_metadata(metadata)
|
||||
if recovery_state is not None:
|
||||
fields["recovery_state"] = recovery_state
|
||||
usage = metadata.get("_last_usage")
|
||||
if isinstance(usage, dict):
|
||||
sanitized_usage: dict[str, int | float] = {}
|
||||
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
|
||||
usage = LLMUsage.from_dict(metadata.get("_last_usage"))
|
||||
if usage is not None:
|
||||
fields["usage"] = usage.to_turn_dict()
|
||||
return fields
|
||||
|
||||
def _detach(self, connection: ServerConnection, chat_id: str) -> None:
|
||||
@@ -518,13 +531,16 @@ class WebSocketChannel(BaseChannel):
|
||||
fork_key: str,
|
||||
) -> None:
|
||||
"""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)
|
||||
await self._send_event(
|
||||
connection,
|
||||
"attached",
|
||||
chat_id=fork_id,
|
||||
**self._attached_model_fields(fork_id),
|
||||
**await self._attached_model_fields_async(fork_id),
|
||||
)
|
||||
await self._send_event(
|
||||
connection,
|
||||
@@ -906,13 +922,13 @@ class WebSocketChannel(BaseChannel):
|
||||
)
|
||||
if scope is None:
|
||||
return
|
||||
self._workspaces.persist_scope(new_id, scope)
|
||||
self._workspaces.stage_scope(new_id, scope)
|
||||
self._attach(connection, new_id)
|
||||
await self._send_event(
|
||||
connection,
|
||||
"attached",
|
||||
chat_id=new_id,
|
||||
**self._attached_model_fields(new_id),
|
||||
**await self._attached_model_fields_async(new_id),
|
||||
)
|
||||
await self._send_event(
|
||||
connection,
|
||||
@@ -968,7 +984,7 @@ class WebSocketChannel(BaseChannel):
|
||||
connection,
|
||||
"attached",
|
||||
chat_id=cid,
|
||||
**self._attached_model_fields(cid),
|
||||
**await self._attached_model_fields_async(cid),
|
||||
)
|
||||
await self._hydrate_after_subscribe(cid)
|
||||
return
|
||||
@@ -1023,7 +1039,7 @@ class WebSocketChannel(BaseChannel):
|
||||
)
|
||||
if scope is None:
|
||||
return
|
||||
self._workspaces.persist_scope(cid, scope)
|
||||
self._workspaces.stage_scope(cid, scope)
|
||||
# Other clients on the same gateway only need an invalidation; they
|
||||
# can reload the authoritative session row without receiving a
|
||||
# local project path that belongs to another connection.
|
||||
@@ -1217,7 +1233,6 @@ class WebSocketChannel(BaseChannel):
|
||||
if session_mentions:
|
||||
metadata["session_mentions"] = session_mentions
|
||||
metadata[WORKSPACE_SCOPE_METADATA_KEY] = scope.metadata()
|
||||
self._workspaces.persist_scope(cid, scope)
|
||||
is_webui = metadata.get("webui") is True
|
||||
queued_owner = None
|
||||
if is_webui and not is_user_shell and builtin_command_starts_agent_turn(content):
|
||||
@@ -1272,6 +1287,7 @@ class WebSocketChannel(BaseChannel):
|
||||
else False
|
||||
),
|
||||
)
|
||||
await asyncio.to_thread(self._workspaces.persist_scope, cid, scope)
|
||||
accepted = True
|
||||
finally:
|
||||
if not accepted and queued_owner is not None:
|
||||
@@ -1561,7 +1577,7 @@ class WebSocketChannel(BaseChannel):
|
||||
turn_id: str | None = None,
|
||||
) -> Any | None:
|
||||
try:
|
||||
return resolver()
|
||||
return await asyncio.to_thread(resolver)
|
||||
except WorkspaceScopeError as exc:
|
||||
await self._send_event(
|
||||
connection,
|
||||
@@ -2019,7 +2035,7 @@ class WebSocketChannel(BaseChannel):
|
||||
latency_ms: int | 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,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
turn_owner: str | None = None,
|
||||
@@ -2034,8 +2050,8 @@ class WebSocketChannel(BaseChannel):
|
||||
body["latency_ms"] = int(latency_ms)
|
||||
if goal_state is not None:
|
||||
body["goal_state"] = goal_state
|
||||
if usage:
|
||||
body["usage"] = usage
|
||||
if usage is not None:
|
||||
body["usage"] = usage.to_turn_dict()
|
||||
if context_window_tokens is not None:
|
||||
body["context_window_tokens"] = int(context_window_tokens)
|
||||
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.schema import Config, ModelPresetConfig
|
||||
from nanobot.providers.base import LLMUsage
|
||||
from nanobot.runtime_context import RUNTIME_CONTEXT_INPUT_META, WEBUI_QUOTE_SOURCE
|
||||
from nanobot.security.workspace_access import WORKSPACE_SCOPE_METADATA_KEY
|
||||
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(
|
||||
conn,
|
||||
"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
|
||||
async def test_workspace_scope_change_invalidates_other_attached_clients(
|
||||
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()
|
||||
|
||||
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])
|
||||
assert payload["event"] == "session_updated"
|
||||
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")
|
||||
assert saved["metadata"]["workspace_scope"] == {
|
||||
"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()
|
||||
await message_task
|
||||
|
||||
saved = sessions.read_session_file(f"websocket:{chat_id}")
|
||||
assert saved["metadata"]["workspace_scope"]["access_mode"] == "restricted"
|
||||
assert sessions.read_session_file(f"websocket:{chat_id}") is None
|
||||
assert channel._workspaces.scope_for_session_key(
|
||||
f"websocket:{chat_id}"
|
||||
).access_mode == "restricted"
|
||||
payload = json.loads(message_conn.send.await_args.args[0])
|
||||
assert payload["event"] == "error"
|
||||
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"]["sandbox_status"]["restrict_to_workspace"] is False
|
||||
assert payload["workspace_scope"]["sandbox_status"]["workspace_root"] == str(project.resolve())
|
||||
saved = sessions.read_session_file("websocket:chat-native")
|
||||
assert saved["metadata"]["workspace_scope"] == {
|
||||
assert sessions.read_session_file("websocket:chat-native") is None
|
||||
assert channel._workspaces.scope_for_session_key(
|
||||
"websocket:chat-native"
|
||||
).metadata() == {
|
||||
"project_path": str(project.resolve()),
|
||||
"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:
|
||||
usage = LLMUsage.reported(input_tokens=120, output_tokens=8, total_tokens=175)
|
||||
manager = MagicMock()
|
||||
manager.read_session_metadata.return_value = {
|
||||
"metadata": {
|
||||
SESSION_MODEL_PRESET_METADATA_KEY: "Deep Research",
|
||||
"_last_usage": {
|
||||
"prompt_tokens": 120,
|
||||
"completion_tokens": 8,
|
||||
"negative": -1,
|
||||
"boolean": True,
|
||||
},
|
||||
"_last_usage": usage.to_dict(),
|
||||
}
|
||||
}
|
||||
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") == {
|
||||
"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))
|
||||
mock_ws = AsyncMock()
|
||||
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(
|
||||
channel="websocket",
|
||||
@@ -3232,7 +3331,7 @@ async def test_send_turn_end_includes_latency_ms_when_present() -> None:
|
||||
content="",
|
||||
event=TurnEndEvent(
|
||||
latency_ms=1500,
|
||||
usage={"prompt_tokens": 80, "completion_tokens": 20, "cached_tokens": 40},
|
||||
usage=usage,
|
||||
context_window_tokens=128_000,
|
||||
),
|
||||
))
|
||||
@@ -3242,7 +3341,19 @@ async def test_send_turn_end_includes_latency_ms_when_present() -> None:
|
||||
"event": "turn_end",
|
||||
"chat_id": "chat-1",
|
||||
"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,
|
||||
},
|
||||
{"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
|
||||
|
||||
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(
|
||||
key="websocket:context-route",
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
metadata={"_last_usage": {"prompt_tokens": 12, "completion_tokens": 3}},
|
||||
metadata={"_last_usage": usage.to_dict()},
|
||||
)
|
||||
manager = MagicMock()
|
||||
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
|
||||
body = json.loads(response.body.decode())
|
||||
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)
|
||||
|
||||
|
||||
|
||||
@@ -4,6 +4,7 @@ import asyncio
|
||||
import json
|
||||
import random
|
||||
import socket
|
||||
import threading
|
||||
import time
|
||||
from contextlib import suppress
|
||||
from pathlib import Path
|
||||
@@ -2576,6 +2577,69 @@ async def test_webui_automations_route_lists_all_jobs_and_allows_user_actions(
|
||||
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
|
||||
async def test_webui_automations_route_manages_local_triggers(
|
||||
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")
|
||||
resp = channel.gateway.http._handle_bootstrap(_LOCAL, _NO_HEADERS)
|
||||
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
|
||||
|
||||
@@ -34,6 +34,7 @@ from nanobot.config.paths import is_default_workspace
|
||||
from nanobot.config.schema import Config
|
||||
from nanobot.gateway.runtime import GatewayInstance
|
||||
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.utils.evaluator import evaluate_response, resolve_evaluator_prompt
|
||||
from nanobot.utils.helpers import sync_workspace_templates
|
||||
@@ -45,6 +46,30 @@ __all__ = ["_run_gateway"]
|
||||
|
||||
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:
|
||||
"""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.session_turns import is_bound_cron_job
|
||||
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 (
|
||||
ProviderSnapshot,
|
||||
build_provider_snapshot,
|
||||
@@ -330,7 +357,6 @@ def _run_gateway(
|
||||
)
|
||||
from nanobot.triggers.local_runner import run_local_trigger_queue
|
||||
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
|
||||
webui_url = _webui_browser_url(config)
|
||||
@@ -361,7 +387,8 @@ def _run_gateway(
|
||||
runtime_events = RuntimeEventBus()
|
||||
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):
|
||||
snapshot.provider.set_fallback_model_observer(fallback_model_observer)
|
||||
return snapshot
|
||||
@@ -371,20 +398,19 @@ def _run_gateway(
|
||||
**kwargs: Any,
|
||||
) -> ProviderSnapshot:
|
||||
try:
|
||||
return _observe_fallback_models(load_provider_snapshot(*args, **kwargs))
|
||||
return _observe_provider(load_provider_snapshot(*args, **kwargs))
|
||||
except ValueError as exc:
|
||||
if unconfigured_provider_error is None:
|
||||
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:
|
||||
provider_snapshot = build_unconfigured_provider_snapshot(
|
||||
config,
|
||||
unconfigured_provider_error,
|
||||
provider_snapshot = _observe_provider(
|
||||
build_unconfigured_provider_snapshot(config, unconfigured_provider_error)
|
||||
)
|
||||
else:
|
||||
try:
|
||||
provider_snapshot = _observe_fallback_models(build_provider_snapshot(config))
|
||||
provider_snapshot = _observe_provider(build_provider_snapshot(config))
|
||||
except ValueError as exc:
|
||||
console.print(f"[red]Error: {exc}[/red]")
|
||||
raise typer.Exit(1) from exc
|
||||
@@ -443,7 +469,6 @@ def _run_gateway(
|
||||
runtime_events=runtime_events,
|
||||
turn_delivery_factory=turn_delivery_factory,
|
||||
provider_signature=provider_snapshot.signature,
|
||||
hooks=[TokenUsageHook(timezone_name=config.agents.defaults.timezone)],
|
||||
local_trigger_store=trigger_store,
|
||||
hook_factories=[create_file_edit_activity_hook],
|
||||
tool_registry=tools,
|
||||
@@ -493,12 +518,22 @@ def _run_gateway(
|
||||
and hasattr(session_manager, "save")
|
||||
):
|
||||
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}
|
||||
if msg.media:
|
||||
extra["media"] = list(msg.media)
|
||||
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)
|
||||
|
||||
message_tool = agent.tools.get("message")
|
||||
@@ -564,25 +599,18 @@ def _run_gateway(
|
||||
except Exception:
|
||||
logger.exception("Dream cron job failed")
|
||||
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)
|
||||
if sha:
|
||||
logger.info("Dream commit: {}", sha)
|
||||
store.compact_history()
|
||||
prune_dream_sessions(agent.sessions)
|
||||
await asyncio.to_thread(prune_dream_sessions, agent.sessions)
|
||||
return None
|
||||
|
||||
# Heartbeat is a system job that checks HEARTBEAT.md for active tasks.
|
||||
if job.name == "heartbeat":
|
||||
heartbeat_file = config.workspace_path / "HEARTBEAT.md"
|
||||
try:
|
||||
content = heartbeat_file.read_text(encoding="utf-8")
|
||||
content = await asyncio.to_thread(heartbeat_file.read_text, encoding="utf-8")
|
||||
except OSError:
|
||||
logger.debug("Heartbeat: HEARTBEAT.md missing")
|
||||
return None
|
||||
@@ -590,7 +618,7 @@ def _run_gateway(
|
||||
logger.debug("Heartbeat: HEARTBEAT.md has no active tasks")
|
||||
return None
|
||||
|
||||
channel, chat_id = _pick_heartbeat_target()
|
||||
channel, chat_id = await _pick_heartbeat_target()
|
||||
if channel == "cli":
|
||||
return None
|
||||
|
||||
@@ -618,9 +646,19 @@ def _run_gateway(
|
||||
message_tool.reset_suppress_delivery(suppress_token)
|
||||
|
||||
# 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)
|
||||
agent.sessions.save(session)
|
||||
await _call_session_manager(
|
||||
agent.sessions,
|
||||
"save_async",
|
||||
agent.sessions.save,
|
||||
session,
|
||||
)
|
||||
|
||||
if not resp or not resp.content:
|
||||
return
|
||||
@@ -630,14 +668,15 @@ def _run_gateway(
|
||||
evaluator_prompt = resolve_evaluator_prompt(config.workspace_path)
|
||||
|
||||
# Fail closed: stay silent on evaluator failure instead of notifying.
|
||||
should_notify = await evaluate_response(
|
||||
response=response,
|
||||
task_context=prompt,
|
||||
provider=agent.provider,
|
||||
model=agent.model,
|
||||
evaluator_prompt=evaluator_prompt,
|
||||
default_notify=False,
|
||||
)
|
||||
with llm_usage_source("cron"):
|
||||
should_notify = await evaluate_response(
|
||||
response=response,
|
||||
task_context=prompt,
|
||||
provider=agent.provider,
|
||||
model=agent.model,
|
||||
evaluator_prompt=evaluator_prompt,
|
||||
default_notify=False,
|
||||
)
|
||||
|
||||
if should_notify:
|
||||
logger.info("Heartbeat: completed, delivering response")
|
||||
@@ -696,17 +735,27 @@ def _run_gateway(
|
||||
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."""
|
||||
sidebar_state = read_webui_sidebar_state()
|
||||
sidebar_state = await asyncio.to_thread(read_webui_sidebar_state)
|
||||
unified_metadata = None
|
||||
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):
|
||||
unified_metadata = record["metadata"]
|
||||
sessions = await _call_session_manager(
|
||||
session_manager,
|
||||
"list_sessions_async",
|
||||
session_manager.list_sessions,
|
||||
)
|
||||
return _pick_heartbeat_target_from_sessions(
|
||||
enabled_channels=channels.enabled_channels,
|
||||
sessions=session_manager.list_sessions(),
|
||||
sessions=sessions,
|
||||
archived_keys=sidebar_state.get("archived_keys", []),
|
||||
unified_session_metadata=unified_metadata,
|
||||
)
|
||||
@@ -913,6 +962,10 @@ def _run_gateway(
|
||||
_monitor_local_clients(),
|
||||
name="nanobot-gateway-client-monitor",
|
||||
),
|
||||
asyncio.create_task(
|
||||
_monitor_event_loop_lag(),
|
||||
name="nanobot-event-loop-lag-monitor",
|
||||
),
|
||||
]
|
||||
if health_server_enabled:
|
||||
tasks.append(asyncio.create_task(
|
||||
@@ -980,7 +1033,11 @@ def _run_gateway(
|
||||
# Flush all cached sessions to durable storage before exit.
|
||||
# This prevents data loss on filesystems with write-back
|
||||
# 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:
|
||||
logger.info("Shutdown: flushed {} session(s) to disk", flushed)
|
||||
finally:
|
||||
|
||||
+72
-25
@@ -3,6 +3,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import inspect
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
@@ -14,6 +15,8 @@ from typing import TYPE_CHECKING, Any, Literal, cast
|
||||
from nanobot import __version__
|
||||
from nanobot.bus.events import INBOUND_META_USER_SHELL, OutboundMessage
|
||||
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.restart import set_restart_notice_to_env
|
||||
from nanobot.utils.workspace_prompts import initialize_workspace_prompt
|
||||
@@ -22,6 +25,7 @@ if TYPE_CHECKING:
|
||||
from nanobot.agent.loop import AgentLoop
|
||||
from nanobot.session.manager import Session
|
||||
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:
|
||||
# - 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())
|
||||
|
||||
|
||||
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:
|
||||
"""Cancel all active tasks and subagents for the session."""
|
||||
loop = ctx.loop
|
||||
@@ -257,8 +307,8 @@ async def cmd_restart(ctx: CommandContext) -> OutboundMessage:
|
||||
async def cmd_status(ctx: CommandContext) -> OutboundMessage:
|
||||
"""Build an outbound status message for a session."""
|
||||
loop = ctx.loop
|
||||
session = ctx.session or loop.sessions.get_or_create(ctx.key)
|
||||
runtime = ctx.runtime or loop.runtime_for_session(session)
|
||||
session = ctx.session or await _get_or_create_session(loop, ctx.key)
|
||||
runtime = ctx.runtime or await _runtime_for_session(loop, session)
|
||||
ctx_est = 0
|
||||
with suppress(Exception):
|
||||
ctx_est, _ = loop.consolidator.estimate_session_prompt_tokens(
|
||||
@@ -266,7 +316,8 @@ async def cmd_status(ctx: CommandContext) -> OutboundMessage:
|
||||
runtime=runtime,
|
||||
)
|
||||
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)
|
||||
search_usage_text: str | None = None
|
||||
@@ -305,29 +356,32 @@ async def cmd_new(ctx: CommandContext) -> OutboundMessage:
|
||||
loop = ctx.loop
|
||||
await loop._cancel_active_tasks(ctx.key) # pyright: ignore[reportPrivateUsage]
|
||||
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)
|
||||
archive_snapshot = None
|
||||
runtime = None
|
||||
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(
|
||||
session,
|
||||
messages=snapshot,
|
||||
metadata=dict(session.metadata),
|
||||
provider_state=None,
|
||||
)
|
||||
session.clear()
|
||||
loop.sessions.save(session)
|
||||
loop.sessions.invalidate(session.key)
|
||||
if archive_snapshot is not None and runtime is not None:
|
||||
loop.schedule_background(
|
||||
loop.consolidator.archive_session( # pyright: ignore[reportUnknownMemberType]
|
||||
archive_snapshot,
|
||||
archive_end=len(snapshot),
|
||||
runtime=runtime,
|
||||
async def reset_and_schedule_archive() -> None:
|
||||
session.clear()
|
||||
await _save_session(loop, session)
|
||||
loop.sessions.invalidate(session.key)
|
||||
if archive_snapshot is not None and runtime is not None:
|
||||
loop.schedule_background(
|
||||
loop.consolidator.archive_session( # pyright: ignore[reportUnknownMemberType]
|
||||
archive_snapshot,
|
||||
archive_end=len(snapshot),
|
||||
runtime=runtime,
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
await shield_and_drain(reset_and_schedule_archive())
|
||||
return OutboundMessage(
|
||||
channel=ctx.msg.channel, chat_id=ctx.msg.chat_id,
|
||||
content="New session started.",
|
||||
@@ -376,7 +430,7 @@ async def cmd_model(ctx: CommandContext) -> OutboundMessage:
|
||||
metadata = {**dict(ctx.msg.metadata or {}), "render_as": "text"}
|
||||
|
||||
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(
|
||||
channel=ctx.msg.channel,
|
||||
chat_id=ctx.msg.chat_id,
|
||||
@@ -386,7 +440,7 @@ async def cmd_model(ctx: CommandContext) -> OutboundMessage:
|
||||
|
||||
name = args
|
||||
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:
|
||||
names = _model_preset_names(loop)
|
||||
return OutboundMessage(
|
||||
@@ -478,13 +532,6 @@ async def cmd_dream(ctx: CommandContext) -> OutboundMessage:
|
||||
elapsed = time.monotonic() - t0
|
||||
content = f"Dream failed after {elapsed:.1f}s: {e}"
|
||||
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():
|
||||
commit_msg = build_dream_commit_message("dream: manual run", diff_body)
|
||||
sha = store.git.auto_commit(commit_msg)
|
||||
@@ -854,7 +901,7 @@ async def cmd_history(ctx: CommandContext) -> OutboundMessage:
|
||||
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)
|
||||
visible = [_format_history_message(m) for m in history]
|
||||
visible = [m for m in visible if m is not None]
|
||||
|
||||
@@ -8,6 +8,7 @@ import time
|
||||
import uuid
|
||||
from typing import TYPE_CHECKING, Any, Protocol
|
||||
|
||||
from nanobot.agent.automation_turns import AutomationTurnAcceptedCancellation
|
||||
from nanobot.agent.tools.cron import CronTool
|
||||
from nanobot.bus.events import InboundMessage, OutboundMessage
|
||||
from nanobot.cron.session_delivery import origin_delivery_context
|
||||
@@ -127,6 +128,15 @@ async def run_bound_cron_job(
|
||||
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:
|
||||
error_text = str(exc) or exc.__class__.__name__
|
||||
cron.write_run_record(
|
||||
|
||||
+230
-87
@@ -11,11 +11,12 @@ from dataclasses import asdict
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from types import EllipsisType
|
||||
from typing import Any, Callable, Coroutine, Literal
|
||||
from typing import Any, Callable, Coroutine, Literal, TypeVar
|
||||
|
||||
from filelock import FileLock
|
||||
from loguru import logger
|
||||
|
||||
from nanobot.agent.automation_turns import AutomationTurnAcceptedCancellation
|
||||
from nanobot.cron.session_turns import is_bound_cron_job
|
||||
from nanobot.cron.types import (
|
||||
CronJob,
|
||||
@@ -25,10 +26,14 @@ from nanobot.cron.types import (
|
||||
CronSchedule,
|
||||
CronStore,
|
||||
)
|
||||
from nanobot.utils.cancellation import shield_and_drain
|
||||
from nanobot.utils.run_records import (
|
||||
write_run_record as write_automation_run_record,
|
||||
)
|
||||
|
||||
_FILE_LOCK_TIMEOUT_SECONDS = 5
|
||||
_T = TypeVar("_T")
|
||||
|
||||
|
||||
class CronJobSkippedError(Exception):
|
||||
"""Raised by cron callbacks when a job was intentionally skipped."""
|
||||
@@ -164,10 +169,16 @@ class CronService:
|
||||
self.store_path = store_path
|
||||
self._action_path = store_path.parent / "action.jsonl"
|
||||
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._store: CronStore | 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._active_executions = 0
|
||||
self._store_dirty = False
|
||||
@@ -451,25 +462,58 @@ class CronService:
|
||||
"""Write an internal audit record for one cron execution."""
|
||||
write_automation_run_record(self._run_records_dir, run_id, record)
|
||||
|
||||
async def start(self) -> None:
|
||||
"""Start the cron service."""
|
||||
self._running = True
|
||||
loaded = 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 [].
|
||||
self._running = False
|
||||
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."
|
||||
async def run_sync(
|
||||
self,
|
||||
operation: Callable[..., _T],
|
||||
/,
|
||||
*args: Any,
|
||||
**kwargs: Any,
|
||||
) -> _T:
|
||||
"""Serialize a complete cron transaction in a worker thread.
|
||||
|
||||
A running thread cannot be cancelled safely. Keep the transaction lock
|
||||
until it exits so cancellation is never reported while that worker can
|
||||
still mutate cron state behind a later operation.
|
||||
"""
|
||||
async with self._operation_lock:
|
||||
return await shield_and_drain(
|
||||
asyncio.to_thread(operation, *args, **kwargs)
|
||||
)
|
||||
self._recompute_next_runs()
|
||||
self._save_store()
|
||||
self._arm_timer()
|
||||
logger.info("Cron service started with {} jobs", len(self._store.jobs if self._store else []))
|
||||
|
||||
async def start(self) -> None:
|
||||
"""Start the cron service and settle accepted work before cancellation."""
|
||||
|
||||
async def settle_start() -> None:
|
||||
self._event_loop = asyncio.get_running_loop()
|
||||
self._running = True
|
||||
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:
|
||||
"""Stop the cron service."""
|
||||
@@ -497,8 +541,22 @@ class CronService:
|
||||
if j.enabled and j.state.next_run_at_ms]
|
||||
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:
|
||||
"""Schedule the next timer tick."""
|
||||
"""Schedule the next timer tick on the owning event loop."""
|
||||
if self._timer_task:
|
||||
self._timer_task.cancel()
|
||||
|
||||
@@ -520,7 +578,7 @@ class CronService:
|
||||
self._timer_task = asyncio.create_task(tick())
|
||||
|
||||
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
|
||||
self._active_executions += 1
|
||||
try:
|
||||
@@ -528,11 +586,17 @@ class CronService:
|
||||
# to persist their advanced schedule. Persist that exact snapshot
|
||||
# before reloading or executing anything else; otherwise the older
|
||||
# disk state can replay the same job.
|
||||
if self._store_dirty:
|
||||
self._save_store()
|
||||
return
|
||||
async with self._operation_lock:
|
||||
if self._store_dirty:
|
||||
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``
|
||||
# may still hold the previous, known-good in-memory snapshot.
|
||||
if store is None:
|
||||
@@ -547,7 +611,8 @@ class CronService:
|
||||
for job in due_jobs:
|
||||
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:
|
||||
# A load/persist failure must not kill the scheduler: keep 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.
|
||||
self._arm_timer()
|
||||
|
||||
async def _execute_job(self, job: CronJob) -> None:
|
||||
"""Execute a single job."""
|
||||
async def _claim_job(self, job_id: str) -> bool:
|
||||
"""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()
|
||||
logger.info("Cron: executing job '{}' ({})", job.name, job.id)
|
||||
status: Literal["ok", "error", "skipped"]
|
||||
error: str | None
|
||||
accepted_cancellation: AutomationTurnAcceptedCancellation | None = None
|
||||
|
||||
try:
|
||||
if self.on_job:
|
||||
await self.on_job(job)
|
||||
try:
|
||||
if self.on_job:
|
||||
await self.on_job(job)
|
||||
status = "ok"
|
||||
error = None
|
||||
logger.info("Cron: job '{}' completed", job.name)
|
||||
except AutomationTurnAcceptedCancellation as exc:
|
||||
# The agent owns this turn now. Advance and persist the schedule
|
||||
# before allowing shutdown cancellation to unwind the timer.
|
||||
status = "ok"
|
||||
error = None
|
||||
accepted_cancellation = exc
|
||||
logger.info("Cron: job '{}' was accepted before cancellation", job.name)
|
||||
except CronJobSkippedError as exc:
|
||||
status = "skipped"
|
||||
error = str(exc) or None
|
||||
logger.warning("Cron: job '{}' skipped: {}", job.name, error or "")
|
||||
except asyncio.CancelledError as exc:
|
||||
current = asyncio.current_task()
|
||||
if current is not None and current.cancelling():
|
||||
raise
|
||||
status = "error"
|
||||
error = str(exc) or exc.__class__.__name__
|
||||
logger.exception("Cron: job '{}' was cancelled", job.name)
|
||||
except Exception as exc:
|
||||
status = "error"
|
||||
error = str(exc)
|
||||
logger.exception("Cron: job '{}' failed", job.name)
|
||||
|
||||
job.state.last_status = "ok"
|
||||
job.state.last_error = None
|
||||
logger.info("Cron: job '{}' completed", job.name)
|
||||
|
||||
except CronJobSkippedError as e:
|
||||
job.state.last_status = "skipped"
|
||||
job.state.last_error = str(e) or None
|
||||
logger.warning("Cron: job '{}' skipped: {}", job.name, job.state.last_error or "")
|
||||
except asyncio.CancelledError as e:
|
||||
current = asyncio.current_task()
|
||||
if current is not None and current.cancelling():
|
||||
raise
|
||||
job.state.last_status = "error"
|
||||
job.state.last_error = str(e) or e.__class__.__name__
|
||||
logger.exception("Cron: job '{}' was cancelled", job.name)
|
||||
except Exception as e:
|
||||
job.state.last_status = "error"
|
||||
job.state.last_error = str(e)
|
||||
logger.exception("Cron: job '{}' failed", job.name)
|
||||
|
||||
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())
|
||||
settlement = asyncio.create_task(
|
||||
self._settle_job_execution(
|
||||
job,
|
||||
start_ms=start_ms,
|
||||
status=status,
|
||||
error=error,
|
||||
persist=accepted_cancellation is not None,
|
||||
)
|
||||
)
|
||||
if accepted_cancellation is not None:
|
||||
await self._drain_settlement_on_cancellation(settlement)
|
||||
raise accepted_cancellation
|
||||
await settlement
|
||||
return True
|
||||
finally:
|
||||
await self._release_job_claim(job.id)
|
||||
|
||||
def _append_action(
|
||||
self,
|
||||
@@ -697,7 +828,7 @@ class CronService:
|
||||
store = self._require_store()
|
||||
store.jobs.append(job)
|
||||
self._save_store()
|
||||
self._arm_timer()
|
||||
self._request_timer_rearm()
|
||||
else:
|
||||
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.append(job)
|
||||
self._save_store()
|
||||
self._arm_timer()
|
||||
self._request_timer_rearm()
|
||||
logger.info("Cron: registered system job '{}' ({})", job.name, job.id)
|
||||
return job
|
||||
|
||||
@@ -726,7 +857,7 @@ class CronService:
|
||||
removed = len(store.jobs) < before
|
||||
if removed:
|
||||
self._save_store()
|
||||
self._arm_timer()
|
||||
self._request_timer_rearm()
|
||||
logger.info("Cron: removed system job {}", job_id)
|
||||
return removed
|
||||
|
||||
@@ -747,7 +878,7 @@ class CronService:
|
||||
if removed:
|
||||
if self._should_persist_store():
|
||||
self._save_store()
|
||||
self._arm_timer()
|
||||
self._request_timer_rearm()
|
||||
else:
|
||||
self._append_action("del", {"job_id": job_id})
|
||||
logger.info("Cron: removed job {}", job_id)
|
||||
@@ -769,7 +900,7 @@ class CronService:
|
||||
job.state.next_run_at_ms = None
|
||||
if self._should_persist_store():
|
||||
self._save_store()
|
||||
self._arm_timer()
|
||||
self._request_timer_rearm()
|
||||
else:
|
||||
self._append_action("update", asdict(job))
|
||||
return job
|
||||
@@ -825,7 +956,7 @@ class CronService:
|
||||
|
||||
if self._should_persist_store():
|
||||
self._save_store()
|
||||
self._arm_timer()
|
||||
self._request_timer_rearm()
|
||||
else:
|
||||
self._append_action("update", asdict(job))
|
||||
|
||||
@@ -840,19 +971,31 @@ class CronService:
|
||||
# A manual run is another side-effecting entrypoint. Do not start
|
||||
# it while the result of a previous timer execution is still only
|
||||
# in memory.
|
||||
if self._store_dirty:
|
||||
self._save_store()
|
||||
store = self._require_store(reload_during_execution=reload_store)
|
||||
async with self._operation_lock:
|
||||
if self._store_dirty:
|
||||
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:
|
||||
if job.id == job_id:
|
||||
if self._is_unbound_agent_job(job):
|
||||
self._enforce_agent_binding(job)
|
||||
self._save_store()
|
||||
async with self._operation_lock:
|
||||
self._enforce_agent_binding(job)
|
||||
await shield_and_drain(
|
||||
asyncio.to_thread(self._save_store)
|
||||
)
|
||||
return False
|
||||
if not force and not job.enabled:
|
||||
return False
|
||||
await self._execute_job(job)
|
||||
self._save_store()
|
||||
executed = await self._execute_job(job)
|
||||
if not executed:
|
||||
return False
|
||||
async with self._operation_lock:
|
||||
await shield_and_drain(asyncio.to_thread(self._save_store))
|
||||
return True
|
||||
return False
|
||||
finally:
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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)
|
||||
@@ -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")
|
||||
@@ -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
@@ -13,6 +13,7 @@ from nanobot.agent.loop import AgentLoop
|
||||
from nanobot.agent.tools.mcp import MCPProvider
|
||||
from nanobot.agent.tools.registry import ToolRegistry
|
||||
from nanobot.config.schema import Config
|
||||
from nanobot.providers.base import LLMUsage
|
||||
from nanobot.providers.image_generation import image_gen_provider_configs
|
||||
from nanobot.sdk.clients import MemoryClient, RuntimeClient, SessionClient
|
||||
from nanobot.sdk.runtime import (
|
||||
@@ -43,6 +44,7 @@ from nanobot.utils.llm_runtime import LLMRuntime
|
||||
|
||||
__all__ = [
|
||||
"Nanobot",
|
||||
"LLMUsage",
|
||||
"RunResult",
|
||||
"RunStream",
|
||||
"SessionInfo",
|
||||
@@ -287,7 +289,7 @@ class Nanobot:
|
||||
type=STREAM_EVENT_RUN_COMPLETED,
|
||||
content=result.content,
|
||||
result=result,
|
||||
usage=dict(result.usage),
|
||||
usage=result.usage,
|
||||
metadata=dict(result.metadata),
|
||||
))
|
||||
return result
|
||||
|
||||
@@ -5,11 +5,12 @@ from __future__ import annotations
|
||||
from importlib import import_module
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from nanobot.providers.base import LLMProvider, LLMResponse
|
||||
from nanobot.providers.base import LLMProvider, LLMResponse, LLMUsage
|
||||
|
||||
__all__ = [
|
||||
"LLMProvider",
|
||||
"LLMResponse",
|
||||
"LLMUsage",
|
||||
"AnthropicProvider",
|
||||
"OpenAICompatProvider",
|
||||
"OpenAICodexProvider",
|
||||
|
||||
@@ -17,6 +17,7 @@ from loguru import logger
|
||||
from nanobot.providers.base import (
|
||||
LLMProvider,
|
||||
LLMResponse,
|
||||
LLMUsage,
|
||||
ToolCallRequest,
|
||||
resolve_stream_idle_timeout_s,
|
||||
tool_arguments_object_for_replay,
|
||||
@@ -90,8 +91,10 @@ class AnthropicProvider(LLMProvider):
|
||||
api_base: str | None = None,
|
||||
default_model: str = "claude-sonnet-4-6",
|
||||
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.extra_headers = extra_headers or {}
|
||||
|
||||
@@ -689,24 +692,25 @@ class AnthropicProvider(LLMProvider):
|
||||
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")
|
||||
|
||||
usage: dict[str, int] = {}
|
||||
usage: LLMUsage | None = None
|
||||
if response.usage:
|
||||
input_tokens = response.usage.input_tokens
|
||||
cache_creation = getattr(response.usage, "cache_creation_input_tokens", 0) or 0
|
||||
cache_read = getattr(response.usage, "cache_read_input_tokens", 0) or 0
|
||||
total_prompt_tokens = input_tokens + cache_creation + cache_read
|
||||
usage = {
|
||||
"prompt_tokens": total_prompt_tokens,
|
||||
"completion_tokens": response.usage.output_tokens,
|
||||
"total_tokens": total_prompt_tokens + response.usage.output_tokens,
|
||||
}
|
||||
for attr in ("cache_creation_input_tokens", "cache_read_input_tokens"):
|
||||
val = getattr(response.usage, attr, 0)
|
||||
if val:
|
||||
usage[attr] = val
|
||||
# Normalize to cached_tokens for downstream consistency.
|
||||
if cache_read:
|
||||
usage["cached_tokens"] = cache_read
|
||||
cache_write_raw = getattr(
|
||||
response.usage,
|
||||
"cache_creation_input_tokens",
|
||||
None,
|
||||
)
|
||||
cache_read_raw = getattr(response.usage, "cache_read_input_tokens", None)
|
||||
cache_write = int(cache_write_raw) if cache_write_raw is not None else None
|
||||
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) + (
|
||||
cache_read or 0
|
||||
)
|
||||
usage = LLMUsage.reported(
|
||||
input_tokens=logical_input,
|
||||
output_tokens=int(response.usage.output_tokens),
|
||||
cache_read_tokens=cache_read,
|
||||
cache_write_tokens=cache_write,
|
||||
)
|
||||
|
||||
return LLMResponse(
|
||||
content="".join(content_parts) or None,
|
||||
|
||||
@@ -106,8 +106,10 @@ class AzureOpenAIProvider(LLMProvider):
|
||||
api_key: str = "",
|
||||
api_base: str = "",
|
||||
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._native_compaction_available = True
|
||||
|
||||
|
||||
+440
-9
@@ -6,6 +6,7 @@ import asyncio
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
from abc import ABC, abstractmethod
|
||||
from collections.abc import Awaitable, Callable
|
||||
from contextlib import suppress
|
||||
@@ -13,19 +14,23 @@ from copy import deepcopy
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timezone
|
||||
from email.utils import parsedate_to_datetime
|
||||
from typing import Any, cast
|
||||
from typing import TYPE_CHECKING, Any, Literal, cast
|
||||
|
||||
import json_repair
|
||||
from loguru import logger
|
||||
|
||||
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"
|
||||
DEFAULT_STREAM_IDLE_TIMEOUT_S = 90.0
|
||||
MAX_STREAM_IDLE_TIMEOUT_S = 3600.0
|
||||
RETRY_AFTER_BUFFER = 1
|
||||
|
||||
RetryEventCallback = Callable[[str], Awaitable[None]]
|
||||
LLMCallObserver = Callable[["LLMCallRecord"], None]
|
||||
|
||||
|
||||
def resolve_stream_idle_timeout_s(
|
||||
@@ -253,13 +258,298 @@ class ProviderCallContext:
|
||||
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
|
||||
class LLMResponse:
|
||||
"""Response from an LLM provider."""
|
||||
content: str | None
|
||||
tool_calls: list[ToolCallRequest] = field(default_factory=list)
|
||||
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
|
||||
# first token and provider retry gaps; ``ttft_ms`` measures the first
|
||||
# streamed reasoning/content delta from request start. They stay separate
|
||||
@@ -383,10 +673,109 @@ class LLMProvider(ABC):
|
||||
|
||||
_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_base = api_base
|
||||
self.provider_name = provider_name
|
||||
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(
|
||||
self,
|
||||
@@ -773,18 +1162,39 @@ class LLMProvider(ABC):
|
||||
|
||||
async def _safe_chat(self, **kwargs: Any) -> LLMResponse:
|
||||
"""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:
|
||||
provider_context = kwargs.pop("provider_context", None)
|
||||
if isinstance(provider_context, ProviderCallContext):
|
||||
return await self.chat_with_context(
|
||||
response = await self.chat_with_context(
|
||||
provider_context=provider_context,
|
||||
**kwargs,
|
||||
)
|
||||
return await self.chat(**kwargs)
|
||||
else:
|
||||
response = await self.chat(**kwargs)
|
||||
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
|
||||
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(
|
||||
self,
|
||||
@@ -847,18 +1257,39 @@ class LLMProvider(ABC):
|
||||
|
||||
async def _safe_chat_stream(self, **kwargs: Any) -> LLMResponse:
|
||||
"""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:
|
||||
provider_context = kwargs.pop("provider_context", None)
|
||||
if isinstance(provider_context, ProviderCallContext):
|
||||
return await self.chat_stream_with_context(
|
||||
response = await self.chat_stream_with_context(
|
||||
provider_context=provider_context,
|
||||
**kwargs,
|
||||
)
|
||||
return await self.chat_stream(**kwargs)
|
||||
else:
|
||||
response = await self.chat_stream(**kwargs)
|
||||
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
|
||||
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(
|
||||
self,
|
||||
|
||||
@@ -14,6 +14,7 @@ from typing import Any, cast
|
||||
from nanobot.providers.base import (
|
||||
LLMProvider,
|
||||
LLMResponse,
|
||||
LLMUsage,
|
||||
ToolCallRequest,
|
||||
parse_tool_arguments,
|
||||
resolve_stream_idle_timeout_s,
|
||||
@@ -60,8 +61,9 @@ class BedrockProvider(LLMProvider):
|
||||
profile: str | None = None,
|
||||
extra_body: dict[str, 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.region = region or os.environ.get("AWS_REGION") or os.environ.get("AWS_DEFAULT_REGION")
|
||||
self.profile = profile
|
||||
@@ -453,25 +455,25 @@ class BedrockProvider(LLMProvider):
|
||||
}.get(stop_reason or "", stop_reason or "stop")
|
||||
|
||||
@staticmethod
|
||||
def _usage(usage: dict[str, Any] | None) -> dict[str, int]:
|
||||
def _usage(usage: dict[str, Any] | None) -> LLMUsage | None:
|
||||
if not usage:
|
||||
return {}
|
||||
prompt = int(usage.get("inputTokens") or 0)
|
||||
completion = int(usage.get("outputTokens") or 0)
|
||||
total = int(usage.get("totalTokens") or prompt + completion)
|
||||
result = {
|
||||
"prompt_tokens": prompt,
|
||||
"completion_tokens": completion,
|
||||
"total_tokens": total,
|
||||
}
|
||||
cache_read = int(usage.get("cacheReadInputTokens") or 0)
|
||||
cache_write = int(usage.get("cacheWriteInputTokens") or 0)
|
||||
if cache_read:
|
||||
result["cached_tokens"] = cache_read
|
||||
result["cache_read_input_tokens"] = cache_read
|
||||
if cache_write:
|
||||
result["cache_creation_input_tokens"] = cache_write
|
||||
return result
|
||||
return None
|
||||
|
||||
def _optional_count(key: str) -> int | None:
|
||||
raw = usage.get(key)
|
||||
return int(raw) if raw is not None else None
|
||||
|
||||
cache_read = _optional_count("cacheReadInputTokens")
|
||||
cache_write = _optional_count("cacheWriteInputTokens")
|
||||
logical_input = int(usage.get("inputTokens") or 0) + (cache_read or 0) + (
|
||||
cache_write or 0
|
||||
)
|
||||
return LLMUsage.reported(
|
||||
input_tokens=logical_input,
|
||||
output_tokens=int(usage.get("outputTokens") or 0),
|
||||
cache_read_tokens=cache_read,
|
||||
cache_write_tokens=cache_write,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _parse_reasoning(block: dict[str, Any]) -> tuple[str | None, dict[str, Any] | None]:
|
||||
|
||||
@@ -172,6 +172,7 @@ def _make_provider_core(
|
||||
default_model=model,
|
||||
proxy=getattr(p, "proxy", None) if p else None,
|
||||
extra_body=p.extra_body if p else None,
|
||||
provider_name=provider_name,
|
||||
)
|
||||
elif backend == "xai_grok":
|
||||
from nanobot.providers.xai_grok_provider import XAIGrokProvider
|
||||
@@ -180,6 +181,7 @@ def _make_provider_core(
|
||||
default_model=model,
|
||||
proxy=getattr(p, "proxy", None) if p else None,
|
||||
extra_body=p.extra_body if p else None,
|
||||
provider_name=provider_name,
|
||||
)
|
||||
elif backend == "azure_openai":
|
||||
from nanobot.providers.azure_openai_provider import AzureOpenAIProvider
|
||||
@@ -190,11 +192,12 @@ def _make_provider_core(
|
||||
api_key=p.api_key or "",
|
||||
api_base=p.api_base,
|
||||
default_model=model,
|
||||
provider_name=provider_name,
|
||||
)
|
||||
elif backend == "github_copilot":
|
||||
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":
|
||||
from nanobot.providers.anthropic_provider import AnthropicProvider
|
||||
|
||||
@@ -203,6 +206,7 @@ def _make_provider_core(
|
||||
api_base=config.get_api_base(model, preset=preset),
|
||||
default_model=model,
|
||||
extra_headers=_provider_extra_headers(spec, p),
|
||||
provider_name=provider_name,
|
||||
)
|
||||
elif backend == "bedrock":
|
||||
from nanobot.providers.bedrock_provider import BedrockProvider
|
||||
@@ -214,6 +218,7 @@ def _make_provider_core(
|
||||
region=getattr(p, "region", None) if p else None,
|
||||
profile=getattr(p, "profile", None) if p else None,
|
||||
extra_body=p.extra_body if p else None,
|
||||
provider_name=provider_name,
|
||||
)
|
||||
else:
|
||||
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",
|
||||
extra_query=p.extra_query if p else None,
|
||||
proxy=p.proxy if p else None,
|
||||
provider_name=provider_name,
|
||||
)
|
||||
|
||||
provider.generation = preset.to_generation_settings()
|
||||
|
||||
@@ -13,6 +13,7 @@ from loguru import logger
|
||||
|
||||
from nanobot.providers.base import (
|
||||
GenerationSettings,
|
||||
LLMCallObserver,
|
||||
LLMProvider,
|
||||
LLMResponse,
|
||||
ProviderCallContext,
|
||||
@@ -124,7 +125,10 @@ class FallbackProvider(LLMProvider):
|
||||
fallback_model_observer: FallbackModelObserver | None = None,
|
||||
primary_context_window_tokens: int | None = None,
|
||||
):
|
||||
primary_generation = primary.generation
|
||||
self._primary = primary
|
||||
super().__init__(provider_name=primary.provider_name)
|
||||
self._primary.generation = primary_generation
|
||||
self._fallback_presets = list(fallback_presets)
|
||||
self._provider_factory = provider_factory
|
||||
self._fallback_model_observer = fallback_model_observer
|
||||
@@ -148,6 +152,11 @@ class FallbackProvider(LLMProvider):
|
||||
"""Attach a process-level observer without changing request call signatures."""
|
||||
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
|
||||
def supports_progress_deltas(self) -> bool:
|
||||
return bool(getattr(self._primary, "supports_progress_deltas", False))
|
||||
@@ -503,6 +512,7 @@ class FallbackProvider(LLMProvider):
|
||||
)
|
||||
try:
|
||||
fallback_provider = self._provider_factory(fallback)
|
||||
fallback_provider.set_llm_call_observer(self._llm_call_observer)
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Failed to create provider for fallback '{}': {}", fallback_model, exc
|
||||
|
||||
@@ -174,7 +174,12 @@ def login_github_copilot(
|
||||
class GitHubCopilotProvider(OpenAICompatProvider):
|
||||
"""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
|
||||
|
||||
self._copilot_access_token: str | None = None
|
||||
@@ -190,6 +195,7 @@ class GitHubCopilotProvider(OpenAICompatProvider):
|
||||
"User-Agent": USER_AGENT,
|
||||
},
|
||||
spec=find_by_name("github_copilot"),
|
||||
provider_name=provider_name,
|
||||
)
|
||||
|
||||
async def _get_copilot_access_token(self) -> str:
|
||||
|
||||
@@ -20,7 +20,7 @@ from nanobot.providers.registry import find_by_name
|
||||
from nanobot.security.network import (
|
||||
PinnedDNSAsyncTransport,
|
||||
UnsafeURLRequestError,
|
||||
resolve_url_target,
|
||||
async_resolve_url_target,
|
||||
)
|
||||
from nanobot.utils.helpers import detect_image_mime
|
||||
|
||||
@@ -174,7 +174,7 @@ async def _download_image_data_url(
|
||||
current_url = url
|
||||
for _ in range(_IMAGE_DOWNLOAD_MAX_REDIRECTS + 1):
|
||||
if proxy:
|
||||
ok, error, _ = resolve_url_target(
|
||||
ok, error, _ = await async_resolve_url_target(
|
||||
current_url,
|
||||
trust_remote_dns=True,
|
||||
)
|
||||
|
||||
@@ -7,6 +7,7 @@ from __future__ import annotations
|
||||
import asyncio
|
||||
import hashlib
|
||||
import json
|
||||
import ssl
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import Any, cast
|
||||
|
||||
@@ -50,12 +51,26 @@ class OpenAICodexProvider(LLMProvider):
|
||||
default_model: str = "openai-codex/gpt-5.6-sol",
|
||||
proxy: str | 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.proxy = proxy or None
|
||||
self._extra_body = dict(extra_body or {})
|
||||
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(
|
||||
self,
|
||||
@@ -129,7 +144,7 @@ class OpenAICodexProvider(LLMProvider):
|
||||
DEFAULT_CODEX_URL,
|
||||
headers,
|
||||
wire_body,
|
||||
verify=True,
|
||||
verify=self._ssl_context(verify=True),
|
||||
proxy=self.proxy,
|
||||
on_content_delta=on_content_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,
|
||||
headers,
|
||||
wire_body,
|
||||
verify=False,
|
||||
verify=self._ssl_context(verify=False),
|
||||
proxy=self.proxy,
|
||||
on_content_delta=on_content_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,
|
||||
headers: dict[str, str],
|
||||
body: dict[str, Any],
|
||||
verify: bool,
|
||||
verify: ssl.SSLContext | bool,
|
||||
proxy: str | None = None,
|
||||
on_content_delta: Callable[[str], Awaitable[None]] | None = None,
|
||||
on_thinking_delta: Callable[[str], Awaitable[None]] | None = None,
|
||||
|
||||
@@ -26,6 +26,7 @@ from pydantic.alias_generators import to_snake
|
||||
from nanobot.providers.base import (
|
||||
LLMProvider,
|
||||
LLMResponse,
|
||||
LLMUsage,
|
||||
ProviderCallContext,
|
||||
ProviderConversationState,
|
||||
ToolCallRequest,
|
||||
@@ -517,8 +518,9 @@ class OpenAICompatProvider(LLMProvider):
|
||||
api_type: str = "auto",
|
||||
extra_query: dict[str, 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.extra_headers = extra_headers or {}
|
||||
self._spec = spec
|
||||
@@ -1428,12 +1430,12 @@ class OpenAICompatProvider(LLMProvider):
|
||||
return "".join(parts) or None
|
||||
|
||||
@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.
|
||||
|
||||
Handles both dict-based (raw JSON) and object-based (SDK Pydantic)
|
||||
responses. Provider-specific ``cached_tokens`` fields are normalised
|
||||
under a single key; see the priority chain inside for details.
|
||||
responses. Provider-specific cache fields are normalized once at
|
||||
this Chat Completions wire boundary.
|
||||
"""
|
||||
# --- resolve usage object ---
|
||||
usage_obj = None
|
||||
@@ -1445,21 +1447,18 @@ class OpenAICompatProvider(LLMProvider):
|
||||
|
||||
usage_map = cls._maybe_mapping(usage_obj)
|
||||
if usage_map is not None:
|
||||
result = {
|
||||
"prompt_tokens": int(usage_map.get("prompt_tokens") or 0),
|
||||
"completion_tokens": int(usage_map.get("completion_tokens") or 0),
|
||||
"total_tokens": int(usage_map.get("total_tokens") or 0),
|
||||
}
|
||||
input_tokens = int(usage_map.get("prompt_tokens") or 0)
|
||||
output_tokens = int(usage_map.get("completion_tokens") or 0)
|
||||
elif usage_obj:
|
||||
result = {
|
||||
"prompt_tokens": getattr(usage_obj, "prompt_tokens", 0) or 0,
|
||||
"completion_tokens": getattr(usage_obj, "completion_tokens", 0) or 0,
|
||||
"total_tokens": getattr(usage_obj, "total_tokens", 0) or 0,
|
||||
}
|
||||
input_tokens = int(getattr(usage_obj, "prompt_tokens", 0) or 0)
|
||||
output_tokens = int(getattr(usage_obj, "completion_tokens", 0) or 0)
|
||||
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).
|
||||
# Priority order ensures the most specific field wins.
|
||||
for path in (
|
||||
@@ -1468,17 +1467,28 @@ class OpenAICompatProvider(LLMProvider):
|
||||
("prompt_cache_hit_tokens",), # DeepSeek/SiliconFlow
|
||||
):
|
||||
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)
|
||||
if cached:
|
||||
result["cached_tokens"] = cached
|
||||
if cached is not None:
|
||||
cache_read = cached
|
||||
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
|
||||
def _get_nested_int(obj: object, path: tuple[str, ...]) -> int:
|
||||
"""Drill into *obj* by *path* segments and return an ``int`` value.
|
||||
def _get_nested_int(obj: object, path: tuple[str, ...]) -> int | None:
|
||||
"""Return a present usage count while preserving explicit zero.
|
||||
|
||||
Supports both dict-key access and attribute access so it works
|
||||
uniformly with raw JSON dicts **and** SDK Pydantic models.
|
||||
@@ -1486,12 +1496,17 @@ class OpenAICompatProvider(LLMProvider):
|
||||
current: object = obj
|
||||
for segment in path:
|
||||
if current is None:
|
||||
return 0
|
||||
return None
|
||||
if isinstance(current, dict):
|
||||
current = cast(dict[str, Any], current).get(segment)
|
||||
else:
|
||||
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:
|
||||
if isinstance(response, str):
|
||||
@@ -1645,7 +1660,7 @@ class OpenAICompatProvider(LLMProvider):
|
||||
reasoning_parts: list[str] = []
|
||||
tc_bufs: dict[int, dict[str, Any]] = {}
|
||||
finish_reason = "stop"
|
||||
usage: dict[str, int] = {}
|
||||
usage: LLMUsage | None = None
|
||||
|
||||
def _accum_tc(tc: Any, idx_hint: int) -> None:
|
||||
"""Accumulate one streaming tool-call delta into *tc_bufs*."""
|
||||
|
||||
@@ -10,7 +10,7 @@ from typing import Any, AsyncGenerator, cast
|
||||
import httpx
|
||||
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
|
||||
|
||||
FINISH_REASON_MAP = {
|
||||
@@ -186,33 +186,40 @@ def _response_finish_reason(
|
||||
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)
|
||||
usage_raw: object = (
|
||||
response_object.get("usage")
|
||||
if response_object is not None
|
||||
else getattr(response, "usage", None)
|
||||
)
|
||||
if not usage_raw:
|
||||
return {}
|
||||
if usage_raw is None:
|
||||
return None
|
||||
usage = _response_object(usage_raw)
|
||||
if usage is None:
|
||||
return {}
|
||||
prompt_tokens = int(usage.get("input_tokens") or usage.get("prompt_tokens") or 0)
|
||||
completion_tokens = int(
|
||||
usage.get("output_tokens") or usage.get("completion_tokens") or 0
|
||||
)
|
||||
total_tokens = int(usage.get("total_tokens") or prompt_tokens + completion_tokens)
|
||||
result = {
|
||||
"prompt_tokens": prompt_tokens,
|
||||
"completion_tokens": completion_tokens,
|
||||
"total_tokens": total_tokens,
|
||||
}
|
||||
return None
|
||||
|
||||
def _usage_int(container: dict[str, Any] | None, key: str) -> int | None:
|
||||
if container is None:
|
||||
return None
|
||||
raw = container.get(key)
|
||||
if raw is None or isinstance(raw, bool):
|
||||
return None
|
||||
try:
|
||||
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"))
|
||||
cached_tokens = int(input_details.get("cached_tokens") or 0) if input_details else 0
|
||||
if cached_tokens > 0:
|
||||
result["cached_tokens"] = cached_tokens
|
||||
return result
|
||||
return LLMUsage.reported(
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
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:
|
||||
@@ -352,14 +359,14 @@ async def consume_sse_with_reasoning(
|
||||
on_reasoning_delta: Callable[[str], Awaitable[None]] | None = None,
|
||||
on_response_event: Callable[[dict[str, Any]], Awaitable[None]] | 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."""
|
||||
content = ""
|
||||
tool_calls: list[ToolCallRequest] = []
|
||||
tool_call_buffers: dict[str, dict[str, Any]] = {}
|
||||
tool_call_args_emitted: set[str] = set()
|
||||
finish_reason = "stop"
|
||||
usage: dict[str, int] = {}
|
||||
usage: LLMUsage | None = None
|
||||
reasoning_content: str | None = None
|
||||
streamed_reasoning = False
|
||||
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_reasoning_delta: Callable[[str], Awaitable[None]] | 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)``."""
|
||||
content = ""
|
||||
tool_calls: list[ToolCallRequest] = []
|
||||
tool_call_buffers: dict[str, dict[str, Any]] = {}
|
||||
tool_call_args_emitted: set[str] = set()
|
||||
finish_reason = "stop"
|
||||
usage: dict[str, int] = {}
|
||||
usage: LLMUsage | None = None
|
||||
reasoning_content: str | None = None
|
||||
streamed_reasoning = False
|
||||
refusal_seen = False
|
||||
@@ -823,20 +830,7 @@ async def consume_sdk_stream(
|
||||
if on_content_delta and remaining_text:
|
||||
await on_content_delta(remaining_text)
|
||||
if resp:
|
||||
usage_obj = getattr(resp, "usage", None)
|
||||
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
|
||||
usage = _usage_from_response_obj(resp) or usage
|
||||
if not reasoning_content:
|
||||
reasoning_content = _extract_reasoning_summary_from_output(
|
||||
getattr(resp, "output", None)
|
||||
|
||||
@@ -7,7 +7,7 @@ from typing import Any, cast
|
||||
|
||||
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
|
||||
|
||||
RESPONSES_STATE_KIND = "openai_responses"
|
||||
@@ -84,7 +84,7 @@ def build_responses_state(
|
||||
model: str,
|
||||
input_items: list[dict[str, Any]],
|
||||
output_items: list[dict[str, Any]],
|
||||
usage: dict[str, int] | None = None,
|
||||
usage: LLMUsage | None = None,
|
||||
) -> ProviderConversationState:
|
||||
"""Create the canonical next state from request input and every output item."""
|
||||
unpruned_items = [*input_items, *output_items]
|
||||
@@ -178,16 +178,8 @@ def _prune_before_latest_output_compaction(
|
||||
return output_items[latest:]
|
||||
|
||||
|
||||
def _context_tokens_from_usage(usage: dict[str, int] | None) -> int:
|
||||
if not usage:
|
||||
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 _context_tokens_from_usage(usage: LLMUsage | None) -> int:
|
||||
return usage.total_tokens if usage is not None else 0
|
||||
|
||||
|
||||
def _state_items(
|
||||
|
||||
@@ -11,7 +11,7 @@ class UnconfiguredProvider(LLMProvider):
|
||||
"""Keep the gateway available for settings before a model is configured."""
|
||||
|
||||
def __init__(self, default_model: str) -> None:
|
||||
super().__init__()
|
||||
super().__init__(provider_name="unconfigured")
|
||||
self._default_model = default_model
|
||||
|
||||
async def chat(
|
||||
|
||||
@@ -18,6 +18,7 @@ from nanobot import __version__
|
||||
from nanobot.providers.base import (
|
||||
LLMProvider,
|
||||
LLMResponse,
|
||||
LLMUsage,
|
||||
ToolCallRequest,
|
||||
resolve_stream_idle_timeout_s,
|
||||
)
|
||||
@@ -69,8 +70,10 @@ class XAIGrokProvider(LLMProvider):
|
||||
default_model: str = DEFAULT_XAI_GROK_MODEL,
|
||||
proxy: str | 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.proxy = proxy or None
|
||||
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_thinking_delta: Callable[[str], 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:
|
||||
hosted_event = _xai_hosted_tool_event(event)
|
||||
if hosted_event is not None and on_tool_call_delta is not None:
|
||||
|
||||
@@ -210,18 +210,20 @@ class RuntimeClient:
|
||||
|
||||
async def compact_session(self, session_key: str) -> SessionSnapshot:
|
||||
"""Run token consolidation for one session."""
|
||||
session = self._loop.sessions.get_or_create(session_key)
|
||||
runtime = self._loop.runtime_for_session(session)
|
||||
session = await self._loop.sessions.get_or_create_async(session_key)
|
||||
runtime = await self._loop.runtime_for_session_async(session)
|
||||
await self._loop.consolidator.maybe_consolidate_by_tokens(
|
||||
session,
|
||||
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:
|
||||
"""Run idle-session compaction for one session and return the summary."""
|
||||
session = self._loop.sessions.get_or_create(session_key)
|
||||
runtime = self._loop.runtime_for_session(session)
|
||||
session = await self._loop.sessions.get_or_create_async(session_key)
|
||||
runtime = await self._loop.runtime_for_session_async(session)
|
||||
return await self._loop.consolidator.compact_idle_session(
|
||||
session_key,
|
||||
runtime=runtime,
|
||||
|
||||
@@ -6,6 +6,7 @@ from copy import deepcopy
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Literal, Mapping, TypeAlias, cast
|
||||
|
||||
from nanobot.providers.base import LLMUsage
|
||||
from nanobot.runtime_context import public_history_messages
|
||||
|
||||
StreamEventType: TypeAlias = Literal[
|
||||
@@ -53,7 +54,7 @@ class RunResult:
|
||||
content: str
|
||||
tools_used: list[str] = 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
|
||||
error: str | None = None
|
||||
metadata: dict[str, Any] = field(default_factory=dict)
|
||||
@@ -72,7 +73,7 @@ class StreamEvent:
|
||||
arguments: dict[str, Any] | None = None
|
||||
iteration: int | None = None
|
||||
resuming: bool | None = None
|
||||
usage: dict[str, int] = field(default_factory=dict)
|
||||
usage: LLMUsage | None = None
|
||||
error: str | None = None
|
||||
metadata: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
+144
-42
@@ -29,6 +29,7 @@ _BLOCKED_NETWORKS = [
|
||||
|
||||
_URL_RE = re.compile(r"https?://[^\s\"'`;|<>]+", re.IGNORECASE)
|
||||
_allowed_networks: list[ipaddress.IPv4Network | ipaddress.IPv6Network] = []
|
||||
_DNS_RESOLUTION_TIMEOUT_SECONDS = 5.0
|
||||
|
||||
|
||||
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)
|
||||
|
||||
|
||||
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(
|
||||
url: str,
|
||||
*,
|
||||
@@ -97,52 +155,43 @@ def resolve_url_target(
|
||||
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.
|
||||
"""
|
||||
try:
|
||||
p = urlparse(url)
|
||||
except Exception as e:
|
||||
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", ()
|
||||
|
||||
hostname, error = _parse_url_hostname(url)
|
||||
if hostname is None:
|
||||
return False, error or "Missing hostname", ()
|
||||
try:
|
||||
infos = socket.getaddrinfo(hostname, None, socket.AF_UNSPEC, socket.SOCK_STREAM)
|
||||
except socket.gaierror:
|
||||
if not trust_remote_dns:
|
||||
return False, f"Cannot resolve hostname: {hostname}", ()
|
||||
return _unresolved_target_result(hostname, trust_remote_dns=trust_remote_dns)
|
||||
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:
|
||||
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)),)
|
||||
|
||||
addrs: list[ipaddress.IPv4Address | ipaddress.IPv6Address] = []
|
||||
for info in infos:
|
||||
try:
|
||||
addr = ipaddress.ip_address(info[4][0])
|
||||
except 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))
|
||||
async def async_resolve_url_target(
|
||||
url: str,
|
||||
*,
|
||||
allow_loopback: bool = False,
|
||||
trust_remote_dns: bool = False,
|
||||
timeout_s: float = _DNS_RESOLUTION_TIMEOUT_SECONDS,
|
||||
) -> tuple[bool, str, tuple[str, ...]]:
|
||||
"""Resolve and validate an HTTP target without blocking the event loop."""
|
||||
hostname, error = _parse_url_hostname(url)
|
||||
if hostname is None:
|
||||
return False, error or "Missing hostname", ()
|
||||
loop = asyncio.get_running_loop()
|
||||
try:
|
||||
infos = await asyncio.wait_for(
|
||||
loop.getaddrinfo(
|
||||
hostname,
|
||||
None,
|
||||
family=socket.AF_UNSPEC,
|
||||
type=socket.SOCK_STREAM,
|
||||
),
|
||||
timeout=timeout_s,
|
||||
)
|
||||
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]:
|
||||
@@ -151,6 +200,16 @@ def validate_url_target(url: str, *, allow_loopback: bool = False) -> tuple[bool
|
||||
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:
|
||||
"""Return True when process proxy settings would proxy this URL."""
|
||||
try:
|
||||
@@ -277,7 +336,10 @@ class PinnedDNSAsyncTransport(httpx.AsyncBaseTransport):
|
||||
|
||||
async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
|
||||
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:
|
||||
raise UnsafeURLRequestError(error, request=request)
|
||||
async with self._resolver_lock:
|
||||
@@ -320,6 +382,46 @@ def validate_resolved_url(url: str) -> tuple[bool, str]:
|
||||
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:
|
||||
"""Return True if the command string contains a URL targeting an internal/private address."""
|
||||
for m in _URL_RE.finditer(command):
|
||||
|
||||
@@ -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
@@ -1,5 +1,6 @@
|
||||
"""Session management for conversation history."""
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import errno
|
||||
import hashlib
|
||||
@@ -28,6 +29,7 @@ from nanobot.runtime_context import (
|
||||
public_history_message,
|
||||
)
|
||||
from nanobot.session.model_selection import SESSION_MODEL_PRESET_METADATA_KEY
|
||||
from nanobot.utils.cancellation import shield_and_drain
|
||||
from nanobot.utils.helpers import (
|
||||
content_with_media_breadcrumbs,
|
||||
ensure_dir,
|
||||
@@ -71,6 +73,7 @@ _WORKSPACE_STATE_DIR = ".nanobot"
|
||||
_WORKSPACE_ID_FILE = "workspace-id"
|
||||
_WORKSPACE_ID_RE = re.compile(r"^[0-9a-f]{32}$")
|
||||
_SESSION_MIGRATION_LOCK_TIMEOUT_SECONDS = 30
|
||||
_SESSION_FILES_LOCK_TIMEOUT_SECONDS = 5
|
||||
_SESSION_FILES_LOCK_FILENAME = ".session-files.lock"
|
||||
_COPY_CHUNK_SIZE = 1024 * 1024
|
||||
|
||||
@@ -560,7 +563,8 @@ class JsonlSessionStore:
|
||||
self.sessions_dir = ensure_dir(root / workspace_id)
|
||||
self.legacy_sessions_dir = get_legacy_sessions_dir()
|
||||
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:
|
||||
self._migrate_from_workspace(canonical_workspace)
|
||||
@@ -1642,6 +1646,7 @@ class SessionManager:
|
||||
self._cache: OrderedDict[str, Session] = OrderedDict()
|
||||
# Preserve identity for sessions held by active callers without retaining idle ones.
|
||||
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._delete_observer: Callable[[str], None] | None = None
|
||||
|
||||
@@ -1741,6 +1746,28 @@ class SessionManager:
|
||||
self._remember(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(
|
||||
self,
|
||||
key: str,
|
||||
@@ -1774,6 +1801,17 @@ class SessionManager:
|
||||
self._store.save(session, fsync=fsync)
|
||||
self._remember(session)
|
||||
|
||||
async def save_async(self, session: Session, *, fsync: bool = False) -> None:
|
||||
"""Persist a session without blocking the caller's event loop."""
|
||||
if not session.policy.persist:
|
||||
return
|
||||
|
||||
async def save_and_remember() -> None:
|
||||
await asyncio.to_thread(self._store.save, session, fsync=fsync)
|
||||
self._remember(session)
|
||||
|
||||
await shield_and_drain(save_and_remember())
|
||||
|
||||
def save_runtime_checkpoint(self, session: Session) -> None:
|
||||
"""Persist volatile recovery state without rewriting long history."""
|
||||
if not session.policy.persist:
|
||||
@@ -1786,6 +1824,23 @@ class SessionManager:
|
||||
# they opt into a dedicated checkpoint primitive.
|
||||
self.save(session)
|
||||
|
||||
async def save_runtime_checkpoint_async(self, session: Session) -> None:
|
||||
"""Persist an in-flight checkpoint without blocking the event loop."""
|
||||
if not session.policy.persist:
|
||||
return
|
||||
|
||||
async def save_and_remember() -> None:
|
||||
if self._store is self._jsonl_store:
|
||||
await asyncio.to_thread(
|
||||
self._jsonl_store.save_runtime_checkpoint,
|
||||
session,
|
||||
)
|
||||
else:
|
||||
await asyncio.to_thread(self._store.save, session)
|
||||
self._remember(session)
|
||||
|
||||
await shield_and_drain(save_and_remember())
|
||||
|
||||
def rename_model_preset(self, old_name: str, new_name: str) -> int:
|
||||
"""Rename a session-scoped model preset across durable and live sessions."""
|
||||
if old_name == new_name:
|
||||
@@ -1827,6 +1882,21 @@ class SessionManager:
|
||||
raise
|
||||
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:
|
||||
"""Re-save every cached session with fsync for durable shutdown.
|
||||
|
||||
@@ -1858,6 +1928,18 @@ class SessionManager:
|
||||
self._delete_observer(key)
|
||||
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:
|
||||
"""Restore session files to the pre-relocation path for an explicit rollback."""
|
||||
return self._jsonl_store.restore_to_workspace()
|
||||
@@ -1930,6 +2012,10 @@ class SessionManager:
|
||||
"""Read session metadata without loading the transcript."""
|
||||
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(
|
||||
self,
|
||||
key: str,
|
||||
@@ -1943,5 +2029,31 @@ class SessionManager:
|
||||
session.metadata.update(deepcopy(updates))
|
||||
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]]:
|
||||
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
@@ -25,6 +25,7 @@ from nanobot.bus.outbound_events import (
|
||||
)
|
||||
from nanobot.bus.queue import MessageBus
|
||||
from nanobot.session import turn_continuation
|
||||
from nanobot.session.async_compat import call_session_manager
|
||||
from nanobot.session.keys import UNIFIED_SESSION_KEY, last_channel_from_metadata
|
||||
from nanobot.session.manager import Session, SessionManager
|
||||
from nanobot.webui.metadata import WEBUI_TURN_METADATA_KEY
|
||||
@@ -460,6 +461,37 @@ class RecoveryCoordinator:
|
||||
repr=False,
|
||||
)
|
||||
|
||||
async def _get_or_create_session(self, key: str) -> Session:
|
||||
return await call_session_manager(
|
||||
self.sessions,
|
||||
"get_or_create_async",
|
||||
self.sessions.get_or_create,
|
||||
key,
|
||||
)
|
||||
|
||||
async def _save_session(self, session: Session) -> None:
|
||||
await call_session_manager(
|
||||
self.sessions,
|
||||
"save_async",
|
||||
self.sessions.save,
|
||||
session,
|
||||
)
|
||||
|
||||
async def _read_session_metadata(self, key: str) -> dict[str, Any] | None:
|
||||
return await call_session_manager(
|
||||
self.sessions,
|
||||
"read_session_metadata_async",
|
||||
self.sessions.read_session_metadata,
|
||||
key,
|
||||
)
|
||||
|
||||
async def _list_sessions(self) -> list[dict[str, Any]]:
|
||||
return await call_session_manager(
|
||||
self.sessions,
|
||||
"list_sessions_async",
|
||||
self.sessions.list_sessions,
|
||||
)
|
||||
|
||||
def register_recovery_task(self, session_key: str, task: asyncio.Task[Any]) -> None:
|
||||
"""Track the task that owns an explicit recovery continuation."""
|
||||
self._active_recovery_tasks[session_key] = task
|
||||
@@ -482,8 +514,8 @@ class RecoveryCoordinator:
|
||||
|
||||
async def scan(self) -> None:
|
||||
"""Recover every interrupted WebUI session once at gateway startup."""
|
||||
for key in self._recovery_candidates():
|
||||
metadata_payload = self.sessions.read_session_metadata(key)
|
||||
for key in await self._recovery_candidates():
|
||||
metadata_payload = await self._read_session_metadata(key)
|
||||
raw_metadata = metadata_payload.get("metadata") if metadata_payload else None
|
||||
metadata = cast(dict[str, Any], raw_metadata) if isinstance(raw_metadata, dict) else {}
|
||||
route = self._websocket_route_for(key, metadata)
|
||||
@@ -492,7 +524,7 @@ class RecoveryCoordinator:
|
||||
unfinished = self._has_unfinished_webui_transcript(key)
|
||||
if not self._needs_recovery(metadata) and not unfinished:
|
||||
continue
|
||||
session = self.sessions.get_or_create(key)
|
||||
session = await self._get_or_create_session(key)
|
||||
try:
|
||||
await self._recover_session(session, route[1])
|
||||
await self._requeue_pending_followups(session)
|
||||
@@ -507,14 +539,14 @@ class RecoveryCoordinator:
|
||||
reason="recovery_failed",
|
||||
can_continue=False,
|
||||
)
|
||||
self.sessions.save(session)
|
||||
await self._save_session(session)
|
||||
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."""
|
||||
candidates = dict.fromkeys(
|
||||
key
|
||||
for item in self.sessions.list_sessions()
|
||||
for item in await self._list_sessions()
|
||||
if isinstance((key := item.get("key")), str)
|
||||
)
|
||||
try:
|
||||
@@ -523,7 +555,7 @@ class RecoveryCoordinator:
|
||||
# duplicating its filename and migration rules here would drift.
|
||||
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")
|
||||
if isinstance(key, str):
|
||||
candidates.setdefault(key, None)
|
||||
@@ -549,7 +581,7 @@ class RecoveryCoordinator:
|
||||
"""Reject stale queued recoveries and let new user input supersede them."""
|
||||
recovery_id = message.metadata.get(RECOVERY_INBOUND_METADATA_KEY)
|
||||
if isinstance(recovery_id, str):
|
||||
session = 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)
|
||||
return bool(
|
||||
state
|
||||
@@ -558,7 +590,7 @@ class RecoveryCoordinator:
|
||||
)
|
||||
if message.channel != "websocket":
|
||||
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)
|
||||
if state and state["status"] in {"resuming", "awaiting_user", "failed"}:
|
||||
await self._cancel_active_recovery(message.session_key)
|
||||
@@ -572,13 +604,13 @@ class RecoveryCoordinator:
|
||||
attempts=cast(int, state.get("attempts", 0)),
|
||||
reason="superseded",
|
||||
)
|
||||
self.sessions.save(session)
|
||||
await self._save_session(session)
|
||||
await self._publish(message.chat_id, recovered)
|
||||
return True
|
||||
|
||||
async def turn_completed(self, session_key: str) -> None:
|
||||
"""Resolve a resuming state after the recovered turn commits."""
|
||||
session = self.sessions.get_or_create(session_key)
|
||||
session = await self._get_or_create_session(session_key)
|
||||
state = recovery_state_from_metadata(session.metadata)
|
||||
if not state or state["status"] != "resuming":
|
||||
return
|
||||
@@ -592,7 +624,7 @@ class RecoveryCoordinator:
|
||||
attempts=cast(int, state.get("attempts", 0)),
|
||||
reason="continued",
|
||||
)
|
||||
self.sessions.save(session)
|
||||
await self._save_session(session)
|
||||
await self._publish(route[1], recovered)
|
||||
|
||||
async def handle_action(self, action: str, payload: dict[str, Any]) -> dict[str, Any]:
|
||||
@@ -603,7 +635,7 @@ class RecoveryCoordinator:
|
||||
raise RecoveryActionError("missing chat_id")
|
||||
if not isinstance(recovery_id, str) or not 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)
|
||||
if not state or state["recovery_id"] != recovery_id:
|
||||
raise RecoveryActionError("recovery state is stale", status=409)
|
||||
@@ -618,7 +650,7 @@ class RecoveryCoordinator:
|
||||
attempts=cast(int, state.get("attempts", 0)),
|
||||
reason="dismissed",
|
||||
)
|
||||
self.sessions.save(session)
|
||||
await self._save_session(session)
|
||||
await self._publish(chat_id, next_state)
|
||||
return next_state
|
||||
if action != "continue":
|
||||
@@ -635,7 +667,7 @@ class RecoveryCoordinator:
|
||||
reason="user_confirmed",
|
||||
resume_message_count=len(session.messages),
|
||||
)
|
||||
self.sessions.save(session)
|
||||
await self._save_session(session)
|
||||
await self._publish(chat_id, next_state)
|
||||
await self._queue_continuation(session, chat_id, next_state)
|
||||
return next_state
|
||||
@@ -668,7 +700,7 @@ class RecoveryCoordinator:
|
||||
attempts=cast(int, state.get("attempts", 1)),
|
||||
reason="loop_guard",
|
||||
)
|
||||
self.sessions.save(session)
|
||||
await self._save_session(session)
|
||||
await self._publish(chat_id, next_state)
|
||||
elif self._has_unfinished_webui_transcript(session.key):
|
||||
# A normal last-client shutdown can materialize the checkpoint
|
||||
@@ -690,7 +722,7 @@ class RecoveryCoordinator:
|
||||
),
|
||||
can_continue=can_continue,
|
||||
)
|
||||
self.sessions.save(session)
|
||||
await self._save_session(session)
|
||||
await self._publish(chat_id, waiting)
|
||||
return
|
||||
if state and state["status"] in {"awaiting_user", "failed"}:
|
||||
@@ -706,7 +738,7 @@ class RecoveryCoordinator:
|
||||
attempts=cast(int, state.get("attempts", 1)),
|
||||
reason="loop_guard",
|
||||
)
|
||||
self.sessions.save(session)
|
||||
await self._save_session(session)
|
||||
await self._publish(chat_id, waiting)
|
||||
return
|
||||
|
||||
@@ -724,7 +756,7 @@ class RecoveryCoordinator:
|
||||
reason="checkpoint_unknown",
|
||||
can_continue=False,
|
||||
)
|
||||
self.sessions.save(session)
|
||||
await self._save_session(session)
|
||||
await self._publish(chat_id, waiting)
|
||||
return
|
||||
if checkpoint is not None and not _runtime_checkpoint_is_well_formed(checkpoint):
|
||||
@@ -738,7 +770,7 @@ class RecoveryCoordinator:
|
||||
reason="checkpoint_invalid",
|
||||
can_continue=False,
|
||||
)
|
||||
self.sessions.save(session)
|
||||
await self._save_session(session)
|
||||
await self._publish(chat_id, waiting)
|
||||
return
|
||||
if phase == "final_response":
|
||||
@@ -750,7 +782,7 @@ class RecoveryCoordinator:
|
||||
attempts=0,
|
||||
reason="answer_restored",
|
||||
)
|
||||
self.sessions.save(session)
|
||||
await self._save_session(session)
|
||||
await self._publish(chat_id, recovered)
|
||||
return
|
||||
if phase in _UNCERTAIN_TOOL_PHASES or pending_calls:
|
||||
@@ -762,7 +794,7 @@ class RecoveryCoordinator:
|
||||
attempts=0,
|
||||
reason="tool_state_unknown",
|
||||
)
|
||||
self.sessions.save(session)
|
||||
await self._save_session(session)
|
||||
await self._publish(chat_id, waiting)
|
||||
return
|
||||
# A gateway restart is a lifecycle boundary. Never enqueue model work
|
||||
@@ -777,7 +809,7 @@ class RecoveryCoordinator:
|
||||
attempts=0,
|
||||
reason="restart_requires_confirmation",
|
||||
)
|
||||
self.sessions.save(session)
|
||||
await self._save_session(session)
|
||||
await self._publish(chat_id, waiting)
|
||||
|
||||
async def _queue_continuation(
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import re
|
||||
import time
|
||||
from collections.abc import Awaitable, Callable
|
||||
@@ -37,7 +38,8 @@ from nanobot.bus.runtime_events import (
|
||||
TurnRuntimeAdmitted,
|
||||
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.runtime_context import public_history_message
|
||||
from nanobot.session.goal_state import goal_state_ws_blob
|
||||
@@ -175,7 +177,7 @@ async def maybe_generate_webui_title(
|
||||
model: str,
|
||||
) -> bool:
|
||||
"""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:
|
||||
return False
|
||||
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 != current_title:
|
||||
session.metadata[WEBUI_TITLE_METADATA_KEY] = cleaned_current_title
|
||||
sessions.save(session)
|
||||
await sessions.save_async(session)
|
||||
return False
|
||||
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)}"
|
||||
|
||||
try:
|
||||
response = await provider.chat_with_retry(
|
||||
[
|
||||
{
|
||||
"role": "system",
|
||||
"content": (
|
||||
"You write short, neutral chat titles. "
|
||||
"Return only the title text."
|
||||
),
|
||||
},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
tools=None,
|
||||
model=model,
|
||||
max_tokens=TITLE_GENERATION_MAX_TOKENS,
|
||||
temperature=0.2,
|
||||
reasoning_effort=TITLE_GENERATION_REASONING_EFFORT,
|
||||
retry_mode="standard",
|
||||
)
|
||||
with llm_usage_source("system"):
|
||||
response = await provider.chat_with_retry(
|
||||
[
|
||||
{
|
||||
"role": "system",
|
||||
"content": (
|
||||
"You write short, neutral chat titles. "
|
||||
"Return only the title text."
|
||||
),
|
||||
},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
tools=None,
|
||||
model=model,
|
||||
max_tokens=TITLE_GENERATION_MAX_TOKENS,
|
||||
temperature=0.2,
|
||||
reasoning_effort=TITLE_GENERATION_REASONING_EFFORT,
|
||||
retry_mode="standard",
|
||||
)
|
||||
except Exception:
|
||||
logger.debug("Failed to generate webui session title for {}", session_key, exc_info=True)
|
||||
return False
|
||||
@@ -239,7 +242,7 @@ async def maybe_generate_webui_title(
|
||||
)
|
||||
return False
|
||||
session.metadata[WEBUI_TITLE_METADATA_KEY] = title
|
||||
sessions.save(session)
|
||||
await sessions.save_async(session)
|
||||
return True
|
||||
|
||||
|
||||
@@ -435,8 +438,8 @@ class WebuiTurnRoutePolicy:
|
||||
)
|
||||
and route.channel == "websocket"
|
||||
):
|
||||
session = self.sessions.get_or_create(session_key)
|
||||
if session.metadata.get(WEBUI_SESSION_METADATA_KEY) is True:
|
||||
session = self.sessions.get_cached(session_key)
|
||||
if session is not None and session.metadata.get(WEBUI_SESSION_METADATA_KEY) is True:
|
||||
metadata = dict(route.metadata)
|
||||
turn_prefix = "session-input" if internal_user_input else "subagent"
|
||||
metadata.update({
|
||||
@@ -578,7 +581,7 @@ class WebuiTurnCoordinator:
|
||||
or not session_key.startswith("websocket:")
|
||||
):
|
||||
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 = (
|
||||
cast(dict[str, Any], metadata_value)
|
||||
@@ -589,7 +592,8 @@ class WebuiTurnCoordinator:
|
||||
return
|
||||
public_metadata = _session_message_public_metadata(envelope)
|
||||
try:
|
||||
append_session_message_input(
|
||||
await asyncio.to_thread(
|
||||
append_session_message_input,
|
||||
session_key,
|
||||
content=event.content,
|
||||
created_at_ms=envelope["created_at_ms"],
|
||||
@@ -614,8 +618,9 @@ class WebuiTurnCoordinator:
|
||||
def _handle_session_turn_started(self, event: SessionTurnStarted) -> None:
|
||||
if not self._is_websocket_event(event.context):
|
||||
return
|
||||
session = self.sessions.get_or_create(event.context.session_key)
|
||||
mark_webui_session(session, event.context.metadata)
|
||||
session = self.sessions.get_cached(event.context.session_key)
|
||||
if session is not None:
|
||||
mark_webui_session(session, event.context.metadata)
|
||||
|
||||
async def _handle_run_status_changed(self, event: TurnRunStatusChanged) -> None:
|
||||
if not self._is_websocket_event(event.context):
|
||||
@@ -695,13 +700,13 @@ class WebuiTurnCoordinator:
|
||||
*,
|
||||
session_key: str,
|
||||
latency_ms: int | None,
|
||||
usage: dict[str, int] | None = None,
|
||||
usage: LLMUsage | None = None,
|
||||
context_window_tokens: int | None = None,
|
||||
) -> None:
|
||||
if msg.channel != "websocket":
|
||||
return
|
||||
|
||||
session = self.sessions.get_or_create(session_key)
|
||||
session = await self.sessions.get_or_create_async(session_key)
|
||||
await self.bus.publish_outbound(
|
||||
outbound_message_for_event(
|
||||
channel=msg.channel,
|
||||
@@ -709,7 +714,7 @@ class WebuiTurnCoordinator:
|
||||
event=TurnEndEvent(
|
||||
latency_ms=latency_ms,
|
||||
goal_state=goal_state_ws_blob(session.metadata),
|
||||
usage=usage or None,
|
||||
usage=usage,
|
||||
context_window_tokens=context_window_tokens,
|
||||
),
|
||||
metadata=msg.metadata,
|
||||
|
||||
@@ -16,7 +16,7 @@ Concrete scenarios showing when and how to use the my tool effectively.
|
||||
→ my(action="check", key="max_iterations")
|
||||
→ 40
|
||||
→ 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."
|
||||
```
|
||||
|
||||
@@ -72,6 +72,6 @@ Concrete scenarios showing when and how to use the my tool effectively.
|
||||
### Token-conscious behavior
|
||||
```
|
||||
→ 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."
|
||||
```
|
||||
|
||||
@@ -5,17 +5,24 @@ from __future__ import annotations
|
||||
import asyncio
|
||||
import uuid
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import Any
|
||||
from contextlib import suppress
|
||||
from typing import Any, TypeVar
|
||||
|
||||
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.triggers.local_session_turns import LOCAL_TRIGGER_META
|
||||
from nanobot.triggers.local_store import LocalTriggerStore
|
||||
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
|
||||
|
||||
_T = TypeVar("_T")
|
||||
|
||||
|
||||
async def run_local_trigger_queue(
|
||||
*,
|
||||
@@ -29,14 +36,16 @@ async def run_local_trigger_queue(
|
||||
if submit_turn is None:
|
||||
raise ValueError("run_local_trigger_queue requires submit_turn")
|
||||
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:
|
||||
logger.warning(
|
||||
"Trigger: recovered {} interrupted delivery file(s) from processing",
|
||||
recovered,
|
||||
)
|
||||
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:
|
||||
await asyncio.sleep(poll_interval_s)
|
||||
continue
|
||||
@@ -49,30 +58,24 @@ async def run_local_trigger_queue(
|
||||
submit_turn=submit_turn,
|
||||
is_channel_enabled=is_channel_enabled,
|
||||
)
|
||||
store.complete_delivery(delivery)
|
||||
except _DeliverySettledOnCancellation:
|
||||
raise
|
||||
except asyncio.CancelledError as exc:
|
||||
store.retry_delivery(delivery, str(exc) or exc.__class__.__name__)
|
||||
_write_delivery_run_record(
|
||||
error = str(exc) or exc.__class__.__name__
|
||||
await shield_and_drain(asyncio.to_thread(store.retry_delivery, delivery, error))
|
||||
await _write_delivery_run_record(
|
||||
store,
|
||||
delivery,
|
||||
status="interrupted",
|
||||
error=str(exc) or exc.__class__.__name__,
|
||||
error=error,
|
||||
)
|
||||
raise
|
||||
except _TerminalDeliveryError as exc:
|
||||
store.record_delivery(
|
||||
delivery.trigger_id,
|
||||
status="error",
|
||||
error=str(exc),
|
||||
run_at_ms=delivery.created_at_ms,
|
||||
await _await_delivery_settlement(
|
||||
_settle_failed_delivery(store, delivery, error=str(exc)),
|
||||
store=store,
|
||||
delivery=delivery,
|
||||
)
|
||||
_write_delivery_run_record(
|
||||
store,
|
||||
delivery,
|
||||
status="error",
|
||||
error=str(exc),
|
||||
)
|
||||
store.complete_delivery(delivery)
|
||||
logger.warning(
|
||||
"Trigger: dropped delivery {} for {}: {}",
|
||||
delivery.id,
|
||||
@@ -81,19 +84,11 @@ async def run_local_trigger_queue(
|
||||
)
|
||||
except AutomationTurnError as exc:
|
||||
error = str(exc) or exc.__class__.__name__
|
||||
store.record_delivery(
|
||||
delivery.trigger_id,
|
||||
status="error",
|
||||
error=error,
|
||||
run_at_ms=delivery.created_at_ms,
|
||||
await _await_delivery_settlement(
|
||||
_settle_failed_delivery(store, delivery, error=error),
|
||||
store=store,
|
||||
delivery=delivery,
|
||||
)
|
||||
_write_delivery_run_record(
|
||||
store,
|
||||
delivery,
|
||||
status="error",
|
||||
error=error,
|
||||
)
|
||||
store.complete_delivery(delivery)
|
||||
logger.warning(
|
||||
"Trigger: delivery {} for {} reached the agent but failed: {}",
|
||||
delivery.id,
|
||||
@@ -102,18 +97,10 @@ async def run_local_trigger_queue(
|
||||
)
|
||||
except Exception as exc:
|
||||
error = str(exc) or exc.__class__.__name__
|
||||
retried = store.retry_delivery(delivery, error)
|
||||
_write_delivery_run_record(
|
||||
store,
|
||||
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,
|
||||
retried = await _await_delivery_settlement(
|
||||
_settle_retryable_delivery(store, delivery, error=error),
|
||||
store=store,
|
||||
delivery=delivery,
|
||||
)
|
||||
logger.exception(
|
||||
"Trigger: failed delivery {} for {}{}",
|
||||
@@ -127,6 +114,10 @@ class _TerminalDeliveryError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
class _DeliverySettledOnCancellation(asyncio.CancelledError):
|
||||
"""Cancellation reported only after an already-submitted delivery is settled."""
|
||||
|
||||
|
||||
async def _deliver_delivery(
|
||||
store: LocalTriggerStore,
|
||||
delivery: TriggerDelivery,
|
||||
@@ -134,7 +125,7 @@ async def _deliver_delivery(
|
||||
submit_turn: Callable[[InboundMessage], Awaitable[OutboundMessage | None]],
|
||||
is_channel_enabled: Callable[[str], bool],
|
||||
) -> None:
|
||||
trigger = store.get(delivery.trigger_id)
|
||||
trigger = await asyncio.to_thread(store.get, delivery.trigger_id)
|
||||
if trigger is None:
|
||||
raise _TerminalDeliveryError("trigger not found")
|
||||
if not trigger.enabled:
|
||||
@@ -142,7 +133,14 @@ async def _deliver_delivery(
|
||||
if not is_channel_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(
|
||||
channel=trigger.channel,
|
||||
sender_id=trigger.sender_id,
|
||||
@@ -151,22 +149,177 @@ async def _deliver_delivery(
|
||||
metadata=_delivery_metadata(trigger, delivery),
|
||||
session_key_override=trigger.session_key,
|
||||
)
|
||||
response = await submit_turn(msg)
|
||||
store.record_delivery(
|
||||
trigger.id,
|
||||
status="ok",
|
||||
run_at_ms=delivery.created_at_ms,
|
||||
try:
|
||||
response = await submit_turn(msg)
|
||||
except AutomationTurnAcceptedCancellation:
|
||||
try:
|
||||
await _await_delivery_settlement(
|
||||
_settle_accepted_delivery(store, delivery, trigger=trigger),
|
||||
store=store,
|
||||
delivery=delivery,
|
||||
)
|
||||
except _DeliverySettledOnCancellation:
|
||||
raise
|
||||
except Exception:
|
||||
logger.exception(
|
||||
"Trigger: failed to persist accepted delivery {}; dropping retry",
|
||||
delivery.id,
|
||||
)
|
||||
with suppress(Exception):
|
||||
await shield_and_drain(asyncio.to_thread(store.complete_delivery, delivery))
|
||||
raise _DeliverySettledOnCancellation from None
|
||||
|
||||
try:
|
||||
await _await_delivery_settlement(
|
||||
_settle_submitted_delivery(store, delivery, trigger=trigger, response=response),
|
||||
store=store,
|
||||
delivery=delivery,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception(
|
||||
"Trigger: failed to persist status for submitted delivery {}; dropping retry",
|
||||
delivery.id,
|
||||
)
|
||||
with suppress(Exception):
|
||||
await shield_and_drain(asyncio.to_thread(store.complete_delivery, delivery))
|
||||
|
||||
|
||||
async def _await_delivery_settlement(
|
||||
operation: Awaitable[_T],
|
||||
*,
|
||||
store: LocalTriggerStore,
|
||||
delivery: TriggerDelivery,
|
||||
) -> _T:
|
||||
settlement = asyncio.ensure_future(operation)
|
||||
try:
|
||||
return await asyncio.shield(settlement)
|
||||
except asyncio.CancelledError:
|
||||
while not settlement.done():
|
||||
try:
|
||||
await asyncio.shield(settlement)
|
||||
except asyncio.CancelledError:
|
||||
continue
|
||||
try:
|
||||
settlement.result()
|
||||
except Exception:
|
||||
logger.exception(
|
||||
"Trigger: failed to settle delivery {} during cancellation",
|
||||
delivery.id,
|
||||
)
|
||||
completion = asyncio.create_task(
|
||||
shield_and_drain(asyncio.to_thread(store.complete_delivery, delivery))
|
||||
)
|
||||
while not completion.done():
|
||||
try:
|
||||
await asyncio.shield(completion)
|
||||
except asyncio.CancelledError:
|
||||
continue
|
||||
with suppress(Exception):
|
||||
completion.result()
|
||||
raise _DeliverySettledOnCancellation from None
|
||||
|
||||
|
||||
async def _settle_failed_delivery(
|
||||
store: LocalTriggerStore,
|
||||
delivery: TriggerDelivery,
|
||||
*,
|
||||
error: str,
|
||||
) -> None:
|
||||
await _write_delivery_run_record(
|
||||
store,
|
||||
delivery,
|
||||
status="error",
|
||||
error=error,
|
||||
)
|
||||
_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,
|
||||
delivery,
|
||||
trigger=trigger,
|
||||
status="ok",
|
||||
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,
|
||||
delivery: TriggerDelivery,
|
||||
*,
|
||||
@@ -176,12 +329,15 @@ def _write_delivery_run_record(
|
||||
response: str | None = None,
|
||||
) -> None:
|
||||
try:
|
||||
store.write_delivery_run_record(
|
||||
delivery,
|
||||
trigger=trigger,
|
||||
status=status,
|
||||
error=error,
|
||||
response=response,
|
||||
await shield_and_drain(
|
||||
asyncio.to_thread(
|
||||
store.write_delivery_run_record,
|
||||
delivery,
|
||||
trigger=trigger,
|
||||
status=status,
|
||||
error=error,
|
||||
response=response,
|
||||
)
|
||||
)
|
||||
except Exception:
|
||||
logger.exception(
|
||||
|
||||
@@ -24,6 +24,7 @@ _MAX_RUN_HISTORY = 20
|
||||
_MAX_DELIVERY_ATTEMPTS = 10
|
||||
_RUN_RECORD_TEXT_MAX_CHARS = 4000
|
||||
_PROCESSING_RECOVERY_ERROR = "delivery was recovered from interrupted processing"
|
||||
_FILE_LOCK_TIMEOUT_SECONDS = 5
|
||||
|
||||
|
||||
class TriggerStoreError(RuntimeError):
|
||||
@@ -49,7 +50,10 @@ class LocalTriggerStore:
|
||||
self.processing_dir = self.root / "processing"
|
||||
self.failed_dir = self.root / "failed"
|
||||
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(
|
||||
self,
|
||||
|
||||
@@ -3,8 +3,48 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Awaitable
|
||||
from typing import TypeVar
|
||||
|
||||
_T = TypeVar("_T")
|
||||
|
||||
|
||||
def task_is_cancelling() -> bool:
|
||||
task = asyncio.current_task()
|
||||
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()
|
||||
|
||||
@@ -1,5 +1,7 @@
|
||||
"""Utility functions for nanobot."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import json
|
||||
import os
|
||||
@@ -12,11 +14,14 @@ from contextlib import suppress
|
||||
from datetime import datetime
|
||||
from functools import lru_cache
|
||||
from pathlib import Path
|
||||
from typing import Any, TypeVar, cast, overload
|
||||
from typing import TYPE_CHECKING, Any, TypeVar, cast, overload
|
||||
|
||||
import tiktoken
|
||||
from loguru import logger
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from nanobot.providers.base import LLMUsage
|
||||
|
||||
_TOOLS_TOKEN_CACHE_MAX_ENTRIES = 64
|
||||
_TOOLS_TOKEN_CACHE: dict[int, tuple[tuple[int, ...], dict[bool, int]]] = {}
|
||||
_T = TypeVar("_T")
|
||||
@@ -793,7 +798,7 @@ def build_status_content(
|
||||
version: str,
|
||||
model: str,
|
||||
start_time: float,
|
||||
last_usage: dict[str, int],
|
||||
last_usage: LLMUsage | None,
|
||||
context_window_tokens: int,
|
||||
session_msg_count: int,
|
||||
context_tokens_estimate: int,
|
||||
@@ -814,9 +819,9 @@ def build_status_content(
|
||||
if uptime_s >= 3600
|
||||
else f"{uptime_s // 60}m {uptime_s % 60}s"
|
||||
)
|
||||
last_in = last_usage.get("prompt_tokens", 0)
|
||||
last_out = last_usage.get("completion_tokens", 0)
|
||||
cached = last_usage.get("cached_tokens", 0)
|
||||
last_in = last_usage.input_tokens if last_usage else 0
|
||||
last_out = last_usage.output_tokens if last_usage else 0
|
||||
cached = last_usage.cache_read_tokens if last_usage else None
|
||||
ctx_total = max(context_window_tokens, 0)
|
||||
# Budget mirrors Consolidator formula: ctx_window - max_completion - _SAFETY_BUFFER
|
||||
ctx_budget = max(ctx_total - int(max_completion_tokens) - 1024, 1)
|
||||
|
||||
+32
-19
@@ -2,6 +2,7 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import re
|
||||
import uuid
|
||||
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.webui_turns import WEBUI_TITLE_METADATA_KEY, clean_generated_title
|
||||
from nanobot.utils.cancellation import shield_and_drain
|
||||
from nanobot.webui.transcript import (
|
||||
append_fork_marker,
|
||||
delete_webui_transcript,
|
||||
@@ -93,24 +95,35 @@ async def handle_webui_fork_chat(
|
||||
await channel.send_webui_protocol_error(connection, "session_manager_unavailable")
|
||||
return
|
||||
|
||||
try:
|
||||
forked = create_webui_chat_fork(
|
||||
session_manager,
|
||||
source_chat_id=source_chat_id,
|
||||
before_user_index=raw_index,
|
||||
title=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")
|
||||
async def create_and_attach() -> None:
|
||||
try:
|
||||
forked = await asyncio.to_thread(
|
||||
create_webui_chat_fork,
|
||||
session_manager,
|
||||
source_chat_id=source_chat_id,
|
||||
before_user_index=raw_index,
|
||||
title=(
|
||||
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
|
||||
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(
|
||||
connection,
|
||||
fork_id=fork_id,
|
||||
fork_key=fork_key,
|
||||
)
|
||||
await channel.attach_webui_fork(
|
||||
connection,
|
||||
fork_id=fork_id,
|
||||
fork_key=fork_key,
|
||||
)
|
||||
|
||||
await shield_and_drain(create_and_attach())
|
||||
|
||||
@@ -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.registry import ToolRegistry
|
||||
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
|
||||
|
||||
McpReload = Callable[[], Awaitable[dict[str, Any]]]
|
||||
@@ -259,7 +259,7 @@ class McpOAuthManager:
|
||||
):
|
||||
flow.error = "The MCP server returned an unsafe authorization URL."
|
||||
raise McpOAuthError(flow.error)
|
||||
ok, _error = validate_url_target(authorization_url)
|
||||
ok, _error = await async_validate_url_target(authorization_url)
|
||||
if not ok:
|
||||
flow.error = "The MCP server returned an unsafe authorization URL."
|
||||
raise McpOAuthError(flow.error)
|
||||
|
||||
@@ -4,6 +4,7 @@ from __future__ import annotations
|
||||
|
||||
from typing import Any, cast
|
||||
|
||||
from nanobot.providers.base import LLMUsage
|
||||
from nanobot.session.manager import Session
|
||||
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 = (
|
||||
estimate_message_tokens({"role": "system", "content": summary}) if summary else 0
|
||||
)
|
||||
raw_usage = session.metadata.get("_last_usage")
|
||||
last_usage = (
|
||||
{
|
||||
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
|
||||
)
|
||||
stored_usage = LLMUsage.from_dict(session.metadata.get("_last_usage"))
|
||||
last_usage = stored_usage.to_turn_dict() if stored_usage is not None else None
|
||||
|
||||
return {
|
||||
"schema_version": 1,
|
||||
|
||||
@@ -27,6 +27,7 @@ from nanobot.providers.image_generation import (
|
||||
)
|
||||
from nanobot.providers.registry import find_by_name
|
||||
from nanobot.security.network import is_loopback_host
|
||||
from nanobot.utils.cancellation import shield_and_drain
|
||||
from nanobot.webui.settings_contracts import (
|
||||
QueryParams,
|
||||
SettingsRequest,
|
||||
@@ -640,7 +641,11 @@ class CapabilitySettingsHandler:
|
||||
) -> SettingsRouteResult:
|
||||
if action == "api-status":
|
||||
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":
|
||||
return await self._start_api(request, operations)
|
||||
@@ -673,17 +678,22 @@ class CapabilitySettingsHandler:
|
||||
return SettingsRouteResult.failure(404, "unknown settings action")
|
||||
|
||||
operation, section, apply_image_reload = mutation
|
||||
try:
|
||||
payload = self.settings.mutate(operation, request.query)
|
||||
except WebUISettingsError as exc:
|
||||
return SettingsRouteResult.failure(exc.status, exc.message)
|
||||
if apply_image_reload:
|
||||
payload, image_restart_cleared = await self.apply_image_runtime_change(
|
||||
|
||||
async def mutate_and_apply() -> tuple[dict[str, Any], bool]:
|
||||
payload = await self.settings.mutate_async(operation, request.query)
|
||||
if not apply_image_reload:
|
||||
return payload, False
|
||||
return await self.apply_image_runtime_change(
|
||||
payload,
|
||||
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(
|
||||
payload,
|
||||
decorate_restart=True,
|
||||
@@ -726,16 +736,17 @@ class CapabilitySettingsHandler:
|
||||
400,
|
||||
"API service API key must be a string",
|
||||
)
|
||||
try:
|
||||
await asyncio.to_thread(
|
||||
self.settings.mutate,
|
||||
allow_install = await self._allow_feature_package_install(request)
|
||||
|
||||
async def mutate_and_start() -> Any:
|
||||
await self.settings.mutate_async(
|
||||
operations.nanobot_features_action,
|
||||
"enable",
|
||||
{"name": ["api"]},
|
||||
allow_install=self._allow_feature_package_install(request),
|
||||
allow_install=allow_install,
|
||||
)
|
||||
self.settings.mutate(operations.update_api, request.query)
|
||||
config = self.settings.config.load()
|
||||
await self.settings.mutate_async(operations.update_api, request.query)
|
||||
config = await self.settings.config.load_async()
|
||||
runtime = operations.api_runtime()
|
||||
options = ApiStartOptions(
|
||||
host=config.api.host,
|
||||
@@ -744,10 +755,13 @@ class CapabilitySettingsHandler:
|
||||
config_path=str(self.settings.config.path),
|
||||
)
|
||||
current = runtime.status()
|
||||
result = await asyncio.to_thread(
|
||||
return await asyncio.to_thread(
|
||||
runtime.restart if current.running else runtime.start_background,
|
||||
options,
|
||||
)
|
||||
|
||||
try:
|
||||
result = await shield_and_drain(mutate_and_start())
|
||||
if not result.ok:
|
||||
return SettingsRouteResult.failure(
|
||||
500,
|
||||
@@ -762,7 +776,8 @@ class CapabilitySettingsHandler:
|
||||
self.logger.exception("failed to start managed API service")
|
||||
return SettingsRouteResult.failure(500, str(exc))
|
||||
return SettingsRouteResult.success(
|
||||
api_service_payload(
|
||||
await asyncio.to_thread(
|
||||
api_service_payload,
|
||||
self.settings,
|
||||
operations.api_runtime(),
|
||||
last_action="started",
|
||||
@@ -775,7 +790,7 @@ class CapabilitySettingsHandler:
|
||||
) -> SettingsRouteResult:
|
||||
runtime = operations.api_runtime()
|
||||
try:
|
||||
result = await asyncio.to_thread(runtime.stop)
|
||||
result = await shield_and_drain(asyncio.to_thread(runtime.stop))
|
||||
except Exception as exc:
|
||||
self.logger.exception("failed to stop managed API service")
|
||||
return SettingsRouteResult.failure(500, str(exc))
|
||||
@@ -785,20 +800,20 @@ class CapabilitySettingsHandler:
|
||||
api_runtime_message(result.message),
|
||||
)
|
||||
return SettingsRouteResult.success(
|
||||
api_service_payload(
|
||||
await asyncio.to_thread(
|
||||
api_service_payload,
|
||||
self.settings,
|
||||
operations.api_runtime(),
|
||||
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:
|
||||
return True
|
||||
try:
|
||||
return bool(
|
||||
self.settings.config.load().tools.webui_allow_remote_package_install
|
||||
)
|
||||
config = await self.settings.config.load_async()
|
||||
return bool(config.tools.webui_allow_remote_package_install)
|
||||
except Exception:
|
||||
self.logger.exception("failed to load remote package install policy")
|
||||
return False
|
||||
|
||||
@@ -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.oauth_guidance import OAUTH_CLI_KIT_MISSING_MESSAGE
|
||||
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 (
|
||||
QueryParams,
|
||||
SettingsRequest,
|
||||
@@ -1651,6 +1652,30 @@ class ModelSettingsHandler:
|
||||
if self.settings.refresh_runtime_config is not None:
|
||||
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(
|
||||
self,
|
||||
action: str,
|
||||
@@ -1659,8 +1684,12 @@ class ModelSettingsHandler:
|
||||
) -> SettingsRouteResult:
|
||||
try:
|
||||
if action == "agent-update":
|
||||
payload = self.settings.mutate(operations.update_agent, request.query)
|
||||
self._refresh_runtime_config()
|
||||
payload = await shield_and_drain(
|
||||
self._mutate_and_refresh(
|
||||
operations.update_agent,
|
||||
request.query,
|
||||
)
|
||||
)
|
||||
return SettingsRouteResult.success(
|
||||
payload,
|
||||
decorate_restart=True,
|
||||
@@ -1668,12 +1697,13 @@ class ModelSettingsHandler:
|
||||
)
|
||||
|
||||
if action == "model-update":
|
||||
payload = self.settings.mutate(
|
||||
operations.update_model,
|
||||
request.query,
|
||||
rename_model_preset=self.settings.rename_model_preset,
|
||||
payload = await shield_and_drain(
|
||||
self._mutate_and_refresh(
|
||||
operations.update_model,
|
||||
request.query,
|
||||
rename_model_preset=self.settings.rename_model_preset,
|
||||
)
|
||||
)
|
||||
self._refresh_runtime_config()
|
||||
return SettingsRouteResult.success(payload, decorate_restart=True)
|
||||
|
||||
mutation = {
|
||||
@@ -1684,19 +1714,19 @@ class ModelSettingsHandler:
|
||||
"provider-create": operations.create_provider,
|
||||
}.get(action)
|
||||
if mutation is not None:
|
||||
payload = self.settings.mutate(mutation, request.query)
|
||||
self._refresh_runtime_config()
|
||||
payload = await shield_and_drain(
|
||||
self._mutate_and_refresh(mutation, request.query)
|
||||
)
|
||||
return SettingsRouteResult.success(payload, decorate_restart=True)
|
||||
|
||||
if action == "provider-update":
|
||||
payload = self.settings.mutate(
|
||||
operations.update_provider,
|
||||
request.query,
|
||||
payload, image_restart_cleared = await shield_and_drain(
|
||||
self._update_provider_and_runtime(
|
||||
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(
|
||||
payload,
|
||||
decorate_restart=True,
|
||||
@@ -1724,11 +1754,13 @@ class ModelSettingsHandler:
|
||||
return SettingsRouteResult.success(payload)
|
||||
|
||||
if action == "oauth-login":
|
||||
payload = await asyncio.to_thread(
|
||||
self.settings.read,
|
||||
operations.oauth_login,
|
||||
request.query,
|
||||
oauth_flows=self.settings.oauth_flows,
|
||||
payload = await shield_and_drain(
|
||||
asyncio.to_thread(
|
||||
self.settings.read,
|
||||
operations.oauth_login,
|
||||
request.query,
|
||||
oauth_flows=self.settings.oauth_flows,
|
||||
)
|
||||
)
|
||||
elif action == "oauth-complete":
|
||||
raw_response = (request.payload or {}).get("authorization_response")
|
||||
@@ -1736,19 +1768,23 @@ class ModelSettingsHandler:
|
||||
raise WebUISettingsError(
|
||||
"OAuth authorization response must be a string"
|
||||
)
|
||||
payload = await asyncio.to_thread(
|
||||
self.settings.read,
|
||||
operations.oauth_complete,
|
||||
request.query,
|
||||
raw_response or None,
|
||||
oauth_flows=self.settings.oauth_flows,
|
||||
payload = await shield_and_drain(
|
||||
asyncio.to_thread(
|
||||
self.settings.read,
|
||||
operations.oauth_complete,
|
||||
request.query,
|
||||
raw_response or None,
|
||||
oauth_flows=self.settings.oauth_flows,
|
||||
)
|
||||
)
|
||||
elif action == "oauth-logout":
|
||||
payload = await asyncio.to_thread(
|
||||
self.settings.read,
|
||||
operations.oauth_logout,
|
||||
request.query,
|
||||
oauth_flows=self.settings.oauth_flows,
|
||||
payload = await shield_and_drain(
|
||||
asyncio.to_thread(
|
||||
self.settings.read,
|
||||
operations.oauth_logout,
|
||||
request.query,
|
||||
oauth_flows=self.settings.oauth_flows,
|
||||
)
|
||||
)
|
||||
else:
|
||||
return SettingsRouteResult.failure(404, "unknown settings action")
|
||||
|
||||
@@ -4,6 +4,7 @@ from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import html
|
||||
import inspect
|
||||
import json
|
||||
from collections.abc import Awaitable, Callable, Mapping
|
||||
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.validation import validate_channel_config
|
||||
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_contracts as contracts
|
||||
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:
|
||||
"""Authenticate and dispatch settings requests to transport-neutral domains."""
|
||||
|
||||
@@ -284,9 +296,9 @@ class WebUISettingsRouter:
|
||||
if not self._authorized(request):
|
||||
return self._unauthorized()
|
||||
if route == ("root", "settings"):
|
||||
return self._handle_settings()
|
||||
return await _call_settings_handler(self._handle_settings)
|
||||
if route == ("root", "usage"):
|
||||
return self._handle_settings_usage()
|
||||
return await _call_settings_handler(self._handle_settings_usage)
|
||||
|
||||
domain, action = route
|
||||
domain_request = self._domain_request(
|
||||
@@ -415,19 +427,18 @@ class WebUISettingsRouter:
|
||||
)
|
||||
return self._json_response(payload)
|
||||
|
||||
def _handle_settings(self) -> Response:
|
||||
return self._json_response(
|
||||
self._with_restart_state(
|
||||
self.settings.read(
|
||||
settings_payload,
|
||||
surface=self._runtime_surface,
|
||||
runtime_capability_overrides=self._runtime_capabilities,
|
||||
)
|
||||
)
|
||||
async def _handle_settings(self) -> Response:
|
||||
payload = await self.settings.read_async(
|
||||
settings_payload,
|
||||
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:
|
||||
return self._json_response(self.settings.read(settings_usage_payload))
|
||||
async def _handle_settings_usage(self) -> Response:
|
||||
return self._json_response(
|
||||
await self.settings.read_async(settings_usage_payload)
|
||||
)
|
||||
|
||||
def _model_operations(self) -> model_domain.ModelSettingsOperations:
|
||||
return model_domain.ModelSettingsOperations(
|
||||
@@ -560,7 +571,7 @@ class WebUISettingsRouter:
|
||||
allow_install=allow_install,
|
||||
)
|
||||
|
||||
def _allow_feature_package_install(
|
||||
async def _allow_feature_package_install(
|
||||
self,
|
||||
connection: Any,
|
||||
request: WsRequest,
|
||||
@@ -570,29 +581,33 @@ class WebUISettingsRouter:
|
||||
request,
|
||||
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:
|
||||
if not self._authorized(request):
|
||||
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")
|
||||
query = self._parse_mcp_settings_query(request)
|
||||
try:
|
||||
name, cfg = await asyncio.to_thread(
|
||||
self.settings.mutate,
|
||||
|
||||
async def mutate_and_start() -> dict[str, Any]:
|
||||
name, cfg = await self.settings.mutate_async(
|
||||
ensure_mcp_oauth_server,
|
||||
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"}
|
||||
payload = await self._mcp_oauth.start(
|
||||
return await self._mcp_oauth.start(
|
||||
name,
|
||||
cfg,
|
||||
redirect_uri,
|
||||
reload_mcp=self._reload_mcp_runtime,
|
||||
reset_credentials=reset,
|
||||
)
|
||||
|
||||
try:
|
||||
payload = await shield_and_drain(mutate_and_start())
|
||||
except Exception as exc:
|
||||
return self._mcp_oauth_error_response(exc, action="start")
|
||||
return self._json_response(payload)
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import threading
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
@@ -12,9 +13,11 @@ from filelock import FileLock
|
||||
|
||||
from nanobot.config.loader import load_config, save_config
|
||||
from nanobot.config.schema import Config
|
||||
from nanobot.utils.cancellation import shield_and_drain
|
||||
|
||||
_T = TypeVar("_T")
|
||||
_WEBUI_OAUTH_MAX_FLOWS = 8
|
||||
_SETTINGS_FILE_LOCK_TIMEOUT_SECONDS = 5
|
||||
|
||||
|
||||
class WebUISettingsConfig:
|
||||
@@ -25,13 +28,20 @@ class WebUISettingsConfig:
|
||||
self.path.parent.mkdir(parents=True, exist_ok=True)
|
||||
self._lock = threading.RLock()
|
||||
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:
|
||||
"""Load this gateway's config without consulting the process-global path."""
|
||||
with self._lock:
|
||||
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:
|
||||
"""Apply and atomically persist one path-scoped read-modify-write operation."""
|
||||
with self._lock, self._file_lock:
|
||||
@@ -40,11 +50,21 @@ class WebUISettingsConfig:
|
||||
save_config(config, self.path)
|
||||
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:
|
||||
"""Run a path-aware read-modify-write operation under the config-file lock."""
|
||||
with self._lock, self._file_lock:
|
||||
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:
|
||||
"""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."""
|
||||
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(
|
||||
self,
|
||||
operation: Callable[..., _T],
|
||||
@@ -161,3 +191,15 @@ class WebUISettingsServices:
|
||||
**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)
|
||||
)
|
||||
|
||||
@@ -20,8 +20,10 @@ from nanobot.channels.contracts import (
|
||||
channel_update_instance_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.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_contracts import (
|
||||
QueryParams,
|
||||
@@ -31,7 +33,6 @@ from nanobot.webui.settings_contracts import (
|
||||
query_first,
|
||||
query_first_alias,
|
||||
)
|
||||
from nanobot.webui.token_usage import token_usage_payload
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from nanobot.webui.settings_services import WebUISettingsServices
|
||||
@@ -121,7 +122,7 @@ def system_settings_payload(
|
||||
},
|
||||
"unified_session": defaults.unified_session,
|
||||
},
|
||||
"usage": token_usage_payload(timezone_name=defaults.timezone),
|
||||
"usage": llm_usage_payload(timezone_name=defaults.timezone),
|
||||
"advanced": {
|
||||
"restrict_to_workspace": config.tools.restrict_to_workspace,
|
||||
"workspace_sandbox": sandbox_status.as_dict(),
|
||||
@@ -139,7 +140,7 @@ def system_settings_payload(
|
||||
|
||||
def settings_usage_payload(config: Config) -> dict[str, Any]:
|
||||
"""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]:
|
||||
@@ -446,12 +447,17 @@ class SystemSettingsHandler:
|
||||
operations: SystemSettingsOperations,
|
||||
) -> SettingsRouteResult:
|
||||
try:
|
||||
payload = await asyncio.to_thread(
|
||||
pending = asyncio.to_thread(
|
||||
operations.cli_apps_action,
|
||||
action,
|
||||
request.query,
|
||||
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:
|
||||
return SettingsRouteResult.failure(exc.status, exc.message)
|
||||
except Exception as exc:
|
||||
@@ -505,17 +511,29 @@ class SystemSettingsHandler:
|
||||
action: str,
|
||||
operations: SystemSettingsOperations,
|
||||
) -> 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(
|
||||
self._nanobot_features_action,
|
||||
action,
|
||||
request.query,
|
||||
operations,
|
||||
allow_install=(
|
||||
action != "enable"
|
||||
or self.allow_feature_package_install(request)
|
||||
),
|
||||
allow_install=allow_install,
|
||||
)
|
||||
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:
|
||||
return SettingsRouteResult.failure(exc.status, exc.message)
|
||||
except Exception as exc:
|
||||
@@ -527,13 +545,6 @@ class SystemSettingsHandler:
|
||||
action,
|
||||
)
|
||||
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(
|
||||
payload,
|
||||
decorate_restart=True,
|
||||
@@ -628,6 +639,15 @@ class SystemSettingsHandler:
|
||||
self,
|
||||
request: SettingsRequest,
|
||||
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:
|
||||
name = (query_first(request.query, "name") or "").strip()
|
||||
instance_id = (
|
||||
@@ -682,7 +702,7 @@ class SystemSettingsHandler:
|
||||
"enable",
|
||||
feature_query,
|
||||
operations,
|
||||
allow_install=self.allow_feature_package_install(request),
|
||||
allow_install=await self.allow_feature_package_install(request),
|
||||
)
|
||||
except OptionalFeatureError as exc:
|
||||
return SettingsRouteResult.failure(
|
||||
@@ -825,6 +845,22 @@ class SystemSettingsHandler:
|
||||
channel_name: str,
|
||||
payload: dict[str, Any],
|
||||
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]:
|
||||
target = {"name": [channel_name]}
|
||||
if payload.get("instance_id"):
|
||||
@@ -835,11 +871,11 @@ class SystemSettingsHandler:
|
||||
"enable",
|
||||
target,
|
||||
operations,
|
||||
allow_install=self.allow_feature_package_install(request),
|
||||
allow_install=await self.allow_feature_package_install(request),
|
||||
)
|
||||
except OptionalFeatureError as exc:
|
||||
features = self.feature_runtime_fallback(
|
||||
self._nanobot_features_payload(operations),
|
||||
await asyncio.to_thread(self._nanobot_features_payload, operations),
|
||||
message=(
|
||||
f"{channel_name} connected, but enabling channel support failed: "
|
||||
f"{exc.message}"
|
||||
@@ -859,13 +895,12 @@ class SystemSettingsHandler:
|
||||
)
|
||||
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:
|
||||
return True
|
||||
try:
|
||||
return bool(
|
||||
self.settings.config.load().tools.webui_allow_remote_package_install
|
||||
)
|
||||
config = await self.settings.config.load_async()
|
||||
return bool(config.tools.webui_allow_remote_package_install)
|
||||
except Exception:
|
||||
self.logger.exception("failed to load remote package install policy")
|
||||
return False
|
||||
@@ -925,13 +960,18 @@ class SystemSettingsHandler:
|
||||
operations: SystemSettingsOperations,
|
||||
) -> SettingsRouteResult:
|
||||
try:
|
||||
payload = await operations.mcp_presets_action(
|
||||
pending = operations.mcp_presets_action(
|
||||
action,
|
||||
request.query,
|
||||
reload_mcp=operations.reload_mcp,
|
||||
mcp_runtime_status=operations.mcp_runtime_status,
|
||||
config=self.settings.config,
|
||||
)
|
||||
payload = (
|
||||
await pending
|
||||
if action is None
|
||||
else await shield_and_drain(pending)
|
||||
)
|
||||
except Exception as exc:
|
||||
status = getattr(exc, "status", 500)
|
||||
message = getattr(exc, "message", str(exc))
|
||||
|
||||
@@ -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")
|
||||
@@ -1813,11 +1813,16 @@ def replay_transcript_to_ui_messages(
|
||||
break
|
||||
content = str(candidate.get("content") or "")
|
||||
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 (
|
||||
candidate.get("reasoningStreaming")
|
||||
or candidate.get("reasoning") is not None
|
||||
or has_answer
|
||||
or candidate.get("isStreaming")
|
||||
or (
|
||||
candidate.get("isStreaming")
|
||||
and candidate.get("reasoning") is None
|
||||
)
|
||||
):
|
||||
prev[i] = {
|
||||
**candidate,
|
||||
@@ -1827,15 +1832,6 @@ def replay_transcript_to_ui_messages(
|
||||
**turn_fields,
|
||||
}
|
||||
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
|
||||
segment = _ensure_activity_segment()
|
||||
prev.append(
|
||||
@@ -1915,19 +1911,6 @@ def replay_transcript_to_ui_messages(
|
||||
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(
|
||||
*,
|
||||
latency_ms: int | None = None,
|
||||
@@ -2442,7 +2425,6 @@ def replay_transcript_to_ui_messages(
|
||||
for i, m in enumerate(messages):
|
||||
if m.get("isStreaming"):
|
||||
messages[i] = {**m, "isStreaming": False}
|
||||
prune_reasoning_only()
|
||||
lat = rec.get("latency_ms")
|
||||
usage = rec.get("usage")
|
||||
sanitized_usage = (
|
||||
|
||||
@@ -5,6 +5,7 @@ from __future__ import annotations
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
from collections import OrderedDict
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any, cast
|
||||
|
||||
@@ -28,6 +29,7 @@ _MAX_STATE_FILE_BYTES = 128 * 1024
|
||||
_DEFAULT_ACCESS_MODES = {"default", "full"}
|
||||
_LEGACY_RESTRICTED_DEFAULT_ACCESS_MODE = "restricted"
|
||||
_WEBUI_SCOPE_CHANNEL = "websocket"
|
||||
_MAX_DRAFT_SCOPES = 128
|
||||
|
||||
|
||||
def _scope_change_is_non_escalating(current: WorkspaceScope, requested: WorkspaceScope) -> bool:
|
||||
@@ -186,6 +188,7 @@ class WebUIWorkspaceController:
|
||||
self._sessions = session_manager
|
||||
self._default_workspace = default_workspace
|
||||
self._default_restrict_to_workspace = default_restrict_to_workspace
|
||||
self._draft_scopes: OrderedDict[str, WorkspaceScope] = OrderedDict()
|
||||
|
||||
def default_scope(self) -> WorkspaceScope:
|
||||
return default_scope_for_webui(
|
||||
@@ -230,6 +233,10 @@ class WebUIWorkspaceController:
|
||||
return self._scope_from_metadata_value(raw_scope, default_scope=default_scope)
|
||||
|
||||
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:
|
||||
return self.default_scope()
|
||||
data = self._sessions.read_session_metadata(session_key)
|
||||
@@ -328,8 +335,24 @@ class WebUIWorkspaceController:
|
||||
return scope
|
||||
|
||||
def persist_scope(self, chat_id: str, scope: WorkspaceScope) -> None:
|
||||
session_key = f"websocket:{chat_id}"
|
||||
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[WORKSPACE_SCOPE_METADATA_KEY] = scope.metadata()
|
||||
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
@@ -34,6 +34,7 @@ from nanobot.session.session_handles import (
|
||||
SessionHandleResolver,
|
||||
)
|
||||
from nanobot.triggers.local_types import LocalTrigger
|
||||
from nanobot.utils.cancellation import shield_and_drain
|
||||
from nanobot.webui.file_preview import (
|
||||
WebUIFilePreviewError,
|
||||
file_preview_availability_payload,
|
||||
@@ -135,6 +136,18 @@ _WEBUI_MUTATION_PAYLOAD_ATTR = "_nanobot_webui_mutation_payload"
|
||||
_WEBUI_MUTATION_REQUEST_ATTR = "_nanobot_webui_mutation_request"
|
||||
_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 = {
|
||||
"automation.enable": "/api/webui/automations/enable",
|
||||
"automation.disable": "/api/webui/automations/disable",
|
||||
@@ -420,7 +433,12 @@ class GatewayHTTPHandler:
|
||||
response = await self._dispatch_resolved(connection, request, got)
|
||||
return response
|
||||
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(
|
||||
self,
|
||||
@@ -556,7 +574,14 @@ class GatewayHTTPHandler:
|
||||
|
||||
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)
|
||||
if elapsed_ms < _SLOW_WEBUI_HTTP_LOG_MS:
|
||||
return
|
||||
@@ -564,9 +589,10 @@ class GatewayHTTPHandler:
|
||||
return
|
||||
status = getattr(response, "status_code", None)
|
||||
self._log.warning(
|
||||
"slow webui http route path={} status={} duration_ms={}",
|
||||
path,
|
||||
"slow webui http operation={} status={} input_chars={} duration_ms={}",
|
||||
_slow_http_operation(path),
|
||||
status if status is not None else "none",
|
||||
input_chars,
|
||||
elapsed_ms,
|
||||
)
|
||||
|
||||
@@ -694,7 +720,11 @@ class GatewayHTTPHandler:
|
||||
async def _dispatch_session_routes(self, request: WsRequest, got: str) -> Response | None:
|
||||
m = re.match(r"^/api/sessions/([^/]+)/webui-thread$", got)
|
||||
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)
|
||||
if m:
|
||||
@@ -702,15 +732,27 @@ class GatewayHTTPHandler:
|
||||
|
||||
m = re.match(r"^/api/sessions/([^/]+)/file-preview$", got)
|
||||
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)
|
||||
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)
|
||||
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
|
||||
|
||||
@@ -957,13 +999,24 @@ class GatewayHTTPHandler:
|
||||
|
||||
# -- 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(
|
||||
self,
|
||||
request: WsRequest,
|
||||
got: str,
|
||||
) -> Response | None:
|
||||
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)
|
||||
if m:
|
||||
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()
|
||||
if not job_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:
|
||||
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:
|
||||
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:
|
||||
return _http_error(404, "automation not found")
|
||||
if job.payload.kind == "system_event":
|
||||
@@ -1044,13 +1108,23 @@ class GatewayHTTPHandler:
|
||||
return _http_error(409, "automation has no linked chat")
|
||||
|
||||
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")
|
||||
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")
|
||||
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":
|
||||
return _http_error(404, "automation not found")
|
||||
if result == "protected":
|
||||
@@ -1068,7 +1142,11 @@ class GatewayHTTPHandler:
|
||||
if isinstance(parsed, str):
|
||||
return _http_error(400, parsed)
|
||||
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:
|
||||
return _http_error(400, str(exc))
|
||||
if result == "not_found":
|
||||
@@ -1078,7 +1156,7 @@ class GatewayHTTPHandler:
|
||||
else:
|
||||
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(
|
||||
self,
|
||||
@@ -1163,9 +1241,17 @@ class GatewayHTTPHandler:
|
||||
if got == "/api/webui/skills/install":
|
||||
return await self._handle_webui_skill_install(connection, request)
|
||||
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":
|
||||
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":
|
||||
return self._handle_webui_skills(request)
|
||||
m = re.match(r"^/api/webui/skills/([^/]+)$", got)
|
||||
@@ -1276,7 +1362,7 @@ class GatewayHTTPHandler:
|
||||
) -> Response:
|
||||
if not self.check_api_token(request):
|
||||
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")
|
||||
if self._skill_install_lock.locked():
|
||||
return _http_error(409, "another skill installation is already in progress")
|
||||
@@ -1308,13 +1394,16 @@ class GatewayHTTPHandler:
|
||||
"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):
|
||||
return True
|
||||
try:
|
||||
return bool(
|
||||
self.settings.config.load().tools.webui_allow_remote_package_install
|
||||
)
|
||||
config = await self.settings.config.load_async()
|
||||
return bool(config.tools.webui_allow_remote_package_install)
|
||||
except Exception:
|
||||
self._log.exception("failed to load remote package install policy")
|
||||
return False
|
||||
|
||||
@@ -128,7 +128,7 @@ async def test_pending_document_attachment_keeps_body_out_of_prompt(
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
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.provider.chat_with_retry = chat_with_retry
|
||||
|
||||
@@ -49,6 +49,29 @@ def test_system_prompt_stays_stable_when_clock_changes(tmp_path, monkeypatch) ->
|
||||
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:
|
||||
workspace = _make_workspace(tmp_path)
|
||||
builder = ContextBuilder(workspace)
|
||||
|
||||
@@ -412,7 +412,7 @@ class TestEphemeralDirect:
|
||||
provider.supports_tools = True
|
||||
provider.generation = MagicMock(max_tokens=4096)
|
||||
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 (
|
||||
@@ -556,9 +556,9 @@ class TestEphemeralDirect:
|
||||
"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(
|
||||
@@ -646,7 +646,7 @@ class TestEphemeralHooks:
|
||||
provider.generation = MagicMock(max_tokens=4096)
|
||||
provider.chat_with_retry = AsyncMock(
|
||||
return_value=LLMResponse(
|
||||
content="done", finish_reason="stop", tool_calls=[], usage={},
|
||||
content="done", finish_reason="stop", tool_calls=[], usage=None,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@@ -12,7 +12,7 @@ from nanobot.utils.evaluator import (
|
||||
|
||||
class DummyProvider(LLMProvider):
|
||||
def __init__(self, responses: list[LLMResponse]):
|
||||
super().__init__()
|
||||
super().__init__(provider_name="dummy")
|
||||
self._responses = list(responses)
|
||||
|
||||
async def chat(self, *args, **kwargs) -> LLMResponse:
|
||||
|
||||
@@ -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:
|
||||
loop = _make_loop(tmp_path, context_window_tokens=32_768)
|
||||
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.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:
|
||||
loop = _make_loop(tmp_path, context_window_tokens=8_000)
|
||||
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.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
|
||||
|
||||
@@ -453,7 +453,7 @@ async def test_agent_loop_extra_hook_receives_calls(tmp_path):
|
||||
|
||||
loop = _make_loop(tmp_path, hooks=[TrackingHook()])
|
||||
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=[])
|
||||
|
||||
@@ -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.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=[])
|
||||
|
||||
@@ -541,7 +541,7 @@ async def test_agent_loop_extra_hook_error_isolation(tmp_path):
|
||||
|
||||
loop = _make_loop(tmp_path, hooks=[BadHook()])
|
||||
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=[])
|
||||
|
||||
@@ -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(
|
||||
content="working",
|
||||
tool_calls=[ToolCallRequest(id="c1", name="list_dir", arguments={"path": "."})],
|
||||
usage={},
|
||||
usage=None,
|
||||
))
|
||||
loop.tools.get_definitions = MagicMock(return_value=[])
|
||||
loop.tools.execute = AsyncMock(return_value="ok")
|
||||
|
||||
@@ -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_with_retry = AsyncMock()
|
||||
|
||||
@@ -48,7 +48,7 @@ async def test_ephemeral_runner_enters_and_restores_turn_scopes(tmp_path):
|
||||
|
||||
async def chat_with_retry(**_kwargs):
|
||||
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.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(
|
||||
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."},
|
||||
)
|
||||
],
|
||||
usage={},
|
||||
usage=None,
|
||||
),
|
||||
LLMResponse(
|
||||
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."},
|
||||
)
|
||||
],
|
||||
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.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.generation = GenerationSettings()
|
||||
provider.chat_with_retry = AsyncMock(side_effect=[
|
||||
LLMResponse(content="first answer", usage={}),
|
||||
LLMResponse(content="second answer", usage={}),
|
||||
LLMResponse(content="first answer", usage=None),
|
||||
LLMResponse(content="second answer", usage=None),
|
||||
])
|
||||
loop = AgentLoop(bus=MessageBus(), provider=provider, workspace=tmp_path, model="test-model")
|
||||
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.get_default_model.return_value = "test-model"
|
||||
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.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=None)
|
||||
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",
|
||||
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.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."},
|
||||
)
|
||||
],
|
||||
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.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):
|
||||
await on_content_delta("<think>hidden")
|
||||
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
|
||||
|
||||
@@ -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):
|
||||
await on_content_delta("Hello <thin")
|
||||
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
|
||||
|
||||
@@ -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):
|
||||
await on_content_delta("Hello <think>")
|
||||
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
|
||||
|
||||
@@ -459,8 +459,8 @@ async def test_loop_retries_think_only_final_response(tmp_path):
|
||||
async def chat_with_retry(**kwargs):
|
||||
call_count["n"] += 1
|
||||
if call_count["n"] == 1:
|
||||
return LLMResponse(content="<think>hidden</think>", tool_calls=[], usage={})
|
||||
return LLMResponse(content="Recovered answer", tool_calls=[], usage={})
|
||||
return LLMResponse(content="<think>hidden</think>", tool_calls=[], usage=None)
|
||||
return LLMResponse(content="Recovered answer", tool_calls=[], usage=None)
|
||||
|
||||
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"
|
||||
loop = AgentLoop(bus=bus, provider=provider, workspace=tmp_path, model="test-model")
|
||||
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_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",
|
||||
arguments={"command": "curl http://169.254.169.254/latest/meta-data/"},
|
||||
)],
|
||||
usage={},
|
||||
usage=None,
|
||||
)
|
||||
responses = iter([
|
||||
tool_call_resp,
|
||||
LLMResponse(
|
||||
content="I cannot access private URLs. Please share the local file.",
|
||||
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.get_default_model.return_value = "test-model"
|
||||
provider.chat_with_retry = AsyncMock(side_effect=[
|
||||
LLMResponse(content="429 rate limit exceeded", finish_reason="error", tool_calls=[], usage={}),
|
||||
LLMResponse(content="Recovered answer", tool_calls=[], usage={}),
|
||||
LLMResponse(content="429 rate limit exceeded", finish_reason="error", tool_calls=[], usage=None),
|
||||
LLMResponse(content="Recovered answer", tool_calls=[], usage=None),
|
||||
])
|
||||
|
||||
loop = AgentLoop(bus=MessageBus(), provider=provider, workspace=tmp_path, model="test-model")
|
||||
|
||||
@@ -20,7 +20,7 @@ from nanobot.bus.outbound_events import (
|
||||
)
|
||||
from nanobot.bus.queue import MessageBus
|
||||
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.runtime_context import (
|
||||
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]
|
||||
|
||||
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 (
|
||||
"done",
|
||||
[],
|
||||
@@ -1898,10 +1898,9 @@ async def test_turn_usage_is_persisted_with_the_saved_session(tmp_path: Path) ->
|
||||
)
|
||||
|
||||
loop.sessions.invalidate("cli:usage")
|
||||
assert loop.sessions.get_or_create("cli:usage").metadata["_last_usage"] == {
|
||||
"prompt_tokens": 64,
|
||||
"completion_tokens": 9,
|
||||
}
|
||||
assert loop.sessions.get_or_create("cli:usage").metadata["_last_usage"] == (
|
||||
LLMUsage.reported(input_tokens=64, output_tokens=9).to_dict()
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
||||
@@ -30,7 +30,7 @@ def _loop(tmp_path, responses: list[str], **kwargs) -> AgentLoop:
|
||||
provider.get_default_model.return_value = "test-model"
|
||||
provider.generation = GenerationSettings()
|
||||
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(
|
||||
bus=MessageBus(),
|
||||
|
||||
@@ -123,21 +123,21 @@ def allow_loopback_mcp_urls(monkeypatch: pytest.MonkeyPatch):
|
||||
_resolver_lock = asyncio.Lock()
|
||||
|
||||
monkeypatch.setattr(mcp_module, "PinnedDNSAsyncTransport", TestPinnedDNSAsyncTransport)
|
||||
monkeypatch.setattr(
|
||||
mcp_module,
|
||||
"validate_url_target",
|
||||
lambda url, *, allow_loopback=False: (True, ""),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
mcp_module,
|
||||
"resolve_url_target",
|
||||
lambda url, *, allow_loopback=False: (True, "", ("127.0.0.1",)),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
security_network,
|
||||
"resolve_url_target",
|
||||
lambda url, *, allow_loopback=False: (True, "", ("127.0.0.1",)),
|
||||
)
|
||||
async def allow_url(url: str, *, allow_loopback: bool = False) -> tuple[bool, str]:
|
||||
return True, ""
|
||||
|
||||
async def resolve_url(
|
||||
url: str,
|
||||
*,
|
||||
allow_loopback: bool = False,
|
||||
trust_remote_dns: bool = False,
|
||||
timeout_s: float = 3.0,
|
||||
) -> tuple[bool, str, tuple[str, ...]]:
|
||||
return True, "", ("127.0.0.1",)
|
||||
|
||||
monkeypatch.setattr(mcp_module, "async_validate_url_target", allow_url)
|
||||
monkeypatch.setattr(mcp_module, "async_resolve_url_target", resolve_url)
|
||||
monkeypatch.setattr(security_network, "async_resolve_url_target", resolve_url)
|
||||
monkeypatch.setattr(
|
||||
mcp_module,
|
||||
"env_proxy_applies_to_url",
|
||||
|
||||
+239
-31
@@ -14,6 +14,7 @@ from nanobot.config.schema import AgentDefaults
|
||||
from nanobot.providers.base import (
|
||||
LLMProvider,
|
||||
LLMResponse,
|
||||
LLMUsage,
|
||||
ProviderCallContext,
|
||||
ProviderConversationState,
|
||||
ToolCallRequest,
|
||||
@@ -22,6 +23,163 @@ from nanobot.providers.base import (
|
||||
_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
|
||||
async def test_runner_preserves_reasoning_fields_and_tool_results():
|
||||
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": "."})],
|
||||
reasoning_content="hidden reasoning",
|
||||
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
|
||||
return LLMResponse(content="done", tool_calls=[], usage={})
|
||||
return LLMResponse(content="done", tool_calls=[], usage=None)
|
||||
|
||||
provider.chat_with_retry = chat_with_retry
|
||||
tools = MagicMock()
|
||||
@@ -441,7 +599,7 @@ async def test_runner_uses_no_tools_finalization_after_max_iterations():
|
||||
return LLMResponse(
|
||||
content="Read the directory twice. More investigation remains.",
|
||||
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
|
||||
@@ -499,6 +657,55 @@ async def test_runner_times_out_hung_llm_request():
|
||||
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
|
||||
async def test_runner_applies_outer_wall_timeout_to_streaming_requests():
|
||||
from nanobot.agent.hook import AgentHook, AgentHookContext
|
||||
@@ -664,10 +871,10 @@ async def test_runner_replaces_empty_tool_result_with_marker():
|
||||
return LLMResponse(
|
||||
content="working",
|
||||
tool_calls=[ToolCallRequest(id="call_1", name="noop", arguments={})],
|
||||
usage={},
|
||||
usage=None,
|
||||
)
|
||||
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
|
||||
tools = MagicMock()
|
||||
@@ -702,12 +909,12 @@ async def test_runner_retries_empty_final_response_with_summary_prompt():
|
||||
return LLMResponse(
|
||||
content=None,
|
||||
tool_calls=[],
|
||||
usage={"prompt_tokens": 5, "completion_tokens": 1},
|
||||
usage=LLMUsage.reported(input_tokens=5, output_tokens=1),
|
||||
)
|
||||
return LLMResponse(
|
||||
content="final answer",
|
||||
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
|
||||
@@ -729,8 +936,9 @@ async def test_runner_retries_empty_final_response_with_summary_prompt():
|
||||
assert calls[0]["tools"] is not None
|
||||
assert calls[1]["tools"] is not None
|
||||
assert calls[2]["tools"] is None
|
||||
assert result.usage["prompt_tokens"] == 13
|
||||
assert result.usage["completion_tokens"] == 9
|
||||
assert result.usage is not None
|
||||
assert result.usage.input_tokens == 13
|
||||
assert result.usage.output_tokens == 9
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -802,7 +1010,7 @@ async def test_runner_uses_specific_message_after_empty_finalization_retry():
|
||||
provider = MagicMock(spec=LLMProvider)
|
||||
|
||||
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
|
||||
tools = MagicMock()
|
||||
@@ -842,14 +1050,14 @@ async def test_empty_finalization_retry_discards_candidate_provider_state():
|
||||
provider = MagicMock(spec=LLMProvider)
|
||||
provider.can_resume_conversation_state.return_value = True
|
||||
provider.chat_with_retry = AsyncMock(side_effect=[
|
||||
LLMResponse(content=None, tool_calls=[], usage={}),
|
||||
LLMResponse(content=None, tool_calls=[], usage={}),
|
||||
LLMResponse(content=None, tool_calls=[], usage=None),
|
||||
LLMResponse(content=None, tool_calls=[], usage=None),
|
||||
LLMResponse(
|
||||
content="finalized without tools",
|
||||
tool_calls=[ToolCallRequest(id="call_1", name="exec", arguments={})],
|
||||
finish_reason="stop",
|
||||
provider_state=candidate,
|
||||
usage={},
|
||||
usage=None,
|
||||
),
|
||||
])
|
||||
tools = MagicMock()
|
||||
@@ -988,20 +1196,20 @@ async def test_runner_empty_response_does_not_break_tool_chain():
|
||||
return LLMResponse(
|
||||
content=None,
|
||||
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:
|
||||
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:
|
||||
return LLMResponse(
|
||||
content=None,
|
||||
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(
|
||||
content="Here are the results.",
|
||||
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
|
||||
@@ -1030,9 +1238,8 @@ async def test_runner_empty_response_does_not_break_tool_chain():
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_runner_accumulates_usage_and_preserves_cached_tokens():
|
||||
"""Runner should accumulate prompt/completion tokens across iterations
|
||||
and preserve cached_tokens from provider responses."""
|
||||
async def test_runner_accumulates_usage_and_preserves_cache_reads():
|
||||
"""Runner accumulates usage across iterations, including cache reads."""
|
||||
from nanobot.agent.runner import AgentRunner
|
||||
|
||||
provider = MagicMock(spec=LLMProvider)
|
||||
@@ -1044,12 +1251,12 @@ async def test_runner_accumulates_usage_and_preserves_cached_tokens():
|
||||
return LLMResponse(
|
||||
content="thinking",
|
||||
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(
|
||||
content="done",
|
||||
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
|
||||
@@ -1067,11 +1274,12 @@ async def test_runner_accumulates_usage_and_preserves_cached_tokens():
|
||||
))
|
||||
|
||||
# Usage should be accumulated across iterations
|
||||
assert result.usage["prompt_tokens"] == 300 # 100 + 200
|
||||
assert result.usage["completion_tokens"] == 30 # 10 + 20
|
||||
assert result.usage["cached_tokens"] == 230 # 80 + 150
|
||||
assert result.usage["context_tokens"] == 200
|
||||
assert result.usage["request_count"] == 2
|
||||
assert result.usage is not None
|
||||
assert result.usage.input_tokens == 300 # 100 + 200
|
||||
assert result.usage.output_tokens == 30 # 10 + 20
|
||||
assert result.usage.cache_read_tokens == 230 # 80 + 150
|
||||
assert result.usage.context_tokens == 200
|
||||
assert result.usage.request_count == 2
|
||||
|
||||
|
||||
@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):
|
||||
captured.update(kwargs)
|
||||
return LLMResponse(content="done", tool_calls=[], usage={})
|
||||
return LLMResponse(content="done", tool_calls=[], usage=None)
|
||||
|
||||
provider = MagicMock(spec=LLMProvider)
|
||||
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):
|
||||
captured.update(kwargs)
|
||||
return LLMResponse(content="done", tool_calls=[], usage={})
|
||||
return LLMResponse(content="done", tool_calls=[], usage=None)
|
||||
|
||||
provider = MagicMock(spec=LLMProvider)
|
||||
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):
|
||||
captured.update(kwargs)
|
||||
return LLMResponse(content="done", tool_calls=[], usage={})
|
||||
return LLMResponse(content="done", tool_calls=[], usage=None)
|
||||
|
||||
provider = MagicMock(spec=LLMProvider)
|
||||
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):
|
||||
captured.update(kwargs)
|
||||
return LLMResponse(content="done", tool_calls=[], usage={})
|
||||
return LLMResponse(content="done", tool_calls=[], usage=None)
|
||||
|
||||
provider = MagicMock(spec=LLMProvider)
|
||||
provider.chat_with_retry = chat_with_retry
|
||||
|
||||
@@ -90,7 +90,7 @@ async def test_llm_error_not_appended_to_session_messages():
|
||||
|
||||
provider = MagicMock(spec=LLMProvider)
|
||||
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.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.",
|
||||
finish_reason=finish_reason,
|
||||
tool_calls=[ToolCallRequest(id="call_1", name="exec", arguments={"command": "echo nope"})],
|
||||
usage={},
|
||||
usage=None,
|
||||
))
|
||||
tools = MagicMock()
|
||||
tools.get_definitions.return_value = []
|
||||
@@ -189,7 +189,7 @@ async def test_runner_tool_error_sets_final_content():
|
||||
return LLMResponse(
|
||||
content="working",
|
||||
tool_calls=[ToolCallRequest(id="call_1", name="read_file", arguments={"path": "x"})],
|
||||
usage={},
|
||||
usage=None,
|
||||
)
|
||||
|
||||
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=[
|
||||
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
|
||||
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="tc2", name="exec", arguments={"cmd": "bad"}),
|
||||
],
|
||||
usage={},
|
||||
usage=None,
|
||||
)
|
||||
|
||||
provider.chat_with_retry = chat_with_retry
|
||||
@@ -333,9 +333,9 @@ async def test_length_finish_with_blank_content_routes_to_length_recovery():
|
||||
content="",
|
||||
finish_reason="length",
|
||||
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.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"
|
||||
)
|
||||
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]
|
||||
|
||||
@@ -78,7 +78,7 @@ class _FakeProvider(LLMProvider):
|
||||
*,
|
||||
responses: list[LLMResponse] | None = None,
|
||||
):
|
||||
super().__init__()
|
||||
super().__init__(provider_name=name)
|
||||
self.name = name
|
||||
self._response = response or _make_response()
|
||||
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
|
||||
|
||||
|
||||
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:
|
||||
from nanobot.config.schema import Config
|
||||
from nanobot.providers.factory import provider_signature
|
||||
|
||||
@@ -25,7 +25,7 @@ async def test_runner_exits_normally_without_predicate():
|
||||
|
||||
provider = MagicMock(spec=LLMProvider)
|
||||
provider.chat_with_retry = AsyncMock(return_value=LLMResponse(
|
||||
content="all done", tool_calls=[], usage={},
|
||||
content="all done", tool_calls=[], usage=None,
|
||||
))
|
||||
tools = MagicMock()
|
||||
tools.get_definitions.return_value = []
|
||||
@@ -50,7 +50,7 @@ async def test_runner_exits_normally_with_inactive_goal():
|
||||
|
||||
provider = MagicMock(spec=LLMProvider)
|
||||
provider.chat_with_retry = AsyncMock(return_value=LLMResponse(
|
||||
content="all done", tool_calls=[], usage={},
|
||||
content="all done", tool_calls=[], usage=None,
|
||||
))
|
||||
tools = MagicMock()
|
||||
tools.get_definitions.return_value = []
|
||||
@@ -82,7 +82,7 @@ async def test_runner_forces_continue_when_goal_active():
|
||||
|
||||
provider = MagicMock(spec=LLMProvider)
|
||||
provider.chat_with_retry = AsyncMock(return_value=LLMResponse(
|
||||
content="still working", tool_calls=[], usage={},
|
||||
content="still working", tool_calls=[], usage=None,
|
||||
))
|
||||
tools = MagicMock()
|
||||
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.chat_with_retry = AsyncMock(return_value=LLMResponse(
|
||||
content="still working", tool_calls=[], usage={},
|
||||
content="still working", tool_calls=[], usage=None,
|
||||
))
|
||||
tools = MagicMock()
|
||||
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.chat_with_retry = AsyncMock(return_value=LLMResponse(
|
||||
content="still working", tool_calls=[], usage={},
|
||||
content="still working", tool_calls=[], usage=None,
|
||||
))
|
||||
tools = MagicMock()
|
||||
tools.get_definitions.return_value = []
|
||||
@@ -165,7 +165,7 @@ async def test_runner_does_not_force_continue_on_error():
|
||||
|
||||
provider = MagicMock(spec=LLMProvider)
|
||||
provider.chat_with_retry = AsyncMock(return_value=LLMResponse(
|
||||
content=None, tool_calls=[], usage={},
|
||||
content=None, tool_calls=[], usage=None,
|
||||
finish_reason="error",
|
||||
))
|
||||
tools = MagicMock()
|
||||
@@ -191,7 +191,7 @@ async def test_runner_uses_custom_goal_continue_message():
|
||||
|
||||
provider = MagicMock(spec=LLMProvider)
|
||||
provider.chat_with_retry = AsyncMock(return_value=LLMResponse(
|
||||
content="still working", tool_calls=[], usage={},
|
||||
content="still working", tool_calls=[], usage=None,
|
||||
))
|
||||
tools = MagicMock()
|
||||
tools.get_definitions.return_value = []
|
||||
@@ -220,7 +220,7 @@ async def test_runner_resolves_goal_continue_message_lazily():
|
||||
|
||||
provider = MagicMock(spec=LLMProvider)
|
||||
provider.chat_with_retry = AsyncMock(return_value=LLMResponse(
|
||||
content="still working", tool_calls=[], usage={},
|
||||
content="still working", tool_calls=[], usage=None,
|
||||
))
|
||||
tools = MagicMock()
|
||||
tools.get_definitions.return_value = []
|
||||
|
||||
@@ -273,7 +273,7 @@ async def test_runner_drops_orphan_tool_results_before_model_request():
|
||||
|
||||
async def chat_with_retry(*, messages, **kwargs):
|
||||
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
|
||||
tools = MagicMock()
|
||||
@@ -312,7 +312,7 @@ async def test_backfill_repairs_model_context_without_shifting_save_turn_boundar
|
||||
|
||||
provider = MagicMock()
|
||||
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_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):
|
||||
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
|
||||
tools = MagicMock()
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user