mirror of
https://github.com/HKUDS/nanobot.git
synced 2026-08-31 16:21:50 +03:00
Compare commits
30
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 | ||
|
|
ffa58aa5ef | ||
|
|
1e9d46fb36 | ||
|
|
ab7351be63 | ||
|
|
cfc872fb52 | ||
|
|
cfc1fae8b5 | ||
|
|
9807e9cf37 | ||
|
|
961b1fdd7d |
@@ -209,7 +209,13 @@ Use `nanobot gateway --background` for the same direct entry point without keepi
|
|||||||
nanobot agent
|
nanobot agent
|
||||||
```
|
```
|
||||||
|
|
||||||
This opens the native terminal client with the configured model and tools, using the launch directory as its workspace. Use `/sessions` to switch saved conversations, `/new-chat` to preserve this conversation and start another one, `/branch` to fork from a completed reply, `/context` to inspect the compacted summary and raw message suffix available to the agent, or `/diff` to review the latest turn's file changes. Type `@` to mention an installed app, configured MCP server, or saved session. While nanobot is working, `Enter` steers the current turn, `Tab` queues a visible follow-up for the next turn, and `Option+Up` on macOS (`Alt+Up` on Windows/Linux) returns the latest queued message for editing. Press `Shift+Enter` to add a newline; `Ctrl+J` is the universal fallback for terminals that cannot distinguish modified Enter keys. Use `PageUp` at the top to load earlier transcript pages. Each launch starts a new session; `--session` selects an existing WebSocket session, while `--workspace` overrides the launch directory. Use `--classic` to resume a session from another channel. The existing nanobot `/new` command keeps its original behavior: it resets the current chat. `nanobot agent` and `nanobot webui` share one on-demand local gateway: either command can start it, each launcher releases only its own client, and the last interactive launcher to exit stops it. Use `/detach` to close the TUI while keeping the gateway and any active agent turn running in the background; after the terminal is restored, nanobot prints the exact `nanobot gateway stop` command for that config and workspace. Use `nanobot gateway --background` to start persistently before opening a client. Type `exit` or press `Ctrl+C` when you are done; after the terminal is restored, nanobot prints a ready-to-run `nanobot agent --session ...` command that resumes the session. Use `nanobot agent --classic` for the legacy Python prompt.
|
This opens the native terminal client with the launch directory as its workspace. It shares saved conversations and the local gateway with the WebUI.
|
||||||
|
|
||||||
|
- Type `/` to discover commands, `/sessions` to switch conversations, or `@` to mention an app, MCP server, or saved session.
|
||||||
|
- Press `Enter` to send or steer, `Tab` to queue a follow-up, and `Shift+Enter` to add a newline (`Ctrl+J` works in terminals that cannot distinguish modified Enter keys).
|
||||||
|
- Use `/detach` to leave the current task running, or start with `nanobot gateway --background` when nanobot should stay online after all local clients exit.
|
||||||
|
|
||||||
|
Each launch starts a new session by default. Use `--session` to resume one and `--workspace` to choose another workspace. See the [CLI reference](./docs/cli-reference.md#agent-cli) for session branching, diffs, history, shortcuts, gateway lifecycle, and compatibility options.
|
||||||
|
|
||||||
For one request and an immediate exit, use:
|
For one request and an immediate exit, use:
|
||||||
|
|
||||||
|
|||||||
@@ -23,6 +23,7 @@ if TYPE_CHECKING:
|
|||||||
STREAM_EVENT_TOOL_FAILED,
|
STREAM_EVENT_TOOL_FAILED,
|
||||||
STREAM_EVENT_TOOL_STARTED,
|
STREAM_EVENT_TOOL_STARTED,
|
||||||
STREAM_EVENT_TYPES,
|
STREAM_EVENT_TYPES,
|
||||||
|
LLMUsage,
|
||||||
Nanobot,
|
Nanobot,
|
||||||
RunResult,
|
RunResult,
|
||||||
RunStream,
|
RunStream,
|
||||||
@@ -56,6 +57,7 @@ __logo__ = "🐈"
|
|||||||
|
|
||||||
_LAZY_EXPORTS = {
|
_LAZY_EXPORTS = {
|
||||||
"Nanobot": ".nanobot",
|
"Nanobot": ".nanobot",
|
||||||
|
"LLMUsage": ".nanobot",
|
||||||
"RunStream": ".nanobot",
|
"RunStream": ".nanobot",
|
||||||
"RunResult": ".nanobot",
|
"RunResult": ".nanobot",
|
||||||
"RequestContext": ".agent.tools.context",
|
"RequestContext": ".agent.tools.context",
|
||||||
@@ -93,6 +95,7 @@ def __getattr__(name: str) -> Any:
|
|||||||
|
|
||||||
__all__ = [
|
__all__ = [
|
||||||
"Nanobot",
|
"Nanobot",
|
||||||
|
"LLMUsage",
|
||||||
"RunResult",
|
"RunResult",
|
||||||
"RequestContext",
|
"RequestContext",
|
||||||
"RuntimeContextBlock",
|
"RuntimeContextBlock",
|
||||||
|
|||||||
@@ -2,7 +2,9 @@
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
from collections.abc import Collection
|
import asyncio
|
||||||
|
import inspect
|
||||||
|
from collections.abc import Awaitable, Collection
|
||||||
from datetime import datetime
|
from datetime import datetime
|
||||||
from typing import TYPE_CHECKING, Any, Callable, Coroutine
|
from typing import TYPE_CHECKING, Any, Callable, Coroutine
|
||||||
|
|
||||||
@@ -47,9 +49,27 @@ class AutoCompact:
|
|||||||
return idle_seconds >= self._ttl * 60
|
return idle_seconds >= self._ttl * 60
|
||||||
|
|
||||||
def _has_unarchived_messages(self, key: str) -> bool:
|
def _has_unarchived_messages(self, key: str) -> bool:
|
||||||
session = self.sessions.get_or_create(key)
|
return self._session_has_unarchived_messages(self.sessions.get_or_create(key))
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _session_has_unarchived_messages(session: Session) -> bool:
|
||||||
return session.last_consolidated < len(session.messages)
|
return session.last_consolidated < len(session.messages)
|
||||||
|
|
||||||
|
def _has_native_async_session_method(self, name: str) -> bool:
|
||||||
|
"""Check the manager's real class, not mock-generated instance attributes."""
|
||||||
|
method = inspect.getattr_static(type(self.sessions), name, None)
|
||||||
|
return inspect.iscoroutinefunction(method)
|
||||||
|
|
||||||
|
async def _list_sessions_nonblocking(self) -> list[dict[str, Any]]:
|
||||||
|
if self._has_native_async_session_method("list_sessions_async"):
|
||||||
|
return await self.sessions.list_sessions_async()
|
||||||
|
return await asyncio.to_thread(self.sessions.list_sessions)
|
||||||
|
|
||||||
|
async def _get_or_create_nonblocking(self, key: str) -> Session:
|
||||||
|
if self._has_native_async_session_method("get_or_create_async"):
|
||||||
|
return await self.sessions.get_or_create_async(key)
|
||||||
|
return await asyncio.to_thread(self.sessions.get_or_create, key)
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def _is_internal_session(cls, key: str) -> bool:
|
def _is_internal_session(cls, key: str) -> bool:
|
||||||
return key.startswith(cls._INTERNAL_SESSION_PREFIXES)
|
return key.startswith(cls._INTERNAL_SESSION_PREFIXES)
|
||||||
@@ -79,6 +99,31 @@ class AutoCompact:
|
|||||||
self._archiving.add(key)
|
self._archiving.add(key)
|
||||||
schedule_background(self._archive(key, runtime=runtime))
|
schedule_background(self._archive(key, runtime=runtime))
|
||||||
|
|
||||||
|
async def check_expired_async(
|
||||||
|
self,
|
||||||
|
schedule_background: Callable[[Coroutine[Any, Any, None]], None],
|
||||||
|
resolve_runtime: Callable[[Session], Awaitable[LLMRuntime]],
|
||||||
|
active_session_keys: Collection[str] = (),
|
||||||
|
) -> None:
|
||||||
|
"""Schedule idle archival without blocking the event loop."""
|
||||||
|
now = datetime.now()
|
||||||
|
active_keys = set(active_session_keys)
|
||||||
|
for info in await self._list_sessions_nonblocking():
|
||||||
|
key = info.get("key", "")
|
||||||
|
if not key or self._is_internal_session(key) or key in self._archiving:
|
||||||
|
continue
|
||||||
|
if key in active_keys or not self._is_expired(info.get("updated_at"), now):
|
||||||
|
continue
|
||||||
|
session = await self._get_or_create_nonblocking(key)
|
||||||
|
if not self._session_has_unarchived_messages(session):
|
||||||
|
continue
|
||||||
|
try:
|
||||||
|
runtime = await resolve_runtime(session)
|
||||||
|
except (KeyError, ValueError):
|
||||||
|
continue
|
||||||
|
self._archiving.add(key)
|
||||||
|
schedule_background(self._archive_async(key, runtime=runtime))
|
||||||
|
|
||||||
async def _archive(self, key: str, *, runtime: LLMRuntime) -> None:
|
async def _archive(self, key: str, *, runtime: LLMRuntime) -> None:
|
||||||
if self._is_internal_session(key):
|
if self._is_internal_session(key):
|
||||||
self._archiving.discard(key)
|
self._archiving.discard(key)
|
||||||
@@ -90,18 +135,38 @@ class AutoCompact:
|
|||||||
max_suffix=self._RECENT_SUFFIX_MESSAGES,
|
max_suffix=self._RECENT_SUFFIX_MESSAGES,
|
||||||
)
|
)
|
||||||
if summary and summary != "(nothing)":
|
if summary and summary != "(nothing)":
|
||||||
session = self.sessions.get_or_create(key)
|
self._record_stored_summary(key, self.sessions.get_or_create(key))
|
||||||
stored = session_summary_from_metadata(
|
|
||||||
session.metadata,
|
|
||||||
fallback_last_active=session.updated_at,
|
|
||||||
)
|
|
||||||
if stored is not None:
|
|
||||||
self._summaries[key] = stored
|
|
||||||
except Exception:
|
except Exception:
|
||||||
logger.exception("Auto-compact: failed for {}", key)
|
logger.exception("Auto-compact: failed for {}", key)
|
||||||
finally:
|
finally:
|
||||||
self._archiving.discard(key)
|
self._archiving.discard(key)
|
||||||
|
|
||||||
|
async def _archive_async(self, key: str, *, runtime: LLMRuntime) -> None:
|
||||||
|
if self._is_internal_session(key):
|
||||||
|
self._archiving.discard(key)
|
||||||
|
return
|
||||||
|
try:
|
||||||
|
summary = await self.consolidator.compact_idle_session(
|
||||||
|
key,
|
||||||
|
runtime=runtime,
|
||||||
|
max_suffix=self._RECENT_SUFFIX_MESSAGES,
|
||||||
|
)
|
||||||
|
if summary and summary != "(nothing)":
|
||||||
|
session = await self._get_or_create_nonblocking(key)
|
||||||
|
self._record_stored_summary(key, session)
|
||||||
|
except Exception:
|
||||||
|
logger.exception("Auto-compact: failed for {}", key)
|
||||||
|
finally:
|
||||||
|
self._archiving.discard(key)
|
||||||
|
|
||||||
|
def _record_stored_summary(self, key: str, session: Session) -> None:
|
||||||
|
stored = session_summary_from_metadata(
|
||||||
|
session.metadata,
|
||||||
|
fallback_last_active=session.updated_at,
|
||||||
|
)
|
||||||
|
if stored is not None:
|
||||||
|
self._summaries[key] = stored
|
||||||
|
|
||||||
def prepare_session(self, session: Session, key: str) -> tuple[Session, SessionSummary | None]:
|
def prepare_session(self, session: Session, key: str) -> tuple[Session, SessionSummary | None]:
|
||||||
if self._is_internal_session(key):
|
if self._is_internal_session(key):
|
||||||
self._archiving.discard(key)
|
self._archiving.discard(key)
|
||||||
@@ -110,6 +175,28 @@ class AutoCompact:
|
|||||||
if key in self._archiving or self._is_expired(session.updated_at):
|
if key in self._archiving or self._is_expired(session.updated_at):
|
||||||
logger.info("Auto-compact: reloading session {} (archiving={})", key, key in self._archiving)
|
logger.info("Auto-compact: reloading session {} (archiving={})", key, key in self._archiving)
|
||||||
session = self.sessions.get_or_create(key)
|
session = self.sessions.get_or_create(key)
|
||||||
|
return self._prepared_summary(session, key)
|
||||||
|
|
||||||
|
async def prepare_session_async(
|
||||||
|
self,
|
||||||
|
session: Session,
|
||||||
|
key: str,
|
||||||
|
) -> tuple[Session, SessionSummary | None]:
|
||||||
|
"""Prepare a session without blocking on a reload."""
|
||||||
|
if self._is_internal_session(key):
|
||||||
|
self._archiving.discard(key)
|
||||||
|
self._summaries.pop(key, None)
|
||||||
|
return session, None
|
||||||
|
if key in self._archiving or self._is_expired(session.updated_at):
|
||||||
|
logger.info("Auto-compact: reloading session {} (archiving={})", key, key in self._archiving)
|
||||||
|
session = await self._get_or_create_nonblocking(key)
|
||||||
|
return self._prepared_summary(session, key)
|
||||||
|
|
||||||
|
def _prepared_summary(
|
||||||
|
self,
|
||||||
|
session: Session,
|
||||||
|
key: str,
|
||||||
|
) -> tuple[Session, SessionSummary | None]:
|
||||||
# Hot path: summary from in-memory dict (process hasn't restarted).
|
# Hot path: summary from in-memory dict (process hasn't restarted).
|
||||||
entry = self._summaries.pop(key, None)
|
entry = self._summaries.pop(key, None)
|
||||||
if entry:
|
if entry:
|
||||||
|
|||||||
@@ -13,6 +13,19 @@ class AutomationTurnError(RuntimeError):
|
|||||||
"""Raised when an automation turn reaches the agent and finishes with an error."""
|
"""Raised when an automation turn reaches the agent and finishes with an error."""
|
||||||
|
|
||||||
|
|
||||||
|
class AutomationTurnAcceptedCancellation(asyncio.CancelledError):
|
||||||
|
"""Cancellation raised after an automation turn was accepted for processing.
|
||||||
|
|
||||||
|
Callers must not replay the turn: the accepted agent work now has independent
|
||||||
|
ownership and may continue after the submitting task is cancelled.
|
||||||
|
"""
|
||||||
|
|
||||||
|
|
||||||
|
def _consume_future_exception(future: asyncio.Future[object]) -> None:
|
||||||
|
if not future.cancelled():
|
||||||
|
future.exception()
|
||||||
|
|
||||||
|
|
||||||
async def publish_next_deferred_turn(
|
async def publish_next_deferred_turn(
|
||||||
*,
|
*,
|
||||||
deferred_queues: dict[str, list[InboundMessage]],
|
deferred_queues: dict[str, list[InboundMessage]],
|
||||||
@@ -70,19 +83,36 @@ class AutomationTurnCoordinator:
|
|||||||
future: asyncio.Future[OutboundMessage | None] = loop.create_future()
|
future: asyncio.Future[OutboundMessage | None] = loop.create_future()
|
||||||
self._waiters[turn_id] = future
|
self._waiters[turn_id] = future
|
||||||
self._pending_messages_by_turn_id[turn_id] = msg
|
self._pending_messages_by_turn_id[turn_id] = msg
|
||||||
|
accepted = False
|
||||||
try:
|
try:
|
||||||
if self._is_running():
|
if self._is_running():
|
||||||
await self._publish_inbound(msg)
|
await self._publish_inbound(msg)
|
||||||
|
accepted = True
|
||||||
else:
|
else:
|
||||||
await self._dispatch(msg)
|
# Direct dispatch is given independent task ownership for the
|
||||||
|
# same reason as publishing to the inbound queue: once admitted,
|
||||||
|
# cancelling this submitter must not cancel and then replay the
|
||||||
|
# already-running agent turn.
|
||||||
|
dispatch_future: asyncio.Future[object] = asyncio.ensure_future(
|
||||||
|
self._dispatch(msg)
|
||||||
|
)
|
||||||
|
dispatch_future.add_done_callback(_consume_future_exception)
|
||||||
|
accepted = True
|
||||||
|
await asyncio.shield(dispatch_future)
|
||||||
try:
|
try:
|
||||||
return await future
|
return await future
|
||||||
except asyncio.CancelledError:
|
except asyncio.CancelledError as exc:
|
||||||
|
if accepted:
|
||||||
|
raise AutomationTurnAcceptedCancellation(*exc.args) from None
|
||||||
raise
|
raise
|
||||||
except AutomationTurnError:
|
except AutomationTurnError:
|
||||||
raise
|
raise
|
||||||
except Exception as exc:
|
except Exception as exc:
|
||||||
raise AutomationTurnError(str(exc) or exc.__class__.__name__) from exc
|
raise AutomationTurnError(str(exc) or exc.__class__.__name__) from exc
|
||||||
|
except asyncio.CancelledError as exc:
|
||||||
|
if accepted and not isinstance(exc, AutomationTurnAcceptedCancellation):
|
||||||
|
raise AutomationTurnAcceptedCancellation(*exc.args) from None
|
||||||
|
raise
|
||||||
finally:
|
finally:
|
||||||
self._waiters.pop(turn_id, None)
|
self._waiters.pop(turn_id, None)
|
||||||
self._pending_messages_by_turn_id.pop(turn_id, None)
|
self._pending_messages_by_turn_id.pop(turn_id, None)
|
||||||
|
|||||||
@@ -112,6 +112,14 @@ class ContextBuilder:
|
|||||||
|
|
||||||
parts.append(render_template("agent/tool_contract.md"))
|
parts.append(render_template("agent/tool_contract.md"))
|
||||||
|
|
||||||
|
project_path = root.expanduser().resolve()
|
||||||
|
if project_path != self.workspace.expanduser().resolve():
|
||||||
|
parts.append(
|
||||||
|
"# Current Project\n\n"
|
||||||
|
f"Working directory: {project_path}\n"
|
||||||
|
"Use it as the default root for project files and relative tool paths."
|
||||||
|
)
|
||||||
|
|
||||||
if include_memory:
|
if include_memory:
|
||||||
memory = self.memory.read_memory()
|
memory = self.memory.read_memory()
|
||||||
if memory and not self._is_template_content(memory, "memory/MEMORY.md"):
|
if memory and not self._is_template_content(memory, "memory/MEMORY.md"):
|
||||||
|
|||||||
@@ -9,7 +9,7 @@ from typing import Any
|
|||||||
|
|
||||||
from loguru import logger
|
from loguru import logger
|
||||||
|
|
||||||
from nanobot.providers.base import LLMResponse, ToolCallRequest
|
from nanobot.providers.base import LLMResponse, LLMUsage, ToolCallRequest
|
||||||
|
|
||||||
|
|
||||||
@dataclass(slots=True)
|
@dataclass(slots=True)
|
||||||
@@ -19,7 +19,7 @@ class AgentHookContext:
|
|||||||
iteration: int
|
iteration: int
|
||||||
messages: list[dict[str, Any]]
|
messages: list[dict[str, Any]]
|
||||||
response: LLMResponse | None = None
|
response: LLMResponse | None = None
|
||||||
usage: dict[str, int] = field(default_factory=dict)
|
usage: LLMUsage | None = None
|
||||||
tool_calls: list[ToolCallRequest] = field(default_factory=list)
|
tool_calls: list[ToolCallRequest] = field(default_factory=list)
|
||||||
tool_results: list[Any] = field(default_factory=list)
|
tool_results: list[Any] = field(default_factory=list)
|
||||||
tool_events: list[dict[str, str]] = field(default_factory=list)
|
tool_events: list[dict[str, str]] = field(default_factory=list)
|
||||||
@@ -39,7 +39,7 @@ class AgentRunHookContext:
|
|||||||
messages: list[dict[str, Any]]
|
messages: list[dict[str, Any]]
|
||||||
final_content: str | None = None
|
final_content: str | None = None
|
||||||
tools_used: list[str] = field(default_factory=list)
|
tools_used: list[str] = field(default_factory=list)
|
||||||
usage: dict[str, int] = field(default_factory=dict)
|
usage: LLMUsage | None = None
|
||||||
stop_reason: str | None = None
|
stop_reason: str | None = None
|
||||||
error: str | None = None
|
error: str | None = None
|
||||||
tool_events: list[dict[str, str]] = field(default_factory=list)
|
tool_events: list[dict[str, str]] = field(default_factory=list)
|
||||||
@@ -284,7 +284,7 @@ class SDKCaptureHook(AgentHook):
|
|||||||
super().__init__()
|
super().__init__()
|
||||||
self.tools_used: list[str] = []
|
self.tools_used: list[str] = []
|
||||||
self.messages: list[dict[str, Any]] = []
|
self.messages: list[dict[str, Any]] = []
|
||||||
self.usage: dict[str, int] = {}
|
self.usage: LLMUsage | None = None
|
||||||
self.stop_reason: str | None = None
|
self.stop_reason: str | None = None
|
||||||
self.error: str | None = None
|
self.error: str | None = None
|
||||||
self.tool_events: list[dict[str, str]] = []
|
self.tool_events: list[dict[str, str]] = []
|
||||||
@@ -294,7 +294,7 @@ class SDKCaptureHook(AgentHook):
|
|||||||
for call in context.tool_calls:
|
for call in context.tool_calls:
|
||||||
self.tools_used.append(call.name)
|
self.tools_used.append(call.name)
|
||||||
self.messages = list(context.messages)
|
self.messages = list(context.messages)
|
||||||
self.usage = dict(context.usage)
|
self.usage = context.usage
|
||||||
self.stop_reason = context.stop_reason
|
self.stop_reason = context.stop_reason
|
||||||
self.error = context.error
|
self.error = context.error
|
||||||
self.tool_events = list(context.tool_events)
|
self.tool_events = list(context.tool_events)
|
||||||
@@ -302,7 +302,7 @@ class SDKCaptureHook(AgentHook):
|
|||||||
async def after_run(self, context: AgentRunHookContext) -> None:
|
async def after_run(self, context: AgentRunHookContext) -> None:
|
||||||
self.tools_used = list(context.tools_used)
|
self.tools_used = list(context.tools_used)
|
||||||
self.messages = list(context.messages)
|
self.messages = list(context.messages)
|
||||||
self.usage = dict(context.usage)
|
self.usage = context.usage
|
||||||
self.stop_reason = context.stop_reason
|
self.stop_reason = context.stop_reason
|
||||||
self.error = context.error
|
self.error = context.error
|
||||||
self.tool_events = list(context.tool_events)
|
self.tool_events = list(context.tool_events)
|
||||||
|
|||||||
+278
-180
@@ -49,7 +49,8 @@ from nanobot.bus.queue import MessageBus
|
|||||||
from nanobot.bus.runtime_events import RuntimeEventBus
|
from nanobot.bus.runtime_events import RuntimeEventBus
|
||||||
from nanobot.command import CommandContext, CommandRouter, register_builtin_commands
|
from nanobot.command import CommandContext, CommandRouter, register_builtin_commands
|
||||||
from nanobot.config.schema import AgentDefaults, ModelPresetConfig
|
from nanobot.config.schema import AgentDefaults, ModelPresetConfig
|
||||||
from nanobot.providers.base import LLMProvider, ProviderConversationState
|
from nanobot.llm_usage.context import source_from_request
|
||||||
|
from nanobot.providers.base import LLMProvider, LLMUsage, ProviderConversationState
|
||||||
from nanobot.providers.factory import ProviderSnapshot
|
from nanobot.providers.factory import ProviderSnapshot
|
||||||
from nanobot.runtime_context import (
|
from nanobot.runtime_context import (
|
||||||
RUNTIME_CONTEXT_HISTORY_META,
|
RUNTIME_CONTEXT_HISTORY_META,
|
||||||
@@ -66,6 +67,7 @@ from nanobot.security.workspace_access import (
|
|||||||
reset_workspace_scope,
|
reset_workspace_scope,
|
||||||
)
|
)
|
||||||
from nanobot.session import turn_continuation
|
from nanobot.session import turn_continuation
|
||||||
|
from nanobot.session.async_compat import call_session_manager
|
||||||
from nanobot.session.automation_turns import automation_history_overrides
|
from nanobot.session.automation_turns import automation_history_overrides
|
||||||
from nanobot.session.goal_state import (
|
from nanobot.session.goal_state import (
|
||||||
goal_state_runtime_lines,
|
goal_state_runtime_lines,
|
||||||
@@ -79,6 +81,15 @@ from nanobot.session.model_selection import (
|
|||||||
SESSION_MODEL_PRESET_METADATA_KEY,
|
SESSION_MODEL_PRESET_METADATA_KEY,
|
||||||
model_preset_from_metadata,
|
model_preset_from_metadata,
|
||||||
)
|
)
|
||||||
|
from nanobot.session.recovery import (
|
||||||
|
PENDING_FOLLOWUP_ID_KEY,
|
||||||
|
RECOVERY_INBOUND_METADATA_KEY,
|
||||||
|
RecoveryAdmission,
|
||||||
|
acknowledge_pending_followups,
|
||||||
|
record_pending_followup,
|
||||||
|
restore_pending_interruption,
|
||||||
|
restore_runtime_checkpoint,
|
||||||
|
)
|
||||||
from nanobot.session.summary import SessionSummary
|
from nanobot.session.summary import SessionSummary
|
||||||
from nanobot.triggers.local_turns import LocalTriggerTurnCoordinator
|
from nanobot.triggers.local_turns import LocalTriggerTurnCoordinator
|
||||||
from nanobot.utils.cancellation import task_is_cancelling
|
from nanobot.utils.cancellation import task_is_cancelling
|
||||||
@@ -158,7 +169,7 @@ class TurnContext:
|
|||||||
turn_wall_started_at: float = field(default_factory=time.time)
|
turn_wall_started_at: float = field(default_factory=time.time)
|
||||||
visible_run_started_at: float | None = None
|
visible_run_started_at: float | None = None
|
||||||
turn_latency_ms: int | None = None
|
turn_latency_ms: int | None = None
|
||||||
usage: dict[str, int] = field(default_factory=dict)
|
usage: LLMUsage | None = None
|
||||||
|
|
||||||
def require_runtime(self) -> LLMRuntime:
|
def require_runtime(self) -> LLMRuntime:
|
||||||
"""Return the runtime established by the BUILD stage."""
|
"""Return the runtime established by the BUILD stage."""
|
||||||
@@ -194,7 +205,7 @@ class AgentLoop:
|
|||||||
return self.tools.tool_names
|
return self.tools.tool_names
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def last_usage(self) -> Mapping[str, int]:
|
def last_usage(self) -> LLMUsage | None:
|
||||||
"""Latest aggregate usage exposed through the runtime-control snapshot."""
|
"""Latest aggregate usage exposed through the runtime-control snapshot."""
|
||||||
return self._last_usage
|
return self._last_usage
|
||||||
|
|
||||||
@@ -291,12 +302,14 @@ class AgentLoop:
|
|||||||
restart_mode: str = "auto",
|
restart_mode: str = "auto",
|
||||||
local_trigger_store: LocalTriggerStore | None = None,
|
local_trigger_store: LocalTriggerStore | None = None,
|
||||||
idle_compact_check_interval_seconds: int = 0,
|
idle_compact_check_interval_seconds: int = 0,
|
||||||
|
recovery_admission: RecoveryAdmission | None = None,
|
||||||
):
|
):
|
||||||
from nanobot.config.schema import ToolsConfig
|
from nanobot.config.schema import ToolsConfig
|
||||||
|
|
||||||
_tc = tools_config or ToolsConfig()
|
_tc = tools_config or ToolsConfig()
|
||||||
defaults = AgentDefaults()
|
defaults = AgentDefaults()
|
||||||
self.bus = bus
|
self.bus = bus
|
||||||
|
self._recovery_admission = recovery_admission
|
||||||
if turn_delivery_factory is not None:
|
if turn_delivery_factory is not None:
|
||||||
if turn_delivery_factory.bus is not bus:
|
if turn_delivery_factory.bus is not bus:
|
||||||
raise ValueError("turn delivery factory must use the agent message bus")
|
raise ValueError("turn delivery factory must use the agent message bus")
|
||||||
@@ -367,7 +380,7 @@ class AgentLoop:
|
|||||||
default_restrict_to_workspace=restrict_to_workspace,
|
default_restrict_to_workspace=restrict_to_workspace,
|
||||||
)
|
)
|
||||||
self._start_time = time.time()
|
self._start_time = time.time()
|
||||||
self._last_usage: dict[str, int] = {}
|
self._last_usage: LLMUsage | None = None
|
||||||
self._extra_hooks: list[AgentHook] = hooks or []
|
self._extra_hooks: list[AgentHook] = hooks or []
|
||||||
self._hook_factories: list[AgentTurnHookFactory] = hook_factories or []
|
self._hook_factories: list[AgentTurnHookFactory] = hook_factories or []
|
||||||
|
|
||||||
@@ -409,6 +422,7 @@ class AgentLoop:
|
|||||||
# When a session has an active task, new messages for that session
|
# When a session has an active task, new messages for that session
|
||||||
# are routed here instead of creating a new task.
|
# are routed here instead of creating a new task.
|
||||||
self._pending_queues: dict[str, asyncio.Queue[InboundMessage]] = {}
|
self._pending_queues: dict[str, asyncio.Queue[InboundMessage]] = {}
|
||||||
|
self._preserve_inflight_turns_on_shutdown = False
|
||||||
self._deferred_automation_turns: dict[str, list[InboundMessage]] = {}
|
self._deferred_automation_turns: dict[str, list[InboundMessage]] = {}
|
||||||
self._cron_turns = CronTurnCoordinator(
|
self._cron_turns = CronTurnCoordinator(
|
||||||
publish_inbound=self.bus.publish_inbound,
|
publish_inbound=self.bus.publish_inbound,
|
||||||
@@ -527,6 +541,33 @@ class AgentLoop:
|
|||||||
**extra,
|
**extra,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
async def _get_or_create_session(self, key: str) -> Session:
|
||||||
|
"""Use native async session loading, with a compatibility fallback."""
|
||||||
|
return await call_session_manager(
|
||||||
|
self.sessions,
|
||||||
|
"get_or_create_async",
|
||||||
|
self.sessions.get_or_create,
|
||||||
|
key,
|
||||||
|
)
|
||||||
|
|
||||||
|
async def _save_session(self, session: Session) -> None:
|
||||||
|
"""Use native async session saving, with a compatibility fallback."""
|
||||||
|
await call_session_manager(
|
||||||
|
self.sessions,
|
||||||
|
"save_async",
|
||||||
|
self.sessions.save,
|
||||||
|
session,
|
||||||
|
)
|
||||||
|
|
||||||
|
async def _save_runtime_checkpoint(self, session: Session) -> None:
|
||||||
|
"""Use native async checkpoint saving, with a compatibility fallback."""
|
||||||
|
await call_session_manager(
|
||||||
|
self.sessions,
|
||||||
|
"save_runtime_checkpoint_async",
|
||||||
|
self.sessions.save_runtime_checkpoint,
|
||||||
|
session,
|
||||||
|
)
|
||||||
|
|
||||||
def _sync_subagent_runtime_limits(self) -> None:
|
def _sync_subagent_runtime_limits(self) -> None:
|
||||||
"""Keep subagent runtime limits aligned with mutable loop settings."""
|
"""Keep subagent runtime limits aligned with mutable loop settings."""
|
||||||
self.subagents.max_iterations = self.max_iterations
|
self.subagents.max_iterations = self.max_iterations
|
||||||
@@ -566,6 +607,30 @@ class AgentLoop:
|
|||||||
self.sessions.save(session)
|
self.sessions.save(session)
|
||||||
return self.llm_runtime()
|
return self.llm_runtime()
|
||||||
|
|
||||||
|
async def runtime_for_session_async(
|
||||||
|
self,
|
||||||
|
session: Session,
|
||||||
|
*,
|
||||||
|
recover_removed: bool = True,
|
||||||
|
) -> LLMRuntime:
|
||||||
|
"""Resolve a session runtime without blocking on recovery persistence."""
|
||||||
|
name = model_preset_from_metadata(session.metadata)
|
||||||
|
if name is None:
|
||||||
|
return self.llm_runtime()
|
||||||
|
try:
|
||||||
|
return self.runtime_resolver.resolve_preset(name)
|
||||||
|
except KeyError:
|
||||||
|
if not recover_removed or name in self.runtime_resolver.model_presets:
|
||||||
|
raise
|
||||||
|
logger.warning(
|
||||||
|
"Session '{}' references removed model preset '{}'; falling back to default",
|
||||||
|
session.key,
|
||||||
|
name,
|
||||||
|
)
|
||||||
|
session.metadata.pop(SESSION_MODEL_PRESET_METADATA_KEY, None)
|
||||||
|
await self._save_session(session)
|
||||||
|
return self.llm_runtime()
|
||||||
|
|
||||||
def set_session_model_preset(
|
def set_session_model_preset(
|
||||||
self,
|
self,
|
||||||
session_key: str,
|
session_key: str,
|
||||||
@@ -578,6 +643,18 @@ class AgentLoop:
|
|||||||
self.sessions.save(session)
|
self.sessions.save(session)
|
||||||
return runtime
|
return runtime
|
||||||
|
|
||||||
|
async def set_session_model_preset_async(
|
||||||
|
self,
|
||||||
|
session_key: str,
|
||||||
|
name: str,
|
||||||
|
) -> LLMRuntime:
|
||||||
|
"""Validate and persist one session's preset selection without blocking."""
|
||||||
|
runtime = self.runtime_resolver.resolve_preset(name)
|
||||||
|
session = await self._get_or_create_session(session_key)
|
||||||
|
session.metadata[SESSION_MODEL_PRESET_METADATA_KEY] = runtime.model_preset
|
||||||
|
await self._save_session(session)
|
||||||
|
return runtime
|
||||||
|
|
||||||
def _publish_runtime_selection(
|
def _publish_runtime_selection(
|
||||||
self,
|
self,
|
||||||
runtime: LLMRuntime,
|
runtime: LLMRuntime,
|
||||||
@@ -690,17 +767,14 @@ class AgentLoop:
|
|||||||
session_key=session_key,
|
session_key=session_key,
|
||||||
)
|
)
|
||||||
|
|
||||||
def _persist_user_message_early(
|
def _stage_user_message_early(
|
||||||
self,
|
self,
|
||||||
msg: InboundMessage,
|
msg: InboundMessage,
|
||||||
session: Session,
|
session: Session,
|
||||||
runtime_context_blocks: list[RuntimeContextBlock] | None = None,
|
runtime_context_blocks: list[RuntimeContextBlock] | None = None,
|
||||||
**kwargs: Any,
|
**kwargs: Any,
|
||||||
) -> bool:
|
) -> bool:
|
||||||
"""Persist the triggering user message before the turn starts.
|
"""Add the triggering user message and recovery markers in memory."""
|
||||||
|
|
||||||
Returns True if the message was persisted.
|
|
||||||
"""
|
|
||||||
if not turn_continuation.should_persist_user_message(msg.metadata):
|
if not turn_continuation.should_persist_user_message(msg.metadata):
|
||||||
return False
|
return False
|
||||||
media_paths = [
|
media_paths = [
|
||||||
@@ -726,10 +800,48 @@ class AgentLoop:
|
|||||||
extra[RUNTIME_CONTEXT_HISTORY_META] = runtime_context_meta
|
extra[RUNTIME_CONTEXT_HISTORY_META] = runtime_context_meta
|
||||||
session.add_message("user", text, **extra)
|
session.add_message("user", text, **extra)
|
||||||
self._mark_pending_user_turn(session)
|
self._mark_pending_user_turn(session)
|
||||||
self.sessions.save(session)
|
followup_id = msg.metadata.get(PENDING_FOLLOWUP_ID_KEY)
|
||||||
|
if isinstance(followup_id, str) and followup_id:
|
||||||
|
acknowledge_pending_followups(session, [followup_id])
|
||||||
return True
|
return True
|
||||||
return False
|
return False
|
||||||
|
|
||||||
|
def _persist_user_message_early(
|
||||||
|
self,
|
||||||
|
msg: InboundMessage,
|
||||||
|
session: Session,
|
||||||
|
runtime_context_blocks: list[RuntimeContextBlock] | None = None,
|
||||||
|
**kwargs: Any,
|
||||||
|
) -> bool:
|
||||||
|
"""Synchronously persist the user message for compatibility callers."""
|
||||||
|
persisted = self._stage_user_message_early(
|
||||||
|
msg,
|
||||||
|
session,
|
||||||
|
runtime_context_blocks,
|
||||||
|
**kwargs,
|
||||||
|
)
|
||||||
|
if persisted:
|
||||||
|
self.sessions.save(session)
|
||||||
|
return persisted
|
||||||
|
|
||||||
|
async def _persist_user_message_early_async(
|
||||||
|
self,
|
||||||
|
msg: InboundMessage,
|
||||||
|
session: Session,
|
||||||
|
runtime_context_blocks: list[RuntimeContextBlock] | None = None,
|
||||||
|
**kwargs: Any,
|
||||||
|
) -> bool:
|
||||||
|
"""Persist the user message without blocking the event loop."""
|
||||||
|
persisted = self._stage_user_message_early(
|
||||||
|
msg,
|
||||||
|
session,
|
||||||
|
runtime_context_blocks,
|
||||||
|
**kwargs,
|
||||||
|
)
|
||||||
|
if persisted:
|
||||||
|
await self._save_session(session)
|
||||||
|
return persisted
|
||||||
|
|
||||||
def _build_initial_messages(self, ctx: TurnContext) -> list[dict[str, Any]]:
|
def _build_initial_messages(self, ctx: TurnContext) -> list[dict[str, Any]]:
|
||||||
"""Build the initial message list for the LLM turn."""
|
"""Build the initial message list for the LLM turn."""
|
||||||
assert ctx.session is not None
|
assert ctx.session is not None
|
||||||
@@ -827,7 +939,7 @@ class AgentLoop:
|
|||||||
if tool is None:
|
if tool is None:
|
||||||
content = "Shell execution is disabled in this nanobot configuration."
|
content = "Shell execution is disabled in this nanobot configuration."
|
||||||
else:
|
else:
|
||||||
session = ctx.session or self.sessions.get_or_create(ctx.key)
|
session = ctx.session or await AgentLoop._get_or_create_session(self, ctx.key)
|
||||||
scope = self.workspace_scopes.for_turn(
|
scope = self.workspace_scopes.for_turn(
|
||||||
channel=ctx.msg.channel,
|
channel=ctx.msg.channel,
|
||||||
message_metadata=metadata,
|
message_metadata=metadata,
|
||||||
@@ -985,7 +1097,7 @@ class AgentLoop:
|
|||||||
public_payload[self._PROVIDER_STATE_CHECKPOINT_VERSION_KEY] = (
|
public_payload[self._PROVIDER_STATE_CHECKPOINT_VERSION_KEY] = (
|
||||||
self._PROVIDER_STATE_CHECKPOINT_VERSION
|
self._PROVIDER_STATE_CHECKPOINT_VERSION
|
||||||
)
|
)
|
||||||
self._set_runtime_checkpoint(session, public_payload)
|
await self._set_runtime_checkpoint_async(session, public_payload)
|
||||||
|
|
||||||
async def _drain_pending(*, limit: int = _MAX_INJECTIONS_PER_TURN) -> list[dict[str, Any]]:
|
async def _drain_pending(*, limit: int = _MAX_INJECTIONS_PER_TURN) -> list[dict[str, Any]]:
|
||||||
"""Drain follow-up messages from the pending queue.
|
"""Drain follow-up messages from the pending queue.
|
||||||
@@ -1061,6 +1173,9 @@ class AgentLoop:
|
|||||||
row["subagent_task_id"] = task_id
|
row["subagent_task_id"] = task_id
|
||||||
row[HIDDEN_HISTORY_META] = subagent_marker
|
row[HIDDEN_HISTORY_META] = subagent_marker
|
||||||
row["injected_event"] = "subagent_result"
|
row["injected_event"] = "subagent_result"
|
||||||
|
followup_id = metadata.get(PENDING_FOLLOWUP_ID_KEY)
|
||||||
|
if isinstance(followup_id, str) and followup_id:
|
||||||
|
row[PENDING_FOLLOWUP_ID_KEY] = followup_id
|
||||||
return row
|
return row
|
||||||
|
|
||||||
items: list[dict[str, Any]] = []
|
items: list[dict[str, Any]] = []
|
||||||
@@ -1184,6 +1299,11 @@ class AgentLoop:
|
|||||||
message_metadata=metadata,
|
message_metadata=metadata,
|
||||||
),
|
),
|
||||||
provider_state=provider_state,
|
provider_state=provider_state,
|
||||||
|
llm_usage_source=source_from_request(
|
||||||
|
active_session_key,
|
||||||
|
channel=channel,
|
||||||
|
metadata=metadata,
|
||||||
|
),
|
||||||
))
|
))
|
||||||
finally:
|
finally:
|
||||||
turn_scope_stack.close()
|
turn_scope_stack.close()
|
||||||
@@ -1215,18 +1335,33 @@ class AgentLoop:
|
|||||||
logger.error("LLM returned error: {}", (result.final_content or "")[:200])
|
logger.error("LLM returned error: {}", (result.final_content or "")[:200])
|
||||||
return result.final_content, result.tools_used, result.messages, result.stop_reason, result.had_injections
|
return result.final_content, result.tools_used, result.messages, result.stop_reason, result.had_injections
|
||||||
|
|
||||||
def _check_expired_sessions_if_due(self) -> None:
|
def _idle_compact_scan_due(self) -> bool:
|
||||||
"""Scan idle sessions no more often than the configured interval."""
|
|
||||||
now = time.monotonic()
|
now = time.monotonic()
|
||||||
if now < self._next_idle_compact_check_at:
|
if now < self._next_idle_compact_check_at:
|
||||||
return
|
return False
|
||||||
self._next_idle_compact_check_at = now + self._idle_compact_check_interval_s
|
self._next_idle_compact_check_at = now + self._idle_compact_check_interval_s
|
||||||
|
return True
|
||||||
|
|
||||||
|
def _check_expired_sessions_if_due(self) -> None:
|
||||||
|
"""Synchronously scan idle sessions for compatibility with direct callers."""
|
||||||
|
if not self._idle_compact_scan_due():
|
||||||
|
return
|
||||||
self.auto_compact.check_expired(
|
self.auto_compact.check_expired(
|
||||||
self.schedule_background,
|
self.schedule_background,
|
||||||
self.runtime_for_session,
|
self.runtime_for_session,
|
||||||
active_session_keys=self._pending_queues.keys(),
|
active_session_keys=self._pending_queues.keys(),
|
||||||
)
|
)
|
||||||
|
|
||||||
|
async def _check_expired_sessions_if_due_async(self) -> None:
|
||||||
|
"""Scan idle sessions without blocking the event loop."""
|
||||||
|
if not self._idle_compact_scan_due():
|
||||||
|
return
|
||||||
|
await self.auto_compact.check_expired_async(
|
||||||
|
self.schedule_background,
|
||||||
|
self.runtime_for_session_async,
|
||||||
|
active_session_keys=self._pending_queues.keys(),
|
||||||
|
)
|
||||||
|
|
||||||
async def run(self) -> None:
|
async def run(self) -> None:
|
||||||
"""Run the agent loop, dispatching messages as tasks to stay responsive to /stop."""
|
"""Run the agent loop, dispatching messages as tasks to stay responsive to /stop."""
|
||||||
self._running = True
|
self._running = True
|
||||||
@@ -1237,7 +1372,7 @@ class AgentLoop:
|
|||||||
try:
|
try:
|
||||||
msg = await asyncio.wait_for(self.bus.consume_inbound(), timeout=1.0)
|
msg = await asyncio.wait_for(self.bus.consume_inbound(), timeout=1.0)
|
||||||
except asyncio.TimeoutError:
|
except asyncio.TimeoutError:
|
||||||
self._check_expired_sessions_if_due()
|
await self._check_expired_sessions_if_due_async()
|
||||||
continue
|
continue
|
||||||
except asyncio.CancelledError:
|
except asyncio.CancelledError:
|
||||||
# Preserve real task cancellation so shutdown can complete cleanly.
|
# Preserve real task cancellation so shutdown can complete cleanly.
|
||||||
@@ -1285,6 +1420,23 @@ class AgentLoop:
|
|||||||
break
|
break
|
||||||
if deferred:
|
if deferred:
|
||||||
continue
|
continue
|
||||||
|
routed_msg = msg
|
||||||
|
if effective_key != msg.session_key:
|
||||||
|
routed_msg = dataclasses.replace(
|
||||||
|
msg,
|
||||||
|
session_key_override=effective_key,
|
||||||
|
)
|
||||||
|
# A newer WebUI message must supersede an explicit recovery
|
||||||
|
# before it is injected into that recovery's pending queue.
|
||||||
|
# Without this admission point, a recovered turn could finish
|
||||||
|
# first and only then observe the user's newer request.
|
||||||
|
if (
|
||||||
|
effective_key in self._pending_queues
|
||||||
|
and msg.channel == "websocket"
|
||||||
|
and self._recovery_admission is not None
|
||||||
|
and not await self._recovery_admission.admit(routed_msg)
|
||||||
|
):
|
||||||
|
continue
|
||||||
# If this session already has an active pending queue (i.e. a task
|
# If this session already has an active pending queue (i.e. a task
|
||||||
# is processing this session), route the message there for mid-turn
|
# is processing this session), route the message there for mid-turn
|
||||||
# injection instead of creating a competing task.
|
# injection instead of creating a competing task.
|
||||||
@@ -1297,12 +1449,18 @@ class AgentLoop:
|
|||||||
self.commands.dispatch,
|
self.commands.dispatch,
|
||||||
)
|
)
|
||||||
continue
|
continue
|
||||||
pending_msg = msg
|
pending_msg = routed_msg
|
||||||
if effective_key != msg.session_key:
|
session = await self._get_or_create_session(effective_key)
|
||||||
|
followup_id = record_pending_followup(session, pending_msg)
|
||||||
|
if followup_id is not None:
|
||||||
pending_msg = dataclasses.replace(
|
pending_msg = dataclasses.replace(
|
||||||
msg,
|
pending_msg,
|
||||||
session_key_override=effective_key,
|
metadata={
|
||||||
|
**pending_msg.metadata,
|
||||||
|
PENDING_FOLLOWUP_ID_KEY: followup_id,
|
||||||
|
},
|
||||||
)
|
)
|
||||||
|
await self._save_session(session)
|
||||||
try:
|
try:
|
||||||
self._pending_queues[effective_key].put_nowait(pending_msg)
|
self._pending_queues[effective_key].put_nowait(pending_msg)
|
||||||
except asyncio.QueueFull:
|
except asyncio.QueueFull:
|
||||||
@@ -1310,6 +1468,7 @@ class AgentLoop:
|
|||||||
"Pending queue full for session {}, falling back to queued task",
|
"Pending queue full for session {}, falling back to queued task",
|
||||||
effective_key,
|
effective_key,
|
||||||
)
|
)
|
||||||
|
msg = pending_msg
|
||||||
else:
|
else:
|
||||||
logger.info(
|
logger.info(
|
||||||
"Routed follow-up message to pending queue for session {}",
|
"Routed follow-up message to pending queue for session {}",
|
||||||
@@ -1319,17 +1478,45 @@ class AgentLoop:
|
|||||||
# Compute the effective session key before dispatching
|
# Compute the effective session key before dispatching
|
||||||
# This ensures /stop command can find tasks correctly when unified session is enabled
|
# This ensures /stop command can find tasks correctly when unified session is enabled
|
||||||
task = asyncio.create_task(self._dispatch(msg))
|
task = asyncio.create_task(self._dispatch(msg))
|
||||||
active_tasks = self._active_tasks.setdefault(effective_key, set())
|
active_tasks: set[asyncio.Task[Any]] = self._active_tasks.setdefault(
|
||||||
|
effective_key,
|
||||||
|
set(),
|
||||||
|
)
|
||||||
active_tasks.add(task)
|
active_tasks.add(task)
|
||||||
task.add_done_callback(active_tasks.discard)
|
task.add_done_callback(active_tasks.discard)
|
||||||
finally:
|
finally:
|
||||||
await self.aclose()
|
await self.aclose()
|
||||||
|
|
||||||
|
def preserve_inflight_turns_on_shutdown(self) -> None:
|
||||||
|
"""Keep durable checkpoints when the owning gateway exits.
|
||||||
|
|
||||||
|
Normal cancellation intentionally materializes partial output so a
|
||||||
|
user-stopped turn leaves a readable conversation. Gateway lifecycle
|
||||||
|
shutdown is different: RecoveryCoordinator needs the checkpoint intact
|
||||||
|
to safely offer the unfinished turn for explicit continuation later.
|
||||||
|
"""
|
||||||
|
self._preserve_inflight_turns_on_shutdown = True
|
||||||
|
|
||||||
async def _dispatch(self, msg: InboundMessage) -> None:
|
async def _dispatch(self, msg: InboundMessage) -> None:
|
||||||
"""Process a message: per-session serial, cross-session concurrent."""
|
"""Process a message: per-session serial, cross-session concurrent."""
|
||||||
session_key = self._effective_session_key(msg)
|
session_key = self._effective_session_key(msg)
|
||||||
if session_key != msg.session_key:
|
if session_key != msg.session_key:
|
||||||
msg = dataclasses.replace(msg, session_key_override=session_key)
|
msg = dataclasses.replace(msg, session_key_override=session_key)
|
||||||
|
recovery_task_registered = False
|
||||||
|
recovery_admission = self._recovery_admission
|
||||||
|
current_task: asyncio.Task[Any] | None = None
|
||||||
|
if recovery_admission is not None:
|
||||||
|
recovery_id = msg.metadata.get(RECOVERY_INBOUND_METADATA_KEY)
|
||||||
|
if isinstance(recovery_id, str) and recovery_id:
|
||||||
|
current_task = asyncio.current_task()
|
||||||
|
if current_task is not None:
|
||||||
|
recovery_admission.register_recovery_task(session_key, current_task)
|
||||||
|
recovery_task_registered = True
|
||||||
|
if not await recovery_admission.admit(msg):
|
||||||
|
logger.info("Skipped stale recovery for session {}", session_key)
|
||||||
|
if recovery_task_registered and current_task is not None:
|
||||||
|
recovery_admission.unregister_recovery_task(session_key, current_task)
|
||||||
|
return
|
||||||
lock = self._get_session_lock(session_key)
|
lock = self._get_session_lock(session_key)
|
||||||
gate = self._concurrency_gate or nullcontext()
|
gate = self._concurrency_gate or nullcontext()
|
||||||
|
|
||||||
@@ -1373,21 +1560,21 @@ class AgentLoop:
|
|||||||
session_key,
|
session_key,
|
||||||
exc_info=True,
|
exc_info=True,
|
||||||
)
|
)
|
||||||
# Preserve partial context from the interrupted turn so
|
# An explicit turn stop materializes partial context so
|
||||||
# the user does not lose tool results and assistant
|
# the next prompt can see completed tool results. Gateway
|
||||||
# messages accumulated before /stop. The checkpoint was
|
# shutdown keeps the durable checkpoint untouched instead,
|
||||||
# already persisted to session metadata by
|
# allowing RecoveryCoordinator to offer Continue safely.
|
||||||
# _emit_checkpoint during tool execution; materializing
|
if (
|
||||||
# it into session history now makes it visible in the
|
session_key in self._discarding_sessions
|
||||||
# next conversation turn.
|
or self._preserve_inflight_turns_on_shutdown
|
||||||
if session_key in self._discarding_sessions:
|
):
|
||||||
raise
|
raise
|
||||||
try:
|
try:
|
||||||
key = self._effective_session_key(msg)
|
key = self._effective_session_key(msg)
|
||||||
session = self.sessions.get_or_create(key)
|
session = await self._get_or_create_session(key)
|
||||||
if self._restore_runtime_checkpoint(session):
|
if self._restore_runtime_checkpoint(session):
|
||||||
self._clear_pending_user_turn(session)
|
self._clear_pending_user_turn(session)
|
||||||
self.sessions.save(session)
|
await self._save_session(session)
|
||||||
logger.info(
|
logger.info(
|
||||||
"Restored partial context for cancelled session {}",
|
"Restored partial context for cancelled session {}",
|
||||||
key,
|
key,
|
||||||
@@ -1437,6 +1624,12 @@ class AgentLoop:
|
|||||||
await delivery.idle()
|
await delivery.idle()
|
||||||
await self._publish_next_deferred_automation_turn(session_key)
|
await self._publish_next_deferred_automation_turn(session_key)
|
||||||
finally:
|
finally:
|
||||||
|
if (
|
||||||
|
recovery_task_registered
|
||||||
|
and current_task is not None
|
||||||
|
and recovery_admission is not None
|
||||||
|
):
|
||||||
|
recovery_admission.unregister_recovery_task(session_key, current_task)
|
||||||
if pending is None:
|
if pending is None:
|
||||||
await delivery.idle()
|
await delivery.idle()
|
||||||
await self._publish_next_deferred_automation_turn(session_key)
|
await self._publish_next_deferred_automation_turn(session_key)
|
||||||
@@ -1706,7 +1899,7 @@ class AgentLoop:
|
|||||||
if ctx.session is None:
|
if ctx.session is None:
|
||||||
raise RuntimeError("required session is not active")
|
raise RuntimeError("required session is not active")
|
||||||
else:
|
else:
|
||||||
ctx.session = self.sessions.get_or_create(ctx.session_key)
|
ctx.session = await self._get_or_create_session(ctx.session_key)
|
||||||
session = ctx.session
|
session = ctx.session
|
||||||
ctx.ephemeral = ctx.ephemeral or not session.policy.persist
|
ctx.ephemeral = ctx.ephemeral or not session.policy.persist
|
||||||
tools = ctx.tools or self.tools
|
tools = ctx.tools or self.tools
|
||||||
@@ -1737,13 +1930,16 @@ class AgentLoop:
|
|||||||
self.workspace_scopes.persist_message_scope(session, msg)
|
self.workspace_scopes.persist_message_scope(session, msg)
|
||||||
|
|
||||||
if self._restore_runtime_checkpoint(session):
|
if self._restore_runtime_checkpoint(session):
|
||||||
self.sessions.save(session)
|
await self._save_session(session)
|
||||||
if self._restore_pending_user_turn(session):
|
if (
|
||||||
self.sessions.save(session)
|
RECOVERY_INBOUND_METADATA_KEY not in msg.metadata
|
||||||
|
and restore_pending_interruption(session)
|
||||||
|
):
|
||||||
|
await self._save_session(session)
|
||||||
|
|
||||||
async def _compact_session(self, ctx: TurnContext) -> None:
|
async def _compact_session(self, ctx: TurnContext) -> None:
|
||||||
session = ctx.require_session()
|
session = ctx.require_session()
|
||||||
ctx.session, pending = self.auto_compact.prepare_session(
|
ctx.session, pending = await self.auto_compact.prepare_session_async(
|
||||||
session,
|
session,
|
||||||
ctx.session_key,
|
ctx.session_key,
|
||||||
)
|
)
|
||||||
@@ -1780,14 +1976,14 @@ class AgentLoop:
|
|||||||
# them out of LLM context. /new is excluded because it
|
# them out of LLM context. /new is excluded because it
|
||||||
# intentionally clears the session.
|
# intentionally clears the session.
|
||||||
if cmd_ctx.raw.lower() != "/new":
|
if cmd_ctx.raw.lower() != "/new":
|
||||||
ctx.input_persisted_early = self._persist_user_message_early(
|
ctx.input_persisted_early = await self._persist_user_message_early_async(
|
||||||
ctx.msg, session, _command=True
|
ctx.msg, session, _command=True
|
||||||
)
|
)
|
||||||
session.add_message(
|
session.add_message(
|
||||||
"assistant", result.content, _command=True
|
"assistant", result.content, _command=True
|
||||||
)
|
)
|
||||||
self._clear_pending_user_turn(session)
|
self._clear_pending_user_turn(session)
|
||||||
self.sessions.save(session)
|
await self._save_session(session)
|
||||||
if not ctx.ephemeral:
|
if not ctx.ephemeral:
|
||||||
await self.runtime_event_publisher.session_turn_persisted(
|
await self.runtime_event_publisher.session_turn_persisted(
|
||||||
ctx.msg,
|
ctx.msg,
|
||||||
@@ -1802,7 +1998,7 @@ class AgentLoop:
|
|||||||
session = ctx.require_session()
|
session = ctx.require_session()
|
||||||
runtime = ctx.runtime
|
runtime = ctx.runtime
|
||||||
if runtime is None:
|
if runtime is None:
|
||||||
runtime = self.runtime_for_session(session)
|
runtime = await self.runtime_for_session_async(session)
|
||||||
ctx.runtime = runtime
|
ctx.runtime = runtime
|
||||||
if ctx.session_key.startswith("dream:"):
|
if ctx.session_key.startswith("dream:"):
|
||||||
logger.info(
|
logger.info(
|
||||||
@@ -1846,7 +2042,7 @@ class AgentLoop:
|
|||||||
# provider compatibility or prompt assembly work. A compatible
|
# provider compatibility or prompt assembly work. A compatible
|
||||||
# staged state replaces this in a second atomic save below.
|
# staged state replaces this in a second atomic save below.
|
||||||
session.provider_state = None
|
session.provider_state = None
|
||||||
self.sessions.save(session)
|
await self._save_session(session)
|
||||||
ctx.input_persisted_early = True
|
ctx.input_persisted_early = True
|
||||||
await ctx.delivery.runtime_admitted(runtime)
|
await ctx.delivery.runtime_admitted(runtime)
|
||||||
|
|
||||||
@@ -1900,7 +2096,7 @@ class AgentLoop:
|
|||||||
elif stored_state is not None:
|
elif stored_state is not None:
|
||||||
session.provider_state = None
|
session.provider_state = None
|
||||||
if ctx.kind is TurnKind.USER:
|
if ctx.kind is TurnKind.USER:
|
||||||
ctx.input_persisted_early = self._persist_user_message_early(
|
ctx.input_persisted_early = await self._persist_user_message_early_async(
|
||||||
ctx.msg,
|
ctx.msg,
|
||||||
session,
|
session,
|
||||||
runtime_context_blocks=ctx.runtime_context_blocks,
|
runtime_context_blocks=ctx.runtime_context_blocks,
|
||||||
@@ -1910,7 +2106,7 @@ class AgentLoop:
|
|||||||
elif subagent_followup_persisted and staged_provider_state:
|
elif subagent_followup_persisted and staged_provider_state:
|
||||||
# Upgrade the replay-safe baseline to the resumable state before
|
# Upgrade the replay-safe baseline to the resumable state before
|
||||||
# prompt assembly and the first model checkpoint.
|
# prompt assembly and the first model checkpoint.
|
||||||
self.sessions.save(session)
|
await self._save_session(session)
|
||||||
ctx.initial_messages = self._build_initial_messages(ctx)
|
ctx.initial_messages = self._build_initial_messages(ctx)
|
||||||
|
|
||||||
if ctx.on_progress is None:
|
if ctx.on_progress is None:
|
||||||
@@ -1952,7 +2148,7 @@ class AgentLoop:
|
|||||||
ctx.all_messages = all_msgs
|
ctx.all_messages = all_msgs
|
||||||
ctx.stop_reason = stop_reason
|
ctx.stop_reason = stop_reason
|
||||||
ctx.had_injections = had_injections
|
ctx.had_injections = had_injections
|
||||||
ctx.usage = dict(self._last_usage)
|
ctx.usage = self._last_usage
|
||||||
ctx.delivery.record_usage(ctx.usage)
|
ctx.delivery.record_usage(ctx.usage)
|
||||||
if ctx.kind is TurnKind.USER:
|
if ctx.kind is TurnKind.USER:
|
||||||
await turn_continuation.maybe_continue_turn(ctx)
|
await turn_continuation.maybe_continue_turn(ctx)
|
||||||
@@ -1979,13 +2175,16 @@ class AgentLoop:
|
|||||||
else ctx.turn_wall_started_at
|
else ctx.turn_wall_started_at
|
||||||
)
|
)
|
||||||
ctx.turn_latency_ms = max(0, int((time.time() - latency_started_at) * 1000))
|
ctx.turn_latency_ms = max(0, int((time.time() - latency_started_at) * 1000))
|
||||||
if ctx.usage and not ctx.ephemeral:
|
if ctx.usage is not None and not ctx.ephemeral:
|
||||||
session.metadata["_last_usage"] = dict(ctx.usage)
|
session.metadata["_last_usage"] = ctx.usage.to_dict()
|
||||||
self._save_turn(
|
self._save_turn(
|
||||||
session, ctx.all_messages, ctx.save_skip,
|
session, ctx.all_messages, ctx.save_skip,
|
||||||
turn_latency_ms=ctx.turn_latency_ms,
|
turn_latency_ms=ctx.turn_latency_ms,
|
||||||
)
|
)
|
||||||
ctx.delivery.record_latency(ctx.turn_latency_ms)
|
ctx.delivery.record_latency(ctx.turn_latency_ms)
|
||||||
|
self._clear_pending_user_turn(session)
|
||||||
|
self._clear_runtime_checkpoint(session)
|
||||||
|
await self._save_session(session)
|
||||||
if not ctx.ephemeral:
|
if not ctx.ephemeral:
|
||||||
self.schedule_background(
|
self.schedule_background(
|
||||||
self.consolidator.maybe_consolidate_by_tokens(
|
self.consolidator.maybe_consolidate_by_tokens(
|
||||||
@@ -1993,10 +2192,6 @@ class AgentLoop:
|
|||||||
runtime=runtime,
|
runtime=runtime,
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
self._clear_pending_user_turn(session)
|
|
||||||
self._clear_runtime_checkpoint(session)
|
|
||||||
self.sessions.save(session)
|
|
||||||
if not ctx.ephemeral:
|
|
||||||
await self.runtime_event_publisher.session_turn_persisted(
|
await self.runtime_event_publisher.session_turn_persisted(
|
||||||
ctx.msg,
|
ctx.msg,
|
||||||
ctx.session_key,
|
ctx.session_key,
|
||||||
@@ -2093,8 +2288,21 @@ class AgentLoop:
|
|||||||
if m.get("role") == "tool" and m.get("tool_call_id")
|
if m.get("role") == "tool" and m.get("tool_call_id")
|
||||||
}
|
}
|
||||||
last_assistant_idx: int | None = None
|
last_assistant_idx: int | None = None
|
||||||
|
saved_followup_ids: set[str] = set()
|
||||||
for m in messages[skip:]:
|
for m in messages[skip:]:
|
||||||
entry = dict(m)
|
entry = dict(m)
|
||||||
|
followup_id_value = cast(object, entry.pop(PENDING_FOLLOWUP_ID_KEY, None))
|
||||||
|
followup_ids = (
|
||||||
|
[followup_id_value]
|
||||||
|
if isinstance(followup_id_value, str)
|
||||||
|
else [
|
||||||
|
followup_id
|
||||||
|
for followup_id in cast(list[object], followup_id_value)
|
||||||
|
if isinstance(followup_id, str)
|
||||||
|
]
|
||||||
|
if isinstance(followup_id_value, list)
|
||||||
|
else []
|
||||||
|
)
|
||||||
internal_meta = cast(object, entry.pop("_meta", None))
|
internal_meta = cast(object, entry.pop("_meta", None))
|
||||||
runtime_context_meta = (
|
runtime_context_meta = (
|
||||||
cast(dict[str, Any], internal_meta).get(
|
cast(dict[str, Any], internal_meta).get(
|
||||||
@@ -2147,6 +2355,8 @@ class AgentLoop:
|
|||||||
entry[RUNTIME_CONTEXT_HISTORY_META] = runtime_context_meta
|
entry[RUNTIME_CONTEXT_HISTORY_META] = runtime_context_meta
|
||||||
entry.setdefault("timestamp", datetime.now().isoformat())
|
entry.setdefault("timestamp", datetime.now().isoformat())
|
||||||
session.messages.append(entry)
|
session.messages.append(entry)
|
||||||
|
if role == "user":
|
||||||
|
saved_followup_ids.update(followup_id for followup_id in followup_ids if followup_id)
|
||||||
if role == "assistant":
|
if role == "assistant":
|
||||||
last_assistant_idx = len(session.messages) - 1
|
last_assistant_idx = len(session.messages) - 1
|
||||||
declared_tool_call_ids.update(
|
declared_tool_call_ids.update(
|
||||||
@@ -2161,6 +2371,8 @@ class AgentLoop:
|
|||||||
)
|
)
|
||||||
if turn_latency_ms is not None and last_assistant_idx is not None:
|
if turn_latency_ms is not None and last_assistant_idx is not None:
|
||||||
session.messages[last_assistant_idx]["latency_ms"] = int(turn_latency_ms)
|
session.messages[last_assistant_idx]["latency_ms"] = int(turn_latency_ms)
|
||||||
|
if saved_followup_ids:
|
||||||
|
acknowledge_pending_followups(session, saved_followup_ids)
|
||||||
session.updated_at = datetime.now()
|
session.updated_at = datetime.now()
|
||||||
|
|
||||||
def _persist_subagent_followup(self, session: Session, msg: InboundMessage) -> bool:
|
def _persist_subagent_followup(self, session: Session, msg: InboundMessage) -> bool:
|
||||||
@@ -2192,10 +2404,23 @@ class AgentLoop:
|
|||||||
)
|
)
|
||||||
return True
|
return True
|
||||||
|
|
||||||
def _set_runtime_checkpoint(self, session: Session, payload: dict[str, Any]) -> None:
|
def _set_runtime_checkpoint(
|
||||||
"""Persist the latest in-flight turn state into session metadata."""
|
self,
|
||||||
|
session: Session,
|
||||||
|
payload: dict[str, Any],
|
||||||
|
) -> None:
|
||||||
|
"""Synchronously persist a checkpoint for compatibility callers."""
|
||||||
session.metadata[self._RUNTIME_CHECKPOINT_KEY] = payload
|
session.metadata[self._RUNTIME_CHECKPOINT_KEY] = payload
|
||||||
self.sessions.save(session)
|
self.sessions.save_runtime_checkpoint(session)
|
||||||
|
|
||||||
|
async def _set_runtime_checkpoint_async(
|
||||||
|
self,
|
||||||
|
session: Session,
|
||||||
|
payload: dict[str, Any],
|
||||||
|
) -> None:
|
||||||
|
"""Persist the latest in-flight turn state without blocking the event loop."""
|
||||||
|
session.metadata[self._RUNTIME_CHECKPOINT_KEY] = payload
|
||||||
|
await self._save_runtime_checkpoint(session)
|
||||||
|
|
||||||
def _mark_pending_user_turn(self, session: Session) -> None:
|
def _mark_pending_user_turn(self, session: Session) -> None:
|
||||||
session.metadata[self._PENDING_USER_TURN_KEY] = True
|
session.metadata[self._PENDING_USER_TURN_KEY] = True
|
||||||
@@ -2207,136 +2432,9 @@ class AgentLoop:
|
|||||||
if self._RUNTIME_CHECKPOINT_KEY in session.metadata:
|
if self._RUNTIME_CHECKPOINT_KEY in session.metadata:
|
||||||
session.metadata.pop(self._RUNTIME_CHECKPOINT_KEY, None)
|
session.metadata.pop(self._RUNTIME_CHECKPOINT_KEY, None)
|
||||||
|
|
||||||
@staticmethod
|
|
||||||
def _checkpoint_message_key(message: dict[str, Any]) -> tuple[Any, ...]:
|
|
||||||
return (
|
|
||||||
message.get("role"),
|
|
||||||
message.get("content"),
|
|
||||||
message.get("tool_call_id"),
|
|
||||||
message.get("name"),
|
|
||||||
message.get("tool_calls"),
|
|
||||||
message.get("reasoning_content"),
|
|
||||||
message.get("thinking_blocks"),
|
|
||||||
)
|
|
||||||
|
|
||||||
def _restore_runtime_checkpoint(self, session: Session) -> bool:
|
def _restore_runtime_checkpoint(self, session: Session) -> bool:
|
||||||
"""Materialize an unfinished turn into session history before a new request."""
|
"""Materialize an unfinished turn into session history before a new request."""
|
||||||
from datetime import datetime
|
return restore_runtime_checkpoint(session)
|
||||||
|
|
||||||
checkpoint = cast(
|
|
||||||
object,
|
|
||||||
session.metadata.get(self._RUNTIME_CHECKPOINT_KEY),
|
|
||||||
)
|
|
||||||
if not isinstance(checkpoint, dict):
|
|
||||||
return False
|
|
||||||
checkpoint_data = cast(dict[str, Any], checkpoint)
|
|
||||||
|
|
||||||
assistant_message = cast(object, checkpoint_data.get("assistant_message"))
|
|
||||||
completed_tool_results = cast(
|
|
||||||
Iterable[object],
|
|
||||||
checkpoint_data.get("completed_tool_results") or [],
|
|
||||||
)
|
|
||||||
pending_tool_calls = cast(
|
|
||||||
Iterable[object],
|
|
||||||
checkpoint_data.get("pending_tool_calls") or [],
|
|
||||||
)
|
|
||||||
|
|
||||||
restored_messages: list[dict[str, Any]] = []
|
|
||||||
if isinstance(assistant_message, dict):
|
|
||||||
restored = dict(cast(dict[str, Any], assistant_message))
|
|
||||||
restored.setdefault("timestamp", datetime.now().isoformat())
|
|
||||||
restored_messages.append(restored)
|
|
||||||
for message in completed_tool_results:
|
|
||||||
if isinstance(message, dict):
|
|
||||||
restored = dict(cast(dict[str, Any], message))
|
|
||||||
restored.setdefault("timestamp", datetime.now().isoformat())
|
|
||||||
restored_messages.append(restored)
|
|
||||||
for tool_call in pending_tool_calls:
|
|
||||||
if not isinstance(tool_call, dict):
|
|
||||||
continue
|
|
||||||
tool_call_data = cast(dict[str, Any], tool_call)
|
|
||||||
tool_id = tool_call_data.get("id")
|
|
||||||
function_data = cast(
|
|
||||||
dict[str, Any],
|
|
||||||
tool_call_data.get("function") or {},
|
|
||||||
)
|
|
||||||
name = function_data.get("name") or "tool"
|
|
||||||
restored_messages.append(
|
|
||||||
{
|
|
||||||
"role": "tool",
|
|
||||||
"tool_call_id": tool_id,
|
|
||||||
"name": name,
|
|
||||||
"content": "Error: Task interrupted before this tool finished.",
|
|
||||||
"timestamp": datetime.now().isoformat(),
|
|
||||||
}
|
|
||||||
)
|
|
||||||
|
|
||||||
overlap = 0
|
|
||||||
max_overlap = min(len(session.messages), len(restored_messages))
|
|
||||||
for size in range(max_overlap, 0, -1):
|
|
||||||
existing = session.messages[-size:]
|
|
||||||
restored = restored_messages[:size]
|
|
||||||
if all(
|
|
||||||
self._checkpoint_message_key(left) == self._checkpoint_message_key(right)
|
|
||||||
for left, right in zip(existing, restored)
|
|
||||||
):
|
|
||||||
overlap = size
|
|
||||||
break
|
|
||||||
appended_messages = restored_messages[overlap:]
|
|
||||||
session.messages.extend(appended_messages)
|
|
||||||
assistant_message_data = (
|
|
||||||
cast(dict[str, Any], assistant_message)
|
|
||||||
if isinstance(assistant_message, dict)
|
|
||||||
else None
|
|
||||||
)
|
|
||||||
provider_state_is_synchronized = (
|
|
||||||
checkpoint_data.get(self._PROVIDER_STATE_CHECKPOINT_VERSION_KEY)
|
|
||||||
== self._PROVIDER_STATE_CHECKPOINT_VERSION
|
|
||||||
)
|
|
||||||
phase = checkpoint_data.get("phase")
|
|
||||||
exact_final_response = (
|
|
||||||
phase == "final_response"
|
|
||||||
and assistant_message_data is not None
|
|
||||||
and assistant_message_data.get("role") == "assistant"
|
|
||||||
and not bool(checkpoint_data.get("completed_tool_results"))
|
|
||||||
and not bool(checkpoint_data.get("pending_tool_calls"))
|
|
||||||
)
|
|
||||||
exact_completed_tools = (
|
|
||||||
phase == "tools_completed"
|
|
||||||
and assistant_message_data is not None
|
|
||||||
and assistant_message_data.get("role") == "assistant"
|
|
||||||
and not bool(checkpoint_data.get("pending_tool_calls"))
|
|
||||||
)
|
|
||||||
if not (
|
|
||||||
provider_state_is_synchronized
|
|
||||||
and (exact_final_response or exact_completed_tools)
|
|
||||||
):
|
|
||||||
session.provider_state = None
|
|
||||||
|
|
||||||
self._clear_pending_user_turn(session)
|
|
||||||
self._clear_runtime_checkpoint(session)
|
|
||||||
return True
|
|
||||||
|
|
||||||
def _restore_pending_user_turn(self, session: Session) -> bool:
|
|
||||||
"""Close a turn that only persisted the user message before crashing."""
|
|
||||||
from datetime import datetime
|
|
||||||
|
|
||||||
if not session.metadata.get(self._PENDING_USER_TURN_KEY):
|
|
||||||
return False
|
|
||||||
|
|
||||||
if session.messages and session.messages[-1].get("role") == "user":
|
|
||||||
session.messages.append(
|
|
||||||
{
|
|
||||||
"role": "assistant",
|
|
||||||
"content": "Error: Task interrupted before a response was generated.",
|
|
||||||
"timestamp": datetime.now().isoformat(),
|
|
||||||
}
|
|
||||||
)
|
|
||||||
session.provider_state = None
|
|
||||||
session.updated_at = datetime.now()
|
|
||||||
|
|
||||||
self._clear_pending_user_turn(session)
|
|
||||||
return True
|
|
||||||
|
|
||||||
async def process_direct(
|
async def process_direct(
|
||||||
self,
|
self,
|
||||||
|
|||||||
+37
-18
@@ -20,7 +20,9 @@ from typing import TYPE_CHECKING, Any, Callable, Iterator, cast
|
|||||||
|
|
||||||
from loguru import logger
|
from loguru import logger
|
||||||
|
|
||||||
|
from nanobot.llm_usage.context import llm_usage_source
|
||||||
from nanobot.runtime_context import public_history_messages
|
from nanobot.runtime_context import public_history_messages
|
||||||
|
from nanobot.session.async_compat import call_session_manager
|
||||||
from nanobot.session.manager import (
|
from nanobot.session.manager import (
|
||||||
MIN_COMPACTED_REPLAY_MESSAGES,
|
MIN_COMPACTED_REPLAY_MESSAGES,
|
||||||
Session,
|
Session,
|
||||||
@@ -823,6 +825,22 @@ class Consolidator:
|
|||||||
weakref.WeakValueDictionary()
|
weakref.WeakValueDictionary()
|
||||||
)
|
)
|
||||||
|
|
||||||
|
async def _get_or_create_session(self, key: str) -> Session:
|
||||||
|
return await call_session_manager(
|
||||||
|
self.sessions,
|
||||||
|
"get_or_create_async",
|
||||||
|
self.sessions.get_or_create,
|
||||||
|
key,
|
||||||
|
)
|
||||||
|
|
||||||
|
async def _save_session(self, session: Session) -> None:
|
||||||
|
await call_session_manager(
|
||||||
|
self.sessions,
|
||||||
|
"save_async",
|
||||||
|
self.sessions.save,
|
||||||
|
session,
|
||||||
|
)
|
||||||
|
|
||||||
def get_lock(self, session_key: str) -> asyncio.Lock:
|
def get_lock(self, session_key: str) -> asyncio.Lock:
|
||||||
"""Return the shared consolidation lock for one session."""
|
"""Return the shared consolidation lock for one session."""
|
||||||
return self._locks.setdefault(session_key, asyncio.Lock())
|
return self._locks.setdefault(session_key, asyncio.Lock())
|
||||||
@@ -858,13 +876,13 @@ class Consolidator:
|
|||||||
return []
|
return []
|
||||||
return session.get_history()
|
return session.get_history()
|
||||||
|
|
||||||
def _persist_last_summary(self, session: Session, summary: str | None) -> None:
|
async def _persist_last_summary(self, session: Session, summary: str | None) -> None:
|
||||||
if summary and summary != "(nothing)":
|
if summary and summary != "(nothing)":
|
||||||
session.metadata["_last_summary"] = {
|
session.metadata["_last_summary"] = {
|
||||||
"text": summary,
|
"text": summary,
|
||||||
"last_active": session.updated_at.isoformat(),
|
"last_active": session.updated_at.isoformat(),
|
||||||
}
|
}
|
||||||
self.sessions.save(session)
|
await self._save_session(session)
|
||||||
|
|
||||||
def estimate_session_prompt_tokens(
|
def estimate_session_prompt_tokens(
|
||||||
self,
|
self,
|
||||||
@@ -915,15 +933,16 @@ class Consolidator:
|
|||||||
if not messages:
|
if not messages:
|
||||||
return None
|
return None
|
||||||
try:
|
try:
|
||||||
response = await runtime.provider.chat_with_retry(
|
with llm_usage_source("dream"):
|
||||||
model=runtime.model,
|
response = await runtime.provider.chat_with_retry(
|
||||||
messages=request_messages,
|
model=runtime.model,
|
||||||
tools=request_tools,
|
messages=request_messages,
|
||||||
tool_choice="none",
|
tools=request_tools,
|
||||||
temperature=runtime.generation.temperature,
|
tool_choice="none",
|
||||||
max_tokens=runtime.generation.max_tokens,
|
temperature=runtime.generation.temperature,
|
||||||
reasoning_effort=runtime.generation.reasoning_effort,
|
max_tokens=runtime.generation.max_tokens,
|
||||||
)
|
reasoning_effort=runtime.generation.reasoning_effort,
|
||||||
|
)
|
||||||
except Exception:
|
except Exception:
|
||||||
logger.warning("Consolidation provider call failed, raw-dumping to history")
|
logger.warning("Consolidation provider call failed, raw-dumping to history")
|
||||||
self.store.raw_archive(messages, session_key=session_key)
|
self.store.raw_archive(messages, session_key=session_key)
|
||||||
@@ -1055,7 +1074,7 @@ class Consolidator:
|
|||||||
lock = self.get_lock(session.key)
|
lock = self.get_lock(session.key)
|
||||||
async with lock:
|
async with lock:
|
||||||
# Refresh session reference: AutoCompact may have replaced it.
|
# Refresh session reference: AutoCompact may have replaced it.
|
||||||
fresh = self.sessions.get_or_create(session.key)
|
fresh = await self._get_or_create_session(session.key)
|
||||||
if fresh is not session:
|
if fresh is not session:
|
||||||
session = fresh
|
session = fresh
|
||||||
if not session.messages:
|
if not session.messages:
|
||||||
@@ -1069,7 +1088,7 @@ class Consolidator:
|
|||||||
runtime=runtime,
|
runtime=runtime,
|
||||||
)
|
)
|
||||||
if estimated <= 0:
|
if estimated <= 0:
|
||||||
self._persist_last_summary(session, last_summary)
|
await self._persist_last_summary(session, last_summary)
|
||||||
return
|
return
|
||||||
if estimated < budget:
|
if estimated < budget:
|
||||||
unconsolidated_count = len(session.messages) - session.last_consolidated
|
unconsolidated_count = len(session.messages) - session.last_consolidated
|
||||||
@@ -1081,7 +1100,7 @@ class Consolidator:
|
|||||||
source,
|
source,
|
||||||
unconsolidated_count,
|
unconsolidated_count,
|
||||||
)
|
)
|
||||||
self._persist_last_summary(session, last_summary)
|
await self._persist_last_summary(session, last_summary)
|
||||||
return
|
return
|
||||||
|
|
||||||
for round_num in range(self._MAX_CONSOLIDATION_ROUNDS):
|
for round_num in range(self._MAX_CONSOLIDATION_ROUNDS):
|
||||||
@@ -1125,7 +1144,7 @@ class Consolidator:
|
|||||||
last_summary = summary
|
last_summary = summary
|
||||||
session.last_consolidated = end_idx
|
session.last_consolidated = end_idx
|
||||||
session.provider_state = None
|
session.provider_state = None
|
||||||
self.sessions.save(session)
|
await self._save_session(session)
|
||||||
if not summary:
|
if not summary:
|
||||||
# LLM is degraded — stop hammering it this call;
|
# LLM is degraded — stop hammering it this call;
|
||||||
# the next invocation can retry a fresh chunk.
|
# the next invocation can retry a fresh chunk.
|
||||||
@@ -1141,7 +1160,7 @@ class Consolidator:
|
|||||||
# Persist the last summary to session metadata so it can be injected
|
# Persist the last summary to session metadata so it can be injected
|
||||||
# into the runtime context on the next prepare_session() call, aligning
|
# into the runtime context on the next prepare_session() call, aligning
|
||||||
# the summary injection strategy with AutoCompact._archive().
|
# the summary injection strategy with AutoCompact._archive().
|
||||||
self._persist_last_summary(session, last_summary)
|
await self._persist_last_summary(session, last_summary)
|
||||||
|
|
||||||
async def compact_idle_session(
|
async def compact_idle_session(
|
||||||
self,
|
self,
|
||||||
@@ -1166,7 +1185,7 @@ class Consolidator:
|
|||||||
lock = self.get_lock(session_key)
|
lock = self.get_lock(session_key)
|
||||||
async with lock:
|
async with lock:
|
||||||
self.sessions.invalidate(session_key)
|
self.sessions.invalidate(session_key)
|
||||||
session = self.sessions.get_or_create(session_key)
|
session = await self._get_or_create_session(session_key)
|
||||||
|
|
||||||
archive_start = session.last_consolidated
|
archive_start = session.last_consolidated
|
||||||
messages_to_archive = list(session.messages[archive_start:])
|
messages_to_archive = list(session.messages[archive_start:])
|
||||||
@@ -1191,7 +1210,7 @@ class Consolidator:
|
|||||||
# through the captured batch so new messages remain eligible next time.
|
# through the captured batch so new messages remain eligible next time.
|
||||||
session.last_consolidated = archive_end
|
session.last_consolidated = archive_end
|
||||||
session.provider_state = None
|
session.provider_state = None
|
||||||
self.sessions.save(session)
|
await self._save_session(session)
|
||||||
|
|
||||||
visible = session.get_history(
|
visible = session.get_history(
|
||||||
max_messages=MIN_COMPACTED_REPLAY_MESSAGES,
|
max_messages=MIN_COMPACTED_REPLAY_MESSAGES,
|
||||||
|
|||||||
@@ -210,12 +210,14 @@ class AgentProgressHook(AgentHook):
|
|||||||
tool_hint=False,
|
tool_hint=False,
|
||||||
tool_events=tool_events,
|
tool_events=tool_events,
|
||||||
)
|
)
|
||||||
u = context.usage or {}
|
u = context.usage
|
||||||
logger.debug(
|
logger.debug(
|
||||||
"LLM usage: prompt={} completion={} cached={}",
|
"LLM usage: input={} output={} cache_read={} cache_write={} source={}",
|
||||||
u.get("prompt_tokens", 0),
|
u.input_tokens if u else 0,
|
||||||
u.get("completion_tokens", 0),
|
u.output_tokens if u else 0,
|
||||||
u.get("cached_tokens", 0),
|
u.cache_read_tokens if u else None,
|
||||||
|
u.cache_write_tokens if u else None,
|
||||||
|
u.source if u else "missing",
|
||||||
)
|
)
|
||||||
|
|
||||||
def finalize_content(self, context: AgentHookContext, content: str | None) -> str | None:
|
def finalize_content(self, context: AgentHookContext, content: str | None) -> str | None:
|
||||||
|
|||||||
+141
-102
@@ -6,7 +6,7 @@ import asyncio
|
|||||||
import inspect
|
import inspect
|
||||||
import os
|
import os
|
||||||
import time
|
import time
|
||||||
from collections.abc import Awaitable, Callable, Iterable
|
from collections.abc import Awaitable, Callable, Iterable, Sized
|
||||||
from copy import deepcopy
|
from copy import deepcopy
|
||||||
from dataclasses import dataclass, field
|
from dataclasses import dataclass, field
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
@@ -20,9 +20,16 @@ from nanobot.agent.context_governance import (
|
|||||||
)
|
)
|
||||||
from nanobot.agent.hook import AgentHook, AgentHookContext, AgentRunHookContext
|
from nanobot.agent.hook import AgentHook, AgentHookContext, AgentRunHookContext
|
||||||
from nanobot.agent.tools.registry import ToolRegistry, is_tool_error_result
|
from nanobot.agent.tools.registry import ToolRegistry, is_tool_error_result
|
||||||
|
from nanobot.llm_usage.context import (
|
||||||
|
LLMUsageSource,
|
||||||
|
bind_llm_usage_source,
|
||||||
|
reset_llm_usage_source,
|
||||||
|
source_from_session_key,
|
||||||
|
)
|
||||||
from nanobot.providers.base import (
|
from nanobot.providers.base import (
|
||||||
LLMProvider,
|
LLMProvider,
|
||||||
LLMResponse,
|
LLMResponse,
|
||||||
|
LLMUsage,
|
||||||
ProviderCallContext,
|
ProviderCallContext,
|
||||||
ProviderConversationState,
|
ProviderConversationState,
|
||||||
ToolCallRequest,
|
ToolCallRequest,
|
||||||
@@ -37,6 +44,7 @@ from nanobot.runtime_context import (
|
|||||||
reattach_runtime_context,
|
reattach_runtime_context,
|
||||||
)
|
)
|
||||||
from nanobot.session.history_visibility import is_hidden_history_message
|
from nanobot.session.history_visibility import is_hidden_history_message
|
||||||
|
from nanobot.session.recovery import PENDING_FOLLOWUP_ID_KEY
|
||||||
from nanobot.utils.helpers import (
|
from nanobot.utils.helpers import (
|
||||||
IncrementalThinkExtractor,
|
IncrementalThinkExtractor,
|
||||||
build_assistant_message,
|
build_assistant_message,
|
||||||
@@ -75,6 +83,22 @@ _MAX_EMPTY_RETRIES = 2
|
|||||||
_MAX_LENGTH_RECOVERIES = 3
|
_MAX_LENGTH_RECOVERIES = 3
|
||||||
_MAX_INJECTIONS_PER_TURN = 3
|
_MAX_INJECTIONS_PER_TURN = 3
|
||||||
_MAX_INJECTION_CYCLES = 5
|
_MAX_INJECTION_CYCLES = 5
|
||||||
|
_SLOW_TOOL_LOG_MS = 1_000
|
||||||
|
|
||||||
|
|
||||||
|
def _tool_input_scale(params: object) -> tuple[int, int]:
|
||||||
|
"""Return bounded structural counts without logging argument content."""
|
||||||
|
if not isinstance(params, dict):
|
||||||
|
return 0, len(params) if isinstance(params, str | bytes) else 0
|
||||||
|
params_dict = cast(dict[object, object], params)
|
||||||
|
items = len(params_dict)
|
||||||
|
chars = 0
|
||||||
|
for value in params_dict.values():
|
||||||
|
if isinstance(value, str | bytes):
|
||||||
|
chars += len(value)
|
||||||
|
elif isinstance(value, list | tuple | set | dict):
|
||||||
|
items += len(cast(Sized, value))
|
||||||
|
return items, chars
|
||||||
|
|
||||||
|
|
||||||
def _restore_outer_whitespace(content: str, original: str | None) -> str:
|
def _restore_outer_whitespace(content: str, original: str | None) -> str:
|
||||||
@@ -116,6 +140,7 @@ class AgentRunSpec:
|
|||||||
goal_continue_message: GoalContinueMessage | None = None
|
goal_continue_message: GoalContinueMessage | None = None
|
||||||
finalize_on_max_iterations: bool = True
|
finalize_on_max_iterations: bool = True
|
||||||
provider_state: ProviderConversationState | None = None
|
provider_state: ProviderConversationState | None = None
|
||||||
|
llm_usage_source: LLMUsageSource | None = None
|
||||||
|
|
||||||
|
|
||||||
@dataclass(slots=True)
|
@dataclass(slots=True)
|
||||||
@@ -125,7 +150,7 @@ class AgentRunResult:
|
|||||||
final_content: str | None
|
final_content: str | None
|
||||||
messages: list[dict[str, Any]]
|
messages: list[dict[str, Any]]
|
||||||
tools_used: list[str] = field(default_factory=list)
|
tools_used: list[str] = field(default_factory=list)
|
||||||
usage: dict[str, int] = field(default_factory=dict)
|
usage: LLMUsage | None = None
|
||||||
stop_reason: str = "completed"
|
stop_reason: str = "completed"
|
||||||
error: str | None = None
|
error: str | None = None
|
||||||
tool_events: list[dict[str, str]] = field(default_factory=list)
|
tool_events: list[dict[str, str]] = field(default_factory=list)
|
||||||
@@ -234,6 +259,23 @@ class AgentRunner:
|
|||||||
merged.get("content"),
|
merged.get("content"),
|
||||||
injection.get("content"),
|
injection.get("content"),
|
||||||
)
|
)
|
||||||
|
followup_id = injection.get(PENDING_FOLLOWUP_ID_KEY)
|
||||||
|
if isinstance(followup_id, str) and followup_id:
|
||||||
|
existing = cast(object, merged.get(PENDING_FOLLOWUP_ID_KEY))
|
||||||
|
followup_ids = (
|
||||||
|
[existing]
|
||||||
|
if isinstance(existing, str)
|
||||||
|
else [
|
||||||
|
item
|
||||||
|
for item in cast(list[object], existing)
|
||||||
|
if isinstance(item, str)
|
||||||
|
]
|
||||||
|
if isinstance(existing, list)
|
||||||
|
else []
|
||||||
|
)
|
||||||
|
if followup_id not in followup_ids:
|
||||||
|
followup_ids.append(followup_id)
|
||||||
|
merged[PENDING_FOLLOWUP_ID_KEY] = followup_ids
|
||||||
messages[-1] = merged
|
messages[-1] = merged
|
||||||
continue
|
continue
|
||||||
messages.append(injection)
|
messages.append(injection)
|
||||||
@@ -373,6 +415,9 @@ class AgentRunner:
|
|||||||
hook = spec.hook or AgentHook()
|
hook = spec.hook or AgentHook()
|
||||||
messages = list(spec.initial_messages)
|
messages = list(spec.initial_messages)
|
||||||
context = AgentRunHookContext(messages=deepcopy(messages))
|
context = AgentRunHookContext(messages=deepcopy(messages))
|
||||||
|
llm_usage_source_token = bind_llm_usage_source(
|
||||||
|
spec.llm_usage_source or source_from_session_key(spec.session_key)
|
||||||
|
)
|
||||||
|
|
||||||
try:
|
try:
|
||||||
await hook.before_run(context)
|
await hook.before_run(context)
|
||||||
@@ -394,7 +439,7 @@ class AgentRunner:
|
|||||||
context.messages = deepcopy(result.messages)
|
context.messages = deepcopy(result.messages)
|
||||||
context.final_content = result.final_content
|
context.final_content = result.final_content
|
||||||
context.tools_used = list(result.tools_used)
|
context.tools_used = list(result.tools_used)
|
||||||
context.usage = dict(result.usage)
|
context.usage = result.usage
|
||||||
context.stop_reason = result.stop_reason
|
context.stop_reason = result.stop_reason
|
||||||
context.error = result.error
|
context.error = result.error
|
||||||
context.tool_events = deepcopy(result.tool_events)
|
context.tool_events = deepcopy(result.tool_events)
|
||||||
@@ -405,17 +450,20 @@ class AgentRunner:
|
|||||||
await hook.after_run(context)
|
await hook.after_run(context)
|
||||||
return result
|
return result
|
||||||
finally:
|
finally:
|
||||||
context.messages = deepcopy(messages)
|
try:
|
||||||
if context.exception is None:
|
context.messages = deepcopy(messages)
|
||||||
await hook.on_finally(context)
|
if context.exception is None:
|
||||||
else:
|
|
||||||
try:
|
|
||||||
await hook.on_finally(context)
|
await hook.on_finally(context)
|
||||||
except Exception:
|
else:
|
||||||
logger.exception(
|
try:
|
||||||
"AgentHook.on_finally error after {}",
|
await hook.on_finally(context)
|
||||||
context.stop_reason or "run exception",
|
except Exception:
|
||||||
)
|
logger.exception(
|
||||||
|
"AgentHook.on_finally error after {}",
|
||||||
|
context.stop_reason or "run exception",
|
||||||
|
)
|
||||||
|
finally:
|
||||||
|
reset_llm_usage_source(llm_usage_source_token)
|
||||||
|
|
||||||
async def _run_core(
|
async def _run_core(
|
||||||
self,
|
self,
|
||||||
@@ -425,7 +473,7 @@ class AgentRunner:
|
|||||||
) -> AgentRunResult:
|
) -> AgentRunResult:
|
||||||
final_content: str | None = None
|
final_content: str | None = None
|
||||||
tools_used: list[str] = []
|
tools_used: list[str] = []
|
||||||
usage = {"prompt_tokens": 0, "completion_tokens": 0}
|
usage: LLMUsage | None = None
|
||||||
error: str | None = None
|
error: str | None = None
|
||||||
stop_reason = "completed"
|
stop_reason = "completed"
|
||||||
tool_events: list[dict[str, str]] = []
|
tool_events: list[dict[str, str]] = []
|
||||||
@@ -501,8 +549,8 @@ class AgentRunner:
|
|||||||
)
|
)
|
||||||
response.content = cleaned_content
|
response.content = cleaned_content
|
||||||
raw_usage = self._usage_or_estimate(spec, messages_for_model, response)
|
raw_usage = self._usage_or_estimate(spec, messages_for_model, response)
|
||||||
context.usage = dict(raw_usage)
|
context.usage = raw_usage
|
||||||
self._accumulate_usage(usage, raw_usage)
|
usage = self._merge_usage(usage, raw_usage)
|
||||||
if reasoning_text and not context.streamed_reasoning:
|
if reasoning_text and not context.streamed_reasoning:
|
||||||
await hook.emit_reasoning(reasoning_text)
|
await hook.emit_reasoning(reasoning_text)
|
||||||
await hook.emit_reasoning_end()
|
await hook.emit_reasoning_end()
|
||||||
@@ -665,10 +713,10 @@ class AgentRunner:
|
|||||||
conversation_state=conversation_state,
|
conversation_state=conversation_state,
|
||||||
)
|
)
|
||||||
retry_usage = self._usage_or_estimate(spec, retry_messages, response)
|
retry_usage = self._usage_or_estimate(spec, retry_messages, response)
|
||||||
self._accumulate_usage(usage, retry_usage)
|
usage = self._merge_usage(usage, retry_usage)
|
||||||
raw_usage = self._merge_usage(raw_usage, retry_usage)
|
raw_usage = self._merge_usage(raw_usage, retry_usage)
|
||||||
context.response = response
|
context.response = response
|
||||||
context.usage = dict(raw_usage)
|
context.usage = raw_usage
|
||||||
context.tool_calls = list(response.tool_calls)
|
context.tool_calls = list(response.tool_calls)
|
||||||
original_content = response.content
|
original_content = response.content
|
||||||
clean = hook.finalize_content(context, response.content)
|
clean = hook.finalize_content(context, response.content)
|
||||||
@@ -841,7 +889,7 @@ class AgentRunner:
|
|||||||
had_injections = True
|
had_injections = True
|
||||||
terminal_content = None
|
terminal_content = None
|
||||||
if spec.finalize_on_max_iterations:
|
if spec.finalize_on_max_iterations:
|
||||||
terminal_content = await self._try_finalize_after_max_iterations(
|
terminal_content, usage = await self._try_finalize_after_max_iterations(
|
||||||
spec,
|
spec,
|
||||||
hook,
|
hook,
|
||||||
messages,
|
messages,
|
||||||
@@ -904,18 +952,7 @@ class AgentRunner:
|
|||||||
conversation_state: ProviderConversationStateController,
|
conversation_state: ProviderConversationStateController,
|
||||||
provider_context: ProviderCallContext | None = None,
|
provider_context: ProviderCallContext | None = None,
|
||||||
) -> LLMResponse:
|
) -> LLMResponse:
|
||||||
timeout_s: float | None = spec.llm_timeout_s
|
timeout_s = self._resolve_llm_timeout_s(spec)
|
||||||
if timeout_s is None:
|
|
||||||
# Default to a finite timeout to avoid per-session lock starvation when an LLM
|
|
||||||
# request hangs indefinitely (e.g. gateway/network stall).
|
|
||||||
# Set NANOBOT_LLM_TIMEOUT_S=0 to disable.
|
|
||||||
raw = os.environ.get("NANOBOT_LLM_TIMEOUT_S", "300").strip()
|
|
||||||
try:
|
|
||||||
timeout_s = float(raw)
|
|
||||||
except (TypeError, ValueError):
|
|
||||||
timeout_s = 300.0
|
|
||||||
if timeout_s <= 0:
|
|
||||||
timeout_s = None
|
|
||||||
|
|
||||||
kwargs = self._build_request_kwargs(
|
kwargs = self._build_request_kwargs(
|
||||||
spec,
|
spec,
|
||||||
@@ -1229,9 +1266,9 @@ class AgentRunner:
|
|||||||
spec: AgentRunSpec,
|
spec: AgentRunSpec,
|
||||||
hook: AgentHook,
|
hook: AgentHook,
|
||||||
messages: list[dict[str, Any]],
|
messages: list[dict[str, Any]],
|
||||||
usage: dict[str, int],
|
usage: LLMUsage | None,
|
||||||
conversation_state: ProviderConversationStateController,
|
conversation_state: ProviderConversationStateController,
|
||||||
) -> str | None:
|
) -> tuple[str | None, LLMUsage | None]:
|
||||||
retry_messages = self._budget_exhausted_finalization_messages(messages)
|
retry_messages = self._budget_exhausted_finalization_messages(messages)
|
||||||
try:
|
try:
|
||||||
response = await self._request_no_tools(
|
response = await self._request_no_tools(
|
||||||
@@ -1246,10 +1283,10 @@ class AgentRunner:
|
|||||||
"Budget-exhausted finalization failed for {}; using fallback",
|
"Budget-exhausted finalization failed for {}; using fallback",
|
||||||
spec.session_key or "default",
|
spec.session_key or "default",
|
||||||
)
|
)
|
||||||
return None
|
return None, usage
|
||||||
|
|
||||||
raw_usage = self._usage_or_estimate(spec, retry_messages, response)
|
raw_usage = self._usage_or_estimate(spec, retry_messages, response)
|
||||||
self._accumulate_usage(usage, raw_usage)
|
usage = self._merge_usage(usage, raw_usage)
|
||||||
if response.finish_reason == "error" or response.has_tool_calls:
|
if response.finish_reason == "error" or response.has_tool_calls:
|
||||||
logger.warning(
|
logger.warning(
|
||||||
"Budget-exhausted finalization returned finish_reason='{}' "
|
"Budget-exhausted finalization returned finish_reason='{}' "
|
||||||
@@ -1258,19 +1295,19 @@ class AgentRunner:
|
|||||||
len(response.tool_calls),
|
len(response.tool_calls),
|
||||||
spec.session_key or "default",
|
spec.session_key or "default",
|
||||||
)
|
)
|
||||||
return None
|
return None, usage
|
||||||
|
|
||||||
context = AgentHookContext(
|
context = AgentHookContext(
|
||||||
iteration=spec.max_iterations,
|
iteration=spec.max_iterations,
|
||||||
messages=messages,
|
messages=messages,
|
||||||
response=response,
|
response=response,
|
||||||
usage=dict(raw_usage),
|
usage=raw_usage,
|
||||||
session_key=spec.session_key,
|
session_key=spec.session_key,
|
||||||
)
|
)
|
||||||
clean = hook.finalize_content(context, response.content)
|
clean = hook.finalize_content(context, response.content)
|
||||||
if is_blank_text(clean):
|
if is_blank_text(clean):
|
||||||
return None
|
return None, usage
|
||||||
return clean
|
return clean, usage
|
||||||
|
|
||||||
async def _request_no_tools(
|
async def _request_no_tools(
|
||||||
self,
|
self,
|
||||||
@@ -1284,10 +1321,38 @@ class AgentRunner:
|
|||||||
messages,
|
messages,
|
||||||
tools=None,
|
tools=None,
|
||||||
)
|
)
|
||||||
return await spec.runtime.provider.chat_with_retry(
|
coro = spec.runtime.provider.chat_with_retry(
|
||||||
**kwargs,
|
**kwargs,
|
||||||
provider_context=provider_context,
|
provider_context=provider_context,
|
||||||
)
|
)
|
||||||
|
timeout_s = self._resolve_llm_timeout_s(spec)
|
||||||
|
try:
|
||||||
|
return (
|
||||||
|
await coro
|
||||||
|
if timeout_s is None
|
||||||
|
else await asyncio.wait_for(coro, timeout=timeout_s)
|
||||||
|
)
|
||||||
|
except asyncio.TimeoutError:
|
||||||
|
return LLMResponse(
|
||||||
|
content=f"Error calling LLM: timed out after {timeout_s:g}s",
|
||||||
|
finish_reason="error",
|
||||||
|
error_kind="timeout",
|
||||||
|
)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _resolve_llm_timeout_s(spec: AgentRunSpec) -> float | None:
|
||||||
|
"""Resolve the wall-clock limit shared by every model request path."""
|
||||||
|
timeout_s = spec.llm_timeout_s
|
||||||
|
if timeout_s is None:
|
||||||
|
# Default to a finite timeout to avoid per-session lock starvation when an LLM
|
||||||
|
# request hangs indefinitely (e.g. gateway/network stall).
|
||||||
|
# Set NANOBOT_LLM_TIMEOUT_S=0 to disable.
|
||||||
|
raw = os.environ.get("NANOBOT_LLM_TIMEOUT_S", "300").strip()
|
||||||
|
try:
|
||||||
|
timeout_s = float(raw)
|
||||||
|
except (TypeError, ValueError):
|
||||||
|
timeout_s = 300.0
|
||||||
|
return timeout_s if timeout_s > 0 else None
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def _budget_exhausted_finalization_messages(
|
def _budget_exhausted_finalization_messages(
|
||||||
@@ -1314,31 +1379,24 @@ class AgentRunner:
|
|||||||
spec: AgentRunSpec,
|
spec: AgentRunSpec,
|
||||||
messages: list[dict[str, Any]],
|
messages: list[dict[str, Any]],
|
||||||
response: LLMResponse,
|
response: LLMResponse,
|
||||||
) -> dict[str, int]:
|
) -> LLMUsage | None:
|
||||||
usage = self._usage_dict(response.usage)
|
usage = response.usage
|
||||||
total = self._usage_total(usage)
|
if response.finish_reason == "error":
|
||||||
if total > 0:
|
if usage is None or usage.total_tokens == 0:
|
||||||
usage["total_tokens"] = total
|
usage = LLMUsage.empty_request()
|
||||||
usage.setdefault("provider_tokens", total)
|
elif usage is None or usage.total_tokens == 0:
|
||||||
elif response.finish_reason == "error":
|
|
||||||
return {}
|
|
||||||
else:
|
|
||||||
usage = self._estimate_response_usage(spec, messages, response)
|
usage = self._estimate_response_usage(spec, messages, response)
|
||||||
completion = usage.get("completion_tokens", 0)
|
return usage.with_timing(
|
||||||
if response.generation_ms is not None and completion > 0:
|
generation_ms=response.generation_ms,
|
||||||
usage["generation_ms"] = response.generation_ms
|
ttft_ms=response.ttft_ms,
|
||||||
usage["measured_completion_tokens"] = completion
|
)
|
||||||
if response.ttft_ms is not None:
|
|
||||||
usage["ttft_ms"] = response.ttft_ms
|
|
||||||
usage["timed_requests"] = 1
|
|
||||||
return usage
|
|
||||||
|
|
||||||
def _estimate_response_usage(
|
def _estimate_response_usage(
|
||||||
self,
|
self,
|
||||||
spec: AgentRunSpec,
|
spec: AgentRunSpec,
|
||||||
messages: list[dict[str, Any]],
|
messages: list[dict[str, Any]],
|
||||||
response: LLMResponse,
|
response: LLMResponse,
|
||||||
) -> dict[str, int]:
|
) -> LLMUsage:
|
||||||
try:
|
try:
|
||||||
tools = spec.tools.get_definitions()
|
tools = spec.tools.get_definitions()
|
||||||
except Exception:
|
except Exception:
|
||||||
@@ -1356,52 +1414,21 @@ class AgentRunner:
|
|||||||
thinking_blocks=response.thinking_blocks,
|
thinking_blocks=response.thinking_blocks,
|
||||||
)
|
)
|
||||||
completion_tokens = estimate_message_tokens(assistant_message)
|
completion_tokens = estimate_message_tokens(assistant_message)
|
||||||
total_tokens = max(0, prompt_tokens) + max(0, completion_tokens)
|
return LLMUsage.estimated(
|
||||||
if total_tokens <= 0:
|
input_tokens=max(0, prompt_tokens),
|
||||||
return {}
|
output_tokens=max(0, completion_tokens),
|
||||||
return {
|
)
|
||||||
"prompt_tokens": max(0, prompt_tokens),
|
|
||||||
"completion_tokens": max(0, completion_tokens),
|
|
||||||
"total_tokens": total_tokens,
|
|
||||||
"estimated_tokens": total_tokens,
|
|
||||||
}
|
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def _usage_dict(usage: dict[str, Any] | None) -> dict[str, int]:
|
def _merge_usage(
|
||||||
if not usage:
|
left: LLMUsage | None,
|
||||||
return {}
|
right: LLMUsage | None,
|
||||||
result: dict[str, int] = {}
|
) -> LLMUsage | None:
|
||||||
for key, value in usage.items():
|
if left is None:
|
||||||
try:
|
return right
|
||||||
result[key] = int(value or 0)
|
if right is None:
|
||||||
except (TypeError, ValueError):
|
return left
|
||||||
continue
|
return left + right
|
||||||
return result
|
|
||||||
|
|
||||||
@staticmethod
|
|
||||||
def _usage_total(usage: dict[str, int]) -> int:
|
|
||||||
return max(0, usage.get("total_tokens", 0) or (
|
|
||||||
usage.get("prompt_tokens", 0) + usage.get("completion_tokens", 0)
|
|
||||||
))
|
|
||||||
|
|
||||||
@staticmethod
|
|
||||||
def _merge_usage(left: dict[str, int], right: dict[str, int]) -> dict[str, int]:
|
|
||||||
merged = dict(left)
|
|
||||||
for key, value in right.items():
|
|
||||||
merged[key] = merged.get(key, 0) + value
|
|
||||||
return merged
|
|
||||||
|
|
||||||
@staticmethod
|
|
||||||
def _accumulate_usage(total: dict[str, int], request: dict[str, int]) -> None:
|
|
||||||
"""Fold one model request into the current turn's usage."""
|
|
||||||
total["request_count"] = total.get("request_count", 0) + 1
|
|
||||||
prompt_tokens = request.get("prompt_tokens")
|
|
||||||
if prompt_tokens is not None and prompt_tokens >= 0:
|
|
||||||
total["context_tokens"] = prompt_tokens
|
|
||||||
for key, value in request.items():
|
|
||||||
if key in {"context_tokens", "request_count"} or value < 0:
|
|
||||||
continue
|
|
||||||
total[key] = total.get(key, 0) + value
|
|
||||||
|
|
||||||
async def _execute_tools(
|
async def _execute_tools(
|
||||||
self,
|
self,
|
||||||
@@ -1510,6 +1537,7 @@ class AgentRunner:
|
|||||||
RuntimeError(prep_error) if spec.fail_on_tool_error else None
|
RuntimeError(prep_error) if spec.fail_on_tool_error else None
|
||||||
)
|
)
|
||||||
await hook.before_execute_tool(context, tool_call, tool, params)
|
await hook.before_execute_tool(context, tool_call, tool, params)
|
||||||
|
tool_started_at = time.perf_counter()
|
||||||
try:
|
try:
|
||||||
if tool is not None:
|
if tool is not None:
|
||||||
result = await tool.execute(**params)
|
result = await tool.execute(**params)
|
||||||
@@ -1538,6 +1566,17 @@ class AgentRunner:
|
|||||||
if spec.fail_on_tool_error:
|
if spec.fail_on_tool_error:
|
||||||
return payload, event, exc
|
return payload, event, exc
|
||||||
return payload, event, None
|
return payload, event, None
|
||||||
|
finally:
|
||||||
|
duration_ms = int((time.perf_counter() - tool_started_at) * 1000)
|
||||||
|
if duration_ms >= _SLOW_TOOL_LOG_MS:
|
||||||
|
input_items, input_chars = _tool_input_scale(params)
|
||||||
|
logger.warning(
|
||||||
|
"slow tool operation={} input_items={} input_chars={} duration_ms={}",
|
||||||
|
tool_call.name,
|
||||||
|
input_items,
|
||||||
|
input_chars,
|
||||||
|
duration_ms,
|
||||||
|
)
|
||||||
|
|
||||||
if is_tool_error_result(result):
|
if is_tool_error_result(result):
|
||||||
await hook.on_execute_tool_error(context, tool_call, tool, params, result)
|
await hook.on_execute_tool_error(context, tool_call, tool, params, result)
|
||||||
|
|||||||
@@ -8,7 +8,7 @@ import warnings
|
|||||||
from collections.abc import Mapping
|
from collections.abc import Mapping
|
||||||
from dataclasses import dataclass, field
|
from dataclasses import dataclass, field
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import Any, Callable, TypedDict
|
from typing import Any, Callable, NotRequired, TypedDict
|
||||||
|
|
||||||
from loguru import logger
|
from loguru import logger
|
||||||
|
|
||||||
@@ -28,7 +28,8 @@ from nanobot.agent.tools.registry import ToolRegistry
|
|||||||
from nanobot.bus.events import InboundMessage
|
from nanobot.bus.events import InboundMessage
|
||||||
from nanobot.bus.queue import MessageBus
|
from nanobot.bus.queue import MessageBus
|
||||||
from nanobot.config.schema import AgentDefaults, ToolsConfig
|
from nanobot.config.schema import AgentDefaults, ToolsConfig
|
||||||
from nanobot.providers.base import LLMProvider
|
from nanobot.llm_usage.context import LLMUsageSource, current_llm_usage_source
|
||||||
|
from nanobot.providers.base import LLMProvider, LLMUsage
|
||||||
from nanobot.security.workspace_access import (
|
from nanobot.security.workspace_access import (
|
||||||
WorkspaceScope,
|
WorkspaceScope,
|
||||||
bind_workspace_scope,
|
bind_workspace_scope,
|
||||||
@@ -43,6 +44,7 @@ class _SubagentOrigin(TypedDict):
|
|||||||
channel: str
|
channel: str
|
||||||
chat_id: str
|
chat_id: str
|
||||||
session_key: str | None
|
session_key: str | None
|
||||||
|
llm_usage_source: NotRequired[LLMUsageSource]
|
||||||
|
|
||||||
|
|
||||||
@dataclass(slots=True)
|
@dataclass(slots=True)
|
||||||
@@ -56,7 +58,7 @@ class SubagentStatus:
|
|||||||
phase: str = "initializing" # initializing | awaiting_tools | tools_completed | final_response | done | error
|
phase: str = "initializing" # initializing | awaiting_tools | tools_completed | final_response | done | error
|
||||||
iteration: int = 0
|
iteration: int = 0
|
||||||
tool_events: list[dict[str, str]] = field(default_factory=list)
|
tool_events: list[dict[str, str]] = field(default_factory=list)
|
||||||
usage: dict[str, int] = field(default_factory=dict)
|
usage: LLMUsage | None = None
|
||||||
stop_reason: str | None = None
|
stop_reason: str | None = None
|
||||||
error: str | None = None
|
error: str | None = None
|
||||||
|
|
||||||
@@ -82,7 +84,7 @@ class _SubagentHook(AgentHook):
|
|||||||
return
|
return
|
||||||
self._status.iteration = context.iteration
|
self._status.iteration = context.iteration
|
||||||
self._status.tool_events = list(context.tool_events)
|
self._status.tool_events = list(context.tool_events)
|
||||||
self._status.usage = dict(context.usage)
|
self._status.usage = context.usage
|
||||||
if context.error:
|
if context.error:
|
||||||
self._status.error = str(context.error)
|
self._status.error = str(context.error)
|
||||||
|
|
||||||
@@ -252,6 +254,7 @@ class SubagentManager:
|
|||||||
"channel": origin_channel,
|
"channel": origin_channel,
|
||||||
"chat_id": origin_chat_id,
|
"chat_id": origin_chat_id,
|
||||||
"session_key": session_key,
|
"session_key": session_key,
|
||||||
|
"llm_usage_source": current_llm_usage_source(),
|
||||||
}
|
}
|
||||||
|
|
||||||
status = SubagentStatus(
|
status = SubagentStatus(
|
||||||
@@ -315,6 +318,7 @@ class SubagentManager:
|
|||||||
"channel": origin_channel,
|
"channel": origin_channel,
|
||||||
"chat_id": origin_chat_id,
|
"chat_id": origin_chat_id,
|
||||||
"session_key": session_key,
|
"session_key": session_key,
|
||||||
|
"llm_usage_source": current_llm_usage_source(),
|
||||||
}
|
}
|
||||||
status = SubagentStatus(
|
status = SubagentStatus(
|
||||||
task_id=task_id,
|
task_id=task_id,
|
||||||
@@ -417,6 +421,10 @@ class SubagentManager:
|
|||||||
session_key=sess_key,
|
session_key=sess_key,
|
||||||
workspace=root,
|
workspace=root,
|
||||||
llm_timeout_s=llm_timeout,
|
llm_timeout_s=llm_timeout,
|
||||||
|
llm_usage_source=origin.get(
|
||||||
|
"llm_usage_source",
|
||||||
|
current_llm_usage_source(),
|
||||||
|
),
|
||||||
))
|
))
|
||||||
finally:
|
finally:
|
||||||
if token is not None:
|
if token is not None:
|
||||||
|
|||||||
@@ -4,7 +4,10 @@
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import inspect
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
from typing import TypedDict
|
||||||
|
|
||||||
from pydantic import Field
|
from pydantic import Field
|
||||||
|
|
||||||
@@ -24,6 +27,14 @@ from nanobot.runtime_context import RuntimeContextBlock, wrap_runtime_context_li
|
|||||||
from nanobot.security.workspace_access import current_tool_workspace
|
from nanobot.security.workspace_access import current_tool_workspace
|
||||||
|
|
||||||
|
|
||||||
|
class _CliAppRunKwargs(TypedDict):
|
||||||
|
args: list[str]
|
||||||
|
json_output: bool
|
||||||
|
working_dir: str | None
|
||||||
|
timeout: int | None
|
||||||
|
restrict_to_workspace: bool
|
||||||
|
|
||||||
|
|
||||||
class CliAppsToolConfig(Base):
|
class CliAppsToolConfig(Base):
|
||||||
"""CLI Apps tool configuration."""
|
"""CLI Apps tool configuration."""
|
||||||
|
|
||||||
@@ -147,14 +158,17 @@ class CliAppsTool(Tool):
|
|||||||
)
|
)
|
||||||
workspace = access.project_path or self.workspace
|
workspace = access.project_path or self.workspace
|
||||||
manager = CliAppManager(workspace=workspace, runtime=self.runtime)
|
manager = CliAppManager(workspace=workspace, runtime=self.runtime)
|
||||||
|
run_kwargs: _CliAppRunKwargs = {
|
||||||
|
"args": args or [],
|
||||||
|
"json_output": bool(json),
|
||||||
|
"working_dir": working_dir,
|
||||||
|
"timeout": timeout,
|
||||||
|
"restrict_to_workspace": access.restrict_to_workspace,
|
||||||
|
}
|
||||||
try:
|
try:
|
||||||
return manager.run(
|
run_async = inspect.getattr_static(type(manager), "run_async", None)
|
||||||
name,
|
if inspect.iscoroutinefunction(run_async):
|
||||||
args=args or [],
|
return await manager.run_async(name, **run_kwargs)
|
||||||
json_output=bool(json),
|
return await asyncio.to_thread(manager.run, name, **run_kwargs)
|
||||||
working_dir=working_dir,
|
|
||||||
timeout=timeout,
|
|
||||||
restrict_to_workspace=access.restrict_to_workspace,
|
|
||||||
)
|
|
||||||
except CliAppError as exc:
|
except CliAppError as exc:
|
||||||
return ToolResult.error(f"Error: {exc.message}")
|
return ToolResult.error(f"Error: {exc.message}")
|
||||||
|
|||||||
@@ -143,14 +143,41 @@ class CronTool(Tool):
|
|||||||
tz: str | None = None,
|
tz: str | None = None,
|
||||||
at: str | None = None,
|
at: str | None = None,
|
||||||
job_id: str | None = None,
|
job_id: str | None = None,
|
||||||
|
) -> str:
|
||||||
|
if action == "add" and self._in_cron_context.get():
|
||||||
|
return ToolResult.error(
|
||||||
|
"Error: cannot schedule new jobs from within a cron job execution"
|
||||||
|
)
|
||||||
|
return await self._cron.run_sync(
|
||||||
|
self._execute_sync,
|
||||||
|
action,
|
||||||
|
name,
|
||||||
|
message,
|
||||||
|
every_seconds,
|
||||||
|
cron_expr,
|
||||||
|
tz,
|
||||||
|
at,
|
||||||
|
job_id,
|
||||||
|
)
|
||||||
|
|
||||||
|
def _execute_sync(
|
||||||
|
self,
|
||||||
|
action: str,
|
||||||
|
name: str | None,
|
||||||
|
message: str,
|
||||||
|
every_seconds: int | None,
|
||||||
|
cron_expr: str | None,
|
||||||
|
tz: str | None,
|
||||||
|
at: str | None,
|
||||||
|
job_id: str | None,
|
||||||
) -> str:
|
) -> str:
|
||||||
if action == "add":
|
if action == "add":
|
||||||
if self._in_cron_context.get():
|
if self._in_cron_context.get():
|
||||||
return ToolResult.error("Error: cannot schedule new jobs from within a cron job execution")
|
return ToolResult.error("Error: cannot schedule new jobs from within a cron job execution")
|
||||||
return self._add_job(name, message, every_seconds, cron_expr, tz, at)
|
return self._add_job(name, message, every_seconds, cron_expr, tz, at)
|
||||||
elif action == "list":
|
if action == "list":
|
||||||
return self._list_jobs()
|
return self._list_jobs()
|
||||||
elif action == "remove":
|
if action == "remove":
|
||||||
return self._remove_job(job_id)
|
return self._remove_job(job_id)
|
||||||
return f"Unknown action: {action}"
|
return f"Unknown action: {action}"
|
||||||
|
|
||||||
|
|||||||
@@ -2,9 +2,11 @@
|
|||||||
|
|
||||||
# pyright: reportPrivateUsage=false, reportUnusedFunction=false
|
# pyright: reportPrivateUsage=false, reportUnusedFunction=false
|
||||||
|
|
||||||
|
import asyncio
|
||||||
import difflib
|
import difflib
|
||||||
import mimetypes
|
import mimetypes
|
||||||
import os
|
import os
|
||||||
|
import threading
|
||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import Any
|
from typing import Any
|
||||||
@@ -21,6 +23,7 @@ from nanobot.agent.tools.schema import (
|
|||||||
)
|
)
|
||||||
from nanobot.config_base import Base
|
from nanobot.config_base import Base
|
||||||
from nanobot.security.workspace_access import current_tool_workspace
|
from nanobot.security.workspace_access import current_tool_workspace
|
||||||
|
from nanobot.utils.cancellation import shield_and_drain
|
||||||
from nanobot.utils.helpers import build_image_content_blocks, detect_image_mime
|
from nanobot.utils.helpers import build_image_content_blocks, detect_image_mime
|
||||||
|
|
||||||
|
|
||||||
@@ -664,22 +667,31 @@ def _match_covers_line(match: _MatchSpan, line: int) -> bool:
|
|||||||
return match.line <= line <= _match_end_line(match)
|
return match.line <= line <= _match_end_line(match)
|
||||||
|
|
||||||
|
|
||||||
def _find_exact_matches(content: str, old_text: str) -> list[_MatchSpan]:
|
def _find_exact_matches(
|
||||||
|
content: str,
|
||||||
|
old_text: str,
|
||||||
|
*,
|
||||||
|
max_matches: int | None = None,
|
||||||
|
) -> list[_MatchSpan]:
|
||||||
matches: list[_MatchSpan] = []
|
matches: list[_MatchSpan] = []
|
||||||
start = 0
|
search_start = 0
|
||||||
while True:
|
line_start = 0
|
||||||
idx = content.find(old_text, start)
|
line = 1
|
||||||
|
while max_matches is None or len(matches) < max_matches:
|
||||||
|
idx = content.find(old_text, search_start)
|
||||||
if idx == -1:
|
if idx == -1:
|
||||||
break
|
break
|
||||||
|
line += content.count("\n", line_start, idx)
|
||||||
matches.append(
|
matches.append(
|
||||||
_MatchSpan(
|
_MatchSpan(
|
||||||
start=idx,
|
start=idx,
|
||||||
end=idx + len(old_text),
|
end=idx + len(old_text),
|
||||||
text=content[idx : idx + len(old_text)],
|
text=content[idx : idx + len(old_text)],
|
||||||
line=content.count("\n", 0, idx) + 1,
|
line=line,
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
start = idx + max(1, len(old_text))
|
line_start = idx
|
||||||
|
search_start = idx + max(1, len(old_text))
|
||||||
return matches
|
return matches
|
||||||
|
|
||||||
|
|
||||||
@@ -735,27 +747,36 @@ def _find_quote_matches(content: str, old_text: str) -> list[_MatchSpan]:
|
|||||||
norm_content = _normalize_quotes(content)
|
norm_content = _normalize_quotes(content)
|
||||||
norm_old = _normalize_quotes(old_text)
|
norm_old = _normalize_quotes(old_text)
|
||||||
matches: list[_MatchSpan] = []
|
matches: list[_MatchSpan] = []
|
||||||
start = 0
|
search_start = 0
|
||||||
|
line_start = 0
|
||||||
|
line = 1
|
||||||
while True:
|
while True:
|
||||||
idx = norm_content.find(norm_old, start)
|
idx = norm_content.find(norm_old, search_start)
|
||||||
if idx == -1:
|
if idx == -1:
|
||||||
break
|
break
|
||||||
|
line += content.count("\n", line_start, idx)
|
||||||
matches.append(
|
matches.append(
|
||||||
_MatchSpan(
|
_MatchSpan(
|
||||||
start=idx,
|
start=idx,
|
||||||
end=idx + len(old_text),
|
end=idx + len(old_text),
|
||||||
text=content[idx : idx + len(old_text)],
|
text=content[idx : idx + len(old_text)],
|
||||||
line=content.count("\n", 0, idx) + 1,
|
line=line,
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
start = idx + max(1, len(norm_old))
|
line_start = idx
|
||||||
|
search_start = idx + max(1, len(norm_old))
|
||||||
return matches
|
return matches
|
||||||
|
|
||||||
|
|
||||||
def _find_matches(content: str, old_text: str) -> list[_MatchSpan]:
|
def _find_matches(
|
||||||
"""Locate all matches using progressively looser strategies."""
|
content: str,
|
||||||
|
old_text: str,
|
||||||
|
*,
|
||||||
|
max_exact_matches: int | None = None,
|
||||||
|
) -> list[_MatchSpan]:
|
||||||
|
"""Locate matches using progressively looser strategies."""
|
||||||
for matcher in (
|
for matcher in (
|
||||||
lambda: _find_exact_matches(content, old_text),
|
lambda: _find_exact_matches(content, old_text, max_matches=max_exact_matches),
|
||||||
lambda: _find_trim_matches(content, old_text),
|
lambda: _find_trim_matches(content, old_text),
|
||||||
lambda: _find_trim_matches(content, old_text, normalize_quotes=True),
|
lambda: _find_trim_matches(content, old_text, normalize_quotes=True),
|
||||||
lambda: _find_quote_matches(content, old_text),
|
lambda: _find_quote_matches(content, old_text),
|
||||||
@@ -869,6 +890,43 @@ class EditFileTool(_FsTool):
|
|||||||
new_text: str | None = None,
|
new_text: str | None = None,
|
||||||
replace_all: bool = False, occurrence: int | None = None,
|
replace_all: bool = False, occurrence: int | None = None,
|
||||||
line_hint: int | None = None, expected_replacements: int | None = None, **kwargs: Any,
|
line_hint: int | None = None, expected_replacements: int | None = None, **kwargs: Any,
|
||||||
|
) -> str:
|
||||||
|
cancelled = threading.Event()
|
||||||
|
commit_lock = threading.Lock()
|
||||||
|
try:
|
||||||
|
return await asyncio.to_thread(
|
||||||
|
self._execute_sync,
|
||||||
|
path=path,
|
||||||
|
old_text=old_text,
|
||||||
|
new_text=new_text,
|
||||||
|
replace_all=replace_all,
|
||||||
|
occurrence=occurrence,
|
||||||
|
line_hint=line_hint,
|
||||||
|
expected_replacements=expected_replacements,
|
||||||
|
cancelled=cancelled,
|
||||||
|
commit_lock=commit_lock,
|
||||||
|
)
|
||||||
|
except asyncio.CancelledError:
|
||||||
|
cancelled.set()
|
||||||
|
# If a commit already started, do not report cancellation until the
|
||||||
|
# file bytes and FileStates record are settled. Otherwise, taking
|
||||||
|
# the lock first guarantees the worker observes ``cancelled`` before
|
||||||
|
# it can mutate the target.
|
||||||
|
await shield_and_drain(
|
||||||
|
asyncio.to_thread(self._wait_for_commit, commit_lock)
|
||||||
|
)
|
||||||
|
raise
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _wait_for_commit(commit_lock: threading.Lock) -> None:
|
||||||
|
with commit_lock:
|
||||||
|
pass
|
||||||
|
|
||||||
|
def _execute_sync(
|
||||||
|
self, *, path: str | None, old_text: str | None,
|
||||||
|
new_text: str | None, replace_all: bool, occurrence: int | None,
|
||||||
|
line_hint: int | None, expected_replacements: int | None,
|
||||||
|
cancelled: threading.Event, commit_lock: threading.Lock,
|
||||||
) -> str:
|
) -> str:
|
||||||
try:
|
try:
|
||||||
if not path:
|
if not path:
|
||||||
@@ -892,9 +950,12 @@ class EditFileTool(_FsTool):
|
|||||||
# Create-file semantics: old_text='' + file doesn't exist → create
|
# Create-file semantics: old_text='' + file doesn't exist → create
|
||||||
if not file_exists:
|
if not file_exists:
|
||||||
if old_text == "":
|
if old_text == "":
|
||||||
fp.parent.mkdir(parents=True, exist_ok=True)
|
with commit_lock:
|
||||||
fp.write_text(new_text, encoding="utf-8")
|
if cancelled.is_set():
|
||||||
self._file_states.record_write(fp)
|
return ToolResult.error("Error: edit_file cancelled.")
|
||||||
|
fp.parent.mkdir(parents=True, exist_ok=True)
|
||||||
|
fp.write_text(new_text, encoding="utf-8")
|
||||||
|
self._file_states.record_write(fp)
|
||||||
return f"Successfully created {fp}"
|
return f"Successfully created {fp}"
|
||||||
return self._file_not_found_msg(path, fp)
|
return self._file_not_found_msg(path, fp)
|
||||||
|
|
||||||
@@ -912,8 +973,11 @@ class EditFileTool(_FsTool):
|
|||||||
content = raw.decode("utf-8")
|
content = raw.decode("utf-8")
|
||||||
if content.strip():
|
if content.strip():
|
||||||
return ToolResult.error(f"Error: Cannot create file — {path} already exists and is not empty.")
|
return ToolResult.error(f"Error: Cannot create file — {path} already exists and is not empty.")
|
||||||
fp.write_text(new_text, encoding="utf-8")
|
with commit_lock:
|
||||||
self._file_states.record_write(fp)
|
if cancelled.is_set():
|
||||||
|
return ToolResult.error("Error: edit_file cancelled.")
|
||||||
|
fp.write_text(new_text, encoding="utf-8")
|
||||||
|
self._file_states.record_write(fp)
|
||||||
return f"Successfully edited {fp}"
|
return f"Successfully edited {fp}"
|
||||||
|
|
||||||
# Read-before-edit check
|
# Read-before-edit check
|
||||||
@@ -923,7 +987,11 @@ class EditFileTool(_FsTool):
|
|||||||
uses_crlf = b"\r\n" in raw
|
uses_crlf = b"\r\n" in raw
|
||||||
content = raw.decode("utf-8").replace("\r\n", "\n")
|
content = raw.decode("utf-8").replace("\r\n", "\n")
|
||||||
norm_old = old_text.replace("\r\n", "\n")
|
norm_old = old_text.replace("\r\n", "\n")
|
||||||
matches = _find_matches(content, norm_old)
|
matches = _find_matches(
|
||||||
|
content,
|
||||||
|
norm_old,
|
||||||
|
max_exact_matches=occurrence,
|
||||||
|
)
|
||||||
|
|
||||||
if not matches:
|
if not matches:
|
||||||
return self._not_found_msg(old_text, content, path)
|
return self._not_found_msg(old_text, content, path)
|
||||||
@@ -1000,8 +1068,11 @@ class EditFileTool(_FsTool):
|
|||||||
if uses_crlf:
|
if uses_crlf:
|
||||||
new_content = new_content.replace("\n", "\r\n")
|
new_content = new_content.replace("\n", "\r\n")
|
||||||
|
|
||||||
fp.write_bytes(new_content.encode("utf-8"))
|
with commit_lock:
|
||||||
self._file_states.record_write(fp)
|
if cancelled.is_set():
|
||||||
|
return ToolResult.error("Error: edit_file cancelled.")
|
||||||
|
fp.write_bytes(new_content.encode("utf-8"))
|
||||||
|
self._file_states.record_write(fp)
|
||||||
msg = f"Successfully edited {fp}"
|
msg = f"Successfully edited {fp}"
|
||||||
if warning:
|
if warning:
|
||||||
msg = f"{warning}\n{msg}"
|
msg = f"{warning}\n{msg}"
|
||||||
|
|||||||
@@ -17,6 +17,7 @@ from nanobot.agent.tools.context import RequestContext, ToolContext, current_req
|
|||||||
from nanobot.agent.tools.schema import StringSchema, tool_parameters_schema
|
from nanobot.agent.tools.schema import StringSchema, tool_parameters_schema
|
||||||
from nanobot.bus.runtime_events import GoalStateChanged, RuntimeEventBus, RuntimeEventContext
|
from nanobot.bus.runtime_events import GoalStateChanged, RuntimeEventBus, RuntimeEventContext
|
||||||
from nanobot.runtime_context import RuntimeContextBlock, wrap_runtime_context_lines
|
from nanobot.runtime_context import RuntimeContextBlock, wrap_runtime_context_lines
|
||||||
|
from nanobot.session.async_compat import call_session_manager
|
||||||
from nanobot.session.goal_state import (
|
from nanobot.session.goal_state import (
|
||||||
GOAL_STATE_KEY,
|
GOAL_STATE_KEY,
|
||||||
MAX_GOAL_OBJECTIVE_CHARS,
|
MAX_GOAL_OBJECTIVE_CHARS,
|
||||||
@@ -28,6 +29,7 @@ from nanobot.session.goal_state import (
|
|||||||
sustained_goal_active,
|
sustained_goal_active,
|
||||||
)
|
)
|
||||||
from nanobot.session.turn_continuation import reset_goal_continuation_rounds
|
from nanobot.session.turn_continuation import reset_goal_continuation_rounds
|
||||||
|
from nanobot.utils.cancellation import shield_and_drain
|
||||||
from nanobot.utils.prompt_templates import render_template
|
from nanobot.utils.prompt_templates import render_template
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
@@ -60,36 +62,68 @@ class _GoalToolsMixin:
|
|||||||
self._sessions = sessions
|
self._sessions = sessions
|
||||||
self._runtime_events = runtime_events
|
self._runtime_events = runtime_events
|
||||||
|
|
||||||
def _session(self):
|
async def _get_or_create_session(self, key: str):
|
||||||
|
return await call_session_manager(
|
||||||
|
self._sessions,
|
||||||
|
"get_or_create_async",
|
||||||
|
self._sessions.get_or_create,
|
||||||
|
key,
|
||||||
|
)
|
||||||
|
|
||||||
|
async def _save_session(self, session: Any) -> None:
|
||||||
|
await call_session_manager(
|
||||||
|
self._sessions,
|
||||||
|
"save_async",
|
||||||
|
self._sessions.save,
|
||||||
|
session,
|
||||||
|
)
|
||||||
|
|
||||||
|
async def _session(self):
|
||||||
request_ctx = current_request_context()
|
request_ctx = current_request_context()
|
||||||
if request_ctx is None:
|
if request_ctx is None:
|
||||||
return None
|
return None
|
||||||
key = request_ctx.session_key
|
key = request_ctx.session_key
|
||||||
if not key:
|
if not key:
|
||||||
return None
|
return None
|
||||||
return self._sessions.get_or_create(key)
|
return await self._get_or_create_session(key)
|
||||||
|
|
||||||
def _goal_mutation_allowed(self) -> bool:
|
def _goal_mutation_allowed(self) -> bool:
|
||||||
return current_request_context() is not None and goal_mutation_allowed()
|
return current_request_context() is not None and goal_mutation_allowed()
|
||||||
|
|
||||||
def _save_goal_state(
|
async def _save_goal_state(
|
||||||
self,
|
self,
|
||||||
sess: Any,
|
sess: Any,
|
||||||
blob: dict[str, Any],
|
blob: dict[str, Any],
|
||||||
*,
|
*,
|
||||||
reset_continuation: bool = False,
|
reset_continuation: bool = False,
|
||||||
|
revoke_permission: bool = False,
|
||||||
) -> None:
|
) -> None:
|
||||||
previous_metadata = deepcopy(sess.metadata)
|
previous_metadata = deepcopy(sess.metadata)
|
||||||
sess.metadata[GOAL_STATE_KEY] = blob
|
saved = False
|
||||||
discard_legacy_goal_state_key(sess.metadata)
|
|
||||||
if reset_continuation:
|
async def save_and_publish() -> None:
|
||||||
reset_goal_continuation_rounds(sess.metadata)
|
nonlocal saved
|
||||||
|
sess.metadata[GOAL_STATE_KEY] = blob
|
||||||
|
discard_legacy_goal_state_key(sess.metadata)
|
||||||
|
if reset_continuation:
|
||||||
|
reset_goal_continuation_rounds(sess.metadata)
|
||||||
|
try:
|
||||||
|
await self._save_session(sess)
|
||||||
|
except BaseException:
|
||||||
|
sess.metadata.clear()
|
||||||
|
sess.metadata.update(previous_metadata)
|
||||||
|
raise
|
||||||
|
saved = True
|
||||||
|
await self._publish_goal_state_changed(sess.metadata)
|
||||||
|
|
||||||
try:
|
try:
|
||||||
self._sessions.save(sess)
|
await shield_and_drain(save_and_publish())
|
||||||
except BaseException:
|
finally:
|
||||||
sess.metadata.clear()
|
# This ContextVar belongs to the caller task, not the settlement task.
|
||||||
sess.metadata.update(previous_metadata)
|
# Apply the post-save permission effect here even when cancellation was
|
||||||
raise
|
# delayed until the durable save and runtime notification completed.
|
||||||
|
if revoke_permission and saved:
|
||||||
|
revoke_goal_mutation_permission()
|
||||||
|
|
||||||
async def _publish_goal_state_changed(self, metadata: dict[str, Any]) -> None:
|
async def _publish_goal_state_changed(self, metadata: dict[str, Any]) -> None:
|
||||||
runtime_events = self._runtime_events
|
runtime_events = self._runtime_events
|
||||||
@@ -175,7 +209,7 @@ class CreateGoalTool(Tool, _GoalToolsMixin):
|
|||||||
) -> RuntimeContextBlock | None:
|
) -> RuntimeContextBlock | None:
|
||||||
if not request.session_key:
|
if not request.session_key:
|
||||||
return None
|
return None
|
||||||
session = self._sessions.get_or_create(request.session_key)
|
session = await self._get_or_create_session(request.session_key)
|
||||||
goal_start_requested = explicit_goal_requested(request.metadata)
|
goal_start_requested = explicit_goal_requested(request.metadata)
|
||||||
goal_active = sustained_goal_active(session.metadata)
|
goal_active = sustained_goal_active(session.metadata)
|
||||||
if not goal_start_requested and not goal_active:
|
if not goal_start_requested and not goal_active:
|
||||||
@@ -197,7 +231,7 @@ class CreateGoalTool(Tool, _GoalToolsMixin):
|
|||||||
ui_summary: str | None = None,
|
ui_summary: str | None = None,
|
||||||
**kwargs: Any,
|
**kwargs: Any,
|
||||||
) -> str:
|
) -> str:
|
||||||
sess = self._session()
|
sess = await self._session()
|
||||||
if sess is None:
|
if sess is None:
|
||||||
return ToolResult.error(
|
return ToolResult.error(
|
||||||
"Error: create_goal requires an active chat session (missing routing context)."
|
"Error: create_goal requires an active chat session (missing routing context)."
|
||||||
@@ -225,8 +259,7 @@ class CreateGoalTool(Tool, _GoalToolsMixin):
|
|||||||
"ui_summary": summary,
|
"ui_summary": summary,
|
||||||
"started_at": _iso_now(),
|
"started_at": _iso_now(),
|
||||||
}
|
}
|
||||||
self._save_goal_state(sess, blob, reset_continuation=True)
|
await self._save_goal_state(sess, blob, reset_continuation=True)
|
||||||
await self._publish_goal_state_changed(sess.metadata)
|
|
||||||
extra = f"\nSummary line: {summary}" if summary else ""
|
extra = f"\nSummary line: {summary}" if summary else ""
|
||||||
return (
|
return (
|
||||||
"Goal recorded. Keep working toward the objective using ordinary tools. "
|
"Goal recorded. Keep working toward the objective using ordinary tools. "
|
||||||
@@ -305,7 +338,7 @@ class UpdateGoalTool(Tool, _GoalToolsMixin):
|
|||||||
ui_summary: str | None = None,
|
ui_summary: str | None = None,
|
||||||
**kwargs: Any,
|
**kwargs: Any,
|
||||||
) -> str:
|
) -> str:
|
||||||
sess = self._session()
|
sess = await self._session()
|
||||||
if sess is None:
|
if sess is None:
|
||||||
return ToolResult.error("Error: update_goal requires an active chat session.")
|
return ToolResult.error("Error: update_goal requires an active chat session.")
|
||||||
prior = parse_goal_state(goal_state_raw(sess.metadata))
|
prior = parse_goal_state(goal_state_raw(sess.metadata))
|
||||||
@@ -340,8 +373,7 @@ class UpdateGoalTool(Tool, _GoalToolsMixin):
|
|||||||
"previous_objective": str(prior.get("objective") or ""),
|
"previous_objective": str(prior.get("objective") or ""),
|
||||||
"recap": (recap or "").strip(),
|
"recap": (recap or "").strip(),
|
||||||
}
|
}
|
||||||
self._save_goal_state(sess, blob, reset_continuation=True)
|
await self._save_goal_state(sess, blob, reset_continuation=True)
|
||||||
await self._publish_goal_state_changed(sess.metadata)
|
|
||||||
extra = f"\nSummary line: {summary}" if summary else ""
|
extra = f"\nSummary line: {summary}" if summary else ""
|
||||||
return "Goal replaced. Continue toward the new objective using ordinary tools." + extra
|
return "Goal replaced. Continue toward the new objective using ordinary tools." + extra
|
||||||
|
|
||||||
@@ -359,9 +391,7 @@ class UpdateGoalTool(Tool, _GoalToolsMixin):
|
|||||||
}
|
}
|
||||||
if normalized == "complete":
|
if normalized == "complete":
|
||||||
blob["completed_at"] = ended
|
blob["completed_at"] = ended
|
||||||
self._save_goal_state(sess, blob)
|
await self._save_goal_state(sess, blob, revoke_permission=True)
|
||||||
revoke_goal_mutation_permission()
|
|
||||||
await self._publish_goal_state_changed(sess.metadata)
|
|
||||||
|
|
||||||
tail = (recap or "").strip()
|
tail = (recap or "").strip()
|
||||||
label = {
|
label = {
|
||||||
|
|||||||
@@ -20,10 +20,10 @@ from nanobot.agent.tools.base import Tool, ToolResult
|
|||||||
from nanobot.agent.tools.registry import ToolRegistry
|
from nanobot.agent.tools.registry import ToolRegistry
|
||||||
from nanobot.security.network import (
|
from nanobot.security.network import (
|
||||||
PinnedDNSAsyncTransport,
|
PinnedDNSAsyncTransport,
|
||||||
|
async_resolve_url_target,
|
||||||
|
async_validate_url_target,
|
||||||
env_proxy_applies_to_url,
|
env_proxy_applies_to_url,
|
||||||
httpx_env_proxy_mounts,
|
httpx_env_proxy_mounts,
|
||||||
resolve_url_target,
|
|
||||||
validate_url_target,
|
|
||||||
)
|
)
|
||||||
from nanobot.utils.cancellation import task_is_cancelling
|
from nanobot.utils.cancellation import task_is_cancelling
|
||||||
|
|
||||||
@@ -249,7 +249,7 @@ async def _probe_http_url(url: str, timeout: float = 3.0) -> bool:
|
|||||||
port = parsed.port
|
port = parsed.port
|
||||||
if not port:
|
if not port:
|
||||||
port = 443 if parsed.scheme == "https" else 80
|
port = 443 if parsed.scheme == "https" else 80
|
||||||
ok, _, resolved_ips = resolve_url_target(url)
|
ok, _, resolved_ips = await async_resolve_url_target(url)
|
||||||
if not ok:
|
if not ok:
|
||||||
return False
|
return False
|
||||||
if env_proxy_applies_to_url(url):
|
if env_proxy_applies_to_url(url):
|
||||||
@@ -298,7 +298,7 @@ def _pinned_transport_kwargs() -> dict[str, Any]:
|
|||||||
|
|
||||||
async def _validate_mcp_request_url(request: httpx.Request) -> None:
|
async def _validate_mcp_request_url(request: httpx.Request) -> None:
|
||||||
"""Validate each outgoing MCP HTTP request, including redirect targets."""
|
"""Validate each outgoing MCP HTTP request, including redirect targets."""
|
||||||
ok, error = validate_url_target(str(request.url))
|
ok, error = await async_validate_url_target(str(request.url))
|
||||||
if not ok:
|
if not ok:
|
||||||
raise httpx.RequestError(
|
raise httpx.RequestError(
|
||||||
f"Blocked unsafe MCP URL {_redact_url(str(request.url))} ({error})",
|
f"Blocked unsafe MCP URL {_redact_url(str(request.url))} ({error})",
|
||||||
@@ -1031,7 +1031,7 @@ async def connect_mcp_servers(
|
|||||||
return False
|
return False
|
||||||
|
|
||||||
if transport_type in {"sse", "streamableHttp"}:
|
if transport_type in {"sse", "streamableHttp"}:
|
||||||
ok, error = validate_url_target(cfg.url)
|
ok, error = await async_validate_url_target(cfg.url)
|
||||||
if not ok:
|
if not ok:
|
||||||
logger.warning(
|
logger.warning(
|
||||||
"MCP server '{}': blocked unsafe URL {} ({})",
|
"MCP server '{}': blocked unsafe URL {} ({})",
|
||||||
|
|||||||
@@ -12,6 +12,7 @@ if TYPE_CHECKING:
|
|||||||
from nanobot.agent.tools.shell import ExecToolConfig
|
from nanobot.agent.tools.shell import ExecToolConfig
|
||||||
from nanobot.agent.tools.web import WebToolsConfig
|
from nanobot.agent.tools.web import WebToolsConfig
|
||||||
from nanobot.config.schema import ModelPresetConfig
|
from nanobot.config.schema import ModelPresetConfig
|
||||||
|
from nanobot.providers.base import LLMUsage
|
||||||
from nanobot.utils.llm_runtime import LLMRuntime
|
from nanobot.utils.llm_runtime import LLMRuntime
|
||||||
|
|
||||||
|
|
||||||
@@ -65,7 +66,7 @@ class RuntimeSnapshot:
|
|||||||
web_config: dict[str, object]
|
web_config: dict[str, object]
|
||||||
exec_config: dict[str, object]
|
exec_config: dict[str, object]
|
||||||
subagent_statuses: dict[str, dict[str, object]]
|
subagent_statuses: dict[str, dict[str, object]]
|
||||||
last_usage: dict[str, int]
|
last_usage: Mapping[str, JsonScalar]
|
||||||
scratchpad: dict[str, JsonValue]
|
scratchpad: dict[str, JsonValue]
|
||||||
|
|
||||||
def as_mapping(self) -> Mapping[str, object]:
|
def as_mapping(self) -> Mapping[str, object]:
|
||||||
@@ -106,6 +107,13 @@ class RuntimeControl(Protocol):
|
|||||||
session_key: str | None,
|
session_key: str | None,
|
||||||
) -> LLMRuntime: ...
|
) -> LLMRuntime: ...
|
||||||
|
|
||||||
|
async def set_model_preset_async(
|
||||||
|
self,
|
||||||
|
name: str,
|
||||||
|
*,
|
||||||
|
session_key: str | None,
|
||||||
|
) -> LLMRuntime: ...
|
||||||
|
|
||||||
def set_max_iterations(self, value: int) -> None: ...
|
def set_max_iterations(self, value: int) -> None: ...
|
||||||
|
|
||||||
def set_context_window_tokens(self, value: int) -> LLMRuntime: ...
|
def set_context_window_tokens(self, value: int) -> LLMRuntime: ...
|
||||||
@@ -151,7 +159,7 @@ class _RuntimeControlTarget(Protocol):
|
|||||||
def tool_names(self) -> list[str]: ...
|
def tool_names(self) -> list[str]: ...
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def last_usage(self) -> Mapping[str, int]: ...
|
def last_usage(self) -> LLMUsage | None: ...
|
||||||
|
|
||||||
def set_runtime_model(self, model: str) -> LLMRuntime: ...
|
def set_runtime_model(self, model: str) -> LLMRuntime: ...
|
||||||
|
|
||||||
@@ -161,6 +169,12 @@ class _RuntimeControlTarget(Protocol):
|
|||||||
|
|
||||||
def set_session_model_preset(self, session_key: str, name: str) -> LLMRuntime: ...
|
def set_session_model_preset(self, session_key: str, name: str) -> LLMRuntime: ...
|
||||||
|
|
||||||
|
async def set_session_model_preset_async(
|
||||||
|
self,
|
||||||
|
session_key: str,
|
||||||
|
name: str,
|
||||||
|
) -> LLMRuntime: ...
|
||||||
|
|
||||||
|
|
||||||
class AgentRuntimeControl:
|
class AgentRuntimeControl:
|
||||||
"""Allowlisted adapter from agent-loop state to ``RuntimeControl``."""
|
"""Allowlisted adapter from agent-loop state to ``RuntimeControl``."""
|
||||||
@@ -190,7 +204,7 @@ class AgentRuntimeControl:
|
|||||||
web_config=_snapshot_web_config(target.web_config),
|
web_config=_snapshot_web_config(target.web_config),
|
||||||
exec_config=_snapshot_exec_config(target.exec_config),
|
exec_config=_snapshot_exec_config(target.exec_config),
|
||||||
subagent_statuses=_snapshot_subagent_statuses(target.subagents),
|
subagent_statuses=_snapshot_subagent_statuses(target.subagents),
|
||||||
last_usage=dict(target.last_usage),
|
last_usage=target.last_usage.to_dict() if target.last_usage is not None else {},
|
||||||
scratchpad=_snapshot_json_mapping(self.__scratchpad),
|
scratchpad=_snapshot_json_mapping(self.__scratchpad),
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -207,6 +221,16 @@ class AgentRuntimeControl:
|
|||||||
return self.__target.set_session_model_preset(session_key, name)
|
return self.__target.set_session_model_preset(session_key, name)
|
||||||
return self.__target.set_model_preset(name)
|
return self.__target.set_model_preset(name)
|
||||||
|
|
||||||
|
async def set_model_preset_async(
|
||||||
|
self,
|
||||||
|
name: str,
|
||||||
|
*,
|
||||||
|
session_key: str | None,
|
||||||
|
) -> LLMRuntime:
|
||||||
|
if session_key is not None:
|
||||||
|
return await self.__target.set_session_model_preset_async(session_key, name)
|
||||||
|
return self.__target.set_model_preset(name)
|
||||||
|
|
||||||
def set_max_iterations(self, value: int) -> None:
|
def set_max_iterations(self, value: int) -> None:
|
||||||
self.__target.max_iterations = value
|
self.__target.max_iterations = value
|
||||||
self.__target.subagents.max_iterations = value
|
self.__target.subagents.max_iterations = value
|
||||||
@@ -297,7 +321,7 @@ def _snapshot_subagent_status(status: SubagentStatus) -> dict[str, object]:
|
|||||||
"phase": status.phase,
|
"phase": status.phase,
|
||||||
"iteration": status.iteration,
|
"iteration": status.iteration,
|
||||||
"tool_events": [dict(event) for event in status.tool_events],
|
"tool_events": [dict(event) for event in status.tool_events],
|
||||||
"usage": dict(status.usage),
|
"usage": status.usage.to_dict() if status.usage is not None else None,
|
||||||
"stop_reason": status.stop_reason,
|
"stop_reason": status.stop_reason,
|
||||||
"error": status.error,
|
"error": status.error,
|
||||||
}
|
}
|
||||||
|
|||||||
+103
-53
@@ -4,9 +4,12 @@
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
import fnmatch
|
import fnmatch
|
||||||
import os
|
import os
|
||||||
import re
|
import re
|
||||||
|
import threading
|
||||||
|
import time
|
||||||
from contextlib import suppress
|
from contextlib import suppress
|
||||||
from pathlib import Path, PurePosixPath
|
from pathlib import Path, PurePosixPath
|
||||||
from typing import Any, Iterable, TypeVar
|
from typing import Any, Iterable, TypeVar
|
||||||
@@ -125,6 +128,8 @@ class _SearchTool(_FsTool):
|
|||||||
class FindFilesTool(_SearchTool):
|
class FindFilesTool(_SearchTool):
|
||||||
"""Find files by path fragment, glob, or type."""
|
"""Find files by path fragment, glob, or type."""
|
||||||
_scopes = {"core", "subagent"}
|
_scopes = {"core", "subagent"}
|
||||||
|
_MAX_SCAN_PATHS = 500_000
|
||||||
|
_MAX_SCAN_SECONDS = 30.0
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def name(self) -> str:
|
def name(self) -> str:
|
||||||
@@ -218,66 +223,111 @@ class FindFilesTool(_SearchTool):
|
|||||||
offset: int = 0,
|
offset: int = 0,
|
||||||
**kwargs: Any,
|
**kwargs: Any,
|
||||||
) -> str:
|
) -> str:
|
||||||
|
cancelled = threading.Event()
|
||||||
try:
|
try:
|
||||||
target = self._resolve(path or ".")
|
return await asyncio.to_thread(
|
||||||
if not target.exists():
|
self._execute_sync,
|
||||||
return ToolResult.error(f"Error: Path not found: {path}")
|
path=path,
|
||||||
if not (target.is_dir() or target.is_file()):
|
query=query,
|
||||||
return ToolResult.error(f"Error: Unsupported path: {path}")
|
glob=glob,
|
||||||
|
file_type=type,
|
||||||
if sort not in {"path", "modified"}:
|
include_dirs=include_dirs,
|
||||||
return ToolResult.error("Error: sort must be 'path' or 'modified'")
|
sort=sort,
|
||||||
|
head_limit=head_limit,
|
||||||
limit = (
|
offset=offset,
|
||||||
_DEFAULT_FILE_HEAD_LIMIT
|
cancelled=cancelled,
|
||||||
if head_limit is None
|
|
||||||
else None if head_limit == 0 else head_limit
|
|
||||||
)
|
)
|
||||||
root = target if target.is_dir() else target.parent
|
except asyncio.CancelledError:
|
||||||
matches: list[tuple[str, float]] = []
|
cancelled.set()
|
||||||
|
raise
|
||||||
for candidate in self._iter_paths(target, include_dirs=include_dirs):
|
|
||||||
if candidate.is_dir() and not include_dirs:
|
|
||||||
continue
|
|
||||||
rel_path = candidate.relative_to(root).as_posix()
|
|
||||||
display_path = self._display_path(candidate, root)
|
|
||||||
name = candidate.name
|
|
||||||
|
|
||||||
if glob and not _match_glob(rel_path, name, glob):
|
|
||||||
continue
|
|
||||||
if candidate.is_file() and not _matches_type(name, type):
|
|
||||||
continue
|
|
||||||
if candidate.is_dir() and type:
|
|
||||||
continue
|
|
||||||
if not _matches_query(display_path, query):
|
|
||||||
continue
|
|
||||||
try:
|
|
||||||
mtime = candidate.stat().st_mtime
|
|
||||||
except OSError:
|
|
||||||
mtime = 0.0
|
|
||||||
suffix = "/" if candidate.is_dir() else ""
|
|
||||||
matches.append((display_path + suffix, mtime))
|
|
||||||
|
|
||||||
if sort == "modified":
|
|
||||||
matches.sort(key=lambda item: (-item[1], item[0]))
|
|
||||||
else:
|
|
||||||
matches.sort(key=lambda item: item[0])
|
|
||||||
|
|
||||||
paths = [item[0] for item in matches]
|
|
||||||
paged, truncated = _paginate(paths, limit, offset)
|
|
||||||
if not paged:
|
|
||||||
return "No files found"
|
|
||||||
|
|
||||||
result = "\n".join(paged)
|
|
||||||
note = _pagination_note(limit, offset, truncated)
|
|
||||||
if note:
|
|
||||||
result += "\n\n" + note
|
|
||||||
return result
|
|
||||||
except PermissionError as e:
|
except PermissionError as e:
|
||||||
return ToolResult.error(f"Error: {e}")
|
return ToolResult.error(f"Error: {e}")
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
return ToolResult.error(f"Error finding files: {e}")
|
return ToolResult.error(f"Error finding files: {e}")
|
||||||
|
|
||||||
|
def _execute_sync(
|
||||||
|
self,
|
||||||
|
*,
|
||||||
|
path: str,
|
||||||
|
query: str | None,
|
||||||
|
glob: str | None,
|
||||||
|
file_type: str | None,
|
||||||
|
include_dirs: bool,
|
||||||
|
sort: str,
|
||||||
|
head_limit: int | None,
|
||||||
|
offset: int,
|
||||||
|
cancelled: threading.Event,
|
||||||
|
) -> str:
|
||||||
|
target = self._resolve(path or ".")
|
||||||
|
if not target.exists():
|
||||||
|
return ToolResult.error(f"Error: Path not found: {path}")
|
||||||
|
if not (target.is_dir() or target.is_file()):
|
||||||
|
return ToolResult.error(f"Error: Unsupported path: {path}")
|
||||||
|
|
||||||
|
if sort not in {"path", "modified"}:
|
||||||
|
return ToolResult.error("Error: sort must be 'path' or 'modified'")
|
||||||
|
|
||||||
|
limit = (
|
||||||
|
_DEFAULT_FILE_HEAD_LIMIT
|
||||||
|
if head_limit is None
|
||||||
|
else None if head_limit == 0 else head_limit
|
||||||
|
)
|
||||||
|
root = target if target.is_dir() else target.parent
|
||||||
|
matches: list[tuple[str, float]] = []
|
||||||
|
deadline = time.monotonic() + self._MAX_SCAN_SECONDS
|
||||||
|
scanned = 0
|
||||||
|
|
||||||
|
for candidate in self._iter_paths(target, include_dirs=include_dirs):
|
||||||
|
if cancelled.is_set():
|
||||||
|
raise RuntimeError("find_files scan cancelled")
|
||||||
|
scanned += 1
|
||||||
|
if scanned > self._MAX_SCAN_PATHS:
|
||||||
|
return ToolResult.error(
|
||||||
|
f"Error: find_files scan exceeded {self._MAX_SCAN_PATHS} paths; "
|
||||||
|
"narrow path, query, glob, or type and retry."
|
||||||
|
)
|
||||||
|
if time.monotonic() > deadline:
|
||||||
|
return ToolResult.error(
|
||||||
|
f"Error: find_files scan exceeded {self._MAX_SCAN_SECONDS:g} seconds; "
|
||||||
|
"narrow path, query, glob, or type and retry."
|
||||||
|
)
|
||||||
|
if candidate.is_dir() and not include_dirs:
|
||||||
|
continue
|
||||||
|
rel_path = candidate.relative_to(root).as_posix()
|
||||||
|
display_path = self._display_path(candidate, root)
|
||||||
|
name = candidate.name
|
||||||
|
|
||||||
|
if glob and not _match_glob(rel_path, name, glob):
|
||||||
|
continue
|
||||||
|
if candidate.is_file() and not _matches_type(name, file_type):
|
||||||
|
continue
|
||||||
|
if candidate.is_dir() and file_type:
|
||||||
|
continue
|
||||||
|
if not _matches_query(display_path, query):
|
||||||
|
continue
|
||||||
|
try:
|
||||||
|
mtime = candidate.stat().st_mtime
|
||||||
|
except OSError:
|
||||||
|
mtime = 0.0
|
||||||
|
suffix = "/" if candidate.is_dir() else ""
|
||||||
|
matches.append((display_path + suffix, mtime))
|
||||||
|
|
||||||
|
if sort == "modified":
|
||||||
|
matches.sort(key=lambda item: (-item[1], item[0]))
|
||||||
|
else:
|
||||||
|
matches.sort(key=lambda item: item[0])
|
||||||
|
|
||||||
|
paths = [item[0] for item in matches]
|
||||||
|
paged, truncated = _paginate(paths, limit, offset)
|
||||||
|
if not paged:
|
||||||
|
return "No files found"
|
||||||
|
|
||||||
|
result = "\n".join(paged)
|
||||||
|
note = _pagination_note(limit, offset, truncated)
|
||||||
|
if note:
|
||||||
|
result += "\n\n" + note
|
||||||
|
return result
|
||||||
|
|
||||||
|
|
||||||
class GrepTool(_SearchTool):
|
class GrepTool(_SearchTool):
|
||||||
"""Search file contents using a regex-like pattern."""
|
"""Search file contents using a regex-like pattern."""
|
||||||
|
|||||||
@@ -150,7 +150,7 @@ class MyTool(Tool):
|
|||||||
"Actions: check, set.\n"
|
"Actions: check, set.\n"
|
||||||
"- check (no key): full config overview — start here.\n"
|
"- check (no key): full config overview — start here.\n"
|
||||||
"- check (key): drill into a value. Dot-paths allowed "
|
"- check (key): drill into a value. Dot-paths allowed "
|
||||||
"(e.g. '_last_usage.prompt_tokens', 'web_config.enable').\n"
|
"(e.g. '_last_usage.input_tokens', 'web_config.enable').\n"
|
||||||
"- set (key, value): change config or store notes in your scratchpad. "
|
"- set (key, value): change config or store notes in your scratchpad. "
|
||||||
"Scratchpad keys persist across turns but not restarts.\n"
|
"Scratchpad keys persist across turns but not restarts.\n"
|
||||||
"Key values: _current_iteration (current progress), "
|
"Key values: _current_iteration (current progress), "
|
||||||
@@ -370,7 +370,7 @@ class MyTool(Tool):
|
|||||||
if not self._modify_allowed:
|
if not self._modify_allowed:
|
||||||
return ToolResult.error("Error: set is disabled (tools.my.allow_set is false)")
|
return ToolResult.error("Error: set is disabled (tools.my.allow_set is false)")
|
||||||
if action in ("modify", "set"):
|
if action in ("modify", "set"):
|
||||||
return self._modify(key, value)
|
return await self._modify_async(key, value)
|
||||||
return f"Unknown action: {action}"
|
return f"Unknown action: {action}"
|
||||||
|
|
||||||
# -- inspect --
|
# -- inspect --
|
||||||
@@ -492,6 +492,11 @@ class MyTool(Tool):
|
|||||||
return ToolResult.error(f"Error: '{key}' is read-only and cannot be modified")
|
return ToolResult.error(f"Error: '{key}' is read-only and cannot be modified")
|
||||||
return self._modify_scratchpad(key, value)
|
return self._modify_scratchpad(key, value)
|
||||||
|
|
||||||
|
async def _modify_async(self, key: str | None, value: Any) -> str:
|
||||||
|
if key == "model_preset":
|
||||||
|
return await self._modify_model_preset_async(value)
|
||||||
|
return self._modify(key, value)
|
||||||
|
|
||||||
def _modify_model_preset(self, value: Any) -> str:
|
def _modify_model_preset(self, value: Any) -> str:
|
||||||
if not isinstance(value, str) or not value.strip():
|
if not isinstance(value, str) or not value.strip():
|
||||||
return ToolResult.error("Error: 'model_preset' must be a non-empty string")
|
return ToolResult.error("Error: 'model_preset' must be a non-empty string")
|
||||||
@@ -520,6 +525,34 @@ class MyTool(Tool):
|
|||||||
f"context_window_tokens is now {runtime.context_window_tokens!r}"
|
f"context_window_tokens is now {runtime.context_window_tokens!r}"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
async def _modify_model_preset_async(self, value: Any) -> str:
|
||||||
|
if not isinstance(value, str) or not value.strip():
|
||||||
|
return ToolResult.error("Error: 'model_preset' must be a non-empty string")
|
||||||
|
name = value.strip()
|
||||||
|
session_key = current_request_session_key()
|
||||||
|
old = self._runtime_control.snapshot().model_preset
|
||||||
|
try:
|
||||||
|
runtime = await self._runtime_control.set_model_preset_async(
|
||||||
|
name,
|
||||||
|
session_key=session_key,
|
||||||
|
)
|
||||||
|
except (KeyError, ValueError) as exc:
|
||||||
|
message = str(exc.args[0]) if exc.args else str(exc)
|
||||||
|
punctuation = "" if message.endswith((".", "!", "?")) else "."
|
||||||
|
return ToolResult.error(f"Error: {message}{punctuation}")
|
||||||
|
if session_key:
|
||||||
|
self._audit("modify", f"model_preset = {name!r}")
|
||||||
|
return (
|
||||||
|
f"Set model_preset = {name!r} for the next turn; "
|
||||||
|
f"model will be {runtime.model!r}; "
|
||||||
|
f"context_window_tokens will be {runtime.context_window_tokens!r}"
|
||||||
|
)
|
||||||
|
self._audit("modify", f"model_preset: {old!r} -> {name!r}")
|
||||||
|
return (
|
||||||
|
f"Set model_preset = {name!r} (was {old!r}); model is now {runtime.model!r}; "
|
||||||
|
f"context_window_tokens is now {runtime.context_window_tokens!r}"
|
||||||
|
)
|
||||||
|
|
||||||
def _modify_restricted(self, key: str, value: Any) -> str:
|
def _modify_restricted(self, key: str, value: Any) -> str:
|
||||||
spec = self.RESTRICTED[key]
|
spec = self.RESTRICTED[key]
|
||||||
expected = cast(type[Any], spec["type"])
|
expected = cast(type[Any], spec["type"])
|
||||||
|
|||||||
@@ -267,6 +267,7 @@ class ExecTool(Tool):
|
|||||||
|
|
||||||
_MAX_TIMEOUT = 600
|
_MAX_TIMEOUT = 600
|
||||||
_MAX_OUTPUT = 10_000
|
_MAX_OUTPUT = 10_000
|
||||||
|
_PREPARE_TIMEOUT_SECONDS = 6.0
|
||||||
|
|
||||||
# Kernel device files safe as stdio redirect targets (#3599).
|
# Kernel device files safe as stdio redirect targets (#3599).
|
||||||
_BENIGN_DEVICE_PATHS: frozenset[str] = frozenset({
|
_BENIGN_DEVICE_PATHS: frozenset[str] = frozenset({
|
||||||
@@ -324,7 +325,20 @@ class ExecTool(Tool):
|
|||||||
if max_output_chars is None:
|
if max_output_chars is None:
|
||||||
max_output_chars = max_output_tokens
|
max_output_chars = max_output_tokens
|
||||||
|
|
||||||
prepared = self._prepare_command(command, working_dir, timeout, shell, login)
|
try:
|
||||||
|
prepared = await asyncio.wait_for(
|
||||||
|
asyncio.to_thread(
|
||||||
|
self._prepare_command,
|
||||||
|
command,
|
||||||
|
working_dir,
|
||||||
|
timeout,
|
||||||
|
shell,
|
||||||
|
login,
|
||||||
|
),
|
||||||
|
timeout=self._PREPARE_TIMEOUT_SECONDS,
|
||||||
|
)
|
||||||
|
except asyncio.TimeoutError:
|
||||||
|
return ToolResult.error("Error: command validation timed out")
|
||||||
if isinstance(prepared, str):
|
if isinstance(prepared, str):
|
||||||
return prepared
|
return prepared
|
||||||
|
|
||||||
@@ -470,14 +484,18 @@ class ExecTool(Tool):
|
|||||||
+ _WORKSPACE_BOUNDARY_NOTE
|
+ _WORKSPACE_BOUNDARY_NOTE
|
||||||
)
|
)
|
||||||
|
|
||||||
guard_error = self._guard_command(
|
# Full access is an explicit trust decision. Keep the application-level
|
||||||
command,
|
# command guard aligned with the selected access mode instead of
|
||||||
cwd,
|
# continuing to block commands after workspace restriction is disabled.
|
||||||
restrict_to_workspace=access.restrict_to_workspace,
|
if access.restrict_to_workspace:
|
||||||
workspace_root=workspace_root,
|
guard_error = self._guard_command(
|
||||||
)
|
command,
|
||||||
if guard_error:
|
cwd,
|
||||||
return guard_error
|
restrict_to_workspace=True,
|
||||||
|
workspace_root=workspace_root,
|
||||||
|
)
|
||||||
|
if guard_error:
|
||||||
|
return guard_error
|
||||||
|
|
||||||
if self.sandbox:
|
if self.sandbox:
|
||||||
if _IS_WINDOWS:
|
if _IS_WINDOWS:
|
||||||
@@ -912,6 +930,15 @@ class ExecTool(Tool):
|
|||||||
if self._is_benign_device_path(expanded):
|
if self._is_benign_device_path(expanded):
|
||||||
continue
|
continue
|
||||||
except Exception:
|
except Exception:
|
||||||
|
# ``Path.expanduser()`` raises when a named user's home
|
||||||
|
# cannot be resolved (notably on Windows). An extracted
|
||||||
|
# home path must fail closed rather than bypass the guard.
|
||||||
|
if raw.strip().startswith("~"):
|
||||||
|
return ToolResult.error(
|
||||||
|
"Error: Command blocked by safety guard "
|
||||||
|
"(path outside working dir)"
|
||||||
|
+ _WORKSPACE_BOUNDARY_NOTE
|
||||||
|
)
|
||||||
continue
|
continue
|
||||||
|
|
||||||
if self._is_benign_device_path(str(p)):
|
if self._is_benign_device_path(str(p)):
|
||||||
|
|||||||
+12
-14
@@ -96,7 +96,7 @@ def _normalize(text: str) -> str:
|
|||||||
|
|
||||||
|
|
||||||
def _validate_url(url: str) -> tuple[bool, str]:
|
def _validate_url(url: str) -> tuple[bool, str]:
|
||||||
"""Validate URL scheme/domain. Does NOT check resolved IPs (use _validate_url_safe for that)."""
|
"""Validate URL scheme/domain. Does not resolve IPs; use the async safe helper for that."""
|
||||||
try:
|
try:
|
||||||
p = urlparse(url)
|
p = urlparse(url)
|
||||||
if p.scheme not in ('http', 'https'):
|
if p.scheme not in ('http', 'https'):
|
||||||
@@ -108,18 +108,16 @@ def _validate_url(url: str) -> tuple[bool, str]:
|
|||||||
return False, str(e)
|
return False, str(e)
|
||||||
|
|
||||||
|
|
||||||
def _validate_url_safe(url: str) -> tuple[bool, str]:
|
async def _async_validate_url_safe(url: str) -> tuple[bool, str]:
|
||||||
"""Validate URL with SSRF protection: scheme, domain, and resolved IP check."""
|
from nanobot.security.network import async_validate_url_target
|
||||||
from nanobot.security.network import validate_url_target
|
|
||||||
|
|
||||||
return validate_url_target(url)
|
return await async_validate_url_target(url)
|
||||||
|
|
||||||
|
|
||||||
def _resolve_url_safe(url: str) -> tuple[bool, str, tuple[str, ...]]:
|
async def _async_resolve_url_safe(url: str) -> tuple[bool, str, tuple[str, ...]]:
|
||||||
"""Validate URL and return the resolved IPs to pin during the request."""
|
from nanobot.security.network import async_resolve_url_target
|
||||||
from nanobot.security.network import resolve_url_target
|
|
||||||
|
|
||||||
return resolve_url_target(url)
|
return await async_resolve_url_target(url)
|
||||||
|
|
||||||
|
|
||||||
def _pinned_dns_transport() -> httpx.AsyncBaseTransport:
|
def _pinned_dns_transport() -> httpx.AsyncBaseTransport:
|
||||||
@@ -209,7 +207,7 @@ async def _get_with_safe_redirects(
|
|||||||
"""GET a URL while validating every redirect target before requesting it."""
|
"""GET a URL while validating every redirect target before requesting it."""
|
||||||
current_url = url
|
current_url = url
|
||||||
for _ in range(MAX_REDIRECTS + 1):
|
for _ in range(MAX_REDIRECTS + 1):
|
||||||
is_valid, error_msg, _ = _resolve_url_safe(current_url)
|
is_valid, error_msg, _ = await _async_resolve_url_safe(current_url)
|
||||||
if not is_valid:
|
if not is_valid:
|
||||||
return None, f"Redirect blocked: {error_msg}"
|
return None, f"Redirect blocked: {error_msg}"
|
||||||
|
|
||||||
@@ -229,7 +227,7 @@ async def _get_with_safe_redirects(
|
|||||||
return response, None
|
return response, None
|
||||||
|
|
||||||
next_url = urljoin(str(response.url), location)
|
next_url = urljoin(str(response.url), location)
|
||||||
is_valid, error_msg = _validate_url_safe(next_url)
|
is_valid, error_msg = await _async_validate_url_safe(next_url)
|
||||||
if not is_valid:
|
if not is_valid:
|
||||||
await response.aclose()
|
await response.aclose()
|
||||||
return None, f"Redirect blocked: {error_msg}"
|
return None, f"Redirect blocked: {error_msg}"
|
||||||
@@ -249,7 +247,7 @@ async def _stream_with_safe_redirects(
|
|||||||
current_url = url
|
current_url = url
|
||||||
chain_carries_credentials = _url_carries_credentials(url)
|
chain_carries_credentials = _url_carries_credentials(url)
|
||||||
for _ in range(MAX_REDIRECTS + 1):
|
for _ in range(MAX_REDIRECTS + 1):
|
||||||
is_valid, error_msg, _ = _resolve_url_safe(current_url)
|
is_valid, error_msg, _ = await _async_resolve_url_safe(current_url)
|
||||||
if not is_valid:
|
if not is_valid:
|
||||||
return None, None, f"Redirect blocked: {error_msg}", chain_carries_credentials
|
return None, None, f"Redirect blocked: {error_msg}", chain_carries_credentials
|
||||||
|
|
||||||
@@ -283,7 +281,7 @@ async def _stream_with_safe_redirects(
|
|||||||
chain_carries_credentials = (
|
chain_carries_credentials = (
|
||||||
chain_carries_credentials or _url_carries_credentials(next_url)
|
chain_carries_credentials or _url_carries_credentials(next_url)
|
||||||
)
|
)
|
||||||
is_valid, error_msg = _validate_url_safe(next_url)
|
is_valid, error_msg = await _async_validate_url_safe(next_url)
|
||||||
if not is_valid:
|
if not is_valid:
|
||||||
await stream.__aexit__(None, None, None)
|
await stream.__aexit__(None, None, None)
|
||||||
return None, None, f"Redirect blocked: {error_msg}", chain_carries_credentials
|
return None, None, f"Redirect blocked: {error_msg}", chain_carries_credentials
|
||||||
@@ -1106,7 +1104,7 @@ class WebFetchTool(Tool):
|
|||||||
url = url.strip(" \t\r\n`\"'")
|
url = url.strip(" \t\r\n`\"'")
|
||||||
extract_mode = kwargs.pop("extractMode", extract_mode)
|
extract_mode = kwargs.pop("extractMode", extract_mode)
|
||||||
max_chars = cast(int, kwargs.pop("maxChars", max_chars) or self.max_chars)
|
max_chars = cast(int, kwargs.pop("maxChars", max_chars) or self.max_chars)
|
||||||
is_valid, error_msg = _validate_url_safe(url)
|
is_valid, error_msg = await _async_validate_url_safe(url)
|
||||||
if not is_valid:
|
if not is_valid:
|
||||||
return json.dumps({"error": f"URL validation failed: {error_msg}", "url": url}, ensure_ascii=False)
|
return json.dumps({"error": f"URL validation failed: {error_msg}", "url": url}, ensure_ascii=False)
|
||||||
|
|
||||||
|
|||||||
@@ -4,7 +4,7 @@ from __future__ import annotations
|
|||||||
|
|
||||||
import dataclasses
|
import dataclasses
|
||||||
import time
|
import time
|
||||||
from collections.abc import Awaitable, Callable, Mapping
|
from collections.abc import Awaitable, Callable
|
||||||
from dataclasses import dataclass, field
|
from dataclasses import dataclass, field
|
||||||
from typing import TYPE_CHECKING, Any, cast
|
from typing import TYPE_CHECKING, Any, cast
|
||||||
|
|
||||||
@@ -19,6 +19,7 @@ from nanobot.bus.outbound_events import (
|
|||||||
from nanobot.bus.progress import build_bus_progress_callback
|
from nanobot.bus.progress import build_bus_progress_callback
|
||||||
from nanobot.bus.queue import MessageBus
|
from nanobot.bus.queue import MessageBus
|
||||||
from nanobot.bus.runtime_events import RuntimeEventBus, RuntimeEventPublisher
|
from nanobot.bus.runtime_events import RuntimeEventBus, RuntimeEventPublisher
|
||||||
|
from nanobot.providers.base import LLMUsage
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from nanobot.utils.llm_runtime import LLMRuntime
|
from nanobot.utils.llm_runtime import LLMRuntime
|
||||||
@@ -203,7 +204,7 @@ class TurnDelivery:
|
|||||||
def record_latency(self, latency_ms: int | None) -> None:
|
def record_latency(self, latency_ms: int | None) -> None:
|
||||||
self.runtime_event_publisher.record_turn_latency(self.session_key, latency_ms)
|
self.runtime_event_publisher.record_turn_latency(self.session_key, latency_ms)
|
||||||
|
|
||||||
def record_usage(self, usage: Mapping[str, int]) -> None:
|
def record_usage(self, usage: LLMUsage | None) -> None:
|
||||||
self.runtime_event_publisher.record_turn_usage(self.session_key, usage)
|
self.runtime_event_publisher.record_turn_usage(self.session_key, usage)
|
||||||
|
|
||||||
def background_response(
|
def background_response(
|
||||||
|
|||||||
@@ -18,6 +18,7 @@ from aiohttp import web
|
|||||||
from loguru import logger
|
from loguru import logger
|
||||||
|
|
||||||
from nanobot.config.paths import get_media_dir
|
from nanobot.config.paths import get_media_dir
|
||||||
|
from nanobot.providers.base import LLMUsage
|
||||||
from nanobot.utils.helpers import safe_filename
|
from nanobot.utils.helpers import safe_filename
|
||||||
from nanobot.utils.media_decode import (
|
from nanobot.utils.media_decode import (
|
||||||
MAX_FILE_SIZE,
|
MAX_FILE_SIZE,
|
||||||
@@ -93,11 +94,11 @@ def _error_json(status: int, message: str, err_type: str = "invalid_request_erro
|
|||||||
def _chat_completion_response(
|
def _chat_completion_response(
|
||||||
content: str,
|
content: str,
|
||||||
model: str,
|
model: str,
|
||||||
usage: dict[str, int] | None = None,
|
usage: LLMUsage | None = None,
|
||||||
) -> dict[str, Any]:
|
) -> dict[str, Any]:
|
||||||
prompt = (usage or {}).get("prompt_tokens", 0)
|
prompt = usage.input_tokens if usage else 0
|
||||||
completion = (usage or {}).get("completion_tokens", 0)
|
completion = usage.output_tokens if usage else 0
|
||||||
total = (usage or {}).get("total_tokens", 0) or prompt + completion
|
total = usage.total_tokens if usage else 0
|
||||||
return {
|
return {
|
||||||
"id": f"chatcmpl-{uuid.uuid4().hex[:12]}",
|
"id": f"chatcmpl-{uuid.uuid4().hex[:12]}",
|
||||||
"object": "chat.completion",
|
"object": "chat.completion",
|
||||||
|
|||||||
+374
-41
@@ -2,15 +2,20 @@
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import ctypes
|
||||||
import json
|
import json
|
||||||
import os
|
import os
|
||||||
import re
|
import re
|
||||||
import shlex
|
import shlex
|
||||||
import shutil
|
import shutil
|
||||||
|
import signal
|
||||||
import subprocess
|
import subprocess
|
||||||
import sys
|
import sys
|
||||||
import time
|
import time
|
||||||
from collections.abc import Iterable
|
from collections.abc import Iterable
|
||||||
|
from contextlib import suppress
|
||||||
|
from ctypes import wintypes
|
||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
from importlib import metadata as importlib_metadata
|
from importlib import metadata as importlib_metadata
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
@@ -97,6 +102,141 @@ class CliAppsRuntimeConfig:
|
|||||||
catalog_ttl_seconds: int = 3600
|
catalog_ttl_seconds: int = 3600
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(slots=True)
|
||||||
|
class _PreparedCliRun:
|
||||||
|
name: str
|
||||||
|
entry: str
|
||||||
|
resolved: str
|
||||||
|
args: list[str]
|
||||||
|
cwd: Path
|
||||||
|
timeout: int
|
||||||
|
env: dict[str, str]
|
||||||
|
artifact_snapshot: dict[Path, tuple[int, int]]
|
||||||
|
|
||||||
|
|
||||||
|
class _JobObjectBasicLimitInformation(ctypes.Structure):
|
||||||
|
_fields_ = [
|
||||||
|
("PerProcessUserTimeLimit", ctypes.c_int64),
|
||||||
|
("PerJobUserTimeLimit", ctypes.c_int64),
|
||||||
|
("LimitFlags", wintypes.DWORD),
|
||||||
|
("MinimumWorkingSetSize", ctypes.c_size_t),
|
||||||
|
("MaximumWorkingSetSize", ctypes.c_size_t),
|
||||||
|
("ActiveProcessLimit", wintypes.DWORD),
|
||||||
|
("Affinity", ctypes.c_size_t),
|
||||||
|
("PriorityClass", wintypes.DWORD),
|
||||||
|
("SchedulingClass", wintypes.DWORD),
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
class _IoCounters(ctypes.Structure):
|
||||||
|
_fields_ = [
|
||||||
|
("ReadOperationCount", ctypes.c_uint64),
|
||||||
|
("WriteOperationCount", ctypes.c_uint64),
|
||||||
|
("OtherOperationCount", ctypes.c_uint64),
|
||||||
|
("ReadTransferCount", ctypes.c_uint64),
|
||||||
|
("WriteTransferCount", ctypes.c_uint64),
|
||||||
|
("OtherTransferCount", ctypes.c_uint64),
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
class _JobObjectExtendedLimitInformation(ctypes.Structure):
|
||||||
|
_fields_ = [
|
||||||
|
("BasicLimitInformation", _JobObjectBasicLimitInformation),
|
||||||
|
("IoInfo", _IoCounters),
|
||||||
|
("ProcessMemoryLimit", ctypes.c_size_t),
|
||||||
|
("JobMemoryLimit", ctypes.c_size_t),
|
||||||
|
("PeakProcessMemoryUsed", ctypes.c_size_t),
|
||||||
|
("PeakJobMemoryUsed", ctypes.c_size_t),
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
class _WindowsJob:
|
||||||
|
"""Best-effort Windows process tree ownership for timeout/cancellation."""
|
||||||
|
|
||||||
|
_KILL_ON_JOB_CLOSE = 0x00002000
|
||||||
|
_EXTENDED_LIMIT_INFORMATION = 9
|
||||||
|
_PROCESS_TERMINATE = 0x0001
|
||||||
|
_PROCESS_SET_QUOTA = 0x0100
|
||||||
|
|
||||||
|
def __init__(self) -> None:
|
||||||
|
win_dll = getattr(ctypes, "WinDLL")
|
||||||
|
self._kernel32 = win_dll("kernel32", use_last_error=True)
|
||||||
|
self._kernel32.CreateJobObjectW.argtypes = [wintypes.LPVOID, wintypes.LPCWSTR]
|
||||||
|
self._kernel32.CreateJobObjectW.restype = wintypes.HANDLE
|
||||||
|
self._kernel32.OpenProcess.argtypes = [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD]
|
||||||
|
self._kernel32.OpenProcess.restype = wintypes.HANDLE
|
||||||
|
self._kernel32.AssignProcessToJobObject.argtypes = [wintypes.HANDLE, wintypes.HANDLE]
|
||||||
|
self._kernel32.AssignProcessToJobObject.restype = wintypes.BOOL
|
||||||
|
self._kernel32.SetInformationJobObject.argtypes = [
|
||||||
|
wintypes.HANDLE,
|
||||||
|
ctypes.c_int,
|
||||||
|
wintypes.LPVOID,
|
||||||
|
wintypes.DWORD,
|
||||||
|
]
|
||||||
|
self._kernel32.SetInformationJobObject.restype = wintypes.BOOL
|
||||||
|
self._kernel32.TerminateJobObject.argtypes = [wintypes.HANDLE, wintypes.UINT]
|
||||||
|
self._kernel32.TerminateJobObject.restype = wintypes.BOOL
|
||||||
|
self._kernel32.CloseHandle.argtypes = [wintypes.HANDLE]
|
||||||
|
self._kernel32.CloseHandle.restype = wintypes.BOOL
|
||||||
|
self._handle: Any = self._kernel32.CreateJobObjectW(None, None)
|
||||||
|
if not self._handle:
|
||||||
|
raise OSError(ctypes.get_last_error(), "CreateJobObjectW failed")
|
||||||
|
try:
|
||||||
|
self._set_kill_on_close(True)
|
||||||
|
except OSError:
|
||||||
|
self._kernel32.CloseHandle(self._handle)
|
||||||
|
self._handle = None
|
||||||
|
raise
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def create(cls) -> _WindowsJob | None:
|
||||||
|
if os.name != "nt":
|
||||||
|
return None
|
||||||
|
try:
|
||||||
|
return cls()
|
||||||
|
except OSError as exc:
|
||||||
|
logger.debug("CLI Apps: Windows job object unavailable: {}", exc)
|
||||||
|
return None
|
||||||
|
|
||||||
|
def _set_kill_on_close(self, enabled: bool) -> None:
|
||||||
|
info = _JobObjectExtendedLimitInformation()
|
||||||
|
info.BasicLimitInformation.LimitFlags = self._KILL_ON_JOB_CLOSE if enabled else 0
|
||||||
|
ok = self._kernel32.SetInformationJobObject(
|
||||||
|
self._handle,
|
||||||
|
self._EXTENDED_LIMIT_INFORMATION,
|
||||||
|
ctypes.byref(info),
|
||||||
|
ctypes.sizeof(info),
|
||||||
|
)
|
||||||
|
if not ok:
|
||||||
|
raise OSError(ctypes.get_last_error(), "SetInformationJobObject failed")
|
||||||
|
|
||||||
|
def assign(self, pid: int) -> bool:
|
||||||
|
process_handle = self._kernel32.OpenProcess(
|
||||||
|
self._PROCESS_TERMINATE | self._PROCESS_SET_QUOTA,
|
||||||
|
False,
|
||||||
|
pid,
|
||||||
|
)
|
||||||
|
if not process_handle:
|
||||||
|
return False
|
||||||
|
try:
|
||||||
|
return bool(self._kernel32.AssignProcessToJobObject(self._handle, process_handle))
|
||||||
|
finally:
|
||||||
|
self._kernel32.CloseHandle(process_handle)
|
||||||
|
|
||||||
|
def terminate(self) -> None:
|
||||||
|
if self._handle and not self._kernel32.TerminateJobObject(self._handle, 1):
|
||||||
|
raise OSError(ctypes.get_last_error(), "TerminateJobObject failed")
|
||||||
|
|
||||||
|
def close(self, *, kill_descendants: bool) -> None:
|
||||||
|
if not self._handle:
|
||||||
|
return
|
||||||
|
if not kill_descendants:
|
||||||
|
with suppress(OSError):
|
||||||
|
self._set_kill_on_close(False)
|
||||||
|
self._kernel32.CloseHandle(self._handle)
|
||||||
|
self._handle = None
|
||||||
|
|
||||||
|
|
||||||
_BRANDS: dict[str, tuple[str, str]] = {
|
_BRANDS: dict[str, tuple[str, str]] = {
|
||||||
"1password-cli": ("1password", "#3B66BC"),
|
"1password-cli": ("1password", "#3B66BC"),
|
||||||
"arcgis": ("arcgis", "#2C7AC3"),
|
"arcgis": ("arcgis", "#2C7AC3"),
|
||||||
@@ -1428,6 +1568,197 @@ Use the `run_cli_app` tool with `name="{name}"` for command execution. Do not in
|
|||||||
lines.append(f"- {rel} ({kind}, {self._format_artifact_size(path)})")
|
lines.append(f"- {rel} ({kind}, {self._format_artifact_size(path)})")
|
||||||
return lines
|
return lines
|
||||||
|
|
||||||
|
def _prepare_run(
|
||||||
|
self,
|
||||||
|
name: str,
|
||||||
|
args: list[str] | None,
|
||||||
|
*,
|
||||||
|
json_output: bool,
|
||||||
|
working_dir: str | None,
|
||||||
|
timeout: int | None,
|
||||||
|
restrict_to_workspace: bool,
|
||||||
|
) -> _PreparedCliRun:
|
||||||
|
app = self.get_app(name)
|
||||||
|
installed = self._load_installed()
|
||||||
|
app_name = str(app["name"])
|
||||||
|
if app_name not in installed:
|
||||||
|
raise CliAppError(f"CLI app '{name}' is not installed")
|
||||||
|
cwd = self._resolve_cwd(working_dir, restrict_to_workspace=restrict_to_workspace)
|
||||||
|
entry = str(installed[app_name].get("entry_point") or app.get("entry_point") or "")
|
||||||
|
resolved = shutil.which(entry)
|
||||||
|
if not entry or not resolved:
|
||||||
|
raise CliAppError(f"{entry or name} is not available on PATH")
|
||||||
|
clean_args = [str(arg) for arg in (args or [])]
|
||||||
|
if json_output and "--json" not in clean_args:
|
||||||
|
clean_args = ["--json", *clean_args]
|
||||||
|
effective_timeout = max(1, min(timeout or self.runtime.run_timeout, 600))
|
||||||
|
return _PreparedCliRun(
|
||||||
|
name=name,
|
||||||
|
entry=entry,
|
||||||
|
resolved=resolved,
|
||||||
|
args=clean_args,
|
||||||
|
cwd=cwd,
|
||||||
|
timeout=effective_timeout,
|
||||||
|
env=self._subprocess_env(),
|
||||||
|
artifact_snapshot=self._artifact_snapshot(cwd),
|
||||||
|
)
|
||||||
|
|
||||||
|
def _format_run_result(
|
||||||
|
self,
|
||||||
|
prepared: _PreparedCliRun,
|
||||||
|
*,
|
||||||
|
returncode: int,
|
||||||
|
stdout: str,
|
||||||
|
stderr: str,
|
||||||
|
) -> str:
|
||||||
|
command = " ".join([prepared.entry, *(shlex.quote(arg) for arg in prepared.args)])
|
||||||
|
output = [
|
||||||
|
f"CLI app '{prepared.name}' exited {returncode}.",
|
||||||
|
f"Command: {command}",
|
||||||
|
]
|
||||||
|
if stdout:
|
||||||
|
output.append("\nSTDOUT:\n" + stdout.rstrip())
|
||||||
|
if stderr:
|
||||||
|
output.append("\nSTDERR:\n" + stderr.rstrip())
|
||||||
|
artifacts = self._changed_artifacts(prepared.cwd, prepared.artifact_snapshot)
|
||||||
|
if artifacts:
|
||||||
|
output.append(
|
||||||
|
"\nArtifacts created or updated:\n"
|
||||||
|
+ "\n".join(self._format_artifact_lines(prepared.cwd, artifacts))
|
||||||
|
)
|
||||||
|
if any(path.suffix.lower() in _INLINE_ARTIFACT_EXTENSIONS for path in artifacts):
|
||||||
|
output.append(
|
||||||
|
"\nTo show a preview in WebUI, reference a raster artifact with Markdown "
|
||||||
|
"using its workspace-relative path, for example ``."
|
||||||
|
)
|
||||||
|
return _truncate("\n".join(output))
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _terminate_run_process_sync(
|
||||||
|
process: subprocess.Popen[str],
|
||||||
|
job: _WindowsJob | None,
|
||||||
|
) -> None:
|
||||||
|
if job is not None:
|
||||||
|
with suppress(OSError):
|
||||||
|
job.terminate()
|
||||||
|
job.close(kill_descendants=True)
|
||||||
|
elif os.name == "nt":
|
||||||
|
with suppress(OSError, subprocess.TimeoutExpired):
|
||||||
|
subprocess.run(
|
||||||
|
["taskkill", "/PID", str(process.pid), "/T", "/F"],
|
||||||
|
check=False,
|
||||||
|
stdout=subprocess.DEVNULL,
|
||||||
|
stderr=subprocess.DEVNULL,
|
||||||
|
timeout=5,
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
with suppress(ProcessLookupError, PermissionError):
|
||||||
|
os.killpg(process.pid, signal.SIGKILL)
|
||||||
|
if process.poll() is None:
|
||||||
|
with suppress(ProcessLookupError):
|
||||||
|
process.kill()
|
||||||
|
with suppress(subprocess.TimeoutExpired):
|
||||||
|
process.wait(timeout=5)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
async def _terminate_run_process(
|
||||||
|
process: asyncio.subprocess.Process,
|
||||||
|
job: _WindowsJob | None,
|
||||||
|
) -> None:
|
||||||
|
if job is not None:
|
||||||
|
with suppress(OSError):
|
||||||
|
await asyncio.to_thread(job.terminate)
|
||||||
|
job.close(kill_descendants=True)
|
||||||
|
elif os.name == "nt":
|
||||||
|
with suppress(OSError, asyncio.TimeoutError):
|
||||||
|
await asyncio.wait_for(
|
||||||
|
asyncio.to_thread(
|
||||||
|
subprocess.run,
|
||||||
|
["taskkill", "/PID", str(process.pid), "/T", "/F"],
|
||||||
|
check=False,
|
||||||
|
stdout=subprocess.DEVNULL,
|
||||||
|
stderr=subprocess.DEVNULL,
|
||||||
|
timeout=5,
|
||||||
|
),
|
||||||
|
timeout=6.0,
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
with suppress(ProcessLookupError, PermissionError):
|
||||||
|
os.killpg(process.pid, signal.SIGKILL)
|
||||||
|
if process.returncode is None:
|
||||||
|
with suppress(ProcessLookupError):
|
||||||
|
process.kill()
|
||||||
|
with suppress(asyncio.TimeoutError, ProcessLookupError):
|
||||||
|
await asyncio.wait_for(process.wait(), timeout=5.0)
|
||||||
|
|
||||||
|
async def run_async(
|
||||||
|
self,
|
||||||
|
name: str,
|
||||||
|
args: list[str] | None = None,
|
||||||
|
*,
|
||||||
|
json_output: bool = False,
|
||||||
|
working_dir: str | None = None,
|
||||||
|
timeout: int | None = None,
|
||||||
|
restrict_to_workspace: bool = False,
|
||||||
|
) -> str:
|
||||||
|
prepared = await asyncio.to_thread(
|
||||||
|
self._prepare_run,
|
||||||
|
name,
|
||||||
|
args,
|
||||||
|
json_output=json_output,
|
||||||
|
working_dir=working_dir,
|
||||||
|
timeout=timeout,
|
||||||
|
restrict_to_workspace=restrict_to_workspace,
|
||||||
|
)
|
||||||
|
process_kwargs: dict[str, Any] = {}
|
||||||
|
if os.name == "nt":
|
||||||
|
process_kwargs["creationflags"] = subprocess.CREATE_NEW_PROCESS_GROUP
|
||||||
|
else:
|
||||||
|
process_kwargs["start_new_session"] = True
|
||||||
|
job = _WindowsJob.create()
|
||||||
|
try:
|
||||||
|
process = await asyncio.create_subprocess_exec(
|
||||||
|
prepared.resolved,
|
||||||
|
*prepared.args,
|
||||||
|
cwd=str(prepared.cwd),
|
||||||
|
stdout=asyncio.subprocess.PIPE,
|
||||||
|
stderr=asyncio.subprocess.PIPE,
|
||||||
|
env=prepared.env,
|
||||||
|
**process_kwargs,
|
||||||
|
)
|
||||||
|
except BaseException:
|
||||||
|
if job is not None:
|
||||||
|
job.close(kill_descendants=False)
|
||||||
|
raise
|
||||||
|
if job is not None and not job.assign(process.pid):
|
||||||
|
job.close(kill_descendants=False)
|
||||||
|
job = None
|
||||||
|
try:
|
||||||
|
stdout_raw, stderr_raw = await asyncio.wait_for(
|
||||||
|
process.communicate(),
|
||||||
|
timeout=prepared.timeout,
|
||||||
|
)
|
||||||
|
except asyncio.TimeoutError:
|
||||||
|
await self._terminate_run_process(process, job)
|
||||||
|
return f"CLI app '{prepared.name}' timed out after {prepared.timeout}s"
|
||||||
|
except asyncio.CancelledError:
|
||||||
|
await self._terminate_run_process(process, job)
|
||||||
|
raise
|
||||||
|
except BaseException:
|
||||||
|
await self._terminate_run_process(process, job)
|
||||||
|
raise
|
||||||
|
if job is not None:
|
||||||
|
job.close(kill_descendants=False)
|
||||||
|
stdout = stdout_raw.decode("utf-8", errors="replace")
|
||||||
|
stderr = stderr_raw.decode("utf-8", errors="replace")
|
||||||
|
return await asyncio.to_thread(
|
||||||
|
self._format_run_result,
|
||||||
|
prepared,
|
||||||
|
returncode=process.returncode or 0,
|
||||||
|
stdout=stdout,
|
||||||
|
stderr=stderr,
|
||||||
|
)
|
||||||
|
|
||||||
def run(
|
def run(
|
||||||
self,
|
self,
|
||||||
name: str,
|
name: str,
|
||||||
@@ -1438,50 +1769,52 @@ Use the `run_cli_app` tool with `name="{name}"` for command execution. Do not in
|
|||||||
timeout: int | None = None,
|
timeout: int | None = None,
|
||||||
restrict_to_workspace: bool = False,
|
restrict_to_workspace: bool = False,
|
||||||
) -> str:
|
) -> str:
|
||||||
app = self.get_app(name)
|
prepared = self._prepare_run(
|
||||||
installed = self._load_installed()
|
name,
|
||||||
if str(app["name"]) not in installed:
|
args,
|
||||||
raise CliAppError(f"CLI app '{name}' is not installed")
|
json_output=json_output,
|
||||||
cwd = self._resolve_cwd(working_dir, restrict_to_workspace=restrict_to_workspace)
|
working_dir=working_dir,
|
||||||
entry = str(installed[str(app["name"])].get("entry_point") or app.get("entry_point") or "")
|
timeout=timeout,
|
||||||
resolved = shutil.which(entry)
|
restrict_to_workspace=restrict_to_workspace,
|
||||||
if not entry or not resolved:
|
)
|
||||||
raise CliAppError(f"{entry or name} is not available on PATH")
|
process_kwargs: dict[str, Any] = {}
|
||||||
clean_args = [str(arg) for arg in (args or [])]
|
if os.name == "nt":
|
||||||
if json_output and "--json" not in clean_args:
|
process_kwargs["creationflags"] = subprocess.CREATE_NEW_PROCESS_GROUP
|
||||||
clean_args = ["--json", *clean_args]
|
else:
|
||||||
effective_timeout = max(1, min(timeout or self.runtime.run_timeout, 600))
|
process_kwargs["start_new_session"] = True
|
||||||
artifact_snapshot = self._artifact_snapshot(cwd)
|
job = _WindowsJob.create()
|
||||||
try:
|
try:
|
||||||
result = subprocess.run(
|
process = subprocess.Popen(
|
||||||
[resolved, *clean_args],
|
[prepared.resolved, *prepared.args],
|
||||||
cwd=str(cwd),
|
cwd=str(prepared.cwd),
|
||||||
capture_output=True,
|
stdout=subprocess.PIPE,
|
||||||
|
stderr=subprocess.PIPE,
|
||||||
text=True,
|
text=True,
|
||||||
encoding="utf-8",
|
encoding="utf-8",
|
||||||
errors="replace",
|
errors="replace",
|
||||||
timeout=effective_timeout,
|
env=prepared.env,
|
||||||
env=self._subprocess_env(),
|
**process_kwargs,
|
||||||
)
|
)
|
||||||
|
except BaseException:
|
||||||
|
if job is not None:
|
||||||
|
job.close(kill_descendants=False)
|
||||||
|
raise
|
||||||
|
if job is not None and not job.assign(process.pid):
|
||||||
|
job.close(kill_descendants=False)
|
||||||
|
job = None
|
||||||
|
try:
|
||||||
|
stdout, stderr = process.communicate(timeout=prepared.timeout)
|
||||||
except subprocess.TimeoutExpired:
|
except subprocess.TimeoutExpired:
|
||||||
return f"CLI app '{name}' timed out after {effective_timeout}s"
|
self._terminate_run_process_sync(process, job)
|
||||||
output = [
|
return f"CLI app '{prepared.name}' timed out after {prepared.timeout}s"
|
||||||
f"CLI app '{name}' exited {result.returncode}.",
|
except BaseException:
|
||||||
f"Command: {entry} {' '.join(shlex.quote(arg) for arg in clean_args)}".rstrip(),
|
self._terminate_run_process_sync(process, job)
|
||||||
]
|
raise
|
||||||
if result.stdout:
|
if job is not None:
|
||||||
output.append("\nSTDOUT:\n" + result.stdout.rstrip())
|
job.close(kill_descendants=False)
|
||||||
if result.stderr:
|
return self._format_run_result(
|
||||||
output.append("\nSTDERR:\n" + result.stderr.rstrip())
|
prepared,
|
||||||
artifacts = self._changed_artifacts(cwd, artifact_snapshot)
|
returncode=process.returncode,
|
||||||
if artifacts:
|
stdout=stdout,
|
||||||
output.append(
|
stderr=stderr,
|
||||||
"\nArtifacts created or updated:\n"
|
)
|
||||||
+ "\n".join(self._format_artifact_lines(cwd, artifacts))
|
|
||||||
)
|
|
||||||
if any(path.suffix.lower() in _INLINE_ARTIFACT_EXTENSIONS for path in artifacts):
|
|
||||||
output.append(
|
|
||||||
"\nTo show a preview in WebUI, reference a raster artifact with Markdown "
|
|
||||||
"using its workspace-relative path, for example ``."
|
|
||||||
)
|
|
||||||
return _truncate("\n".join(output))
|
|
||||||
|
|||||||
@@ -12,6 +12,7 @@ from dataclasses import dataclass, replace
|
|||||||
from typing import Any, cast
|
from typing import Any, cast
|
||||||
|
|
||||||
from nanobot.bus.events import OutboundMessage
|
from nanobot.bus.events import OutboundMessage
|
||||||
|
from nanobot.providers.base import LLMUsage
|
||||||
|
|
||||||
|
|
||||||
class OutboundEvent:
|
class OutboundEvent:
|
||||||
@@ -58,10 +59,19 @@ class StreamedResponseEvent(OutboundEvent):
|
|||||||
class TurnEndEvent(OutboundEvent):
|
class TurnEndEvent(OutboundEvent):
|
||||||
latency_ms: int | None = None
|
latency_ms: int | None = None
|
||||||
goal_state: dict[str, Any] | None = None
|
goal_state: dict[str, Any] | None = None
|
||||||
usage: dict[str, int] | None = None
|
usage: LLMUsage | None = None
|
||||||
context_window_tokens: int | None = None
|
context_window_tokens: int | None = None
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True)
|
||||||
|
class RecoveryStateEvent(OutboundEvent):
|
||||||
|
status: str
|
||||||
|
recovery_id: str
|
||||||
|
reason: str | None = None
|
||||||
|
attempts: int = 0
|
||||||
|
can_continue: bool | None = None
|
||||||
|
|
||||||
|
|
||||||
@dataclass(frozen=True)
|
@dataclass(frozen=True)
|
||||||
class GoalStatusEvent(OutboundEvent):
|
class GoalStatusEvent(OutboundEvent):
|
||||||
status: str
|
status: str
|
||||||
@@ -188,11 +198,6 @@ def _legacy_event_from_metadata(msg: OutboundMessage) -> OutboundEvent | None:
|
|||||||
return TurnEndEvent(
|
return TurnEndEvent(
|
||||||
latency_ms=_metadata_int(meta, "latency_ms"),
|
latency_ms=_metadata_int(meta, "latency_ms"),
|
||||||
goal_state=cast(dict[str, Any], goal_state) if isinstance(goal_state, dict) else None,
|
goal_state=cast(dict[str, Any], goal_state) if isinstance(goal_state, dict) else None,
|
||||||
usage=(
|
|
||||||
cast(dict[str, int], meta.get("usage"))
|
|
||||||
if isinstance(meta.get("usage"), dict)
|
|
||||||
else None
|
|
||||||
),
|
|
||||||
context_window_tokens=_metadata_int(meta, "context_window_tokens"),
|
context_window_tokens=_metadata_int(meta, "context_window_tokens"),
|
||||||
)
|
)
|
||||||
if meta.get("_session_updated"):
|
if meta.get("_session_updated"):
|
||||||
|
|||||||
@@ -10,13 +10,14 @@ from __future__ import annotations
|
|||||||
import asyncio
|
import asyncio
|
||||||
import contextlib
|
import contextlib
|
||||||
import inspect
|
import inspect
|
||||||
from collections.abc import Awaitable, Callable, Mapping
|
from collections.abc import Awaitable, Callable
|
||||||
from dataclasses import dataclass, field
|
from dataclasses import dataclass, field
|
||||||
from typing import TYPE_CHECKING, Any
|
from typing import TYPE_CHECKING, Any
|
||||||
|
|
||||||
from loguru import logger
|
from loguru import logger
|
||||||
|
|
||||||
from nanobot.bus.events import InboundMessage
|
from nanobot.bus.events import InboundMessage
|
||||||
|
from nanobot.providers.base import LLMUsage
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from nanobot.utils.llm_runtime import LLMRuntime
|
from nanobot.utils.llm_runtime import LLMRuntime
|
||||||
@@ -72,7 +73,7 @@ class TurnCompleted:
|
|||||||
context: RuntimeEventContext
|
context: RuntimeEventContext
|
||||||
latency_ms: int | None = None
|
latency_ms: int | None = None
|
||||||
runtime: LLMRuntime | None = None
|
runtime: LLMRuntime | None = None
|
||||||
usage: dict[str, int] = field(default_factory=dict)
|
usage: LLMUsage | None = None
|
||||||
|
|
||||||
|
|
||||||
@dataclass(frozen=True)
|
@dataclass(frozen=True)
|
||||||
@@ -180,7 +181,7 @@ class RuntimeEventPublisher:
|
|||||||
self.bus = bus or RuntimeEventBus()
|
self.bus = bus or RuntimeEventBus()
|
||||||
self._turn_latency_ms: dict[str, int] = {}
|
self._turn_latency_ms: dict[str, int] = {}
|
||||||
self._turn_runtime: dict[str, LLMRuntime] = {}
|
self._turn_runtime: dict[str, LLMRuntime] = {}
|
||||||
self._turn_usage: dict[str, dict[str, int]] = {}
|
self._turn_usage: dict[str, LLMUsage] = {}
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def _context(
|
def _context(
|
||||||
@@ -206,12 +207,9 @@ class RuntimeEventPublisher:
|
|||||||
if latency_ms is not None:
|
if latency_ms is not None:
|
||||||
self._turn_latency_ms[session_key] = int(latency_ms)
|
self._turn_latency_ms[session_key] = int(latency_ms)
|
||||||
|
|
||||||
def record_turn_usage(self, session_key: str, usage: Mapping[str, int]) -> None:
|
def record_turn_usage(self, session_key: str, usage: LLMUsage | None) -> None:
|
||||||
self._turn_usage[session_key] = {
|
if usage is not None:
|
||||||
key: int(value)
|
self._turn_usage[session_key] = usage
|
||||||
for key, value in usage.items()
|
|
||||||
if type(value) is int and value >= 0
|
|
||||||
}
|
|
||||||
|
|
||||||
def clear_turn(self, session_key: str) -> None:
|
def clear_turn(self, session_key: str) -> None:
|
||||||
self._turn_latency_ms.pop(session_key, None)
|
self._turn_latency_ms.pop(session_key, None)
|
||||||
@@ -332,7 +330,7 @@ class RuntimeEventPublisher:
|
|||||||
),
|
),
|
||||||
latency_ms=self._turn_latency_ms.pop(session_key, None),
|
latency_ms=self._turn_latency_ms.pop(session_key, None),
|
||||||
runtime=self._turn_runtime.pop(session_key, None),
|
runtime=self._turn_runtime.pop(session_key, None),
|
||||||
usage=self._turn_usage.pop(session_key, {}),
|
usage=self._turn_usage.pop(session_key, None),
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
@@ -21,7 +21,10 @@ from nanobot.bus.events import OutboundMessage
|
|||||||
from nanobot.bus.queue import MessageBus
|
from nanobot.bus.queue import MessageBus
|
||||||
from nanobot.channels.base import BaseChannel
|
from nanobot.channels.base import BaseChannel
|
||||||
from nanobot.config.schema import Base
|
from nanobot.config.schema import Base
|
||||||
from nanobot.security.network import validate_resolved_url, validate_url_target
|
from nanobot.security.network import (
|
||||||
|
async_validate_resolved_url,
|
||||||
|
async_validate_url_target,
|
||||||
|
)
|
||||||
|
|
||||||
DINGTALK_MAX_REMOTE_MEDIA_BYTES = 20 * 1024 * 1024
|
DINGTALK_MAX_REMOTE_MEDIA_BYTES = 20 * 1024 * 1024
|
||||||
DINGTALK_MAX_REMOTE_MEDIA_REDIRECTS = 3
|
DINGTALK_MAX_REMOTE_MEDIA_REDIRECTS = 3
|
||||||
@@ -417,8 +420,8 @@ class DingTalkChannel(BaseChannel):
|
|||||||
return self._zip_bytes(filename, data)
|
return self._zip_bytes(filename, data)
|
||||||
return data, filename, content_type
|
return data, filename, content_type
|
||||||
|
|
||||||
def _validate_remote_media_url(self, media_ref: str) -> bool:
|
async def _validate_remote_media_url(self, media_ref: str) -> bool:
|
||||||
ok, err = validate_url_target(media_ref)
|
ok, err = await async_validate_url_target(media_ref)
|
||||||
if not ok:
|
if not ok:
|
||||||
self.logger.warning("remote media URL blocked ref={} reason={}", media_ref, err)
|
self.logger.warning("remote media URL blocked ref={} reason={}", media_ref, err)
|
||||||
return False
|
return False
|
||||||
@@ -434,7 +437,11 @@ class DingTalkChannel(BaseChannel):
|
|||||||
allowed_hosts = {host.lower() for host in self.config.remote_media_redirect_allowed_hosts}
|
allowed_hosts = {host.lower() for host in self.config.remote_media_redirect_allowed_hosts}
|
||||||
return next_host in allowed_hosts
|
return next_host in allowed_hosts
|
||||||
|
|
||||||
def _next_remote_media_url(self, current_url: str, location: str | None) -> str | None:
|
async def _next_remote_media_url(
|
||||||
|
self,
|
||||||
|
current_url: str,
|
||||||
|
location: str | None,
|
||||||
|
) -> str | None:
|
||||||
if not self.config.allow_remote_media_redirects:
|
if not self.config.allow_remote_media_redirects:
|
||||||
self.logger.warning("media download redirect refused ref={}", current_url)
|
self.logger.warning("media download redirect refused ref={}", current_url)
|
||||||
return None
|
return None
|
||||||
@@ -449,7 +456,7 @@ class DingTalkChannel(BaseChannel):
|
|||||||
next_url,
|
next_url,
|
||||||
)
|
)
|
||||||
return None
|
return None
|
||||||
if not self._validate_remote_media_url(next_url):
|
if not await self._validate_remote_media_url(next_url):
|
||||||
return None
|
return None
|
||||||
return next_url
|
return next_url
|
||||||
|
|
||||||
@@ -461,7 +468,7 @@ class DingTalkChannel(BaseChannel):
|
|||||||
if not self._http:
|
if not self._http:
|
||||||
return None, None
|
return None, None
|
||||||
|
|
||||||
if not self._validate_remote_media_url(media_ref):
|
if not await self._validate_remote_media_url(media_ref):
|
||||||
return None, None
|
return None, None
|
||||||
|
|
||||||
try:
|
try:
|
||||||
@@ -473,7 +480,7 @@ class DingTalkChannel(BaseChannel):
|
|||||||
current_url = media_ref
|
current_url = media_ref
|
||||||
for _ in range(DINGTALK_MAX_REMOTE_MEDIA_REDIRECTS + 1):
|
for _ in range(DINGTALK_MAX_REMOTE_MEDIA_REDIRECTS + 1):
|
||||||
async with stream("GET", current_url, follow_redirects=False) as resp:
|
async with stream("GET", current_url, follow_redirects=False) as resp:
|
||||||
final_ok, final_err = validate_resolved_url(str(resp.url))
|
final_ok, final_err = await async_validate_resolved_url(str(resp.url))
|
||||||
if not final_ok:
|
if not final_ok:
|
||||||
self.logger.warning(
|
self.logger.warning(
|
||||||
"remote media redirect blocked ref={} final={} reason={}",
|
"remote media redirect blocked ref={} final={} reason={}",
|
||||||
@@ -483,7 +490,7 @@ class DingTalkChannel(BaseChannel):
|
|||||||
)
|
)
|
||||||
return None, None
|
return None, None
|
||||||
if 300 <= resp.status_code < 400:
|
if 300 <= resp.status_code < 400:
|
||||||
next_url = self._next_remote_media_url(
|
next_url = await self._next_remote_media_url(
|
||||||
str(resp.url), resp.headers.get("location")
|
str(resp.url), resp.headers.get("location")
|
||||||
)
|
)
|
||||||
if not next_url:
|
if not next_url:
|
||||||
@@ -516,7 +523,9 @@ class DingTalkChannel(BaseChannel):
|
|||||||
current_url = media_ref
|
current_url = media_ref
|
||||||
for _ in range(DINGTALK_MAX_REMOTE_MEDIA_REDIRECTS + 1):
|
for _ in range(DINGTALK_MAX_REMOTE_MEDIA_REDIRECTS + 1):
|
||||||
resp = await self._http.get(current_url, follow_redirects=False)
|
resp = await self._http.get(current_url, follow_redirects=False)
|
||||||
final_ok, final_err = validate_resolved_url(str(getattr(resp, "url", current_url)))
|
final_ok, final_err = await async_validate_resolved_url(
|
||||||
|
str(getattr(resp, "url", current_url))
|
||||||
|
)
|
||||||
if not final_ok:
|
if not final_ok:
|
||||||
self.logger.warning(
|
self.logger.warning(
|
||||||
"remote media redirect blocked ref={} final={} reason={}",
|
"remote media redirect blocked ref={} final={} reason={}",
|
||||||
@@ -526,7 +535,7 @@ class DingTalkChannel(BaseChannel):
|
|||||||
)
|
)
|
||||||
return None, None
|
return None, None
|
||||||
if 300 <= resp.status_code < 400:
|
if 300 <= resp.status_code < 400:
|
||||||
next_url = self._next_remote_media_url(
|
next_url = await self._next_remote_media_url(
|
||||||
str(getattr(resp, "url", current_url)), resp.headers.get("location")
|
str(getattr(resp, "url", current_url)), resp.headers.get("location")
|
||||||
)
|
)
|
||||||
if not next_url:
|
if not next_url:
|
||||||
|
|||||||
@@ -145,14 +145,9 @@ class _FakeChannel:
|
|||||||
class _FakeInteractionResponse:
|
class _FakeInteractionResponse:
|
||||||
def __init__(self) -> None:
|
def __init__(self) -> None:
|
||||||
self.messages: list[dict] = []
|
self.messages: list[dict] = []
|
||||||
self._done = False
|
|
||||||
|
|
||||||
async def send_message(self, content: str, *, ephemeral: bool = False) -> None:
|
async def send_message(self, content: str, *, ephemeral: bool = False) -> None:
|
||||||
self.messages.append({"content": content, "ephemeral": ephemeral})
|
self.messages.append({"content": content, "ephemeral": ephemeral})
|
||||||
self._done = True
|
|
||||||
|
|
||||||
def is_done(self) -> bool:
|
|
||||||
return self._done
|
|
||||||
|
|
||||||
|
|
||||||
def _make_interaction(
|
def _make_interaction(
|
||||||
|
|||||||
@@ -104,6 +104,9 @@ class ChannelManager:
|
|||||||
webui_mcp_runtime_status: Callable[[], Mapping[str, str]] | None = None,
|
webui_mcp_runtime_status: Callable[[], Mapping[str, str]] | None = None,
|
||||||
webui_mcp_reload: Callable[[], Awaitable[dict[str, Any]]] | None = None,
|
webui_mcp_reload: Callable[[], Awaitable[dict[str, Any]]] | None = None,
|
||||||
webui_skill_state_action: Callable[[set[str]], None] | None = None,
|
webui_skill_state_action: Callable[[set[str]], None] | None = None,
|
||||||
|
webui_recovery_action: (
|
||||||
|
Callable[[str, dict[str, Any]], Awaitable[dict[str, Any]]] | None
|
||||||
|
) = None,
|
||||||
config_path: Path | None = None,
|
config_path: Path | None = None,
|
||||||
):
|
):
|
||||||
if config_path is None:
|
if config_path is None:
|
||||||
@@ -126,6 +129,7 @@ class ChannelManager:
|
|||||||
self._webui_mcp_runtime_status = webui_mcp_runtime_status
|
self._webui_mcp_runtime_status = webui_mcp_runtime_status
|
||||||
self._webui_mcp_reload = webui_mcp_reload
|
self._webui_mcp_reload = webui_mcp_reload
|
||||||
self._webui_skill_state_action = webui_skill_state_action
|
self._webui_skill_state_action = webui_skill_state_action
|
||||||
|
self._webui_recovery_action = webui_recovery_action
|
||||||
self.channels: dict[str, BaseChannel] = {}
|
self.channels: dict[str, BaseChannel] = {}
|
||||||
self._channel_owners: dict[str, str] = {}
|
self._channel_owners: dict[str, str] = {}
|
||||||
self._channel_runtime_specs: dict[str, tuple[str, str]] = {}
|
self._channel_runtime_specs: dict[str, tuple[str, str]] = {}
|
||||||
@@ -197,6 +201,7 @@ class ChannelManager:
|
|||||||
mcp_runtime_status=self._webui_mcp_runtime_status,
|
mcp_runtime_status=self._webui_mcp_runtime_status,
|
||||||
mcp_reload=self._webui_mcp_reload,
|
mcp_reload=self._webui_mcp_reload,
|
||||||
skill_state_action=self._webui_skill_state_action,
|
skill_state_action=self._webui_skill_state_action,
|
||||||
|
recovery_action=self._webui_recovery_action,
|
||||||
logger=logger,
|
logger=logger,
|
||||||
)
|
)
|
||||||
kwargs["gateway"] = gateway
|
kwargs["gateway"] = gateway
|
||||||
@@ -615,6 +620,12 @@ class ChannelManager:
|
|||||||
if target is None:
|
if target is None:
|
||||||
logger.warning("Restart notice target channel is not enabled: {}", notice.channel)
|
logger.warning("Restart notice target channel is not enabled: {}", notice.channel)
|
||||||
return
|
return
|
||||||
|
if notice.channel == "websocket":
|
||||||
|
# Reconnect and recovery are already represented by WebSocket
|
||||||
|
# protocol state. A generic restart-complete notice must not
|
||||||
|
# masquerade as a recovery transition and overwrite a real
|
||||||
|
# awaiting-user checkpoint in connected clients.
|
||||||
|
return
|
||||||
|
|
||||||
while not target.is_running:
|
while not target.is_running:
|
||||||
remaining = deadline - loop.time()
|
remaining = deadline - loop.time()
|
||||||
|
|||||||
@@ -53,7 +53,6 @@ class MattermostConfig(Base):
|
|||||||
include_thread_context: bool = True
|
include_thread_context: bool = True
|
||||||
thread_context_limit: int = 20
|
thread_context_limit: int = 20
|
||||||
streaming: bool = True
|
streaming: bool = True
|
||||||
streaming_max_chars: int = 16000
|
|
||||||
react_emoji: str = "eyes"
|
react_emoji: str = "eyes"
|
||||||
done_emoji: str = "white_check_mark"
|
done_emoji: str = "white_check_mark"
|
||||||
send_progress: bool = True
|
send_progress: bool = True
|
||||||
@@ -106,7 +105,6 @@ class MattermostChannel(BaseChannel):
|
|||||||
self._ws_task: asyncio.Task[None] | None = None
|
self._ws_task: asyncio.Task[None] | None = None
|
||||||
self._self_id: str | None = None
|
self._self_id: str | None = None
|
||||||
self._self_username: str | None = None
|
self._self_username: str | None = None
|
||||||
self._self_email: str | None = None
|
|
||||||
self._usernames: dict[str, str] = {}
|
self._usernames: dict[str, str] = {}
|
||||||
self._user_emails: dict[str, str] = {}
|
self._user_emails: dict[str, str] = {}
|
||||||
self._channel_types: dict[str, str] = {}
|
self._channel_types: dict[str, str] = {}
|
||||||
@@ -138,7 +136,6 @@ class MattermostChannel(BaseChannel):
|
|||||||
me = cast(dict[str, Any], resp.json())
|
me = cast(dict[str, Any], resp.json())
|
||||||
self._self_id = me.get("id")
|
self._self_id = me.get("id")
|
||||||
self._self_username = me.get("username")
|
self._self_username = me.get("username")
|
||||||
self._self_email = me.get("email", "")
|
|
||||||
self.logger.info("bot @{} connected", self._self_username)
|
self.logger.info("bot @{} connected", self._self_username)
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
self.logger.error("Failed to identify bot user: {}", e)
|
self.logger.error("Failed to identify bot user: {}", e)
|
||||||
|
|||||||
@@ -31,8 +31,6 @@ class _FakeHTTPClient:
|
|||||||
self.delete_calls: list[dict[str, Any]] = []
|
self.delete_calls: list[dict[str, Any]] = []
|
||||||
self._get_responses: dict[str, Any] = {}
|
self._get_responses: dict[str, Any] = {}
|
||||||
self._post_responses: dict[str, Any] = {}
|
self._post_responses: dict[str, Any] = {}
|
||||||
self._put_responses: dict[str, Any] = {}
|
|
||||||
self._delete_status: int | None = None
|
|
||||||
|
|
||||||
def _req(self, method: str, path: str) -> httpx.Request:
|
def _req(self, method: str, path: str) -> httpx.Request:
|
||||||
return httpx.Request(method, f"https://chat.example.com{path}")
|
return httpx.Request(method, f"https://chat.example.com{path}")
|
||||||
@@ -46,12 +44,6 @@ class _FakeHTTPClient:
|
|||||||
def set_post_response(self, path: str, data: Any) -> None:
|
def set_post_response(self, path: str, data: Any) -> None:
|
||||||
self._post_responses[path] = data
|
self._post_responses[path] = data
|
||||||
|
|
||||||
def set_put_response(self, path: str, data: Any) -> None:
|
|
||||||
self._put_responses[path] = data
|
|
||||||
|
|
||||||
def set_delete_status(self, status: int) -> None:
|
|
||||||
self._delete_status = status
|
|
||||||
|
|
||||||
async def get(self, path: str, **kwargs) -> httpx.Response:
|
async def get(self, path: str, **kwargs) -> httpx.Response:
|
||||||
self.get_calls.append({"path": path, **kwargs})
|
self.get_calls.append({"path": path, **kwargs})
|
||||||
data = self._get_responses.get(path, {"id": "resp_" + path.split("/")[-1]})
|
data = self._get_responses.get(path, {"id": "resp_" + path.split("/")[-1]})
|
||||||
@@ -71,13 +63,11 @@ class _FakeHTTPClient:
|
|||||||
|
|
||||||
async def put(self, path: str, *, json: dict[str, Any] | None = None, **kwargs) -> httpx.Response:
|
async def put(self, path: str, *, json: dict[str, Any] | None = None, **kwargs) -> httpx.Response:
|
||||||
self.put_calls.append({"path": path, "json": json})
|
self.put_calls.append({"path": path, "json": json})
|
||||||
data = self._put_responses.get(path, {"id": path.split("/")[-1]})
|
return self._resp(200, {"id": path.split("/")[-1]}, "PUT", path)
|
||||||
return self._resp(200, data, "PUT", path)
|
|
||||||
|
|
||||||
async def delete(self, path: str, **kwargs) -> httpx.Response:
|
async def delete(self, path: str, **kwargs) -> httpx.Response:
|
||||||
self.delete_calls.append({"path": path})
|
self.delete_calls.append({"path": path})
|
||||||
status = self._delete_status if self._delete_status is not None else 200
|
return self._resp(200, {}, "DELETE", path)
|
||||||
return self._resp(status, {}, "DELETE", path)
|
|
||||||
|
|
||||||
async def aclose(self) -> None:
|
async def aclose(self) -> None:
|
||||||
pass
|
pass
|
||||||
@@ -119,7 +109,6 @@ def test_config_defaults():
|
|||||||
assert config.server_url == ""
|
assert config.server_url == ""
|
||||||
assert config.token == ""
|
assert config.token == ""
|
||||||
assert config.streaming is True
|
assert config.streaming is True
|
||||||
assert config.streaming_max_chars == 16000
|
|
||||||
assert config.send_tool_hints is True
|
assert config.send_tool_hints is True
|
||||||
assert config.dm.enabled is True
|
assert config.dm.enabled is True
|
||||||
assert config.dm.policy == "open"
|
assert config.dm.policy == "open"
|
||||||
@@ -150,7 +139,6 @@ def test_config_camelcase_aliases():
|
|||||||
"serverUrl": "https://mm.example.com",
|
"serverUrl": "https://mm.example.com",
|
||||||
"token": "abc123",
|
"token": "abc123",
|
||||||
"allowFromMatchMode": "username",
|
"allowFromMatchMode": "username",
|
||||||
"streamingMaxChars": 8000,
|
|
||||||
"replyInThread": False,
|
"replyInThread": False,
|
||||||
"sendToolHints": False,
|
"sendToolHints": False,
|
||||||
}
|
}
|
||||||
@@ -158,7 +146,6 @@ def test_config_camelcase_aliases():
|
|||||||
assert config.server_url == "https://mm.example.com"
|
assert config.server_url == "https://mm.example.com"
|
||||||
assert config.token == "abc123"
|
assert config.token == "abc123"
|
||||||
assert config.allow_from_match_mode == "username"
|
assert config.allow_from_match_mode == "username"
|
||||||
assert config.streaming_max_chars == 8000
|
|
||||||
assert config.reply_in_thread is False
|
assert config.reply_in_thread is False
|
||||||
assert config.send_tool_hints is False
|
assert config.send_tool_hints is False
|
||||||
|
|
||||||
@@ -194,7 +181,6 @@ async def test_start_identifies_bot():
|
|||||||
|
|
||||||
assert channel._self_id == "botuserid123"
|
assert channel._self_id == "botuserid123"
|
||||||
assert channel._self_username == "nanobot"
|
assert channel._self_username == "nanobot"
|
||||||
assert channel._self_email == "bot@example.com"
|
|
||||||
assert not start_task.done()
|
assert not start_task.done()
|
||||||
user_me_calls = [c for c in fake.get_calls[calls_before:] if "/api/v4/users/me" in c["path"]]
|
user_me_calls = [c for c in fake.get_calls[calls_before:] if "/api/v4/users/me" in c["path"]]
|
||||||
assert len(user_me_calls) == 1
|
assert len(user_me_calls) == 1
|
||||||
@@ -674,7 +660,7 @@ async def test_stream_end_adds_done_emoji():
|
|||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_stream_chunk_boundary_finalizes_and_creates_new():
|
async def test_stream_chunk_boundary_finalizes_and_creates_new():
|
||||||
channel, fake = _make_channel({"streamingMaxChars": 10})
|
channel, fake = _make_channel()
|
||||||
channel._self_id = "bot_id"
|
channel._self_id = "bot_id"
|
||||||
fake.set_post_response("/api/v4/posts", {"id": "post_1"})
|
fake.set_post_response("/api/v4/posts", {"id": "post_1"})
|
||||||
|
|
||||||
|
|||||||
@@ -277,7 +277,7 @@ class MochatChannel(BaseChannel):
|
|||||||
self.config: MochatConfig = config
|
self.config: MochatConfig = config
|
||||||
self._http: httpx.AsyncClient | None = None
|
self._http: httpx.AsyncClient | None = None
|
||||||
self._socket: Any = None
|
self._socket: Any = None
|
||||||
self._ws_connected = self._ws_ready = False
|
self._ws_ready = False
|
||||||
|
|
||||||
self._state_dir = get_runtime_subdir("mochat")
|
self._state_dir = get_runtime_subdir("mochat")
|
||||||
self._cursor_path = self._state_dir / "session_cursors.json"
|
self._cursor_path = self._state_dir / "session_cursors.json"
|
||||||
@@ -346,7 +346,7 @@ class MochatChannel(BaseChannel):
|
|||||||
if self._http:
|
if self._http:
|
||||||
await self._http.aclose()
|
await self._http.aclose()
|
||||||
self._http = None
|
self._http = None
|
||||||
self._ws_connected = self._ws_ready = False
|
self._ws_ready = False
|
||||||
|
|
||||||
async def send(self, msg: OutboundMessage) -> None:
|
async def send(self, msg: OutboundMessage) -> None:
|
||||||
"""Send outbound message to session or panel."""
|
"""Send outbound message to session or panel."""
|
||||||
@@ -422,7 +422,7 @@ class MochatChannel(BaseChannel):
|
|||||||
)
|
)
|
||||||
|
|
||||||
async def connect() -> None:
|
async def connect() -> None:
|
||||||
self._ws_connected, self._ws_ready = True, False
|
self._ws_ready = False
|
||||||
self.logger.info("websocket connected")
|
self.logger.info("websocket connected")
|
||||||
subscribed = await self._subscribe_all()
|
subscribed = await self._subscribe_all()
|
||||||
self._ws_ready = subscribed
|
self._ws_ready = subscribed
|
||||||
@@ -431,7 +431,7 @@ class MochatChannel(BaseChannel):
|
|||||||
async def disconnect() -> None:
|
async def disconnect() -> None:
|
||||||
if not self._running:
|
if not self._running:
|
||||||
return
|
return
|
||||||
self._ws_connected = self._ws_ready = False
|
self._ws_ready = False
|
||||||
self.logger.warning("websocket disconnected")
|
self.logger.warning("websocket disconnected")
|
||||||
await self._ensure_fallback_workers()
|
await self._ensure_fallback_workers()
|
||||||
|
|
||||||
|
|||||||
@@ -24,7 +24,7 @@ from nanobot.bus.queue import MessageBus
|
|||||||
from nanobot.channels.base import BaseChannel
|
from nanobot.channels.base import BaseChannel
|
||||||
from nanobot.config.paths import get_media_dir
|
from nanobot.config.paths import get_media_dir
|
||||||
from nanobot.config.schema import Base
|
from nanobot.config.schema import Base
|
||||||
from nanobot.security.network import validate_url_target
|
from nanobot.security.network import async_validate_url_target
|
||||||
from nanobot.utils.helpers import safe_filename
|
from nanobot.utils.helpers import safe_filename
|
||||||
|
|
||||||
_DOWNLOAD_TIMEOUT = aiohttp.ClientTimeout(total=60)
|
_DOWNLOAD_TIMEOUT = aiohttp.ClientTimeout(total=60)
|
||||||
@@ -473,7 +473,7 @@ class NapcatChannel(BaseChannel):
|
|||||||
if not ref:
|
if not ref:
|
||||||
return None
|
return None
|
||||||
if ref.startswith(("http://", "https://")):
|
if ref.startswith(("http://", "https://")):
|
||||||
ok, err = validate_url_target(ref)
|
ok, err = await async_validate_url_target(ref)
|
||||||
if not ok:
|
if not ok:
|
||||||
logger.warning("napcat: rejected remote image '{}': {}", ref, err)
|
logger.warning("napcat: rejected remote image '{}': {}", ref, err)
|
||||||
return None
|
return None
|
||||||
@@ -525,7 +525,7 @@ class NapcatChannel(BaseChannel):
|
|||||||
# logger.debug("napcat: downloading image from {}", url)
|
# logger.debug("napcat: downloading image from {}", url)
|
||||||
if self._http is None:
|
if self._http is None:
|
||||||
return None
|
return None
|
||||||
ok, err = validate_url_target(url)
|
ok, err = await async_validate_url_target(url)
|
||||||
if not ok:
|
if not ok:
|
||||||
logger.warning("napcat: skip image '{}': {}", url, err)
|
logger.warning("napcat: skip image '{}': {}", url, err)
|
||||||
return None
|
return None
|
||||||
|
|||||||
@@ -149,9 +149,13 @@ async def test_download_image_rejects_redirects(tmp_path, monkeypatch) -> None:
|
|||||||
channel = _channel()
|
channel = _channel()
|
||||||
channel._media_root = tmp_path
|
channel._media_root = tmp_path
|
||||||
channel._http = _FakeHttp(_FakeResponse(status=302))
|
channel._http = _FakeHttp(_FakeResponse(status=302))
|
||||||
|
|
||||||
|
async def allow_url(_url: str) -> tuple[bool, str]:
|
||||||
|
return True, ""
|
||||||
|
|
||||||
monkeypatch.setattr(
|
monkeypatch.setattr(
|
||||||
"nanobot.channels.napcat.runtime.validate_url_target",
|
"nanobot.channels.napcat.runtime.async_validate_url_target",
|
||||||
lambda _url: (True, ""),
|
allow_url,
|
||||||
)
|
)
|
||||||
|
|
||||||
result = await channel._download_image({"url": "https://example.com/a.png", "file": "a.png"})
|
result = await channel._download_image({"url": "https://example.com/a.png", "file": "a.png"})
|
||||||
|
|||||||
@@ -40,7 +40,7 @@ from nanobot.bus.events import OutboundMessage
|
|||||||
from nanobot.bus.queue import MessageBus
|
from nanobot.bus.queue import MessageBus
|
||||||
from nanobot.channels.base import BaseChannel
|
from nanobot.channels.base import BaseChannel
|
||||||
from nanobot.config.schema import Base
|
from nanobot.config.schema import Base
|
||||||
from nanobot.security.network import validate_url_target
|
from nanobot.security.network import async_validate_url_target
|
||||||
from nanobot.utils.logging_bridge import redirect_lib_logging
|
from nanobot.utils.logging_bridge import redirect_lib_logging
|
||||||
|
|
||||||
try:
|
try:
|
||||||
@@ -458,7 +458,7 @@ class QQChannel(BaseChannel):
|
|||||||
return None, None
|
return None, None
|
||||||
|
|
||||||
# Remote URL
|
# Remote URL
|
||||||
ok, err = validate_url_target(media_ref)
|
ok, err = await async_validate_url_target(media_ref)
|
||||||
if not ok:
|
if not ok:
|
||||||
self.logger.warning("outbound media URL validation failed url={} err={}", media_ref, err)
|
self.logger.warning("outbound media URL validation failed url={} err={}", media_ref, err)
|
||||||
return None, None
|
return None, None
|
||||||
|
|||||||
@@ -363,13 +363,6 @@ def test_reported_daily_brief_pattern():
|
|||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
|
||||||
def _resolve_chunk_styles(text: str, max_len: int) -> tuple[list[str], list[list[str]]]:
|
|
||||||
"""Helper: full markdown → signal pipeline, including chunking."""
|
|
||||||
plain, styles = _markdown_to_signal(text)
|
|
||||||
chunks = split_message(plain, max_len) if plain else [""]
|
|
||||||
return chunks, _partition_styles(plain, chunks, styles)
|
|
||||||
|
|
||||||
|
|
||||||
def test_partition_styles_single_chunk_passthrough():
|
def test_partition_styles_single_chunk_passthrough():
|
||||||
plain, styles = _markdown_to_signal("**bold** plain *it*")
|
plain, styles = _markdown_to_signal("**bold** plain *it*")
|
||||||
parts = _partition_styles(plain, [plain], styles)
|
parts = _partition_styles(plain, [plain], styles)
|
||||||
|
|||||||
@@ -23,8 +23,8 @@ from nanobot.config.schema import Base
|
|||||||
from nanobot.pairing import is_approved
|
from nanobot.pairing import is_approved
|
||||||
from nanobot.security.network import (
|
from nanobot.security.network import (
|
||||||
PinnedDNSAsyncTransport,
|
PinnedDNSAsyncTransport,
|
||||||
|
async_validate_url_target,
|
||||||
httpx_env_proxy_mounts,
|
httpx_env_proxy_mounts,
|
||||||
validate_url_target,
|
|
||||||
)
|
)
|
||||||
from nanobot.utils.helpers import safe_filename, split_message
|
from nanobot.utils.helpers import safe_filename, split_message
|
||||||
|
|
||||||
@@ -69,7 +69,6 @@ class SlackConfig(Base):
|
|||||||
webhook_path: str = "/slack/events"
|
webhook_path: str = "/slack/events"
|
||||||
bot_token: str = ""
|
bot_token: str = ""
|
||||||
app_token: str = ""
|
app_token: str = ""
|
||||||
user_token_read_only: bool = True
|
|
||||||
reply_in_thread: bool = True
|
reply_in_thread: bool = True
|
||||||
react_emoji: str = "eyes"
|
react_emoji: str = "eyes"
|
||||||
done_emoji: str = "white_check_mark"
|
done_emoji: str = "white_check_mark"
|
||||||
@@ -96,7 +95,7 @@ _HTML_DOWNLOAD_PREFIXES = (b"<!doctype html", b"<html")
|
|||||||
|
|
||||||
async def _validate_slack_download_request(request: httpx.Request) -> None:
|
async def _validate_slack_download_request(request: httpx.Request) -> None:
|
||||||
"""Validate every Slack file request, including redirects, before transport."""
|
"""Validate every Slack file request, including redirects, before transport."""
|
||||||
ok, error = validate_url_target(str(request.url))
|
ok, error = await async_validate_url_target(str(request.url))
|
||||||
if not ok:
|
if not ok:
|
||||||
raise httpx.RequestError(f"unsafe Slack file URL: {error}", request=request)
|
raise httpx.RequestError(f"unsafe Slack file URL: {error}", request=request)
|
||||||
|
|
||||||
|
|||||||
@@ -859,13 +859,13 @@ def _patch_download_validation(
|
|||||||
monkeypatch: pytest.MonkeyPatch,
|
monkeypatch: pytest.MonkeyPatch,
|
||||||
validated: list[str],
|
validated: list[str],
|
||||||
) -> None:
|
) -> None:
|
||||||
def validate(url: str) -> tuple[bool, str]:
|
async def validate(url: str) -> tuple[bool, str]:
|
||||||
validated.append(url)
|
validated.append(url)
|
||||||
if "169.254.169.254" in url:
|
if "169.254.169.254" in url:
|
||||||
return False, "blocked metadata address"
|
return False, "blocked metadata address"
|
||||||
return True, ""
|
return True, ""
|
||||||
|
|
||||||
monkeypatch.setattr("nanobot.channels.slack.runtime.validate_url_target", validate)
|
monkeypatch.setattr("nanobot.channels.slack.runtime.async_validate_url_target", validate)
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
|
|||||||
@@ -36,7 +36,7 @@ from nanobot.channels.base import BaseChannel
|
|||||||
from nanobot.command.builtin import build_help_text
|
from nanobot.command.builtin import build_help_text
|
||||||
from nanobot.config.paths import get_media_dir
|
from nanobot.config.paths import get_media_dir
|
||||||
from nanobot.config.schema import Base
|
from nanobot.config.schema import Base
|
||||||
from nanobot.security.network import validate_url_target
|
from nanobot.security.network import async_validate_url_target
|
||||||
from nanobot.utils.helpers import split_message
|
from nanobot.utils.helpers import split_message
|
||||||
from nanobot.utils.logging_bridge import redirect_lib_logging
|
from nanobot.utils.logging_bridge import redirect_lib_logging
|
||||||
|
|
||||||
@@ -956,7 +956,7 @@ class TelegramChannel(BaseChannel):
|
|||||||
|
|
||||||
# Telegram Bot API accepts HTTP(S) URLs directly for media params.
|
# Telegram Bot API accepts HTTP(S) URLs directly for media params.
|
||||||
if self._is_remote_media_url(media_path):
|
if self._is_remote_media_url(media_path):
|
||||||
ok, error = validate_url_target(media_path)
|
ok, error = await async_validate_url_target(media_path)
|
||||||
if not ok:
|
if not ok:
|
||||||
raise ValueError(f"unsafe media URL: {error}")
|
raise ValueError(f"unsafe media URL: {error}")
|
||||||
await self._call_with_retry(
|
await self._call_with_retry(
|
||||||
|
|||||||
@@ -1488,7 +1488,14 @@ async def test_send_remote_media_url_after_security_validation(monkeypatch) -> N
|
|||||||
MessageBus(),
|
MessageBus(),
|
||||||
)
|
)
|
||||||
_install_ready_app(channel)
|
_install_ready_app(channel)
|
||||||
monkeypatch.setattr("nanobot.channels.telegram.runtime.validate_url_target", lambda url: (True, ""))
|
|
||||||
|
async def allow_url(_url: str) -> tuple[bool, str]:
|
||||||
|
return True, ""
|
||||||
|
|
||||||
|
monkeypatch.setattr(
|
||||||
|
"nanobot.channels.telegram.runtime.async_validate_url_target",
|
||||||
|
allow_url,
|
||||||
|
)
|
||||||
|
|
||||||
await channel.send(
|
await channel.send(
|
||||||
OutboundMessage(
|
OutboundMessage(
|
||||||
@@ -1546,9 +1553,13 @@ async def test_send_blocks_unsafe_remote_media_url(monkeypatch) -> None:
|
|||||||
MessageBus(),
|
MessageBus(),
|
||||||
)
|
)
|
||||||
_install_ready_app(channel)
|
_install_ready_app(channel)
|
||||||
|
|
||||||
|
async def deny_url(_url: str) -> tuple[bool, str]:
|
||||||
|
return False, "Blocked: example.com resolves to private/internal address 127.0.0.1"
|
||||||
|
|
||||||
monkeypatch.setattr(
|
monkeypatch.setattr(
|
||||||
"nanobot.channels.telegram.runtime.validate_url_target",
|
"nanobot.channels.telegram.runtime.async_validate_url_target",
|
||||||
lambda url: (False, "Blocked: example.com resolves to private/internal address 127.0.0.1"),
|
deny_url,
|
||||||
)
|
)
|
||||||
|
|
||||||
await channel.send(
|
await channel.send(
|
||||||
|
|||||||
@@ -32,6 +32,7 @@ from nanobot.bus.outbound_events import (
|
|||||||
GoalStateSyncEvent,
|
GoalStateSyncEvent,
|
||||||
GoalStatusEvent,
|
GoalStatusEvent,
|
||||||
ProgressEvent,
|
ProgressEvent,
|
||||||
|
RecoveryStateEvent,
|
||||||
RuntimeModelUpdatedEvent,
|
RuntimeModelUpdatedEvent,
|
||||||
SessionUpdatedEvent,
|
SessionUpdatedEvent,
|
||||||
TurnEndEvent,
|
TurnEndEvent,
|
||||||
@@ -43,6 +44,7 @@ from nanobot.bus.queue import MessageBus
|
|||||||
from nanobot.channels.base import BaseChannel
|
from nanobot.channels.base import BaseChannel
|
||||||
from nanobot.command.builtin import USER_SHELL_COMMAND, builtin_command_starts_agent_turn
|
from nanobot.command.builtin import USER_SHELL_COMMAND, builtin_command_starts_agent_turn
|
||||||
from nanobot.config.schema import Base
|
from nanobot.config.schema import Base
|
||||||
|
from nanobot.providers.base import LLMUsage
|
||||||
from nanobot.runtime_context import (
|
from nanobot.runtime_context import (
|
||||||
RUNTIME_CONTEXT_INPUT_META,
|
RUNTIME_CONTEXT_INPUT_META,
|
||||||
WEBUI_QUOTE_METADATA,
|
WEBUI_QUOTE_METADATA,
|
||||||
@@ -53,8 +55,10 @@ from nanobot.security.workspace_access import (
|
|||||||
WORKSPACE_SCOPE_METADATA_KEY,
|
WORKSPACE_SCOPE_METADATA_KEY,
|
||||||
WorkspaceScopeError,
|
WorkspaceScopeError,
|
||||||
)
|
)
|
||||||
|
from nanobot.session.async_compat import call_session_manager
|
||||||
from nanobot.session.goal_state import goal_state_ws_blob
|
from nanobot.session.goal_state import goal_state_ws_blob
|
||||||
from nanobot.session.model_selection import model_preset_from_metadata
|
from nanobot.session.model_selection import model_preset_from_metadata
|
||||||
|
from nanobot.session.recovery import recovery_state_from_metadata
|
||||||
from nanobot.session.webui_turns import (
|
from nanobot.session.webui_turns import (
|
||||||
clear_websocket_turn_if_current,
|
clear_websocket_turn_if_current,
|
||||||
clear_websocket_turns,
|
clear_websocket_turns,
|
||||||
@@ -444,6 +448,26 @@ class WebSocketChannel(BaseChannel):
|
|||||||
if sessions is None:
|
if sessions is None:
|
||||||
return {}
|
return {}
|
||||||
snapshot = sessions.read_session_metadata(f"websocket:{chat_id}")
|
snapshot = sessions.read_session_metadata(f"websocket:{chat_id}")
|
||||||
|
return self._attached_model_fields_from_snapshot(chat_id, snapshot)
|
||||||
|
|
||||||
|
async def _attached_model_fields_async(self, chat_id: str) -> dict[str, Any]:
|
||||||
|
"""Build attach fields without blocking the gateway event loop."""
|
||||||
|
sessions = self.gateway.session_manager
|
||||||
|
if sessions is None:
|
||||||
|
return {}
|
||||||
|
snapshot = await call_session_manager(
|
||||||
|
sessions,
|
||||||
|
"read_session_metadata_async",
|
||||||
|
sessions.read_session_metadata,
|
||||||
|
f"websocket:{chat_id}",
|
||||||
|
)
|
||||||
|
return self._attached_model_fields_from_snapshot(chat_id, snapshot)
|
||||||
|
|
||||||
|
def _attached_model_fields_from_snapshot(
|
||||||
|
self,
|
||||||
|
chat_id: str,
|
||||||
|
snapshot: dict[str, Any] | None,
|
||||||
|
) -> dict[str, Any]:
|
||||||
raw_metadata = snapshot.get("metadata") if snapshot is not None else None
|
raw_metadata = snapshot.get("metadata") if snapshot is not None else None
|
||||||
metadata = cast(dict[str, object], raw_metadata) if isinstance(raw_metadata, dict) else None
|
metadata = cast(dict[str, object], raw_metadata) if isinstance(raw_metadata, dict) else None
|
||||||
fields: dict[str, Any] = {}
|
fields: dict[str, Any] = {}
|
||||||
@@ -453,18 +477,12 @@ class WebSocketChannel(BaseChannel):
|
|||||||
self.logger.warning("ignoring invalid model preset metadata for chat_id={}", chat_id)
|
self.logger.warning("ignoring invalid model preset metadata for chat_id={}", chat_id)
|
||||||
fields["model_preset"] = None
|
fields["model_preset"] = None
|
||||||
if isinstance(metadata, dict):
|
if isinstance(metadata, dict):
|
||||||
usage = metadata.get("_last_usage")
|
recovery_state = recovery_state_from_metadata(metadata)
|
||||||
if isinstance(usage, dict):
|
if recovery_state is not None:
|
||||||
sanitized_usage: dict[str, int | float] = {}
|
fields["recovery_state"] = recovery_state
|
||||||
for key, value in cast(dict[object, object], usage).items():
|
usage = LLMUsage.from_dict(metadata.get("_last_usage"))
|
||||||
if (
|
if usage is not None:
|
||||||
isinstance(key, str)
|
fields["usage"] = usage.to_turn_dict()
|
||||||
and isinstance(value, (int, float))
|
|
||||||
and not isinstance(value, bool)
|
|
||||||
and value >= 0
|
|
||||||
):
|
|
||||||
sanitized_usage[key] = value
|
|
||||||
fields["usage"] = sanitized_usage
|
|
||||||
return fields
|
return fields
|
||||||
|
|
||||||
def _detach(self, connection: ServerConnection, chat_id: str) -> None:
|
def _detach(self, connection: ServerConnection, chat_id: str) -> None:
|
||||||
@@ -513,13 +531,16 @@ class WebSocketChannel(BaseChannel):
|
|||||||
fork_key: str,
|
fork_key: str,
|
||||||
) -> None:
|
) -> None:
|
||||||
"""Attach and hydrate a newly created WebUI chat fork."""
|
"""Attach and hydrate a newly created WebUI chat fork."""
|
||||||
scope = self._workspaces.scope_for_session_key(fork_key)
|
scope = await asyncio.to_thread(
|
||||||
|
self._workspaces.scope_for_session_key,
|
||||||
|
fork_key,
|
||||||
|
)
|
||||||
self._attach(connection, fork_id)
|
self._attach(connection, fork_id)
|
||||||
await self._send_event(
|
await self._send_event(
|
||||||
connection,
|
connection,
|
||||||
"attached",
|
"attached",
|
||||||
chat_id=fork_id,
|
chat_id=fork_id,
|
||||||
**self._attached_model_fields(fork_id),
|
**await self._attached_model_fields_async(fork_id),
|
||||||
)
|
)
|
||||||
await self._send_event(
|
await self._send_event(
|
||||||
connection,
|
connection,
|
||||||
@@ -901,13 +922,13 @@ class WebSocketChannel(BaseChannel):
|
|||||||
)
|
)
|
||||||
if scope is None:
|
if scope is None:
|
||||||
return
|
return
|
||||||
self._workspaces.persist_scope(new_id, scope)
|
self._workspaces.stage_scope(new_id, scope)
|
||||||
self._attach(connection, new_id)
|
self._attach(connection, new_id)
|
||||||
await self._send_event(
|
await self._send_event(
|
||||||
connection,
|
connection,
|
||||||
"attached",
|
"attached",
|
||||||
chat_id=new_id,
|
chat_id=new_id,
|
||||||
**self._attached_model_fields(new_id),
|
**await self._attached_model_fields_async(new_id),
|
||||||
)
|
)
|
||||||
await self._send_event(
|
await self._send_event(
|
||||||
connection,
|
connection,
|
||||||
@@ -963,7 +984,7 @@ class WebSocketChannel(BaseChannel):
|
|||||||
connection,
|
connection,
|
||||||
"attached",
|
"attached",
|
||||||
chat_id=cid,
|
chat_id=cid,
|
||||||
**self._attached_model_fields(cid),
|
**await self._attached_model_fields_async(cid),
|
||||||
)
|
)
|
||||||
await self._hydrate_after_subscribe(cid)
|
await self._hydrate_after_subscribe(cid)
|
||||||
return
|
return
|
||||||
@@ -1018,7 +1039,7 @@ class WebSocketChannel(BaseChannel):
|
|||||||
)
|
)
|
||||||
if scope is None:
|
if scope is None:
|
||||||
return
|
return
|
||||||
self._workspaces.persist_scope(cid, scope)
|
self._workspaces.stage_scope(cid, scope)
|
||||||
# Other clients on the same gateway only need an invalidation; they
|
# Other clients on the same gateway only need an invalidation; they
|
||||||
# can reload the authoritative session row without receiving a
|
# can reload the authoritative session row without receiving a
|
||||||
# local project path that belongs to another connection.
|
# local project path that belongs to another connection.
|
||||||
@@ -1212,7 +1233,6 @@ class WebSocketChannel(BaseChannel):
|
|||||||
if session_mentions:
|
if session_mentions:
|
||||||
metadata["session_mentions"] = session_mentions
|
metadata["session_mentions"] = session_mentions
|
||||||
metadata[WORKSPACE_SCOPE_METADATA_KEY] = scope.metadata()
|
metadata[WORKSPACE_SCOPE_METADATA_KEY] = scope.metadata()
|
||||||
self._workspaces.persist_scope(cid, scope)
|
|
||||||
is_webui = metadata.get("webui") is True
|
is_webui = metadata.get("webui") is True
|
||||||
queued_owner = None
|
queued_owner = None
|
||||||
if is_webui and not is_user_shell and builtin_command_starts_agent_turn(content):
|
if is_webui and not is_user_shell and builtin_command_starts_agent_turn(content):
|
||||||
@@ -1267,6 +1287,7 @@ class WebSocketChannel(BaseChannel):
|
|||||||
else False
|
else False
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
|
await asyncio.to_thread(self._workspaces.persist_scope, cid, scope)
|
||||||
accepted = True
|
accepted = True
|
||||||
finally:
|
finally:
|
||||||
if not accepted and queued_owner is not None:
|
if not accepted and queued_owner is not None:
|
||||||
@@ -1556,7 +1577,7 @@ class WebSocketChannel(BaseChannel):
|
|||||||
turn_id: str | None = None,
|
turn_id: str | None = None,
|
||||||
) -> Any | None:
|
) -> Any | None:
|
||||||
try:
|
try:
|
||||||
return resolver()
|
return await asyncio.to_thread(resolver)
|
||||||
except WorkspaceScopeError as exc:
|
except WorkspaceScopeError as exc:
|
||||||
await self._send_event(
|
await self._send_event(
|
||||||
connection,
|
connection,
|
||||||
@@ -1740,6 +1761,10 @@ class WebSocketChannel(BaseChannel):
|
|||||||
provenance=event.provenance,
|
provenance=event.provenance,
|
||||||
)
|
)
|
||||||
return
|
return
|
||||||
|
if isinstance(event, RecoveryStateEvent):
|
||||||
|
if conns:
|
||||||
|
await self.send_recovery_state(msg.chat_id, event)
|
||||||
|
return
|
||||||
if isinstance(event, GoalStateSyncEvent):
|
if isinstance(event, GoalStateSyncEvent):
|
||||||
if conns:
|
if conns:
|
||||||
await self.send_goal_state(msg.chat_id, event.goal_state or {"active": False})
|
await self.send_goal_state(msg.chat_id, event.goal_state or {"active": False})
|
||||||
@@ -2010,7 +2035,7 @@ class WebSocketChannel(BaseChannel):
|
|||||||
latency_ms: int | None = None,
|
latency_ms: int | None = None,
|
||||||
*,
|
*,
|
||||||
goal_state: dict[str, Any] | None = None,
|
goal_state: dict[str, Any] | None = None,
|
||||||
usage: dict[str, int] | None = None,
|
usage: LLMUsage | None = None,
|
||||||
context_window_tokens: int | None = None,
|
context_window_tokens: int | None = None,
|
||||||
metadata: dict[str, Any] | None = None,
|
metadata: dict[str, Any] | None = None,
|
||||||
turn_owner: str | None = None,
|
turn_owner: str | None = None,
|
||||||
@@ -2025,8 +2050,8 @@ class WebSocketChannel(BaseChannel):
|
|||||||
body["latency_ms"] = int(latency_ms)
|
body["latency_ms"] = int(latency_ms)
|
||||||
if goal_state is not None:
|
if goal_state is not None:
|
||||||
body["goal_state"] = goal_state
|
body["goal_state"] = goal_state
|
||||||
if usage:
|
if usage is not None:
|
||||||
body["usage"] = usage
|
body["usage"] = usage.to_turn_dict()
|
||||||
if context_window_tokens is not None:
|
if context_window_tokens is not None:
|
||||||
body["context_window_tokens"] = int(context_window_tokens)
|
body["context_window_tokens"] = int(context_window_tokens)
|
||||||
canonical_webui_turn = (metadata or {}).get("webui") is True
|
canonical_webui_turn = (metadata or {}).get("webui") is True
|
||||||
@@ -2057,6 +2082,27 @@ class WebSocketChannel(BaseChannel):
|
|||||||
for connection in conns:
|
for connection in conns:
|
||||||
await self._safe_send_to(connection, raw, label=" turn_end ")
|
await self._safe_send_to(connection, raw, label=" turn_end ")
|
||||||
|
|
||||||
|
async def send_recovery_state(
|
||||||
|
self,
|
||||||
|
chat_id: str,
|
||||||
|
event: RecoveryStateEvent,
|
||||||
|
) -> None:
|
||||||
|
"""Publish one structured recovery transition without chat pollution."""
|
||||||
|
body: dict[str, Any] = {
|
||||||
|
"event": "recovery_state",
|
||||||
|
"chat_id": chat_id,
|
||||||
|
"status": event.status,
|
||||||
|
"recovery_id": event.recovery_id,
|
||||||
|
"attempts": event.attempts,
|
||||||
|
}
|
||||||
|
if event.reason:
|
||||||
|
body["reason"] = event.reason
|
||||||
|
if event.can_continue is not None:
|
||||||
|
body["can_continue"] = event.can_continue
|
||||||
|
raw = json.dumps(body, ensure_ascii=False)
|
||||||
|
for connection in list(self._subs.get(chat_id, ())):
|
||||||
|
await self._safe_send_to(connection, raw, label=" recovery_state ")
|
||||||
|
|
||||||
async def send_goal_state(self, chat_id: str, blob: dict[str, Any]) -> None:
|
async def send_goal_state(self, chat_id: str, blob: dict[str, Any]) -> None:
|
||||||
"""Push persisted goal-state snapshot for *chat_id* (multi-chat isolation)."""
|
"""Push persisted goal-state snapshot for *chat_id* (multi-chat isolation)."""
|
||||||
conns = list(self._subs.get(chat_id, ()))
|
conns = list(self._subs.get(chat_id, ()))
|
||||||
|
|||||||
@@ -27,6 +27,7 @@ from nanobot.bus.outbound_events import (
|
|||||||
GoalStateSyncEvent,
|
GoalStateSyncEvent,
|
||||||
GoalStatusEvent,
|
GoalStatusEvent,
|
||||||
ProgressEvent,
|
ProgressEvent,
|
||||||
|
RecoveryStateEvent,
|
||||||
RuntimeModelUpdatedEvent,
|
RuntimeModelUpdatedEvent,
|
||||||
SessionUpdatedEvent,
|
SessionUpdatedEvent,
|
||||||
TurnEndEvent,
|
TurnEndEvent,
|
||||||
@@ -43,6 +44,7 @@ from nanobot.channels.websocket.runtime import (
|
|||||||
)
|
)
|
||||||
from nanobot.config.loader import load_config, save_config
|
from nanobot.config.loader import load_config, save_config
|
||||||
from nanobot.config.schema import Config, ModelPresetConfig
|
from nanobot.config.schema import Config, ModelPresetConfig
|
||||||
|
from nanobot.providers.base import LLMUsage
|
||||||
from nanobot.runtime_context import RUNTIME_CONTEXT_INPUT_META, WEBUI_QUOTE_SOURCE
|
from nanobot.runtime_context import RUNTIME_CONTEXT_INPUT_META, WEBUI_QUOTE_SOURCE
|
||||||
from nanobot.security.workspace_access import WORKSPACE_SCOPE_METADATA_KEY
|
from nanobot.security.workspace_access import WORKSPACE_SCOPE_METADATA_KEY
|
||||||
from nanobot.session import webui_turns as wth
|
from nanobot.session import webui_turns as wth
|
||||||
@@ -1510,6 +1512,7 @@ async def test_webui_message_scope_inherits_persisted_session_scope(
|
|||||||
},
|
},
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
|
assert sessions.list_sessions() == []
|
||||||
await channel._dispatch_envelope(
|
await channel._dispatch_envelope(
|
||||||
conn,
|
conn,
|
||||||
"webui-client",
|
"webui-client",
|
||||||
@@ -1523,6 +1526,87 @@ async def test_webui_message_scope_inherits_persisted_session_scope(
|
|||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_new_chat_without_message_does_not_create_session(
|
||||||
|
bus: MagicMock,
|
||||||
|
tmp_path,
|
||||||
|
) -> None:
|
||||||
|
sessions = SessionManager(tmp_path / "sessions")
|
||||||
|
channel = WebSocketChannel(
|
||||||
|
{"enabled": True, "allowFrom": ["*"], "host": "127.0.0.1"},
|
||||||
|
bus,
|
||||||
|
gateway=_basic_handler(bus, session_manager=sessions, workspace_path=tmp_path),
|
||||||
|
)
|
||||||
|
conn = AsyncMock()
|
||||||
|
conn.remote_address = ("127.0.0.1", 50123)
|
||||||
|
|
||||||
|
await channel._dispatch_envelope(
|
||||||
|
conn,
|
||||||
|
"tui-client",
|
||||||
|
{
|
||||||
|
"type": "new_chat",
|
||||||
|
"workspace_scope": {
|
||||||
|
"project_path": str(tmp_path),
|
||||||
|
"access_mode": "full",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
|
attached = json.loads(conn.send.await_args_list[0].args[0])
|
||||||
|
assert attached["event"] == "attached"
|
||||||
|
assert sessions.list_sessions() == []
|
||||||
|
assert channel._workspaces.scope_for_session_key(
|
||||||
|
f"websocket:{attached['chat_id']}"
|
||||||
|
).access_mode == "full"
|
||||||
|
|
||||||
|
await channel._cleanup_connection(conn)
|
||||||
|
|
||||||
|
assert sessions.list_sessions() == []
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_failed_first_message_does_not_persist_draft_session(
|
||||||
|
bus: MagicMock,
|
||||||
|
tmp_path,
|
||||||
|
) -> None:
|
||||||
|
sessions = SessionManager(tmp_path / "sessions")
|
||||||
|
channel = WebSocketChannel(
|
||||||
|
{"enabled": True, "allowFrom": ["*"], "host": "127.0.0.1"},
|
||||||
|
bus,
|
||||||
|
gateway=_basic_handler(bus, session_manager=sessions, workspace_path=tmp_path),
|
||||||
|
)
|
||||||
|
conn = AsyncMock()
|
||||||
|
conn.remote_address = ("127.0.0.1", 50123)
|
||||||
|
|
||||||
|
await channel._dispatch_envelope(
|
||||||
|
conn,
|
||||||
|
"tui-client",
|
||||||
|
{
|
||||||
|
"type": "new_chat",
|
||||||
|
"workspace_scope": {
|
||||||
|
"project_path": str(tmp_path),
|
||||||
|
"access_mode": "full",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
)
|
||||||
|
chat_id = json.loads(conn.send.await_args_list[0].args[0])["chat_id"]
|
||||||
|
bus.publish_inbound.side_effect = RuntimeError("queue unavailable")
|
||||||
|
|
||||||
|
with pytest.raises(RuntimeError, match="queue unavailable"):
|
||||||
|
await channel._dispatch_envelope(
|
||||||
|
conn,
|
||||||
|
"tui-client",
|
||||||
|
{
|
||||||
|
"type": "message",
|
||||||
|
"chat_id": chat_id,
|
||||||
|
"content": "hello",
|
||||||
|
"webui": True,
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
|
assert sessions.list_sessions() == []
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_workspace_scope_change_invalidates_other_attached_clients(
|
async def test_workspace_scope_change_invalidates_other_attached_clients(
|
||||||
bus: MagicMock,
|
bus: MagicMock,
|
||||||
@@ -1729,6 +1813,10 @@ async def test_webui_set_workspace_scope_rejects_running_chat(bus: MagicMock, tm
|
|||||||
},
|
},
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
|
channel._workspaces.persist_scope(
|
||||||
|
"chat-running",
|
||||||
|
channel._workspaces.scope_for_session_key("websocket:chat-running"),
|
||||||
|
)
|
||||||
conn.send.reset_mock()
|
conn.send.reset_mock()
|
||||||
|
|
||||||
wth._WEBSOCKET_TURN_WALL_STARTED_AT["chat-running"] = 123.0
|
wth._WEBSOCKET_TURN_WALL_STARTED_AT["chat-running"] = 123.0
|
||||||
@@ -1795,6 +1883,13 @@ async def test_remote_webui_scope_allows_access_reduction(
|
|||||||
payload = json.loads(conn.send.await_args.args[0])
|
payload = json.loads(conn.send.await_args.args[0])
|
||||||
assert payload["event"] == "session_updated"
|
assert payload["event"] == "session_updated"
|
||||||
assert payload["workspace_scope"]["access_mode"] == "restricted"
|
assert payload["workspace_scope"]["access_mode"] == "restricted"
|
||||||
|
assert sessions.list_sessions() == []
|
||||||
|
|
||||||
|
await channel._dispatch_envelope(
|
||||||
|
conn,
|
||||||
|
"webui-client",
|
||||||
|
{"type": "message", "chat_id": "chat-remote", "content": "hello", "webui": True},
|
||||||
|
)
|
||||||
saved = sessions.read_session_file("websocket:chat-remote")
|
saved = sessions.read_session_file("websocket:chat-remote")
|
||||||
assert saved["metadata"]["workspace_scope"] == {
|
assert saved["metadata"]["workspace_scope"] == {
|
||||||
"project_path": str(default_workspace.resolve()),
|
"project_path": str(default_workspace.resolve()),
|
||||||
@@ -1864,8 +1959,10 @@ async def test_remote_access_reduction_rejects_stale_in_flight_message_scope(
|
|||||||
release_hydrate.set()
|
release_hydrate.set()
|
||||||
await message_task
|
await message_task
|
||||||
|
|
||||||
saved = sessions.read_session_file(f"websocket:{chat_id}")
|
assert sessions.read_session_file(f"websocket:{chat_id}") is None
|
||||||
assert saved["metadata"]["workspace_scope"]["access_mode"] == "restricted"
|
assert channel._workspaces.scope_for_session_key(
|
||||||
|
f"websocket:{chat_id}"
|
||||||
|
).access_mode == "restricted"
|
||||||
payload = json.loads(message_conn.send.await_args.args[0])
|
payload = json.loads(message_conn.send.await_args.args[0])
|
||||||
assert payload["event"] == "error"
|
assert payload["event"] == "error"
|
||||||
assert payload["detail"] == "workspace_scope_rejected"
|
assert payload["detail"] == "workspace_scope_rejected"
|
||||||
@@ -1953,8 +2050,10 @@ async def test_native_webui_scope_allows_custom_scope_without_loopback(
|
|||||||
assert payload["workspace_scope"]["restrict_to_workspace"] is False
|
assert payload["workspace_scope"]["restrict_to_workspace"] is False
|
||||||
assert payload["workspace_scope"]["sandbox_status"]["restrict_to_workspace"] is False
|
assert payload["workspace_scope"]["sandbox_status"]["restrict_to_workspace"] is False
|
||||||
assert payload["workspace_scope"]["sandbox_status"]["workspace_root"] == str(project.resolve())
|
assert payload["workspace_scope"]["sandbox_status"]["workspace_root"] == str(project.resolve())
|
||||||
saved = sessions.read_session_file("websocket:chat-native")
|
assert sessions.read_session_file("websocket:chat-native") is None
|
||||||
assert saved["metadata"]["workspace_scope"] == {
|
assert channel._workspaces.scope_for_session_key(
|
||||||
|
"websocket:chat-native"
|
||||||
|
).metadata() == {
|
||||||
"project_path": str(project.resolve()),
|
"project_path": str(project.resolve()),
|
||||||
"access_mode": "full",
|
"access_mode": "full",
|
||||||
}
|
}
|
||||||
@@ -2092,16 +2191,12 @@ async def test_send_scopes_turn_model_updates_to_the_subscribed_chat() -> None:
|
|||||||
|
|
||||||
|
|
||||||
def test_attach_fields_restore_the_session_model_and_latest_usage() -> None:
|
def test_attach_fields_restore_the_session_model_and_latest_usage() -> None:
|
||||||
|
usage = LLMUsage.reported(input_tokens=120, output_tokens=8, total_tokens=175)
|
||||||
manager = MagicMock()
|
manager = MagicMock()
|
||||||
manager.read_session_metadata.return_value = {
|
manager.read_session_metadata.return_value = {
|
||||||
"metadata": {
|
"metadata": {
|
||||||
SESSION_MODEL_PRESET_METADATA_KEY: "Deep Research",
|
SESSION_MODEL_PRESET_METADATA_KEY: "Deep Research",
|
||||||
"_last_usage": {
|
"_last_usage": usage.to_dict(),
|
||||||
"prompt_tokens": 120,
|
|
||||||
"completion_tokens": 8,
|
|
||||||
"negative": -1,
|
|
||||||
"boolean": True,
|
|
||||||
},
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
bus = MagicMock()
|
bus = MagicMock()
|
||||||
@@ -2113,7 +2208,7 @@ def test_attach_fields_restore_the_session_model_and_latest_usage() -> None:
|
|||||||
|
|
||||||
assert channel._attached_model_fields("chat-1") == {
|
assert channel._attached_model_fields("chat-1") == {
|
||||||
"model_preset": "Deep Research",
|
"model_preset": "Deep Research",
|
||||||
"usage": {"prompt_tokens": 120, "completion_tokens": 8},
|
"usage": usage.to_turn_dict(),
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
@@ -2720,6 +2815,39 @@ async def test_send_turn_end_emits_turn_end_event() -> None:
|
|||||||
]
|
]
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_recovery_state_is_a_structured_event_not_assistant_text() -> None:
|
||||||
|
bus = MagicMock()
|
||||||
|
channel = WebSocketChannel(
|
||||||
|
{"enabled": True, "allowFrom": ["*"]},
|
||||||
|
bus,
|
||||||
|
gateway=_basic_handler(bus),
|
||||||
|
)
|
||||||
|
mock_ws = AsyncMock()
|
||||||
|
channel._attach(mock_ws, "chat-1")
|
||||||
|
|
||||||
|
await channel.send(OutboundMessage(
|
||||||
|
channel="websocket",
|
||||||
|
chat_id="chat-1",
|
||||||
|
content="",
|
||||||
|
event=RecoveryStateEvent(
|
||||||
|
status="awaiting_user",
|
||||||
|
recovery_id="recovery-1",
|
||||||
|
reason="tool_state_unknown",
|
||||||
|
attempts=1,
|
||||||
|
),
|
||||||
|
))
|
||||||
|
|
||||||
|
assert _sent_ws_payloads(mock_ws) == [{
|
||||||
|
"event": "recovery_state",
|
||||||
|
"chat_id": "chat-1",
|
||||||
|
"status": "awaiting_user",
|
||||||
|
"recovery_id": "recovery-1",
|
||||||
|
"reason": "tool_state_unknown",
|
||||||
|
"attempts": 1,
|
||||||
|
}]
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_system_command_turn_end_only_refreshes_session_metadata() -> None:
|
async def test_system_command_turn_end_only_refreshes_session_metadata() -> None:
|
||||||
bus = MagicMock()
|
bus = MagicMock()
|
||||||
@@ -3191,6 +3319,11 @@ async def test_send_turn_end_includes_latency_ms_when_present() -> None:
|
|||||||
channel = WebSocketChannel({"enabled": True, "allowFrom": ["*"]}, bus, gateway=_basic_handler(bus))
|
channel = WebSocketChannel({"enabled": True, "allowFrom": ["*"]}, bus, gateway=_basic_handler(bus))
|
||||||
mock_ws = AsyncMock()
|
mock_ws = AsyncMock()
|
||||||
channel._attach(mock_ws, "chat-1")
|
channel._attach(mock_ws, "chat-1")
|
||||||
|
usage = LLMUsage.reported(
|
||||||
|
input_tokens=80,
|
||||||
|
output_tokens=20,
|
||||||
|
cache_read_tokens=40,
|
||||||
|
).with_timing(generation_ms=500, ttft_ms=125)
|
||||||
|
|
||||||
await channel.send(OutboundMessage(
|
await channel.send(OutboundMessage(
|
||||||
channel="websocket",
|
channel="websocket",
|
||||||
@@ -3198,7 +3331,7 @@ async def test_send_turn_end_includes_latency_ms_when_present() -> None:
|
|||||||
content="",
|
content="",
|
||||||
event=TurnEndEvent(
|
event=TurnEndEvent(
|
||||||
latency_ms=1500,
|
latency_ms=1500,
|
||||||
usage={"prompt_tokens": 80, "completion_tokens": 20, "cached_tokens": 40},
|
usage=usage,
|
||||||
context_window_tokens=128_000,
|
context_window_tokens=128_000,
|
||||||
),
|
),
|
||||||
))
|
))
|
||||||
@@ -3208,7 +3341,19 @@ async def test_send_turn_end_includes_latency_ms_when_present() -> None:
|
|||||||
"event": "turn_end",
|
"event": "turn_end",
|
||||||
"chat_id": "chat-1",
|
"chat_id": "chat-1",
|
||||||
"latency_ms": 1500,
|
"latency_ms": 1500,
|
||||||
"usage": {"prompt_tokens": 80, "completion_tokens": 20, "cached_tokens": 40},
|
"usage": {
|
||||||
|
"prompt_tokens": 80,
|
||||||
|
"completion_tokens": 20,
|
||||||
|
"total_tokens": 100,
|
||||||
|
"context_tokens": 80,
|
||||||
|
"cached_tokens": 40,
|
||||||
|
"request_count": 1,
|
||||||
|
"estimated_tokens": 0,
|
||||||
|
"generation_ms": 500,
|
||||||
|
"measured_completion_tokens": 20,
|
||||||
|
"ttft_ms": 125,
|
||||||
|
"timed_requests": 1,
|
||||||
|
},
|
||||||
"context_window_tokens": 128_000,
|
"context_window_tokens": 128_000,
|
||||||
},
|
},
|
||||||
{"event": "session_updated", "chat_id": "chat-1", "scope": "thread"},
|
{"event": "session_updated", "chat_id": "chat-1", "scope": "thread"},
|
||||||
@@ -5175,10 +5320,16 @@ async def test_handle_session_context_get_reads_detached_session() -> None:
|
|||||||
|
|
||||||
from nanobot.session import Session
|
from nanobot.session import Session
|
||||||
|
|
||||||
|
usage = LLMUsage.reported(
|
||||||
|
input_tokens=12,
|
||||||
|
output_tokens=3,
|
||||||
|
total_tokens=175,
|
||||||
|
cache_read_tokens=6,
|
||||||
|
).with_timing(generation_ms=300, ttft_ms=45)
|
||||||
session = Session(
|
session = Session(
|
||||||
key="websocket:context-route",
|
key="websocket:context-route",
|
||||||
messages=[{"role": "user", "content": "hello"}],
|
messages=[{"role": "user", "content": "hello"}],
|
||||||
metadata={"_last_usage": {"prompt_tokens": 12, "completion_tokens": 3}},
|
metadata={"_last_usage": usage.to_dict()},
|
||||||
)
|
)
|
||||||
manager = MagicMock()
|
manager = MagicMock()
|
||||||
manager.read_session_snapshot.return_value = session
|
manager.read_session_snapshot.return_value = session
|
||||||
@@ -5195,7 +5346,19 @@ async def test_handle_session_context_get_reads_detached_session() -> None:
|
|||||||
assert response.status_code == 200
|
assert response.status_code == 200
|
||||||
body = json.loads(response.body.decode())
|
body = json.loads(response.body.decode())
|
||||||
assert body["replay_messages"] == 1
|
assert body["replay_messages"] == 1
|
||||||
assert body["last_usage"] == {"prompt_tokens": 12, "completion_tokens": 3}
|
assert body["last_usage"] == {
|
||||||
|
"prompt_tokens": 12,
|
||||||
|
"completion_tokens": 3,
|
||||||
|
"total_tokens": 175,
|
||||||
|
"context_tokens": 12,
|
||||||
|
"cached_tokens": 6,
|
||||||
|
"request_count": 1,
|
||||||
|
"estimated_tokens": 0,
|
||||||
|
"generation_ms": 300,
|
||||||
|
"measured_completion_tokens": 3,
|
||||||
|
"ttft_ms": 45,
|
||||||
|
"timed_requests": 1,
|
||||||
|
}
|
||||||
manager.read_session_snapshot.assert_called_once_with(session.key)
|
manager.read_session_snapshot.assert_called_once_with(session.key)
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -4,6 +4,7 @@ import asyncio
|
|||||||
import json
|
import json
|
||||||
import random
|
import random
|
||||||
import socket
|
import socket
|
||||||
|
import threading
|
||||||
import time
|
import time
|
||||||
from contextlib import suppress
|
from contextlib import suppress
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
@@ -83,6 +84,7 @@ def _make_handler(
|
|||||||
channel_feature_action: Any | None = None,
|
channel_feature_action: Any | None = None,
|
||||||
channel_runtime_status: Any | None = None,
|
channel_runtime_status: Any | None = None,
|
||||||
mcp_reload: Any | None = None,
|
mcp_reload: Any | None = None,
|
||||||
|
recovery_action: Any | None = None,
|
||||||
) -> GatewayServices:
|
) -> GatewayServices:
|
||||||
config = WebSocketConfig.model_validate(cfg) if isinstance(cfg, dict) else cfg
|
config = WebSocketConfig.model_validate(cfg) if isinstance(cfg, dict) else cfg
|
||||||
workspace = workspace_path or Path.cwd()
|
workspace = workspace_path or Path.cwd()
|
||||||
@@ -103,6 +105,7 @@ def _make_handler(
|
|||||||
channel_feature_action=channel_feature_action,
|
channel_feature_action=channel_feature_action,
|
||||||
channel_runtime_status=channel_runtime_status,
|
channel_runtime_status=channel_runtime_status,
|
||||||
mcp_reload=mcp_reload,
|
mcp_reload=mcp_reload,
|
||||||
|
recovery_action=recovery_action,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
@@ -121,6 +124,7 @@ def _ch(
|
|||||||
channel_feature_action: Any | None = None,
|
channel_feature_action: Any | None = None,
|
||||||
channel_runtime_status: Any | None = None,
|
channel_runtime_status: Any | None = None,
|
||||||
mcp_reload: Any | None = None,
|
mcp_reload: Any | None = None,
|
||||||
|
recovery_action: Any | None = None,
|
||||||
**extra: Any,
|
**extra: Any,
|
||||||
) -> WebSocketChannel:
|
) -> WebSocketChannel:
|
||||||
cfg: dict[str, Any] = {
|
cfg: dict[str, Any] = {
|
||||||
@@ -145,6 +149,7 @@ def _ch(
|
|||||||
channel_feature_action=channel_feature_action,
|
channel_feature_action=channel_feature_action,
|
||||||
channel_runtime_status=channel_runtime_status,
|
channel_runtime_status=channel_runtime_status,
|
||||||
mcp_reload=mcp_reload,
|
mcp_reload=mcp_reload,
|
||||||
|
recovery_action=recovery_action,
|
||||||
)
|
)
|
||||||
return InProcessHttpChannel(cfg, bus, gateway=gateway)
|
return InProcessHttpChannel(cfg, bus, gateway=gateway)
|
||||||
|
|
||||||
@@ -1244,39 +1249,6 @@ async def test_pairing_routes_require_token_and_approve_or_deny(
|
|||||||
assert "Missing pairing code" in missing_code.text
|
assert "Missing pairing code" in missing_code.text
|
||||||
|
|
||||||
|
|
||||||
def test_api_service_settings_read_api_key_from_webui_payload(bus: MagicMock) -> None:
|
|
||||||
channel = _ch(bus)
|
|
||||||
request = _FakeReq(path="/api/settings/api-service/start")
|
|
||||||
setattr(
|
|
||||||
request,
|
|
||||||
"_nanobot_webui_mutation_payload",
|
|
||||||
{"host": "0.0.0.0", "port": 8900, "timeout": 120, "api_key": "secret-token"},
|
|
||||||
)
|
|
||||||
|
|
||||||
query = channel.gateway.http.settings_routes._parse_api_service_settings_query(request)
|
|
||||||
|
|
||||||
assert query == {
|
|
||||||
"host": ["0.0.0.0"],
|
|
||||||
"port": ["8900"],
|
|
||||||
"timeout": ["120"],
|
|
||||||
"api_key": ["secret-token"],
|
|
||||||
}
|
|
||||||
|
|
||||||
|
|
||||||
def test_api_service_settings_reject_non_string_api_key(bus: MagicMock) -> None:
|
|
||||||
from nanobot.webui.settings_api import WebUISettingsError
|
|
||||||
|
|
||||||
channel = _ch(bus)
|
|
||||||
request = _FakeReq(path="/api/settings/api-service/start")
|
|
||||||
setattr(
|
|
||||||
request,
|
|
||||||
"_nanobot_webui_mutation_payload",
|
|
||||||
{"host": "127.0.0.1", "api_key": 123},
|
|
||||||
)
|
|
||||||
|
|
||||||
with pytest.raises(WebUISettingsError, match="API key must be a string"):
|
|
||||||
channel.gateway.http.settings_routes._parse_api_service_settings_query(request)
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_nanobot_feature_remote_install_requires_opt_in(
|
async def test_nanobot_feature_remote_install_requires_opt_in(
|
||||||
bus: MagicMock,
|
bus: MagicMock,
|
||||||
@@ -2605,6 +2577,69 @@ async def test_webui_automations_route_lists_all_jobs_and_allows_user_actions(
|
|||||||
await server_task
|
await server_task
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_webui_cron_update_rearms_started_service_on_owner_loop(
|
||||||
|
bus: MagicMock,
|
||||||
|
tmp_path: Path,
|
||||||
|
monkeypatch: pytest.MonkeyPatch,
|
||||||
|
) -> None:
|
||||||
|
store_path = tmp_path / "cron" / "jobs.json"
|
||||||
|
cron = CronService(store_path, max_sleep_ms=60_000)
|
||||||
|
job = cron.add_job(
|
||||||
|
name="Before update",
|
||||||
|
schedule=CronSchedule(kind="every", every_ms=86_400_000),
|
||||||
|
message="Check the repo status",
|
||||||
|
session_key="websocket:abc",
|
||||||
|
origin_channel="websocket",
|
||||||
|
origin_chat_id="abc",
|
||||||
|
)
|
||||||
|
await cron.start()
|
||||||
|
|
||||||
|
owner_thread_id = threading.get_ident()
|
||||||
|
initial_timer = cron._timer_task
|
||||||
|
request_thread_ids: list[int] = []
|
||||||
|
arm_thread_ids: list[int] = []
|
||||||
|
timer_rearmed = asyncio.Event()
|
||||||
|
original_request_timer_rearm = cron._request_timer_rearm
|
||||||
|
original_arm_timer = cron._arm_timer
|
||||||
|
|
||||||
|
def tracked_request_timer_rearm() -> None:
|
||||||
|
request_thread_ids.append(threading.get_ident())
|
||||||
|
original_request_timer_rearm()
|
||||||
|
|
||||||
|
def tracked_arm_timer() -> None:
|
||||||
|
arm_thread_ids.append(threading.get_ident())
|
||||||
|
original_arm_timer()
|
||||||
|
timer_rearmed.set()
|
||||||
|
|
||||||
|
monkeypatch.setattr(cron, "_request_timer_rearm", tracked_request_timer_rearm)
|
||||||
|
monkeypatch.setattr(cron, "_arm_timer", tracked_arm_timer)
|
||||||
|
channel = _ch(bus, cron_service=cron, port=_free_port())
|
||||||
|
|
||||||
|
try:
|
||||||
|
response = await _webui_mutate(
|
||||||
|
channel,
|
||||||
|
"automation.update",
|
||||||
|
{"id": job.id, "values": {"name": "After update"}},
|
||||||
|
)
|
||||||
|
await asyncio.wait_for(timer_rearmed.wait(), timeout=1)
|
||||||
|
|
||||||
|
assert response.status_code == 200
|
||||||
|
assert request_thread_ids
|
||||||
|
assert all(thread_id != owner_thread_id for thread_id in request_thread_ids)
|
||||||
|
assert arm_thread_ids and set(arm_thread_ids) == {owner_thread_id}
|
||||||
|
assert cron._timer_task is not None
|
||||||
|
assert cron._timer_task is not initial_timer
|
||||||
|
assert not cron._timer_task.done()
|
||||||
|
|
||||||
|
stored = json.loads(store_path.read_text(encoding="utf-8"))
|
||||||
|
assert len(stored["jobs"]) == 1
|
||||||
|
assert stored["jobs"][0]["name"] == "After update"
|
||||||
|
finally:
|
||||||
|
cron.stop()
|
||||||
|
await asyncio.sleep(0)
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_webui_automations_route_manages_local_triggers(
|
async def test_webui_automations_route_manages_local_triggers(
|
||||||
bus: MagicMock, tmp_path: Path
|
bus: MagicMock, tmp_path: Path
|
||||||
@@ -3275,6 +3310,28 @@ async def _webui_mutate(
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_recovery_mutation_uses_authenticated_websocket_action(bus: MagicMock) -> None:
|
||||||
|
recovery_action = AsyncMock(return_value={
|
||||||
|
"status": "resuming",
|
||||||
|
"recovery_id": "recovery-1",
|
||||||
|
})
|
||||||
|
channel = _ch(bus, recovery_action=recovery_action)
|
||||||
|
|
||||||
|
response = await _webui_mutate(
|
||||||
|
channel,
|
||||||
|
"recovery.continue",
|
||||||
|
{"chat_id": "chat-1", "recovery_id": "recovery-1"},
|
||||||
|
)
|
||||||
|
|
||||||
|
assert response.status_code == 200
|
||||||
|
assert response.json()["status"] == "resuming"
|
||||||
|
recovery_action.assert_awaited_once_with(
|
||||||
|
"continue",
|
||||||
|
{"chat_id": "chat-1", "recovery_id": "recovery-1"},
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_workspace_folder_picker_is_local_authenticated_mutation(
|
async def test_workspace_folder_picker_is_local_authenticated_mutation(
|
||||||
bus: MagicMock,
|
bus: MagicMock,
|
||||||
@@ -3776,3 +3833,77 @@ def test_bootstrap_secret_also_enforced_on_localhost(bus: MagicMock) -> None:
|
|||||||
channel = _ch(bus, host="0.0.0.0", tokenIssueSecret="s3cret")
|
channel = _ch(bus, host="0.0.0.0", tokenIssueSecret="s3cret")
|
||||||
resp = channel.gateway.http._handle_bootstrap(_LOCAL, _NO_HEADERS)
|
resp = channel.gateway.http._handle_bootstrap(_LOCAL, _NO_HEADERS)
|
||||||
assert resp.status_code == 401
|
assert resp.status_code == 401
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_webui_skill_update_cancellation_waits_for_config_and_runtime_state(
|
||||||
|
bus: MagicMock,
|
||||||
|
tmp_path: Path,
|
||||||
|
monkeypatch: pytest.MonkeyPatch,
|
||||||
|
) -> None:
|
||||||
|
from nanobot.webui import ws_http
|
||||||
|
|
||||||
|
skill_dir = tmp_path / "skills" / "cancel-safe-skill"
|
||||||
|
skill_dir.mkdir(parents=True)
|
||||||
|
(skill_dir / "SKILL.md").write_text(
|
||||||
|
"---\nname: cancel-safe-skill\ndescription: Cancellation test skill.\n---\n",
|
||||||
|
encoding="utf-8",
|
||||||
|
)
|
||||||
|
mutation_started = threading.Event()
|
||||||
|
release_mutation = threading.Event()
|
||||||
|
original_update = ws_http.set_webui_skill_enabled
|
||||||
|
update_calls = 0
|
||||||
|
|
||||||
|
def blocked_update(*args: Any, **kwargs: Any) -> dict[str, Any]:
|
||||||
|
nonlocal update_calls
|
||||||
|
update_calls += 1
|
||||||
|
mutation_started.set()
|
||||||
|
assert release_mutation.wait(timeout=1)
|
||||||
|
return original_update(*args, **kwargs)
|
||||||
|
|
||||||
|
monkeypatch.setattr(ws_http, "set_webui_skill_enabled", blocked_update)
|
||||||
|
channel = _ch(
|
||||||
|
bus,
|
||||||
|
session_manager=_seed_session(tmp_path),
|
||||||
|
workspace_path=tmp_path,
|
||||||
|
port=_free_port(),
|
||||||
|
)
|
||||||
|
runtime_states: list[set[str]] = []
|
||||||
|
channel.gateway.http.skill_state_action = runtime_states.append
|
||||||
|
task = asyncio.create_task(
|
||||||
|
_webui_mutate(
|
||||||
|
channel,
|
||||||
|
"skill.update",
|
||||||
|
{"name": "cancel-safe-skill", "enabled": False},
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
assert await asyncio.to_thread(mutation_started.wait, 1)
|
||||||
|
try:
|
||||||
|
task.cancel()
|
||||||
|
await asyncio.sleep(0)
|
||||||
|
assert not task.done()
|
||||||
|
assert runtime_states == []
|
||||||
|
finally:
|
||||||
|
release_mutation.set()
|
||||||
|
|
||||||
|
with pytest.raises(asyncio.CancelledError):
|
||||||
|
await asyncio.wait_for(task, timeout=1)
|
||||||
|
|
||||||
|
assert update_calls == 1
|
||||||
|
assert "cancel-safe-skill" in channel.gateway.http.disabled_skills
|
||||||
|
assert runtime_states == [{"cancel-safe-skill"}]
|
||||||
|
saved = load_config(channel.gateway.settings.config.path)
|
||||||
|
assert "cancel-safe-skill" in saved.agents.defaults.disabled_skills
|
||||||
|
|
||||||
|
settled_state = (
|
||||||
|
update_calls,
|
||||||
|
set(channel.gateway.http.disabled_skills),
|
||||||
|
list(runtime_states),
|
||||||
|
)
|
||||||
|
await asyncio.sleep(0.05)
|
||||||
|
assert (
|
||||||
|
update_calls,
|
||||||
|
set(channel.gateway.http.disabled_skills),
|
||||||
|
runtime_states,
|
||||||
|
) == settled_state
|
||||||
|
|||||||
@@ -202,12 +202,6 @@ class WsTestClient:
|
|||||||
assert msg.event == "delta", f"Expected 'delta' event, got '{msg.event}'"
|
assert msg.event == "delta", f"Expected 'delta' event, got '{msg.event}'"
|
||||||
return msg
|
return msg
|
||||||
|
|
||||||
async def recv_stream_end(self, timeout: float = 10.0) -> WsMessage:
|
|
||||||
"""Receive and validate a 'stream_end' event."""
|
|
||||||
msg = await self.recv(timeout)
|
|
||||||
assert msg.event == "stream_end", f"Expected 'stream_end' event, got '{msg.event}'"
|
|
||||||
return msg
|
|
||||||
|
|
||||||
async def collect_stream(self, timeout: float = 10.0) -> list[WsMessage]:
|
async def collect_stream(self, timeout: float = 10.0) -> list[WsMessage]:
|
||||||
"""Collect all deltas and the final stream_end into a list."""
|
"""Collect all deltas and the final stream_end into a list."""
|
||||||
messages: list[WsMessage] = []
|
messages: list[WsMessage] = []
|
||||||
@@ -232,10 +226,6 @@ class WsTestClient:
|
|||||||
"""Send a JSON frame."""
|
"""Send a JSON frame."""
|
||||||
await self.ws.send(json.dumps(data, ensure_ascii=False))
|
await self.ws.send(json.dumps(data, ensure_ascii=False))
|
||||||
|
|
||||||
async def send_content(self, content: str) -> None:
|
|
||||||
"""Send content in the preferred JSON format ``{"content": ...}``."""
|
|
||||||
await self.send_json({"content": content})
|
|
||||||
|
|
||||||
# -- Connection introspection -----------------------------------------
|
# -- Connection introspection -----------------------------------------
|
||||||
|
|
||||||
@property
|
@property
|
||||||
|
|||||||
@@ -56,6 +56,7 @@ from nanobot.cli.agent import agent # noqa: E402
|
|||||||
from nanobot.cli.gateway import create_gateway_app # noqa: E402
|
from nanobot.cli.gateway import create_gateway_app # noqa: E402
|
||||||
from nanobot.cli.gateway_runtime import _run_gateway # noqa: E402
|
from nanobot.cli.gateway_runtime import _run_gateway # noqa: E402
|
||||||
from nanobot.cli.log_control import _set_nanobot_logs # noqa: E402
|
from nanobot.cli.log_control import _set_nanobot_logs # noqa: E402
|
||||||
|
from nanobot.cli.process_identity import set_cli_process_identity # noqa: E402
|
||||||
from nanobot.cli.provider import provider_app # noqa: E402
|
from nanobot.cli.provider import provider_app # noqa: E402
|
||||||
from nanobot.cli.runtime_config import ( # noqa: E402
|
from nanobot.cli.runtime_config import ( # noqa: E402
|
||||||
_load_inspection_config,
|
_load_inspection_config,
|
||||||
@@ -99,12 +100,17 @@ def version_callback(value: bool):
|
|||||||
|
|
||||||
@app.callback()
|
@app.callback()
|
||||||
def main(
|
def main(
|
||||||
|
ctx: typer.Context,
|
||||||
version: bool = typer.Option(
|
version: bool = typer.Option(
|
||||||
None, "--version", "-v", callback=version_callback, is_eager=True
|
None, "--version", "-v", callback=version_callback, is_eager=True
|
||||||
),
|
),
|
||||||
):
|
):
|
||||||
"""nanobot - Personal AI Assistant."""
|
"""nanobot - Personal AI Assistant."""
|
||||||
pass
|
# Editable/source installs can retain an older generated console script that
|
||||||
|
# imports this Typer app directly instead of ``nanobot.cli.entry``. Keep the
|
||||||
|
# role identity correct until that launcher is regenerated.
|
||||||
|
command = ctx.invoked_subcommand
|
||||||
|
set_cli_process_identity([command] if command else sys.argv[1:])
|
||||||
|
|
||||||
|
|
||||||
# ============================================================================
|
# ============================================================================
|
||||||
|
|||||||
@@ -6,6 +6,8 @@ import os
|
|||||||
import sys
|
import sys
|
||||||
from contextlib import suppress
|
from contextlib import suppress
|
||||||
|
|
||||||
|
from nanobot.cli.process_identity import set_cli_process_identity
|
||||||
|
|
||||||
|
|
||||||
def _native_tui_candidate(args: list[str]) -> bool:
|
def _native_tui_candidate(args: list[str]) -> bool:
|
||||||
"""Return whether ``agent`` can start without the classic agent stack."""
|
"""Return whether ``agent`` can start without the classic agent stack."""
|
||||||
@@ -34,6 +36,7 @@ def _configure_windows_console() -> None:
|
|||||||
|
|
||||||
def main() -> None:
|
def main() -> None:
|
||||||
"""Dispatch native TUI startup without importing the complete CLI graph."""
|
"""Dispatch native TUI startup without importing the complete CLI graph."""
|
||||||
|
set_cli_process_identity(sys.argv[1:])
|
||||||
_configure_windows_console()
|
_configure_windows_console()
|
||||||
if _native_tui_candidate(sys.argv[1:]):
|
if _native_tui_candidate(sys.argv[1:]):
|
||||||
import typer
|
import typer
|
||||||
|
|||||||
+117
-36
@@ -34,6 +34,7 @@ from nanobot.config.paths import is_default_workspace
|
|||||||
from nanobot.config.schema import Config
|
from nanobot.config.schema import Config
|
||||||
from nanobot.gateway.runtime import GatewayInstance
|
from nanobot.gateway.runtime import GatewayInstance
|
||||||
from nanobot.security.network import is_loopback_host
|
from nanobot.security.network import is_loopback_host
|
||||||
|
from nanobot.session.async_compat import call_session_manager as _call_session_manager
|
||||||
from nanobot.session.keys import UNIFIED_SESSION_KEY, last_channel_from_metadata
|
from nanobot.session.keys import UNIFIED_SESSION_KEY, last_channel_from_metadata
|
||||||
from nanobot.utils.evaluator import evaluate_response, resolve_evaluator_prompt
|
from nanobot.utils.evaluator import evaluate_response, resolve_evaluator_prompt
|
||||||
from nanobot.utils.helpers import sync_workspace_templates
|
from nanobot.utils.helpers import sync_workspace_templates
|
||||||
@@ -45,6 +46,30 @@ __all__ = ["_run_gateway"]
|
|||||||
|
|
||||||
console = Console()
|
console = Console()
|
||||||
|
|
||||||
|
_EVENT_LOOP_LAG_INTERVAL_S = 0.5
|
||||||
|
_EVENT_LOOP_LAG_WARNING_S = 0.25
|
||||||
|
|
||||||
|
|
||||||
|
async def _monitor_event_loop_lag(
|
||||||
|
*,
|
||||||
|
interval_s: float = _EVENT_LOOP_LAG_INTERVAL_S,
|
||||||
|
warning_threshold_s: float = _EVENT_LOOP_LAG_WARNING_S,
|
||||||
|
log: Any | None = None,
|
||||||
|
) -> None:
|
||||||
|
"""Log scheduler drift so gateway-wide stalls have direct evidence."""
|
||||||
|
loop = asyncio.get_running_loop()
|
||||||
|
lag_log = log or logger
|
||||||
|
while True:
|
||||||
|
expected = loop.time() + interval_s
|
||||||
|
await asyncio.sleep(interval_s)
|
||||||
|
lag_s = max(0.0, loop.time() - expected)
|
||||||
|
if lag_s >= warning_threshold_s:
|
||||||
|
lag_log.warning(
|
||||||
|
"event loop lag operation=gateway duration_ms={} interval_ms={}",
|
||||||
|
int(lag_s * 1000),
|
||||||
|
int(interval_s * 1000),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
def _http_endpoint_responding(url: str, *, timeout_s: float = 0.25) -> bool:
|
def _http_endpoint_responding(url: str, *, timeout_s: float = 0.25) -> bool:
|
||||||
"""Return whether an HTTP endpoint responds, including with an auth error."""
|
"""Return whether an HTTP endpoint responds, including with an auth error."""
|
||||||
@@ -313,6 +338,8 @@ def _run_gateway(
|
|||||||
from nanobot.cron.service import CronJobSkippedError, CronService
|
from nanobot.cron.service import CronJobSkippedError, CronService
|
||||||
from nanobot.cron.session_turns import is_bound_cron_job
|
from nanobot.cron.session_turns import is_bound_cron_job
|
||||||
from nanobot.cron.types import CronJob
|
from nanobot.cron.types import CronJob
|
||||||
|
from nanobot.llm_usage import record_llm_call
|
||||||
|
from nanobot.llm_usage.context import llm_usage_source
|
||||||
from nanobot.providers.factory import (
|
from nanobot.providers.factory import (
|
||||||
ProviderSnapshot,
|
ProviderSnapshot,
|
||||||
build_provider_snapshot,
|
build_provider_snapshot,
|
||||||
@@ -322,6 +349,7 @@ def _run_gateway(
|
|||||||
from nanobot.providers.fallback_provider import FallbackProvider
|
from nanobot.providers.fallback_provider import FallbackProvider
|
||||||
from nanobot.providers.image_generation import image_gen_provider_configs
|
from nanobot.providers.image_generation import image_gen_provider_configs
|
||||||
from nanobot.session.manager import SessionManager
|
from nanobot.session.manager import SessionManager
|
||||||
|
from nanobot.session.recovery import RecoveryCoordinator
|
||||||
from nanobot.session.webui_turns import (
|
from nanobot.session.webui_turns import (
|
||||||
WebuiTurnCoordinator,
|
WebuiTurnCoordinator,
|
||||||
WebuiTurnRoutePolicy,
|
WebuiTurnRoutePolicy,
|
||||||
@@ -329,7 +357,6 @@ def _run_gateway(
|
|||||||
)
|
)
|
||||||
from nanobot.triggers.local_runner import run_local_trigger_queue
|
from nanobot.triggers.local_runner import run_local_trigger_queue
|
||||||
from nanobot.triggers.local_store import LocalTriggerStore
|
from nanobot.triggers.local_store import LocalTriggerStore
|
||||||
from nanobot.webui.token_usage import TokenUsageHook
|
|
||||||
|
|
||||||
port = port if port is not None else config.gateway.port
|
port = port if port is not None else config.gateway.port
|
||||||
webui_url = _webui_browser_url(config)
|
webui_url = _webui_browser_url(config)
|
||||||
@@ -360,7 +387,8 @@ def _run_gateway(
|
|||||||
runtime_events = RuntimeEventBus()
|
runtime_events = RuntimeEventBus()
|
||||||
fallback_model_observer = build_webui_fallback_model_observer(bus)
|
fallback_model_observer = build_webui_fallback_model_observer(bus)
|
||||||
|
|
||||||
def _observe_fallback_models(snapshot: ProviderSnapshot) -> ProviderSnapshot:
|
def _observe_provider(snapshot: ProviderSnapshot) -> ProviderSnapshot:
|
||||||
|
snapshot.provider.set_llm_call_observer(record_llm_call)
|
||||||
if isinstance(snapshot.provider, FallbackProvider):
|
if isinstance(snapshot.provider, FallbackProvider):
|
||||||
snapshot.provider.set_fallback_model_observer(fallback_model_observer)
|
snapshot.provider.set_fallback_model_observer(fallback_model_observer)
|
||||||
return snapshot
|
return snapshot
|
||||||
@@ -370,20 +398,19 @@ def _run_gateway(
|
|||||||
**kwargs: Any,
|
**kwargs: Any,
|
||||||
) -> ProviderSnapshot:
|
) -> ProviderSnapshot:
|
||||||
try:
|
try:
|
||||||
return _observe_fallback_models(load_provider_snapshot(*args, **kwargs))
|
return _observe_provider(load_provider_snapshot(*args, **kwargs))
|
||||||
except ValueError as exc:
|
except ValueError as exc:
|
||||||
if unconfigured_provider_error is None:
|
if unconfigured_provider_error is None:
|
||||||
raise
|
raise
|
||||||
return build_unconfigured_provider_snapshot(config, str(exc))
|
return _observe_provider(build_unconfigured_provider_snapshot(config, str(exc)))
|
||||||
|
|
||||||
if unconfigured_provider_error is not None:
|
if unconfigured_provider_error is not None:
|
||||||
provider_snapshot = build_unconfigured_provider_snapshot(
|
provider_snapshot = _observe_provider(
|
||||||
config,
|
build_unconfigured_provider_snapshot(config, unconfigured_provider_error)
|
||||||
unconfigured_provider_error,
|
|
||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
try:
|
try:
|
||||||
provider_snapshot = _observe_fallback_models(build_provider_snapshot(config))
|
provider_snapshot = _observe_provider(build_provider_snapshot(config))
|
||||||
except ValueError as exc:
|
except ValueError as exc:
|
||||||
console.print(f"[red]Error: {exc}[/red]")
|
console.print(f"[red]Error: {exc}[/red]")
|
||||||
raise typer.Exit(1) from exc
|
raise typer.Exit(1) from exc
|
||||||
@@ -422,6 +449,12 @@ def _run_gateway(
|
|||||||
tools = ToolRegistry()
|
tools = ToolRegistry()
|
||||||
mcp_provider = MCPProvider.from_config(config, tools)
|
mcp_provider = MCPProvider.from_config(config, tools)
|
||||||
|
|
||||||
|
recovery = RecoveryCoordinator(
|
||||||
|
sessions=session_manager,
|
||||||
|
bus=bus,
|
||||||
|
unified_session=config.agents.defaults.unified_session,
|
||||||
|
)
|
||||||
|
|
||||||
# Create agent with cron service
|
# Create agent with cron service
|
||||||
agent = AgentLoop.from_config(
|
agent = AgentLoop.from_config(
|
||||||
config, bus,
|
config, bus,
|
||||||
@@ -436,10 +469,10 @@ def _run_gateway(
|
|||||||
runtime_events=runtime_events,
|
runtime_events=runtime_events,
|
||||||
turn_delivery_factory=turn_delivery_factory,
|
turn_delivery_factory=turn_delivery_factory,
|
||||||
provider_signature=provider_snapshot.signature,
|
provider_signature=provider_snapshot.signature,
|
||||||
hooks=[TokenUsageHook(timezone_name=config.agents.defaults.timezone)],
|
|
||||||
local_trigger_store=trigger_store,
|
local_trigger_store=trigger_store,
|
||||||
hook_factories=[create_file_edit_activity_hook],
|
hook_factories=[create_file_edit_activity_hook],
|
||||||
tool_registry=tools,
|
tool_registry=tools,
|
||||||
|
recovery_admission=recovery,
|
||||||
)
|
)
|
||||||
def _schedule_webui_background(awaitable: Awaitable[None]) -> None:
|
def _schedule_webui_background(awaitable: Awaitable[None]) -> None:
|
||||||
agent.schedule_background(cast(Coroutine[Any, Any, None], awaitable))
|
agent.schedule_background(cast(Coroutine[Any, Any, None], awaitable))
|
||||||
@@ -448,6 +481,7 @@ def _run_gateway(
|
|||||||
bus=bus,
|
bus=bus,
|
||||||
sessions=session_manager,
|
sessions=session_manager,
|
||||||
schedule_background=_schedule_webui_background,
|
schedule_background=_schedule_webui_background,
|
||||||
|
recovery=recovery,
|
||||||
)
|
)
|
||||||
webui_turn_coordinator.subscribe(runtime_events)
|
webui_turn_coordinator.subscribe(runtime_events)
|
||||||
from nanobot.bus.events import OutboundMessage
|
from nanobot.bus.events import OutboundMessage
|
||||||
@@ -484,12 +518,22 @@ def _run_gateway(
|
|||||||
and hasattr(session_manager, "save")
|
and hasattr(session_manager, "save")
|
||||||
):
|
):
|
||||||
key = session_key or _channel_session_key(msg.channel, msg.chat_id)
|
key = session_key or _channel_session_key(msg.channel, msg.chat_id)
|
||||||
session = session_manager.get_or_create(key)
|
session = await _call_session_manager(
|
||||||
|
session_manager,
|
||||||
|
"get_or_create_async",
|
||||||
|
session_manager.get_or_create,
|
||||||
|
key,
|
||||||
|
)
|
||||||
extra: dict[str, Any] = {"_channel_delivery": True}
|
extra: dict[str, Any] = {"_channel_delivery": True}
|
||||||
if msg.media:
|
if msg.media:
|
||||||
extra["media"] = list(msg.media)
|
extra["media"] = list(msg.media)
|
||||||
session.add_message("assistant", msg.content, **extra)
|
session.add_message("assistant", msg.content, **extra)
|
||||||
session_manager.save(session)
|
await _call_session_manager(
|
||||||
|
session_manager,
|
||||||
|
"save_async",
|
||||||
|
session_manager.save,
|
||||||
|
session,
|
||||||
|
)
|
||||||
await bus.publish_outbound(msg)
|
await bus.publish_outbound(msg)
|
||||||
|
|
||||||
message_tool = agent.tools.get("message")
|
message_tool = agent.tools.get("message")
|
||||||
@@ -555,25 +599,18 @@ def _run_gateway(
|
|||||||
except Exception:
|
except Exception:
|
||||||
logger.exception("Dream cron job failed")
|
logger.exception("Dream cron job failed")
|
||||||
finally:
|
finally:
|
||||||
from nanobot.webui.token_usage import record_response_token_usage
|
|
||||||
|
|
||||||
record_response_token_usage(
|
|
||||||
resp,
|
|
||||||
source="dream",
|
|
||||||
timezone_name=config.agents.defaults.timezone,
|
|
||||||
)
|
|
||||||
sha = _commit_dream_changes(store)
|
sha = _commit_dream_changes(store)
|
||||||
if sha:
|
if sha:
|
||||||
logger.info("Dream commit: {}", sha)
|
logger.info("Dream commit: {}", sha)
|
||||||
store.compact_history()
|
store.compact_history()
|
||||||
prune_dream_sessions(agent.sessions)
|
await asyncio.to_thread(prune_dream_sessions, agent.sessions)
|
||||||
return None
|
return None
|
||||||
|
|
||||||
# Heartbeat is a system job that checks HEARTBEAT.md for active tasks.
|
# Heartbeat is a system job that checks HEARTBEAT.md for active tasks.
|
||||||
if job.name == "heartbeat":
|
if job.name == "heartbeat":
|
||||||
heartbeat_file = config.workspace_path / "HEARTBEAT.md"
|
heartbeat_file = config.workspace_path / "HEARTBEAT.md"
|
||||||
try:
|
try:
|
||||||
content = heartbeat_file.read_text(encoding="utf-8")
|
content = await asyncio.to_thread(heartbeat_file.read_text, encoding="utf-8")
|
||||||
except OSError:
|
except OSError:
|
||||||
logger.debug("Heartbeat: HEARTBEAT.md missing")
|
logger.debug("Heartbeat: HEARTBEAT.md missing")
|
||||||
return None
|
return None
|
||||||
@@ -581,7 +618,7 @@ def _run_gateway(
|
|||||||
logger.debug("Heartbeat: HEARTBEAT.md has no active tasks")
|
logger.debug("Heartbeat: HEARTBEAT.md has no active tasks")
|
||||||
return None
|
return None
|
||||||
|
|
||||||
channel, chat_id = _pick_heartbeat_target()
|
channel, chat_id = await _pick_heartbeat_target()
|
||||||
if channel == "cli":
|
if channel == "cli":
|
||||||
return None
|
return None
|
||||||
|
|
||||||
@@ -609,9 +646,19 @@ def _run_gateway(
|
|||||||
message_tool.reset_suppress_delivery(suppress_token)
|
message_tool.reset_suppress_delivery(suppress_token)
|
||||||
|
|
||||||
# Keep a small tail of heartbeat history so the loop stays bounded.
|
# Keep a small tail of heartbeat history so the loop stays bounded.
|
||||||
session = agent.sessions.get_or_create("heartbeat")
|
session = await _call_session_manager(
|
||||||
|
agent.sessions,
|
||||||
|
"get_or_create_async",
|
||||||
|
agent.sessions.get_or_create,
|
||||||
|
"heartbeat",
|
||||||
|
)
|
||||||
session.retain_recent_legal_suffix(hb_cfg.keep_recent_messages)
|
session.retain_recent_legal_suffix(hb_cfg.keep_recent_messages)
|
||||||
agent.sessions.save(session)
|
await _call_session_manager(
|
||||||
|
agent.sessions,
|
||||||
|
"save_async",
|
||||||
|
agent.sessions.save,
|
||||||
|
session,
|
||||||
|
)
|
||||||
|
|
||||||
if not resp or not resp.content:
|
if not resp or not resp.content:
|
||||||
return
|
return
|
||||||
@@ -621,14 +668,15 @@ def _run_gateway(
|
|||||||
evaluator_prompt = resolve_evaluator_prompt(config.workspace_path)
|
evaluator_prompt = resolve_evaluator_prompt(config.workspace_path)
|
||||||
|
|
||||||
# Fail closed: stay silent on evaluator failure instead of notifying.
|
# Fail closed: stay silent on evaluator failure instead of notifying.
|
||||||
should_notify = await evaluate_response(
|
with llm_usage_source("cron"):
|
||||||
response=response,
|
should_notify = await evaluate_response(
|
||||||
task_context=prompt,
|
response=response,
|
||||||
provider=agent.provider,
|
task_context=prompt,
|
||||||
model=agent.model,
|
provider=agent.provider,
|
||||||
evaluator_prompt=evaluator_prompt,
|
model=agent.model,
|
||||||
default_notify=False,
|
evaluator_prompt=evaluator_prompt,
|
||||||
)
|
default_notify=False,
|
||||||
|
)
|
||||||
|
|
||||||
if should_notify:
|
if should_notify:
|
||||||
logger.info("Heartbeat: completed, delivering response")
|
logger.info("Heartbeat: completed, delivering response")
|
||||||
@@ -683,20 +731,31 @@ def _run_gateway(
|
|||||||
webui_mcp_runtime_status=mcp_provider.runtime_status,
|
webui_mcp_runtime_status=mcp_provider.runtime_status,
|
||||||
webui_mcp_reload=mcp_provider.reload,
|
webui_mcp_reload=mcp_provider.reload,
|
||||||
webui_skill_state_action=_webui_skill_state_action,
|
webui_skill_state_action=_webui_skill_state_action,
|
||||||
|
webui_recovery_action=recovery.handle_action,
|
||||||
config_path=Path(config_path),
|
config_path=Path(config_path),
|
||||||
)
|
)
|
||||||
|
|
||||||
def _pick_heartbeat_target() -> tuple[str, str]:
|
async def _pick_heartbeat_target() -> tuple[str, str]:
|
||||||
"""Pick a routable channel/chat target for heartbeat-triggered messages."""
|
"""Pick a routable channel/chat target for heartbeat-triggered messages."""
|
||||||
sidebar_state = read_webui_sidebar_state()
|
sidebar_state = await asyncio.to_thread(read_webui_sidebar_state)
|
||||||
unified_metadata = None
|
unified_metadata = None
|
||||||
if config.agents.defaults.unified_session:
|
if config.agents.defaults.unified_session:
|
||||||
record = session_manager.read_session_metadata(UNIFIED_SESSION_KEY)
|
record = await _call_session_manager(
|
||||||
|
session_manager,
|
||||||
|
"read_session_metadata_async",
|
||||||
|
session_manager.read_session_metadata,
|
||||||
|
UNIFIED_SESSION_KEY,
|
||||||
|
)
|
||||||
if isinstance(record, dict) and isinstance(record.get("metadata"), dict):
|
if isinstance(record, dict) and isinstance(record.get("metadata"), dict):
|
||||||
unified_metadata = record["metadata"]
|
unified_metadata = record["metadata"]
|
||||||
|
sessions = await _call_session_manager(
|
||||||
|
session_manager,
|
||||||
|
"list_sessions_async",
|
||||||
|
session_manager.list_sessions,
|
||||||
|
)
|
||||||
return _pick_heartbeat_target_from_sessions(
|
return _pick_heartbeat_target_from_sessions(
|
||||||
enabled_channels=channels.enabled_channels,
|
enabled_channels=channels.enabled_channels,
|
||||||
sessions=session_manager.list_sessions(),
|
sessions=sessions,
|
||||||
archived_keys=sidebar_state.get("archived_keys", []),
|
archived_keys=sidebar_state.get("archived_keys", []),
|
||||||
unified_session_metadata=unified_metadata,
|
unified_session_metadata=unified_metadata,
|
||||||
)
|
)
|
||||||
@@ -849,6 +908,7 @@ def _run_gateway(
|
|||||||
tasks: list[asyncio.Task[Any]] = []
|
tasks: list[asyncio.Task[Any]] = []
|
||||||
shutdown_task: asyncio.Task[Any] | None = None
|
shutdown_task: asyncio.Task[Any] | None = None
|
||||||
runtime_tasks: asyncio.Future[list[Any]] | None = None
|
runtime_tasks: asyncio.Future[list[Any]] | None = None
|
||||||
|
startup_complete = False
|
||||||
shutdown_event = asyncio.Event()
|
shutdown_event = asyncio.Event()
|
||||||
cli_terminal._ensure_interactive_tty_mode()
|
cli_terminal._ensure_interactive_tty_mode()
|
||||||
restore_shutdown_handlers = _install_gateway_shutdown_handlers(
|
restore_shutdown_handlers = _install_gateway_shutdown_handlers(
|
||||||
@@ -861,6 +921,10 @@ def _run_gateway(
|
|||||||
await cron.start()
|
await cron.start()
|
||||||
# Re-read once on first admission to close the watcher subscription window.
|
# Re-read once on first admission to close the watcher subscription window.
|
||||||
agent.runtime_resolver.invalidate()
|
agent.runtime_resolver.invalidate()
|
||||||
|
# Recovery must finish before WebSocket and other channels begin
|
||||||
|
# accepting new input. That makes a new user message reliably
|
||||||
|
# supersede an old recoverable turn instead of racing its queue.
|
||||||
|
await recovery.scan()
|
||||||
async def _run_agent() -> None:
|
async def _run_agent() -> None:
|
||||||
try:
|
try:
|
||||||
await mcp_provider.connect()
|
await mcp_provider.connect()
|
||||||
@@ -898,6 +962,10 @@ def _run_gateway(
|
|||||||
_monitor_local_clients(),
|
_monitor_local_clients(),
|
||||||
name="nanobot-gateway-client-monitor",
|
name="nanobot-gateway-client-monitor",
|
||||||
),
|
),
|
||||||
|
asyncio.create_task(
|
||||||
|
_monitor_event_loop_lag(),
|
||||||
|
name="nanobot-event-loop-lag-monitor",
|
||||||
|
),
|
||||||
]
|
]
|
||||||
if health_server_enabled:
|
if health_server_enabled:
|
||||||
tasks.append(asyncio.create_task(
|
tasks.append(asyncio.create_task(
|
||||||
@@ -915,6 +983,7 @@ def _run_gateway(
|
|||||||
name="nanobot-webui-dev-server",
|
name="nanobot-webui-dev-server",
|
||||||
))
|
))
|
||||||
runtime_tasks = asyncio.gather(*tasks)
|
runtime_tasks = asyncio.gather(*tasks)
|
||||||
|
startup_complete = True
|
||||||
shutdown_task = asyncio.create_task(
|
shutdown_task = asyncio.create_task(
|
||||||
shutdown_event.wait(),
|
shutdown_event.wait(),
|
||||||
name="nanobot-gateway-shutdown",
|
name="nanobot-gateway-shutdown",
|
||||||
@@ -936,6 +1005,10 @@ def _run_gateway(
|
|||||||
|
|
||||||
console.print("\n[red]Error: Gateway crashed unexpectedly[/red]")
|
console.print("\n[red]Error: Gateway crashed unexpectedly[/red]")
|
||||||
console.print(traceback.format_exc())
|
console.print(traceback.format_exc())
|
||||||
|
if not startup_complete:
|
||||||
|
# Do not report a successful gateway command when startup
|
||||||
|
# failed before any runtime task or listener was created.
|
||||||
|
raise typer.Exit(1)
|
||||||
finally:
|
finally:
|
||||||
try:
|
try:
|
||||||
if shutdown_task and not shutdown_task.done():
|
if shutdown_task and not shutdown_task.done():
|
||||||
@@ -943,6 +1016,10 @@ def _run_gateway(
|
|||||||
with suppress(asyncio.CancelledError):
|
with suppress(asyncio.CancelledError):
|
||||||
await shutdown_task
|
await shutdown_task
|
||||||
cron.stop()
|
cron.stop()
|
||||||
|
# A gateway exit interrupts ownership of active turns; it is
|
||||||
|
# not the same as the user stopping a turn. Keep checkpoints
|
||||||
|
# so the next gateway can offer an explicit Continue action.
|
||||||
|
agent.preserve_inflight_turns_on_shutdown()
|
||||||
agent.stop()
|
agent.stop()
|
||||||
# Cancel runtime tasks first, then deterministically close
|
# Cancel runtime tasks first, then deterministically close
|
||||||
# exec/MCP resources while the event loop is still alive.
|
# exec/MCP resources while the event loop is still alive.
|
||||||
@@ -956,7 +1033,11 @@ def _run_gateway(
|
|||||||
# Flush all cached sessions to durable storage before exit.
|
# Flush all cached sessions to durable storage before exit.
|
||||||
# This prevents data loss on filesystems with write-back
|
# This prevents data loss on filesystems with write-back
|
||||||
# caching (rclone VFS, NFS, FUSE mounts, etc.).
|
# caching (rclone VFS, NFS, FUSE mounts, etc.).
|
||||||
flushed = agent.sessions.flush_all()
|
flushed = await _call_session_manager(
|
||||||
|
agent.sessions,
|
||||||
|
"flush_all_async",
|
||||||
|
agent.sessions.flush_all,
|
||||||
|
)
|
||||||
if flushed:
|
if flushed:
|
||||||
logger.info("Shutdown: flushed {} session(s) to disk", flushed)
|
logger.info("Shutdown: flushed {} session(s) to disk", flushed)
|
||||||
finally:
|
finally:
|
||||||
|
|||||||
@@ -0,0 +1,52 @@
|
|||||||
|
"""Give nanobot processes recognizable operating-system names."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import hashlib
|
||||||
|
import os
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import Final
|
||||||
|
|
||||||
|
_ROLES: Final = {"agent", "gateway", "webui"}
|
||||||
|
|
||||||
|
|
||||||
|
def _set_process_title(title: str) -> None:
|
||||||
|
# Process titles are short; do not trade Linux /proc environment visibility for
|
||||||
|
# extra title storage. setproctitle reads this switch when it is imported.
|
||||||
|
os.environ.setdefault("SPT_NOENV", "1")
|
||||||
|
from setproctitle import setproctitle
|
||||||
|
|
||||||
|
setproctitle(title)
|
||||||
|
|
||||||
|
|
||||||
|
def set_cli_process_identity(args: list[str]) -> None:
|
||||||
|
"""Name this CLI process after the nanobot role it is running."""
|
||||||
|
if os.name == "nt":
|
||||||
|
# Windows process managers use the console launcher's executable name,
|
||||||
|
# which packaging already generates as ``nanobot.exe``.
|
||||||
|
return
|
||||||
|
role = args[0] if args and args[0] in _ROLES else None
|
||||||
|
_set_process_title(f"nanobot-{role}" if role else "nanobot")
|
||||||
|
|
||||||
|
|
||||||
|
def named_executable(executable: str, *, name: str, directory: Path) -> str:
|
||||||
|
"""Return a stable POSIX symlink whose basename identifies a child process."""
|
||||||
|
if os.name == "nt":
|
||||||
|
return executable
|
||||||
|
try:
|
||||||
|
target = Path(executable).resolve(strict=True)
|
||||||
|
digest = hashlib.sha256(os.fsencode(target)).hexdigest()[:12]
|
||||||
|
link_dir = directory / digest
|
||||||
|
link = link_dir / name
|
||||||
|
link_dir.mkdir(mode=0o700, parents=True, exist_ok=True)
|
||||||
|
if link.is_symlink() and link.resolve(strict=False) == target:
|
||||||
|
return str(link)
|
||||||
|
if link.exists():
|
||||||
|
return executable
|
||||||
|
pending = link.with_name(f".{name}.{os.getpid()}")
|
||||||
|
pending.unlink(missing_ok=True)
|
||||||
|
pending.symlink_to(target)
|
||||||
|
os.replace(pending, link)
|
||||||
|
except OSError:
|
||||||
|
return executable
|
||||||
|
return str(link)
|
||||||
@@ -17,6 +17,7 @@ from pathlib import Path
|
|||||||
from typing import TYPE_CHECKING, Any, cast
|
from typing import TYPE_CHECKING, Any, cast
|
||||||
|
|
||||||
from nanobot import __version__
|
from nanobot import __version__
|
||||||
|
from nanobot.cli.process_identity import named_executable
|
||||||
from nanobot.cli.runtime_config import _model_display
|
from nanobot.cli.runtime_config import _model_display
|
||||||
from nanobot.cli.webui_support import (
|
from nanobot.cli.webui_support import (
|
||||||
_gateway_health_ready,
|
_gateway_health_ready,
|
||||||
@@ -229,7 +230,12 @@ def _resolve_source_tui_command(source_dir: Path, bun: str) -> list[str]:
|
|||||||
detail = (install.stderr or install.stdout).strip().splitlines()
|
detail = (install.stderr or install.stdout).strip().splitlines()
|
||||||
suffix = f": {detail[-1]}" if detail else ""
|
suffix = f": {detail[-1]}" if detail else ""
|
||||||
raise TuiUnavailableError(f"could not install TUI dependencies{suffix}")
|
raise TuiUnavailableError(f"could not install TUI dependencies{suffix}")
|
||||||
return [bun, str(source_dir / "src" / "index.ts")]
|
executable = named_executable(
|
||||||
|
bun,
|
||||||
|
name="nanobot-tui",
|
||||||
|
directory=get_data_dir() / "run" / "executables",
|
||||||
|
)
|
||||||
|
return [executable, str(source_dir / "src" / "index.ts")]
|
||||||
|
|
||||||
|
|
||||||
def _download_release_tui(asset: str) -> Path | None:
|
def _download_release_tui(asset: str) -> Path | None:
|
||||||
|
|||||||
+72
-25
@@ -3,6 +3,7 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import asyncio
|
import asyncio
|
||||||
|
import inspect
|
||||||
import os
|
import os
|
||||||
import subprocess
|
import subprocess
|
||||||
import sys
|
import sys
|
||||||
@@ -14,6 +15,8 @@ from typing import TYPE_CHECKING, Any, Literal, cast
|
|||||||
from nanobot import __version__
|
from nanobot import __version__
|
||||||
from nanobot.bus.events import INBOUND_META_USER_SHELL, OutboundMessage
|
from nanobot.bus.events import INBOUND_META_USER_SHELL, OutboundMessage
|
||||||
from nanobot.command.router import CommandContext, CommandRouter, normalize_command_text
|
from nanobot.command.router import CommandContext, CommandRouter, normalize_command_text
|
||||||
|
from nanobot.session.async_compat import call_session_manager
|
||||||
|
from nanobot.utils.cancellation import shield_and_drain
|
||||||
from nanobot.utils.helpers import build_status_content
|
from nanobot.utils.helpers import build_status_content
|
||||||
from nanobot.utils.restart import set_restart_notice_to_env
|
from nanobot.utils.restart import set_restart_notice_to_env
|
||||||
from nanobot.utils.workspace_prompts import initialize_workspace_prompt
|
from nanobot.utils.workspace_prompts import initialize_workspace_prompt
|
||||||
@@ -22,6 +25,7 @@ if TYPE_CHECKING:
|
|||||||
from nanobot.agent.loop import AgentLoop
|
from nanobot.agent.loop import AgentLoop
|
||||||
from nanobot.session.manager import Session
|
from nanobot.session.manager import Session
|
||||||
from nanobot.utils.gitstore import CommitInfo
|
from nanobot.utils.gitstore import CommitInfo
|
||||||
|
from nanobot.utils.llm_runtime import LLMRuntime
|
||||||
|
|
||||||
# WebUI protocol contract for how a slash command participates in turn state:
|
# WebUI protocol contract for how a slash command participates in turn state:
|
||||||
# - side_channel: returns control text without starting or ending an agent turn.
|
# - side_channel: returns control text without starting or ending an agent turn.
|
||||||
@@ -201,6 +205,52 @@ def builtin_command_starts_agent_turn(text: str) -> bool:
|
|||||||
return spec.lifecycle == "agent_turn_with_args" and bool(args.strip())
|
return spec.lifecycle == "agent_turn_with_args" and bool(args.strip())
|
||||||
|
|
||||||
|
|
||||||
|
def _has_native_coroutine_method(target: object, name: str) -> bool:
|
||||||
|
"""Check the real target class without trusting dynamic mock attributes."""
|
||||||
|
method = inspect.getattr_static(type(target), name, None)
|
||||||
|
return inspect.iscoroutinefunction(method)
|
||||||
|
|
||||||
|
|
||||||
|
async def _get_or_create_session(loop: AgentLoop, key: str) -> Session:
|
||||||
|
sessions = loop.sessions
|
||||||
|
return await call_session_manager(
|
||||||
|
sessions,
|
||||||
|
"get_or_create_async",
|
||||||
|
sessions.get_or_create,
|
||||||
|
key,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
async def _save_session(loop: AgentLoop, session: Session) -> None:
|
||||||
|
sessions = loop.sessions
|
||||||
|
await call_session_manager(
|
||||||
|
sessions,
|
||||||
|
"save_async",
|
||||||
|
sessions.save,
|
||||||
|
session,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
async def _runtime_for_session(loop: AgentLoop, session: Session) -> LLMRuntime:
|
||||||
|
if _has_native_coroutine_method(loop, "runtime_for_session_async"):
|
||||||
|
return await loop.runtime_for_session_async(session)
|
||||||
|
return await shield_and_drain(
|
||||||
|
asyncio.to_thread(loop.runtime_for_session, session)
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
async def _set_session_model_preset(
|
||||||
|
loop: AgentLoop,
|
||||||
|
session_key: str,
|
||||||
|
name: str,
|
||||||
|
) -> LLMRuntime:
|
||||||
|
if _has_native_coroutine_method(loop, "set_session_model_preset_async"):
|
||||||
|
return await loop.set_session_model_preset_async(session_key, name)
|
||||||
|
return await shield_and_drain(
|
||||||
|
asyncio.to_thread(loop.set_session_model_preset, session_key, name)
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
async def cmd_stop(ctx: CommandContext) -> OutboundMessage:
|
async def cmd_stop(ctx: CommandContext) -> OutboundMessage:
|
||||||
"""Cancel all active tasks and subagents for the session."""
|
"""Cancel all active tasks and subagents for the session."""
|
||||||
loop = ctx.loop
|
loop = ctx.loop
|
||||||
@@ -257,8 +307,8 @@ async def cmd_restart(ctx: CommandContext) -> OutboundMessage:
|
|||||||
async def cmd_status(ctx: CommandContext) -> OutboundMessage:
|
async def cmd_status(ctx: CommandContext) -> OutboundMessage:
|
||||||
"""Build an outbound status message for a session."""
|
"""Build an outbound status message for a session."""
|
||||||
loop = ctx.loop
|
loop = ctx.loop
|
||||||
session = ctx.session or loop.sessions.get_or_create(ctx.key)
|
session = ctx.session or await _get_or_create_session(loop, ctx.key)
|
||||||
runtime = ctx.runtime or loop.runtime_for_session(session)
|
runtime = ctx.runtime or await _runtime_for_session(loop, session)
|
||||||
ctx_est = 0
|
ctx_est = 0
|
||||||
with suppress(Exception):
|
with suppress(Exception):
|
||||||
ctx_est, _ = loop.consolidator.estimate_session_prompt_tokens(
|
ctx_est, _ = loop.consolidator.estimate_session_prompt_tokens(
|
||||||
@@ -266,7 +316,8 @@ async def cmd_status(ctx: CommandContext) -> OutboundMessage:
|
|||||||
runtime=runtime,
|
runtime=runtime,
|
||||||
)
|
)
|
||||||
if ctx_est <= 0:
|
if ctx_est <= 0:
|
||||||
ctx_est = loop._last_usage.get("prompt_tokens", 0) # pyright: ignore[reportPrivateUsage]
|
last_usage = loop._last_usage # pyright: ignore[reportPrivateUsage]
|
||||||
|
ctx_est = last_usage.input_tokens if last_usage is not None else 0
|
||||||
|
|
||||||
# Fetch web search provider usage (best-effort, never blocks the response)
|
# Fetch web search provider usage (best-effort, never blocks the response)
|
||||||
search_usage_text: str | None = None
|
search_usage_text: str | None = None
|
||||||
@@ -305,29 +356,32 @@ async def cmd_new(ctx: CommandContext) -> OutboundMessage:
|
|||||||
loop = ctx.loop
|
loop = ctx.loop
|
||||||
await loop._cancel_active_tasks(ctx.key) # pyright: ignore[reportPrivateUsage]
|
await loop._cancel_active_tasks(ctx.key) # pyright: ignore[reportPrivateUsage]
|
||||||
loop.discard_session_file_state(ctx.key)
|
loop.discard_session_file_state(ctx.key)
|
||||||
session = ctx.session or loop.sessions.get_or_create(ctx.key)
|
session = ctx.session or await _get_or_create_session(loop, ctx.key)
|
||||||
snapshot = list(session.messages)
|
snapshot = list(session.messages)
|
||||||
archive_snapshot = None
|
archive_snapshot = None
|
||||||
runtime = None
|
runtime = None
|
||||||
if session.last_consolidated < len(snapshot):
|
if session.last_consolidated < len(snapshot):
|
||||||
runtime = ctx.runtime or loop.runtime_for_session(session)
|
runtime = ctx.runtime or await _runtime_for_session(loop, session)
|
||||||
archive_snapshot = replace(
|
archive_snapshot = replace(
|
||||||
session,
|
session,
|
||||||
messages=snapshot,
|
messages=snapshot,
|
||||||
metadata=dict(session.metadata),
|
metadata=dict(session.metadata),
|
||||||
provider_state=None,
|
provider_state=None,
|
||||||
)
|
)
|
||||||
session.clear()
|
async def reset_and_schedule_archive() -> None:
|
||||||
loop.sessions.save(session)
|
session.clear()
|
||||||
loop.sessions.invalidate(session.key)
|
await _save_session(loop, session)
|
||||||
if archive_snapshot is not None and runtime is not None:
|
loop.sessions.invalidate(session.key)
|
||||||
loop.schedule_background(
|
if archive_snapshot is not None and runtime is not None:
|
||||||
loop.consolidator.archive_session( # pyright: ignore[reportUnknownMemberType]
|
loop.schedule_background(
|
||||||
archive_snapshot,
|
loop.consolidator.archive_session( # pyright: ignore[reportUnknownMemberType]
|
||||||
archive_end=len(snapshot),
|
archive_snapshot,
|
||||||
runtime=runtime,
|
archive_end=len(snapshot),
|
||||||
|
runtime=runtime,
|
||||||
|
)
|
||||||
)
|
)
|
||||||
)
|
|
||||||
|
await shield_and_drain(reset_and_schedule_archive())
|
||||||
return OutboundMessage(
|
return OutboundMessage(
|
||||||
channel=ctx.msg.channel, chat_id=ctx.msg.chat_id,
|
channel=ctx.msg.channel, chat_id=ctx.msg.chat_id,
|
||||||
content="New session started.",
|
content="New session started.",
|
||||||
@@ -376,7 +430,7 @@ async def cmd_model(ctx: CommandContext) -> OutboundMessage:
|
|||||||
metadata = {**dict(ctx.msg.metadata or {}), "render_as": "text"}
|
metadata = {**dict(ctx.msg.metadata or {}), "render_as": "text"}
|
||||||
|
|
||||||
if not args:
|
if not args:
|
||||||
session = ctx.session or loop.sessions.get_or_create(ctx.key)
|
session = ctx.session or await _get_or_create_session(loop, ctx.key)
|
||||||
return OutboundMessage(
|
return OutboundMessage(
|
||||||
channel=ctx.msg.channel,
|
channel=ctx.msg.channel,
|
||||||
chat_id=ctx.msg.chat_id,
|
chat_id=ctx.msg.chat_id,
|
||||||
@@ -386,7 +440,7 @@ async def cmd_model(ctx: CommandContext) -> OutboundMessage:
|
|||||||
|
|
||||||
name = args
|
name = args
|
||||||
try:
|
try:
|
||||||
runtime = loop.set_session_model_preset(ctx.key, name)
|
runtime = await _set_session_model_preset(loop, ctx.key, name)
|
||||||
except (KeyError, ValueError) as exc:
|
except (KeyError, ValueError) as exc:
|
||||||
names = _model_preset_names(loop)
|
names = _model_preset_names(loop)
|
||||||
return OutboundMessage(
|
return OutboundMessage(
|
||||||
@@ -478,13 +532,6 @@ async def cmd_dream(ctx: CommandContext) -> OutboundMessage:
|
|||||||
elapsed = time.monotonic() - t0
|
elapsed = time.monotonic() - t0
|
||||||
content = f"Dream failed after {elapsed:.1f}s: {e}"
|
content = f"Dream failed after {elapsed:.1f}s: {e}"
|
||||||
finally:
|
finally:
|
||||||
from nanobot.webui.token_usage import record_response_token_usage
|
|
||||||
|
|
||||||
record_response_token_usage(
|
|
||||||
resp,
|
|
||||||
source="dream",
|
|
||||||
timezone_name=getattr(loop.context, "timezone", None),
|
|
||||||
)
|
|
||||||
if store.git.is_initialized():
|
if store.git.is_initialized():
|
||||||
commit_msg = build_dream_commit_message("dream: manual run", diff_body)
|
commit_msg = build_dream_commit_message("dream: manual run", diff_body)
|
||||||
sha = store.git.auto_commit(commit_msg)
|
sha = store.git.auto_commit(commit_msg)
|
||||||
@@ -854,7 +901,7 @@ async def cmd_history(ctx: CommandContext) -> OutboundMessage:
|
|||||||
metadata=dict(ctx.msg.metadata or {}),
|
metadata=dict(ctx.msg.metadata or {}),
|
||||||
)
|
)
|
||||||
|
|
||||||
session = ctx.session or ctx.loop.sessions.get_or_create(ctx.key)
|
session = ctx.session or await _get_or_create_session(ctx.loop, ctx.key)
|
||||||
history = session.get_history(max_messages=0, include_runtime_context=False)
|
history = session.get_history(max_messages=0, include_runtime_context=False)
|
||||||
visible = [_format_history_message(m) for m in history]
|
visible = [_format_history_message(m) for m in history]
|
||||||
visible = [m for m in visible if m is not None]
|
visible = [m for m in visible if m is not None]
|
||||||
|
|||||||
@@ -8,6 +8,7 @@ import time
|
|||||||
import uuid
|
import uuid
|
||||||
from typing import TYPE_CHECKING, Any, Protocol
|
from typing import TYPE_CHECKING, Any, Protocol
|
||||||
|
|
||||||
|
from nanobot.agent.automation_turns import AutomationTurnAcceptedCancellation
|
||||||
from nanobot.agent.tools.cron import CronTool
|
from nanobot.agent.tools.cron import CronTool
|
||||||
from nanobot.bus.events import InboundMessage, OutboundMessage
|
from nanobot.bus.events import InboundMessage, OutboundMessage
|
||||||
from nanobot.cron.session_delivery import origin_delivery_context
|
from nanobot.cron.session_delivery import origin_delivery_context
|
||||||
@@ -127,6 +128,15 @@ async def run_bound_cron_job(
|
|||||||
session_key_override=session_key,
|
session_key_override=session_key,
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
except AutomationTurnAcceptedCancellation:
|
||||||
|
cron.write_run_record(
|
||||||
|
run_id,
|
||||||
|
{
|
||||||
|
**run_record_base,
|
||||||
|
"status": "accepted",
|
||||||
|
},
|
||||||
|
)
|
||||||
|
raise
|
||||||
except (Exception, asyncio.CancelledError) as exc:
|
except (Exception, asyncio.CancelledError) as exc:
|
||||||
error_text = str(exc) or exc.__class__.__name__
|
error_text = str(exc) or exc.__class__.__name__
|
||||||
cron.write_run_record(
|
cron.write_run_record(
|
||||||
|
|||||||
+230
-87
@@ -11,11 +11,12 @@ from dataclasses import asdict
|
|||||||
from datetime import datetime
|
from datetime import datetime
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from types import EllipsisType
|
from types import EllipsisType
|
||||||
from typing import Any, Callable, Coroutine, Literal
|
from typing import Any, Callable, Coroutine, Literal, TypeVar
|
||||||
|
|
||||||
from filelock import FileLock
|
from filelock import FileLock
|
||||||
from loguru import logger
|
from loguru import logger
|
||||||
|
|
||||||
|
from nanobot.agent.automation_turns import AutomationTurnAcceptedCancellation
|
||||||
from nanobot.cron.session_turns import is_bound_cron_job
|
from nanobot.cron.session_turns import is_bound_cron_job
|
||||||
from nanobot.cron.types import (
|
from nanobot.cron.types import (
|
||||||
CronJob,
|
CronJob,
|
||||||
@@ -25,10 +26,14 @@ from nanobot.cron.types import (
|
|||||||
CronSchedule,
|
CronSchedule,
|
||||||
CronStore,
|
CronStore,
|
||||||
)
|
)
|
||||||
|
from nanobot.utils.cancellation import shield_and_drain
|
||||||
from nanobot.utils.run_records import (
|
from nanobot.utils.run_records import (
|
||||||
write_run_record as write_automation_run_record,
|
write_run_record as write_automation_run_record,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
_FILE_LOCK_TIMEOUT_SECONDS = 5
|
||||||
|
_T = TypeVar("_T")
|
||||||
|
|
||||||
|
|
||||||
class CronJobSkippedError(Exception):
|
class CronJobSkippedError(Exception):
|
||||||
"""Raised by cron callbacks when a job was intentionally skipped."""
|
"""Raised by cron callbacks when a job was intentionally skipped."""
|
||||||
@@ -164,10 +169,16 @@ class CronService:
|
|||||||
self.store_path = store_path
|
self.store_path = store_path
|
||||||
self._action_path = store_path.parent / "action.jsonl"
|
self._action_path = store_path.parent / "action.jsonl"
|
||||||
self._run_records_dir = store_path.parent / "runs"
|
self._run_records_dir = store_path.parent / "runs"
|
||||||
self._lock = FileLock(str(self._action_path.parent) + ".lock")
|
self._lock = FileLock(
|
||||||
|
str(self._action_path.parent) + ".lock",
|
||||||
|
timeout=_FILE_LOCK_TIMEOUT_SECONDS,
|
||||||
|
)
|
||||||
self.on_job = on_job
|
self.on_job = on_job
|
||||||
self._store: CronStore | None = None
|
self._store: CronStore | None = None
|
||||||
self._timer_task: asyncio.Task[None] | None = None
|
self._timer_task: asyncio.Task[None] | None = None
|
||||||
|
self._operation_lock = asyncio.Lock()
|
||||||
|
self._claimed_job_ids: set[str] = set()
|
||||||
|
self._event_loop: asyncio.AbstractEventLoop | None = None
|
||||||
self._running = False
|
self._running = False
|
||||||
self._active_executions = 0
|
self._active_executions = 0
|
||||||
self._store_dirty = False
|
self._store_dirty = False
|
||||||
@@ -451,25 +462,58 @@ class CronService:
|
|||||||
"""Write an internal audit record for one cron execution."""
|
"""Write an internal audit record for one cron execution."""
|
||||||
write_automation_run_record(self._run_records_dir, run_id, record)
|
write_automation_run_record(self._run_records_dir, run_id, record)
|
||||||
|
|
||||||
async def start(self) -> None:
|
async def run_sync(
|
||||||
"""Start the cron service."""
|
self,
|
||||||
self._running = True
|
operation: Callable[..., _T],
|
||||||
loaded = self._load_store()
|
/,
|
||||||
if loaded is None:
|
*args: Any,
|
||||||
# Store file existed but was corrupt and has been preserved with
|
**kwargs: Any,
|
||||||
# a ``.corrupt-<ts>`` suffix. Bail out instead of starting with
|
) -> _T:
|
||||||
# an empty store; that would call ``_save_store`` and overwrite
|
"""Serialize a complete cron transaction in a worker thread.
|
||||||
# the now-renamed (but still recoverable) data with [].
|
|
||||||
self._running = False
|
A running thread cannot be cancelled safely. Keep the transaction lock
|
||||||
raise RuntimeError(
|
until it exits so cancellation is never reported while that worker can
|
||||||
f"cron store at {self.store_path} is corrupt and was preserved; "
|
still mutate cron state behind a later operation.
|
||||||
"refusing to start with an empty job list. "
|
"""
|
||||||
"Inspect the .corrupt-<ts> backup and restore manually."
|
async with self._operation_lock:
|
||||||
|
return await shield_and_drain(
|
||||||
|
asyncio.to_thread(operation, *args, **kwargs)
|
||||||
)
|
)
|
||||||
self._recompute_next_runs()
|
|
||||||
self._save_store()
|
async def start(self) -> None:
|
||||||
self._arm_timer()
|
"""Start the cron service and settle accepted work before cancellation."""
|
||||||
logger.info("Cron service started with {} jobs", len(self._store.jobs if self._store else []))
|
|
||||||
|
async def settle_start() -> None:
|
||||||
|
self._event_loop = asyncio.get_running_loop()
|
||||||
|
self._running = True
|
||||||
|
try:
|
||||||
|
async with self._operation_lock:
|
||||||
|
loaded = await asyncio.to_thread(self._load_store)
|
||||||
|
if loaded is None:
|
||||||
|
# Store file existed but was corrupt and has been preserved with
|
||||||
|
# a ``.corrupt-<ts>`` suffix. Bail out instead of starting with
|
||||||
|
# an empty store; that would call ``_save_store`` and overwrite
|
||||||
|
# the now-renamed (but still recoverable) data with [].
|
||||||
|
raise RuntimeError(
|
||||||
|
f"cron store at {self.store_path} is corrupt and was preserved; "
|
||||||
|
"refusing to start with an empty job list. "
|
||||||
|
"Inspect the .corrupt-<ts> backup and restore manually."
|
||||||
|
)
|
||||||
|
self._recompute_next_runs()
|
||||||
|
await asyncio.to_thread(self._save_store)
|
||||||
|
self._arm_timer()
|
||||||
|
logger.info(
|
||||||
|
"Cron service started with {} jobs",
|
||||||
|
len(self._store.jobs if self._store else []),
|
||||||
|
)
|
||||||
|
except BaseException:
|
||||||
|
# A failed start must not retain ownership without a timer. Caller
|
||||||
|
# cancellation is shielded until this composite either reaches the
|
||||||
|
# fully started state above or rolls back here.
|
||||||
|
self.stop()
|
||||||
|
raise
|
||||||
|
|
||||||
|
await shield_and_drain(settle_start())
|
||||||
|
|
||||||
def stop(self) -> None:
|
def stop(self) -> None:
|
||||||
"""Stop the cron service."""
|
"""Stop the cron service."""
|
||||||
@@ -497,8 +541,22 @@ class CronService:
|
|||||||
if j.enabled and j.state.next_run_at_ms]
|
if j.enabled and j.state.next_run_at_ms]
|
||||||
return min(times) if times else None
|
return min(times) if times else None
|
||||||
|
|
||||||
|
def _request_timer_rearm(self) -> None:
|
||||||
|
"""Re-arm on the owning event loop, including from persistence workers."""
|
||||||
|
if not self._running:
|
||||||
|
return
|
||||||
|
loop = self._event_loop
|
||||||
|
try:
|
||||||
|
current_loop = asyncio.get_running_loop()
|
||||||
|
except RuntimeError:
|
||||||
|
current_loop = None
|
||||||
|
if current_loop is loop:
|
||||||
|
self._arm_timer()
|
||||||
|
elif loop is not None and loop.is_running():
|
||||||
|
loop.call_soon_threadsafe(self._arm_timer)
|
||||||
|
|
||||||
def _arm_timer(self) -> None:
|
def _arm_timer(self) -> None:
|
||||||
"""Schedule the next timer tick."""
|
"""Schedule the next timer tick on the owning event loop."""
|
||||||
if self._timer_task:
|
if self._timer_task:
|
||||||
self._timer_task.cancel()
|
self._timer_task.cancel()
|
||||||
|
|
||||||
@@ -520,7 +578,7 @@ class CronService:
|
|||||||
self._timer_task = asyncio.create_task(tick())
|
self._timer_task = asyncio.create_task(tick())
|
||||||
|
|
||||||
async def _on_timer(self) -> None:
|
async def _on_timer(self) -> None:
|
||||||
"""Handle timer tick - run due jobs."""
|
"""Run due jobs while keeping persistence transactions serialized."""
|
||||||
reload_store = self._active_executions == 0
|
reload_store = self._active_executions == 0
|
||||||
self._active_executions += 1
|
self._active_executions += 1
|
||||||
try:
|
try:
|
||||||
@@ -528,11 +586,17 @@ class CronService:
|
|||||||
# to persist their advanced schedule. Persist that exact snapshot
|
# to persist their advanced schedule. Persist that exact snapshot
|
||||||
# before reloading or executing anything else; otherwise the older
|
# before reloading or executing anything else; otherwise the older
|
||||||
# disk state can replay the same job.
|
# disk state can replay the same job.
|
||||||
if self._store_dirty:
|
async with self._operation_lock:
|
||||||
self._save_store()
|
if self._store_dirty:
|
||||||
return
|
await shield_and_drain(asyncio.to_thread(self._save_store))
|
||||||
|
return
|
||||||
|
|
||||||
store = self._load_store(reload_during_execution=reload_store)
|
store = await shield_and_drain(
|
||||||
|
asyncio.to_thread(
|
||||||
|
self._load_store,
|
||||||
|
reload_during_execution=reload_store,
|
||||||
|
)
|
||||||
|
)
|
||||||
# If a hot reload found a corrupt store on disk, ``self._store``
|
# If a hot reload found a corrupt store on disk, ``self._store``
|
||||||
# may still hold the previous, known-good in-memory snapshot.
|
# may still hold the previous, known-good in-memory snapshot.
|
||||||
if store is None:
|
if store is None:
|
||||||
@@ -547,7 +611,8 @@ class CronService:
|
|||||||
for job in due_jobs:
|
for job in due_jobs:
|
||||||
await self._execute_job(job)
|
await self._execute_job(job)
|
||||||
|
|
||||||
self._save_store()
|
async with self._operation_lock:
|
||||||
|
await shield_and_drain(asyncio.to_thread(self._save_store))
|
||||||
except Exception:
|
except Exception:
|
||||||
# A load/persist failure must not kill the scheduler: keep the
|
# A load/persist failure must not kill the scheduler: keep the
|
||||||
# in-memory store and retry on the next tick. This mirrors the
|
# in-memory store and retry on the next tick. This mirrors the
|
||||||
@@ -564,58 +629,124 @@ class CronService:
|
|||||||
# single bad tick cannot silently stop all future jobs.
|
# single bad tick cannot silently stop all future jobs.
|
||||||
self._arm_timer()
|
self._arm_timer()
|
||||||
|
|
||||||
async def _execute_job(self, job: CronJob) -> None:
|
async def _claim_job(self, job_id: str) -> bool:
|
||||||
"""Execute a single job."""
|
"""Claim one job without serializing callbacks for different jobs."""
|
||||||
|
async with self._operation_lock:
|
||||||
|
if job_id in self._claimed_job_ids:
|
||||||
|
return False
|
||||||
|
self._claimed_job_ids.add(job_id)
|
||||||
|
return True
|
||||||
|
|
||||||
|
async def _release_job_claim(self, job_id: str) -> None:
|
||||||
|
async with self._operation_lock:
|
||||||
|
self._claimed_job_ids.discard(job_id)
|
||||||
|
|
||||||
|
async def _settle_job_execution(
|
||||||
|
self,
|
||||||
|
job: CronJob,
|
||||||
|
*,
|
||||||
|
start_ms: int,
|
||||||
|
status: Literal["ok", "error", "skipped"],
|
||||||
|
error: str | None,
|
||||||
|
persist: bool = False,
|
||||||
|
) -> None:
|
||||||
|
end_ms = _now_ms()
|
||||||
|
async with self._operation_lock:
|
||||||
|
job.state.last_status = status
|
||||||
|
job.state.last_error = error
|
||||||
|
job.state.last_run_at_ms = start_ms
|
||||||
|
job.updated_at_ms = end_ms
|
||||||
|
job.state.run_history.append(CronRunRecord(
|
||||||
|
run_at_ms=start_ms,
|
||||||
|
status=status,
|
||||||
|
duration_ms=end_ms - start_ms,
|
||||||
|
error=error,
|
||||||
|
))
|
||||||
|
job.state.run_history = job.state.run_history[-self._MAX_RUN_HISTORY:]
|
||||||
|
|
||||||
|
if job.schedule.kind == "at":
|
||||||
|
if job.delete_after_run:
|
||||||
|
store = await shield_and_drain(
|
||||||
|
asyncio.to_thread(self._require_store)
|
||||||
|
)
|
||||||
|
store.jobs = [item for item in store.jobs if item.id != job.id]
|
||||||
|
else:
|
||||||
|
job.enabled = False
|
||||||
|
job.state.next_run_at_ms = None
|
||||||
|
else:
|
||||||
|
job.state.next_run_at_ms = _compute_next_run(job.schedule, _now_ms())
|
||||||
|
|
||||||
|
if persist:
|
||||||
|
await shield_and_drain(asyncio.to_thread(self._save_store))
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
async def _drain_settlement_on_cancellation(settlement: asyncio.Task[None]) -> None:
|
||||||
|
"""Finish a short durable settlement despite repeated cancellation."""
|
||||||
|
while not settlement.done():
|
||||||
|
try:
|
||||||
|
await asyncio.shield(settlement)
|
||||||
|
except asyncio.CancelledError:
|
||||||
|
continue
|
||||||
|
settlement.result()
|
||||||
|
|
||||||
|
async def _execute_job(self, job: CronJob) -> bool:
|
||||||
|
"""Execute a claimed job and serialize its in-memory settlement."""
|
||||||
|
if not await self._claim_job(job.id):
|
||||||
|
logger.info("Cron: job '{}' ({}) is already running", job.name, job.id)
|
||||||
|
return False
|
||||||
|
|
||||||
start_ms = _now_ms()
|
start_ms = _now_ms()
|
||||||
logger.info("Cron: executing job '{}' ({})", job.name, job.id)
|
logger.info("Cron: executing job '{}' ({})", job.name, job.id)
|
||||||
|
status: Literal["ok", "error", "skipped"]
|
||||||
|
error: str | None
|
||||||
|
accepted_cancellation: AutomationTurnAcceptedCancellation | None = None
|
||||||
|
|
||||||
try:
|
try:
|
||||||
if self.on_job:
|
try:
|
||||||
await self.on_job(job)
|
if self.on_job:
|
||||||
|
await self.on_job(job)
|
||||||
|
status = "ok"
|
||||||
|
error = None
|
||||||
|
logger.info("Cron: job '{}' completed", job.name)
|
||||||
|
except AutomationTurnAcceptedCancellation as exc:
|
||||||
|
# The agent owns this turn now. Advance and persist the schedule
|
||||||
|
# before allowing shutdown cancellation to unwind the timer.
|
||||||
|
status = "ok"
|
||||||
|
error = None
|
||||||
|
accepted_cancellation = exc
|
||||||
|
logger.info("Cron: job '{}' was accepted before cancellation", job.name)
|
||||||
|
except CronJobSkippedError as exc:
|
||||||
|
status = "skipped"
|
||||||
|
error = str(exc) or None
|
||||||
|
logger.warning("Cron: job '{}' skipped: {}", job.name, error or "")
|
||||||
|
except asyncio.CancelledError as exc:
|
||||||
|
current = asyncio.current_task()
|
||||||
|
if current is not None and current.cancelling():
|
||||||
|
raise
|
||||||
|
status = "error"
|
||||||
|
error = str(exc) or exc.__class__.__name__
|
||||||
|
logger.exception("Cron: job '{}' was cancelled", job.name)
|
||||||
|
except Exception as exc:
|
||||||
|
status = "error"
|
||||||
|
error = str(exc)
|
||||||
|
logger.exception("Cron: job '{}' failed", job.name)
|
||||||
|
|
||||||
job.state.last_status = "ok"
|
settlement = asyncio.create_task(
|
||||||
job.state.last_error = None
|
self._settle_job_execution(
|
||||||
logger.info("Cron: job '{}' completed", job.name)
|
job,
|
||||||
|
start_ms=start_ms,
|
||||||
except CronJobSkippedError as e:
|
status=status,
|
||||||
job.state.last_status = "skipped"
|
error=error,
|
||||||
job.state.last_error = str(e) or None
|
persist=accepted_cancellation is not None,
|
||||||
logger.warning("Cron: job '{}' skipped: {}", job.name, job.state.last_error or "")
|
)
|
||||||
except asyncio.CancelledError as e:
|
)
|
||||||
current = asyncio.current_task()
|
if accepted_cancellation is not None:
|
||||||
if current is not None and current.cancelling():
|
await self._drain_settlement_on_cancellation(settlement)
|
||||||
raise
|
raise accepted_cancellation
|
||||||
job.state.last_status = "error"
|
await settlement
|
||||||
job.state.last_error = str(e) or e.__class__.__name__
|
return True
|
||||||
logger.exception("Cron: job '{}' was cancelled", job.name)
|
finally:
|
||||||
except Exception as e:
|
await self._release_job_claim(job.id)
|
||||||
job.state.last_status = "error"
|
|
||||||
job.state.last_error = str(e)
|
|
||||||
logger.exception("Cron: job '{}' failed", job.name)
|
|
||||||
|
|
||||||
end_ms = _now_ms()
|
|
||||||
job.state.last_run_at_ms = start_ms
|
|
||||||
job.updated_at_ms = end_ms
|
|
||||||
|
|
||||||
job.state.run_history.append(CronRunRecord(
|
|
||||||
run_at_ms=start_ms,
|
|
||||||
status=job.state.last_status,
|
|
||||||
duration_ms=end_ms - start_ms,
|
|
||||||
error=job.state.last_error,
|
|
||||||
))
|
|
||||||
job.state.run_history = job.state.run_history[-self._MAX_RUN_HISTORY:]
|
|
||||||
|
|
||||||
# Handle one-shot jobs
|
|
||||||
if job.schedule.kind == "at":
|
|
||||||
if job.delete_after_run:
|
|
||||||
store = self._require_store()
|
|
||||||
store.jobs = [item for item in store.jobs if item.id != job.id]
|
|
||||||
else:
|
|
||||||
job.enabled = False
|
|
||||||
job.state.next_run_at_ms = None
|
|
||||||
else:
|
|
||||||
# Compute next run
|
|
||||||
job.state.next_run_at_ms = _compute_next_run(job.schedule, _now_ms())
|
|
||||||
|
|
||||||
def _append_action(
|
def _append_action(
|
||||||
self,
|
self,
|
||||||
@@ -697,7 +828,7 @@ class CronService:
|
|||||||
store = self._require_store()
|
store = self._require_store()
|
||||||
store.jobs.append(job)
|
store.jobs.append(job)
|
||||||
self._save_store()
|
self._save_store()
|
||||||
self._arm_timer()
|
self._request_timer_rearm()
|
||||||
else:
|
else:
|
||||||
self._append_action("add", asdict(job))
|
self._append_action("add", asdict(job))
|
||||||
|
|
||||||
@@ -714,7 +845,7 @@ class CronService:
|
|||||||
store.jobs = [j for j in store.jobs if j.id != job.id]
|
store.jobs = [j for j in store.jobs if j.id != job.id]
|
||||||
store.jobs.append(job)
|
store.jobs.append(job)
|
||||||
self._save_store()
|
self._save_store()
|
||||||
self._arm_timer()
|
self._request_timer_rearm()
|
||||||
logger.info("Cron: registered system job '{}' ({})", job.name, job.id)
|
logger.info("Cron: registered system job '{}' ({})", job.name, job.id)
|
||||||
return job
|
return job
|
||||||
|
|
||||||
@@ -726,7 +857,7 @@ class CronService:
|
|||||||
removed = len(store.jobs) < before
|
removed = len(store.jobs) < before
|
||||||
if removed:
|
if removed:
|
||||||
self._save_store()
|
self._save_store()
|
||||||
self._arm_timer()
|
self._request_timer_rearm()
|
||||||
logger.info("Cron: removed system job {}", job_id)
|
logger.info("Cron: removed system job {}", job_id)
|
||||||
return removed
|
return removed
|
||||||
|
|
||||||
@@ -747,7 +878,7 @@ class CronService:
|
|||||||
if removed:
|
if removed:
|
||||||
if self._should_persist_store():
|
if self._should_persist_store():
|
||||||
self._save_store()
|
self._save_store()
|
||||||
self._arm_timer()
|
self._request_timer_rearm()
|
||||||
else:
|
else:
|
||||||
self._append_action("del", {"job_id": job_id})
|
self._append_action("del", {"job_id": job_id})
|
||||||
logger.info("Cron: removed job {}", job_id)
|
logger.info("Cron: removed job {}", job_id)
|
||||||
@@ -769,7 +900,7 @@ class CronService:
|
|||||||
job.state.next_run_at_ms = None
|
job.state.next_run_at_ms = None
|
||||||
if self._should_persist_store():
|
if self._should_persist_store():
|
||||||
self._save_store()
|
self._save_store()
|
||||||
self._arm_timer()
|
self._request_timer_rearm()
|
||||||
else:
|
else:
|
||||||
self._append_action("update", asdict(job))
|
self._append_action("update", asdict(job))
|
||||||
return job
|
return job
|
||||||
@@ -825,7 +956,7 @@ class CronService:
|
|||||||
|
|
||||||
if self._should_persist_store():
|
if self._should_persist_store():
|
||||||
self._save_store()
|
self._save_store()
|
||||||
self._arm_timer()
|
self._request_timer_rearm()
|
||||||
else:
|
else:
|
||||||
self._append_action("update", asdict(job))
|
self._append_action("update", asdict(job))
|
||||||
|
|
||||||
@@ -840,19 +971,31 @@ class CronService:
|
|||||||
# A manual run is another side-effecting entrypoint. Do not start
|
# A manual run is another side-effecting entrypoint. Do not start
|
||||||
# it while the result of a previous timer execution is still only
|
# it while the result of a previous timer execution is still only
|
||||||
# in memory.
|
# in memory.
|
||||||
if self._store_dirty:
|
async with self._operation_lock:
|
||||||
self._save_store()
|
if self._store_dirty:
|
||||||
store = self._require_store(reload_during_execution=reload_store)
|
await shield_and_drain(asyncio.to_thread(self._save_store))
|
||||||
|
store = await shield_and_drain(
|
||||||
|
asyncio.to_thread(
|
||||||
|
self._require_store,
|
||||||
|
reload_during_execution=reload_store,
|
||||||
|
)
|
||||||
|
)
|
||||||
for job in store.jobs:
|
for job in store.jobs:
|
||||||
if job.id == job_id:
|
if job.id == job_id:
|
||||||
if self._is_unbound_agent_job(job):
|
if self._is_unbound_agent_job(job):
|
||||||
self._enforce_agent_binding(job)
|
async with self._operation_lock:
|
||||||
self._save_store()
|
self._enforce_agent_binding(job)
|
||||||
|
await shield_and_drain(
|
||||||
|
asyncio.to_thread(self._save_store)
|
||||||
|
)
|
||||||
return False
|
return False
|
||||||
if not force and not job.enabled:
|
if not force and not job.enabled:
|
||||||
return False
|
return False
|
||||||
await self._execute_job(job)
|
executed = await self._execute_job(job)
|
||||||
self._save_store()
|
if not executed:
|
||||||
|
return False
|
||||||
|
async with self._operation_lock:
|
||||||
|
await shield_and_drain(asyncio.to_thread(self._save_store))
|
||||||
return True
|
return True
|
||||||
return False
|
return False
|
||||||
finally:
|
finally:
|
||||||
|
|||||||
@@ -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.mcp import MCPProvider
|
||||||
from nanobot.agent.tools.registry import ToolRegistry
|
from nanobot.agent.tools.registry import ToolRegistry
|
||||||
from nanobot.config.schema import Config
|
from nanobot.config.schema import Config
|
||||||
|
from nanobot.providers.base import LLMUsage
|
||||||
from nanobot.providers.image_generation import image_gen_provider_configs
|
from nanobot.providers.image_generation import image_gen_provider_configs
|
||||||
from nanobot.sdk.clients import MemoryClient, RuntimeClient, SessionClient
|
from nanobot.sdk.clients import MemoryClient, RuntimeClient, SessionClient
|
||||||
from nanobot.sdk.runtime import (
|
from nanobot.sdk.runtime import (
|
||||||
@@ -43,6 +44,7 @@ from nanobot.utils.llm_runtime import LLMRuntime
|
|||||||
|
|
||||||
__all__ = [
|
__all__ = [
|
||||||
"Nanobot",
|
"Nanobot",
|
||||||
|
"LLMUsage",
|
||||||
"RunResult",
|
"RunResult",
|
||||||
"RunStream",
|
"RunStream",
|
||||||
"SessionInfo",
|
"SessionInfo",
|
||||||
@@ -287,7 +289,7 @@ class Nanobot:
|
|||||||
type=STREAM_EVENT_RUN_COMPLETED,
|
type=STREAM_EVENT_RUN_COMPLETED,
|
||||||
content=result.content,
|
content=result.content,
|
||||||
result=result,
|
result=result,
|
||||||
usage=dict(result.usage),
|
usage=result.usage,
|
||||||
metadata=dict(result.metadata),
|
metadata=dict(result.metadata),
|
||||||
))
|
))
|
||||||
return result
|
return result
|
||||||
|
|||||||
@@ -96,22 +96,6 @@ class ManagedProcessRuntime(Generic[_StartOptionsT]):
|
|||||||
# it; poll() both reaps it and reports the real lifecycle state.
|
# it; poll() both reaps it and reports the real lifecycle state.
|
||||||
self._owned_process: Any | None = None
|
self._owned_process: Any | None = None
|
||||||
|
|
||||||
@classmethod
|
|
||||||
def refresh_state_pid(cls, *, paths: ProcessRuntimePaths) -> None:
|
|
||||||
"""Update a managed state file after the recorded process restarts."""
|
|
||||||
if not paths.state_path.exists():
|
|
||||||
return
|
|
||||||
try:
|
|
||||||
state = json.loads(paths.state_path.read_text(encoding="utf-8"))
|
|
||||||
except (json.JSONDecodeError, OSError):
|
|
||||||
return
|
|
||||||
state["pid"] = os.getpid()
|
|
||||||
runtime = cls(paths=paths)
|
|
||||||
state.pop("stable_identity", None)
|
|
||||||
state.update(runtime.process_identity_record(os.getpid()))
|
|
||||||
state["started_at"] = _utc_now()
|
|
||||||
runtime._write_state(state)
|
|
||||||
|
|
||||||
def start_background(self, options: _StartOptionsT) -> ProcessResult:
|
def start_background(self, options: _StartOptionsT) -> ProcessResult:
|
||||||
"""Start the configured command as a detached process."""
|
"""Start the configured command as a detached process."""
|
||||||
with self._lifecycle_lock():
|
with self._lifecycle_lock():
|
||||||
|
|||||||
@@ -5,11 +5,12 @@ from __future__ import annotations
|
|||||||
from importlib import import_module
|
from importlib import import_module
|
||||||
from typing import TYPE_CHECKING
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
from nanobot.providers.base import LLMProvider, LLMResponse
|
from nanobot.providers.base import LLMProvider, LLMResponse, LLMUsage
|
||||||
|
|
||||||
__all__ = [
|
__all__ = [
|
||||||
"LLMProvider",
|
"LLMProvider",
|
||||||
"LLMResponse",
|
"LLMResponse",
|
||||||
|
"LLMUsage",
|
||||||
"AnthropicProvider",
|
"AnthropicProvider",
|
||||||
"OpenAICompatProvider",
|
"OpenAICompatProvider",
|
||||||
"OpenAICodexProvider",
|
"OpenAICodexProvider",
|
||||||
|
|||||||
@@ -17,6 +17,7 @@ from loguru import logger
|
|||||||
from nanobot.providers.base import (
|
from nanobot.providers.base import (
|
||||||
LLMProvider,
|
LLMProvider,
|
||||||
LLMResponse,
|
LLMResponse,
|
||||||
|
LLMUsage,
|
||||||
ToolCallRequest,
|
ToolCallRequest,
|
||||||
resolve_stream_idle_timeout_s,
|
resolve_stream_idle_timeout_s,
|
||||||
tool_arguments_object_for_replay,
|
tool_arguments_object_for_replay,
|
||||||
@@ -90,8 +91,10 @@ class AnthropicProvider(LLMProvider):
|
|||||||
api_base: str | None = None,
|
api_base: str | None = None,
|
||||||
default_model: str = "claude-sonnet-4-6",
|
default_model: str = "claude-sonnet-4-6",
|
||||||
extra_headers: dict[str, str] | None = None,
|
extra_headers: dict[str, str] | None = None,
|
||||||
|
*,
|
||||||
|
provider_name: str = "anthropic",
|
||||||
):
|
):
|
||||||
super().__init__(api_key, api_base)
|
super().__init__(api_key, api_base, provider_name=provider_name)
|
||||||
self.default_model = default_model
|
self.default_model = default_model
|
||||||
self.extra_headers = extra_headers or {}
|
self.extra_headers = extra_headers or {}
|
||||||
|
|
||||||
@@ -689,24 +692,25 @@ class AnthropicProvider(LLMProvider):
|
|||||||
stop_map = {"tool_use": "tool_calls", "end_turn": "stop", "max_tokens": "length"}
|
stop_map = {"tool_use": "tool_calls", "end_turn": "stop", "max_tokens": "length"}
|
||||||
finish_reason = stop_map.get(response.stop_reason or "", response.stop_reason or "stop")
|
finish_reason = stop_map.get(response.stop_reason or "", response.stop_reason or "stop")
|
||||||
|
|
||||||
usage: dict[str, int] = {}
|
usage: LLMUsage | None = None
|
||||||
if response.usage:
|
if response.usage:
|
||||||
input_tokens = response.usage.input_tokens
|
cache_write_raw = getattr(
|
||||||
cache_creation = getattr(response.usage, "cache_creation_input_tokens", 0) or 0
|
response.usage,
|
||||||
cache_read = getattr(response.usage, "cache_read_input_tokens", 0) or 0
|
"cache_creation_input_tokens",
|
||||||
total_prompt_tokens = input_tokens + cache_creation + cache_read
|
None,
|
||||||
usage = {
|
)
|
||||||
"prompt_tokens": total_prompt_tokens,
|
cache_read_raw = getattr(response.usage, "cache_read_input_tokens", None)
|
||||||
"completion_tokens": response.usage.output_tokens,
|
cache_write = int(cache_write_raw) if cache_write_raw is not None else None
|
||||||
"total_tokens": total_prompt_tokens + response.usage.output_tokens,
|
cache_read = int(cache_read_raw) if cache_read_raw is not None else None
|
||||||
}
|
logical_input = int(response.usage.input_tokens) + (cache_write or 0) + (
|
||||||
for attr in ("cache_creation_input_tokens", "cache_read_input_tokens"):
|
cache_read or 0
|
||||||
val = getattr(response.usage, attr, 0)
|
)
|
||||||
if val:
|
usage = LLMUsage.reported(
|
||||||
usage[attr] = val
|
input_tokens=logical_input,
|
||||||
# Normalize to cached_tokens for downstream consistency.
|
output_tokens=int(response.usage.output_tokens),
|
||||||
if cache_read:
|
cache_read_tokens=cache_read,
|
||||||
usage["cached_tokens"] = cache_read
|
cache_write_tokens=cache_write,
|
||||||
|
)
|
||||||
|
|
||||||
return LLMResponse(
|
return LLMResponse(
|
||||||
content="".join(content_parts) or None,
|
content="".join(content_parts) or None,
|
||||||
|
|||||||
@@ -106,8 +106,10 @@ class AzureOpenAIProvider(LLMProvider):
|
|||||||
api_key: str = "",
|
api_key: str = "",
|
||||||
api_base: str = "",
|
api_base: str = "",
|
||||||
default_model: str = "gpt-5.2-chat",
|
default_model: str = "gpt-5.2-chat",
|
||||||
|
*,
|
||||||
|
provider_name: str = "azure_openai",
|
||||||
):
|
):
|
||||||
super().__init__(api_key, api_base)
|
super().__init__(api_key, api_base, provider_name=provider_name)
|
||||||
self.default_model = default_model
|
self.default_model = default_model
|
||||||
self._native_compaction_available = True
|
self._native_compaction_available = True
|
||||||
|
|
||||||
|
|||||||
+440
-9
@@ -6,6 +6,7 @@ import asyncio
|
|||||||
import json
|
import json
|
||||||
import os
|
import os
|
||||||
import re
|
import re
|
||||||
|
import time
|
||||||
from abc import ABC, abstractmethod
|
from abc import ABC, abstractmethod
|
||||||
from collections.abc import Awaitable, Callable
|
from collections.abc import Awaitable, Callable
|
||||||
from contextlib import suppress
|
from contextlib import suppress
|
||||||
@@ -13,19 +14,23 @@ from copy import deepcopy
|
|||||||
from dataclasses import dataclass, field
|
from dataclasses import dataclass, field
|
||||||
from datetime import datetime, timezone
|
from datetime import datetime, timezone
|
||||||
from email.utils import parsedate_to_datetime
|
from email.utils import parsedate_to_datetime
|
||||||
from typing import Any, cast
|
from typing import TYPE_CHECKING, Any, Literal, cast
|
||||||
|
|
||||||
import json_repair
|
import json_repair
|
||||||
from loguru import logger
|
from loguru import logger
|
||||||
|
|
||||||
from nanobot.utils.helpers import sanitize_surrogates_deep
|
from nanobot.utils.helpers import sanitize_surrogates_deep
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from nanobot.llm_usage.models import LLMCallRecord
|
||||||
|
|
||||||
STREAM_IDLE_TIMEOUT_ENV = "NANOBOT_STREAM_IDLE_TIMEOUT_S"
|
STREAM_IDLE_TIMEOUT_ENV = "NANOBOT_STREAM_IDLE_TIMEOUT_S"
|
||||||
DEFAULT_STREAM_IDLE_TIMEOUT_S = 90.0
|
DEFAULT_STREAM_IDLE_TIMEOUT_S = 90.0
|
||||||
MAX_STREAM_IDLE_TIMEOUT_S = 3600.0
|
MAX_STREAM_IDLE_TIMEOUT_S = 3600.0
|
||||||
RETRY_AFTER_BUFFER = 1
|
RETRY_AFTER_BUFFER = 1
|
||||||
|
|
||||||
RetryEventCallback = Callable[[str], Awaitable[None]]
|
RetryEventCallback = Callable[[str], Awaitable[None]]
|
||||||
|
LLMCallObserver = Callable[["LLMCallRecord"], None]
|
||||||
|
|
||||||
|
|
||||||
def resolve_stream_idle_timeout_s(
|
def resolve_stream_idle_timeout_s(
|
||||||
@@ -253,13 +258,298 @@ class ProviderCallContext:
|
|||||||
context_window_tokens: int | None = None
|
context_window_tokens: int | None = None
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True, slots=True)
|
||||||
|
class LLMUsage:
|
||||||
|
"""Canonical token usage reported by, or estimated for, one or more LLM calls.
|
||||||
|
|
||||||
|
``input_tokens`` is the logical input total and therefore includes cache reads
|
||||||
|
and writes. ``None`` cache counts mean the wire protocol did not report that
|
||||||
|
metric, while zero means it explicitly reported no cache activity.
|
||||||
|
|
||||||
|
``total_tokens`` preserves a provider-reported total when it exceeds the
|
||||||
|
visible input plus output (for example, hidden reasoning or tool usage). It
|
||||||
|
must be at least ``input_tokens + output_tokens``. The reported and estimated
|
||||||
|
totals partition it exactly, including after multi-call aggregation.
|
||||||
|
"""
|
||||||
|
|
||||||
|
input_tokens: int
|
||||||
|
output_tokens: int
|
||||||
|
total_tokens: int
|
||||||
|
cache_read_tokens: int | None = None
|
||||||
|
cache_write_tokens: int | None = None
|
||||||
|
reported_tokens: int = 0
|
||||||
|
estimated_tokens: int = 0
|
||||||
|
generation_ms: int = 0
|
||||||
|
measured_output_tokens: int = 0
|
||||||
|
ttft_ms: int = 0
|
||||||
|
timed_requests: int = 0
|
||||||
|
context_tokens: int | None = None
|
||||||
|
request_count: int = 0
|
||||||
|
|
||||||
|
def __post_init__(self) -> None:
|
||||||
|
token_fields = {
|
||||||
|
"input_tokens": self.input_tokens,
|
||||||
|
"output_tokens": self.output_tokens,
|
||||||
|
"total_tokens": self.total_tokens,
|
||||||
|
"reported_tokens": self.reported_tokens,
|
||||||
|
"estimated_tokens": self.estimated_tokens,
|
||||||
|
"generation_ms": self.generation_ms,
|
||||||
|
"measured_output_tokens": self.measured_output_tokens,
|
||||||
|
"ttft_ms": self.ttft_ms,
|
||||||
|
"timed_requests": self.timed_requests,
|
||||||
|
"request_count": self.request_count,
|
||||||
|
}
|
||||||
|
for name, value in token_fields.items():
|
||||||
|
runtime_value = cast(object, value)
|
||||||
|
if (
|
||||||
|
not isinstance(runtime_value, int)
|
||||||
|
or isinstance(runtime_value, bool)
|
||||||
|
or runtime_value < 0
|
||||||
|
):
|
||||||
|
raise ValueError(f"{name} must be a non-negative integer")
|
||||||
|
for name, value in (
|
||||||
|
("cache_read_tokens", self.cache_read_tokens),
|
||||||
|
("cache_write_tokens", self.cache_write_tokens),
|
||||||
|
("context_tokens", self.context_tokens),
|
||||||
|
):
|
||||||
|
runtime_value = cast(object, value)
|
||||||
|
if runtime_value is not None and (
|
||||||
|
not isinstance(runtime_value, int)
|
||||||
|
or isinstance(runtime_value, bool)
|
||||||
|
or runtime_value < 0
|
||||||
|
):
|
||||||
|
raise ValueError(f"{name} must be None or a non-negative integer")
|
||||||
|
|
||||||
|
visible_total = self.input_tokens + self.output_tokens
|
||||||
|
if self.total_tokens < visible_total:
|
||||||
|
raise ValueError("total_tokens must be at least input_tokens + output_tokens")
|
||||||
|
if self.reported_tokens + self.estimated_tokens != self.total_tokens:
|
||||||
|
raise ValueError("reported_tokens + estimated_tokens must equal total_tokens")
|
||||||
|
cache_total = (self.cache_read_tokens or 0) + (self.cache_write_tokens or 0)
|
||||||
|
if cache_total > self.input_tokens:
|
||||||
|
raise ValueError("cache token counts cannot exceed logical input_tokens")
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def reported(
|
||||||
|
cls,
|
||||||
|
*,
|
||||||
|
input_tokens: int,
|
||||||
|
output_tokens: int,
|
||||||
|
total_tokens: int | None = None,
|
||||||
|
cache_read_tokens: int | None = None,
|
||||||
|
cache_write_tokens: int | None = None,
|
||||||
|
) -> LLMUsage:
|
||||||
|
"""Build usage normalized from a provider response."""
|
||||||
|
visible_total = input_tokens + output_tokens
|
||||||
|
normalized_total = (
|
||||||
|
visible_total if total_tokens is None else max(visible_total, total_tokens)
|
||||||
|
)
|
||||||
|
return cls(
|
||||||
|
input_tokens=input_tokens,
|
||||||
|
output_tokens=output_tokens,
|
||||||
|
total_tokens=normalized_total,
|
||||||
|
cache_read_tokens=cache_read_tokens,
|
||||||
|
cache_write_tokens=cache_write_tokens,
|
||||||
|
reported_tokens=normalized_total,
|
||||||
|
context_tokens=input_tokens,
|
||||||
|
request_count=1,
|
||||||
|
)
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def estimated(cls, *, input_tokens: int, output_tokens: int) -> LLMUsage:
|
||||||
|
"""Build usage estimated locally because the provider omitted it."""
|
||||||
|
return cls(
|
||||||
|
input_tokens=input_tokens,
|
||||||
|
output_tokens=output_tokens,
|
||||||
|
total_tokens=input_tokens + output_tokens,
|
||||||
|
estimated_tokens=input_tokens + output_tokens,
|
||||||
|
context_tokens=input_tokens,
|
||||||
|
request_count=1,
|
||||||
|
)
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def empty_request(cls) -> LLMUsage:
|
||||||
|
"""Represent a completed model request with no measurable token usage."""
|
||||||
|
return cls(
|
||||||
|
input_tokens=0,
|
||||||
|
output_tokens=0,
|
||||||
|
total_tokens=0,
|
||||||
|
request_count=1,
|
||||||
|
)
|
||||||
|
|
||||||
|
@property
|
||||||
|
def source(self) -> Literal["reported", "estimated", "mixed"]:
|
||||||
|
if self.estimated_tokens == 0:
|
||||||
|
return "reported"
|
||||||
|
if self.reported_tokens == 0:
|
||||||
|
return "estimated"
|
||||||
|
return "mixed"
|
||||||
|
|
||||||
|
def with_timing(
|
||||||
|
self,
|
||||||
|
*,
|
||||||
|
generation_ms: int | None,
|
||||||
|
ttft_ms: int | None,
|
||||||
|
) -> LLMUsage:
|
||||||
|
"""Attach locally measured streaming telemetry to this usage value."""
|
||||||
|
return LLMUsage(
|
||||||
|
input_tokens=self.input_tokens,
|
||||||
|
output_tokens=self.output_tokens,
|
||||||
|
total_tokens=self.total_tokens,
|
||||||
|
cache_read_tokens=self.cache_read_tokens,
|
||||||
|
cache_write_tokens=self.cache_write_tokens,
|
||||||
|
reported_tokens=self.reported_tokens,
|
||||||
|
estimated_tokens=self.estimated_tokens,
|
||||||
|
generation_ms=max(0, generation_ms or 0),
|
||||||
|
measured_output_tokens=self.output_tokens if generation_ms is not None else 0,
|
||||||
|
ttft_ms=max(0, ttft_ms or 0),
|
||||||
|
timed_requests=1 if ttft_ms is not None else 0,
|
||||||
|
context_tokens=self.context_tokens,
|
||||||
|
request_count=self.request_count,
|
||||||
|
)
|
||||||
|
|
||||||
|
def __add__(self, other: LLMUsage) -> LLMUsage:
|
||||||
|
"""Aggregate calls without turning partially reported cache data into a count."""
|
||||||
|
|
||||||
|
def _sum_cache(left: int | None, right: int | None) -> int | None:
|
||||||
|
return left + right if left is not None and right is not None else None
|
||||||
|
|
||||||
|
return LLMUsage(
|
||||||
|
input_tokens=self.input_tokens + other.input_tokens,
|
||||||
|
output_tokens=self.output_tokens + other.output_tokens,
|
||||||
|
total_tokens=self.total_tokens + other.total_tokens,
|
||||||
|
cache_read_tokens=_sum_cache(self.cache_read_tokens, other.cache_read_tokens),
|
||||||
|
cache_write_tokens=_sum_cache(self.cache_write_tokens, other.cache_write_tokens),
|
||||||
|
reported_tokens=self.reported_tokens + other.reported_tokens,
|
||||||
|
estimated_tokens=self.estimated_tokens + other.estimated_tokens,
|
||||||
|
generation_ms=self.generation_ms + other.generation_ms,
|
||||||
|
measured_output_tokens=(
|
||||||
|
self.measured_output_tokens + other.measured_output_tokens
|
||||||
|
),
|
||||||
|
ttft_ms=self.ttft_ms + other.ttft_ms,
|
||||||
|
timed_requests=self.timed_requests + other.timed_requests,
|
||||||
|
context_tokens=(
|
||||||
|
other.context_tokens
|
||||||
|
if other.context_tokens is not None
|
||||||
|
else self.context_tokens
|
||||||
|
),
|
||||||
|
request_count=self.request_count + other.request_count,
|
||||||
|
)
|
||||||
|
|
||||||
|
def to_dict(self) -> dict[str, int | str | None]:
|
||||||
|
"""Serialize the canonical contract at JSON/persistence boundaries."""
|
||||||
|
return {
|
||||||
|
"input_tokens": self.input_tokens,
|
||||||
|
"output_tokens": self.output_tokens,
|
||||||
|
"total_tokens": self.total_tokens,
|
||||||
|
"cache_read_tokens": self.cache_read_tokens,
|
||||||
|
"cache_write_tokens": self.cache_write_tokens,
|
||||||
|
"reported_tokens": self.reported_tokens,
|
||||||
|
"estimated_tokens": self.estimated_tokens,
|
||||||
|
"source": self.source,
|
||||||
|
"generation_ms": self.generation_ms,
|
||||||
|
"measured_output_tokens": self.measured_output_tokens,
|
||||||
|
"ttft_ms": self.ttft_ms,
|
||||||
|
"timed_requests": self.timed_requests,
|
||||||
|
"context_tokens": self.context_tokens,
|
||||||
|
"request_count": self.request_count,
|
||||||
|
}
|
||||||
|
|
||||||
|
def to_turn_dict(self) -> dict[str, int]:
|
||||||
|
"""Project canonical usage into the compact WebUI/TUI per-turn shape."""
|
||||||
|
result: dict[str, int] = {
|
||||||
|
"prompt_tokens": self.input_tokens,
|
||||||
|
"completion_tokens": self.output_tokens,
|
||||||
|
"total_tokens": self.total_tokens,
|
||||||
|
"request_count": self.request_count,
|
||||||
|
"estimated_tokens": self.estimated_tokens,
|
||||||
|
}
|
||||||
|
if self.context_tokens is not None:
|
||||||
|
result["context_tokens"] = self.context_tokens
|
||||||
|
if self.cache_read_tokens is not None:
|
||||||
|
result["cached_tokens"] = self.cache_read_tokens
|
||||||
|
if self.cache_write_tokens is not None:
|
||||||
|
result["cache_write_tokens"] = self.cache_write_tokens
|
||||||
|
if self.generation_ms > 0 and self.measured_output_tokens > 0:
|
||||||
|
result["generation_ms"] = self.generation_ms
|
||||||
|
result["measured_completion_tokens"] = self.measured_output_tokens
|
||||||
|
if self.timed_requests > 0:
|
||||||
|
result["ttft_ms"] = self.ttft_ms
|
||||||
|
result["timed_requests"] = self.timed_requests
|
||||||
|
return result
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def from_dict(cls, value: object) -> LLMUsage | None:
|
||||||
|
"""Validate the exact first-party serialized contract."""
|
||||||
|
if not isinstance(value, dict):
|
||||||
|
return None
|
||||||
|
data = cast(dict[object, object], value)
|
||||||
|
integer_fields = (
|
||||||
|
"input_tokens",
|
||||||
|
"output_tokens",
|
||||||
|
"reported_tokens",
|
||||||
|
"estimated_tokens",
|
||||||
|
"generation_ms",
|
||||||
|
"measured_output_tokens",
|
||||||
|
"ttft_ms",
|
||||||
|
"timed_requests",
|
||||||
|
"request_count",
|
||||||
|
)
|
||||||
|
serialized_fields = {
|
||||||
|
*integer_fields,
|
||||||
|
"total_tokens",
|
||||||
|
"cache_read_tokens",
|
||||||
|
"cache_write_tokens",
|
||||||
|
"context_tokens",
|
||||||
|
"source",
|
||||||
|
}
|
||||||
|
if set(data) != serialized_fields:
|
||||||
|
return None
|
||||||
|
if any(
|
||||||
|
not isinstance(item := data.get(name), int) or isinstance(item, bool)
|
||||||
|
for name in integer_fields
|
||||||
|
):
|
||||||
|
return None
|
||||||
|
cache_read = data.get("cache_read_tokens")
|
||||||
|
cache_write = data.get("cache_write_tokens")
|
||||||
|
context_tokens = data.get("context_tokens")
|
||||||
|
total = data.get("total_tokens")
|
||||||
|
source = data.get("source")
|
||||||
|
if any(
|
||||||
|
item is not None and (not isinstance(item, int) or isinstance(item, bool))
|
||||||
|
for item in (cache_read, cache_write, context_tokens)
|
||||||
|
) or not isinstance(total, int) or isinstance(total, bool):
|
||||||
|
return None
|
||||||
|
try:
|
||||||
|
usage = cls(
|
||||||
|
input_tokens=cast(int, data["input_tokens"]),
|
||||||
|
output_tokens=cast(int, data["output_tokens"]),
|
||||||
|
total_tokens=total,
|
||||||
|
cache_read_tokens=cast(int | None, cache_read),
|
||||||
|
cache_write_tokens=cast(int | None, cache_write),
|
||||||
|
reported_tokens=cast(int, data["reported_tokens"]),
|
||||||
|
estimated_tokens=cast(int, data["estimated_tokens"]),
|
||||||
|
generation_ms=cast(int, data["generation_ms"]),
|
||||||
|
measured_output_tokens=cast(int, data["measured_output_tokens"]),
|
||||||
|
ttft_ms=cast(int, data["ttft_ms"]),
|
||||||
|
timed_requests=cast(int, data["timed_requests"]),
|
||||||
|
context_tokens=cast(int | None, context_tokens),
|
||||||
|
request_count=cast(int, data["request_count"]),
|
||||||
|
)
|
||||||
|
except (KeyError, ValueError):
|
||||||
|
return None
|
||||||
|
if source != usage.source:
|
||||||
|
return None
|
||||||
|
return usage
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
class LLMResponse:
|
class LLMResponse:
|
||||||
"""Response from an LLM provider."""
|
"""Response from an LLM provider."""
|
||||||
content: str | None
|
content: str | None
|
||||||
tool_calls: list[ToolCallRequest] = field(default_factory=list)
|
tool_calls: list[ToolCallRequest] = field(default_factory=list)
|
||||||
finish_reason: str = "stop"
|
finish_reason: str = "stop"
|
||||||
usage: dict[str, int] = field(default_factory=dict)
|
usage: LLMUsage | None = None
|
||||||
# Locally measured streaming telemetry. ``generation_ms`` excludes time to
|
# Locally measured streaming telemetry. ``generation_ms`` excludes time to
|
||||||
# first token and provider retry gaps; ``ttft_ms`` measures the first
|
# first token and provider retry gaps; ``ttft_ms`` measures the first
|
||||||
# streamed reasoning/content delta from request start. They stay separate
|
# streamed reasoning/content delta from request start. They stay separate
|
||||||
@@ -383,10 +673,109 @@ class LLMProvider(ABC):
|
|||||||
|
|
||||||
_SENTINEL = object()
|
_SENTINEL = object()
|
||||||
|
|
||||||
def __init__(self, api_key: str | None = None, api_base: str | None = None):
|
def __init__(
|
||||||
|
self,
|
||||||
|
api_key: str | None = None,
|
||||||
|
api_base: str | None = None,
|
||||||
|
*,
|
||||||
|
provider_name: str,
|
||||||
|
):
|
||||||
|
runtime_provider_name = cast(object, provider_name)
|
||||||
|
if not isinstance(runtime_provider_name, str) or not runtime_provider_name.strip():
|
||||||
|
raise ValueError("provider_name must be a non-empty configured identity")
|
||||||
self.api_key = api_key
|
self.api_key = api_key
|
||||||
self.api_base = api_base
|
self.api_base = api_base
|
||||||
|
self.provider_name = provider_name
|
||||||
self.generation: GenerationSettings = GenerationSettings()
|
self.generation: GenerationSettings = GenerationSettings()
|
||||||
|
self._llm_call_observer: LLMCallObserver | None = None
|
||||||
|
|
||||||
|
def set_llm_call_observer(self, observer: LLMCallObserver | None) -> None:
|
||||||
|
"""Attach a fail-open observer for each physical retry-managed call."""
|
||||||
|
self._llm_call_observer = observer
|
||||||
|
|
||||||
|
def _usage_for_call(
|
||||||
|
self,
|
||||||
|
response: LLMResponse,
|
||||||
|
kwargs: dict[str, Any],
|
||||||
|
) -> LLMUsage | None:
|
||||||
|
usage = response.usage
|
||||||
|
if usage is None or usage.total_tokens == 0:
|
||||||
|
if response.finish_reason in {"error", "cancelled"}:
|
||||||
|
return None
|
||||||
|
messages = kwargs.get("messages")
|
||||||
|
if not isinstance(messages, list):
|
||||||
|
return usage
|
||||||
|
tools_value = kwargs.get("tools")
|
||||||
|
tools = cast(list[dict[str, Any]], tools_value) if isinstance(tools_value, list) else None
|
||||||
|
model_value = kwargs.get("model")
|
||||||
|
model = model_value if isinstance(model_value, str) else self.get_default_model()
|
||||||
|
try:
|
||||||
|
from nanobot.utils.helpers import (
|
||||||
|
build_assistant_message,
|
||||||
|
estimate_message_tokens,
|
||||||
|
estimate_prompt_tokens_chain,
|
||||||
|
)
|
||||||
|
|
||||||
|
input_tokens, _ = estimate_prompt_tokens_chain(
|
||||||
|
self,
|
||||||
|
model,
|
||||||
|
cast(list[dict[str, Any]], messages),
|
||||||
|
tools,
|
||||||
|
)
|
||||||
|
assistant_message = build_assistant_message(
|
||||||
|
response.content or "",
|
||||||
|
tool_calls=[call.to_openai_tool_call() for call in response.tool_calls],
|
||||||
|
reasoning_content=response.reasoning_content,
|
||||||
|
thinking_blocks=response.thinking_blocks,
|
||||||
|
)
|
||||||
|
usage = LLMUsage.estimated(
|
||||||
|
input_tokens=max(0, input_tokens),
|
||||||
|
output_tokens=max(0, estimate_message_tokens(assistant_message)),
|
||||||
|
)
|
||||||
|
except Exception:
|
||||||
|
logger.exception("failed to estimate usage for {}", self.provider_name)
|
||||||
|
return usage
|
||||||
|
return usage.with_timing(
|
||||||
|
generation_ms=response.generation_ms,
|
||||||
|
ttft_ms=response.ttft_ms,
|
||||||
|
)
|
||||||
|
|
||||||
|
def _observe_llm_call(
|
||||||
|
self,
|
||||||
|
response: LLMResponse,
|
||||||
|
kwargs: dict[str, Any],
|
||||||
|
*,
|
||||||
|
started_at_ms: int,
|
||||||
|
started_at_ns: int,
|
||||||
|
stream: bool,
|
||||||
|
) -> LLMResponse:
|
||||||
|
observer = self._llm_call_observer
|
||||||
|
if observer is None:
|
||||||
|
return response
|
||||||
|
usage = self._usage_for_call(response, kwargs)
|
||||||
|
if usage is not None:
|
||||||
|
response.usage = usage
|
||||||
|
model_value = kwargs.get("model")
|
||||||
|
model = model_value if isinstance(model_value, str) and model_value else self.get_default_model()
|
||||||
|
try:
|
||||||
|
from nanobot.llm_usage.context import current_llm_usage_source
|
||||||
|
from nanobot.llm_usage.models import LLMCallRecord
|
||||||
|
|
||||||
|
observer(LLMCallRecord(
|
||||||
|
started_at_ms=started_at_ms,
|
||||||
|
duration_ms=max(0, (time.monotonic_ns() - started_at_ns) // 1_000_000),
|
||||||
|
provider=self.provider_name,
|
||||||
|
model=model,
|
||||||
|
source=current_llm_usage_source(),
|
||||||
|
stream=stream,
|
||||||
|
finish_reason=response.finish_reason,
|
||||||
|
usage=usage,
|
||||||
|
error_status_code=response.error_status_code,
|
||||||
|
error_kind=response.error_kind,
|
||||||
|
))
|
||||||
|
except Exception:
|
||||||
|
logger.exception("LLM call observer failed for {}", self.provider_name)
|
||||||
|
return response
|
||||||
|
|
||||||
def can_resume_conversation_state(
|
def can_resume_conversation_state(
|
||||||
self,
|
self,
|
||||||
@@ -773,18 +1162,39 @@ class LLMProvider(ABC):
|
|||||||
|
|
||||||
async def _safe_chat(self, **kwargs: Any) -> LLMResponse:
|
async def _safe_chat(self, **kwargs: Any) -> LLMResponse:
|
||||||
"""Call chat() and convert unexpected exceptions to error responses."""
|
"""Call chat() and convert unexpected exceptions to error responses."""
|
||||||
|
started_at_ms = time.time_ns() // 1_000_000
|
||||||
|
started_at_ns = time.monotonic_ns()
|
||||||
try:
|
try:
|
||||||
provider_context = kwargs.pop("provider_context", None)
|
provider_context = kwargs.pop("provider_context", None)
|
||||||
if isinstance(provider_context, ProviderCallContext):
|
if isinstance(provider_context, ProviderCallContext):
|
||||||
return await self.chat_with_context(
|
response = await self.chat_with_context(
|
||||||
provider_context=provider_context,
|
provider_context=provider_context,
|
||||||
**kwargs,
|
**kwargs,
|
||||||
)
|
)
|
||||||
return await self.chat(**kwargs)
|
else:
|
||||||
|
response = await self.chat(**kwargs)
|
||||||
except asyncio.CancelledError:
|
except asyncio.CancelledError:
|
||||||
|
self._observe_llm_call(
|
||||||
|
LLMResponse(
|
||||||
|
content=None,
|
||||||
|
finish_reason="cancelled",
|
||||||
|
error_kind="cancelled",
|
||||||
|
),
|
||||||
|
kwargs,
|
||||||
|
started_at_ms=started_at_ms,
|
||||||
|
started_at_ns=started_at_ns,
|
||||||
|
stream=False,
|
||||||
|
)
|
||||||
raise
|
raise
|
||||||
except Exception as exc:
|
except Exception as exc:
|
||||||
return LLMResponse(content=f"Error calling LLM: {exc}", finish_reason="error")
|
response = LLMResponse(content=f"Error calling LLM: {exc}", finish_reason="error")
|
||||||
|
return self._observe_llm_call(
|
||||||
|
response,
|
||||||
|
kwargs,
|
||||||
|
started_at_ms=started_at_ms,
|
||||||
|
started_at_ns=started_at_ns,
|
||||||
|
stream=False,
|
||||||
|
)
|
||||||
|
|
||||||
async def chat_stream(
|
async def chat_stream(
|
||||||
self,
|
self,
|
||||||
@@ -847,18 +1257,39 @@ class LLMProvider(ABC):
|
|||||||
|
|
||||||
async def _safe_chat_stream(self, **kwargs: Any) -> LLMResponse:
|
async def _safe_chat_stream(self, **kwargs: Any) -> LLMResponse:
|
||||||
"""Call chat_stream() and convert unexpected exceptions to error responses."""
|
"""Call chat_stream() and convert unexpected exceptions to error responses."""
|
||||||
|
started_at_ms = time.time_ns() // 1_000_000
|
||||||
|
started_at_ns = time.monotonic_ns()
|
||||||
try:
|
try:
|
||||||
provider_context = kwargs.pop("provider_context", None)
|
provider_context = kwargs.pop("provider_context", None)
|
||||||
if isinstance(provider_context, ProviderCallContext):
|
if isinstance(provider_context, ProviderCallContext):
|
||||||
return await self.chat_stream_with_context(
|
response = await self.chat_stream_with_context(
|
||||||
provider_context=provider_context,
|
provider_context=provider_context,
|
||||||
**kwargs,
|
**kwargs,
|
||||||
)
|
)
|
||||||
return await self.chat_stream(**kwargs)
|
else:
|
||||||
|
response = await self.chat_stream(**kwargs)
|
||||||
except asyncio.CancelledError:
|
except asyncio.CancelledError:
|
||||||
|
self._observe_llm_call(
|
||||||
|
LLMResponse(
|
||||||
|
content=None,
|
||||||
|
finish_reason="cancelled",
|
||||||
|
error_kind="cancelled",
|
||||||
|
),
|
||||||
|
kwargs,
|
||||||
|
started_at_ms=started_at_ms,
|
||||||
|
started_at_ns=started_at_ns,
|
||||||
|
stream=True,
|
||||||
|
)
|
||||||
raise
|
raise
|
||||||
except Exception as exc:
|
except Exception as exc:
|
||||||
return LLMResponse(content=f"Error calling LLM: {exc}", finish_reason="error")
|
response = LLMResponse(content=f"Error calling LLM: {exc}", finish_reason="error")
|
||||||
|
return self._observe_llm_call(
|
||||||
|
response,
|
||||||
|
kwargs,
|
||||||
|
started_at_ms=started_at_ms,
|
||||||
|
started_at_ns=started_at_ns,
|
||||||
|
stream=True,
|
||||||
|
)
|
||||||
|
|
||||||
async def chat_stream_with_retry(
|
async def chat_stream_with_retry(
|
||||||
self,
|
self,
|
||||||
|
|||||||
@@ -14,6 +14,7 @@ from typing import Any, cast
|
|||||||
from nanobot.providers.base import (
|
from nanobot.providers.base import (
|
||||||
LLMProvider,
|
LLMProvider,
|
||||||
LLMResponse,
|
LLMResponse,
|
||||||
|
LLMUsage,
|
||||||
ToolCallRequest,
|
ToolCallRequest,
|
||||||
parse_tool_arguments,
|
parse_tool_arguments,
|
||||||
resolve_stream_idle_timeout_s,
|
resolve_stream_idle_timeout_s,
|
||||||
@@ -60,8 +61,9 @@ class BedrockProvider(LLMProvider):
|
|||||||
profile: str | None = None,
|
profile: str | None = None,
|
||||||
extra_body: dict[str, Any] | None = None,
|
extra_body: dict[str, Any] | None = None,
|
||||||
client: Any | None = None,
|
client: Any | None = None,
|
||||||
|
provider_name: str = "bedrock",
|
||||||
):
|
):
|
||||||
super().__init__(api_key, api_base)
|
super().__init__(api_key, api_base, provider_name=provider_name)
|
||||||
self.default_model = default_model
|
self.default_model = default_model
|
||||||
self.region = region or os.environ.get("AWS_REGION") or os.environ.get("AWS_DEFAULT_REGION")
|
self.region = region or os.environ.get("AWS_REGION") or os.environ.get("AWS_DEFAULT_REGION")
|
||||||
self.profile = profile
|
self.profile = profile
|
||||||
@@ -453,25 +455,25 @@ class BedrockProvider(LLMProvider):
|
|||||||
}.get(stop_reason or "", stop_reason or "stop")
|
}.get(stop_reason or "", stop_reason or "stop")
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def _usage(usage: dict[str, Any] | None) -> dict[str, int]:
|
def _usage(usage: dict[str, Any] | None) -> LLMUsage | None:
|
||||||
if not usage:
|
if not usage:
|
||||||
return {}
|
return None
|
||||||
prompt = int(usage.get("inputTokens") or 0)
|
|
||||||
completion = int(usage.get("outputTokens") or 0)
|
def _optional_count(key: str) -> int | None:
|
||||||
total = int(usage.get("totalTokens") or prompt + completion)
|
raw = usage.get(key)
|
||||||
result = {
|
return int(raw) if raw is not None else None
|
||||||
"prompt_tokens": prompt,
|
|
||||||
"completion_tokens": completion,
|
cache_read = _optional_count("cacheReadInputTokens")
|
||||||
"total_tokens": total,
|
cache_write = _optional_count("cacheWriteInputTokens")
|
||||||
}
|
logical_input = int(usage.get("inputTokens") or 0) + (cache_read or 0) + (
|
||||||
cache_read = int(usage.get("cacheReadInputTokens") or 0)
|
cache_write or 0
|
||||||
cache_write = int(usage.get("cacheWriteInputTokens") or 0)
|
)
|
||||||
if cache_read:
|
return LLMUsage.reported(
|
||||||
result["cached_tokens"] = cache_read
|
input_tokens=logical_input,
|
||||||
result["cache_read_input_tokens"] = cache_read
|
output_tokens=int(usage.get("outputTokens") or 0),
|
||||||
if cache_write:
|
cache_read_tokens=cache_read,
|
||||||
result["cache_creation_input_tokens"] = cache_write
|
cache_write_tokens=cache_write,
|
||||||
return result
|
)
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def _parse_reasoning(block: dict[str, Any]) -> tuple[str | None, dict[str, Any] | None]:
|
def _parse_reasoning(block: dict[str, Any]) -> tuple[str | None, dict[str, Any] | None]:
|
||||||
|
|||||||
@@ -172,6 +172,7 @@ def _make_provider_core(
|
|||||||
default_model=model,
|
default_model=model,
|
||||||
proxy=getattr(p, "proxy", None) if p else None,
|
proxy=getattr(p, "proxy", None) if p else None,
|
||||||
extra_body=p.extra_body if p else None,
|
extra_body=p.extra_body if p else None,
|
||||||
|
provider_name=provider_name,
|
||||||
)
|
)
|
||||||
elif backend == "xai_grok":
|
elif backend == "xai_grok":
|
||||||
from nanobot.providers.xai_grok_provider import XAIGrokProvider
|
from nanobot.providers.xai_grok_provider import XAIGrokProvider
|
||||||
@@ -180,6 +181,7 @@ def _make_provider_core(
|
|||||||
default_model=model,
|
default_model=model,
|
||||||
proxy=getattr(p, "proxy", None) if p else None,
|
proxy=getattr(p, "proxy", None) if p else None,
|
||||||
extra_body=p.extra_body if p else None,
|
extra_body=p.extra_body if p else None,
|
||||||
|
provider_name=provider_name,
|
||||||
)
|
)
|
||||||
elif backend == "azure_openai":
|
elif backend == "azure_openai":
|
||||||
from nanobot.providers.azure_openai_provider import AzureOpenAIProvider
|
from nanobot.providers.azure_openai_provider import AzureOpenAIProvider
|
||||||
@@ -190,11 +192,12 @@ def _make_provider_core(
|
|||||||
api_key=p.api_key or "",
|
api_key=p.api_key or "",
|
||||||
api_base=p.api_base,
|
api_base=p.api_base,
|
||||||
default_model=model,
|
default_model=model,
|
||||||
|
provider_name=provider_name,
|
||||||
)
|
)
|
||||||
elif backend == "github_copilot":
|
elif backend == "github_copilot":
|
||||||
from nanobot.providers.github_copilot_provider import GitHubCopilotProvider
|
from nanobot.providers.github_copilot_provider import GitHubCopilotProvider
|
||||||
|
|
||||||
provider = GitHubCopilotProvider(default_model=model)
|
provider = GitHubCopilotProvider(default_model=model, provider_name=provider_name)
|
||||||
elif backend == "anthropic":
|
elif backend == "anthropic":
|
||||||
from nanobot.providers.anthropic_provider import AnthropicProvider
|
from nanobot.providers.anthropic_provider import AnthropicProvider
|
||||||
|
|
||||||
@@ -203,6 +206,7 @@ def _make_provider_core(
|
|||||||
api_base=config.get_api_base(model, preset=preset),
|
api_base=config.get_api_base(model, preset=preset),
|
||||||
default_model=model,
|
default_model=model,
|
||||||
extra_headers=_provider_extra_headers(spec, p),
|
extra_headers=_provider_extra_headers(spec, p),
|
||||||
|
provider_name=provider_name,
|
||||||
)
|
)
|
||||||
elif backend == "bedrock":
|
elif backend == "bedrock":
|
||||||
from nanobot.providers.bedrock_provider import BedrockProvider
|
from nanobot.providers.bedrock_provider import BedrockProvider
|
||||||
@@ -214,6 +218,7 @@ def _make_provider_core(
|
|||||||
region=getattr(p, "region", None) if p else None,
|
region=getattr(p, "region", None) if p else None,
|
||||||
profile=getattr(p, "profile", None) if p else None,
|
profile=getattr(p, "profile", None) if p else None,
|
||||||
extra_body=p.extra_body if p else None,
|
extra_body=p.extra_body if p else None,
|
||||||
|
provider_name=provider_name,
|
||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
from nanobot.providers.openai_compat_provider import OpenAICompatProvider
|
from nanobot.providers.openai_compat_provider import OpenAICompatProvider
|
||||||
@@ -228,6 +233,7 @@ def _make_provider_core(
|
|||||||
api_type=p.api_type if p and provider_name == "openai" else "auto",
|
api_type=p.api_type if p and provider_name == "openai" else "auto",
|
||||||
extra_query=p.extra_query if p else None,
|
extra_query=p.extra_query if p else None,
|
||||||
proxy=p.proxy if p else None,
|
proxy=p.proxy if p else None,
|
||||||
|
provider_name=provider_name,
|
||||||
)
|
)
|
||||||
|
|
||||||
provider.generation = preset.to_generation_settings()
|
provider.generation = preset.to_generation_settings()
|
||||||
|
|||||||
@@ -13,6 +13,7 @@ from loguru import logger
|
|||||||
|
|
||||||
from nanobot.providers.base import (
|
from nanobot.providers.base import (
|
||||||
GenerationSettings,
|
GenerationSettings,
|
||||||
|
LLMCallObserver,
|
||||||
LLMProvider,
|
LLMProvider,
|
||||||
LLMResponse,
|
LLMResponse,
|
||||||
ProviderCallContext,
|
ProviderCallContext,
|
||||||
@@ -124,7 +125,10 @@ class FallbackProvider(LLMProvider):
|
|||||||
fallback_model_observer: FallbackModelObserver | None = None,
|
fallback_model_observer: FallbackModelObserver | None = None,
|
||||||
primary_context_window_tokens: int | None = None,
|
primary_context_window_tokens: int | None = None,
|
||||||
):
|
):
|
||||||
|
primary_generation = primary.generation
|
||||||
self._primary = primary
|
self._primary = primary
|
||||||
|
super().__init__(provider_name=primary.provider_name)
|
||||||
|
self._primary.generation = primary_generation
|
||||||
self._fallback_presets = list(fallback_presets)
|
self._fallback_presets = list(fallback_presets)
|
||||||
self._provider_factory = provider_factory
|
self._provider_factory = provider_factory
|
||||||
self._fallback_model_observer = fallback_model_observer
|
self._fallback_model_observer = fallback_model_observer
|
||||||
@@ -148,6 +152,11 @@ class FallbackProvider(LLMProvider):
|
|||||||
"""Attach a process-level observer without changing request call signatures."""
|
"""Attach a process-level observer without changing request call signatures."""
|
||||||
self._fallback_model_observer = observer
|
self._fallback_model_observer = observer
|
||||||
|
|
||||||
|
def set_llm_call_observer(self, observer: LLMCallObserver | None) -> None:
|
||||||
|
"""Attach usage recording to the primary and future fallback leaves."""
|
||||||
|
super().set_llm_call_observer(observer)
|
||||||
|
self._primary.set_llm_call_observer(observer)
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def supports_progress_deltas(self) -> bool:
|
def supports_progress_deltas(self) -> bool:
|
||||||
return bool(getattr(self._primary, "supports_progress_deltas", False))
|
return bool(getattr(self._primary, "supports_progress_deltas", False))
|
||||||
@@ -503,6 +512,7 @@ class FallbackProvider(LLMProvider):
|
|||||||
)
|
)
|
||||||
try:
|
try:
|
||||||
fallback_provider = self._provider_factory(fallback)
|
fallback_provider = self._provider_factory(fallback)
|
||||||
|
fallback_provider.set_llm_call_observer(self._llm_call_observer)
|
||||||
except Exception as exc:
|
except Exception as exc:
|
||||||
logger.warning(
|
logger.warning(
|
||||||
"Failed to create provider for fallback '{}': {}", fallback_model, exc
|
"Failed to create provider for fallback '{}': {}", fallback_model, exc
|
||||||
|
|||||||
@@ -174,7 +174,12 @@ def login_github_copilot(
|
|||||||
class GitHubCopilotProvider(OpenAICompatProvider):
|
class GitHubCopilotProvider(OpenAICompatProvider):
|
||||||
"""Provider that exchanges a stored GitHub OAuth token for Copilot access tokens."""
|
"""Provider that exchanges a stored GitHub OAuth token for Copilot access tokens."""
|
||||||
|
|
||||||
def __init__(self, default_model: str = "github-copilot/gpt-4.1"):
|
def __init__(
|
||||||
|
self,
|
||||||
|
default_model: str = "github-copilot/gpt-4.1",
|
||||||
|
*,
|
||||||
|
provider_name: str = "github_copilot",
|
||||||
|
):
|
||||||
from nanobot.providers.registry import find_by_name
|
from nanobot.providers.registry import find_by_name
|
||||||
|
|
||||||
self._copilot_access_token: str | None = None
|
self._copilot_access_token: str | None = None
|
||||||
@@ -190,6 +195,7 @@ class GitHubCopilotProvider(OpenAICompatProvider):
|
|||||||
"User-Agent": USER_AGENT,
|
"User-Agent": USER_AGENT,
|
||||||
},
|
},
|
||||||
spec=find_by_name("github_copilot"),
|
spec=find_by_name("github_copilot"),
|
||||||
|
provider_name=provider_name,
|
||||||
)
|
)
|
||||||
|
|
||||||
async def _get_copilot_access_token(self) -> str:
|
async def _get_copilot_access_token(self) -> str:
|
||||||
|
|||||||
@@ -20,7 +20,7 @@ from nanobot.providers.registry import find_by_name
|
|||||||
from nanobot.security.network import (
|
from nanobot.security.network import (
|
||||||
PinnedDNSAsyncTransport,
|
PinnedDNSAsyncTransport,
|
||||||
UnsafeURLRequestError,
|
UnsafeURLRequestError,
|
||||||
resolve_url_target,
|
async_resolve_url_target,
|
||||||
)
|
)
|
||||||
from nanobot.utils.helpers import detect_image_mime
|
from nanobot.utils.helpers import detect_image_mime
|
||||||
|
|
||||||
@@ -174,7 +174,7 @@ async def _download_image_data_url(
|
|||||||
current_url = url
|
current_url = url
|
||||||
for _ in range(_IMAGE_DOWNLOAD_MAX_REDIRECTS + 1):
|
for _ in range(_IMAGE_DOWNLOAD_MAX_REDIRECTS + 1):
|
||||||
if proxy:
|
if proxy:
|
||||||
ok, error, _ = resolve_url_target(
|
ok, error, _ = await async_resolve_url_target(
|
||||||
current_url,
|
current_url,
|
||||||
trust_remote_dns=True,
|
trust_remote_dns=True,
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -7,6 +7,7 @@ from __future__ import annotations
|
|||||||
import asyncio
|
import asyncio
|
||||||
import hashlib
|
import hashlib
|
||||||
import json
|
import json
|
||||||
|
import ssl
|
||||||
from collections.abc import Awaitable, Callable
|
from collections.abc import Awaitable, Callable
|
||||||
from typing import Any, cast
|
from typing import Any, cast
|
||||||
|
|
||||||
@@ -50,12 +51,26 @@ class OpenAICodexProvider(LLMProvider):
|
|||||||
default_model: str = "openai-codex/gpt-5.6-sol",
|
default_model: str = "openai-codex/gpt-5.6-sol",
|
||||||
proxy: str | None = None,
|
proxy: str | None = None,
|
||||||
extra_body: dict[str, Any] | None = None,
|
extra_body: dict[str, Any] | None = None,
|
||||||
|
*,
|
||||||
|
provider_name: str = "openai_codex",
|
||||||
):
|
):
|
||||||
super().__init__(api_key=None, api_base=None)
|
super().__init__(api_key=None, api_base=None, provider_name=provider_name)
|
||||||
self.default_model = default_model
|
self.default_model = default_model
|
||||||
self.proxy = proxy or None
|
self.proxy = proxy or None
|
||||||
self._extra_body = dict(extra_body or {})
|
self._extra_body = dict(extra_body or {})
|
||||||
self._native_compaction_available = True
|
self._native_compaction_available = True
|
||||||
|
self._ssl_contexts: dict[bool, ssl.SSLContext] = {}
|
||||||
|
|
||||||
|
def _ssl_context(self, *, verify: bool) -> ssl.SSLContext:
|
||||||
|
"""Reuse synchronous TLS setup across requests on the shared event loop."""
|
||||||
|
context = self._ssl_contexts.get(verify)
|
||||||
|
if context is None:
|
||||||
|
context = httpx.create_ssl_context(
|
||||||
|
verify=verify,
|
||||||
|
trust_env=self.proxy is None,
|
||||||
|
)
|
||||||
|
self._ssl_contexts[verify] = context
|
||||||
|
return context
|
||||||
|
|
||||||
async def _call_codex(
|
async def _call_codex(
|
||||||
self,
|
self,
|
||||||
@@ -129,7 +144,7 @@ class OpenAICodexProvider(LLMProvider):
|
|||||||
DEFAULT_CODEX_URL,
|
DEFAULT_CODEX_URL,
|
||||||
headers,
|
headers,
|
||||||
wire_body,
|
wire_body,
|
||||||
verify=True,
|
verify=self._ssl_context(verify=True),
|
||||||
proxy=self.proxy,
|
proxy=self.proxy,
|
||||||
on_content_delta=on_content_delta if emit_deltas else None,
|
on_content_delta=on_content_delta if emit_deltas else None,
|
||||||
on_thinking_delta=on_thinking_delta if emit_deltas else None,
|
on_thinking_delta=on_thinking_delta if emit_deltas else None,
|
||||||
@@ -145,7 +160,7 @@ class OpenAICodexProvider(LLMProvider):
|
|||||||
DEFAULT_CODEX_URL,
|
DEFAULT_CODEX_URL,
|
||||||
headers,
|
headers,
|
||||||
wire_body,
|
wire_body,
|
||||||
verify=False,
|
verify=self._ssl_context(verify=False),
|
||||||
proxy=self.proxy,
|
proxy=self.proxy,
|
||||||
on_content_delta=on_content_delta if emit_deltas else None,
|
on_content_delta=on_content_delta if emit_deltas else None,
|
||||||
on_thinking_delta=on_thinking_delta if emit_deltas else None,
|
on_thinking_delta=on_thinking_delta if emit_deltas else None,
|
||||||
@@ -411,7 +426,7 @@ async def _request_codex(
|
|||||||
url: str,
|
url: str,
|
||||||
headers: dict[str, str],
|
headers: dict[str, str],
|
||||||
body: dict[str, Any],
|
body: dict[str, Any],
|
||||||
verify: bool,
|
verify: ssl.SSLContext | bool,
|
||||||
proxy: str | None = None,
|
proxy: str | None = None,
|
||||||
on_content_delta: Callable[[str], Awaitable[None]] | None = None,
|
on_content_delta: Callable[[str], Awaitable[None]] | None = None,
|
||||||
on_thinking_delta: Callable[[str], Awaitable[None]] | None = None,
|
on_thinking_delta: Callable[[str], Awaitable[None]] | None = None,
|
||||||
|
|||||||
@@ -26,6 +26,7 @@ from pydantic.alias_generators import to_snake
|
|||||||
from nanobot.providers.base import (
|
from nanobot.providers.base import (
|
||||||
LLMProvider,
|
LLMProvider,
|
||||||
LLMResponse,
|
LLMResponse,
|
||||||
|
LLMUsage,
|
||||||
ProviderCallContext,
|
ProviderCallContext,
|
||||||
ProviderConversationState,
|
ProviderConversationState,
|
||||||
ToolCallRequest,
|
ToolCallRequest,
|
||||||
@@ -517,8 +518,9 @@ class OpenAICompatProvider(LLMProvider):
|
|||||||
api_type: str = "auto",
|
api_type: str = "auto",
|
||||||
extra_query: dict[str, str] | None = None,
|
extra_query: dict[str, str] | None = None,
|
||||||
proxy: str | None = None,
|
proxy: str | None = None,
|
||||||
|
provider_name: str = "openai",
|
||||||
):
|
):
|
||||||
super().__init__(api_key, api_base)
|
super().__init__(api_key, api_base, provider_name=provider_name)
|
||||||
self.default_model = default_model
|
self.default_model = default_model
|
||||||
self.extra_headers = extra_headers or {}
|
self.extra_headers = extra_headers or {}
|
||||||
self._spec = spec
|
self._spec = spec
|
||||||
@@ -1428,12 +1430,12 @@ class OpenAICompatProvider(LLMProvider):
|
|||||||
return "".join(parts) or None
|
return "".join(parts) or None
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def _extract_usage(cls, response: Any) -> dict[str, int]:
|
def _extract_usage(cls, response: Any) -> LLMUsage | None:
|
||||||
"""Extract token usage from an OpenAI-compatible response.
|
"""Extract token usage from an OpenAI-compatible response.
|
||||||
|
|
||||||
Handles both dict-based (raw JSON) and object-based (SDK Pydantic)
|
Handles both dict-based (raw JSON) and object-based (SDK Pydantic)
|
||||||
responses. Provider-specific ``cached_tokens`` fields are normalised
|
responses. Provider-specific cache fields are normalized once at
|
||||||
under a single key; see the priority chain inside for details.
|
this Chat Completions wire boundary.
|
||||||
"""
|
"""
|
||||||
# --- resolve usage object ---
|
# --- resolve usage object ---
|
||||||
usage_obj = None
|
usage_obj = None
|
||||||
@@ -1445,21 +1447,18 @@ class OpenAICompatProvider(LLMProvider):
|
|||||||
|
|
||||||
usage_map = cls._maybe_mapping(usage_obj)
|
usage_map = cls._maybe_mapping(usage_obj)
|
||||||
if usage_map is not None:
|
if usage_map is not None:
|
||||||
result = {
|
input_tokens = int(usage_map.get("prompt_tokens") or 0)
|
||||||
"prompt_tokens": int(usage_map.get("prompt_tokens") or 0),
|
output_tokens = int(usage_map.get("completion_tokens") or 0)
|
||||||
"completion_tokens": int(usage_map.get("completion_tokens") or 0),
|
|
||||||
"total_tokens": int(usage_map.get("total_tokens") or 0),
|
|
||||||
}
|
|
||||||
elif usage_obj:
|
elif usage_obj:
|
||||||
result = {
|
input_tokens = int(getattr(usage_obj, "prompt_tokens", 0) or 0)
|
||||||
"prompt_tokens": getattr(usage_obj, "prompt_tokens", 0) or 0,
|
output_tokens = int(getattr(usage_obj, "completion_tokens", 0) or 0)
|
||||||
"completion_tokens": getattr(usage_obj, "completion_tokens", 0) or 0,
|
|
||||||
"total_tokens": getattr(usage_obj, "total_tokens", 0) or 0,
|
|
||||||
}
|
|
||||||
else:
|
else:
|
||||||
return {}
|
return None
|
||||||
|
|
||||||
# --- cached_tokens (normalised across providers) ---
|
wire_total = cls._get_nested_int(usage_obj, ("total_tokens",))
|
||||||
|
|
||||||
|
cache_read: int | None = None
|
||||||
|
# --- cached_tokens (normalised across Chat-compatible providers) ---
|
||||||
# Try nested paths first (dict), fall back to attribute (SDK object).
|
# Try nested paths first (dict), fall back to attribute (SDK object).
|
||||||
# Priority order ensures the most specific field wins.
|
# Priority order ensures the most specific field wins.
|
||||||
for path in (
|
for path in (
|
||||||
@@ -1468,17 +1467,28 @@ class OpenAICompatProvider(LLMProvider):
|
|||||||
("prompt_cache_hit_tokens",), # DeepSeek/SiliconFlow
|
("prompt_cache_hit_tokens",), # DeepSeek/SiliconFlow
|
||||||
):
|
):
|
||||||
cached = cls._get_nested_int(usage_map, path)
|
cached = cls._get_nested_int(usage_map, path)
|
||||||
if not cached and usage_obj:
|
if cached is None and usage_obj:
|
||||||
cached = cls._get_nested_int(usage_obj, path)
|
cached = cls._get_nested_int(usage_obj, path)
|
||||||
if cached:
|
if cached is not None:
|
||||||
result["cached_tokens"] = cached
|
cache_read = cached
|
||||||
break
|
break
|
||||||
|
|
||||||
return result
|
cache_write = cls._get_nested_int(
|
||||||
|
usage_obj,
|
||||||
|
("prompt_tokens_details", "cache_write_tokens"),
|
||||||
|
)
|
||||||
|
|
||||||
|
return LLMUsage.reported(
|
||||||
|
input_tokens=input_tokens,
|
||||||
|
output_tokens=output_tokens,
|
||||||
|
total_tokens=wire_total,
|
||||||
|
cache_read_tokens=cache_read,
|
||||||
|
cache_write_tokens=cache_write,
|
||||||
|
)
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def _get_nested_int(obj: object, path: tuple[str, ...]) -> int:
|
def _get_nested_int(obj: object, path: tuple[str, ...]) -> int | None:
|
||||||
"""Drill into *obj* by *path* segments and return an ``int`` value.
|
"""Return a present usage count while preserving explicit zero.
|
||||||
|
|
||||||
Supports both dict-key access and attribute access so it works
|
Supports both dict-key access and attribute access so it works
|
||||||
uniformly with raw JSON dicts **and** SDK Pydantic models.
|
uniformly with raw JSON dicts **and** SDK Pydantic models.
|
||||||
@@ -1486,12 +1496,17 @@ class OpenAICompatProvider(LLMProvider):
|
|||||||
current: object = obj
|
current: object = obj
|
||||||
for segment in path:
|
for segment in path:
|
||||||
if current is None:
|
if current is None:
|
||||||
return 0
|
return None
|
||||||
if isinstance(current, dict):
|
if isinstance(current, dict):
|
||||||
current = cast(dict[str, Any], current).get(segment)
|
current = cast(dict[str, Any], current).get(segment)
|
||||||
else:
|
else:
|
||||||
current = getattr(current, segment, None)
|
current = getattr(current, segment, None)
|
||||||
return int(cast(Any, current) or 0) if current is not None else 0
|
if current is None or isinstance(current, bool):
|
||||||
|
return None
|
||||||
|
try:
|
||||||
|
return int(cast(Any, current))
|
||||||
|
except (TypeError, ValueError):
|
||||||
|
return None
|
||||||
|
|
||||||
def _parse(self, response: Any) -> LLMResponse:
|
def _parse(self, response: Any) -> LLMResponse:
|
||||||
if isinstance(response, str):
|
if isinstance(response, str):
|
||||||
@@ -1645,7 +1660,7 @@ class OpenAICompatProvider(LLMProvider):
|
|||||||
reasoning_parts: list[str] = []
|
reasoning_parts: list[str] = []
|
||||||
tc_bufs: dict[int, dict[str, Any]] = {}
|
tc_bufs: dict[int, dict[str, Any]] = {}
|
||||||
finish_reason = "stop"
|
finish_reason = "stop"
|
||||||
usage: dict[str, int] = {}
|
usage: LLMUsage | None = None
|
||||||
|
|
||||||
def _accum_tc(tc: Any, idx_hint: int) -> None:
|
def _accum_tc(tc: Any, idx_hint: int) -> None:
|
||||||
"""Accumulate one streaming tool-call delta into *tc_bufs*."""
|
"""Accumulate one streaming tool-call delta into *tc_bufs*."""
|
||||||
|
|||||||
@@ -10,7 +10,7 @@ from typing import Any, AsyncGenerator, cast
|
|||||||
import httpx
|
import httpx
|
||||||
from loguru import logger
|
from loguru import logger
|
||||||
|
|
||||||
from nanobot.providers.base import LLMResponse, ToolCallRequest, parse_tool_arguments
|
from nanobot.providers.base import LLMResponse, LLMUsage, ToolCallRequest, parse_tool_arguments
|
||||||
from nanobot.providers.openai_responses.state import build_responses_state
|
from nanobot.providers.openai_responses.state import build_responses_state
|
||||||
|
|
||||||
FINISH_REASON_MAP = {
|
FINISH_REASON_MAP = {
|
||||||
@@ -186,33 +186,40 @@ def _response_finish_reason(
|
|||||||
return map_finish_reason(terminal_status)
|
return map_finish_reason(terminal_status)
|
||||||
|
|
||||||
|
|
||||||
def _usage_from_response_obj(response: object) -> dict[str, int]:
|
def _usage_from_response_obj(response: object) -> LLMUsage | None:
|
||||||
response_object = _response_object(response)
|
response_object = _response_object(response)
|
||||||
usage_raw: object = (
|
usage_raw: object = (
|
||||||
response_object.get("usage")
|
response_object.get("usage")
|
||||||
if response_object is not None
|
if response_object is not None
|
||||||
else getattr(response, "usage", None)
|
else getattr(response, "usage", None)
|
||||||
)
|
)
|
||||||
if not usage_raw:
|
if usage_raw is None:
|
||||||
return {}
|
return None
|
||||||
usage = _response_object(usage_raw)
|
usage = _response_object(usage_raw)
|
||||||
if usage is None:
|
if usage is None:
|
||||||
return {}
|
return None
|
||||||
prompt_tokens = int(usage.get("input_tokens") or usage.get("prompt_tokens") or 0)
|
|
||||||
completion_tokens = int(
|
def _usage_int(container: dict[str, Any] | None, key: str) -> int | None:
|
||||||
usage.get("output_tokens") or usage.get("completion_tokens") or 0
|
if container is None:
|
||||||
)
|
return None
|
||||||
total_tokens = int(usage.get("total_tokens") or prompt_tokens + completion_tokens)
|
raw = container.get(key)
|
||||||
result = {
|
if raw is None or isinstance(raw, bool):
|
||||||
"prompt_tokens": prompt_tokens,
|
return None
|
||||||
"completion_tokens": completion_tokens,
|
try:
|
||||||
"total_tokens": total_tokens,
|
return int(raw)
|
||||||
}
|
except (TypeError, ValueError):
|
||||||
|
return None
|
||||||
|
|
||||||
|
input_tokens = _usage_int(usage, "input_tokens") or 0
|
||||||
|
output_tokens = _usage_int(usage, "output_tokens") or 0
|
||||||
input_details = _response_object(usage.get("input_tokens_details"))
|
input_details = _response_object(usage.get("input_tokens_details"))
|
||||||
cached_tokens = int(input_details.get("cached_tokens") or 0) if input_details else 0
|
return LLMUsage.reported(
|
||||||
if cached_tokens > 0:
|
input_tokens=input_tokens,
|
||||||
result["cached_tokens"] = cached_tokens
|
output_tokens=output_tokens,
|
||||||
return result
|
total_tokens=_usage_int(usage, "total_tokens"),
|
||||||
|
cache_read_tokens=_usage_int(input_details, "cached_tokens"),
|
||||||
|
cache_write_tokens=_usage_int(input_details, "cache_write_tokens"),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
def _parse_tool_call_arguments(args_raw: Any, name: str | None) -> Any:
|
def _parse_tool_call_arguments(args_raw: Any, name: str | None) -> Any:
|
||||||
@@ -352,14 +359,14 @@ async def consume_sse_with_reasoning(
|
|||||||
on_reasoning_delta: Callable[[str], Awaitable[None]] | None = None,
|
on_reasoning_delta: Callable[[str], Awaitable[None]] | None = None,
|
||||||
on_response_event: Callable[[dict[str, Any]], Awaitable[None]] | None = None,
|
on_response_event: Callable[[dict[str, Any]], Awaitable[None]] | None = None,
|
||||||
capture: ResponsesStreamCapture | None = None,
|
capture: ResponsesStreamCapture | None = None,
|
||||||
) -> tuple[str, list[ToolCallRequest], str, dict[str, int], str | None]:
|
) -> tuple[str, list[ToolCallRequest], str, LLMUsage | None, str | None]:
|
||||||
"""Consume a Responses API SSE stream, including visible reasoning summaries."""
|
"""Consume a Responses API SSE stream, including visible reasoning summaries."""
|
||||||
content = ""
|
content = ""
|
||||||
tool_calls: list[ToolCallRequest] = []
|
tool_calls: list[ToolCallRequest] = []
|
||||||
tool_call_buffers: dict[str, dict[str, Any]] = {}
|
tool_call_buffers: dict[str, dict[str, Any]] = {}
|
||||||
tool_call_args_emitted: set[str] = set()
|
tool_call_args_emitted: set[str] = set()
|
||||||
finish_reason = "stop"
|
finish_reason = "stop"
|
||||||
usage: dict[str, int] = {}
|
usage: LLMUsage | None = None
|
||||||
reasoning_content: str | None = None
|
reasoning_content: str | None = None
|
||||||
streamed_reasoning = False
|
streamed_reasoning = False
|
||||||
reasoning_summary_key: tuple[str | None, int] | None = None
|
reasoning_summary_key: tuple[str | None, int] | None = None
|
||||||
@@ -657,14 +664,14 @@ async def consume_sdk_stream(
|
|||||||
on_tool_call_delta: Callable[[dict[str, Any]], Awaitable[None]] | None = None,
|
on_tool_call_delta: Callable[[dict[str, Any]], Awaitable[None]] | None = None,
|
||||||
on_reasoning_delta: Callable[[str], Awaitable[None]] | None = None,
|
on_reasoning_delta: Callable[[str], Awaitable[None]] | None = None,
|
||||||
capture: ResponsesStreamCapture | None = None,
|
capture: ResponsesStreamCapture | None = None,
|
||||||
) -> tuple[str, list[ToolCallRequest], str, dict[str, int], str | None]:
|
) -> tuple[str, list[ToolCallRequest], str, LLMUsage | None, str | None]:
|
||||||
"""Consume an SDK async stream from ``client.responses.create(stream=True)``."""
|
"""Consume an SDK async stream from ``client.responses.create(stream=True)``."""
|
||||||
content = ""
|
content = ""
|
||||||
tool_calls: list[ToolCallRequest] = []
|
tool_calls: list[ToolCallRequest] = []
|
||||||
tool_call_buffers: dict[str, dict[str, Any]] = {}
|
tool_call_buffers: dict[str, dict[str, Any]] = {}
|
||||||
tool_call_args_emitted: set[str] = set()
|
tool_call_args_emitted: set[str] = set()
|
||||||
finish_reason = "stop"
|
finish_reason = "stop"
|
||||||
usage: dict[str, int] = {}
|
usage: LLMUsage | None = None
|
||||||
reasoning_content: str | None = None
|
reasoning_content: str | None = None
|
||||||
streamed_reasoning = False
|
streamed_reasoning = False
|
||||||
refusal_seen = False
|
refusal_seen = False
|
||||||
@@ -823,20 +830,7 @@ async def consume_sdk_stream(
|
|||||||
if on_content_delta and remaining_text:
|
if on_content_delta and remaining_text:
|
||||||
await on_content_delta(remaining_text)
|
await on_content_delta(remaining_text)
|
||||||
if resp:
|
if resp:
|
||||||
usage_obj = getattr(resp, "usage", None)
|
usage = _usage_from_response_obj(resp) or usage
|
||||||
if usage_obj:
|
|
||||||
usage = {
|
|
||||||
"prompt_tokens": int(getattr(usage_obj, "input_tokens", 0) or 0),
|
|
||||||
"completion_tokens": int(getattr(usage_obj, "output_tokens", 0) or 0),
|
|
||||||
"total_tokens": int(getattr(usage_obj, "total_tokens", 0) or 0),
|
|
||||||
}
|
|
||||||
usage_data = _response_object(usage_obj) or {}
|
|
||||||
input_details = _response_object(usage_data.get("input_tokens_details"))
|
|
||||||
cached_tokens = (
|
|
||||||
int(input_details.get("cached_tokens") or 0) if input_details else 0
|
|
||||||
)
|
|
||||||
if cached_tokens > 0:
|
|
||||||
usage["cached_tokens"] = cached_tokens
|
|
||||||
if not reasoning_content:
|
if not reasoning_content:
|
||||||
reasoning_content = _extract_reasoning_summary_from_output(
|
reasoning_content = _extract_reasoning_summary_from_output(
|
||||||
getattr(resp, "output", None)
|
getattr(resp, "output", None)
|
||||||
|
|||||||
@@ -7,7 +7,7 @@ from typing import Any, cast
|
|||||||
|
|
||||||
from loguru import logger
|
from loguru import logger
|
||||||
|
|
||||||
from nanobot.providers.base import ProviderConversationState
|
from nanobot.providers.base import LLMUsage, ProviderConversationState
|
||||||
from nanobot.providers.openai_responses.converters import convert_messages
|
from nanobot.providers.openai_responses.converters import convert_messages
|
||||||
|
|
||||||
RESPONSES_STATE_KIND = "openai_responses"
|
RESPONSES_STATE_KIND = "openai_responses"
|
||||||
@@ -84,7 +84,7 @@ def build_responses_state(
|
|||||||
model: str,
|
model: str,
|
||||||
input_items: list[dict[str, Any]],
|
input_items: list[dict[str, Any]],
|
||||||
output_items: list[dict[str, Any]],
|
output_items: list[dict[str, Any]],
|
||||||
usage: dict[str, int] | None = None,
|
usage: LLMUsage | None = None,
|
||||||
) -> ProviderConversationState:
|
) -> ProviderConversationState:
|
||||||
"""Create the canonical next state from request input and every output item."""
|
"""Create the canonical next state from request input and every output item."""
|
||||||
unpruned_items = [*input_items, *output_items]
|
unpruned_items = [*input_items, *output_items]
|
||||||
@@ -178,16 +178,8 @@ def _prune_before_latest_output_compaction(
|
|||||||
return output_items[latest:]
|
return output_items[latest:]
|
||||||
|
|
||||||
|
|
||||||
def _context_tokens_from_usage(usage: dict[str, int] | None) -> int:
|
def _context_tokens_from_usage(usage: LLMUsage | None) -> int:
|
||||||
if not usage:
|
return usage.total_tokens if usage is not None else 0
|
||||||
return 0
|
|
||||||
prompt_tokens = usage.get("prompt_tokens", 0)
|
|
||||||
completion_tokens = usage.get("completion_tokens", 0)
|
|
||||||
total_tokens = usage.get("total_tokens", 0)
|
|
||||||
values = (prompt_tokens, completion_tokens, total_tokens)
|
|
||||||
if any(isinstance(value, bool) for value in values):
|
|
||||||
return 0
|
|
||||||
return max(0, total_tokens or prompt_tokens + completion_tokens)
|
|
||||||
|
|
||||||
|
|
||||||
def _state_items(
|
def _state_items(
|
||||||
|
|||||||
@@ -11,7 +11,7 @@ class UnconfiguredProvider(LLMProvider):
|
|||||||
"""Keep the gateway available for settings before a model is configured."""
|
"""Keep the gateway available for settings before a model is configured."""
|
||||||
|
|
||||||
def __init__(self, default_model: str) -> None:
|
def __init__(self, default_model: str) -> None:
|
||||||
super().__init__()
|
super().__init__(provider_name="unconfigured")
|
||||||
self._default_model = default_model
|
self._default_model = default_model
|
||||||
|
|
||||||
async def chat(
|
async def chat(
|
||||||
|
|||||||
@@ -18,6 +18,7 @@ from nanobot import __version__
|
|||||||
from nanobot.providers.base import (
|
from nanobot.providers.base import (
|
||||||
LLMProvider,
|
LLMProvider,
|
||||||
LLMResponse,
|
LLMResponse,
|
||||||
|
LLMUsage,
|
||||||
ToolCallRequest,
|
ToolCallRequest,
|
||||||
resolve_stream_idle_timeout_s,
|
resolve_stream_idle_timeout_s,
|
||||||
)
|
)
|
||||||
@@ -69,8 +70,10 @@ class XAIGrokProvider(LLMProvider):
|
|||||||
default_model: str = DEFAULT_XAI_GROK_MODEL,
|
default_model: str = DEFAULT_XAI_GROK_MODEL,
|
||||||
proxy: str | None = None,
|
proxy: str | None = None,
|
||||||
extra_body: dict[str, Any] | None = None,
|
extra_body: dict[str, Any] | None = None,
|
||||||
|
*,
|
||||||
|
provider_name: str = "xai_grok",
|
||||||
):
|
):
|
||||||
super().__init__(api_key=None, api_base=None)
|
super().__init__(api_key=None, api_base=None, provider_name=provider_name)
|
||||||
self.default_model = default_model
|
self.default_model = default_model
|
||||||
self.proxy = proxy or None
|
self.proxy = proxy or None
|
||||||
self._extra_body = dict(extra_body or {})
|
self._extra_body = dict(extra_body or {})
|
||||||
@@ -436,7 +439,7 @@ async def _request_xai(
|
|||||||
on_content_delta: Callable[[str], Awaitable[None]] | None = None,
|
on_content_delta: Callable[[str], Awaitable[None]] | None = None,
|
||||||
on_thinking_delta: Callable[[str], Awaitable[None]] | None = None,
|
on_thinking_delta: Callable[[str], Awaitable[None]] | None = None,
|
||||||
on_tool_call_delta: Callable[[dict[str, Any]], Awaitable[None]] | None = None,
|
on_tool_call_delta: Callable[[dict[str, Any]], Awaitable[None]] | None = None,
|
||||||
) -> tuple[str, list[ToolCallRequest], str, dict[str, int], str | None]:
|
) -> tuple[str, list[ToolCallRequest], str, LLMUsage | None, str | None]:
|
||||||
async def _on_response_event(event: dict[str, Any]) -> None:
|
async def _on_response_event(event: dict[str, Any]) -> None:
|
||||||
hosted_event = _xai_hosted_tool_event(event)
|
hosted_event = _xai_hosted_tool_event(event)
|
||||||
if hosted_event is not None and on_tool_call_delta is not None:
|
if hosted_event is not None and on_tool_call_delta is not None:
|
||||||
|
|||||||
@@ -210,18 +210,20 @@ class RuntimeClient:
|
|||||||
|
|
||||||
async def compact_session(self, session_key: str) -> SessionSnapshot:
|
async def compact_session(self, session_key: str) -> SessionSnapshot:
|
||||||
"""Run token consolidation for one session."""
|
"""Run token consolidation for one session."""
|
||||||
session = self._loop.sessions.get_or_create(session_key)
|
session = await self._loop.sessions.get_or_create_async(session_key)
|
||||||
runtime = self._loop.runtime_for_session(session)
|
runtime = await self._loop.runtime_for_session_async(session)
|
||||||
await self._loop.consolidator.maybe_consolidate_by_tokens(
|
await self._loop.consolidator.maybe_consolidate_by_tokens(
|
||||||
session,
|
session,
|
||||||
runtime=runtime,
|
runtime=runtime,
|
||||||
)
|
)
|
||||||
return snapshot_from_session(self._loop.sessions.get_or_create(session_key))
|
return snapshot_from_session(
|
||||||
|
await self._loop.sessions.get_or_create_async(session_key)
|
||||||
|
)
|
||||||
|
|
||||||
async def compact_idle_session(self, session_key: str, *, max_suffix: int = 8) -> str | None:
|
async def compact_idle_session(self, session_key: str, *, max_suffix: int = 8) -> str | None:
|
||||||
"""Run idle-session compaction for one session and return the summary."""
|
"""Run idle-session compaction for one session and return the summary."""
|
||||||
session = self._loop.sessions.get_or_create(session_key)
|
session = await self._loop.sessions.get_or_create_async(session_key)
|
||||||
runtime = self._loop.runtime_for_session(session)
|
runtime = await self._loop.runtime_for_session_async(session)
|
||||||
return await self._loop.consolidator.compact_idle_session(
|
return await self._loop.consolidator.compact_idle_session(
|
||||||
session_key,
|
session_key,
|
||||||
runtime=runtime,
|
runtime=runtime,
|
||||||
|
|||||||
@@ -6,6 +6,7 @@ from copy import deepcopy
|
|||||||
from dataclasses import dataclass, field
|
from dataclasses import dataclass, field
|
||||||
from typing import Any, Literal, Mapping, TypeAlias, cast
|
from typing import Any, Literal, Mapping, TypeAlias, cast
|
||||||
|
|
||||||
|
from nanobot.providers.base import LLMUsage
|
||||||
from nanobot.runtime_context import public_history_messages
|
from nanobot.runtime_context import public_history_messages
|
||||||
|
|
||||||
StreamEventType: TypeAlias = Literal[
|
StreamEventType: TypeAlias = Literal[
|
||||||
@@ -53,7 +54,7 @@ class RunResult:
|
|||||||
content: str
|
content: str
|
||||||
tools_used: list[str] = field(default_factory=list)
|
tools_used: list[str] = field(default_factory=list)
|
||||||
messages: list[dict[str, Any]] = field(default_factory=list)
|
messages: list[dict[str, Any]] = field(default_factory=list)
|
||||||
usage: dict[str, int] = field(default_factory=dict)
|
usage: LLMUsage | None = None
|
||||||
stop_reason: str | None = None
|
stop_reason: str | None = None
|
||||||
error: str | None = None
|
error: str | None = None
|
||||||
metadata: dict[str, Any] = field(default_factory=dict)
|
metadata: dict[str, Any] = field(default_factory=dict)
|
||||||
@@ -72,7 +73,7 @@ class StreamEvent:
|
|||||||
arguments: dict[str, Any] | None = None
|
arguments: dict[str, Any] | None = None
|
||||||
iteration: int | None = None
|
iteration: int | None = None
|
||||||
resuming: bool | None = None
|
resuming: bool | None = None
|
||||||
usage: dict[str, int] = field(default_factory=dict)
|
usage: LLMUsage | None = None
|
||||||
error: str | None = None
|
error: str | None = None
|
||||||
metadata: dict[str, Any] = field(default_factory=dict)
|
metadata: dict[str, Any] = field(default_factory=dict)
|
||||||
|
|
||||||
|
|||||||
+144
-42
@@ -29,6 +29,7 @@ _BLOCKED_NETWORKS = [
|
|||||||
|
|
||||||
_URL_RE = re.compile(r"https?://[^\s\"'`;|<>]+", re.IGNORECASE)
|
_URL_RE = re.compile(r"https?://[^\s\"'`;|<>]+", re.IGNORECASE)
|
||||||
_allowed_networks: list[ipaddress.IPv4Network | ipaddress.IPv6Network] = []
|
_allowed_networks: list[ipaddress.IPv4Network | ipaddress.IPv6Network] = []
|
||||||
|
_DNS_RESOLUTION_TIMEOUT_SECONDS = 5.0
|
||||||
|
|
||||||
|
|
||||||
def is_loopback_host(host: str) -> bool:
|
def is_loopback_host(host: str) -> bool:
|
||||||
@@ -75,6 +76,63 @@ def _is_private(addr: ipaddress.IPv4Address | ipaddress.IPv6Address) -> bool:
|
|||||||
return any(normalized in net for net in _BLOCKED_NETWORKS)
|
return any(normalized in net for net in _BLOCKED_NETWORKS)
|
||||||
|
|
||||||
|
|
||||||
|
def _parse_url_hostname(url: str) -> tuple[str | None, str | None]:
|
||||||
|
try:
|
||||||
|
parsed = urlparse(url)
|
||||||
|
except Exception as exc:
|
||||||
|
return None, str(exc)
|
||||||
|
if parsed.scheme not in ("http", "https"):
|
||||||
|
return None, f"Only http/https allowed, got '{parsed.scheme or 'none'}'"
|
||||||
|
if not parsed.netloc:
|
||||||
|
return None, "Missing domain"
|
||||||
|
if not parsed.hostname:
|
||||||
|
return None, "Missing hostname"
|
||||||
|
return parsed.hostname, None
|
||||||
|
|
||||||
|
|
||||||
|
def _unresolved_target_result(
|
||||||
|
hostname: str,
|
||||||
|
*,
|
||||||
|
trust_remote_dns: bool,
|
||||||
|
) -> tuple[bool, str, tuple[str, ...]]:
|
||||||
|
if not trust_remote_dns:
|
||||||
|
return False, f"Cannot resolve hostname: {hostname}", ()
|
||||||
|
|
||||||
|
normalized_hostname = hostname.rstrip(".").lower()
|
||||||
|
if normalized_hostname == "localhost" or normalized_hostname.endswith(".localhost"):
|
||||||
|
return False, f"Blocked local/internal hostname: {hostname}", ()
|
||||||
|
|
||||||
|
try:
|
||||||
|
literal_addr = ipaddress.ip_address(normalized_hostname)
|
||||||
|
except ValueError:
|
||||||
|
return True, "", ()
|
||||||
|
if _is_private(literal_addr):
|
||||||
|
return False, f"Blocked private/internal address: {literal_addr}", ()
|
||||||
|
return True, "", (str(_normalize_addr(literal_addr)),)
|
||||||
|
|
||||||
|
|
||||||
|
def _resolved_target_result(
|
||||||
|
hostname: str,
|
||||||
|
infos: list[Any],
|
||||||
|
*,
|
||||||
|
allow_loopback: bool,
|
||||||
|
) -> tuple[bool, str, tuple[str, ...]]:
|
||||||
|
addrs: list[ipaddress.IPv4Address | ipaddress.IPv6Address] = []
|
||||||
|
for info in infos:
|
||||||
|
try:
|
||||||
|
addr = ipaddress.ip_address(info[4][0])
|
||||||
|
except (IndexError, TypeError, ValueError):
|
||||||
|
continue
|
||||||
|
addrs.append(addr)
|
||||||
|
if allow_loopback and _is_allowed_loopback_target(hostname, addrs):
|
||||||
|
return True, "", tuple(dict.fromkeys(str(_normalize_addr(addr)) for addr in addrs))
|
||||||
|
for addr in addrs:
|
||||||
|
if _is_private(addr):
|
||||||
|
return False, f"Blocked: {hostname} resolves to private/internal address {addr}", ()
|
||||||
|
|
||||||
|
return True, "", tuple(dict.fromkeys(str(_normalize_addr(addr)) for addr in addrs))
|
||||||
|
|
||||||
|
|
||||||
def resolve_url_target(
|
def resolve_url_target(
|
||||||
url: str,
|
url: str,
|
||||||
*,
|
*,
|
||||||
@@ -97,52 +155,43 @@ def resolve_url_target(
|
|||||||
resolved_ips contains the public IPs that were validated for this URL, or
|
resolved_ips contains the public IPs that were validated for this URL, or
|
||||||
is empty when an unresolved hostname is delegated to a trusted proxy.
|
is empty when an unresolved hostname is delegated to a trusted proxy.
|
||||||
"""
|
"""
|
||||||
try:
|
hostname, error = _parse_url_hostname(url)
|
||||||
p = urlparse(url)
|
if hostname is None:
|
||||||
except Exception as e:
|
return False, error or "Missing hostname", ()
|
||||||
return False, str(e), ()
|
|
||||||
|
|
||||||
if p.scheme not in ("http", "https"):
|
|
||||||
return False, f"Only http/https allowed, got '{p.scheme or 'none'}'", ()
|
|
||||||
if not p.netloc:
|
|
||||||
return False, "Missing domain", ()
|
|
||||||
|
|
||||||
hostname = p.hostname
|
|
||||||
if not hostname:
|
|
||||||
return False, "Missing hostname", ()
|
|
||||||
|
|
||||||
try:
|
try:
|
||||||
infos = socket.getaddrinfo(hostname, None, socket.AF_UNSPEC, socket.SOCK_STREAM)
|
infos = socket.getaddrinfo(hostname, None, socket.AF_UNSPEC, socket.SOCK_STREAM)
|
||||||
except socket.gaierror:
|
except socket.gaierror:
|
||||||
if not trust_remote_dns:
|
return _unresolved_target_result(hostname, trust_remote_dns=trust_remote_dns)
|
||||||
return False, f"Cannot resolve hostname: {hostname}", ()
|
return _resolved_target_result(hostname, infos, allow_loopback=allow_loopback)
|
||||||
|
|
||||||
normalized_hostname = hostname.rstrip(".").lower()
|
|
||||||
if normalized_hostname == "localhost" or normalized_hostname.endswith(".localhost"):
|
|
||||||
return False, f"Blocked local/internal hostname: {hostname}", ()
|
|
||||||
|
|
||||||
try:
|
async def async_resolve_url_target(
|
||||||
literal_addr = ipaddress.ip_address(normalized_hostname)
|
url: str,
|
||||||
except ValueError:
|
*,
|
||||||
return True, "", ()
|
allow_loopback: bool = False,
|
||||||
if _is_private(literal_addr):
|
trust_remote_dns: bool = False,
|
||||||
return False, f"Blocked private/internal address: {literal_addr}", ()
|
timeout_s: float = _DNS_RESOLUTION_TIMEOUT_SECONDS,
|
||||||
return True, "", (str(_normalize_addr(literal_addr)),)
|
) -> tuple[bool, str, tuple[str, ...]]:
|
||||||
|
"""Resolve and validate an HTTP target without blocking the event loop."""
|
||||||
addrs: list[ipaddress.IPv4Address | ipaddress.IPv6Address] = []
|
hostname, error = _parse_url_hostname(url)
|
||||||
for info in infos:
|
if hostname is None:
|
||||||
try:
|
return False, error or "Missing hostname", ()
|
||||||
addr = ipaddress.ip_address(info[4][0])
|
loop = asyncio.get_running_loop()
|
||||||
except ValueError:
|
try:
|
||||||
continue
|
infos = await asyncio.wait_for(
|
||||||
addrs.append(addr)
|
loop.getaddrinfo(
|
||||||
if allow_loopback and _is_allowed_loopback_target(hostname, addrs):
|
hostname,
|
||||||
return True, "", tuple(dict.fromkeys(str(_normalize_addr(addr)) for addr in addrs))
|
None,
|
||||||
for addr in addrs:
|
family=socket.AF_UNSPEC,
|
||||||
if _is_private(addr):
|
type=socket.SOCK_STREAM,
|
||||||
return False, f"Blocked: {hostname} resolves to private/internal address {addr}", ()
|
),
|
||||||
|
timeout=timeout_s,
|
||||||
return True, "", tuple(dict.fromkeys(str(_normalize_addr(addr)) for addr in addrs))
|
)
|
||||||
|
except asyncio.TimeoutError:
|
||||||
|
return False, f"Timed out resolving hostname: {hostname}", ()
|
||||||
|
except socket.gaierror:
|
||||||
|
return _unresolved_target_result(hostname, trust_remote_dns=trust_remote_dns)
|
||||||
|
return _resolved_target_result(hostname, infos, allow_loopback=allow_loopback)
|
||||||
|
|
||||||
|
|
||||||
def validate_url_target(url: str, *, allow_loopback: bool = False) -> tuple[bool, str]:
|
def validate_url_target(url: str, *, allow_loopback: bool = False) -> tuple[bool, str]:
|
||||||
@@ -151,6 +200,16 @@ def validate_url_target(url: str, *, allow_loopback: bool = False) -> tuple[bool
|
|||||||
return ok, error
|
return ok, error
|
||||||
|
|
||||||
|
|
||||||
|
async def async_validate_url_target(
|
||||||
|
url: str,
|
||||||
|
*,
|
||||||
|
allow_loopback: bool = False,
|
||||||
|
) -> tuple[bool, str]:
|
||||||
|
"""Validate a URL using the event loop's asynchronous resolver."""
|
||||||
|
ok, error, _ = await async_resolve_url_target(url, allow_loopback=allow_loopback)
|
||||||
|
return ok, error
|
||||||
|
|
||||||
|
|
||||||
def env_proxy_applies_to_url(url: str) -> bool:
|
def env_proxy_applies_to_url(url: str) -> bool:
|
||||||
"""Return True when process proxy settings would proxy this URL."""
|
"""Return True when process proxy settings would proxy this URL."""
|
||||||
try:
|
try:
|
||||||
@@ -277,7 +336,10 @@ class PinnedDNSAsyncTransport(httpx.AsyncBaseTransport):
|
|||||||
|
|
||||||
async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
|
async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
|
||||||
url = str(request.url)
|
url = str(request.url)
|
||||||
ok, error, resolved_ips = resolve_url_target(url, allow_loopback=self._allow_loopback)
|
ok, error, resolved_ips = await async_resolve_url_target(
|
||||||
|
url,
|
||||||
|
allow_loopback=self._allow_loopback,
|
||||||
|
)
|
||||||
if not ok:
|
if not ok:
|
||||||
raise UnsafeURLRequestError(error, request=request)
|
raise UnsafeURLRequestError(error, request=request)
|
||||||
async with self._resolver_lock:
|
async with self._resolver_lock:
|
||||||
@@ -320,6 +382,46 @@ def validate_resolved_url(url: str) -> tuple[bool, str]:
|
|||||||
return True, ""
|
return True, ""
|
||||||
|
|
||||||
|
|
||||||
|
async def async_validate_resolved_url(url: str) -> tuple[bool, str]:
|
||||||
|
"""Validate a redirect target without blocking on domain resolution."""
|
||||||
|
try:
|
||||||
|
parsed = urlparse(url)
|
||||||
|
except Exception:
|
||||||
|
return True, ""
|
||||||
|
hostname = parsed.hostname
|
||||||
|
if not hostname:
|
||||||
|
return True, ""
|
||||||
|
try:
|
||||||
|
addr = ipaddress.ip_address(hostname)
|
||||||
|
except ValueError:
|
||||||
|
loop = asyncio.get_running_loop()
|
||||||
|
try:
|
||||||
|
infos = await asyncio.wait_for(
|
||||||
|
loop.getaddrinfo(
|
||||||
|
hostname,
|
||||||
|
None,
|
||||||
|
family=socket.AF_UNSPEC,
|
||||||
|
type=socket.SOCK_STREAM,
|
||||||
|
),
|
||||||
|
timeout=_DNS_RESOLUTION_TIMEOUT_SECONDS,
|
||||||
|
)
|
||||||
|
except asyncio.TimeoutError:
|
||||||
|
return False, f"Timed out resolving redirect hostname: {hostname}"
|
||||||
|
except socket.gaierror:
|
||||||
|
return True, ""
|
||||||
|
for info in infos:
|
||||||
|
try:
|
||||||
|
addr = ipaddress.ip_address(info[4][0])
|
||||||
|
except (IndexError, TypeError, ValueError):
|
||||||
|
continue
|
||||||
|
if _is_private(addr):
|
||||||
|
return False, f"Redirect target {hostname} resolves to private address {addr}"
|
||||||
|
return True, ""
|
||||||
|
if _is_private(addr):
|
||||||
|
return False, f"Redirect target is a private address: {addr}"
|
||||||
|
return True, ""
|
||||||
|
|
||||||
|
|
||||||
def contains_internal_url(command: str, *, allow_loopback: bool = False) -> bool:
|
def contains_internal_url(command: str, *, allow_loopback: bool = False) -> bool:
|
||||||
"""Return True if the command string contains a URL targeting an internal/private address."""
|
"""Return True if the command string contains a URL targeting an internal/private address."""
|
||||||
for m in _URL_RE.finditer(command):
|
for m in _URL_RE.finditer(command):
|
||||||
|
|||||||
@@ -0,0 +1,29 @@
|
|||||||
|
"""Compatibility bridge for asynchronous SessionManager operations."""
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import inspect
|
||||||
|
from collections.abc import Awaitable, Callable
|
||||||
|
from typing import Any, TypeVar, cast
|
||||||
|
|
||||||
|
from nanobot.utils.cancellation import shield_and_drain
|
||||||
|
|
||||||
|
_SessionResult = TypeVar("_SessionResult")
|
||||||
|
|
||||||
|
|
||||||
|
async def call_session_manager(
|
||||||
|
manager: object,
|
||||||
|
async_method_name: str,
|
||||||
|
sync_method: Callable[..., _SessionResult],
|
||||||
|
/,
|
||||||
|
*args: Any,
|
||||||
|
**kwargs: Any,
|
||||||
|
) -> _SessionResult:
|
||||||
|
"""Prefer a class-declared coroutine, or offload the established sync contract."""
|
||||||
|
class_async_method = inspect.getattr_static(type(manager), async_method_name, None)
|
||||||
|
if inspect.iscoroutinefunction(class_async_method):
|
||||||
|
async_method = cast(
|
||||||
|
Callable[..., Awaitable[_SessionResult]],
|
||||||
|
getattr(manager, async_method_name),
|
||||||
|
)
|
||||||
|
return await async_method(*args, **kwargs)
|
||||||
|
return await shield_and_drain(asyncio.to_thread(sync_method, *args, **kwargs))
|
||||||
+248
-3
@@ -1,5 +1,6 @@
|
|||||||
"""Session management for conversation history."""
|
"""Session management for conversation history."""
|
||||||
|
|
||||||
|
import asyncio
|
||||||
import base64
|
import base64
|
||||||
import errno
|
import errno
|
||||||
import hashlib
|
import hashlib
|
||||||
@@ -28,6 +29,7 @@ from nanobot.runtime_context import (
|
|||||||
public_history_message,
|
public_history_message,
|
||||||
)
|
)
|
||||||
from nanobot.session.model_selection import SESSION_MODEL_PRESET_METADATA_KEY
|
from nanobot.session.model_selection import SESSION_MODEL_PRESET_METADATA_KEY
|
||||||
|
from nanobot.utils.cancellation import shield_and_drain
|
||||||
from nanobot.utils.helpers import (
|
from nanobot.utils.helpers import (
|
||||||
content_with_media_breadcrumbs,
|
content_with_media_breadcrumbs,
|
||||||
ensure_dir,
|
ensure_dir,
|
||||||
@@ -48,15 +50,21 @@ _SESSION_PREVIEW_MAX_CHARS = 120
|
|||||||
_SESSION_LIST_PREVIEW_MAX_RECORDS = 200
|
_SESSION_LIST_PREVIEW_MAX_RECORDS = 200
|
||||||
_SESSION_LIST_PREVIEW_MAX_CHARS = 1_000_000
|
_SESSION_LIST_PREVIEW_MAX_CHARS = 1_000_000
|
||||||
_SESSION_DATA_ERRORS = (ValueError, TypeError, AttributeError, KeyError)
|
_SESSION_DATA_ERRORS = (ValueError, TypeError, AttributeError, KeyError)
|
||||||
|
_RUNTIME_CHECKPOINT_DATA_ERRORS = (OSError, *_SESSION_DATA_ERRORS)
|
||||||
_PROVIDER_STATE_RECORD_TYPE = "provider_state"
|
_PROVIDER_STATE_RECORD_TYPE = "provider_state"
|
||||||
_PROVIDER_STATE_RECORD_PREFIX_RE = re.compile(
|
_PROVIDER_STATE_RECORD_PREFIX_RE = re.compile(
|
||||||
r'^\s*\{\s*"_type"\s*:\s*"provider_state"\s*(?:,|\})'
|
r'^\s*\{\s*"_type"\s*:\s*"provider_state"\s*(?:,|\})'
|
||||||
)
|
)
|
||||||
|
_RUNTIME_CHECKPOINT_KEY = "runtime_checkpoint"
|
||||||
|
_RUNTIME_CHECKPOINT_VERSION = 1
|
||||||
|
_RUNTIME_CHECKPOINT_SUFFIX = ".checkpoint.json"
|
||||||
_FORK_VOLATILE_METADATA_KEYS = {
|
_FORK_VOLATILE_METADATA_KEYS = {
|
||||||
"goal_state",
|
"goal_state",
|
||||||
"pending_user_turn",
|
"pending_user_turn",
|
||||||
|
"pending_user_followups",
|
||||||
"runtime_checkpoint",
|
"runtime_checkpoint",
|
||||||
"session_handle",
|
"session_handle",
|
||||||
|
"webui_recovery",
|
||||||
"thread_goal",
|
"thread_goal",
|
||||||
"title",
|
"title",
|
||||||
"title_user_edited",
|
"title_user_edited",
|
||||||
@@ -65,6 +73,7 @@ _WORKSPACE_STATE_DIR = ".nanobot"
|
|||||||
_WORKSPACE_ID_FILE = "workspace-id"
|
_WORKSPACE_ID_FILE = "workspace-id"
|
||||||
_WORKSPACE_ID_RE = re.compile(r"^[0-9a-f]{32}$")
|
_WORKSPACE_ID_RE = re.compile(r"^[0-9a-f]{32}$")
|
||||||
_SESSION_MIGRATION_LOCK_TIMEOUT_SECONDS = 30
|
_SESSION_MIGRATION_LOCK_TIMEOUT_SECONDS = 30
|
||||||
|
_SESSION_FILES_LOCK_TIMEOUT_SECONDS = 5
|
||||||
_SESSION_FILES_LOCK_FILENAME = ".session-files.lock"
|
_SESSION_FILES_LOCK_FILENAME = ".session-files.lock"
|
||||||
_COPY_CHUNK_SIZE = 1024 * 1024
|
_COPY_CHUNK_SIZE = 1024 * 1024
|
||||||
|
|
||||||
@@ -554,7 +563,8 @@ class JsonlSessionStore:
|
|||||||
self.sessions_dir = ensure_dir(root / workspace_id)
|
self.sessions_dir = ensure_dir(root / workspace_id)
|
||||||
self.legacy_sessions_dir = get_legacy_sessions_dir()
|
self.legacy_sessions_dir = get_legacy_sessions_dir()
|
||||||
self._session_files_lock = FileLock(
|
self._session_files_lock = FileLock(
|
||||||
str(self.sessions_dir / _SESSION_FILES_LOCK_FILENAME)
|
str(self.sessions_dir / _SESSION_FILES_LOCK_FILENAME),
|
||||||
|
timeout=_SESSION_FILES_LOCK_TIMEOUT_SECONDS,
|
||||||
)
|
)
|
||||||
with self._session_files_lock:
|
with self._session_files_lock:
|
||||||
self._migrate_from_workspace(canonical_workspace)
|
self._migrate_from_workspace(canonical_workspace)
|
||||||
@@ -1001,6 +1011,9 @@ class JsonlSessionStore:
|
|||||||
def get_session_path(self, key: str) -> Path:
|
def get_session_path(self, key: str) -> Path:
|
||||||
return self.sessions_dir / f"{self.storage_key(key)}.jsonl"
|
return self.sessions_dir / f"{self.storage_key(key)}.jsonl"
|
||||||
|
|
||||||
|
def get_runtime_checkpoint_path(self, key: str) -> Path:
|
||||||
|
return self.sessions_dir / f"{self.storage_key(key)}{_RUNTIME_CHECKPOINT_SUFFIX}"
|
||||||
|
|
||||||
def get_legacy_lossy_path(self, key: str) -> Path:
|
def get_legacy_lossy_path(self, key: str) -> Path:
|
||||||
return self.sessions_dir / f"{safe_filename(key.replace(':', '_'))}.jsonl"
|
return self.sessions_dir / f"{safe_filename(key.replace(':', '_'))}.jsonl"
|
||||||
|
|
||||||
@@ -1066,7 +1079,7 @@ class JsonlSessionStore:
|
|||||||
else:
|
else:
|
||||||
messages.append(data)
|
messages.append(data)
|
||||||
|
|
||||||
return Session(
|
session = Session(
|
||||||
key=key,
|
key=key,
|
||||||
messages=messages,
|
messages=messages,
|
||||||
created_at=created_at or datetime.now(),
|
created_at=created_at or datetime.now(),
|
||||||
@@ -1075,6 +1088,8 @@ class JsonlSessionStore:
|
|||||||
last_consolidated=last_consolidated,
|
last_consolidated=last_consolidated,
|
||||||
provider_state=provider_state,
|
provider_state=provider_state,
|
||||||
)
|
)
|
||||||
|
self._overlay_runtime_checkpoint_unlocked(session, path)
|
||||||
|
return session
|
||||||
except _SESSION_DATA_ERRORS as e:
|
except _SESSION_DATA_ERRORS as e:
|
||||||
logger.warning("Failed to load session {}: {}", key, e)
|
logger.warning("Failed to load session {}: {}", key, e)
|
||||||
repaired = self._repair_unlocked(key)
|
repaired = self._repair_unlocked(key)
|
||||||
@@ -1159,7 +1174,7 @@ class JsonlSessionStore:
|
|||||||
if not messages and not metadata and provider_state is None:
|
if not messages and not metadata and provider_state is None:
|
||||||
return None
|
return None
|
||||||
|
|
||||||
return Session(
|
session = Session(
|
||||||
key=key,
|
key=key,
|
||||||
messages=messages,
|
messages=messages,
|
||||||
created_at=created_at or datetime.now(),
|
created_at=created_at or datetime.now(),
|
||||||
@@ -1168,6 +1183,8 @@ class JsonlSessionStore:
|
|||||||
last_consolidated=last_consolidated,
|
last_consolidated=last_consolidated,
|
||||||
provider_state=provider_state,
|
provider_state=provider_state,
|
||||||
)
|
)
|
||||||
|
self._overlay_runtime_checkpoint_unlocked(session, path)
|
||||||
|
return session
|
||||||
except _SESSION_DATA_ERRORS as e:
|
except _SESSION_DATA_ERRORS as e:
|
||||||
logger.warning("Repair failed for session {}: {}", key, e)
|
logger.warning("Repair failed for session {}: {}", key, e)
|
||||||
return None
|
return None
|
||||||
@@ -1186,6 +1203,105 @@ class JsonlSessionStore:
|
|||||||
with self._session_files_lock:
|
with self._session_files_lock:
|
||||||
self._save_unlocked(session, fsync=fsync)
|
self._save_unlocked(session, fsync=fsync)
|
||||||
|
|
||||||
|
def save_runtime_checkpoint(self, session: Session) -> None:
|
||||||
|
"""Atomically persist only the volatile in-flight turn state.
|
||||||
|
|
||||||
|
A checkpoint is written several times during a tool-heavy turn. Keeping it
|
||||||
|
beside the append history avoids copying the full transcript at each safe
|
||||||
|
recovery boundary.
|
||||||
|
"""
|
||||||
|
with self._session_files_lock:
|
||||||
|
path = self.get_session_path(session.key)
|
||||||
|
if not path.exists():
|
||||||
|
# A user turn normally creates the session first. Internal callers
|
||||||
|
# may checkpoint a fresh session, so establish the durable base once.
|
||||||
|
self._save_unlocked(session)
|
||||||
|
return
|
||||||
|
|
||||||
|
checkpoint = session.metadata.get(_RUNTIME_CHECKPOINT_KEY)
|
||||||
|
if not isinstance(checkpoint, dict):
|
||||||
|
self.get_runtime_checkpoint_path(session.key).unlink(missing_ok=True)
|
||||||
|
return
|
||||||
|
|
||||||
|
payload: dict[str, Any] = {
|
||||||
|
"version": _RUNTIME_CHECKPOINT_VERSION,
|
||||||
|
"session_key": session.key,
|
||||||
|
"base_updated_at": session.updated_at.isoformat(),
|
||||||
|
"base_message_count": len(session.messages),
|
||||||
|
"checkpoint": checkpoint,
|
||||||
|
"provider_state": (
|
||||||
|
session.provider_state.to_private_record()
|
||||||
|
if session.provider_state is not None
|
||||||
|
else None
|
||||||
|
),
|
||||||
|
}
|
||||||
|
target = self.get_runtime_checkpoint_path(session.key)
|
||||||
|
tmp = target.with_name(f".{target.name}.{secrets.token_hex(8)}.tmp")
|
||||||
|
try:
|
||||||
|
with open(tmp, "x", encoding="utf-8") as handle:
|
||||||
|
os.chmod(tmp, 0o600)
|
||||||
|
json.dump(
|
||||||
|
payload,
|
||||||
|
handle,
|
||||||
|
ensure_ascii=False,
|
||||||
|
separators=(",", ":"),
|
||||||
|
)
|
||||||
|
os.replace(tmp, target)
|
||||||
|
finally:
|
||||||
|
tmp.unlink(missing_ok=True)
|
||||||
|
|
||||||
|
def _overlay_runtime_checkpoint_unlocked(self, session: Session, main_path: Path) -> None:
|
||||||
|
checkpoint_path = self.get_runtime_checkpoint_path(session.key)
|
||||||
|
try:
|
||||||
|
checkpoint_stat = checkpoint_path.lstat()
|
||||||
|
if not stat.S_ISREG(checkpoint_stat.st_mode):
|
||||||
|
logger.warning(
|
||||||
|
"Ignoring non-regular runtime checkpoint for session {}",
|
||||||
|
session.key,
|
||||||
|
)
|
||||||
|
return
|
||||||
|
# A complete session save supersedes an older sidecar. This comparison
|
||||||
|
# closes the small crash window between replacing the JSONL and unlinking
|
||||||
|
# its previous checkpoint.
|
||||||
|
if main_path.stat().st_mtime_ns > checkpoint_stat.st_mtime_ns:
|
||||||
|
checkpoint_path.unlink(missing_ok=True)
|
||||||
|
return
|
||||||
|
raw = _json_object(json.loads(checkpoint_path.read_text(encoding="utf-8")))
|
||||||
|
if (
|
||||||
|
raw.get("version") != _RUNTIME_CHECKPOINT_VERSION
|
||||||
|
or raw.get("session_key") != session.key
|
||||||
|
or raw.get("base_updated_at") != session.updated_at.isoformat()
|
||||||
|
or raw.get("base_message_count") != len(session.messages)
|
||||||
|
or not isinstance(raw.get("checkpoint"), dict)
|
||||||
|
):
|
||||||
|
checkpoint_path.unlink(missing_ok=True)
|
||||||
|
return
|
||||||
|
provider_record = raw.get("provider_state")
|
||||||
|
provider_state = (
|
||||||
|
None
|
||||||
|
if provider_record is None
|
||||||
|
else ProviderConversationState.from_private_record(provider_record)
|
||||||
|
)
|
||||||
|
if provider_record is not None and provider_state is None:
|
||||||
|
raise ValueError("invalid checkpoint provider state")
|
||||||
|
session.metadata[_RUNTIME_CHECKPOINT_KEY] = cast(
|
||||||
|
dict[str, Any], raw["checkpoint"]
|
||||||
|
)
|
||||||
|
session.provider_state = provider_state
|
||||||
|
except FileNotFoundError:
|
||||||
|
return
|
||||||
|
except _RUNTIME_CHECKPOINT_DATA_ERRORS as exc:
|
||||||
|
logger.warning(
|
||||||
|
"Ignoring invalid runtime checkpoint for session {}: {}",
|
||||||
|
session.key,
|
||||||
|
exc,
|
||||||
|
)
|
||||||
|
# Atomic writes mean a malformed target cannot become valid later.
|
||||||
|
# Remove it once so future loads do not repeatedly parse and log it.
|
||||||
|
with suppress(OSError):
|
||||||
|
if checkpoint_path.is_file() and not checkpoint_path.is_symlink():
|
||||||
|
checkpoint_path.unlink()
|
||||||
|
|
||||||
def _save_unlocked(self, session: Session, *, fsync: bool = False) -> None:
|
def _save_unlocked(self, session: Session, *, fsync: bool = False) -> None:
|
||||||
path = self.get_session_path(session.key)
|
path = self.get_session_path(session.key)
|
||||||
tmp_path = path.with_name(f".{path.name}.{secrets.token_hex(8)}.tmp")
|
tmp_path = path.with_name(f".{path.name}.{secrets.token_hex(8)}.tmp")
|
||||||
@@ -1215,6 +1331,10 @@ class JsonlSessionStore:
|
|||||||
|
|
||||||
os.replace(tmp_path, path)
|
os.replace(tmp_path, path)
|
||||||
|
|
||||||
|
# The full record now contains the authoritative checkpoint state (or
|
||||||
|
# its removal), so an older volatile overlay is no longer needed.
|
||||||
|
self.get_runtime_checkpoint_path(session.key).unlink(missing_ok=True)
|
||||||
|
|
||||||
if fsync:
|
if fsync:
|
||||||
with suppress(PermissionError):
|
with suppress(PermissionError):
|
||||||
fd = os.open(str(path.parent), os.O_RDONLY)
|
fd = os.open(str(path.parent), os.O_RDONLY)
|
||||||
@@ -1278,6 +1398,7 @@ class JsonlSessionStore:
|
|||||||
def _delete_unlocked(self, key: str) -> bool:
|
def _delete_unlocked(self, key: str) -> bool:
|
||||||
paths = [
|
paths = [
|
||||||
self.get_session_path(key),
|
self.get_session_path(key),
|
||||||
|
self.get_runtime_checkpoint_path(key),
|
||||||
self.get_legacy_lossy_path(key),
|
self.get_legacy_lossy_path(key),
|
||||||
self.get_legacy_session_path(key),
|
self.get_legacy_session_path(key),
|
||||||
]
|
]
|
||||||
@@ -1525,6 +1646,7 @@ class SessionManager:
|
|||||||
self._cache: OrderedDict[str, Session] = OrderedDict()
|
self._cache: OrderedDict[str, Session] = OrderedDict()
|
||||||
# Preserve identity for sessions held by active callers without retaining idle ones.
|
# Preserve identity for sessions held by active callers without retaining idle ones.
|
||||||
self._overflow_cache: WeakValueDictionary[str, Session] = WeakValueDictionary()
|
self._overflow_cache: WeakValueDictionary[str, Session] = WeakValueDictionary()
|
||||||
|
self._async_session_locks: WeakValueDictionary[str, asyncio.Lock] = WeakValueDictionary()
|
||||||
self._max_cached_sessions = SESSION_CACHE_MAX_SIZE
|
self._max_cached_sessions = SESSION_CACHE_MAX_SIZE
|
||||||
self._delete_observer: Callable[[str], None] | None = None
|
self._delete_observer: Callable[[str], None] | None = None
|
||||||
|
|
||||||
@@ -1585,6 +1707,10 @@ class SessionManager:
|
|||||||
"""Get the collision-resistant workspace path for a session."""
|
"""Get the collision-resistant workspace path for a session."""
|
||||||
return self._jsonl_store.get_session_path(key)
|
return self._jsonl_store.get_session_path(key)
|
||||||
|
|
||||||
|
def _get_runtime_checkpoint_path(self, key: str) -> Path:
|
||||||
|
"""Get the private in-flight checkpoint path for a session."""
|
||||||
|
return self._jsonl_store.get_runtime_checkpoint_path(key)
|
||||||
|
|
||||||
def _get_legacy_lossy_path(self, key: str) -> Path:
|
def _get_legacy_lossy_path(self, key: str) -> Path:
|
||||||
"""Previous workspace session path using lossy ':' to '_' replacement."""
|
"""Previous workspace session path using lossy ':' to '_' replacement."""
|
||||||
return self._jsonl_store.get_legacy_lossy_path(key)
|
return self._jsonl_store.get_legacy_lossy_path(key)
|
||||||
@@ -1620,6 +1746,28 @@ class SessionManager:
|
|||||||
self._remember(session)
|
self._remember(session)
|
||||||
return session
|
return session
|
||||||
|
|
||||||
|
def _async_session_lock(self, key: str) -> asyncio.Lock:
|
||||||
|
lock = self._async_session_locks.get(key)
|
||||||
|
if lock is None:
|
||||||
|
lock = asyncio.Lock()
|
||||||
|
self._async_session_locks[key] = lock
|
||||||
|
return lock
|
||||||
|
|
||||||
|
async def get_or_create_async(self, key: str) -> Session:
|
||||||
|
"""Load a session without running file I/O or lock waits on the event loop."""
|
||||||
|
cached = self.get_cached(key)
|
||||||
|
if cached is not None:
|
||||||
|
return cached
|
||||||
|
async with self._async_session_lock(key):
|
||||||
|
cached = self.get_cached(key)
|
||||||
|
if cached is not None:
|
||||||
|
return cached
|
||||||
|
session = await asyncio.to_thread(self._load, key)
|
||||||
|
if session is None:
|
||||||
|
session = Session(key=key)
|
||||||
|
self._remember(session)
|
||||||
|
return session
|
||||||
|
|
||||||
def get_or_create_transient(
|
def get_or_create_transient(
|
||||||
self,
|
self,
|
||||||
key: str,
|
key: str,
|
||||||
@@ -1653,6 +1801,46 @@ class SessionManager:
|
|||||||
self._store.save(session, fsync=fsync)
|
self._store.save(session, fsync=fsync)
|
||||||
self._remember(session)
|
self._remember(session)
|
||||||
|
|
||||||
|
async def save_async(self, session: Session, *, fsync: bool = False) -> None:
|
||||||
|
"""Persist a session without blocking the caller's event loop."""
|
||||||
|
if not session.policy.persist:
|
||||||
|
return
|
||||||
|
|
||||||
|
async def save_and_remember() -> None:
|
||||||
|
await asyncio.to_thread(self._store.save, session, fsync=fsync)
|
||||||
|
self._remember(session)
|
||||||
|
|
||||||
|
await shield_and_drain(save_and_remember())
|
||||||
|
|
||||||
|
def save_runtime_checkpoint(self, session: Session) -> None:
|
||||||
|
"""Persist volatile recovery state without rewriting long history."""
|
||||||
|
if not session.policy.persist:
|
||||||
|
return
|
||||||
|
if self._store is self._jsonl_store:
|
||||||
|
self._jsonl_store.save_runtime_checkpoint(session)
|
||||||
|
self._remember(session)
|
||||||
|
return
|
||||||
|
# Third-party stores keep their existing all-or-nothing semantics until
|
||||||
|
# they opt into a dedicated checkpoint primitive.
|
||||||
|
self.save(session)
|
||||||
|
|
||||||
|
async def save_runtime_checkpoint_async(self, session: Session) -> None:
|
||||||
|
"""Persist an in-flight checkpoint without blocking the event loop."""
|
||||||
|
if not session.policy.persist:
|
||||||
|
return
|
||||||
|
|
||||||
|
async def save_and_remember() -> None:
|
||||||
|
if self._store is self._jsonl_store:
|
||||||
|
await asyncio.to_thread(
|
||||||
|
self._jsonl_store.save_runtime_checkpoint,
|
||||||
|
session,
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
await asyncio.to_thread(self._store.save, session)
|
||||||
|
self._remember(session)
|
||||||
|
|
||||||
|
await shield_and_drain(save_and_remember())
|
||||||
|
|
||||||
def rename_model_preset(self, old_name: str, new_name: str) -> int:
|
def rename_model_preset(self, old_name: str, new_name: str) -> int:
|
||||||
"""Rename a session-scoped model preset across durable and live sessions."""
|
"""Rename a session-scoped model preset across durable and live sessions."""
|
||||||
if old_name == new_name:
|
if old_name == new_name:
|
||||||
@@ -1694,6 +1882,21 @@ class SessionManager:
|
|||||||
raise
|
raise
|
||||||
return len(changed)
|
return len(changed)
|
||||||
|
|
||||||
|
async def flush_all_async(self) -> int:
|
||||||
|
"""Re-save every cached session without blocking the event loop."""
|
||||||
|
cached = dict(self._overflow_cache.items())
|
||||||
|
cached.update(self._cache)
|
||||||
|
flushed = 0
|
||||||
|
for key, session in cached.items():
|
||||||
|
try:
|
||||||
|
await shield_and_drain(
|
||||||
|
asyncio.to_thread(self._store.save, session, fsync=True)
|
||||||
|
)
|
||||||
|
flushed += 1
|
||||||
|
except Exception:
|
||||||
|
logger.warning("Failed to flush session {}", key, exc_info=True)
|
||||||
|
return flushed
|
||||||
|
|
||||||
def flush_all(self) -> int:
|
def flush_all(self) -> int:
|
||||||
"""Re-save every cached session with fsync for durable shutdown.
|
"""Re-save every cached session with fsync for durable shutdown.
|
||||||
|
|
||||||
@@ -1725,6 +1928,18 @@ class SessionManager:
|
|||||||
self._delete_observer(key)
|
self._delete_observer(key)
|
||||||
return deleted
|
return deleted
|
||||||
|
|
||||||
|
async def delete_session_async(self, key: str) -> bool:
|
||||||
|
"""Delete a session without blocking the event loop."""
|
||||||
|
|
||||||
|
async def delete_and_notify() -> bool:
|
||||||
|
self.invalidate(key)
|
||||||
|
deleted = await asyncio.to_thread(self._store.delete, key)
|
||||||
|
if self._delete_observer is not None:
|
||||||
|
self._delete_observer(key)
|
||||||
|
return deleted
|
||||||
|
|
||||||
|
return await shield_and_drain(delete_and_notify())
|
||||||
|
|
||||||
def restore_sessions_to_workspace(self) -> SessionRestoreResult:
|
def restore_sessions_to_workspace(self) -> SessionRestoreResult:
|
||||||
"""Restore session files to the pre-relocation path for an explicit rollback."""
|
"""Restore session files to the pre-relocation path for an explicit rollback."""
|
||||||
return self._jsonl_store.restore_to_workspace()
|
return self._jsonl_store.restore_to_workspace()
|
||||||
@@ -1797,6 +2012,10 @@ class SessionManager:
|
|||||||
"""Read session metadata without loading the transcript."""
|
"""Read session metadata without loading the transcript."""
|
||||||
return cast(dict[str, Any] | None, self._store.read_metadata(key))
|
return cast(dict[str, Any] | None, self._store.read_metadata(key))
|
||||||
|
|
||||||
|
async def read_session_metadata_async(self, key: str) -> dict[str, Any] | None:
|
||||||
|
"""Read session metadata without blocking the event loop."""
|
||||||
|
return await asyncio.to_thread(self.read_session_metadata, key)
|
||||||
|
|
||||||
def update_session_metadata(
|
def update_session_metadata(
|
||||||
self,
|
self,
|
||||||
key: str,
|
key: str,
|
||||||
@@ -1810,5 +2029,31 @@ class SessionManager:
|
|||||||
session.metadata.update(deepcopy(updates))
|
session.metadata.update(deepcopy(updates))
|
||||||
return updated
|
return updated
|
||||||
|
|
||||||
|
async def update_session_metadata_async(
|
||||||
|
self,
|
||||||
|
key: str,
|
||||||
|
updates: dict[str, Any],
|
||||||
|
*,
|
||||||
|
fsync: bool = False,
|
||||||
|
) -> bool:
|
||||||
|
"""Update metadata without blocking the event loop."""
|
||||||
|
|
||||||
|
async def update_and_refresh_cache() -> bool:
|
||||||
|
updated = await asyncio.to_thread(
|
||||||
|
self._store.update_metadata,
|
||||||
|
key,
|
||||||
|
updates,
|
||||||
|
fsync=fsync,
|
||||||
|
)
|
||||||
|
if updated and (session := self.get_cached(key)) is not None:
|
||||||
|
session.metadata.update(deepcopy(updates))
|
||||||
|
return updated
|
||||||
|
|
||||||
|
return await shield_and_drain(update_and_refresh_cache())
|
||||||
|
|
||||||
def list_sessions(self) -> list[dict[str, Any]]:
|
def list_sessions(self) -> list[dict[str, Any]]:
|
||||||
return cast(list[dict[str, Any]], self._store.list_sessions())
|
return cast(list[dict[str, Any]], self._store.list_sessions())
|
||||||
|
|
||||||
|
async def list_sessions_async(self) -> list[dict[str, Any]]:
|
||||||
|
"""List persisted sessions without blocking the event loop."""
|
||||||
|
return await asyncio.to_thread(self.list_sessions)
|
||||||
|
|||||||
@@ -0,0 +1,971 @@
|
|||||||
|
"""Durable, side-effect-safe recovery for interrupted WebUI turns.
|
||||||
|
|
||||||
|
The coordinator owns restart policy. AgentLoop only exposes checkpoint
|
||||||
|
materialization and an admission hook, so transport code never has to guess
|
||||||
|
whether an interrupted tool call is safe to replay.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import dataclasses
|
||||||
|
import json
|
||||||
|
from collections.abc import Iterable, Mapping
|
||||||
|
from datetime import datetime
|
||||||
|
from typing import Any, Protocol, cast
|
||||||
|
from uuid import uuid4
|
||||||
|
|
||||||
|
from loguru import logger
|
||||||
|
|
||||||
|
from nanobot.bus.events import InboundMessage
|
||||||
|
from nanobot.bus.outbound_events import (
|
||||||
|
RecoveryStateEvent,
|
||||||
|
SessionUpdatedEvent,
|
||||||
|
outbound_message_for_event,
|
||||||
|
)
|
||||||
|
from nanobot.bus.queue import MessageBus
|
||||||
|
from nanobot.session import turn_continuation
|
||||||
|
from nanobot.session.async_compat import call_session_manager
|
||||||
|
from nanobot.session.keys import UNIFIED_SESSION_KEY, last_channel_from_metadata
|
||||||
|
from nanobot.session.manager import Session, SessionManager
|
||||||
|
from nanobot.webui.metadata import WEBUI_TURN_METADATA_KEY
|
||||||
|
|
||||||
|
RUNTIME_CHECKPOINT_KEY = "runtime_checkpoint"
|
||||||
|
PENDING_USER_TURN_KEY = "pending_user_turn"
|
||||||
|
RECOVERY_METADATA_KEY = "webui_recovery"
|
||||||
|
RECOVERY_INBOUND_METADATA_KEY = "_webui_recovery_id"
|
||||||
|
PENDING_FOLLOWUPS_KEY = "pending_user_followups"
|
||||||
|
PENDING_FOLLOWUP_ID_KEY = "_recovery_followup_id"
|
||||||
|
PROVIDER_STATE_CHECKPOINT_VERSION_KEY = "provider_state_checkpoint_version"
|
||||||
|
PROVIDER_STATE_CHECKPOINT_VERSION = "v1"
|
||||||
|
|
||||||
|
_RECOVERY_STATUSES = frozenset({"resuming", "awaiting_user", "recovered", "failed"})
|
||||||
|
_UNCERTAIN_TOOL_PHASES = frozenset({"awaiting_tools"})
|
||||||
|
_KNOWN_CHECKPOINT_PHASES = frozenset(
|
||||||
|
{"final_response", "tools_completed", "awaiting_tools", "error"}
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class RecoveryActionError(ValueError):
|
||||||
|
"""A stale or malformed recovery action from an authenticated WebUI."""
|
||||||
|
|
||||||
|
def __init__(self, message: str, *, status: int = 400) -> None:
|
||||||
|
super().__init__(message)
|
||||||
|
self.status = status
|
||||||
|
|
||||||
|
|
||||||
|
class RecoveryAdmission(Protocol):
|
||||||
|
"""Narrow AgentLoop boundary for explicit recovery validation."""
|
||||||
|
|
||||||
|
async def admit(self, message: InboundMessage) -> bool: ...
|
||||||
|
|
||||||
|
def register_recovery_task(self, session_key: str, task: asyncio.Task[Any]) -> None: ...
|
||||||
|
|
||||||
|
def unregister_recovery_task(self, session_key: str, task: asyncio.Task[Any]) -> None: ...
|
||||||
|
|
||||||
|
|
||||||
|
def record_pending_followup(session: Session, message: InboundMessage) -> str | None:
|
||||||
|
"""Durably journal a WebUI follow-up before injecting it into a live turn."""
|
||||||
|
if message.channel != "websocket":
|
||||||
|
return None
|
||||||
|
try:
|
||||||
|
metadata_value: object = json.loads(json.dumps(message.metadata))
|
||||||
|
except (TypeError, ValueError):
|
||||||
|
logger.warning("Skipping non-serializable WebUI follow-up for recovery")
|
||||||
|
return None
|
||||||
|
if not isinstance(metadata_value, dict):
|
||||||
|
return None
|
||||||
|
metadata = cast(dict[str, Any], metadata_value)
|
||||||
|
existing_id = metadata.pop(PENDING_FOLLOWUP_ID_KEY, None)
|
||||||
|
followup_id = (
|
||||||
|
existing_id
|
||||||
|
if isinstance(existing_id, str) and existing_id
|
||||||
|
else uuid4().hex
|
||||||
|
)
|
||||||
|
records = _pending_followup_records(session)
|
||||||
|
if any(record.get("id") == followup_id for record in records):
|
||||||
|
return followup_id
|
||||||
|
records.append(
|
||||||
|
{
|
||||||
|
"id": followup_id,
|
||||||
|
"sender_id": message.sender_id,
|
||||||
|
"chat_id": message.chat_id,
|
||||||
|
"content": message.content,
|
||||||
|
"media": list(message.media or []),
|
||||||
|
"metadata": metadata,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
# This journal is the recovery source of truth, not a mirror of the
|
||||||
|
# bounded in-memory injection queue. A queued turn can receive more
|
||||||
|
# follow-ups than the live queue accepts; dropping older journal entries
|
||||||
|
# would make those acknowledged user messages unrecoverable after a
|
||||||
|
# gateway restart. Entries are removed only once their user rows are
|
||||||
|
# committed by ``acknowledge_pending_followups``.
|
||||||
|
session.metadata[PENDING_FOLLOWUPS_KEY] = records
|
||||||
|
session.updated_at = datetime.now()
|
||||||
|
return followup_id
|
||||||
|
|
||||||
|
|
||||||
|
def pending_followups(session: Session) -> list[InboundMessage]:
|
||||||
|
"""Decode still-unacknowledged follow-ups from durable session metadata."""
|
||||||
|
messages: list[InboundMessage] = []
|
||||||
|
for record in _pending_followup_records(session):
|
||||||
|
followup_id = cast(object, record.get("id"))
|
||||||
|
sender_id = cast(object, record.get("sender_id"))
|
||||||
|
chat_id = cast(object, record.get("chat_id"))
|
||||||
|
content = cast(object, record.get("content"))
|
||||||
|
metadata = cast(object, record.get("metadata"))
|
||||||
|
if (
|
||||||
|
not isinstance(followup_id, str)
|
||||||
|
or not followup_id
|
||||||
|
or not isinstance(sender_id, str)
|
||||||
|
or not sender_id
|
||||||
|
or not isinstance(chat_id, str)
|
||||||
|
or not chat_id
|
||||||
|
):
|
||||||
|
continue
|
||||||
|
if not isinstance(content, str) or not isinstance(metadata, dict):
|
||||||
|
continue
|
||||||
|
media_value = cast(object, record.get("media"))
|
||||||
|
media = (
|
||||||
|
[item for item in cast(list[object], media_value) if isinstance(item, str)]
|
||||||
|
if isinstance(media_value, list)
|
||||||
|
else []
|
||||||
|
)
|
||||||
|
messages.append(
|
||||||
|
InboundMessage(
|
||||||
|
channel="websocket",
|
||||||
|
sender_id=sender_id,
|
||||||
|
chat_id=chat_id,
|
||||||
|
content=content,
|
||||||
|
media=media,
|
||||||
|
metadata={**cast(dict[str, Any], metadata), PENDING_FOLLOWUP_ID_KEY: followup_id},
|
||||||
|
session_key_override=session.key,
|
||||||
|
require_existing_session=True,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
return messages
|
||||||
|
|
||||||
|
|
||||||
|
def acknowledge_pending_followups(session: Session, followup_ids: Iterable[str]) -> None:
|
||||||
|
"""Remove journal entries whose user rows were committed to history."""
|
||||||
|
acknowledged = set(followup_ids)
|
||||||
|
if not acknowledged:
|
||||||
|
return
|
||||||
|
records = [record for record in _pending_followup_records(session) if record.get("id") not in acknowledged]
|
||||||
|
if records:
|
||||||
|
session.metadata[PENDING_FOLLOWUPS_KEY] = records
|
||||||
|
else:
|
||||||
|
session.metadata.pop(PENDING_FOLLOWUPS_KEY, None)
|
||||||
|
|
||||||
|
|
||||||
|
def _pending_followup_records(session: Session) -> list[dict[str, Any]]:
|
||||||
|
raw = cast(object, session.metadata.get(PENDING_FOLLOWUPS_KEY))
|
||||||
|
if not isinstance(raw, list):
|
||||||
|
return []
|
||||||
|
values = cast(list[object], raw)
|
||||||
|
return [cast(dict[str, Any], value) for value in values if isinstance(value, dict)]
|
||||||
|
|
||||||
|
|
||||||
|
def _checkpoint_message_key(message: Mapping[str, Any]) -> tuple[Any, ...]:
|
||||||
|
return (
|
||||||
|
message.get("role"),
|
||||||
|
message.get("content"),
|
||||||
|
message.get("tool_call_id"),
|
||||||
|
message.get("name"),
|
||||||
|
message.get("tool_calls"),
|
||||||
|
message.get("reasoning_content"),
|
||||||
|
message.get("thinking_blocks"),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _checkpoint_tool_call_ids(
|
||||||
|
value: object,
|
||||||
|
*,
|
||||||
|
result_rows: bool = False,
|
||||||
|
) -> list[str] | None:
|
||||||
|
"""Validate checkpoint tool rows and return their stable IDs."""
|
||||||
|
if not isinstance(value, list):
|
||||||
|
return None
|
||||||
|
ids: list[str] = []
|
||||||
|
for raw in cast(list[object], value):
|
||||||
|
if not isinstance(raw, dict):
|
||||||
|
return None
|
||||||
|
row = cast(dict[str, Any], raw)
|
||||||
|
id_key = "tool_call_id" if result_rows else "id"
|
||||||
|
call_id = cast(object, row.get(id_key))
|
||||||
|
if not isinstance(call_id, str) or not call_id:
|
||||||
|
return None
|
||||||
|
if result_rows:
|
||||||
|
if row.get("role") != "tool":
|
||||||
|
return None
|
||||||
|
else:
|
||||||
|
function_value = cast(object, row.get("function"))
|
||||||
|
if not isinstance(function_value, dict):
|
||||||
|
return None
|
||||||
|
function = cast(dict[str, Any], function_value)
|
||||||
|
name = cast(object, function.get("name"))
|
||||||
|
if not isinstance(name, str) or not name:
|
||||||
|
return None
|
||||||
|
ids.append(call_id)
|
||||||
|
return ids if len(ids) == len(set(ids)) else None
|
||||||
|
|
||||||
|
|
||||||
|
def _runtime_checkpoint_is_well_formed(checkpoint: Mapping[str, Any]) -> bool:
|
||||||
|
"""Return whether a checkpoint is safe to offer for continuation.
|
||||||
|
|
||||||
|
Restoration stays tolerant so Dismiss can always clear corrupt state.
|
||||||
|
Continue is stricter: silently dropping a malformed tool result could make
|
||||||
|
the model repeat an external side effect.
|
||||||
|
"""
|
||||||
|
assistant_value = cast(object, checkpoint.get("assistant_message"))
|
||||||
|
if not isinstance(assistant_value, dict):
|
||||||
|
return False
|
||||||
|
assistant = cast(dict[str, Any], assistant_value)
|
||||||
|
if assistant.get("role") != "assistant":
|
||||||
|
return False
|
||||||
|
|
||||||
|
completed_ids = _checkpoint_tool_call_ids(
|
||||||
|
cast(object, checkpoint.get("completed_tool_results")),
|
||||||
|
result_rows=True,
|
||||||
|
)
|
||||||
|
pending_ids = _checkpoint_tool_call_ids(
|
||||||
|
cast(object, checkpoint.get("pending_tool_calls")),
|
||||||
|
)
|
||||||
|
if completed_ids is None or pending_ids is None:
|
||||||
|
return False
|
||||||
|
assistant_calls_value = cast(object, assistant.get("tool_calls"))
|
||||||
|
assistant_call_ids = (
|
||||||
|
[]
|
||||||
|
if assistant_calls_value is None
|
||||||
|
else _checkpoint_tool_call_ids(assistant_calls_value)
|
||||||
|
)
|
||||||
|
if assistant_call_ids is None:
|
||||||
|
return False
|
||||||
|
|
||||||
|
phase = checkpoint.get("phase")
|
||||||
|
if phase == "final_response":
|
||||||
|
content = cast(object, assistant.get("content"))
|
||||||
|
return (
|
||||||
|
isinstance(content, str)
|
||||||
|
and bool(content.strip())
|
||||||
|
and not assistant_call_ids
|
||||||
|
and not completed_ids
|
||||||
|
and not pending_ids
|
||||||
|
)
|
||||||
|
if phase == "awaiting_tools":
|
||||||
|
return (
|
||||||
|
bool(assistant_call_ids)
|
||||||
|
and not completed_ids
|
||||||
|
and len(assistant_call_ids) == len(pending_ids)
|
||||||
|
and set(assistant_call_ids) == set(pending_ids)
|
||||||
|
)
|
||||||
|
if phase == "tools_completed":
|
||||||
|
return (
|
||||||
|
bool(assistant_call_ids)
|
||||||
|
and not pending_ids
|
||||||
|
and len(assistant_call_ids) == len(completed_ids)
|
||||||
|
and set(assistant_call_ids) == set(completed_ids)
|
||||||
|
)
|
||||||
|
# Error checkpoints have no current producer contract. Treat legacy or
|
||||||
|
# future instances as review-only until their exact persisted shape is
|
||||||
|
# specified; guessing here could make a partial side effect repeat.
|
||||||
|
return False
|
||||||
|
|
||||||
|
|
||||||
|
def restore_runtime_checkpoint(session: Session) -> bool:
|
||||||
|
"""Materialize the durable checkpoint exactly once and clear it.
|
||||||
|
|
||||||
|
Pending tool calls become explicit interrupted tool results. They are
|
||||||
|
never executed here. Provider-native state is retained only for the two
|
||||||
|
checkpoint shapes known to be synchronized with persisted history.
|
||||||
|
"""
|
||||||
|
checkpoint = cast(object, session.metadata.get(RUNTIME_CHECKPOINT_KEY))
|
||||||
|
if not isinstance(checkpoint, dict):
|
||||||
|
return False
|
||||||
|
data = cast(dict[str, Any], checkpoint)
|
||||||
|
assistant = cast(object, data.get("assistant_message"))
|
||||||
|
completed_value = cast(object, data.get("completed_tool_results"))
|
||||||
|
pending_value = cast(object, data.get("pending_tool_calls"))
|
||||||
|
completed = cast(list[object], completed_value) if isinstance(completed_value, list) else []
|
||||||
|
pending = cast(list[object], pending_value) if isinstance(pending_value, list) else []
|
||||||
|
|
||||||
|
restored: list[dict[str, Any]] = []
|
||||||
|
if isinstance(assistant, dict):
|
||||||
|
assistant_row = cast(dict[str, Any], assistant)
|
||||||
|
else:
|
||||||
|
assistant_row = {}
|
||||||
|
if assistant_row.get("role") == "assistant":
|
||||||
|
row = dict(assistant_row)
|
||||||
|
row.setdefault("timestamp", datetime.now().isoformat())
|
||||||
|
restored.append(row)
|
||||||
|
for value in completed:
|
||||||
|
if not isinstance(value, dict):
|
||||||
|
continue
|
||||||
|
tool_result = cast(dict[str, Any], value)
|
||||||
|
if tool_result.get("role") != "tool":
|
||||||
|
continue
|
||||||
|
row = dict(tool_result)
|
||||||
|
row.setdefault("timestamp", datetime.now().isoformat())
|
||||||
|
restored.append(row)
|
||||||
|
for value in pending:
|
||||||
|
if not isinstance(value, dict):
|
||||||
|
continue
|
||||||
|
tool_call = cast(dict[str, Any], value)
|
||||||
|
tool_call_id = tool_call.get("id")
|
||||||
|
function_value = cast(object, tool_call.get("function"))
|
||||||
|
if not isinstance(tool_call_id, str) or not tool_call_id:
|
||||||
|
continue
|
||||||
|
function = (
|
||||||
|
cast(dict[str, Any], function_value)
|
||||||
|
if isinstance(function_value, dict)
|
||||||
|
else {}
|
||||||
|
)
|
||||||
|
name = function.get("name")
|
||||||
|
restored.append(
|
||||||
|
{
|
||||||
|
"role": "tool",
|
||||||
|
"tool_call_id": tool_call_id,
|
||||||
|
"name": name if isinstance(name, str) and name else "tool",
|
||||||
|
"content": "Error: Task interrupted before this tool finished.",
|
||||||
|
"timestamp": datetime.now().isoformat(),
|
||||||
|
"_recovery_interrupted": True,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
overlap = 0
|
||||||
|
for size in range(min(len(session.messages), len(restored)), 0, -1):
|
||||||
|
if all(
|
||||||
|
_checkpoint_message_key(left) == _checkpoint_message_key(right)
|
||||||
|
for left, right in zip(session.messages[-size:], restored[:size])
|
||||||
|
):
|
||||||
|
overlap = size
|
||||||
|
break
|
||||||
|
session.messages.extend(restored[overlap:])
|
||||||
|
|
||||||
|
assistant_data = cast(dict[str, Any], assistant) if isinstance(assistant, dict) else None
|
||||||
|
synchronized = (
|
||||||
|
data.get(PROVIDER_STATE_CHECKPOINT_VERSION_KEY)
|
||||||
|
== PROVIDER_STATE_CHECKPOINT_VERSION
|
||||||
|
)
|
||||||
|
phase = data.get("phase")
|
||||||
|
exact_final = (
|
||||||
|
phase == "final_response"
|
||||||
|
and assistant_data is not None
|
||||||
|
and assistant_data.get("role") == "assistant"
|
||||||
|
and not data.get("completed_tool_results")
|
||||||
|
and not data.get("pending_tool_calls")
|
||||||
|
)
|
||||||
|
exact_tools = (
|
||||||
|
phase == "tools_completed"
|
||||||
|
and assistant_data is not None
|
||||||
|
and assistant_data.get("role") == "assistant"
|
||||||
|
and not data.get("pending_tool_calls")
|
||||||
|
)
|
||||||
|
if not (synchronized and (exact_final or exact_tools)):
|
||||||
|
session.provider_state = None
|
||||||
|
|
||||||
|
session.metadata.pop(PENDING_USER_TURN_KEY, None)
|
||||||
|
session.metadata.pop(RUNTIME_CHECKPOINT_KEY, None)
|
||||||
|
session.updated_at = datetime.now()
|
||||||
|
return True
|
||||||
|
|
||||||
|
|
||||||
|
def _discard_runtime_checkpoint(session: Session) -> bool:
|
||||||
|
"""Drop checkpoint state that cannot be projected into valid history."""
|
||||||
|
if RUNTIME_CHECKPOINT_KEY not in session.metadata:
|
||||||
|
return False
|
||||||
|
session.metadata.pop(RUNTIME_CHECKPOINT_KEY, None)
|
||||||
|
session.provider_state = None
|
||||||
|
session.updated_at = datetime.now()
|
||||||
|
return True
|
||||||
|
|
||||||
|
|
||||||
|
def restore_pending_interruption(session: Session, *, superseded: bool = False) -> bool:
|
||||||
|
"""Close a persisted user-only turn without pretending it was answered."""
|
||||||
|
if not session.metadata.get(PENDING_USER_TURN_KEY):
|
||||||
|
return False
|
||||||
|
if session.messages and session.messages[-1].get("role") == "user":
|
||||||
|
content = (
|
||||||
|
"Task recovery was superseded by a newer message."
|
||||||
|
if superseded
|
||||||
|
else "Error: Task interrupted before a response was generated."
|
||||||
|
)
|
||||||
|
session.messages.append(
|
||||||
|
{
|
||||||
|
"role": "assistant",
|
||||||
|
"content": content,
|
||||||
|
"timestamp": datetime.now().isoformat(),
|
||||||
|
"_recovery_interrupted": True,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
session.provider_state = None
|
||||||
|
session.updated_at = datetime.now()
|
||||||
|
session.metadata.pop(PENDING_USER_TURN_KEY, None)
|
||||||
|
return True
|
||||||
|
|
||||||
|
|
||||||
|
def append_recovery_interruption(session: Session, *, superseded: bool = False) -> None:
|
||||||
|
"""Close a restored partial turn whose last durable row is not the user message."""
|
||||||
|
if session.messages and session.messages[-1].get("_recovery_interrupted") is True:
|
||||||
|
return
|
||||||
|
session.messages.append(
|
||||||
|
{
|
||||||
|
"role": "assistant",
|
||||||
|
"content": (
|
||||||
|
"Task recovery was superseded by a newer message."
|
||||||
|
if superseded
|
||||||
|
else "Error: Task recovery was interrupted before completion."
|
||||||
|
),
|
||||||
|
"timestamp": datetime.now().isoformat(),
|
||||||
|
"_recovery_interrupted": True,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
session.provider_state = None
|
||||||
|
session.updated_at = datetime.now()
|
||||||
|
|
||||||
|
|
||||||
|
def recovery_state_from_metadata(metadata: Mapping[str, Any] | None) -> dict[str, Any] | None:
|
||||||
|
"""Return a sanitized recovery state suitable for the WebSocket wire."""
|
||||||
|
value = metadata.get(RECOVERY_METADATA_KEY) if metadata else None
|
||||||
|
if not isinstance(value, dict):
|
||||||
|
return None
|
||||||
|
state = cast(dict[str, Any], value)
|
||||||
|
status = state.get("status")
|
||||||
|
recovery_id = state.get("recovery_id")
|
||||||
|
if status not in _RECOVERY_STATUSES or not isinstance(recovery_id, str):
|
||||||
|
return None
|
||||||
|
payload: dict[str, Any] = {"status": status, "recovery_id": recovery_id}
|
||||||
|
reason = state.get("reason")
|
||||||
|
if isinstance(reason, str) and reason:
|
||||||
|
payload["reason"] = reason
|
||||||
|
attempts = state.get("attempts")
|
||||||
|
if isinstance(attempts, int) and attempts >= 0:
|
||||||
|
payload["attempts"] = attempts
|
||||||
|
can_continue = state.get("can_continue")
|
||||||
|
if isinstance(can_continue, bool):
|
||||||
|
payload["can_continue"] = can_continue
|
||||||
|
return payload
|
||||||
|
|
||||||
|
|
||||||
|
@dataclasses.dataclass(slots=True)
|
||||||
|
class RecoveryCoordinator:
|
||||||
|
"""Classify, announce, and gate durable WebUI turn recovery."""
|
||||||
|
|
||||||
|
sessions: SessionManager
|
||||||
|
bus: MessageBus
|
||||||
|
unified_session: bool = False
|
||||||
|
_active_recovery_tasks: dict[str, asyncio.Task[Any]] = dataclasses.field(
|
||||||
|
default_factory=dict,
|
||||||
|
init=False,
|
||||||
|
repr=False,
|
||||||
|
)
|
||||||
|
|
||||||
|
async def _get_or_create_session(self, key: str) -> Session:
|
||||||
|
return await call_session_manager(
|
||||||
|
self.sessions,
|
||||||
|
"get_or_create_async",
|
||||||
|
self.sessions.get_or_create,
|
||||||
|
key,
|
||||||
|
)
|
||||||
|
|
||||||
|
async def _save_session(self, session: Session) -> None:
|
||||||
|
await call_session_manager(
|
||||||
|
self.sessions,
|
||||||
|
"save_async",
|
||||||
|
self.sessions.save,
|
||||||
|
session,
|
||||||
|
)
|
||||||
|
|
||||||
|
async def _read_session_metadata(self, key: str) -> dict[str, Any] | None:
|
||||||
|
return await call_session_manager(
|
||||||
|
self.sessions,
|
||||||
|
"read_session_metadata_async",
|
||||||
|
self.sessions.read_session_metadata,
|
||||||
|
key,
|
||||||
|
)
|
||||||
|
|
||||||
|
async def _list_sessions(self) -> list[dict[str, Any]]:
|
||||||
|
return await call_session_manager(
|
||||||
|
self.sessions,
|
||||||
|
"list_sessions_async",
|
||||||
|
self.sessions.list_sessions,
|
||||||
|
)
|
||||||
|
|
||||||
|
def register_recovery_task(self, session_key: str, task: asyncio.Task[Any]) -> None:
|
||||||
|
"""Track the task that owns an explicit recovery continuation."""
|
||||||
|
self._active_recovery_tasks[session_key] = task
|
||||||
|
|
||||||
|
def unregister_recovery_task(self, session_key: str, task: asyncio.Task[Any]) -> None:
|
||||||
|
"""Drop a recovery task without removing a newer task for the same session."""
|
||||||
|
if self._active_recovery_tasks.get(session_key) is task:
|
||||||
|
self._active_recovery_tasks.pop(session_key, None)
|
||||||
|
|
||||||
|
async def _cancel_active_recovery(self, session_key: str) -> None:
|
||||||
|
"""Stop an explicit continuation before accepting newer user input."""
|
||||||
|
task = self._active_recovery_tasks.get(session_key)
|
||||||
|
if task is None or task is asyncio.current_task() or task.done():
|
||||||
|
return
|
||||||
|
task.cancel()
|
||||||
|
# AgentLoop's cancellation path materializes any partial checkpoint and
|
||||||
|
# releases its pending queue. Wait for that ownership to be released
|
||||||
|
# before the newer message is routed.
|
||||||
|
await asyncio.gather(task, return_exceptions=True)
|
||||||
|
|
||||||
|
async def scan(self) -> None:
|
||||||
|
"""Recover every interrupted WebUI session once at gateway startup."""
|
||||||
|
for key in await self._recovery_candidates():
|
||||||
|
metadata_payload = await self._read_session_metadata(key)
|
||||||
|
raw_metadata = metadata_payload.get("metadata") if metadata_payload else None
|
||||||
|
metadata = cast(dict[str, Any], raw_metadata) if isinstance(raw_metadata, dict) else {}
|
||||||
|
route = self._websocket_route_for(key, metadata)
|
||||||
|
if route is None:
|
||||||
|
continue
|
||||||
|
unfinished = self._has_unfinished_webui_transcript(key)
|
||||||
|
if not self._needs_recovery(metadata) and not unfinished:
|
||||||
|
continue
|
||||||
|
session = await self._get_or_create_session(key)
|
||||||
|
try:
|
||||||
|
await self._recover_session(session, route[1])
|
||||||
|
await self._requeue_pending_followups(session)
|
||||||
|
except Exception:
|
||||||
|
logger.exception("failed to recover interrupted WebUI session {}", session.key)
|
||||||
|
state = recovery_state_from_metadata(session.metadata)
|
||||||
|
failed = self._set_state(
|
||||||
|
session,
|
||||||
|
status="failed",
|
||||||
|
recovery_id=cast(str, state["recovery_id"]) if state else uuid4().hex,
|
||||||
|
attempts=cast(int, state.get("attempts", 0)) if state else 0,
|
||||||
|
reason="recovery_failed",
|
||||||
|
can_continue=False,
|
||||||
|
)
|
||||||
|
await self._save_session(session)
|
||||||
|
await self._publish(route[1], failed)
|
||||||
|
|
||||||
|
async def _recovery_candidates(self) -> list[str]:
|
||||||
|
"""Discover canonical and transcript-only WebUI sessions cheaply."""
|
||||||
|
candidates = dict.fromkeys(
|
||||||
|
key
|
||||||
|
for item in await self._list_sessions()
|
||||||
|
if isinstance((key := item.get("key")), str)
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
# Imported lazily because the sidebar index also projects recovery
|
||||||
|
# metadata. The index is the owner of transcript-only discovery;
|
||||||
|
# duplicating its filename and migration rules here would drift.
|
||||||
|
from nanobot.webui.session_list_index import list_webui_sessions
|
||||||
|
|
||||||
|
for item in await asyncio.to_thread(list_webui_sessions, self.sessions):
|
||||||
|
key = item.get("key")
|
||||||
|
if isinstance(key, str):
|
||||||
|
candidates.setdefault(key, None)
|
||||||
|
except Exception:
|
||||||
|
# Canonical checkpoint recovery remains available even if the
|
||||||
|
# optional display-history index is corrupt or unavailable.
|
||||||
|
logger.exception("failed to discover transcript-only WebUI sessions")
|
||||||
|
return list(candidates)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _needs_recovery(metadata: Mapping[str, Any]) -> bool:
|
||||||
|
if metadata.get(PENDING_USER_TURN_KEY) is True:
|
||||||
|
return True
|
||||||
|
if isinstance(metadata.get(RUNTIME_CHECKPOINT_KEY), dict):
|
||||||
|
return True
|
||||||
|
followups = metadata.get(PENDING_FOLLOWUPS_KEY)
|
||||||
|
if isinstance(followups, list) and len(cast(list[object], followups)) > 0:
|
||||||
|
return True
|
||||||
|
state = recovery_state_from_metadata(metadata)
|
||||||
|
return bool(state and state["status"] in {"resuming", "awaiting_user", "failed"})
|
||||||
|
|
||||||
|
async def admit(self, message: InboundMessage) -> bool:
|
||||||
|
"""Reject stale queued recoveries and let new user input supersede them."""
|
||||||
|
recovery_id = message.metadata.get(RECOVERY_INBOUND_METADATA_KEY)
|
||||||
|
if isinstance(recovery_id, str):
|
||||||
|
session = await self._get_or_create_session(message.session_key)
|
||||||
|
state = recovery_state_from_metadata(session.metadata)
|
||||||
|
return bool(
|
||||||
|
state
|
||||||
|
and state["status"] == "resuming"
|
||||||
|
and state["recovery_id"] == recovery_id
|
||||||
|
)
|
||||||
|
if message.channel != "websocket":
|
||||||
|
return True
|
||||||
|
session = await self._get_or_create_session(message.session_key)
|
||||||
|
state = recovery_state_from_metadata(session.metadata)
|
||||||
|
if state and state["status"] in {"resuming", "awaiting_user", "failed"}:
|
||||||
|
await self._cancel_active_recovery(message.session_key)
|
||||||
|
restore_runtime_checkpoint(session)
|
||||||
|
if not restore_pending_interruption(session, superseded=True):
|
||||||
|
append_recovery_interruption(session, superseded=True)
|
||||||
|
recovered = self._set_state(
|
||||||
|
session,
|
||||||
|
status="recovered",
|
||||||
|
recovery_id=cast(str, state["recovery_id"]),
|
||||||
|
attempts=cast(int, state.get("attempts", 0)),
|
||||||
|
reason="superseded",
|
||||||
|
)
|
||||||
|
await self._save_session(session)
|
||||||
|
await self._publish(message.chat_id, recovered)
|
||||||
|
return True
|
||||||
|
|
||||||
|
async def turn_completed(self, session_key: str) -> None:
|
||||||
|
"""Resolve a resuming state after the recovered turn commits."""
|
||||||
|
session = await self._get_or_create_session(session_key)
|
||||||
|
state = recovery_state_from_metadata(session.metadata)
|
||||||
|
if not state or state["status"] != "resuming":
|
||||||
|
return
|
||||||
|
route = self._websocket_route(session)
|
||||||
|
if route is None:
|
||||||
|
return
|
||||||
|
recovered = self._set_state(
|
||||||
|
session,
|
||||||
|
status="recovered",
|
||||||
|
recovery_id=cast(str, state["recovery_id"]),
|
||||||
|
attempts=cast(int, state.get("attempts", 0)),
|
||||||
|
reason="continued",
|
||||||
|
)
|
||||||
|
await self._save_session(session)
|
||||||
|
await self._publish(route[1], recovered)
|
||||||
|
|
||||||
|
async def handle_action(self, action: str, payload: dict[str, Any]) -> dict[str, Any]:
|
||||||
|
"""Apply an authenticated continue/dismiss operation."""
|
||||||
|
chat_id = payload.get("chat_id")
|
||||||
|
recovery_id = payload.get("recovery_id")
|
||||||
|
if not isinstance(chat_id, str) or not chat_id:
|
||||||
|
raise RecoveryActionError("missing chat_id")
|
||||||
|
if not isinstance(recovery_id, str) or not recovery_id:
|
||||||
|
raise RecoveryActionError("missing recovery_id")
|
||||||
|
session = await self._get_or_create_session(self._session_key(chat_id))
|
||||||
|
state = recovery_state_from_metadata(session.metadata)
|
||||||
|
if not state or state["recovery_id"] != recovery_id:
|
||||||
|
raise RecoveryActionError("recovery state is stale", status=409)
|
||||||
|
|
||||||
|
if action == "dismiss":
|
||||||
|
restore_runtime_checkpoint(session)
|
||||||
|
restore_pending_interruption(session)
|
||||||
|
next_state = self._set_state(
|
||||||
|
session,
|
||||||
|
status="recovered",
|
||||||
|
recovery_id=recovery_id,
|
||||||
|
attempts=cast(int, state.get("attempts", 0)),
|
||||||
|
reason="dismissed",
|
||||||
|
)
|
||||||
|
await self._save_session(session)
|
||||||
|
await self._publish(chat_id, next_state)
|
||||||
|
return next_state
|
||||||
|
if action != "continue":
|
||||||
|
raise RecoveryActionError("unknown recovery action")
|
||||||
|
if state["status"] not in {"awaiting_user", "failed"}:
|
||||||
|
raise RecoveryActionError("recovery is not waiting for confirmation", status=409)
|
||||||
|
if state.get("can_continue") is False:
|
||||||
|
raise RecoveryActionError("recovery context is unavailable", status=409)
|
||||||
|
next_state = self._set_state(
|
||||||
|
session,
|
||||||
|
status="resuming",
|
||||||
|
recovery_id=recovery_id,
|
||||||
|
attempts=cast(int, state.get("attempts", 0)) + 1,
|
||||||
|
reason="user_confirmed",
|
||||||
|
resume_message_count=len(session.messages),
|
||||||
|
)
|
||||||
|
await self._save_session(session)
|
||||||
|
await self._publish(chat_id, next_state)
|
||||||
|
await self._queue_continuation(session, chat_id, next_state)
|
||||||
|
return next_state
|
||||||
|
|
||||||
|
async def _recover_session(self, session: Session, chat_id: str) -> None:
|
||||||
|
checkpoint_value = cast(object, session.metadata.get(RUNTIME_CHECKPOINT_KEY))
|
||||||
|
checkpoint = (
|
||||||
|
cast(dict[str, Any], checkpoint_value)
|
||||||
|
if isinstance(checkpoint_value, dict)
|
||||||
|
else None
|
||||||
|
)
|
||||||
|
pending = session.metadata.get(PENDING_USER_TURN_KEY) is True
|
||||||
|
state = recovery_state_from_metadata(session.metadata)
|
||||||
|
if not pending and checkpoint is None:
|
||||||
|
if state and state["status"] == "resuming":
|
||||||
|
resume_count = self._resume_message_count(session)
|
||||||
|
if resume_count is not None and len(session.messages) > resume_count:
|
||||||
|
next_state = self._set_state(
|
||||||
|
session,
|
||||||
|
status="recovered",
|
||||||
|
recovery_id=cast(str, state["recovery_id"]),
|
||||||
|
attempts=cast(int, state.get("attempts", 0)),
|
||||||
|
reason="committed",
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
next_state = self._set_state(
|
||||||
|
session,
|
||||||
|
status="awaiting_user",
|
||||||
|
recovery_id=cast(str, state["recovery_id"]),
|
||||||
|
attempts=cast(int, state.get("attempts", 1)),
|
||||||
|
reason="loop_guard",
|
||||||
|
)
|
||||||
|
await self._save_session(session)
|
||||||
|
await self._publish(chat_id, next_state)
|
||||||
|
elif self._has_unfinished_webui_transcript(session.key):
|
||||||
|
# A normal last-client shutdown can materialize the checkpoint
|
||||||
|
# before the process exits. In that path there is no pending
|
||||||
|
# marker left to classify, but the append-only transcript still
|
||||||
|
# contains an activity row without a turn_end. Treat it as an
|
||||||
|
# interrupted turn instead of letting the UI resurrect it as a
|
||||||
|
# forever-running spinner.
|
||||||
|
can_continue = self._has_saved_continuation_context(session)
|
||||||
|
waiting = self._set_state(
|
||||||
|
session,
|
||||||
|
status="awaiting_user",
|
||||||
|
recovery_id=uuid4().hex,
|
||||||
|
attempts=0,
|
||||||
|
reason=(
|
||||||
|
"interrupted_with_saved_context"
|
||||||
|
if can_continue
|
||||||
|
else "interrupted_without_checkpoint"
|
||||||
|
),
|
||||||
|
can_continue=can_continue,
|
||||||
|
)
|
||||||
|
await self._save_session(session)
|
||||||
|
await self._publish(chat_id, waiting)
|
||||||
|
return
|
||||||
|
if state and state["status"] in {"awaiting_user", "failed"}:
|
||||||
|
await self._publish(chat_id, state)
|
||||||
|
return
|
||||||
|
if state and state["status"] == "resuming":
|
||||||
|
restore_runtime_checkpoint(session)
|
||||||
|
restore_pending_interruption(session)
|
||||||
|
waiting = self._set_state(
|
||||||
|
session,
|
||||||
|
status="awaiting_user",
|
||||||
|
recovery_id=cast(str, state["recovery_id"]),
|
||||||
|
attempts=cast(int, state.get("attempts", 1)),
|
||||||
|
reason="loop_guard",
|
||||||
|
)
|
||||||
|
await self._save_session(session)
|
||||||
|
await self._publish(chat_id, waiting)
|
||||||
|
return
|
||||||
|
|
||||||
|
recovery_id = uuid4().hex
|
||||||
|
phase = checkpoint.get("phase") if checkpoint is not None else None
|
||||||
|
pending_calls = checkpoint.get("pending_tool_calls") if checkpoint is not None else None
|
||||||
|
if checkpoint is not None and phase not in _KNOWN_CHECKPOINT_PHASES:
|
||||||
|
_discard_runtime_checkpoint(session)
|
||||||
|
restore_pending_interruption(session)
|
||||||
|
waiting = self._set_state(
|
||||||
|
session,
|
||||||
|
status="awaiting_user",
|
||||||
|
recovery_id=recovery_id,
|
||||||
|
attempts=0,
|
||||||
|
reason="checkpoint_unknown",
|
||||||
|
can_continue=False,
|
||||||
|
)
|
||||||
|
await self._save_session(session)
|
||||||
|
await self._publish(chat_id, waiting)
|
||||||
|
return
|
||||||
|
if checkpoint is not None and not _runtime_checkpoint_is_well_formed(checkpoint):
|
||||||
|
_discard_runtime_checkpoint(session)
|
||||||
|
restore_pending_interruption(session)
|
||||||
|
waiting = self._set_state(
|
||||||
|
session,
|
||||||
|
status="awaiting_user",
|
||||||
|
recovery_id=recovery_id,
|
||||||
|
attempts=0,
|
||||||
|
reason="checkpoint_invalid",
|
||||||
|
can_continue=False,
|
||||||
|
)
|
||||||
|
await self._save_session(session)
|
||||||
|
await self._publish(chat_id, waiting)
|
||||||
|
return
|
||||||
|
if phase == "final_response":
|
||||||
|
restore_runtime_checkpoint(session)
|
||||||
|
recovered = self._set_state(
|
||||||
|
session,
|
||||||
|
status="recovered",
|
||||||
|
recovery_id=recovery_id,
|
||||||
|
attempts=0,
|
||||||
|
reason="answer_restored",
|
||||||
|
)
|
||||||
|
await self._save_session(session)
|
||||||
|
await self._publish(chat_id, recovered)
|
||||||
|
return
|
||||||
|
if phase in _UNCERTAIN_TOOL_PHASES or pending_calls:
|
||||||
|
restore_runtime_checkpoint(session)
|
||||||
|
waiting = self._set_state(
|
||||||
|
session,
|
||||||
|
status="awaiting_user",
|
||||||
|
recovery_id=recovery_id,
|
||||||
|
attempts=0,
|
||||||
|
reason="tool_state_unknown",
|
||||||
|
)
|
||||||
|
await self._save_session(session)
|
||||||
|
await self._publish(chat_id, waiting)
|
||||||
|
return
|
||||||
|
# A gateway restart is a lifecycle boundary. Never enqueue model work
|
||||||
|
# implicitly: even a synchronized checkpoint may sit next to an
|
||||||
|
# external side effect that the user should review first. The final
|
||||||
|
# answer path above only restores persisted output; it never executes.
|
||||||
|
restore_runtime_checkpoint(session)
|
||||||
|
waiting = self._set_state(
|
||||||
|
session,
|
||||||
|
status="awaiting_user",
|
||||||
|
recovery_id=recovery_id,
|
||||||
|
attempts=0,
|
||||||
|
reason="restart_requires_confirmation",
|
||||||
|
)
|
||||||
|
await self._save_session(session)
|
||||||
|
await self._publish(chat_id, waiting)
|
||||||
|
|
||||||
|
async def _queue_continuation(
|
||||||
|
self,
|
||||||
|
session: Session,
|
||||||
|
chat_id: str,
|
||||||
|
state: Mapping[str, Any],
|
||||||
|
) -> None:
|
||||||
|
recovery_id = cast(str, state["recovery_id"])
|
||||||
|
await self.bus.publish_inbound(
|
||||||
|
InboundMessage(
|
||||||
|
channel="websocket",
|
||||||
|
sender_id="system:recovery",
|
||||||
|
chat_id=chat_id,
|
||||||
|
content=(
|
||||||
|
"Continue the interrupted request from the saved conversation context. "
|
||||||
|
"Do not repeat completed work or mention the restart unless it affects the answer."
|
||||||
|
),
|
||||||
|
metadata={
|
||||||
|
"webui": True,
|
||||||
|
"_wants_stream": True,
|
||||||
|
WEBUI_TURN_METADATA_KEY: f"recovery:{recovery_id}",
|
||||||
|
RECOVERY_INBOUND_METADATA_KEY: recovery_id,
|
||||||
|
turn_continuation.INTERNAL_CONTINUATION_META: True,
|
||||||
|
turn_continuation.SKIP_USER_PERSIST_META: True,
|
||||||
|
},
|
||||||
|
session_key_override=session.key,
|
||||||
|
require_existing_session=True,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
async def _requeue_pending_followups(self, session: Session) -> None:
|
||||||
|
"""Return durable live-turn follow-ups to the bus after a restart."""
|
||||||
|
for message in pending_followups(session):
|
||||||
|
await self.bus.publish_inbound(message)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _resume_message_count(session: Session) -> int | None:
|
||||||
|
raw_value = cast(object, session.metadata.get(RECOVERY_METADATA_KEY))
|
||||||
|
value = cast(dict[str, Any], raw_value) if isinstance(raw_value, dict) else None
|
||||||
|
if value is None:
|
||||||
|
return None
|
||||||
|
count = value.get("resume_message_count")
|
||||||
|
return count if isinstance(count, int) and count >= 0 else None
|
||||||
|
|
||||||
|
async def _publish(
|
||||||
|
self,
|
||||||
|
chat_id: str,
|
||||||
|
state: Mapping[str, Any],
|
||||||
|
) -> None:
|
||||||
|
"""Publish the recovery state and invalidate its sidebar projection."""
|
||||||
|
await self.bus.publish_outbound(
|
||||||
|
outbound_message_for_event(
|
||||||
|
channel="websocket",
|
||||||
|
chat_id=chat_id,
|
||||||
|
event=RecoveryStateEvent(
|
||||||
|
status=cast(str, state["status"]),
|
||||||
|
recovery_id=cast(str, state["recovery_id"]),
|
||||||
|
reason=cast(str | None, state.get("reason")),
|
||||||
|
attempts=cast(int, state.get("attempts", 0)),
|
||||||
|
can_continue=cast(bool | None, state.get("can_continue")),
|
||||||
|
),
|
||||||
|
)
|
||||||
|
)
|
||||||
|
await self.bus.publish_outbound(
|
||||||
|
outbound_message_for_event(
|
||||||
|
channel="websocket",
|
||||||
|
chat_id=chat_id,
|
||||||
|
event=SessionUpdatedEvent(scope="thread"),
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _set_state(
|
||||||
|
session: Session,
|
||||||
|
*,
|
||||||
|
status: str,
|
||||||
|
recovery_id: str,
|
||||||
|
attempts: int,
|
||||||
|
reason: str,
|
||||||
|
resume_message_count: int | None = None,
|
||||||
|
can_continue: bool = True,
|
||||||
|
) -> dict[str, Any]:
|
||||||
|
state = {
|
||||||
|
"status": status,
|
||||||
|
"recovery_id": recovery_id,
|
||||||
|
"attempts": max(0, attempts),
|
||||||
|
"reason": reason,
|
||||||
|
"updated_at": datetime.now().isoformat(),
|
||||||
|
}
|
||||||
|
if not can_continue:
|
||||||
|
state["can_continue"] = False
|
||||||
|
if resume_message_count is not None:
|
||||||
|
state["resume_message_count"] = max(0, resume_message_count)
|
||||||
|
session.metadata[RECOVERY_METADATA_KEY] = state
|
||||||
|
session.updated_at = datetime.now()
|
||||||
|
return state
|
||||||
|
|
||||||
|
def _session_key(self, chat_id: str) -> str:
|
||||||
|
return UNIFIED_SESSION_KEY if self.unified_session else f"websocket:{chat_id}"
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _has_unfinished_webui_transcript(session_key: str) -> bool:
|
||||||
|
"""Detect a stale WebUI activity tail after an unclean gateway stop.
|
||||||
|
|
||||||
|
The transcript is intentionally consulted only as a last-resort signal:
|
||||||
|
a durable pending turn or runtime checkpoint always takes precedence.
|
||||||
|
This keeps browser disconnects harmless while preventing a materialized
|
||||||
|
partial turn from being presented as active forever after a restart.
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
from nanobot.webui.transcript import has_unfinished_transcript_tail
|
||||||
|
|
||||||
|
return has_unfinished_transcript_tail(session_key)
|
||||||
|
except (OSError, ValueError, TypeError):
|
||||||
|
# Recovery must fail closed if the optional display transcript is
|
||||||
|
# corrupt or unavailable; the normal checkpoint path still applies.
|
||||||
|
return False
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _has_saved_continuation_context(session: Session) -> bool:
|
||||||
|
"""Whether an interrupted turn left model-visible context to continue from."""
|
||||||
|
last_user = next(
|
||||||
|
(
|
||||||
|
index
|
||||||
|
for index in range(len(session.messages) - 1, -1, -1)
|
||||||
|
if session.messages[index].get("role") == "user"
|
||||||
|
),
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
if last_user is None:
|
||||||
|
return False
|
||||||
|
tail = session.messages[last_user + 1 :]
|
||||||
|
return bool(tail) and (
|
||||||
|
tail[-1].get("role") == "tool"
|
||||||
|
or any(message.get("_recovery_interrupted") is True for message in tail)
|
||||||
|
or any(
|
||||||
|
message.get("role") == "assistant" and bool(message.get("tool_calls"))
|
||||||
|
for message in tail
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _websocket_route(session: Session) -> tuple[str, str] | None:
|
||||||
|
return RecoveryCoordinator._websocket_route_for(session.key, session.metadata)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _websocket_route_for(
|
||||||
|
session_key: str,
|
||||||
|
metadata: Mapping[str, Any],
|
||||||
|
) -> tuple[str, str] | None:
|
||||||
|
if session_key.startswith("websocket:"):
|
||||||
|
chat_id = session_key.split(":", 1)[1]
|
||||||
|
return ("websocket", chat_id) if chat_id else None
|
||||||
|
if session_key == UNIFIED_SESSION_KEY:
|
||||||
|
route = last_channel_from_metadata(metadata)
|
||||||
|
if route and route[0] == "websocket":
|
||||||
|
return route
|
||||||
|
return None
|
||||||
@@ -2,6 +2,7 @@
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
import re
|
import re
|
||||||
import time
|
import time
|
||||||
from collections.abc import Awaitable, Callable
|
from collections.abc import Awaitable, Callable
|
||||||
@@ -37,12 +38,14 @@ from nanobot.bus.runtime_events import (
|
|||||||
TurnRuntimeAdmitted,
|
TurnRuntimeAdmitted,
|
||||||
UserInputAccepted,
|
UserInputAccepted,
|
||||||
)
|
)
|
||||||
from nanobot.providers.base import LLMProvider
|
from nanobot.llm_usage.context import llm_usage_source
|
||||||
|
from nanobot.providers.base import LLMProvider, LLMUsage
|
||||||
from nanobot.providers.fallback_provider import FallbackModelObserver
|
from nanobot.providers.fallback_provider import FallbackModelObserver
|
||||||
from nanobot.runtime_context import public_history_message
|
from nanobot.runtime_context import public_history_message
|
||||||
from nanobot.session.goal_state import goal_state_ws_blob
|
from nanobot.session.goal_state import goal_state_ws_blob
|
||||||
from nanobot.session.history_visibility import is_hidden_history_message
|
from nanobot.session.history_visibility import is_hidden_history_message
|
||||||
from nanobot.session.manager import Session, SessionManager
|
from nanobot.session.manager import Session, SessionManager
|
||||||
|
from nanobot.session.recovery import RecoveryCoordinator
|
||||||
from nanobot.session.session_handles import session_handle_for_name
|
from nanobot.session.session_handles import session_handle_for_name
|
||||||
from nanobot.session.session_messages import (
|
from nanobot.session.session_messages import (
|
||||||
SessionMessageEnvelope,
|
SessionMessageEnvelope,
|
||||||
@@ -174,7 +177,7 @@ async def maybe_generate_webui_title(
|
|||||||
model: str,
|
model: str,
|
||||||
) -> bool:
|
) -> bool:
|
||||||
"""Generate and persist a short title for WebUI-owned sessions only."""
|
"""Generate and persist a short title for WebUI-owned sessions only."""
|
||||||
session = sessions.get_or_create(session_key)
|
session = await sessions.get_or_create_async(session_key)
|
||||||
if session.metadata.get(WEBUI_SESSION_METADATA_KEY) is not True:
|
if session.metadata.get(WEBUI_SESSION_METADATA_KEY) is not True:
|
||||||
return False
|
return False
|
||||||
if session.metadata.get(WEBUI_TITLE_USER_EDITED_METADATA_KEY) is True:
|
if session.metadata.get(WEBUI_TITLE_USER_EDITED_METADATA_KEY) is True:
|
||||||
@@ -185,7 +188,7 @@ async def maybe_generate_webui_title(
|
|||||||
if cleaned_current_title:
|
if cleaned_current_title:
|
||||||
if cleaned_current_title != current_title:
|
if cleaned_current_title != current_title:
|
||||||
session.metadata[WEBUI_TITLE_METADATA_KEY] = cleaned_current_title
|
session.metadata[WEBUI_TITLE_METADATA_KEY] = cleaned_current_title
|
||||||
sessions.save(session)
|
await sessions.save_async(session)
|
||||||
return False
|
return False
|
||||||
session.metadata.pop(WEBUI_TITLE_METADATA_KEY, None)
|
session.metadata.pop(WEBUI_TITLE_METADATA_KEY, None)
|
||||||
|
|
||||||
@@ -207,24 +210,25 @@ async def maybe_generate_webui_title(
|
|||||||
prompt += f"\nAssistant: {truncate_text(assistant_text, 1_000)}"
|
prompt += f"\nAssistant: {truncate_text(assistant_text, 1_000)}"
|
||||||
|
|
||||||
try:
|
try:
|
||||||
response = await provider.chat_with_retry(
|
with llm_usage_source("system"):
|
||||||
[
|
response = await provider.chat_with_retry(
|
||||||
{
|
[
|
||||||
"role": "system",
|
{
|
||||||
"content": (
|
"role": "system",
|
||||||
"You write short, neutral chat titles. "
|
"content": (
|
||||||
"Return only the title text."
|
"You write short, neutral chat titles. "
|
||||||
),
|
"Return only the title text."
|
||||||
},
|
),
|
||||||
{"role": "user", "content": prompt},
|
},
|
||||||
],
|
{"role": "user", "content": prompt},
|
||||||
tools=None,
|
],
|
||||||
model=model,
|
tools=None,
|
||||||
max_tokens=TITLE_GENERATION_MAX_TOKENS,
|
model=model,
|
||||||
temperature=0.2,
|
max_tokens=TITLE_GENERATION_MAX_TOKENS,
|
||||||
reasoning_effort=TITLE_GENERATION_REASONING_EFFORT,
|
temperature=0.2,
|
||||||
retry_mode="standard",
|
reasoning_effort=TITLE_GENERATION_REASONING_EFFORT,
|
||||||
)
|
retry_mode="standard",
|
||||||
|
)
|
||||||
except Exception:
|
except Exception:
|
||||||
logger.debug("Failed to generate webui session title for {}", session_key, exc_info=True)
|
logger.debug("Failed to generate webui session title for {}", session_key, exc_info=True)
|
||||||
return False
|
return False
|
||||||
@@ -238,7 +242,7 @@ async def maybe_generate_webui_title(
|
|||||||
)
|
)
|
||||||
return False
|
return False
|
||||||
session.metadata[WEBUI_TITLE_METADATA_KEY] = title
|
session.metadata[WEBUI_TITLE_METADATA_KEY] = title
|
||||||
sessions.save(session)
|
await sessions.save_async(session)
|
||||||
return True
|
return True
|
||||||
|
|
||||||
|
|
||||||
@@ -434,8 +438,8 @@ class WebuiTurnRoutePolicy:
|
|||||||
)
|
)
|
||||||
and route.channel == "websocket"
|
and route.channel == "websocket"
|
||||||
):
|
):
|
||||||
session = self.sessions.get_or_create(session_key)
|
session = self.sessions.get_cached(session_key)
|
||||||
if session.metadata.get(WEBUI_SESSION_METADATA_KEY) is True:
|
if session is not None and session.metadata.get(WEBUI_SESSION_METADATA_KEY) is True:
|
||||||
metadata = dict(route.metadata)
|
metadata = dict(route.metadata)
|
||||||
turn_prefix = "session-input" if internal_user_input else "subagent"
|
turn_prefix = "session-input" if internal_user_input else "subagent"
|
||||||
metadata.update({
|
metadata.update({
|
||||||
@@ -511,6 +515,7 @@ class WebuiTurnCoordinator:
|
|||||||
bus: MessageBus
|
bus: MessageBus
|
||||||
sessions: SessionManager
|
sessions: SessionManager
|
||||||
schedule_background: Callable[[Awaitable[None]], None]
|
schedule_background: Callable[[Awaitable[None]], None]
|
||||||
|
recovery: RecoveryCoordinator | None = None
|
||||||
|
|
||||||
def subscribe(self, runtime_events: RuntimeEventBus) -> Callable[[], None]:
|
def subscribe(self, runtime_events: RuntimeEventBus) -> Callable[[], None]:
|
||||||
"""Subscribe this coordinator to runtime events."""
|
"""Subscribe this coordinator to runtime events."""
|
||||||
@@ -576,7 +581,7 @@ class WebuiTurnCoordinator:
|
|||||||
or not session_key.startswith("websocket:")
|
or not session_key.startswith("websocket:")
|
||||||
):
|
):
|
||||||
return
|
return
|
||||||
persisted = self.sessions.read_session_metadata(session_key)
|
persisted = await self.sessions.read_session_metadata_async(session_key)
|
||||||
metadata_value: object = persisted.get("metadata") if persisted is not None else None
|
metadata_value: object = persisted.get("metadata") if persisted is not None else None
|
||||||
metadata = (
|
metadata = (
|
||||||
cast(dict[str, Any], metadata_value)
|
cast(dict[str, Any], metadata_value)
|
||||||
@@ -587,7 +592,8 @@ class WebuiTurnCoordinator:
|
|||||||
return
|
return
|
||||||
public_metadata = _session_message_public_metadata(envelope)
|
public_metadata = _session_message_public_metadata(envelope)
|
||||||
try:
|
try:
|
||||||
append_session_message_input(
|
await asyncio.to_thread(
|
||||||
|
append_session_message_input,
|
||||||
session_key,
|
session_key,
|
||||||
content=event.content,
|
content=event.content,
|
||||||
created_at_ms=envelope["created_at_ms"],
|
created_at_ms=envelope["created_at_ms"],
|
||||||
@@ -612,8 +618,9 @@ class WebuiTurnCoordinator:
|
|||||||
def _handle_session_turn_started(self, event: SessionTurnStarted) -> None:
|
def _handle_session_turn_started(self, event: SessionTurnStarted) -> None:
|
||||||
if not self._is_websocket_event(event.context):
|
if not self._is_websocket_event(event.context):
|
||||||
return
|
return
|
||||||
session = self.sessions.get_or_create(event.context.session_key)
|
session = self.sessions.get_cached(event.context.session_key)
|
||||||
mark_webui_session(session, event.context.metadata)
|
if session is not None:
|
||||||
|
mark_webui_session(session, event.context.metadata)
|
||||||
|
|
||||||
async def _handle_run_status_changed(self, event: TurnRunStatusChanged) -> None:
|
async def _handle_run_status_changed(self, event: TurnRunStatusChanged) -> None:
|
||||||
if not self._is_websocket_event(event.context):
|
if not self._is_websocket_event(event.context):
|
||||||
@@ -654,6 +661,8 @@ class WebuiTurnCoordinator:
|
|||||||
event.runtime.context_window_tokens if event.runtime is not None else None
|
event.runtime.context_window_tokens if event.runtime is not None else None
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
|
if self.recovery is not None:
|
||||||
|
await self.recovery.turn_completed(event.context.session_key)
|
||||||
self._schedule_title_update_from_event(event)
|
self._schedule_title_update_from_event(event)
|
||||||
|
|
||||||
async def _handle_goal_state_changed(self, event: GoalStateChanged) -> None:
|
async def _handle_goal_state_changed(self, event: GoalStateChanged) -> None:
|
||||||
@@ -685,28 +694,19 @@ class WebuiTurnCoordinator:
|
|||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|
||||||
async def publish_run_status(
|
|
||||||
self,
|
|
||||||
msg: InboundMessage,
|
|
||||||
status: str,
|
|
||||||
*,
|
|
||||||
started_at: float | None = None,
|
|
||||||
) -> None:
|
|
||||||
await publish_turn_run_status(self.bus, msg, status, started_at=started_at)
|
|
||||||
|
|
||||||
async def handle_turn_end(
|
async def handle_turn_end(
|
||||||
self,
|
self,
|
||||||
msg: InboundMessage,
|
msg: InboundMessage,
|
||||||
*,
|
*,
|
||||||
session_key: str,
|
session_key: str,
|
||||||
latency_ms: int | None,
|
latency_ms: int | None,
|
||||||
usage: dict[str, int] | None = None,
|
usage: LLMUsage | None = None,
|
||||||
context_window_tokens: int | None = None,
|
context_window_tokens: int | None = None,
|
||||||
) -> None:
|
) -> None:
|
||||||
if msg.channel != "websocket":
|
if msg.channel != "websocket":
|
||||||
return
|
return
|
||||||
|
|
||||||
session = self.sessions.get_or_create(session_key)
|
session = await self.sessions.get_or_create_async(session_key)
|
||||||
await self.bus.publish_outbound(
|
await self.bus.publish_outbound(
|
||||||
outbound_message_for_event(
|
outbound_message_for_event(
|
||||||
channel=msg.channel,
|
channel=msg.channel,
|
||||||
@@ -714,7 +714,7 @@ class WebuiTurnCoordinator:
|
|||||||
event=TurnEndEvent(
|
event=TurnEndEvent(
|
||||||
latency_ms=latency_ms,
|
latency_ms=latency_ms,
|
||||||
goal_state=goal_state_ws_blob(session.metadata),
|
goal_state=goal_state_ws_blob(session.metadata),
|
||||||
usage=usage or None,
|
usage=usage,
|
||||||
context_window_tokens=context_window_tokens,
|
context_window_tokens=context_window_tokens,
|
||||||
),
|
),
|
||||||
metadata=msg.metadata,
|
metadata=msg.metadata,
|
||||||
|
|||||||
@@ -16,7 +16,7 @@ Concrete scenarios showing when and how to use the my tool effectively.
|
|||||||
→ my(action="check", key="max_iterations")
|
→ my(action="check", key="max_iterations")
|
||||||
→ 40
|
→ 40
|
||||||
→ my(action="check", key="_last_usage")
|
→ my(action="check", key="_last_usage")
|
||||||
→ {"prompt_tokens": 62000, "completion_tokens": 3000}
|
→ {"input_tokens": 62000, "output_tokens": 3000}
|
||||||
→ "I hit the iteration limit (40). The task was complex. I can ask the user if they want to increase it."
|
→ "I hit the iteration limit (40). The task was complex. I can ask the user if they want to increase it."
|
||||||
```
|
```
|
||||||
|
|
||||||
@@ -72,6 +72,6 @@ Concrete scenarios showing when and how to use the my tool effectively.
|
|||||||
### Token-conscious behavior
|
### Token-conscious behavior
|
||||||
```
|
```
|
||||||
→ my(action="check", key="_last_usage")
|
→ my(action="check", key="_last_usage")
|
||||||
→ {"prompt_tokens": 58000, "completion_tokens": 12000}
|
→ {"input_tokens": 58000, "output_tokens": 12000}
|
||||||
→ "I've consumed ~70k tokens. I'll keep my remaining responses focused."
|
→ "I've consumed ~70k tokens. I'll keep my remaining responses focused."
|
||||||
```
|
```
|
||||||
|
|||||||
@@ -5,17 +5,24 @@ from __future__ import annotations
|
|||||||
import asyncio
|
import asyncio
|
||||||
import uuid
|
import uuid
|
||||||
from collections.abc import Awaitable, Callable
|
from collections.abc import Awaitable, Callable
|
||||||
from typing import Any
|
from contextlib import suppress
|
||||||
|
from typing import Any, TypeVar
|
||||||
|
|
||||||
from loguru import logger
|
from loguru import logger
|
||||||
|
|
||||||
from nanobot.agent.automation_turns import AutomationTurnError
|
from nanobot.agent.automation_turns import (
|
||||||
|
AutomationTurnAcceptedCancellation,
|
||||||
|
AutomationTurnError,
|
||||||
|
)
|
||||||
from nanobot.bus.events import InboundMessage, OutboundMessage
|
from nanobot.bus.events import InboundMessage, OutboundMessage
|
||||||
from nanobot.triggers.local_session_turns import LOCAL_TRIGGER_META
|
from nanobot.triggers.local_session_turns import LOCAL_TRIGGER_META
|
||||||
from nanobot.triggers.local_store import LocalTriggerStore
|
from nanobot.triggers.local_store import LocalTriggerStore
|
||||||
from nanobot.triggers.local_types import LocalTrigger, TriggerDelivery
|
from nanobot.triggers.local_types import LocalTrigger, TriggerDelivery
|
||||||
|
from nanobot.utils.cancellation import shield_and_drain
|
||||||
from nanobot.webui.metadata import WEBUI_MESSAGE_SOURCE_METADATA_KEY, WEBUI_TURN_METADATA_KEY
|
from nanobot.webui.metadata import WEBUI_MESSAGE_SOURCE_METADATA_KEY, WEBUI_TURN_METADATA_KEY
|
||||||
|
|
||||||
|
_T = TypeVar("_T")
|
||||||
|
|
||||||
|
|
||||||
async def run_local_trigger_queue(
|
async def run_local_trigger_queue(
|
||||||
*,
|
*,
|
||||||
@@ -29,14 +36,16 @@ async def run_local_trigger_queue(
|
|||||||
if submit_turn is None:
|
if submit_turn is None:
|
||||||
raise ValueError("run_local_trigger_queue requires submit_turn")
|
raise ValueError("run_local_trigger_queue requires submit_turn")
|
||||||
logger.info("Local trigger queue started")
|
logger.info("Local trigger queue started")
|
||||||
recovered = store.recover_processing_deliveries()
|
recovered = await shield_and_drain(asyncio.to_thread(store.recover_processing_deliveries))
|
||||||
if recovered:
|
if recovered:
|
||||||
logger.warning(
|
logger.warning(
|
||||||
"Trigger: recovered {} interrupted delivery file(s) from processing",
|
"Trigger: recovered {} interrupted delivery file(s) from processing",
|
||||||
recovered,
|
recovered,
|
||||||
)
|
)
|
||||||
while True:
|
while True:
|
||||||
deliveries = store.claim_deliveries(limit=batch_size)
|
deliveries = await shield_and_drain(
|
||||||
|
asyncio.to_thread(store.claim_deliveries, limit=batch_size)
|
||||||
|
)
|
||||||
if not deliveries:
|
if not deliveries:
|
||||||
await asyncio.sleep(poll_interval_s)
|
await asyncio.sleep(poll_interval_s)
|
||||||
continue
|
continue
|
||||||
@@ -49,30 +58,24 @@ async def run_local_trigger_queue(
|
|||||||
submit_turn=submit_turn,
|
submit_turn=submit_turn,
|
||||||
is_channel_enabled=is_channel_enabled,
|
is_channel_enabled=is_channel_enabled,
|
||||||
)
|
)
|
||||||
store.complete_delivery(delivery)
|
except _DeliverySettledOnCancellation:
|
||||||
|
raise
|
||||||
except asyncio.CancelledError as exc:
|
except asyncio.CancelledError as exc:
|
||||||
store.retry_delivery(delivery, str(exc) or exc.__class__.__name__)
|
error = str(exc) or exc.__class__.__name__
|
||||||
_write_delivery_run_record(
|
await shield_and_drain(asyncio.to_thread(store.retry_delivery, delivery, error))
|
||||||
|
await _write_delivery_run_record(
|
||||||
store,
|
store,
|
||||||
delivery,
|
delivery,
|
||||||
status="interrupted",
|
status="interrupted",
|
||||||
error=str(exc) or exc.__class__.__name__,
|
error=error,
|
||||||
)
|
)
|
||||||
raise
|
raise
|
||||||
except _TerminalDeliveryError as exc:
|
except _TerminalDeliveryError as exc:
|
||||||
store.record_delivery(
|
await _await_delivery_settlement(
|
||||||
delivery.trigger_id,
|
_settle_failed_delivery(store, delivery, error=str(exc)),
|
||||||
status="error",
|
store=store,
|
||||||
error=str(exc),
|
delivery=delivery,
|
||||||
run_at_ms=delivery.created_at_ms,
|
|
||||||
)
|
)
|
||||||
_write_delivery_run_record(
|
|
||||||
store,
|
|
||||||
delivery,
|
|
||||||
status="error",
|
|
||||||
error=str(exc),
|
|
||||||
)
|
|
||||||
store.complete_delivery(delivery)
|
|
||||||
logger.warning(
|
logger.warning(
|
||||||
"Trigger: dropped delivery {} for {}: {}",
|
"Trigger: dropped delivery {} for {}: {}",
|
||||||
delivery.id,
|
delivery.id,
|
||||||
@@ -81,19 +84,11 @@ async def run_local_trigger_queue(
|
|||||||
)
|
)
|
||||||
except AutomationTurnError as exc:
|
except AutomationTurnError as exc:
|
||||||
error = str(exc) or exc.__class__.__name__
|
error = str(exc) or exc.__class__.__name__
|
||||||
store.record_delivery(
|
await _await_delivery_settlement(
|
||||||
delivery.trigger_id,
|
_settle_failed_delivery(store, delivery, error=error),
|
||||||
status="error",
|
store=store,
|
||||||
error=error,
|
delivery=delivery,
|
||||||
run_at_ms=delivery.created_at_ms,
|
|
||||||
)
|
)
|
||||||
_write_delivery_run_record(
|
|
||||||
store,
|
|
||||||
delivery,
|
|
||||||
status="error",
|
|
||||||
error=error,
|
|
||||||
)
|
|
||||||
store.complete_delivery(delivery)
|
|
||||||
logger.warning(
|
logger.warning(
|
||||||
"Trigger: delivery {} for {} reached the agent but failed: {}",
|
"Trigger: delivery {} for {} reached the agent but failed: {}",
|
||||||
delivery.id,
|
delivery.id,
|
||||||
@@ -102,18 +97,10 @@ async def run_local_trigger_queue(
|
|||||||
)
|
)
|
||||||
except Exception as exc:
|
except Exception as exc:
|
||||||
error = str(exc) or exc.__class__.__name__
|
error = str(exc) or exc.__class__.__name__
|
||||||
retried = store.retry_delivery(delivery, error)
|
retried = await _await_delivery_settlement(
|
||||||
_write_delivery_run_record(
|
_settle_retryable_delivery(store, delivery, error=error),
|
||||||
store,
|
store=store,
|
||||||
delivery,
|
delivery=delivery,
|
||||||
status="retrying" if retried else "error",
|
|
||||||
error=error,
|
|
||||||
)
|
|
||||||
store.record_delivery(
|
|
||||||
delivery.trigger_id,
|
|
||||||
status="error",
|
|
||||||
error=error,
|
|
||||||
run_at_ms=delivery.created_at_ms,
|
|
||||||
)
|
)
|
||||||
logger.exception(
|
logger.exception(
|
||||||
"Trigger: failed delivery {} for {}{}",
|
"Trigger: failed delivery {} for {}{}",
|
||||||
@@ -127,6 +114,10 @@ class _TerminalDeliveryError(RuntimeError):
|
|||||||
pass
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
class _DeliverySettledOnCancellation(asyncio.CancelledError):
|
||||||
|
"""Cancellation reported only after an already-submitted delivery is settled."""
|
||||||
|
|
||||||
|
|
||||||
async def _deliver_delivery(
|
async def _deliver_delivery(
|
||||||
store: LocalTriggerStore,
|
store: LocalTriggerStore,
|
||||||
delivery: TriggerDelivery,
|
delivery: TriggerDelivery,
|
||||||
@@ -134,7 +125,7 @@ async def _deliver_delivery(
|
|||||||
submit_turn: Callable[[InboundMessage], Awaitable[OutboundMessage | None]],
|
submit_turn: Callable[[InboundMessage], Awaitable[OutboundMessage | None]],
|
||||||
is_channel_enabled: Callable[[str], bool],
|
is_channel_enabled: Callable[[str], bool],
|
||||||
) -> None:
|
) -> None:
|
||||||
trigger = store.get(delivery.trigger_id)
|
trigger = await asyncio.to_thread(store.get, delivery.trigger_id)
|
||||||
if trigger is None:
|
if trigger is None:
|
||||||
raise _TerminalDeliveryError("trigger not found")
|
raise _TerminalDeliveryError("trigger not found")
|
||||||
if not trigger.enabled:
|
if not trigger.enabled:
|
||||||
@@ -142,7 +133,14 @@ async def _deliver_delivery(
|
|||||||
if not is_channel_enabled(trigger.channel):
|
if not is_channel_enabled(trigger.channel):
|
||||||
raise _TerminalDeliveryError(f"target channel is not enabled: {trigger.channel}")
|
raise _TerminalDeliveryError(f"target channel is not enabled: {trigger.channel}")
|
||||||
|
|
||||||
store.write_delivery_run_record(delivery, trigger=trigger, status="processing")
|
await shield_and_drain(
|
||||||
|
asyncio.to_thread(
|
||||||
|
store.write_delivery_run_record,
|
||||||
|
delivery,
|
||||||
|
trigger=trigger,
|
||||||
|
status="processing",
|
||||||
|
)
|
||||||
|
)
|
||||||
msg = InboundMessage(
|
msg = InboundMessage(
|
||||||
channel=trigger.channel,
|
channel=trigger.channel,
|
||||||
sender_id=trigger.sender_id,
|
sender_id=trigger.sender_id,
|
||||||
@@ -151,22 +149,177 @@ async def _deliver_delivery(
|
|||||||
metadata=_delivery_metadata(trigger, delivery),
|
metadata=_delivery_metadata(trigger, delivery),
|
||||||
session_key_override=trigger.session_key,
|
session_key_override=trigger.session_key,
|
||||||
)
|
)
|
||||||
response = await submit_turn(msg)
|
try:
|
||||||
store.record_delivery(
|
response = await submit_turn(msg)
|
||||||
trigger.id,
|
except AutomationTurnAcceptedCancellation:
|
||||||
status="ok",
|
try:
|
||||||
run_at_ms=delivery.created_at_ms,
|
await _await_delivery_settlement(
|
||||||
|
_settle_accepted_delivery(store, delivery, trigger=trigger),
|
||||||
|
store=store,
|
||||||
|
delivery=delivery,
|
||||||
|
)
|
||||||
|
except _DeliverySettledOnCancellation:
|
||||||
|
raise
|
||||||
|
except Exception:
|
||||||
|
logger.exception(
|
||||||
|
"Trigger: failed to persist accepted delivery {}; dropping retry",
|
||||||
|
delivery.id,
|
||||||
|
)
|
||||||
|
with suppress(Exception):
|
||||||
|
await shield_and_drain(asyncio.to_thread(store.complete_delivery, delivery))
|
||||||
|
raise _DeliverySettledOnCancellation from None
|
||||||
|
|
||||||
|
try:
|
||||||
|
await _await_delivery_settlement(
|
||||||
|
_settle_submitted_delivery(store, delivery, trigger=trigger, response=response),
|
||||||
|
store=store,
|
||||||
|
delivery=delivery,
|
||||||
|
)
|
||||||
|
except Exception:
|
||||||
|
logger.exception(
|
||||||
|
"Trigger: failed to persist status for submitted delivery {}; dropping retry",
|
||||||
|
delivery.id,
|
||||||
|
)
|
||||||
|
with suppress(Exception):
|
||||||
|
await shield_and_drain(asyncio.to_thread(store.complete_delivery, delivery))
|
||||||
|
|
||||||
|
|
||||||
|
async def _await_delivery_settlement(
|
||||||
|
operation: Awaitable[_T],
|
||||||
|
*,
|
||||||
|
store: LocalTriggerStore,
|
||||||
|
delivery: TriggerDelivery,
|
||||||
|
) -> _T:
|
||||||
|
settlement = asyncio.ensure_future(operation)
|
||||||
|
try:
|
||||||
|
return await asyncio.shield(settlement)
|
||||||
|
except asyncio.CancelledError:
|
||||||
|
while not settlement.done():
|
||||||
|
try:
|
||||||
|
await asyncio.shield(settlement)
|
||||||
|
except asyncio.CancelledError:
|
||||||
|
continue
|
||||||
|
try:
|
||||||
|
settlement.result()
|
||||||
|
except Exception:
|
||||||
|
logger.exception(
|
||||||
|
"Trigger: failed to settle delivery {} during cancellation",
|
||||||
|
delivery.id,
|
||||||
|
)
|
||||||
|
completion = asyncio.create_task(
|
||||||
|
shield_and_drain(asyncio.to_thread(store.complete_delivery, delivery))
|
||||||
|
)
|
||||||
|
while not completion.done():
|
||||||
|
try:
|
||||||
|
await asyncio.shield(completion)
|
||||||
|
except asyncio.CancelledError:
|
||||||
|
continue
|
||||||
|
with suppress(Exception):
|
||||||
|
completion.result()
|
||||||
|
raise _DeliverySettledOnCancellation from None
|
||||||
|
|
||||||
|
|
||||||
|
async def _settle_failed_delivery(
|
||||||
|
store: LocalTriggerStore,
|
||||||
|
delivery: TriggerDelivery,
|
||||||
|
*,
|
||||||
|
error: str,
|
||||||
|
) -> None:
|
||||||
|
await _write_delivery_run_record(
|
||||||
|
store,
|
||||||
|
delivery,
|
||||||
|
status="error",
|
||||||
|
error=error,
|
||||||
)
|
)
|
||||||
_write_delivery_run_record(
|
await shield_and_drain(asyncio.to_thread(store.complete_delivery, delivery))
|
||||||
|
# Publish the terminal status only after the durable delivery state is settled.
|
||||||
|
await shield_and_drain(
|
||||||
|
asyncio.to_thread(
|
||||||
|
store.record_delivery,
|
||||||
|
delivery.trigger_id,
|
||||||
|
status="error",
|
||||||
|
error=error,
|
||||||
|
run_at_ms=delivery.created_at_ms,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
async def _settle_retryable_delivery(
|
||||||
|
store: LocalTriggerStore,
|
||||||
|
delivery: TriggerDelivery,
|
||||||
|
*,
|
||||||
|
error: str,
|
||||||
|
) -> bool:
|
||||||
|
retried = await shield_and_drain(asyncio.to_thread(store.retry_delivery, delivery, error))
|
||||||
|
await _write_delivery_run_record(
|
||||||
|
store,
|
||||||
|
delivery,
|
||||||
|
status="retrying" if retried else "error",
|
||||||
|
error=error,
|
||||||
|
)
|
||||||
|
await shield_and_drain(
|
||||||
|
asyncio.to_thread(
|
||||||
|
store.record_delivery,
|
||||||
|
delivery.trigger_id,
|
||||||
|
status="error",
|
||||||
|
error=error,
|
||||||
|
run_at_ms=delivery.created_at_ms,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
return retried
|
||||||
|
|
||||||
|
|
||||||
|
async def _settle_accepted_delivery(
|
||||||
|
store: LocalTriggerStore,
|
||||||
|
delivery: TriggerDelivery,
|
||||||
|
*,
|
||||||
|
trigger: LocalTrigger,
|
||||||
|
) -> None:
|
||||||
|
"""Commit an accepted delivery without claiming the agent turn completed."""
|
||||||
|
await _write_delivery_run_record(
|
||||||
|
store,
|
||||||
|
delivery,
|
||||||
|
trigger=trigger,
|
||||||
|
status="accepted",
|
||||||
|
)
|
||||||
|
await shield_and_drain(asyncio.to_thread(store.complete_delivery, delivery))
|
||||||
|
await shield_and_drain(
|
||||||
|
asyncio.to_thread(
|
||||||
|
store.record_delivery,
|
||||||
|
trigger.id,
|
||||||
|
status="ok",
|
||||||
|
run_at_ms=delivery.created_at_ms,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
async def _settle_submitted_delivery(
|
||||||
|
store: LocalTriggerStore,
|
||||||
|
delivery: TriggerDelivery,
|
||||||
|
*,
|
||||||
|
trigger: LocalTrigger,
|
||||||
|
response: OutboundMessage | None,
|
||||||
|
) -> None:
|
||||||
|
await _write_delivery_run_record(
|
||||||
store,
|
store,
|
||||||
delivery,
|
delivery,
|
||||||
trigger=trigger,
|
trigger=trigger,
|
||||||
status="ok",
|
status="ok",
|
||||||
response=response.content if response else "",
|
response=response.content if response else "",
|
||||||
)
|
)
|
||||||
|
await shield_and_drain(asyncio.to_thread(store.complete_delivery, delivery))
|
||||||
|
# last_status is the externally visible commit marker for a settled delivery.
|
||||||
|
await shield_and_drain(
|
||||||
|
asyncio.to_thread(
|
||||||
|
store.record_delivery,
|
||||||
|
trigger.id,
|
||||||
|
status="ok",
|
||||||
|
run_at_ms=delivery.created_at_ms,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
def _write_delivery_run_record(
|
async def _write_delivery_run_record(
|
||||||
store: LocalTriggerStore,
|
store: LocalTriggerStore,
|
||||||
delivery: TriggerDelivery,
|
delivery: TriggerDelivery,
|
||||||
*,
|
*,
|
||||||
@@ -176,12 +329,15 @@ def _write_delivery_run_record(
|
|||||||
response: str | None = None,
|
response: str | None = None,
|
||||||
) -> None:
|
) -> None:
|
||||||
try:
|
try:
|
||||||
store.write_delivery_run_record(
|
await shield_and_drain(
|
||||||
delivery,
|
asyncio.to_thread(
|
||||||
trigger=trigger,
|
store.write_delivery_run_record,
|
||||||
status=status,
|
delivery,
|
||||||
error=error,
|
trigger=trigger,
|
||||||
response=response,
|
status=status,
|
||||||
|
error=error,
|
||||||
|
response=response,
|
||||||
|
)
|
||||||
)
|
)
|
||||||
except Exception:
|
except Exception:
|
||||||
logger.exception(
|
logger.exception(
|
||||||
|
|||||||
@@ -24,6 +24,7 @@ _MAX_RUN_HISTORY = 20
|
|||||||
_MAX_DELIVERY_ATTEMPTS = 10
|
_MAX_DELIVERY_ATTEMPTS = 10
|
||||||
_RUN_RECORD_TEXT_MAX_CHARS = 4000
|
_RUN_RECORD_TEXT_MAX_CHARS = 4000
|
||||||
_PROCESSING_RECOVERY_ERROR = "delivery was recovered from interrupted processing"
|
_PROCESSING_RECOVERY_ERROR = "delivery was recovered from interrupted processing"
|
||||||
|
_FILE_LOCK_TIMEOUT_SECONDS = 5
|
||||||
|
|
||||||
|
|
||||||
class TriggerStoreError(RuntimeError):
|
class TriggerStoreError(RuntimeError):
|
||||||
@@ -49,7 +50,10 @@ class LocalTriggerStore:
|
|||||||
self.processing_dir = self.root / "processing"
|
self.processing_dir = self.root / "processing"
|
||||||
self.failed_dir = self.root / "failed"
|
self.failed_dir = self.root / "failed"
|
||||||
self.runs_dir = self.root / "runs"
|
self.runs_dir = self.root / "runs"
|
||||||
self._lock = FileLock(str(self.root / ".lock"))
|
self._lock = FileLock(
|
||||||
|
str(self.root / ".lock"),
|
||||||
|
timeout=_FILE_LOCK_TIMEOUT_SECONDS,
|
||||||
|
)
|
||||||
|
|
||||||
def create(
|
def create(
|
||||||
self,
|
self,
|
||||||
|
|||||||
@@ -3,8 +3,48 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import asyncio
|
import asyncio
|
||||||
|
from collections.abc import Awaitable
|
||||||
|
from typing import TypeVar
|
||||||
|
|
||||||
|
_T = TypeVar("_T")
|
||||||
|
|
||||||
|
|
||||||
def task_is_cancelling() -> bool:
|
def task_is_cancelling() -> bool:
|
||||||
task = asyncio.current_task()
|
task = asyncio.current_task()
|
||||||
return task is not None and task.cancelling() > 0
|
return task is not None and task.cancelling() > 0
|
||||||
|
|
||||||
|
|
||||||
|
async def shield_and_drain(awaitable: Awaitable[_T]) -> _T:
|
||||||
|
"""Delay caller cancellation until an accepted operation has fully settled.
|
||||||
|
|
||||||
|
``asyncio.to_thread`` cannot stop a worker that has already started. Shielding
|
||||||
|
keeps cancellation from detaching that worker, and draining also lets any
|
||||||
|
post-write in-memory settlement in ``awaitable`` finish. Cancellation is still
|
||||||
|
re-raised as soon as the accepted operation is done.
|
||||||
|
"""
|
||||||
|
settlement = asyncio.ensure_future(awaitable)
|
||||||
|
cancellation: asyncio.CancelledError | None = None
|
||||||
|
|
||||||
|
while not settlement.done():
|
||||||
|
try:
|
||||||
|
result = await asyncio.shield(settlement)
|
||||||
|
except asyncio.CancelledError as exc:
|
||||||
|
if cancellation is None:
|
||||||
|
cancellation = exc
|
||||||
|
except BaseException:
|
||||||
|
if cancellation is None:
|
||||||
|
raise
|
||||||
|
break
|
||||||
|
else:
|
||||||
|
if cancellation is not None:
|
||||||
|
raise cancellation
|
||||||
|
return result
|
||||||
|
|
||||||
|
if cancellation is not None:
|
||||||
|
try:
|
||||||
|
settlement.result()
|
||||||
|
except BaseException:
|
||||||
|
# The caller's cancellation wins once settlement has been observed.
|
||||||
|
pass
|
||||||
|
raise cancellation
|
||||||
|
return settlement.result()
|
||||||
|
|||||||
@@ -1,5 +1,7 @@
|
|||||||
"""Utility functions for nanobot."""
|
"""Utility functions for nanobot."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
import base64
|
import base64
|
||||||
import json
|
import json
|
||||||
import os
|
import os
|
||||||
@@ -12,11 +14,14 @@ from contextlib import suppress
|
|||||||
from datetime import datetime
|
from datetime import datetime
|
||||||
from functools import lru_cache
|
from functools import lru_cache
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import Any, TypeVar, cast, overload
|
from typing import TYPE_CHECKING, Any, TypeVar, cast, overload
|
||||||
|
|
||||||
import tiktoken
|
import tiktoken
|
||||||
from loguru import logger
|
from loguru import logger
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from nanobot.providers.base import LLMUsage
|
||||||
|
|
||||||
_TOOLS_TOKEN_CACHE_MAX_ENTRIES = 64
|
_TOOLS_TOKEN_CACHE_MAX_ENTRIES = 64
|
||||||
_TOOLS_TOKEN_CACHE: dict[int, tuple[tuple[int, ...], dict[bool, int]]] = {}
|
_TOOLS_TOKEN_CACHE: dict[int, tuple[tuple[int, ...], dict[bool, int]]] = {}
|
||||||
_T = TypeVar("_T")
|
_T = TypeVar("_T")
|
||||||
@@ -793,7 +798,7 @@ def build_status_content(
|
|||||||
version: str,
|
version: str,
|
||||||
model: str,
|
model: str,
|
||||||
start_time: float,
|
start_time: float,
|
||||||
last_usage: dict[str, int],
|
last_usage: LLMUsage | None,
|
||||||
context_window_tokens: int,
|
context_window_tokens: int,
|
||||||
session_msg_count: int,
|
session_msg_count: int,
|
||||||
context_tokens_estimate: int,
|
context_tokens_estimate: int,
|
||||||
@@ -814,9 +819,9 @@ def build_status_content(
|
|||||||
if uptime_s >= 3600
|
if uptime_s >= 3600
|
||||||
else f"{uptime_s // 60}m {uptime_s % 60}s"
|
else f"{uptime_s // 60}m {uptime_s % 60}s"
|
||||||
)
|
)
|
||||||
last_in = last_usage.get("prompt_tokens", 0)
|
last_in = last_usage.input_tokens if last_usage else 0
|
||||||
last_out = last_usage.get("completion_tokens", 0)
|
last_out = last_usage.output_tokens if last_usage else 0
|
||||||
cached = last_usage.get("cached_tokens", 0)
|
cached = last_usage.cache_read_tokens if last_usage else None
|
||||||
ctx_total = max(context_window_tokens, 0)
|
ctx_total = max(context_window_tokens, 0)
|
||||||
# Budget mirrors Consolidator formula: ctx_window - max_completion - _SAFETY_BUFFER
|
# Budget mirrors Consolidator formula: ctx_window - max_completion - _SAFETY_BUFFER
|
||||||
ctx_budget = max(ctx_total - int(max_completion_tokens) - 1024, 1)
|
ctx_budget = max(ctx_total - int(max_completion_tokens) - 1024, 1)
|
||||||
|
|||||||
+32
-19
@@ -2,6 +2,7 @@
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
import re
|
import re
|
||||||
import uuid
|
import uuid
|
||||||
from collections.abc import Mapping
|
from collections.abc import Mapping
|
||||||
@@ -9,6 +10,7 @@ from typing import TYPE_CHECKING, Any, TypeGuard
|
|||||||
|
|
||||||
from nanobot.session.manager import SessionManager
|
from nanobot.session.manager import SessionManager
|
||||||
from nanobot.session.webui_turns import WEBUI_TITLE_METADATA_KEY, clean_generated_title
|
from nanobot.session.webui_turns import WEBUI_TITLE_METADATA_KEY, clean_generated_title
|
||||||
|
from nanobot.utils.cancellation import shield_and_drain
|
||||||
from nanobot.webui.transcript import (
|
from nanobot.webui.transcript import (
|
||||||
append_fork_marker,
|
append_fork_marker,
|
||||||
delete_webui_transcript,
|
delete_webui_transcript,
|
||||||
@@ -93,24 +95,35 @@ async def handle_webui_fork_chat(
|
|||||||
await channel.send_webui_protocol_error(connection, "session_manager_unavailable")
|
await channel.send_webui_protocol_error(connection, "session_manager_unavailable")
|
||||||
return
|
return
|
||||||
|
|
||||||
try:
|
async def create_and_attach() -> None:
|
||||||
forked = create_webui_chat_fork(
|
try:
|
||||||
session_manager,
|
forked = await asyncio.to_thread(
|
||||||
source_chat_id=source_chat_id,
|
create_webui_chat_fork,
|
||||||
before_user_index=raw_index,
|
session_manager,
|
||||||
title=envelope.get("title") if isinstance(envelope.get("title"), str) else None,
|
source_chat_id=source_chat_id,
|
||||||
)
|
before_user_index=raw_index,
|
||||||
if forked is None:
|
title=(
|
||||||
await channel.send_webui_protocol_error(connection, "invalid fork source or index")
|
envelope.get("title")
|
||||||
|
if isinstance(envelope.get("title"), str)
|
||||||
|
else None
|
||||||
|
),
|
||||||
|
)
|
||||||
|
if forked is None:
|
||||||
|
await channel.send_webui_protocol_error(
|
||||||
|
connection,
|
||||||
|
"invalid fork source or index",
|
||||||
|
)
|
||||||
|
return
|
||||||
|
fork_id, fork_key = forked
|
||||||
|
except Exception as exc:
|
||||||
|
channel.logger.warning("fork_chat failed: {}", exc)
|
||||||
|
await channel.send_webui_protocol_error(connection, "fork_chat_failed")
|
||||||
return
|
return
|
||||||
fork_id, fork_key = forked
|
|
||||||
except Exception as exc:
|
|
||||||
channel.logger.warning("fork_chat failed: {}", exc)
|
|
||||||
await channel.send_webui_protocol_error(connection, "fork_chat_failed")
|
|
||||||
return
|
|
||||||
|
|
||||||
await channel.attach_webui_fork(
|
await channel.attach_webui_fork(
|
||||||
connection,
|
connection,
|
||||||
fork_id=fork_id,
|
fork_id=fork_id,
|
||||||
fork_key=fork_key,
|
fork_key=fork_key,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
await shield_and_drain(create_and_attach())
|
||||||
|
|||||||
@@ -69,6 +69,7 @@ def build_gateway_services(
|
|||||||
mcp_runtime_status: Callable[[], Mapping[str, str]] | None = None,
|
mcp_runtime_status: Callable[[], Mapping[str, str]] | None = None,
|
||||||
mcp_reload: Callable[[], Awaitable[dict[str, Any]]] | None = None,
|
mcp_reload: Callable[[], Awaitable[dict[str, Any]]] | None = None,
|
||||||
skill_state_action: Callable[[set[str]], None] | None = None,
|
skill_state_action: Callable[[set[str]], None] | None = None,
|
||||||
|
recovery_action: Callable[[str, dict[str, Any]], Awaitable[dict[str, Any]]] | None = None,
|
||||||
logger: Any = default_logger,
|
logger: Any = default_logger,
|
||||||
) -> GatewayServices:
|
) -> GatewayServices:
|
||||||
settings = WebUISettingsServices.create(
|
settings = WebUISettingsServices.create(
|
||||||
@@ -131,6 +132,7 @@ def build_gateway_services(
|
|||||||
mcp_runtime_status=mcp_runtime_status,
|
mcp_runtime_status=mcp_runtime_status,
|
||||||
mcp_reload=mcp_reload,
|
mcp_reload=mcp_reload,
|
||||||
skill_state_action=skill_state_action,
|
skill_state_action=skill_state_action,
|
||||||
|
recovery_action=recovery_action,
|
||||||
log=logger,
|
log=logger,
|
||||||
)
|
)
|
||||||
return GatewayServices(
|
return GatewayServices(
|
||||||
|
|||||||
@@ -63,9 +63,6 @@ class GatewayTokenStore:
|
|||||||
self.api_tokens[token_value] = expiry
|
self.api_tokens[token_value] = expiry
|
||||||
return token_value
|
return token_value
|
||||||
|
|
||||||
def take_issued_token_if_valid(self, token_value: str | None) -> bool:
|
|
||||||
return self.take_issued_token_audience(token_value) is not None
|
|
||||||
|
|
||||||
def take_issued_token_audience(
|
def take_issued_token_audience(
|
||||||
self,
|
self,
|
||||||
token_value: str | None,
|
token_value: str | None,
|
||||||
|
|||||||
@@ -16,7 +16,7 @@ from nanobot.agent.tools.mcp import MCPConnection, connect_mcp_servers
|
|||||||
from nanobot.agent.tools.mcp_oauth import MCP_OAUTH_CALLBACK_PATH, MCPOAuthHandlers
|
from nanobot.agent.tools.mcp_oauth import MCP_OAUTH_CALLBACK_PATH, MCPOAuthHandlers
|
||||||
from nanobot.agent.tools.registry import ToolRegistry
|
from nanobot.agent.tools.registry import ToolRegistry
|
||||||
from nanobot.config.schema import MCPServerConfig
|
from nanobot.config.schema import MCPServerConfig
|
||||||
from nanobot.security.network import validate_url_target
|
from nanobot.security.network import async_validate_url_target
|
||||||
from nanobot.webui.http_utils import is_loopback_host
|
from nanobot.webui.http_utils import is_loopback_host
|
||||||
|
|
||||||
McpReload = Callable[[], Awaitable[dict[str, Any]]]
|
McpReload = Callable[[], Awaitable[dict[str, Any]]]
|
||||||
@@ -259,7 +259,7 @@ class McpOAuthManager:
|
|||||||
):
|
):
|
||||||
flow.error = "The MCP server returned an unsafe authorization URL."
|
flow.error = "The MCP server returned an unsafe authorization URL."
|
||||||
raise McpOAuthError(flow.error)
|
raise McpOAuthError(flow.error)
|
||||||
ok, _error = validate_url_target(authorization_url)
|
ok, _error = await async_validate_url_target(authorization_url)
|
||||||
if not ok:
|
if not ok:
|
||||||
flow.error = "The MCP server returned an unsafe authorization URL."
|
flow.error = "The MCP server returned an unsafe authorization URL."
|
||||||
raise McpOAuthError(flow.error)
|
raise McpOAuthError(flow.error)
|
||||||
|
|||||||
@@ -4,6 +4,7 @@ from __future__ import annotations
|
|||||||
|
|
||||||
from typing import Any, cast
|
from typing import Any, cast
|
||||||
|
|
||||||
|
from nanobot.providers.base import LLMUsage
|
||||||
from nanobot.session.manager import Session
|
from nanobot.session.manager import Session
|
||||||
from nanobot.utils.helpers import estimate_message_tokens, truncate_text
|
from nanobot.utils.helpers import estimate_message_tokens, truncate_text
|
||||||
|
|
||||||
@@ -36,18 +37,8 @@ def session_context_payload(session: Session) -> dict[str, Any]:
|
|||||||
summary_tokens = (
|
summary_tokens = (
|
||||||
estimate_message_tokens({"role": "system", "content": summary}) if summary else 0
|
estimate_message_tokens({"role": "system", "content": summary}) if summary else 0
|
||||||
)
|
)
|
||||||
raw_usage = session.metadata.get("_last_usage")
|
stored_usage = LLMUsage.from_dict(session.metadata.get("_last_usage"))
|
||||||
last_usage = (
|
last_usage = stored_usage.to_turn_dict() if stored_usage is not None else None
|
||||||
{
|
|
||||||
key: value
|
|
||||||
for key, value in cast(dict[object, object], raw_usage).items()
|
|
||||||
if isinstance(key, str)
|
|
||||||
and type(value) is int
|
|
||||||
and value >= 0
|
|
||||||
}
|
|
||||||
if isinstance(raw_usage, dict)
|
|
||||||
else None
|
|
||||||
)
|
|
||||||
|
|
||||||
return {
|
return {
|
||||||
"schema_version": 1,
|
"schema_version": 1,
|
||||||
|
|||||||
@@ -31,8 +31,9 @@ from nanobot.session.manager import (
|
|||||||
_metadata_title, # pyright: ignore[reportPrivateUsage]
|
_metadata_title, # pyright: ignore[reportPrivateUsage]
|
||||||
)
|
)
|
||||||
from nanobot.session.model_selection import model_preset_from_metadata
|
from nanobot.session.model_selection import model_preset_from_metadata
|
||||||
|
from nanobot.session.recovery import recovery_state_from_metadata
|
||||||
|
|
||||||
_INDEX_VERSION = 7
|
_INDEX_VERSION = 8
|
||||||
_INDEX_FILENAME = ".webui_session_index.json"
|
_INDEX_FILENAME = ".webui_session_index.json"
|
||||||
_MODEL_PRESET_FIELD = "model_preset"
|
_MODEL_PRESET_FIELD = "model_preset"
|
||||||
_ROW_SOURCE_FIELD = "_source"
|
_ROW_SOURCE_FIELD = "_source"
|
||||||
@@ -245,6 +246,7 @@ def _public_row(sessions_dir: Path, webui_dir: Path, row: dict[str, Any]) -> dic
|
|||||||
"title": row.get("title", ""),
|
"title": row.get("title", ""),
|
||||||
"preview": row.get("preview", ""),
|
"preview": row.get("preview", ""),
|
||||||
_MODEL_PRESET_FIELD: row.get(_MODEL_PRESET_FIELD),
|
_MODEL_PRESET_FIELD: row.get(_MODEL_PRESET_FIELD),
|
||||||
|
"recovery_state": row.get("recovery_state"),
|
||||||
_WORKSPACE_SCOPE_PRESENT_FIELD: row.get(_WORKSPACE_SCOPE_PRESENT_FIELD, False),
|
_WORKSPACE_SCOPE_PRESENT_FIELD: row.get(_WORKSPACE_SCOPE_PRESENT_FIELD, False),
|
||||||
_WORKSPACE_SCOPE_VALUE_FIELD: row.get(_WORKSPACE_SCOPE_VALUE_FIELD),
|
_WORKSPACE_SCOPE_VALUE_FIELD: row.get(_WORKSPACE_SCOPE_VALUE_FIELD),
|
||||||
"path": str(path),
|
"path": str(path),
|
||||||
@@ -485,6 +487,7 @@ def _indexed_row_for_session(session: Session, path: Path, webui_dir: Path) -> d
|
|||||||
"title": _metadata_title(session.metadata),
|
"title": _metadata_title(session.metadata),
|
||||||
"preview": _preview_from_messages(session.messages),
|
"preview": _preview_from_messages(session.messages),
|
||||||
_MODEL_PRESET_FIELD: model_preset_from_metadata(session.metadata),
|
_MODEL_PRESET_FIELD: model_preset_from_metadata(session.metadata),
|
||||||
|
"recovery_state": recovery_state_from_metadata(session.metadata),
|
||||||
**_indexed_workspace_scope_fields(session.metadata),
|
**_indexed_workspace_scope_fields(session.metadata),
|
||||||
_ROW_SOURCE_FIELD: _SESSION_SOURCE,
|
_ROW_SOURCE_FIELD: _SESSION_SOURCE,
|
||||||
"file": path.name,
|
"file": path.name,
|
||||||
@@ -601,6 +604,7 @@ def _scan_transcript_row(
|
|||||||
"title": "",
|
"title": "",
|
||||||
"preview": preview or fallback_preview,
|
"preview": preview or fallback_preview,
|
||||||
_MODEL_PRESET_FIELD: None,
|
_MODEL_PRESET_FIELD: None,
|
||||||
|
"recovery_state": None,
|
||||||
**_indexed_workspace_scope_fields({}),
|
**_indexed_workspace_scope_fields({}),
|
||||||
_ROW_SOURCE_FIELD: _TRANSCRIPT_SOURCE,
|
_ROW_SOURCE_FIELD: _TRANSCRIPT_SOURCE,
|
||||||
"file": stem,
|
"file": stem,
|
||||||
@@ -687,6 +691,7 @@ def _scan_session_row(
|
|||||||
"title": _metadata_title(metadata),
|
"title": _metadata_title(metadata),
|
||||||
"preview": preview or fallback_preview,
|
"preview": preview or fallback_preview,
|
||||||
_MODEL_PRESET_FIELD: model_preset_from_metadata(metadata),
|
_MODEL_PRESET_FIELD: model_preset_from_metadata(metadata),
|
||||||
|
"recovery_state": recovery_state_from_metadata(metadata),
|
||||||
**_indexed_workspace_scope_fields(metadata),
|
**_indexed_workspace_scope_fields(metadata),
|
||||||
_ROW_SOURCE_FIELD: _SESSION_SOURCE,
|
_ROW_SOURCE_FIELD: _SESSION_SOURCE,
|
||||||
"file": path.name,
|
"file": path.name,
|
||||||
|
|||||||
@@ -27,6 +27,7 @@ from nanobot.providers.image_generation import (
|
|||||||
)
|
)
|
||||||
from nanobot.providers.registry import find_by_name
|
from nanobot.providers.registry import find_by_name
|
||||||
from nanobot.security.network import is_loopback_host
|
from nanobot.security.network import is_loopback_host
|
||||||
|
from nanobot.utils.cancellation import shield_and_drain
|
||||||
from nanobot.webui.settings_contracts import (
|
from nanobot.webui.settings_contracts import (
|
||||||
QueryParams,
|
QueryParams,
|
||||||
SettingsRequest,
|
SettingsRequest,
|
||||||
@@ -640,7 +641,11 @@ class CapabilitySettingsHandler:
|
|||||||
) -> SettingsRouteResult:
|
) -> SettingsRouteResult:
|
||||||
if action == "api-status":
|
if action == "api-status":
|
||||||
return SettingsRouteResult.success(
|
return SettingsRouteResult.success(
|
||||||
api_service_payload(self.settings, operations.api_runtime())
|
await asyncio.to_thread(
|
||||||
|
api_service_payload,
|
||||||
|
self.settings,
|
||||||
|
operations.api_runtime(),
|
||||||
|
)
|
||||||
)
|
)
|
||||||
if action == "api-start":
|
if action == "api-start":
|
||||||
return await self._start_api(request, operations)
|
return await self._start_api(request, operations)
|
||||||
@@ -673,17 +678,22 @@ class CapabilitySettingsHandler:
|
|||||||
return SettingsRouteResult.failure(404, "unknown settings action")
|
return SettingsRouteResult.failure(404, "unknown settings action")
|
||||||
|
|
||||||
operation, section, apply_image_reload = mutation
|
operation, section, apply_image_reload = mutation
|
||||||
try:
|
|
||||||
payload = self.settings.mutate(operation, request.query)
|
async def mutate_and_apply() -> tuple[dict[str, Any], bool]:
|
||||||
except WebUISettingsError as exc:
|
payload = await self.settings.mutate_async(operation, request.query)
|
||||||
return SettingsRouteResult.failure(exc.status, exc.message)
|
if not apply_image_reload:
|
||||||
if apply_image_reload:
|
return payload, False
|
||||||
payload, image_restart_cleared = await self.apply_image_runtime_change(
|
return await self.apply_image_runtime_change(
|
||||||
payload,
|
payload,
|
||||||
operations.reload_image,
|
operations.reload_image,
|
||||||
)
|
)
|
||||||
else:
|
|
||||||
image_restart_cleared = False
|
try:
|
||||||
|
payload, image_restart_cleared = await shield_and_drain(
|
||||||
|
mutate_and_apply()
|
||||||
|
)
|
||||||
|
except WebUISettingsError as exc:
|
||||||
|
return SettingsRouteResult.failure(exc.status, exc.message)
|
||||||
return SettingsRouteResult.success(
|
return SettingsRouteResult.success(
|
||||||
payload,
|
payload,
|
||||||
decorate_restart=True,
|
decorate_restart=True,
|
||||||
@@ -726,16 +736,17 @@ class CapabilitySettingsHandler:
|
|||||||
400,
|
400,
|
||||||
"API service API key must be a string",
|
"API service API key must be a string",
|
||||||
)
|
)
|
||||||
try:
|
allow_install = await self._allow_feature_package_install(request)
|
||||||
await asyncio.to_thread(
|
|
||||||
self.settings.mutate,
|
async def mutate_and_start() -> Any:
|
||||||
|
await self.settings.mutate_async(
|
||||||
operations.nanobot_features_action,
|
operations.nanobot_features_action,
|
||||||
"enable",
|
"enable",
|
||||||
{"name": ["api"]},
|
{"name": ["api"]},
|
||||||
allow_install=self._allow_feature_package_install(request),
|
allow_install=allow_install,
|
||||||
)
|
)
|
||||||
self.settings.mutate(operations.update_api, request.query)
|
await self.settings.mutate_async(operations.update_api, request.query)
|
||||||
config = self.settings.config.load()
|
config = await self.settings.config.load_async()
|
||||||
runtime = operations.api_runtime()
|
runtime = operations.api_runtime()
|
||||||
options = ApiStartOptions(
|
options = ApiStartOptions(
|
||||||
host=config.api.host,
|
host=config.api.host,
|
||||||
@@ -744,10 +755,13 @@ class CapabilitySettingsHandler:
|
|||||||
config_path=str(self.settings.config.path),
|
config_path=str(self.settings.config.path),
|
||||||
)
|
)
|
||||||
current = runtime.status()
|
current = runtime.status()
|
||||||
result = await asyncio.to_thread(
|
return await asyncio.to_thread(
|
||||||
runtime.restart if current.running else runtime.start_background,
|
runtime.restart if current.running else runtime.start_background,
|
||||||
options,
|
options,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
try:
|
||||||
|
result = await shield_and_drain(mutate_and_start())
|
||||||
if not result.ok:
|
if not result.ok:
|
||||||
return SettingsRouteResult.failure(
|
return SettingsRouteResult.failure(
|
||||||
500,
|
500,
|
||||||
@@ -762,7 +776,8 @@ class CapabilitySettingsHandler:
|
|||||||
self.logger.exception("failed to start managed API service")
|
self.logger.exception("failed to start managed API service")
|
||||||
return SettingsRouteResult.failure(500, str(exc))
|
return SettingsRouteResult.failure(500, str(exc))
|
||||||
return SettingsRouteResult.success(
|
return SettingsRouteResult.success(
|
||||||
api_service_payload(
|
await asyncio.to_thread(
|
||||||
|
api_service_payload,
|
||||||
self.settings,
|
self.settings,
|
||||||
operations.api_runtime(),
|
operations.api_runtime(),
|
||||||
last_action="started",
|
last_action="started",
|
||||||
@@ -775,7 +790,7 @@ class CapabilitySettingsHandler:
|
|||||||
) -> SettingsRouteResult:
|
) -> SettingsRouteResult:
|
||||||
runtime = operations.api_runtime()
|
runtime = operations.api_runtime()
|
||||||
try:
|
try:
|
||||||
result = await asyncio.to_thread(runtime.stop)
|
result = await shield_and_drain(asyncio.to_thread(runtime.stop))
|
||||||
except Exception as exc:
|
except Exception as exc:
|
||||||
self.logger.exception("failed to stop managed API service")
|
self.logger.exception("failed to stop managed API service")
|
||||||
return SettingsRouteResult.failure(500, str(exc))
|
return SettingsRouteResult.failure(500, str(exc))
|
||||||
@@ -785,20 +800,20 @@ class CapabilitySettingsHandler:
|
|||||||
api_runtime_message(result.message),
|
api_runtime_message(result.message),
|
||||||
)
|
)
|
||||||
return SettingsRouteResult.success(
|
return SettingsRouteResult.success(
|
||||||
api_service_payload(
|
await asyncio.to_thread(
|
||||||
|
api_service_payload,
|
||||||
self.settings,
|
self.settings,
|
||||||
operations.api_runtime(),
|
operations.api_runtime(),
|
||||||
last_action="stopped",
|
last_action="stopped",
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|
||||||
def _allow_feature_package_install(self, request: SettingsRequest) -> bool:
|
async def _allow_feature_package_install(self, request: SettingsRequest) -> bool:
|
||||||
if request.local_browser:
|
if request.local_browser:
|
||||||
return True
|
return True
|
||||||
try:
|
try:
|
||||||
return bool(
|
config = await self.settings.config.load_async()
|
||||||
self.settings.config.load().tools.webui_allow_remote_package_install
|
return bool(config.tools.webui_allow_remote_package_install)
|
||||||
)
|
|
||||||
except Exception:
|
except Exception:
|
||||||
self.logger.exception("failed to load remote package install policy")
|
self.logger.exception("failed to load remote package install policy")
|
||||||
return False
|
return False
|
||||||
|
|||||||
@@ -29,6 +29,7 @@ from nanobot.config.schema import Config, FallbackCandidate, ModelPresetConfig,
|
|||||||
from nanobot.providers.image_generation import get_image_gen_provider
|
from nanobot.providers.image_generation import get_image_gen_provider
|
||||||
from nanobot.providers.oauth_guidance import OAUTH_CLI_KIT_MISSING_MESSAGE
|
from nanobot.providers.oauth_guidance import OAUTH_CLI_KIT_MISSING_MESSAGE
|
||||||
from nanobot.providers.registry import PROVIDERS, create_dynamic_spec, find_by_name
|
from nanobot.providers.registry import PROVIDERS, create_dynamic_spec, find_by_name
|
||||||
|
from nanobot.utils.cancellation import shield_and_drain
|
||||||
from nanobot.webui.settings_contracts import (
|
from nanobot.webui.settings_contracts import (
|
||||||
QueryParams,
|
QueryParams,
|
||||||
SettingsRequest,
|
SettingsRequest,
|
||||||
@@ -1651,6 +1652,30 @@ class ModelSettingsHandler:
|
|||||||
if self.settings.refresh_runtime_config is not None:
|
if self.settings.refresh_runtime_config is not None:
|
||||||
self.settings.refresh_runtime_config()
|
self.settings.refresh_runtime_config()
|
||||||
|
|
||||||
|
async def _mutate_and_refresh(
|
||||||
|
self,
|
||||||
|
operation: SettingsOperation,
|
||||||
|
query: QueryParams,
|
||||||
|
**kwargs: Any,
|
||||||
|
) -> dict[str, Any]:
|
||||||
|
payload = await self.settings.mutate_async(operation, query, **kwargs)
|
||||||
|
self._refresh_runtime_config()
|
||||||
|
return payload
|
||||||
|
|
||||||
|
async def _update_provider_and_runtime(
|
||||||
|
self,
|
||||||
|
operation: SettingsOperation,
|
||||||
|
query: QueryParams,
|
||||||
|
apply_image_runtime_change: Callable[
|
||||||
|
[dict[str, Any]],
|
||||||
|
Awaitable[tuple[dict[str, Any], bool]],
|
||||||
|
],
|
||||||
|
) -> tuple[dict[str, Any], bool]:
|
||||||
|
payload = await self.settings.mutate_async(operation, query)
|
||||||
|
payload, image_restart_cleared = await apply_image_runtime_change(payload)
|
||||||
|
self._refresh_runtime_config()
|
||||||
|
return payload, image_restart_cleared
|
||||||
|
|
||||||
async def handle(
|
async def handle(
|
||||||
self,
|
self,
|
||||||
action: str,
|
action: str,
|
||||||
@@ -1659,8 +1684,12 @@ class ModelSettingsHandler:
|
|||||||
) -> SettingsRouteResult:
|
) -> SettingsRouteResult:
|
||||||
try:
|
try:
|
||||||
if action == "agent-update":
|
if action == "agent-update":
|
||||||
payload = self.settings.mutate(operations.update_agent, request.query)
|
payload = await shield_and_drain(
|
||||||
self._refresh_runtime_config()
|
self._mutate_and_refresh(
|
||||||
|
operations.update_agent,
|
||||||
|
request.query,
|
||||||
|
)
|
||||||
|
)
|
||||||
return SettingsRouteResult.success(
|
return SettingsRouteResult.success(
|
||||||
payload,
|
payload,
|
||||||
decorate_restart=True,
|
decorate_restart=True,
|
||||||
@@ -1668,12 +1697,13 @@ class ModelSettingsHandler:
|
|||||||
)
|
)
|
||||||
|
|
||||||
if action == "model-update":
|
if action == "model-update":
|
||||||
payload = self.settings.mutate(
|
payload = await shield_and_drain(
|
||||||
operations.update_model,
|
self._mutate_and_refresh(
|
||||||
request.query,
|
operations.update_model,
|
||||||
rename_model_preset=self.settings.rename_model_preset,
|
request.query,
|
||||||
|
rename_model_preset=self.settings.rename_model_preset,
|
||||||
|
)
|
||||||
)
|
)
|
||||||
self._refresh_runtime_config()
|
|
||||||
return SettingsRouteResult.success(payload, decorate_restart=True)
|
return SettingsRouteResult.success(payload, decorate_restart=True)
|
||||||
|
|
||||||
mutation = {
|
mutation = {
|
||||||
@@ -1684,19 +1714,19 @@ class ModelSettingsHandler:
|
|||||||
"provider-create": operations.create_provider,
|
"provider-create": operations.create_provider,
|
||||||
}.get(action)
|
}.get(action)
|
||||||
if mutation is not None:
|
if mutation is not None:
|
||||||
payload = self.settings.mutate(mutation, request.query)
|
payload = await shield_and_drain(
|
||||||
self._refresh_runtime_config()
|
self._mutate_and_refresh(mutation, request.query)
|
||||||
|
)
|
||||||
return SettingsRouteResult.success(payload, decorate_restart=True)
|
return SettingsRouteResult.success(payload, decorate_restart=True)
|
||||||
|
|
||||||
if action == "provider-update":
|
if action == "provider-update":
|
||||||
payload = self.settings.mutate(
|
payload, image_restart_cleared = await shield_and_drain(
|
||||||
operations.update_provider,
|
self._update_provider_and_runtime(
|
||||||
request.query,
|
operations.update_provider,
|
||||||
|
request.query,
|
||||||
|
operations.apply_image_runtime_change,
|
||||||
|
)
|
||||||
)
|
)
|
||||||
payload, image_restart_cleared = await operations.apply_image_runtime_change(
|
|
||||||
payload
|
|
||||||
)
|
|
||||||
self._refresh_runtime_config()
|
|
||||||
return SettingsRouteResult.success(
|
return SettingsRouteResult.success(
|
||||||
payload,
|
payload,
|
||||||
decorate_restart=True,
|
decorate_restart=True,
|
||||||
@@ -1724,11 +1754,13 @@ class ModelSettingsHandler:
|
|||||||
return SettingsRouteResult.success(payload)
|
return SettingsRouteResult.success(payload)
|
||||||
|
|
||||||
if action == "oauth-login":
|
if action == "oauth-login":
|
||||||
payload = await asyncio.to_thread(
|
payload = await shield_and_drain(
|
||||||
self.settings.read,
|
asyncio.to_thread(
|
||||||
operations.oauth_login,
|
self.settings.read,
|
||||||
request.query,
|
operations.oauth_login,
|
||||||
oauth_flows=self.settings.oauth_flows,
|
request.query,
|
||||||
|
oauth_flows=self.settings.oauth_flows,
|
||||||
|
)
|
||||||
)
|
)
|
||||||
elif action == "oauth-complete":
|
elif action == "oauth-complete":
|
||||||
raw_response = (request.payload or {}).get("authorization_response")
|
raw_response = (request.payload or {}).get("authorization_response")
|
||||||
@@ -1736,19 +1768,23 @@ class ModelSettingsHandler:
|
|||||||
raise WebUISettingsError(
|
raise WebUISettingsError(
|
||||||
"OAuth authorization response must be a string"
|
"OAuth authorization response must be a string"
|
||||||
)
|
)
|
||||||
payload = await asyncio.to_thread(
|
payload = await shield_and_drain(
|
||||||
self.settings.read,
|
asyncio.to_thread(
|
||||||
operations.oauth_complete,
|
self.settings.read,
|
||||||
request.query,
|
operations.oauth_complete,
|
||||||
raw_response or None,
|
request.query,
|
||||||
oauth_flows=self.settings.oauth_flows,
|
raw_response or None,
|
||||||
|
oauth_flows=self.settings.oauth_flows,
|
||||||
|
)
|
||||||
)
|
)
|
||||||
elif action == "oauth-logout":
|
elif action == "oauth-logout":
|
||||||
payload = await asyncio.to_thread(
|
payload = await shield_and_drain(
|
||||||
self.settings.read,
|
asyncio.to_thread(
|
||||||
operations.oauth_logout,
|
self.settings.read,
|
||||||
request.query,
|
operations.oauth_logout,
|
||||||
oauth_flows=self.settings.oauth_flows,
|
request.query,
|
||||||
|
oauth_flows=self.settings.oauth_flows,
|
||||||
|
)
|
||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
return SettingsRouteResult.failure(404, "unknown settings action")
|
return SettingsRouteResult.failure(404, "unknown settings action")
|
||||||
|
|||||||
@@ -4,6 +4,7 @@ from __future__ import annotations
|
|||||||
|
|
||||||
import asyncio
|
import asyncio
|
||||||
import html
|
import html
|
||||||
|
import inspect
|
||||||
import json
|
import json
|
||||||
from collections.abc import Awaitable, Callable, Mapping
|
from collections.abc import Awaitable, Callable, Mapping
|
||||||
from typing import Any, cast
|
from typing import Any, cast
|
||||||
@@ -18,6 +19,7 @@ from nanobot.bus.queue import MessageBus
|
|||||||
from nanobot.channels.registry import load_channel_plugin
|
from nanobot.channels.registry import load_channel_plugin
|
||||||
from nanobot.channels.validation import validate_channel_config
|
from nanobot.channels.validation import validate_channel_config
|
||||||
from nanobot.pairing import approve_code, deny_code, list_pending
|
from nanobot.pairing import approve_code, deny_code, list_pending
|
||||||
|
from nanobot.utils.cancellation import shield_and_drain
|
||||||
from nanobot.webui import settings_capabilities as capability_domain
|
from nanobot.webui import settings_capabilities as capability_domain
|
||||||
from nanobot.webui import settings_contracts as contracts
|
from nanobot.webui import settings_contracts as contracts
|
||||||
from nanobot.webui import settings_models as model_domain
|
from nanobot.webui import settings_models as model_domain
|
||||||
@@ -36,7 +38,6 @@ from nanobot.webui.nanobot_features_api import (
|
|||||||
nanobot_features_payload,
|
nanobot_features_payload,
|
||||||
)
|
)
|
||||||
from nanobot.webui.settings_api import (
|
from nanobot.webui.settings_api import (
|
||||||
WebUISettingsError,
|
|
||||||
complete_oauth_provider,
|
complete_oauth_provider,
|
||||||
create_model_configuration,
|
create_model_configuration,
|
||||||
create_provider_settings,
|
create_provider_settings,
|
||||||
@@ -209,6 +210,16 @@ def _payload_query(payload: dict[str, Any]) -> QueryParams:
|
|||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
async def _call_settings_handler(
|
||||||
|
handler: Callable[[], Response | Awaitable[Response]],
|
||||||
|
) -> Response:
|
||||||
|
"""Keep synchronous handlers off-loop while supporting native async handlers."""
|
||||||
|
result = await asyncio.to_thread(handler)
|
||||||
|
if inspect.isawaitable(result):
|
||||||
|
return await result
|
||||||
|
return result
|
||||||
|
|
||||||
|
|
||||||
class WebUISettingsRouter:
|
class WebUISettingsRouter:
|
||||||
"""Authenticate and dispatch settings requests to transport-neutral domains."""
|
"""Authenticate and dispatch settings requests to transport-neutral domains."""
|
||||||
|
|
||||||
@@ -285,9 +296,9 @@ class WebUISettingsRouter:
|
|||||||
if not self._authorized(request):
|
if not self._authorized(request):
|
||||||
return self._unauthorized()
|
return self._unauthorized()
|
||||||
if route == ("root", "settings"):
|
if route == ("root", "settings"):
|
||||||
return self._handle_settings()
|
return await _call_settings_handler(self._handle_settings)
|
||||||
if route == ("root", "usage"):
|
if route == ("root", "usage"):
|
||||||
return self._handle_settings_usage()
|
return await _call_settings_handler(self._handle_settings_usage)
|
||||||
|
|
||||||
domain, action = route
|
domain, action = route
|
||||||
domain_request = self._domain_request(
|
domain_request = self._domain_request(
|
||||||
@@ -416,19 +427,18 @@ class WebUISettingsRouter:
|
|||||||
)
|
)
|
||||||
return self._json_response(payload)
|
return self._json_response(payload)
|
||||||
|
|
||||||
def _handle_settings(self) -> Response:
|
async def _handle_settings(self) -> Response:
|
||||||
return self._json_response(
|
payload = await self.settings.read_async(
|
||||||
self._with_restart_state(
|
settings_payload,
|
||||||
self.settings.read(
|
surface=self._runtime_surface,
|
||||||
settings_payload,
|
runtime_capability_overrides=self._runtime_capabilities,
|
||||||
surface=self._runtime_surface,
|
|
||||||
runtime_capability_overrides=self._runtime_capabilities,
|
|
||||||
)
|
|
||||||
)
|
|
||||||
)
|
)
|
||||||
|
return self._json_response(self._with_restart_state(payload))
|
||||||
|
|
||||||
def _handle_settings_usage(self) -> Response:
|
async def _handle_settings_usage(self) -> Response:
|
||||||
return self._json_response(self.settings.read(settings_usage_payload))
|
return self._json_response(
|
||||||
|
await self.settings.read_async(settings_usage_payload)
|
||||||
|
)
|
||||||
|
|
||||||
def _model_operations(self) -> model_domain.ModelSettingsOperations:
|
def _model_operations(self) -> model_domain.ModelSettingsOperations:
|
||||||
return model_domain.ModelSettingsOperations(
|
return model_domain.ModelSettingsOperations(
|
||||||
@@ -490,17 +500,6 @@ class WebUISettingsRouter:
|
|||||||
lambda: request_image_generation_reload(self.bus),
|
lambda: request_image_generation_reload(self.bus),
|
||||||
)
|
)
|
||||||
|
|
||||||
async def _apply_image_generation_runtime_change(
|
|
||||||
self,
|
|
||||||
payload: dict[str, Any],
|
|
||||||
) -> dict[str, Any]:
|
|
||||||
updated, restart_cleared = (
|
|
||||||
await self._apply_image_generation_runtime_change_result(payload)
|
|
||||||
)
|
|
||||||
if restart_cleared:
|
|
||||||
self._restart_sections.discard("image")
|
|
||||||
return updated
|
|
||||||
|
|
||||||
async def _reload_mcp_runtime(self) -> dict[str, Any]:
|
async def _reload_mcp_runtime(self) -> dict[str, Any]:
|
||||||
if self._mcp_reload is None:
|
if self._mcp_reload is None:
|
||||||
return {
|
return {
|
||||||
@@ -531,47 +530,9 @@ class WebUISettingsRouter:
|
|||||||
def _parse_mcp_settings_query(self, request: WsRequest) -> QueryParams:
|
def _parse_mcp_settings_query(self, request: WsRequest) -> QueryParams:
|
||||||
return self._query(request)
|
return self._query(request)
|
||||||
|
|
||||||
def _parse_provider_settings_query(self, request: WsRequest) -> QueryParams:
|
|
||||||
return self._query(request)
|
|
||||||
|
|
||||||
def _parse_api_service_settings_query(self, request: WsRequest) -> QueryParams:
|
|
||||||
payload = _mutation_payload(request)
|
|
||||||
if payload is not None:
|
|
||||||
api_key = payload.get("api_key")
|
|
||||||
if api_key is not None and not isinstance(api_key, str):
|
|
||||||
raise WebUISettingsError("API service API key must be a string")
|
|
||||||
return self._query(request)
|
|
||||||
|
|
||||||
def _api_runtime(self) -> ApiRuntime:
|
def _api_runtime(self) -> ApiRuntime:
|
||||||
return ApiRuntime(paths=api_runtime_paths(self.settings.config.path))
|
return ApiRuntime(paths=api_runtime_paths(self.settings.config.path))
|
||||||
|
|
||||||
def _api_service_payload(
|
|
||||||
self,
|
|
||||||
*,
|
|
||||||
last_action: str | None = None,
|
|
||||||
) -> dict[str, Any]:
|
|
||||||
return capability_domain.api_service_payload(
|
|
||||||
self.settings,
|
|
||||||
self._api_runtime(),
|
|
||||||
last_action=last_action,
|
|
||||||
)
|
|
||||||
|
|
||||||
@staticmethod
|
|
||||||
def _masked_secret(value: str) -> str | None:
|
|
||||||
return capability_domain.masked_api_secret(value)
|
|
||||||
|
|
||||||
@staticmethod
|
|
||||||
def _api_runtime_message(message: str) -> str:
|
|
||||||
return capability_domain.api_runtime_message(message)
|
|
||||||
|
|
||||||
def _parse_channel_values(self, request: WsRequest) -> dict[str, Any]:
|
|
||||||
return self._system.parse_channel_values(
|
|
||||||
SettingsRequest(
|
|
||||||
query=self._query(request),
|
|
||||||
payload=_mutation_payload(request),
|
|
||||||
)
|
|
||||||
)
|
|
||||||
|
|
||||||
def _save_channel_config_values(
|
def _save_channel_config_values(
|
||||||
self,
|
self,
|
||||||
name: str,
|
name: str,
|
||||||
@@ -610,18 +571,7 @@ class WebUISettingsRouter:
|
|||||||
allow_install=allow_install,
|
allow_install=allow_install,
|
||||||
)
|
)
|
||||||
|
|
||||||
@staticmethod
|
async def _allow_feature_package_install(
|
||||||
def _feature_runtime_fallback(
|
|
||||||
payload: dict[str, Any],
|
|
||||||
*,
|
|
||||||
message: str,
|
|
||||||
) -> dict[str, Any]:
|
|
||||||
return system_domain.SystemSettingsHandler.feature_runtime_fallback(
|
|
||||||
payload,
|
|
||||||
message=message,
|
|
||||||
)
|
|
||||||
|
|
||||||
def _allow_feature_package_install(
|
|
||||||
self,
|
self,
|
||||||
connection: Any,
|
connection: Any,
|
||||||
request: WsRequest,
|
request: WsRequest,
|
||||||
@@ -631,29 +581,33 @@ class WebUISettingsRouter:
|
|||||||
request,
|
request,
|
||||||
needs_local_browser=True,
|
needs_local_browser=True,
|
||||||
)
|
)
|
||||||
return self._system.allow_feature_package_install(domain_request)
|
return await self._system.allow_feature_package_install(domain_request)
|
||||||
|
|
||||||
async def _handle_mcp_oauth_start(self, request: WsRequest) -> Response:
|
async def _handle_mcp_oauth_start(self, request: WsRequest) -> Response:
|
||||||
if not self._authorized(request):
|
if not self._authorized(request):
|
||||||
return self._unauthorized()
|
return self._unauthorized()
|
||||||
if self._mcp_oauth_redirect_uri is None:
|
redirect_uri_for_request = self._mcp_oauth_redirect_uri
|
||||||
|
if redirect_uri_for_request is None:
|
||||||
return self._error_response(500, "MCP OAuth callback is not configured")
|
return self._error_response(500, "MCP OAuth callback is not configured")
|
||||||
query = self._parse_mcp_settings_query(request)
|
query = self._parse_mcp_settings_query(request)
|
||||||
try:
|
|
||||||
name, cfg = await asyncio.to_thread(
|
async def mutate_and_start() -> dict[str, Any]:
|
||||||
self.settings.mutate,
|
name, cfg = await self.settings.mutate_async(
|
||||||
ensure_mcp_oauth_server,
|
ensure_mcp_oauth_server,
|
||||||
query,
|
query,
|
||||||
)
|
)
|
||||||
redirect_uri = self._mcp_oauth_redirect_uri(request)
|
redirect_uri = redirect_uri_for_request(request)
|
||||||
reset = (_query_first(query, "reset") or "").lower() in {"1", "true", "yes"}
|
reset = (_query_first(query, "reset") or "").lower() in {"1", "true", "yes"}
|
||||||
payload = await self._mcp_oauth.start(
|
return await self._mcp_oauth.start(
|
||||||
name,
|
name,
|
||||||
cfg,
|
cfg,
|
||||||
redirect_uri,
|
redirect_uri,
|
||||||
reload_mcp=self._reload_mcp_runtime,
|
reload_mcp=self._reload_mcp_runtime,
|
||||||
reset_credentials=reset,
|
reset_credentials=reset,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
try:
|
||||||
|
payload = await shield_and_drain(mutate_and_start())
|
||||||
except Exception as exc:
|
except Exception as exc:
|
||||||
return self._mcp_oauth_error_response(exc, action="start")
|
return self._mcp_oauth_error_response(exc, action="start")
|
||||||
return self._json_response(payload)
|
return self._json_response(payload)
|
||||||
|
|||||||
@@ -2,6 +2,7 @@
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
import threading
|
import threading
|
||||||
from collections.abc import Callable
|
from collections.abc import Callable
|
||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
@@ -12,9 +13,11 @@ from filelock import FileLock
|
|||||||
|
|
||||||
from nanobot.config.loader import load_config, save_config
|
from nanobot.config.loader import load_config, save_config
|
||||||
from nanobot.config.schema import Config
|
from nanobot.config.schema import Config
|
||||||
|
from nanobot.utils.cancellation import shield_and_drain
|
||||||
|
|
||||||
_T = TypeVar("_T")
|
_T = TypeVar("_T")
|
||||||
_WEBUI_OAUTH_MAX_FLOWS = 8
|
_WEBUI_OAUTH_MAX_FLOWS = 8
|
||||||
|
_SETTINGS_FILE_LOCK_TIMEOUT_SECONDS = 5
|
||||||
|
|
||||||
|
|
||||||
class WebUISettingsConfig:
|
class WebUISettingsConfig:
|
||||||
@@ -25,13 +28,20 @@ class WebUISettingsConfig:
|
|||||||
self.path.parent.mkdir(parents=True, exist_ok=True)
|
self.path.parent.mkdir(parents=True, exist_ok=True)
|
||||||
self._lock = threading.RLock()
|
self._lock = threading.RLock()
|
||||||
lock_path = self.path.with_suffix(f"{self.path.suffix}.lock")
|
lock_path = self.path.with_suffix(f"{self.path.suffix}.lock")
|
||||||
self._file_lock = FileLock(str(lock_path))
|
self._file_lock = FileLock(
|
||||||
|
str(lock_path),
|
||||||
|
timeout=_SETTINGS_FILE_LOCK_TIMEOUT_SECONDS,
|
||||||
|
)
|
||||||
|
|
||||||
def load(self) -> Config:
|
def load(self) -> Config:
|
||||||
"""Load this gateway's config without consulting the process-global path."""
|
"""Load this gateway's config without consulting the process-global path."""
|
||||||
with self._lock:
|
with self._lock:
|
||||||
return load_config(self.path)
|
return load_config(self.path)
|
||||||
|
|
||||||
|
async def load_async(self) -> Config:
|
||||||
|
"""Load config without running file I/O or lock waits on the event loop."""
|
||||||
|
return await asyncio.to_thread(self.load)
|
||||||
|
|
||||||
def update(self, mutation: Callable[[Config], _T]) -> _T:
|
def update(self, mutation: Callable[[Config], _T]) -> _T:
|
||||||
"""Apply and atomically persist one path-scoped read-modify-write operation."""
|
"""Apply and atomically persist one path-scoped read-modify-write operation."""
|
||||||
with self._lock, self._file_lock:
|
with self._lock, self._file_lock:
|
||||||
@@ -40,11 +50,21 @@ class WebUISettingsConfig:
|
|||||||
save_config(config, self.path)
|
save_config(config, self.path)
|
||||||
return result
|
return result
|
||||||
|
|
||||||
|
async def update_async(self, mutation: Callable[[Config], _T]) -> _T:
|
||||||
|
"""Update config without blocking the event loop."""
|
||||||
|
return await shield_and_drain(asyncio.to_thread(self.update, mutation))
|
||||||
|
|
||||||
def run_serialized(self, operation: Callable[[Path], _T]) -> _T:
|
def run_serialized(self, operation: Callable[[Path], _T]) -> _T:
|
||||||
"""Run a path-aware read-modify-write operation under the config-file lock."""
|
"""Run a path-aware read-modify-write operation under the config-file lock."""
|
||||||
with self._lock, self._file_lock:
|
with self._lock, self._file_lock:
|
||||||
return operation(self.path)
|
return operation(self.path)
|
||||||
|
|
||||||
|
async def run_serialized_async(self, operation: Callable[[Path], _T]) -> _T:
|
||||||
|
"""Run a serialized config operation without blocking the event loop."""
|
||||||
|
return await shield_and_drain(
|
||||||
|
asyncio.to_thread(self.run_serialized, operation)
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
class WebUIOAuthFlowRegistry:
|
class WebUIOAuthFlowRegistry:
|
||||||
"""Bounded, thread-safe OAuth flows owned by one gateway instance."""
|
"""Bounded, thread-safe OAuth flows owned by one gateway instance."""
|
||||||
@@ -146,6 +166,16 @@ class WebUISettingsServices:
|
|||||||
"""Run a settings read against this gateway's explicit config path."""
|
"""Run a settings read against this gateway's explicit config path."""
|
||||||
return operation(*args, config_path=self.config.path, **kwargs)
|
return operation(*args, config_path=self.config.path, **kwargs)
|
||||||
|
|
||||||
|
async def read_async(
|
||||||
|
self,
|
||||||
|
operation: Callable[..., _T],
|
||||||
|
/,
|
||||||
|
*args: Any,
|
||||||
|
**kwargs: Any,
|
||||||
|
) -> _T:
|
||||||
|
"""Run a settings read without blocking the event loop."""
|
||||||
|
return await asyncio.to_thread(self.read, operation, *args, **kwargs)
|
||||||
|
|
||||||
def mutate(
|
def mutate(
|
||||||
self,
|
self,
|
||||||
operation: Callable[..., _T],
|
operation: Callable[..., _T],
|
||||||
@@ -161,3 +191,15 @@ class WebUISettingsServices:
|
|||||||
**kwargs,
|
**kwargs,
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|
||||||
|
async def mutate_async(
|
||||||
|
self,
|
||||||
|
operation: Callable[..., _T],
|
||||||
|
/,
|
||||||
|
*args: Any,
|
||||||
|
**kwargs: Any,
|
||||||
|
) -> _T:
|
||||||
|
"""Mutate settings without blocking the event loop."""
|
||||||
|
return await shield_and_drain(
|
||||||
|
asyncio.to_thread(self.mutate, operation, *args, **kwargs)
|
||||||
|
)
|
||||||
|
|||||||
@@ -20,8 +20,10 @@ from nanobot.channels.contracts import (
|
|||||||
channel_update_instance_config,
|
channel_update_instance_config,
|
||||||
)
|
)
|
||||||
from nanobot.config.schema import Config
|
from nanobot.config.schema import Config
|
||||||
|
from nanobot.llm_usage import llm_usage_payload
|
||||||
from nanobot.optional_features import OptionalFeatureError, with_channel_runtime_status
|
from nanobot.optional_features import OptionalFeatureError, with_channel_runtime_status
|
||||||
from nanobot.security.workspace_access import workspace_sandbox_status
|
from nanobot.security.workspace_access import workspace_sandbox_status
|
||||||
|
from nanobot.utils.cancellation import shield_and_drain
|
||||||
from nanobot.webui.settings_capabilities import network_safety_payload
|
from nanobot.webui.settings_capabilities import network_safety_payload
|
||||||
from nanobot.webui.settings_contracts import (
|
from nanobot.webui.settings_contracts import (
|
||||||
QueryParams,
|
QueryParams,
|
||||||
@@ -31,7 +33,6 @@ from nanobot.webui.settings_contracts import (
|
|||||||
query_first,
|
query_first,
|
||||||
query_first_alias,
|
query_first_alias,
|
||||||
)
|
)
|
||||||
from nanobot.webui.token_usage import token_usage_payload
|
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from nanobot.webui.settings_services import WebUISettingsServices
|
from nanobot.webui.settings_services import WebUISettingsServices
|
||||||
@@ -121,7 +122,7 @@ def system_settings_payload(
|
|||||||
},
|
},
|
||||||
"unified_session": defaults.unified_session,
|
"unified_session": defaults.unified_session,
|
||||||
},
|
},
|
||||||
"usage": token_usage_payload(timezone_name=defaults.timezone),
|
"usage": llm_usage_payload(timezone_name=defaults.timezone),
|
||||||
"advanced": {
|
"advanced": {
|
||||||
"restrict_to_workspace": config.tools.restrict_to_workspace,
|
"restrict_to_workspace": config.tools.restrict_to_workspace,
|
||||||
"workspace_sandbox": sandbox_status.as_dict(),
|
"workspace_sandbox": sandbox_status.as_dict(),
|
||||||
@@ -139,7 +140,7 @@ def system_settings_payload(
|
|||||||
|
|
||||||
def settings_usage_payload(config: Config) -> dict[str, Any]:
|
def settings_usage_payload(config: Config) -> dict[str, Any]:
|
||||||
"""Return the lightweight token usage slice for Overview refreshes."""
|
"""Return the lightweight token usage slice for Overview refreshes."""
|
||||||
return token_usage_payload(timezone_name=config.agents.defaults.timezone)
|
return llm_usage_payload(timezone_name=config.agents.defaults.timezone)
|
||||||
|
|
||||||
|
|
||||||
def update_agent_system_settings(config: Config, query: QueryParams) -> tuple[bool, bool]:
|
def update_agent_system_settings(config: Config, query: QueryParams) -> tuple[bool, bool]:
|
||||||
@@ -446,12 +447,17 @@ class SystemSettingsHandler:
|
|||||||
operations: SystemSettingsOperations,
|
operations: SystemSettingsOperations,
|
||||||
) -> SettingsRouteResult:
|
) -> SettingsRouteResult:
|
||||||
try:
|
try:
|
||||||
payload = await asyncio.to_thread(
|
pending = asyncio.to_thread(
|
||||||
operations.cli_apps_action,
|
operations.cli_apps_action,
|
||||||
action,
|
action,
|
||||||
request.query,
|
request.query,
|
||||||
config_path=self.settings.config.path,
|
config_path=self.settings.config.path,
|
||||||
)
|
)
|
||||||
|
payload = (
|
||||||
|
await shield_and_drain(pending)
|
||||||
|
if action in {"install", "update", "uninstall"}
|
||||||
|
else await pending
|
||||||
|
)
|
||||||
except WebUISettingsError as exc:
|
except WebUISettingsError as exc:
|
||||||
return SettingsRouteResult.failure(exc.status, exc.message)
|
return SettingsRouteResult.failure(exc.status, exc.message)
|
||||||
except Exception as exc:
|
except Exception as exc:
|
||||||
@@ -505,17 +511,29 @@ class SystemSettingsHandler:
|
|||||||
action: str,
|
action: str,
|
||||||
operations: SystemSettingsOperations,
|
operations: SystemSettingsOperations,
|
||||||
) -> SettingsRouteResult:
|
) -> SettingsRouteResult:
|
||||||
try:
|
allow_install = (
|
||||||
|
action != "enable"
|
||||||
|
or await self.allow_feature_package_install(request)
|
||||||
|
)
|
||||||
|
|
||||||
|
async def mutate_and_apply() -> dict[str, Any]:
|
||||||
payload = await asyncio.to_thread(
|
payload = await asyncio.to_thread(
|
||||||
self._nanobot_features_action,
|
self._nanobot_features_action,
|
||||||
action,
|
action,
|
||||||
request.query,
|
request.query,
|
||||||
operations,
|
operations,
|
||||||
allow_install=(
|
allow_install=allow_install,
|
||||||
action != "enable"
|
|
||||||
or self.allow_feature_package_install(request)
|
|
||||||
),
|
|
||||||
)
|
)
|
||||||
|
payload = await self._apply_feature_runtime_change(
|
||||||
|
action,
|
||||||
|
request.query,
|
||||||
|
payload,
|
||||||
|
operations,
|
||||||
|
)
|
||||||
|
return self._with_channel_runtime_status(payload, operations)
|
||||||
|
|
||||||
|
try:
|
||||||
|
payload = await shield_and_drain(mutate_and_apply())
|
||||||
except OptionalFeatureError as exc:
|
except OptionalFeatureError as exc:
|
||||||
return SettingsRouteResult.failure(exc.status, exc.message)
|
return SettingsRouteResult.failure(exc.status, exc.message)
|
||||||
except Exception as exc:
|
except Exception as exc:
|
||||||
@@ -527,13 +545,6 @@ class SystemSettingsHandler:
|
|||||||
action,
|
action,
|
||||||
)
|
)
|
||||||
return SettingsRouteResult.failure(status, message)
|
return SettingsRouteResult.failure(status, message)
|
||||||
payload = await self._apply_feature_runtime_change(
|
|
||||||
action,
|
|
||||||
request.query,
|
|
||||||
payload,
|
|
||||||
operations,
|
|
||||||
)
|
|
||||||
payload = self._with_channel_runtime_status(payload, operations)
|
|
||||||
return SettingsRouteResult.success(
|
return SettingsRouteResult.success(
|
||||||
payload,
|
payload,
|
||||||
decorate_restart=True,
|
decorate_restart=True,
|
||||||
@@ -628,6 +639,15 @@ class SystemSettingsHandler:
|
|||||||
self,
|
self,
|
||||||
request: SettingsRequest,
|
request: SettingsRequest,
|
||||||
operations: SystemSettingsOperations,
|
operations: SystemSettingsOperations,
|
||||||
|
) -> SettingsRouteResult:
|
||||||
|
return await shield_and_drain(
|
||||||
|
self._channel_configure_settled(request, operations)
|
||||||
|
)
|
||||||
|
|
||||||
|
async def _channel_configure_settled(
|
||||||
|
self,
|
||||||
|
request: SettingsRequest,
|
||||||
|
operations: SystemSettingsOperations,
|
||||||
) -> SettingsRouteResult:
|
) -> SettingsRouteResult:
|
||||||
name = (query_first(request.query, "name") or "").strip()
|
name = (query_first(request.query, "name") or "").strip()
|
||||||
instance_id = (
|
instance_id = (
|
||||||
@@ -682,7 +702,7 @@ class SystemSettingsHandler:
|
|||||||
"enable",
|
"enable",
|
||||||
feature_query,
|
feature_query,
|
||||||
operations,
|
operations,
|
||||||
allow_install=self.allow_feature_package_install(request),
|
allow_install=await self.allow_feature_package_install(request),
|
||||||
)
|
)
|
||||||
except OptionalFeatureError as exc:
|
except OptionalFeatureError as exc:
|
||||||
return SettingsRouteResult.failure(
|
return SettingsRouteResult.failure(
|
||||||
@@ -825,6 +845,22 @@ class SystemSettingsHandler:
|
|||||||
channel_name: str,
|
channel_name: str,
|
||||||
payload: dict[str, Any],
|
payload: dict[str, Any],
|
||||||
operations: SystemSettingsOperations,
|
operations: SystemSettingsOperations,
|
||||||
|
) -> dict[str, Any]:
|
||||||
|
return await shield_and_drain(
|
||||||
|
self._settle_channel_connect_success(
|
||||||
|
request,
|
||||||
|
channel_name,
|
||||||
|
payload,
|
||||||
|
operations,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
async def _settle_channel_connect_success(
|
||||||
|
self,
|
||||||
|
request: SettingsRequest,
|
||||||
|
channel_name: str,
|
||||||
|
payload: dict[str, Any],
|
||||||
|
operations: SystemSettingsOperations,
|
||||||
) -> dict[str, Any]:
|
) -> dict[str, Any]:
|
||||||
target = {"name": [channel_name]}
|
target = {"name": [channel_name]}
|
||||||
if payload.get("instance_id"):
|
if payload.get("instance_id"):
|
||||||
@@ -835,11 +871,11 @@ class SystemSettingsHandler:
|
|||||||
"enable",
|
"enable",
|
||||||
target,
|
target,
|
||||||
operations,
|
operations,
|
||||||
allow_install=self.allow_feature_package_install(request),
|
allow_install=await self.allow_feature_package_install(request),
|
||||||
)
|
)
|
||||||
except OptionalFeatureError as exc:
|
except OptionalFeatureError as exc:
|
||||||
features = self.feature_runtime_fallback(
|
features = self.feature_runtime_fallback(
|
||||||
self._nanobot_features_payload(operations),
|
await asyncio.to_thread(self._nanobot_features_payload, operations),
|
||||||
message=(
|
message=(
|
||||||
f"{channel_name} connected, but enabling channel support failed: "
|
f"{channel_name} connected, but enabling channel support failed: "
|
||||||
f"{exc.message}"
|
f"{exc.message}"
|
||||||
@@ -859,13 +895,12 @@ class SystemSettingsHandler:
|
|||||||
)
|
)
|
||||||
return updated
|
return updated
|
||||||
|
|
||||||
def allow_feature_package_install(self, request: SettingsRequest) -> bool:
|
async def allow_feature_package_install(self, request: SettingsRequest) -> bool:
|
||||||
if request.local_browser:
|
if request.local_browser:
|
||||||
return True
|
return True
|
||||||
try:
|
try:
|
||||||
return bool(
|
config = await self.settings.config.load_async()
|
||||||
self.settings.config.load().tools.webui_allow_remote_package_install
|
return bool(config.tools.webui_allow_remote_package_install)
|
||||||
)
|
|
||||||
except Exception:
|
except Exception:
|
||||||
self.logger.exception("failed to load remote package install policy")
|
self.logger.exception("failed to load remote package install policy")
|
||||||
return False
|
return False
|
||||||
@@ -925,13 +960,18 @@ class SystemSettingsHandler:
|
|||||||
operations: SystemSettingsOperations,
|
operations: SystemSettingsOperations,
|
||||||
) -> SettingsRouteResult:
|
) -> SettingsRouteResult:
|
||||||
try:
|
try:
|
||||||
payload = await operations.mcp_presets_action(
|
pending = operations.mcp_presets_action(
|
||||||
action,
|
action,
|
||||||
request.query,
|
request.query,
|
||||||
reload_mcp=operations.reload_mcp,
|
reload_mcp=operations.reload_mcp,
|
||||||
mcp_runtime_status=operations.mcp_runtime_status,
|
mcp_runtime_status=operations.mcp_runtime_status,
|
||||||
config=self.settings.config,
|
config=self.settings.config,
|
||||||
)
|
)
|
||||||
|
payload = (
|
||||||
|
await pending
|
||||||
|
if action is None
|
||||||
|
else await shield_and_drain(pending)
|
||||||
|
)
|
||||||
except Exception as exc:
|
except Exception as exc:
|
||||||
status = getattr(exc, "status", 500)
|
status = getattr(exc, "status", 500)
|
||||||
message = getattr(exc, "message", str(exc))
|
message = getattr(exc, "message", str(exc))
|
||||||
|
|||||||
@@ -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")
|
|
||||||
+32
-42
@@ -1813,11 +1813,16 @@ def replay_transcript_to_ui_messages(
|
|||||||
break
|
break
|
||||||
content = str(candidate.get("content") or "")
|
content = str(candidate.get("content") or "")
|
||||||
has_answer = len(content) > 0
|
has_answer = len(content) > 0
|
||||||
|
if has_answer:
|
||||||
|
break
|
||||||
|
# A completed reasoning field is closed even while its assistant
|
||||||
|
# placeholder remains streaming for the rest of the turn.
|
||||||
if (
|
if (
|
||||||
candidate.get("reasoningStreaming")
|
candidate.get("reasoningStreaming")
|
||||||
or candidate.get("reasoning") is not None
|
or (
|
||||||
or has_answer
|
candidate.get("isStreaming")
|
||||||
or candidate.get("isStreaming")
|
and candidate.get("reasoning") is None
|
||||||
|
)
|
||||||
):
|
):
|
||||||
prev[i] = {
|
prev[i] = {
|
||||||
**candidate,
|
**candidate,
|
||||||
@@ -1827,15 +1832,6 @@ def replay_transcript_to_ui_messages(
|
|||||||
**turn_fields,
|
**turn_fields,
|
||||||
}
|
}
|
||||||
return
|
return
|
||||||
if not has_answer and candidate.get("isStreaming"):
|
|
||||||
prev[i] = {
|
|
||||||
**candidate,
|
|
||||||
"reasoning": chunk,
|
|
||||||
"reasoningStreaming": True,
|
|
||||||
"activitySegmentId": candidate.get("activitySegmentId") or _ensure_activity_segment(),
|
|
||||||
**turn_fields,
|
|
||||||
}
|
|
||||||
return
|
|
||||||
break
|
break
|
||||||
segment = _ensure_activity_segment()
|
segment = _ensure_activity_segment()
|
||||||
prev.append(
|
prev.append(
|
||||||
@@ -1870,7 +1866,14 @@ def replay_transcript_to_ui_messages(
|
|||||||
return None
|
return None
|
||||||
return str(last.get("id"))
|
return str(last.get("id"))
|
||||||
|
|
||||||
def demote_interrupted_assistant(segment: str) -> None:
|
def close_interrupted_assistant() -> None:
|
||||||
|
"""Close an answer segment before tool activity without changing its semantics.
|
||||||
|
|
||||||
|
The wire protocol already marks answer, reasoning, and activity phases.
|
||||||
|
A later tool event does not turn previously emitted answer text into
|
||||||
|
reasoning; preserving ``content`` also keeps live and replay projections
|
||||||
|
equivalent.
|
||||||
|
"""
|
||||||
nonlocal buffer_message_id, buffer_parts
|
nonlocal buffer_message_id, buffer_parts
|
||||||
for i in range(len(messages) - 1, -1, -1):
|
for i in range(len(messages) - 1, -1, -1):
|
||||||
candidate = messages[i]
|
candidate = messages[i]
|
||||||
@@ -1886,19 +1889,7 @@ def replay_transcript_to_ui_messages(
|
|||||||
or candidate.get("media")
|
or candidate.get("media")
|
||||||
):
|
):
|
||||||
continue
|
continue
|
||||||
reasoning_parts = [
|
messages[i] = {**candidate, "isStreaming": False}
|
||||||
part
|
|
||||||
for part in (candidate.get("reasoning"), content)
|
|
||||||
if isinstance(part, str) and part.strip()
|
|
||||||
]
|
|
||||||
messages[i] = {
|
|
||||||
**candidate,
|
|
||||||
"content": "",
|
|
||||||
"reasoning": "\n\n".join(reasoning_parts),
|
|
||||||
"reasoningStreaming": False,
|
|
||||||
"isStreaming": False,
|
|
||||||
"activitySegmentId": candidate.get("activitySegmentId") or segment,
|
|
||||||
}
|
|
||||||
if buffer_message_id == candidate.get("id"):
|
if buffer_message_id == candidate.get("id"):
|
||||||
buffer_message_id = None
|
buffer_message_id = None
|
||||||
buffer_parts = []
|
buffer_parts = []
|
||||||
@@ -1920,19 +1911,6 @@ def replay_transcript_to_ui_messages(
|
|||||||
and not m.get("media")
|
and not m.get("media")
|
||||||
)
|
)
|
||||||
|
|
||||||
def is_tool_trace_at(index: int) -> bool:
|
|
||||||
m = messages[index] if 0 <= index < len(messages) else None
|
|
||||||
return bool(m and m.get("kind") == "trace")
|
|
||||||
|
|
||||||
def prune_reasoning_only() -> None:
|
|
||||||
nonlocal messages
|
|
||||||
kept: list[dict[str, Any]] = []
|
|
||||||
for i, m in enumerate(messages):
|
|
||||||
if is_reasoning_only_placeholder(m) and not is_tool_trace_at(i + 1):
|
|
||||||
continue
|
|
||||||
kept.append(m)
|
|
||||||
messages = kept
|
|
||||||
|
|
||||||
def stamp_completion(
|
def stamp_completion(
|
||||||
*,
|
*,
|
||||||
latency_ms: int | None = None,
|
latency_ms: int | None = None,
|
||||||
@@ -2069,7 +2047,7 @@ def replay_transcript_to_ui_messages(
|
|||||||
if not segment:
|
if not segment:
|
||||||
segment = _new_activity_segment(activate=False)
|
segment = _new_activity_segment(activate=False)
|
||||||
active_file_edit_segment_id = segment
|
active_file_edit_segment_id = segment
|
||||||
demote_interrupted_assistant(segment)
|
close_interrupted_assistant()
|
||||||
strip_covered_file_edit_tool_hints_from_recent_messages(edits, turn_fields)
|
strip_covered_file_edit_tool_hints_from_recent_messages(edits, turn_fields)
|
||||||
target_index = find_file_edit_trace_index(segment, edits)
|
target_index = find_file_edit_trace_index(segment, edits)
|
||||||
if target_index is not None:
|
if target_index is not None:
|
||||||
@@ -2363,7 +2341,7 @@ def replay_transcript_to_ui_messages(
|
|||||||
if not trace_lines:
|
if not trace_lines:
|
||||||
continue
|
continue
|
||||||
segment = _ensure_activity_segment()
|
segment = _ensure_activity_segment()
|
||||||
demote_interrupted_assistant(segment)
|
close_interrupted_assistant()
|
||||||
last = messages[-1] if messages else None
|
last = messages[-1] if messages else None
|
||||||
if (
|
if (
|
||||||
last
|
last
|
||||||
@@ -2447,7 +2425,6 @@ def replay_transcript_to_ui_messages(
|
|||||||
for i, m in enumerate(messages):
|
for i, m in enumerate(messages):
|
||||||
if m.get("isStreaming"):
|
if m.get("isStreaming"):
|
||||||
messages[i] = {**m, "isStreaming": False}
|
messages[i] = {**m, "isStreaming": False}
|
||||||
prune_reasoning_only()
|
|
||||||
lat = rec.get("latency_ms")
|
lat = rec.get("latency_ms")
|
||||||
usage = rec.get("usage")
|
usage = rec.get("usage")
|
||||||
sanitized_usage = (
|
sanitized_usage = (
|
||||||
@@ -2546,6 +2523,19 @@ def has_pending_tool_calls(
|
|||||||
return False
|
return False
|
||||||
|
|
||||||
|
|
||||||
|
def has_unfinished_transcript_tail(session_key: str) -> bool:
|
||||||
|
"""Return whether the active transcript ends in an unfinished turn.
|
||||||
|
|
||||||
|
Recovery runs at gateway startup and only needs the newest, still-active
|
||||||
|
turn. Completed turns are rotated into immutable segment files, so reading
|
||||||
|
every historical segment here would make restart cost grow with the full
|
||||||
|
conversation history.
|
||||||
|
"""
|
||||||
|
return has_pending_tool_calls(
|
||||||
|
_read_transcript_file(webui_transcript_path(session_key))
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
def completed_turn_ids(lines: list[dict[str, Any]]) -> list[str]:
|
def completed_turn_ids(lines: list[dict[str, Any]]) -> list[str]:
|
||||||
"""Return stable identities for turns with an explicitly persisted completion."""
|
"""Return stable identities for turns with an explicitly persisted completion."""
|
||||||
completed: list[str] = []
|
completed: list[str] = []
|
||||||
|
|||||||
@@ -5,6 +5,7 @@ from __future__ import annotations
|
|||||||
import json
|
import json
|
||||||
import os
|
import os
|
||||||
import time
|
import time
|
||||||
|
from collections import OrderedDict
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import TYPE_CHECKING, Any, cast
|
from typing import TYPE_CHECKING, Any, cast
|
||||||
|
|
||||||
@@ -28,6 +29,7 @@ _MAX_STATE_FILE_BYTES = 128 * 1024
|
|||||||
_DEFAULT_ACCESS_MODES = {"default", "full"}
|
_DEFAULT_ACCESS_MODES = {"default", "full"}
|
||||||
_LEGACY_RESTRICTED_DEFAULT_ACCESS_MODE = "restricted"
|
_LEGACY_RESTRICTED_DEFAULT_ACCESS_MODE = "restricted"
|
||||||
_WEBUI_SCOPE_CHANNEL = "websocket"
|
_WEBUI_SCOPE_CHANNEL = "websocket"
|
||||||
|
_MAX_DRAFT_SCOPES = 128
|
||||||
|
|
||||||
|
|
||||||
def _scope_change_is_non_escalating(current: WorkspaceScope, requested: WorkspaceScope) -> bool:
|
def _scope_change_is_non_escalating(current: WorkspaceScope, requested: WorkspaceScope) -> bool:
|
||||||
@@ -186,6 +188,7 @@ class WebUIWorkspaceController:
|
|||||||
self._sessions = session_manager
|
self._sessions = session_manager
|
||||||
self._default_workspace = default_workspace
|
self._default_workspace = default_workspace
|
||||||
self._default_restrict_to_workspace = default_restrict_to_workspace
|
self._default_restrict_to_workspace = default_restrict_to_workspace
|
||||||
|
self._draft_scopes: OrderedDict[str, WorkspaceScope] = OrderedDict()
|
||||||
|
|
||||||
def default_scope(self) -> WorkspaceScope:
|
def default_scope(self) -> WorkspaceScope:
|
||||||
return default_scope_for_webui(
|
return default_scope_for_webui(
|
||||||
@@ -230,6 +233,10 @@ class WebUIWorkspaceController:
|
|||||||
return self._scope_from_metadata_value(raw_scope, default_scope=default_scope)
|
return self._scope_from_metadata_value(raw_scope, default_scope=default_scope)
|
||||||
|
|
||||||
def scope_for_session_key(self, session_key: str) -> WorkspaceScope:
|
def scope_for_session_key(self, session_key: str) -> WorkspaceScope:
|
||||||
|
draft = self._draft_scopes.get(session_key)
|
||||||
|
if draft is not None:
|
||||||
|
self._draft_scopes.move_to_end(session_key)
|
||||||
|
return draft
|
||||||
if self._sessions is None:
|
if self._sessions is None:
|
||||||
return self.default_scope()
|
return self.default_scope()
|
||||||
data = self._sessions.read_session_metadata(session_key)
|
data = self._sessions.read_session_metadata(session_key)
|
||||||
@@ -328,8 +335,24 @@ class WebUIWorkspaceController:
|
|||||||
return scope
|
return scope
|
||||||
|
|
||||||
def persist_scope(self, chat_id: str, scope: WorkspaceScope) -> None:
|
def persist_scope(self, chat_id: str, scope: WorkspaceScope) -> None:
|
||||||
|
session_key = f"websocket:{chat_id}"
|
||||||
if self._sessions is not None:
|
if self._sessions is not None:
|
||||||
session = self._sessions.get_or_create(f"websocket:{chat_id}")
|
session = self._sessions.get_or_create(session_key)
|
||||||
session.metadata["webui"] = True
|
session.metadata["webui"] = True
|
||||||
session.metadata[WORKSPACE_SCOPE_METADATA_KEY] = scope.metadata()
|
session.metadata[WORKSPACE_SCOPE_METADATA_KEY] = scope.metadata()
|
||||||
self._sessions.save(session)
|
self._sessions.save(session)
|
||||||
|
self._draft_scopes.pop(session_key, None)
|
||||||
|
|
||||||
|
def stage_scope(self, chat_id: str, scope: WorkspaceScope) -> None:
|
||||||
|
"""Keep a new chat's scope transient until its first accepted message."""
|
||||||
|
session_key = f"websocket:{chat_id}"
|
||||||
|
if (
|
||||||
|
self._sessions is not None
|
||||||
|
and self._sessions.read_session_metadata(session_key) is not None
|
||||||
|
):
|
||||||
|
self.persist_scope(chat_id, scope)
|
||||||
|
return
|
||||||
|
self._draft_scopes[session_key] = scope
|
||||||
|
self._draft_scopes.move_to_end(session_key)
|
||||||
|
while len(self._draft_scopes) > _MAX_DRAFT_SCOPES:
|
||||||
|
self._draft_scopes.popitem(last=False)
|
||||||
|
|||||||
+152
-24
@@ -29,10 +29,12 @@ from nanobot.cron.session_turns import is_bound_cron_job
|
|||||||
from nanobot.cron.types import CronJob, CronSchedule
|
from nanobot.cron.types import CronJob, CronSchedule
|
||||||
from nanobot.security.workspace_access import WorkspaceScope
|
from nanobot.security.workspace_access import WorkspaceScope
|
||||||
from nanobot.session.manager import SessionManager
|
from nanobot.session.manager import SessionManager
|
||||||
|
from nanobot.session.recovery import RecoveryActionError
|
||||||
from nanobot.session.session_handles import (
|
from nanobot.session.session_handles import (
|
||||||
SessionHandleResolver,
|
SessionHandleResolver,
|
||||||
)
|
)
|
||||||
from nanobot.triggers.local_types import LocalTrigger
|
from nanobot.triggers.local_types import LocalTrigger
|
||||||
|
from nanobot.utils.cancellation import shield_and_drain
|
||||||
from nanobot.webui.file_preview import (
|
from nanobot.webui.file_preview import (
|
||||||
WebUIFilePreviewError,
|
WebUIFilePreviewError,
|
||||||
file_preview_availability_payload,
|
file_preview_availability_payload,
|
||||||
@@ -134,6 +136,18 @@ _WEBUI_MUTATION_PAYLOAD_ATTR = "_nanobot_webui_mutation_payload"
|
|||||||
_WEBUI_MUTATION_REQUEST_ATTR = "_nanobot_webui_mutation_request"
|
_WEBUI_MUTATION_REQUEST_ATTR = "_nanobot_webui_mutation_request"
|
||||||
_NO_STORE_HEADERS = [("Cache-Control", "no-store")]
|
_NO_STORE_HEADERS = [("Cache-Control", "no-store")]
|
||||||
|
|
||||||
|
|
||||||
|
def _slow_http_operation(path: str) -> str:
|
||||||
|
"""Return a route family without logging user-controlled path/query values."""
|
||||||
|
clean_path = path.split("?", 1)[0]
|
||||||
|
if clean_path == "/webui/bootstrap":
|
||||||
|
return clean_path
|
||||||
|
parts = [part for part in clean_path.split("/") if part]
|
||||||
|
if len(parts) >= 2 and parts[0] == "api":
|
||||||
|
return f"/api/{parts[1]}"
|
||||||
|
return "/webui"
|
||||||
|
|
||||||
|
|
||||||
_WEBUI_MUTATION_PATHS = {
|
_WEBUI_MUTATION_PATHS = {
|
||||||
"automation.enable": "/api/webui/automations/enable",
|
"automation.enable": "/api/webui/automations/enable",
|
||||||
"automation.disable": "/api/webui/automations/disable",
|
"automation.disable": "/api/webui/automations/disable",
|
||||||
@@ -145,6 +159,8 @@ _WEBUI_MUTATION_PATHS = {
|
|||||||
"skill.delete": "/api/webui/skills/delete",
|
"skill.delete": "/api/webui/skills/delete",
|
||||||
"sidebar.update": "/api/webui/sidebar-state/update",
|
"sidebar.update": "/api/webui/sidebar-state/update",
|
||||||
"workspace.pick_folder": "/api/workspaces/pick-folder",
|
"workspace.pick_folder": "/api/workspaces/pick-folder",
|
||||||
|
"recovery.continue": "/api/webui/recovery/continue",
|
||||||
|
"recovery.dismiss": "/api/webui/recovery/dismiss",
|
||||||
"settings.agent.update": "/api/settings/update",
|
"settings.agent.update": "/api/settings/update",
|
||||||
"settings.model_configuration.create": "/api/settings/model-configurations/create",
|
"settings.model_configuration.create": "/api/settings/model-configurations/create",
|
||||||
"settings.model_configuration.update": "/api/settings/model-configurations/update",
|
"settings.model_configuration.update": "/api/settings/model-configurations/update",
|
||||||
@@ -323,6 +339,9 @@ class GatewayHTTPHandler:
|
|||||||
mcp_runtime_status: Callable[[], Mapping[str, str]] | None = None,
|
mcp_runtime_status: Callable[[], Mapping[str, str]] | None = None,
|
||||||
mcp_reload: Callable[[], Awaitable[dict[str, Any]]] | None = None,
|
mcp_reload: Callable[[], Awaitable[dict[str, Any]]] | None = None,
|
||||||
skill_state_action: Callable[[set[str]], None] | None = None,
|
skill_state_action: Callable[[set[str]], None] | None = None,
|
||||||
|
recovery_action: (
|
||||||
|
Callable[[str, dict[str, Any]], Awaitable[dict[str, Any]]] | None
|
||||||
|
) = None,
|
||||||
log: Any = logger,
|
log: Any = logger,
|
||||||
) -> None:
|
) -> None:
|
||||||
self.config = config
|
self.config = config
|
||||||
@@ -340,6 +359,7 @@ class GatewayHTTPHandler:
|
|||||||
disabled_skills if disabled_skills is not None else set()
|
disabled_skills if disabled_skills is not None else set()
|
||||||
)
|
)
|
||||||
self.skill_state_action = skill_state_action
|
self.skill_state_action = skill_state_action
|
||||||
|
self.recovery_action = recovery_action
|
||||||
self._skill_install_lock = asyncio.Lock()
|
self._skill_install_lock = asyncio.Lock()
|
||||||
self._folder_picker_lock = asyncio.Lock()
|
self._folder_picker_lock = asyncio.Lock()
|
||||||
self.cron_service = cron_service
|
self.cron_service = cron_service
|
||||||
@@ -413,7 +433,12 @@ class GatewayHTTPHandler:
|
|||||||
response = await self._dispatch_resolved(connection, request, got)
|
response = await self._dispatch_resolved(connection, request, got)
|
||||||
return response
|
return response
|
||||||
finally:
|
finally:
|
||||||
self._log_slow_http(got, response, started)
|
self._log_slow_http(
|
||||||
|
got,
|
||||||
|
response,
|
||||||
|
started,
|
||||||
|
input_chars=len(request.path),
|
||||||
|
)
|
||||||
|
|
||||||
async def dispatch_webui_mutation(
|
async def dispatch_webui_mutation(
|
||||||
self,
|
self,
|
||||||
@@ -454,6 +479,8 @@ class GatewayHTTPHandler:
|
|||||||
return True
|
return True
|
||||||
if re.match(r"^/api/webui/automations/(enable|disable|delete|run|update)$", path):
|
if re.match(r"^/api/webui/automations/(enable|disable|delete|run|update)$", path):
|
||||||
return True
|
return True
|
||||||
|
if path in {"/api/webui/recovery/continue", "/api/webui/recovery/dismiss"}:
|
||||||
|
return True
|
||||||
return path in {
|
return path in {
|
||||||
"/api/webui/skills/install",
|
"/api/webui/skills/install",
|
||||||
"/api/webui/skills/update",
|
"/api/webui/skills/update",
|
||||||
@@ -507,6 +534,11 @@ class GatewayHTTPHandler:
|
|||||||
if response is not None:
|
if response is not None:
|
||||||
return response
|
return response
|
||||||
|
|
||||||
|
# Recovery routes
|
||||||
|
response = await self._dispatch_recovery_route(request, got)
|
||||||
|
if response is not None:
|
||||||
|
return response
|
||||||
|
|
||||||
# Session routes
|
# Session routes
|
||||||
response = await self._dispatch_session_routes(request, got)
|
response = await self._dispatch_session_routes(request, got)
|
||||||
if response is not None:
|
if response is not None:
|
||||||
@@ -542,7 +574,14 @@ class GatewayHTTPHandler:
|
|||||||
|
|
||||||
return connection.respond(404, "Not Found")
|
return connection.respond(404, "Not Found")
|
||||||
|
|
||||||
def _log_slow_http(self, path: str, response: Any | None, started: float) -> None:
|
def _log_slow_http(
|
||||||
|
self,
|
||||||
|
path: str,
|
||||||
|
response: Any | None,
|
||||||
|
started: float,
|
||||||
|
*,
|
||||||
|
input_chars: int,
|
||||||
|
) -> None:
|
||||||
elapsed_ms = int((time.perf_counter() - started) * 1000)
|
elapsed_ms = int((time.perf_counter() - started) * 1000)
|
||||||
if elapsed_ms < _SLOW_WEBUI_HTTP_LOG_MS:
|
if elapsed_ms < _SLOW_WEBUI_HTTP_LOG_MS:
|
||||||
return
|
return
|
||||||
@@ -550,9 +589,10 @@ class GatewayHTTPHandler:
|
|||||||
return
|
return
|
||||||
status = getattr(response, "status_code", None)
|
status = getattr(response, "status_code", None)
|
||||||
self._log.warning(
|
self._log.warning(
|
||||||
"slow webui http route path={} status={} duration_ms={}",
|
"slow webui http operation={} status={} input_chars={} duration_ms={}",
|
||||||
path,
|
_slow_http_operation(path),
|
||||||
status if status is not None else "none",
|
status if status is not None else "none",
|
||||||
|
input_chars,
|
||||||
elapsed_ms,
|
elapsed_ms,
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -680,7 +720,11 @@ class GatewayHTTPHandler:
|
|||||||
async def _dispatch_session_routes(self, request: WsRequest, got: str) -> Response | None:
|
async def _dispatch_session_routes(self, request: WsRequest, got: str) -> Response | None:
|
||||||
m = re.match(r"^/api/sessions/([^/]+)/webui-thread$", got)
|
m = re.match(r"^/api/sessions/([^/]+)/webui-thread$", got)
|
||||||
if m:
|
if m:
|
||||||
return self._handle_webui_thread_get(request, m.group(1))
|
return await asyncio.to_thread(
|
||||||
|
self._handle_webui_thread_get,
|
||||||
|
request,
|
||||||
|
m.group(1),
|
||||||
|
)
|
||||||
|
|
||||||
m = re.match(r"^/api/sessions/([^/]+)/context$", got)
|
m = re.match(r"^/api/sessions/([^/]+)/context$", got)
|
||||||
if m:
|
if m:
|
||||||
@@ -688,18 +732,51 @@ class GatewayHTTPHandler:
|
|||||||
|
|
||||||
m = re.match(r"^/api/sessions/([^/]+)/file-preview$", got)
|
m = re.match(r"^/api/sessions/([^/]+)/file-preview$", got)
|
||||||
if m:
|
if m:
|
||||||
return self._handle_file_preview(request, m.group(1))
|
return await asyncio.to_thread(
|
||||||
|
self._handle_file_preview,
|
||||||
|
request,
|
||||||
|
m.group(1),
|
||||||
|
)
|
||||||
|
|
||||||
m = re.match(r"^/api/sessions/([^/]+)/automations$", got)
|
m = re.match(r"^/api/sessions/([^/]+)/automations$", got)
|
||||||
if m:
|
if m:
|
||||||
return self._handle_session_automations(request, m.group(1))
|
return await self._run_cron_transaction(
|
||||||
|
self._handle_session_automations,
|
||||||
|
request,
|
||||||
|
m.group(1),
|
||||||
|
)
|
||||||
|
|
||||||
m = re.match(r"^/api/sessions/([^/]+)/delete$", got)
|
m = re.match(r"^/api/sessions/([^/]+)/delete$", got)
|
||||||
if m:
|
if m:
|
||||||
return self._handle_session_delete(request, m.group(1))
|
return await self._run_cron_transaction(
|
||||||
|
self._handle_session_delete,
|
||||||
|
request,
|
||||||
|
m.group(1),
|
||||||
|
)
|
||||||
|
|
||||||
return None
|
return None
|
||||||
|
|
||||||
|
async def _dispatch_recovery_route(
|
||||||
|
self,
|
||||||
|
request: WsRequest,
|
||||||
|
path: str,
|
||||||
|
) -> Response | None:
|
||||||
|
match = re.fullmatch(r"/api/webui/recovery/(continue|dismiss)", path)
|
||||||
|
if match is None:
|
||||||
|
return None
|
||||||
|
if not getattr(request, _WEBUI_MUTATION_REQUEST_ATTR, False):
|
||||||
|
return _http_error(405, "WebUI recovery actions require an authenticated WebSocket")
|
||||||
|
if self.recovery_action is None:
|
||||||
|
return _http_error(503, "WebUI recovery is unavailable")
|
||||||
|
payload = _mutation_payload(request)
|
||||||
|
if payload is None:
|
||||||
|
return _http_error(400, "invalid recovery payload")
|
||||||
|
try:
|
||||||
|
result = await self.recovery_action(match.group(1), payload)
|
||||||
|
except RecoveryActionError as exc:
|
||||||
|
return _http_error(exc.status, str(exc))
|
||||||
|
return _http_json_response(result)
|
||||||
|
|
||||||
async def _handle_session_context_get(self, request: WsRequest, key: str) -> Response:
|
async def _handle_session_context_get(self, request: WsRequest, key: str) -> Response:
|
||||||
if not self.check_api_token(request):
|
if not self.check_api_token(request):
|
||||||
return _http_error(401, "Unauthorized")
|
return _http_error(401, "Unauthorized")
|
||||||
@@ -746,6 +823,10 @@ class GatewayHTTPHandler:
|
|||||||
for k, v in s.items()
|
for k, v in s.items()
|
||||||
if k != "path" and k not in WEBUI_SESSION_INDEX_INTERNAL_FIELDS
|
if k != "path" and k not in WEBUI_SESSION_INDEX_INTERNAL_FIELDS
|
||||||
}
|
}
|
||||||
|
# Keep the additive recovery field absent for ordinary sessions so
|
||||||
|
# older clients and compact list responses stay unchanged.
|
||||||
|
if row.get("recovery_state") is None:
|
||||||
|
row.pop("recovery_state", None)
|
||||||
chat_id = key.split(":", 1)[1]
|
chat_id = key.split(":", 1)[1]
|
||||||
started_at = websocket_turn_wall_started_at(chat_id)
|
started_at = websocket_turn_wall_started_at(chat_id)
|
||||||
if started_at is not None:
|
if started_at is not None:
|
||||||
@@ -918,13 +999,24 @@ class GatewayHTTPHandler:
|
|||||||
|
|
||||||
# -- Automation routes --------------------------------------------------
|
# -- Automation routes --------------------------------------------------
|
||||||
|
|
||||||
|
async def _run_cron_transaction(
|
||||||
|
self,
|
||||||
|
operation: Callable[..., Any],
|
||||||
|
/,
|
||||||
|
*args: Any,
|
||||||
|
**kwargs: Any,
|
||||||
|
) -> Any:
|
||||||
|
if self.cron_service is not None:
|
||||||
|
return await self.cron_service.run_sync(operation, *args, **kwargs)
|
||||||
|
return await shield_and_drain(asyncio.to_thread(operation, *args, **kwargs))
|
||||||
|
|
||||||
async def _dispatch_automation_routes(
|
async def _dispatch_automation_routes(
|
||||||
self,
|
self,
|
||||||
request: WsRequest,
|
request: WsRequest,
|
||||||
got: str,
|
got: str,
|
||||||
) -> Response | None:
|
) -> Response | None:
|
||||||
if got == "/api/webui/automations":
|
if got == "/api/webui/automations":
|
||||||
return self._handle_webui_automations(request)
|
return await self._run_cron_transaction(self._handle_webui_automations, request)
|
||||||
m = re.match(r"^/api/webui/automations/(enable|disable|delete|run|update)$", got)
|
m = re.match(r"^/api/webui/automations/(enable|disable|delete|run|update)$", got)
|
||||||
if m:
|
if m:
|
||||||
return await self._handle_webui_automation_action(request, m.group(1))
|
return await self._handle_webui_automation_action(request, m.group(1))
|
||||||
@@ -990,13 +1082,24 @@ class GatewayHTTPHandler:
|
|||||||
job_id = (_query_first(query, "id") or _query_first(query, "job_id") or "").strip()
|
job_id = (_query_first(query, "id") or _query_first(query, "job_id") or "").strip()
|
||||||
if not job_id:
|
if not job_id:
|
||||||
return _http_error(400, "missing automation id")
|
return _http_error(400, "missing automation id")
|
||||||
trigger = self.local_trigger_store.get(job_id) if self.local_trigger_store else None
|
trigger = (
|
||||||
|
await asyncio.to_thread(self.local_trigger_store.get, job_id)
|
||||||
|
if self.local_trigger_store
|
||||||
|
else None
|
||||||
|
)
|
||||||
if trigger is not None:
|
if trigger is not None:
|
||||||
return self._handle_local_trigger_action(request, action, trigger)
|
return await shield_and_drain(
|
||||||
|
asyncio.to_thread(
|
||||||
|
self._handle_local_trigger_action,
|
||||||
|
request,
|
||||||
|
action,
|
||||||
|
trigger,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
if self.cron_service is None:
|
if self.cron_service is None:
|
||||||
return _http_error(404, "automation not found")
|
return _http_error(404, "automation not found")
|
||||||
job = self.cron_service.get_job(job_id)
|
job = await self.cron_service.run_sync(self.cron_service.get_job, job_id)
|
||||||
if job is None:
|
if job is None:
|
||||||
return _http_error(404, "automation not found")
|
return _http_error(404, "automation not found")
|
||||||
if job.payload.kind == "system_event":
|
if job.payload.kind == "system_event":
|
||||||
@@ -1005,13 +1108,23 @@ class GatewayHTTPHandler:
|
|||||||
return _http_error(409, "automation has no linked chat")
|
return _http_error(409, "automation has no linked chat")
|
||||||
|
|
||||||
if action == "enable":
|
if action == "enable":
|
||||||
if self.cron_service.enable_job(job_id, enabled=True) is None:
|
result = await self.cron_service.run_sync(
|
||||||
|
self.cron_service.enable_job,
|
||||||
|
job_id,
|
||||||
|
enabled=True,
|
||||||
|
)
|
||||||
|
if result is None:
|
||||||
return _http_error(404, "automation not found")
|
return _http_error(404, "automation not found")
|
||||||
elif action == "disable":
|
elif action == "disable":
|
||||||
if self.cron_service.enable_job(job_id, enabled=False) is None:
|
result = await self.cron_service.run_sync(
|
||||||
|
self.cron_service.enable_job,
|
||||||
|
job_id,
|
||||||
|
enabled=False,
|
||||||
|
)
|
||||||
|
if result is None:
|
||||||
return _http_error(404, "automation not found")
|
return _http_error(404, "automation not found")
|
||||||
elif action == "delete":
|
elif action == "delete":
|
||||||
result = self.cron_service.remove_job(job_id)
|
result = await self.cron_service.run_sync(self.cron_service.remove_job, job_id)
|
||||||
if result == "not_found":
|
if result == "not_found":
|
||||||
return _http_error(404, "automation not found")
|
return _http_error(404, "automation not found")
|
||||||
if result == "protected":
|
if result == "protected":
|
||||||
@@ -1029,7 +1142,11 @@ class GatewayHTTPHandler:
|
|||||||
if isinstance(parsed, str):
|
if isinstance(parsed, str):
|
||||||
return _http_error(400, parsed)
|
return _http_error(400, parsed)
|
||||||
try:
|
try:
|
||||||
result = self.cron_service.update_job(job_id, **parsed)
|
result = await self.cron_service.run_sync(
|
||||||
|
self.cron_service.update_job,
|
||||||
|
job_id,
|
||||||
|
**parsed,
|
||||||
|
)
|
||||||
except ValueError as exc:
|
except ValueError as exc:
|
||||||
return _http_error(400, str(exc))
|
return _http_error(400, str(exc))
|
||||||
if result == "not_found":
|
if result == "not_found":
|
||||||
@@ -1039,7 +1156,7 @@ class GatewayHTTPHandler:
|
|||||||
else:
|
else:
|
||||||
return _http_error(404, "unknown automation action")
|
return _http_error(404, "unknown automation action")
|
||||||
|
|
||||||
return self._handle_webui_automations(request)
|
return await self._run_cron_transaction(self._handle_webui_automations, request)
|
||||||
|
|
||||||
def _handle_local_trigger_action(
|
def _handle_local_trigger_action(
|
||||||
self,
|
self,
|
||||||
@@ -1124,9 +1241,17 @@ class GatewayHTTPHandler:
|
|||||||
if got == "/api/webui/skills/install":
|
if got == "/api/webui/skills/install":
|
||||||
return await self._handle_webui_skill_install(connection, request)
|
return await self._handle_webui_skill_install(connection, request)
|
||||||
if got == "/api/webui/skills/update":
|
if got == "/api/webui/skills/update":
|
||||||
return self._handle_webui_skill_update(request)
|
return await shield_and_drain(
|
||||||
|
asyncio.to_thread(self._handle_webui_skill_update, request)
|
||||||
|
)
|
||||||
if got == "/api/webui/skills/delete":
|
if got == "/api/webui/skills/delete":
|
||||||
return self._handle_webui_skill_delete(connection, request)
|
return await shield_and_drain(
|
||||||
|
asyncio.to_thread(
|
||||||
|
self._handle_webui_skill_delete,
|
||||||
|
connection,
|
||||||
|
request,
|
||||||
|
)
|
||||||
|
)
|
||||||
if got == "/api/webui/skills":
|
if got == "/api/webui/skills":
|
||||||
return self._handle_webui_skills(request)
|
return self._handle_webui_skills(request)
|
||||||
m = re.match(r"^/api/webui/skills/([^/]+)$", got)
|
m = re.match(r"^/api/webui/skills/([^/]+)$", got)
|
||||||
@@ -1237,7 +1362,7 @@ class GatewayHTTPHandler:
|
|||||||
) -> Response:
|
) -> Response:
|
||||||
if not self.check_api_token(request):
|
if not self.check_api_token(request):
|
||||||
return _http_error(401, "Unauthorized")
|
return _http_error(401, "Unauthorized")
|
||||||
if not self._allow_webui_package_install(connection, request):
|
if not await self._allow_webui_package_install(connection, request):
|
||||||
return _http_error(403, "remote skill installation is disabled")
|
return _http_error(403, "remote skill installation is disabled")
|
||||||
if self._skill_install_lock.locked():
|
if self._skill_install_lock.locked():
|
||||||
return _http_error(409, "another skill installation is already in progress")
|
return _http_error(409, "another skill installation is already in progress")
|
||||||
@@ -1269,13 +1394,16 @@ class GatewayHTTPHandler:
|
|||||||
"last_action": action,
|
"last_action": action,
|
||||||
})
|
})
|
||||||
|
|
||||||
def _allow_webui_package_install(self, connection: Any, request: WsRequest) -> bool:
|
async def _allow_webui_package_install(
|
||||||
|
self,
|
||||||
|
connection: Any,
|
||||||
|
request: WsRequest,
|
||||||
|
) -> bool:
|
||||||
if _is_local_browser_request(connection, request.headers):
|
if _is_local_browser_request(connection, request.headers):
|
||||||
return True
|
return True
|
||||||
try:
|
try:
|
||||||
return bool(
|
config = await self.settings.config.load_async()
|
||||||
self.settings.config.load().tools.webui_allow_remote_package_install
|
return bool(config.tools.webui_allow_remote_package_install)
|
||||||
)
|
|
||||||
except Exception:
|
except Exception:
|
||||||
self._log.exception("failed to load remote package install policy")
|
self._log.exception("failed to load remote package install policy")
|
||||||
return False
|
return False
|
||||||
|
|||||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user