fix(gateway): keep event loop responsive (NAN-33)

Move blocking filesystem, persistence, subprocess, media, and DNS work off the gateway event loop while preserving existing contracts. Add bounded cancellation and responsiveness regression coverage.
This commit is contained in:
chengyongru
2026-08-25 10:18:13 +08:00
parent 7fb0811fbb
commit be3a42ebac
80 changed files with 5186 additions and 771 deletions
+96 -9
View File
@@ -2,7 +2,9 @@
from __future__ import annotations
from collections.abc import Collection
import asyncio
import inspect
from collections.abc import Awaitable, Collection
from datetime import datetime
from typing import TYPE_CHECKING, Any, Callable, Coroutine
@@ -47,9 +49,27 @@ class AutoCompact:
return idle_seconds >= self._ttl * 60
def _has_unarchived_messages(self, key: str) -> bool:
session = self.sessions.get_or_create(key)
return self._session_has_unarchived_messages(self.sessions.get_or_create(key))
@staticmethod
def _session_has_unarchived_messages(session: Session) -> bool:
return session.last_consolidated < len(session.messages)
def _has_native_async_session_method(self, name: str) -> bool:
"""Check the manager's real class, not mock-generated instance attributes."""
method = inspect.getattr_static(type(self.sessions), name, None)
return inspect.iscoroutinefunction(method)
async def _list_sessions_nonblocking(self) -> list[dict[str, Any]]:
if self._has_native_async_session_method("list_sessions_async"):
return await self.sessions.list_sessions_async()
return await asyncio.to_thread(self.sessions.list_sessions)
async def _get_or_create_nonblocking(self, key: str) -> Session:
if self._has_native_async_session_method("get_or_create_async"):
return await self.sessions.get_or_create_async(key)
return await asyncio.to_thread(self.sessions.get_or_create, key)
@classmethod
def _is_internal_session(cls, key: str) -> bool:
return key.startswith(cls._INTERNAL_SESSION_PREFIXES)
@@ -79,6 +99,31 @@ class AutoCompact:
self._archiving.add(key)
schedule_background(self._archive(key, runtime=runtime))
async def check_expired_async(
self,
schedule_background: Callable[[Coroutine[Any, Any, None]], None],
resolve_runtime: Callable[[Session], Awaitable[LLMRuntime]],
active_session_keys: Collection[str] = (),
) -> None:
"""Schedule idle archival without blocking the event loop."""
now = datetime.now()
active_keys = set(active_session_keys)
for info in await self._list_sessions_nonblocking():
key = info.get("key", "")
if not key or self._is_internal_session(key) or key in self._archiving:
continue
if key in active_keys or not self._is_expired(info.get("updated_at"), now):
continue
session = await self._get_or_create_nonblocking(key)
if not self._session_has_unarchived_messages(session):
continue
try:
runtime = await resolve_runtime(session)
except (KeyError, ValueError):
continue
self._archiving.add(key)
schedule_background(self._archive_async(key, runtime=runtime))
async def _archive(self, key: str, *, runtime: LLMRuntime) -> None:
if self._is_internal_session(key):
self._archiving.discard(key)
@@ -90,18 +135,38 @@ class AutoCompact:
max_suffix=self._RECENT_SUFFIX_MESSAGES,
)
if summary and summary != "(nothing)":
session = self.sessions.get_or_create(key)
stored = session_summary_from_metadata(
session.metadata,
fallback_last_active=session.updated_at,
)
if stored is not None:
self._summaries[key] = stored
self._record_stored_summary(key, self.sessions.get_or_create(key))
except Exception:
logger.exception("Auto-compact: failed for {}", key)
finally:
self._archiving.discard(key)
async def _archive_async(self, key: str, *, runtime: LLMRuntime) -> None:
if self._is_internal_session(key):
self._archiving.discard(key)
return
try:
summary = await self.consolidator.compact_idle_session(
key,
runtime=runtime,
max_suffix=self._RECENT_SUFFIX_MESSAGES,
)
if summary and summary != "(nothing)":
session = await self._get_or_create_nonblocking(key)
self._record_stored_summary(key, session)
except Exception:
logger.exception("Auto-compact: failed for {}", key)
finally:
self._archiving.discard(key)
def _record_stored_summary(self, key: str, session: Session) -> None:
stored = session_summary_from_metadata(
session.metadata,
fallback_last_active=session.updated_at,
)
if stored is not None:
self._summaries[key] = stored
def prepare_session(self, session: Session, key: str) -> tuple[Session, SessionSummary | None]:
if self._is_internal_session(key):
self._archiving.discard(key)
@@ -110,6 +175,28 @@ class AutoCompact:
if key in self._archiving or self._is_expired(session.updated_at):
logger.info("Auto-compact: reloading session {} (archiving={})", key, key in self._archiving)
session = self.sessions.get_or_create(key)
return self._prepared_summary(session, key)
async def prepare_session_async(
self,
session: Session,
key: str,
) -> tuple[Session, SessionSummary | None]:
"""Prepare a session without blocking on a reload."""
if self._is_internal_session(key):
self._archiving.discard(key)
self._summaries.pop(key, None)
return session, None
if key in self._archiving or self._is_expired(session.updated_at):
logger.info("Auto-compact: reloading session {} (archiving={})", key, key in self._archiving)
session = await self._get_or_create_nonblocking(key)
return self._prepared_summary(session, key)
def _prepared_summary(
self,
session: Session,
key: str,
) -> tuple[Session, SessionSummary | None]:
# Hot path: summary from in-memory dict (process hasn't restarted).
entry = self._summaries.pop(key, None)
if entry:
+32 -2
View File
@@ -13,6 +13,19 @@ class AutomationTurnError(RuntimeError):
"""Raised when an automation turn reaches the agent and finishes with an error."""
class AutomationTurnAcceptedCancellation(asyncio.CancelledError):
"""Cancellation raised after an automation turn was accepted for processing.
Callers must not replay the turn: the accepted agent work now has independent
ownership and may continue after the submitting task is cancelled.
"""
def _consume_future_exception(future: asyncio.Future[object]) -> None:
if not future.cancelled():
future.exception()
async def publish_next_deferred_turn(
*,
deferred_queues: dict[str, list[InboundMessage]],
@@ -70,19 +83,36 @@ class AutomationTurnCoordinator:
future: asyncio.Future[OutboundMessage | None] = loop.create_future()
self._waiters[turn_id] = future
self._pending_messages_by_turn_id[turn_id] = msg
accepted = False
try:
if self._is_running():
await self._publish_inbound(msg)
accepted = True
else:
await self._dispatch(msg)
# Direct dispatch is given independent task ownership for the
# same reason as publishing to the inbound queue: once admitted,
# cancelling this submitter must not cancel and then replay the
# already-running agent turn.
dispatch_future: asyncio.Future[object] = asyncio.ensure_future(
self._dispatch(msg)
)
dispatch_future.add_done_callback(_consume_future_exception)
accepted = True
await asyncio.shield(dispatch_future)
try:
return await future
except asyncio.CancelledError:
except asyncio.CancelledError as exc:
if accepted:
raise AutomationTurnAcceptedCancellation(*exc.args) from None
raise
except AutomationTurnError:
raise
except Exception as exc:
raise AutomationTurnError(str(exc) or exc.__class__.__name__) from exc
except asyncio.CancelledError as exc:
if accepted and not isinstance(exc, AutomationTurnAcceptedCancellation):
raise AutomationTurnAcceptedCancellation(*exc.args) from None
raise
finally:
self._waiters.pop(turn_id, None)
self._pending_messages_by_turn_id.pop(turn_id, None)
+155 -32
View File
@@ -67,6 +67,7 @@ from nanobot.security.workspace_access import (
reset_workspace_scope,
)
from nanobot.session import turn_continuation
from nanobot.session.async_compat import call_session_manager
from nanobot.session.automation_turns import automation_history_overrides
from nanobot.session.goal_state import (
goal_state_runtime_lines,
@@ -540,6 +541,33 @@ class AgentLoop:
**extra,
)
async def _get_or_create_session(self, key: str) -> Session:
"""Use native async session loading, with a compatibility fallback."""
return await call_session_manager(
self.sessions,
"get_or_create_async",
self.sessions.get_or_create,
key,
)
async def _save_session(self, session: Session) -> None:
"""Use native async session saving, with a compatibility fallback."""
await call_session_manager(
self.sessions,
"save_async",
self.sessions.save,
session,
)
async def _save_runtime_checkpoint(self, session: Session) -> None:
"""Use native async checkpoint saving, with a compatibility fallback."""
await call_session_manager(
self.sessions,
"save_runtime_checkpoint_async",
self.sessions.save_runtime_checkpoint,
session,
)
def _sync_subagent_runtime_limits(self) -> None:
"""Keep subagent runtime limits aligned with mutable loop settings."""
self.subagents.max_iterations = self.max_iterations
@@ -579,6 +607,30 @@ class AgentLoop:
self.sessions.save(session)
return self.llm_runtime()
async def runtime_for_session_async(
self,
session: Session,
*,
recover_removed: bool = True,
) -> LLMRuntime:
"""Resolve a session runtime without blocking on recovery persistence."""
name = model_preset_from_metadata(session.metadata)
if name is None:
return self.llm_runtime()
try:
return self.runtime_resolver.resolve_preset(name)
except KeyError:
if not recover_removed or name in self.runtime_resolver.model_presets:
raise
logger.warning(
"Session '{}' references removed model preset '{}'; falling back to default",
session.key,
name,
)
session.metadata.pop(SESSION_MODEL_PRESET_METADATA_KEY, None)
await self._save_session(session)
return self.llm_runtime()
def set_session_model_preset(
self,
session_key: str,
@@ -591,6 +643,18 @@ class AgentLoop:
self.sessions.save(session)
return runtime
async def set_session_model_preset_async(
self,
session_key: str,
name: str,
) -> LLMRuntime:
"""Validate and persist one session's preset selection without blocking."""
runtime = self.runtime_resolver.resolve_preset(name)
session = await self._get_or_create_session(session_key)
session.metadata[SESSION_MODEL_PRESET_METADATA_KEY] = runtime.model_preset
await self._save_session(session)
return runtime
def _publish_runtime_selection(
self,
runtime: LLMRuntime,
@@ -703,17 +767,14 @@ class AgentLoop:
session_key=session_key,
)
def _persist_user_message_early(
def _stage_user_message_early(
self,
msg: InboundMessage,
session: Session,
runtime_context_blocks: list[RuntimeContextBlock] | None = None,
**kwargs: Any,
) -> bool:
"""Persist the triggering user message before the turn starts.
Returns True if the message was persisted.
"""
"""Add the triggering user message and recovery markers in memory."""
if not turn_continuation.should_persist_user_message(msg.metadata):
return False
media_paths = [
@@ -742,10 +803,45 @@ class AgentLoop:
followup_id = msg.metadata.get(PENDING_FOLLOWUP_ID_KEY)
if isinstance(followup_id, str) and followup_id:
acknowledge_pending_followups(session, [followup_id])
self.sessions.save(session)
return True
return False
def _persist_user_message_early(
self,
msg: InboundMessage,
session: Session,
runtime_context_blocks: list[RuntimeContextBlock] | None = None,
**kwargs: Any,
) -> bool:
"""Synchronously persist the user message for compatibility callers."""
persisted = self._stage_user_message_early(
msg,
session,
runtime_context_blocks,
**kwargs,
)
if persisted:
self.sessions.save(session)
return persisted
async def _persist_user_message_early_async(
self,
msg: InboundMessage,
session: Session,
runtime_context_blocks: list[RuntimeContextBlock] | None = None,
**kwargs: Any,
) -> bool:
"""Persist the user message without blocking the event loop."""
persisted = self._stage_user_message_early(
msg,
session,
runtime_context_blocks,
**kwargs,
)
if persisted:
await self._save_session(session)
return persisted
def _build_initial_messages(self, ctx: TurnContext) -> list[dict[str, Any]]:
"""Build the initial message list for the LLM turn."""
assert ctx.session is not None
@@ -843,7 +939,7 @@ class AgentLoop:
if tool is None:
content = "Shell execution is disabled in this nanobot configuration."
else:
session = ctx.session or self.sessions.get_or_create(ctx.key)
session = ctx.session or await AgentLoop._get_or_create_session(self, ctx.key)
scope = self.workspace_scopes.for_turn(
channel=ctx.msg.channel,
message_metadata=metadata,
@@ -1001,7 +1097,7 @@ class AgentLoop:
public_payload[self._PROVIDER_STATE_CHECKPOINT_VERSION_KEY] = (
self._PROVIDER_STATE_CHECKPOINT_VERSION
)
self._set_runtime_checkpoint(session, public_payload)
await self._set_runtime_checkpoint_async(session, public_payload)
async def _drain_pending(*, limit: int = _MAX_INJECTIONS_PER_TURN) -> list[dict[str, Any]]:
"""Drain follow-up messages from the pending queue.
@@ -1239,18 +1335,33 @@ class AgentLoop:
logger.error("LLM returned error: {}", (result.final_content or "")[:200])
return result.final_content, result.tools_used, result.messages, result.stop_reason, result.had_injections
def _check_expired_sessions_if_due(self) -> None:
"""Scan idle sessions no more often than the configured interval."""
def _idle_compact_scan_due(self) -> bool:
now = time.monotonic()
if now < self._next_idle_compact_check_at:
return
return False
self._next_idle_compact_check_at = now + self._idle_compact_check_interval_s
return True
def _check_expired_sessions_if_due(self) -> None:
"""Synchronously scan idle sessions for compatibility with direct callers."""
if not self._idle_compact_scan_due():
return
self.auto_compact.check_expired(
self.schedule_background,
self.runtime_for_session,
active_session_keys=self._pending_queues.keys(),
)
async def _check_expired_sessions_if_due_async(self) -> None:
"""Scan idle sessions without blocking the event loop."""
if not self._idle_compact_scan_due():
return
await self.auto_compact.check_expired_async(
self.schedule_background,
self.runtime_for_session_async,
active_session_keys=self._pending_queues.keys(),
)
async def run(self) -> None:
"""Run the agent loop, dispatching messages as tasks to stay responsive to /stop."""
self._running = True
@@ -1261,7 +1372,7 @@ class AgentLoop:
try:
msg = await asyncio.wait_for(self.bus.consume_inbound(), timeout=1.0)
except asyncio.TimeoutError:
self._check_expired_sessions_if_due()
await self._check_expired_sessions_if_due_async()
continue
except asyncio.CancelledError:
# Preserve real task cancellation so shutdown can complete cleanly.
@@ -1339,7 +1450,7 @@ class AgentLoop:
)
continue
pending_msg = routed_msg
session = self.sessions.get_or_create(effective_key)
session = await self._get_or_create_session(effective_key)
followup_id = record_pending_followup(session, pending_msg)
if followup_id is not None:
pending_msg = dataclasses.replace(
@@ -1349,7 +1460,7 @@ class AgentLoop:
PENDING_FOLLOWUP_ID_KEY: followup_id,
},
)
self.sessions.save(session)
await self._save_session(session)
try:
self._pending_queues[effective_key].put_nowait(pending_msg)
except asyncio.QueueFull:
@@ -1460,10 +1571,10 @@ class AgentLoop:
raise
try:
key = self._effective_session_key(msg)
session = self.sessions.get_or_create(key)
session = await self._get_or_create_session(key)
if self._restore_runtime_checkpoint(session):
self._clear_pending_user_turn(session)
self.sessions.save(session)
await self._save_session(session)
logger.info(
"Restored partial context for cancelled session {}",
key,
@@ -1788,7 +1899,7 @@ class AgentLoop:
if ctx.session is None:
raise RuntimeError("required session is not active")
else:
ctx.session = self.sessions.get_or_create(ctx.session_key)
ctx.session = await self._get_or_create_session(ctx.session_key)
session = ctx.session
ctx.ephemeral = ctx.ephemeral or not session.policy.persist
tools = ctx.tools or self.tools
@@ -1819,16 +1930,16 @@ class AgentLoop:
self.workspace_scopes.persist_message_scope(session, msg)
if self._restore_runtime_checkpoint(session):
self.sessions.save(session)
await self._save_session(session)
if (
RECOVERY_INBOUND_METADATA_KEY not in msg.metadata
and restore_pending_interruption(session)
):
self.sessions.save(session)
await self._save_session(session)
async def _compact_session(self, ctx: TurnContext) -> None:
session = ctx.require_session()
ctx.session, pending = self.auto_compact.prepare_session(
ctx.session, pending = await self.auto_compact.prepare_session_async(
session,
ctx.session_key,
)
@@ -1865,14 +1976,14 @@ class AgentLoop:
# them out of LLM context. /new is excluded because it
# intentionally clears the session.
if cmd_ctx.raw.lower() != "/new":
ctx.input_persisted_early = self._persist_user_message_early(
ctx.input_persisted_early = await self._persist_user_message_early_async(
ctx.msg, session, _command=True
)
session.add_message(
"assistant", result.content, _command=True
)
self._clear_pending_user_turn(session)
self.sessions.save(session)
await self._save_session(session)
if not ctx.ephemeral:
await self.runtime_event_publisher.session_turn_persisted(
ctx.msg,
@@ -1887,7 +1998,7 @@ class AgentLoop:
session = ctx.require_session()
runtime = ctx.runtime
if runtime is None:
runtime = self.runtime_for_session(session)
runtime = await self.runtime_for_session_async(session)
ctx.runtime = runtime
if ctx.session_key.startswith("dream:"):
logger.info(
@@ -1931,7 +2042,7 @@ class AgentLoop:
# provider compatibility or prompt assembly work. A compatible
# staged state replaces this in a second atomic save below.
session.provider_state = None
self.sessions.save(session)
await self._save_session(session)
ctx.input_persisted_early = True
await ctx.delivery.runtime_admitted(runtime)
@@ -1985,7 +2096,7 @@ class AgentLoop:
elif stored_state is not None:
session.provider_state = None
if ctx.kind is TurnKind.USER:
ctx.input_persisted_early = self._persist_user_message_early(
ctx.input_persisted_early = await self._persist_user_message_early_async(
ctx.msg,
session,
runtime_context_blocks=ctx.runtime_context_blocks,
@@ -1995,7 +2106,7 @@ class AgentLoop:
elif subagent_followup_persisted and staged_provider_state:
# Upgrade the replay-safe baseline to the resumable state before
# prompt assembly and the first model checkpoint.
self.sessions.save(session)
await self._save_session(session)
ctx.initial_messages = self._build_initial_messages(ctx)
if ctx.on_progress is None:
@@ -2071,6 +2182,9 @@ class AgentLoop:
turn_latency_ms=ctx.turn_latency_ms,
)
ctx.delivery.record_latency(ctx.turn_latency_ms)
self._clear_pending_user_turn(session)
self._clear_runtime_checkpoint(session)
await self._save_session(session)
if not ctx.ephemeral:
self.schedule_background(
self.consolidator.maybe_consolidate_by_tokens(
@@ -2078,10 +2192,6 @@ class AgentLoop:
runtime=runtime,
)
)
self._clear_pending_user_turn(session)
self._clear_runtime_checkpoint(session)
self.sessions.save(session)
if not ctx.ephemeral:
await self.runtime_event_publisher.session_turn_persisted(
ctx.msg,
ctx.session_key,
@@ -2294,11 +2404,24 @@ class AgentLoop:
)
return True
def _set_runtime_checkpoint(self, session: Session, payload: dict[str, Any]) -> None:
"""Persist the latest in-flight turn state into session metadata."""
def _set_runtime_checkpoint(
self,
session: Session,
payload: dict[str, Any],
) -> None:
"""Synchronously persist a checkpoint for compatibility callers."""
session.metadata[self._RUNTIME_CHECKPOINT_KEY] = payload
self.sessions.save_runtime_checkpoint(session)
async def _set_runtime_checkpoint_async(
self,
session: Session,
payload: dict[str, Any],
) -> None:
"""Persist the latest in-flight turn state without blocking the event loop."""
session.metadata[self._RUNTIME_CHECKPOINT_KEY] = payload
await self._save_runtime_checkpoint(session)
def _mark_pending_user_turn(self, session: Session) -> None:
session.metadata[self._PENDING_USER_TURN_KEY] = True
+26 -9
View File
@@ -22,6 +22,7 @@ from loguru import logger
from nanobot.llm_usage.context import llm_usage_source
from nanobot.runtime_context import public_history_messages
from nanobot.session.async_compat import call_session_manager
from nanobot.session.manager import (
MIN_COMPACTED_REPLAY_MESSAGES,
Session,
@@ -824,6 +825,22 @@ class Consolidator:
weakref.WeakValueDictionary()
)
async def _get_or_create_session(self, key: str) -> Session:
return await call_session_manager(
self.sessions,
"get_or_create_async",
self.sessions.get_or_create,
key,
)
async def _save_session(self, session: Session) -> None:
await call_session_manager(
self.sessions,
"save_async",
self.sessions.save,
session,
)
def get_lock(self, session_key: str) -> asyncio.Lock:
"""Return the shared consolidation lock for one session."""
return self._locks.setdefault(session_key, asyncio.Lock())
@@ -859,13 +876,13 @@ class Consolidator:
return []
return session.get_history()
def _persist_last_summary(self, session: Session, summary: str | None) -> None:
async def _persist_last_summary(self, session: Session, summary: str | None) -> None:
if summary and summary != "(nothing)":
session.metadata["_last_summary"] = {
"text": summary,
"last_active": session.updated_at.isoformat(),
}
self.sessions.save(session)
await self._save_session(session)
def estimate_session_prompt_tokens(
self,
@@ -1057,7 +1074,7 @@ class Consolidator:
lock = self.get_lock(session.key)
async with lock:
# Refresh session reference: AutoCompact may have replaced it.
fresh = self.sessions.get_or_create(session.key)
fresh = await self._get_or_create_session(session.key)
if fresh is not session:
session = fresh
if not session.messages:
@@ -1071,7 +1088,7 @@ class Consolidator:
runtime=runtime,
)
if estimated <= 0:
self._persist_last_summary(session, last_summary)
await self._persist_last_summary(session, last_summary)
return
if estimated < budget:
unconsolidated_count = len(session.messages) - session.last_consolidated
@@ -1083,7 +1100,7 @@ class Consolidator:
source,
unconsolidated_count,
)
self._persist_last_summary(session, last_summary)
await self._persist_last_summary(session, last_summary)
return
for round_num in range(self._MAX_CONSOLIDATION_ROUNDS):
@@ -1127,7 +1144,7 @@ class Consolidator:
last_summary = summary
session.last_consolidated = end_idx
session.provider_state = None
self.sessions.save(session)
await self._save_session(session)
if not summary:
# LLM is degraded — stop hammering it this call;
# the next invocation can retry a fresh chunk.
@@ -1143,7 +1160,7 @@ class Consolidator:
# Persist the last summary to session metadata so it can be injected
# into the runtime context on the next prepare_session() call, aligning
# the summary injection strategy with AutoCompact._archive().
self._persist_last_summary(session, last_summary)
await self._persist_last_summary(session, last_summary)
async def compact_idle_session(
self,
@@ -1168,7 +1185,7 @@ class Consolidator:
lock = self.get_lock(session_key)
async with lock:
self.sessions.invalidate(session_key)
session = self.sessions.get_or_create(session_key)
session = await self._get_or_create_session(session_key)
archive_start = session.last_consolidated
messages_to_archive = list(session.messages[archive_start:])
@@ -1193,7 +1210,7 @@ class Consolidator:
# through the captured batch so new messages remain eligible next time.
session.last_consolidated = archive_end
session.provider_state = None
self.sessions.save(session)
await self._save_session(session)
visible = session.get_history(
max_messages=MIN_COMPACTED_REPLAY_MESSAGES,
+29 -1
View File
@@ -6,7 +6,7 @@ import asyncio
import inspect
import os
import time
from collections.abc import Awaitable, Callable, Iterable
from collections.abc import Awaitable, Callable, Iterable, Sized
from copy import deepcopy
from dataclasses import dataclass, field
from pathlib import Path
@@ -83,6 +83,22 @@ _MAX_EMPTY_RETRIES = 2
_MAX_LENGTH_RECOVERIES = 3
_MAX_INJECTIONS_PER_TURN = 3
_MAX_INJECTION_CYCLES = 5
_SLOW_TOOL_LOG_MS = 1_000
def _tool_input_scale(params: object) -> tuple[int, int]:
"""Return bounded structural counts without logging argument content."""
if not isinstance(params, dict):
return 0, len(params) if isinstance(params, str | bytes) else 0
params_dict = cast(dict[object, object], params)
items = len(params_dict)
chars = 0
for value in params_dict.values():
if isinstance(value, str | bytes):
chars += len(value)
elif isinstance(value, list | tuple | set | dict):
items += len(cast(Sized, value))
return items, chars
def _restore_outer_whitespace(content: str, original: str | None) -> str:
@@ -1521,6 +1537,7 @@ class AgentRunner:
RuntimeError(prep_error) if spec.fail_on_tool_error else None
)
await hook.before_execute_tool(context, tool_call, tool, params)
tool_started_at = time.perf_counter()
try:
if tool is not None:
result = await tool.execute(**params)
@@ -1549,6 +1566,17 @@ class AgentRunner:
if spec.fail_on_tool_error:
return payload, event, exc
return payload, event, None
finally:
duration_ms = int((time.perf_counter() - tool_started_at) * 1000)
if duration_ms >= _SLOW_TOOL_LOG_MS:
input_items, input_chars = _tool_input_scale(params)
logger.warning(
"slow tool operation={} input_items={} input_chars={} duration_ms={}",
tool_call.name,
input_items,
input_chars,
duration_ms,
)
if is_tool_error_result(result):
await hook.on_execute_tool_error(context, tool_call, tool, params, result)
+22 -8
View File
@@ -4,7 +4,10 @@
from __future__ import annotations
import asyncio
import inspect
from pathlib import Path
from typing import TypedDict
from pydantic import Field
@@ -24,6 +27,14 @@ from nanobot.runtime_context import RuntimeContextBlock, wrap_runtime_context_li
from nanobot.security.workspace_access import current_tool_workspace
class _CliAppRunKwargs(TypedDict):
args: list[str]
json_output: bool
working_dir: str | None
timeout: int | None
restrict_to_workspace: bool
class CliAppsToolConfig(Base):
"""CLI Apps tool configuration."""
@@ -147,14 +158,17 @@ class CliAppsTool(Tool):
)
workspace = access.project_path or self.workspace
manager = CliAppManager(workspace=workspace, runtime=self.runtime)
run_kwargs: _CliAppRunKwargs = {
"args": args or [],
"json_output": bool(json),
"working_dir": working_dir,
"timeout": timeout,
"restrict_to_workspace": access.restrict_to_workspace,
}
try:
return manager.run(
name,
args=args or [],
json_output=bool(json),
working_dir=working_dir,
timeout=timeout,
restrict_to_workspace=access.restrict_to_workspace,
)
run_async = inspect.getattr_static(type(manager), "run_async", None)
if inspect.iscoroutinefunction(run_async):
return await manager.run_async(name, **run_kwargs)
return await asyncio.to_thread(manager.run, name, **run_kwargs)
except CliAppError as exc:
return ToolResult.error(f"Error: {exc.message}")
+29 -2
View File
@@ -143,14 +143,41 @@ class CronTool(Tool):
tz: str | None = None,
at: str | None = None,
job_id: str | None = None,
) -> str:
if action == "add" and self._in_cron_context.get():
return ToolResult.error(
"Error: cannot schedule new jobs from within a cron job execution"
)
return await self._cron.run_sync(
self._execute_sync,
action,
name,
message,
every_seconds,
cron_expr,
tz,
at,
job_id,
)
def _execute_sync(
self,
action: str,
name: str | None,
message: str,
every_seconds: int | None,
cron_expr: str | None,
tz: str | None,
at: str | None,
job_id: str | None,
) -> str:
if action == "add":
if self._in_cron_context.get():
return ToolResult.error("Error: cannot schedule new jobs from within a cron job execution")
return self._add_job(name, message, every_seconds, cron_expr, tz, at)
elif action == "list":
if action == "list":
return self._list_jobs()
elif action == "remove":
if action == "remove":
return self._remove_job(job_id)
return f"Unknown action: {action}"
+92 -21
View File
@@ -2,9 +2,11 @@
# pyright: reportPrivateUsage=false, reportUnusedFunction=false
import asyncio
import difflib
import mimetypes
import os
import threading
from dataclasses import dataclass
from pathlib import Path
from typing import Any
@@ -21,6 +23,7 @@ from nanobot.agent.tools.schema import (
)
from nanobot.config_base import Base
from nanobot.security.workspace_access import current_tool_workspace
from nanobot.utils.cancellation import shield_and_drain
from nanobot.utils.helpers import build_image_content_blocks, detect_image_mime
@@ -664,22 +667,31 @@ def _match_covers_line(match: _MatchSpan, line: int) -> bool:
return match.line <= line <= _match_end_line(match)
def _find_exact_matches(content: str, old_text: str) -> list[_MatchSpan]:
def _find_exact_matches(
content: str,
old_text: str,
*,
max_matches: int | None = None,
) -> list[_MatchSpan]:
matches: list[_MatchSpan] = []
start = 0
while True:
idx = content.find(old_text, start)
search_start = 0
line_start = 0
line = 1
while max_matches is None or len(matches) < max_matches:
idx = content.find(old_text, search_start)
if idx == -1:
break
line += content.count("\n", line_start, idx)
matches.append(
_MatchSpan(
start=idx,
end=idx + len(old_text),
text=content[idx : idx + len(old_text)],
line=content.count("\n", 0, idx) + 1,
line=line,
)
)
start = idx + max(1, len(old_text))
line_start = idx
search_start = idx + max(1, len(old_text))
return matches
@@ -735,27 +747,36 @@ def _find_quote_matches(content: str, old_text: str) -> list[_MatchSpan]:
norm_content = _normalize_quotes(content)
norm_old = _normalize_quotes(old_text)
matches: list[_MatchSpan] = []
start = 0
search_start = 0
line_start = 0
line = 1
while True:
idx = norm_content.find(norm_old, start)
idx = norm_content.find(norm_old, search_start)
if idx == -1:
break
line += content.count("\n", line_start, idx)
matches.append(
_MatchSpan(
start=idx,
end=idx + len(old_text),
text=content[idx : idx + len(old_text)],
line=content.count("\n", 0, idx) + 1,
line=line,
)
)
start = idx + max(1, len(norm_old))
line_start = idx
search_start = idx + max(1, len(norm_old))
return matches
def _find_matches(content: str, old_text: str) -> list[_MatchSpan]:
"""Locate all matches using progressively looser strategies."""
def _find_matches(
content: str,
old_text: str,
*,
max_exact_matches: int | None = None,
) -> list[_MatchSpan]:
"""Locate matches using progressively looser strategies."""
for matcher in (
lambda: _find_exact_matches(content, old_text),
lambda: _find_exact_matches(content, old_text, max_matches=max_exact_matches),
lambda: _find_trim_matches(content, old_text),
lambda: _find_trim_matches(content, old_text, normalize_quotes=True),
lambda: _find_quote_matches(content, old_text),
@@ -869,6 +890,43 @@ class EditFileTool(_FsTool):
new_text: str | None = None,
replace_all: bool = False, occurrence: int | None = None,
line_hint: int | None = None, expected_replacements: int | None = None, **kwargs: Any,
) -> str:
cancelled = threading.Event()
commit_lock = threading.Lock()
try:
return await asyncio.to_thread(
self._execute_sync,
path=path,
old_text=old_text,
new_text=new_text,
replace_all=replace_all,
occurrence=occurrence,
line_hint=line_hint,
expected_replacements=expected_replacements,
cancelled=cancelled,
commit_lock=commit_lock,
)
except asyncio.CancelledError:
cancelled.set()
# If a commit already started, do not report cancellation until the
# file bytes and FileStates record are settled. Otherwise, taking
# the lock first guarantees the worker observes ``cancelled`` before
# it can mutate the target.
await shield_and_drain(
asyncio.to_thread(self._wait_for_commit, commit_lock)
)
raise
@staticmethod
def _wait_for_commit(commit_lock: threading.Lock) -> None:
with commit_lock:
pass
def _execute_sync(
self, *, path: str | None, old_text: str | None,
new_text: str | None, replace_all: bool, occurrence: int | None,
line_hint: int | None, expected_replacements: int | None,
cancelled: threading.Event, commit_lock: threading.Lock,
) -> str:
try:
if not path:
@@ -892,9 +950,12 @@ class EditFileTool(_FsTool):
# Create-file semantics: old_text='' + file doesn't exist → create
if not file_exists:
if old_text == "":
fp.parent.mkdir(parents=True, exist_ok=True)
fp.write_text(new_text, encoding="utf-8")
self._file_states.record_write(fp)
with commit_lock:
if cancelled.is_set():
return ToolResult.error("Error: edit_file cancelled.")
fp.parent.mkdir(parents=True, exist_ok=True)
fp.write_text(new_text, encoding="utf-8")
self._file_states.record_write(fp)
return f"Successfully created {fp}"
return self._file_not_found_msg(path, fp)
@@ -912,8 +973,11 @@ class EditFileTool(_FsTool):
content = raw.decode("utf-8")
if content.strip():
return ToolResult.error(f"Error: Cannot create file — {path} already exists and is not empty.")
fp.write_text(new_text, encoding="utf-8")
self._file_states.record_write(fp)
with commit_lock:
if cancelled.is_set():
return ToolResult.error("Error: edit_file cancelled.")
fp.write_text(new_text, encoding="utf-8")
self._file_states.record_write(fp)
return f"Successfully edited {fp}"
# Read-before-edit check
@@ -923,7 +987,11 @@ class EditFileTool(_FsTool):
uses_crlf = b"\r\n" in raw
content = raw.decode("utf-8").replace("\r\n", "\n")
norm_old = old_text.replace("\r\n", "\n")
matches = _find_matches(content, norm_old)
matches = _find_matches(
content,
norm_old,
max_exact_matches=occurrence,
)
if not matches:
return self._not_found_msg(old_text, content, path)
@@ -1000,8 +1068,11 @@ class EditFileTool(_FsTool):
if uses_crlf:
new_content = new_content.replace("\n", "\r\n")
fp.write_bytes(new_content.encode("utf-8"))
self._file_states.record_write(fp)
with commit_lock:
if cancelled.is_set():
return ToolResult.error("Error: edit_file cancelled.")
fp.write_bytes(new_content.encode("utf-8"))
self._file_states.record_write(fp)
msg = f"Successfully edited {fp}"
if warning:
msg = f"{warning}\n{msg}"
+52 -22
View File
@@ -17,6 +17,7 @@ from nanobot.agent.tools.context import RequestContext, ToolContext, current_req
from nanobot.agent.tools.schema import StringSchema, tool_parameters_schema
from nanobot.bus.runtime_events import GoalStateChanged, RuntimeEventBus, RuntimeEventContext
from nanobot.runtime_context import RuntimeContextBlock, wrap_runtime_context_lines
from nanobot.session.async_compat import call_session_manager
from nanobot.session.goal_state import (
GOAL_STATE_KEY,
MAX_GOAL_OBJECTIVE_CHARS,
@@ -28,6 +29,7 @@ from nanobot.session.goal_state import (
sustained_goal_active,
)
from nanobot.session.turn_continuation import reset_goal_continuation_rounds
from nanobot.utils.cancellation import shield_and_drain
from nanobot.utils.prompt_templates import render_template
if TYPE_CHECKING:
@@ -60,36 +62,68 @@ class _GoalToolsMixin:
self._sessions = sessions
self._runtime_events = runtime_events
def _session(self):
async def _get_or_create_session(self, key: str):
return await call_session_manager(
self._sessions,
"get_or_create_async",
self._sessions.get_or_create,
key,
)
async def _save_session(self, session: Any) -> None:
await call_session_manager(
self._sessions,
"save_async",
self._sessions.save,
session,
)
async def _session(self):
request_ctx = current_request_context()
if request_ctx is None:
return None
key = request_ctx.session_key
if not key:
return None
return self._sessions.get_or_create(key)
return await self._get_or_create_session(key)
def _goal_mutation_allowed(self) -> bool:
return current_request_context() is not None and goal_mutation_allowed()
def _save_goal_state(
async def _save_goal_state(
self,
sess: Any,
blob: dict[str, Any],
*,
reset_continuation: bool = False,
revoke_permission: bool = False,
) -> None:
previous_metadata = deepcopy(sess.metadata)
sess.metadata[GOAL_STATE_KEY] = blob
discard_legacy_goal_state_key(sess.metadata)
if reset_continuation:
reset_goal_continuation_rounds(sess.metadata)
saved = False
async def save_and_publish() -> None:
nonlocal saved
sess.metadata[GOAL_STATE_KEY] = blob
discard_legacy_goal_state_key(sess.metadata)
if reset_continuation:
reset_goal_continuation_rounds(sess.metadata)
try:
await self._save_session(sess)
except BaseException:
sess.metadata.clear()
sess.metadata.update(previous_metadata)
raise
saved = True
await self._publish_goal_state_changed(sess.metadata)
try:
self._sessions.save(sess)
except BaseException:
sess.metadata.clear()
sess.metadata.update(previous_metadata)
raise
await shield_and_drain(save_and_publish())
finally:
# This ContextVar belongs to the caller task, not the settlement task.
# Apply the post-save permission effect here even when cancellation was
# delayed until the durable save and runtime notification completed.
if revoke_permission and saved:
revoke_goal_mutation_permission()
async def _publish_goal_state_changed(self, metadata: dict[str, Any]) -> None:
runtime_events = self._runtime_events
@@ -175,7 +209,7 @@ class CreateGoalTool(Tool, _GoalToolsMixin):
) -> RuntimeContextBlock | None:
if not request.session_key:
return None
session = self._sessions.get_or_create(request.session_key)
session = await self._get_or_create_session(request.session_key)
goal_start_requested = explicit_goal_requested(request.metadata)
goal_active = sustained_goal_active(session.metadata)
if not goal_start_requested and not goal_active:
@@ -197,7 +231,7 @@ class CreateGoalTool(Tool, _GoalToolsMixin):
ui_summary: str | None = None,
**kwargs: Any,
) -> str:
sess = self._session()
sess = await self._session()
if sess is None:
return ToolResult.error(
"Error: create_goal requires an active chat session (missing routing context)."
@@ -225,8 +259,7 @@ class CreateGoalTool(Tool, _GoalToolsMixin):
"ui_summary": summary,
"started_at": _iso_now(),
}
self._save_goal_state(sess, blob, reset_continuation=True)
await self._publish_goal_state_changed(sess.metadata)
await self._save_goal_state(sess, blob, reset_continuation=True)
extra = f"\nSummary line: {summary}" if summary else ""
return (
"Goal recorded. Keep working toward the objective using ordinary tools. "
@@ -305,7 +338,7 @@ class UpdateGoalTool(Tool, _GoalToolsMixin):
ui_summary: str | None = None,
**kwargs: Any,
) -> str:
sess = self._session()
sess = await self._session()
if sess is None:
return ToolResult.error("Error: update_goal requires an active chat session.")
prior = parse_goal_state(goal_state_raw(sess.metadata))
@@ -340,8 +373,7 @@ class UpdateGoalTool(Tool, _GoalToolsMixin):
"previous_objective": str(prior.get("objective") or ""),
"recap": (recap or "").strip(),
}
self._save_goal_state(sess, blob, reset_continuation=True)
await self._publish_goal_state_changed(sess.metadata)
await self._save_goal_state(sess, blob, reset_continuation=True)
extra = f"\nSummary line: {summary}" if summary else ""
return "Goal replaced. Continue toward the new objective using ordinary tools." + extra
@@ -359,9 +391,7 @@ class UpdateGoalTool(Tool, _GoalToolsMixin):
}
if normalized == "complete":
blob["completed_at"] = ended
self._save_goal_state(sess, blob)
revoke_goal_mutation_permission()
await self._publish_goal_state_changed(sess.metadata)
await self._save_goal_state(sess, blob, revoke_permission=True)
tail = (recap or "").strip()
label = {
+5 -5
View File
@@ -20,10 +20,10 @@ from nanobot.agent.tools.base import Tool, ToolResult
from nanobot.agent.tools.registry import ToolRegistry
from nanobot.security.network import (
PinnedDNSAsyncTransport,
async_resolve_url_target,
async_validate_url_target,
env_proxy_applies_to_url,
httpx_env_proxy_mounts,
resolve_url_target,
validate_url_target,
)
from nanobot.utils.cancellation import task_is_cancelling
@@ -249,7 +249,7 @@ async def _probe_http_url(url: str, timeout: float = 3.0) -> bool:
port = parsed.port
if not port:
port = 443 if parsed.scheme == "https" else 80
ok, _, resolved_ips = resolve_url_target(url)
ok, _, resolved_ips = await async_resolve_url_target(url)
if not ok:
return False
if env_proxy_applies_to_url(url):
@@ -298,7 +298,7 @@ def _pinned_transport_kwargs() -> dict[str, Any]:
async def _validate_mcp_request_url(request: httpx.Request) -> None:
"""Validate each outgoing MCP HTTP request, including redirect targets."""
ok, error = validate_url_target(str(request.url))
ok, error = await async_validate_url_target(str(request.url))
if not ok:
raise httpx.RequestError(
f"Blocked unsafe MCP URL {_redact_url(str(request.url))} ({error})",
@@ -1031,7 +1031,7 @@ async def connect_mcp_servers(
return False
if transport_type in {"sse", "streamableHttp"}:
ok, error = validate_url_target(cfg.url)
ok, error = await async_validate_url_target(cfg.url)
if not ok:
logger.warning(
"MCP server '{}': blocked unsafe URL {} ({})",
+23
View File
@@ -107,6 +107,13 @@ class RuntimeControl(Protocol):
session_key: str | None,
) -> LLMRuntime: ...
async def set_model_preset_async(
self,
name: str,
*,
session_key: str | None,
) -> LLMRuntime: ...
def set_max_iterations(self, value: int) -> None: ...
def set_context_window_tokens(self, value: int) -> LLMRuntime: ...
@@ -162,6 +169,12 @@ class _RuntimeControlTarget(Protocol):
def set_session_model_preset(self, session_key: str, name: str) -> LLMRuntime: ...
async def set_session_model_preset_async(
self,
session_key: str,
name: str,
) -> LLMRuntime: ...
class AgentRuntimeControl:
"""Allowlisted adapter from agent-loop state to ``RuntimeControl``."""
@@ -208,6 +221,16 @@ class AgentRuntimeControl:
return self.__target.set_session_model_preset(session_key, name)
return self.__target.set_model_preset(name)
async def set_model_preset_async(
self,
name: str,
*,
session_key: str | None,
) -> LLMRuntime:
if session_key is not None:
return await self.__target.set_session_model_preset_async(session_key, name)
return self.__target.set_model_preset(name)
def set_max_iterations(self, value: int) -> None:
self.__target.max_iterations = value
self.__target.subagents.max_iterations = value
+103 -53
View File
@@ -4,9 +4,12 @@
from __future__ import annotations
import asyncio
import fnmatch
import os
import re
import threading
import time
from contextlib import suppress
from pathlib import Path, PurePosixPath
from typing import Any, Iterable, TypeVar
@@ -125,6 +128,8 @@ class _SearchTool(_FsTool):
class FindFilesTool(_SearchTool):
"""Find files by path fragment, glob, or type."""
_scopes = {"core", "subagent"}
_MAX_SCAN_PATHS = 500_000
_MAX_SCAN_SECONDS = 30.0
@property
def name(self) -> str:
@@ -218,66 +223,111 @@ class FindFilesTool(_SearchTool):
offset: int = 0,
**kwargs: Any,
) -> str:
cancelled = threading.Event()
try:
target = self._resolve(path or ".")
if not target.exists():
return ToolResult.error(f"Error: Path not found: {path}")
if not (target.is_dir() or target.is_file()):
return ToolResult.error(f"Error: Unsupported path: {path}")
if sort not in {"path", "modified"}:
return ToolResult.error("Error: sort must be 'path' or 'modified'")
limit = (
_DEFAULT_FILE_HEAD_LIMIT
if head_limit is None
else None if head_limit == 0 else head_limit
return await asyncio.to_thread(
self._execute_sync,
path=path,
query=query,
glob=glob,
file_type=type,
include_dirs=include_dirs,
sort=sort,
head_limit=head_limit,
offset=offset,
cancelled=cancelled,
)
root = target if target.is_dir() else target.parent
matches: list[tuple[str, float]] = []
for candidate in self._iter_paths(target, include_dirs=include_dirs):
if candidate.is_dir() and not include_dirs:
continue
rel_path = candidate.relative_to(root).as_posix()
display_path = self._display_path(candidate, root)
name = candidate.name
if glob and not _match_glob(rel_path, name, glob):
continue
if candidate.is_file() and not _matches_type(name, type):
continue
if candidate.is_dir() and type:
continue
if not _matches_query(display_path, query):
continue
try:
mtime = candidate.stat().st_mtime
except OSError:
mtime = 0.0
suffix = "/" if candidate.is_dir() else ""
matches.append((display_path + suffix, mtime))
if sort == "modified":
matches.sort(key=lambda item: (-item[1], item[0]))
else:
matches.sort(key=lambda item: item[0])
paths = [item[0] for item in matches]
paged, truncated = _paginate(paths, limit, offset)
if not paged:
return "No files found"
result = "\n".join(paged)
note = _pagination_note(limit, offset, truncated)
if note:
result += "\n\n" + note
return result
except asyncio.CancelledError:
cancelled.set()
raise
except PermissionError as e:
return ToolResult.error(f"Error: {e}")
except Exception as e:
return ToolResult.error(f"Error finding files: {e}")
def _execute_sync(
self,
*,
path: str,
query: str | None,
glob: str | None,
file_type: str | None,
include_dirs: bool,
sort: str,
head_limit: int | None,
offset: int,
cancelled: threading.Event,
) -> str:
target = self._resolve(path or ".")
if not target.exists():
return ToolResult.error(f"Error: Path not found: {path}")
if not (target.is_dir() or target.is_file()):
return ToolResult.error(f"Error: Unsupported path: {path}")
if sort not in {"path", "modified"}:
return ToolResult.error("Error: sort must be 'path' or 'modified'")
limit = (
_DEFAULT_FILE_HEAD_LIMIT
if head_limit is None
else None if head_limit == 0 else head_limit
)
root = target if target.is_dir() else target.parent
matches: list[tuple[str, float]] = []
deadline = time.monotonic() + self._MAX_SCAN_SECONDS
scanned = 0
for candidate in self._iter_paths(target, include_dirs=include_dirs):
if cancelled.is_set():
raise RuntimeError("find_files scan cancelled")
scanned += 1
if scanned > self._MAX_SCAN_PATHS:
return ToolResult.error(
f"Error: find_files scan exceeded {self._MAX_SCAN_PATHS} paths; "
"narrow path, query, glob, or type and retry."
)
if time.monotonic() > deadline:
return ToolResult.error(
f"Error: find_files scan exceeded {self._MAX_SCAN_SECONDS:g} seconds; "
"narrow path, query, glob, or type and retry."
)
if candidate.is_dir() and not include_dirs:
continue
rel_path = candidate.relative_to(root).as_posix()
display_path = self._display_path(candidate, root)
name = candidate.name
if glob and not _match_glob(rel_path, name, glob):
continue
if candidate.is_file() and not _matches_type(name, file_type):
continue
if candidate.is_dir() and file_type:
continue
if not _matches_query(display_path, query):
continue
try:
mtime = candidate.stat().st_mtime
except OSError:
mtime = 0.0
suffix = "/" if candidate.is_dir() else ""
matches.append((display_path + suffix, mtime))
if sort == "modified":
matches.sort(key=lambda item: (-item[1], item[0]))
else:
matches.sort(key=lambda item: item[0])
paths = [item[0] for item in matches]
paged, truncated = _paginate(paths, limit, offset)
if not paged:
return "No files found"
result = "\n".join(paged)
note = _pagination_note(limit, offset, truncated)
if note:
result += "\n\n" + note
return result
class GrepTool(_SearchTool):
"""Search file contents using a regex-like pattern."""
+34 -1
View File
@@ -370,7 +370,7 @@ class MyTool(Tool):
if not self._modify_allowed:
return ToolResult.error("Error: set is disabled (tools.my.allow_set is false)")
if action in ("modify", "set"):
return self._modify(key, value)
return await self._modify_async(key, value)
return f"Unknown action: {action}"
# -- inspect --
@@ -492,6 +492,11 @@ class MyTool(Tool):
return ToolResult.error(f"Error: '{key}' is read-only and cannot be modified")
return self._modify_scratchpad(key, value)
async def _modify_async(self, key: str | None, value: Any) -> str:
if key == "model_preset":
return await self._modify_model_preset_async(value)
return self._modify(key, value)
def _modify_model_preset(self, value: Any) -> str:
if not isinstance(value, str) or not value.strip():
return ToolResult.error("Error: 'model_preset' must be a non-empty string")
@@ -520,6 +525,34 @@ class MyTool(Tool):
f"context_window_tokens is now {runtime.context_window_tokens!r}"
)
async def _modify_model_preset_async(self, value: Any) -> str:
if not isinstance(value, str) or not value.strip():
return ToolResult.error("Error: 'model_preset' must be a non-empty string")
name = value.strip()
session_key = current_request_session_key()
old = self._runtime_control.snapshot().model_preset
try:
runtime = await self._runtime_control.set_model_preset_async(
name,
session_key=session_key,
)
except (KeyError, ValueError) as exc:
message = str(exc.args[0]) if exc.args else str(exc)
punctuation = "" if message.endswith((".", "!", "?")) else "."
return ToolResult.error(f"Error: {message}{punctuation}")
if session_key:
self._audit("modify", f"model_preset = {name!r}")
return (
f"Set model_preset = {name!r} for the next turn; "
f"model will be {runtime.model!r}; "
f"context_window_tokens will be {runtime.context_window_tokens!r}"
)
self._audit("modify", f"model_preset: {old!r} -> {name!r}")
return (
f"Set model_preset = {name!r} (was {old!r}); model is now {runtime.model!r}; "
f"context_window_tokens is now {runtime.context_window_tokens!r}"
)
def _modify_restricted(self, key: str, value: Any) -> str:
spec = self.RESTRICTED[key]
expected = cast(type[Any], spec["type"])
+24 -1
View File
@@ -267,6 +267,7 @@ class ExecTool(Tool):
_MAX_TIMEOUT = 600
_MAX_OUTPUT = 10_000
_PREPARE_TIMEOUT_SECONDS = 6.0
# Kernel device files safe as stdio redirect targets (#3599).
_BENIGN_DEVICE_PATHS: frozenset[str] = frozenset({
@@ -324,7 +325,20 @@ class ExecTool(Tool):
if max_output_chars is None:
max_output_chars = max_output_tokens
prepared = self._prepare_command(command, working_dir, timeout, shell, login)
try:
prepared = await asyncio.wait_for(
asyncio.to_thread(
self._prepare_command,
command,
working_dir,
timeout,
shell,
login,
),
timeout=self._PREPARE_TIMEOUT_SECONDS,
)
except asyncio.TimeoutError:
return ToolResult.error("Error: command validation timed out")
if isinstance(prepared, str):
return prepared
@@ -916,6 +930,15 @@ class ExecTool(Tool):
if self._is_benign_device_path(expanded):
continue
except Exception:
# ``Path.expanduser()`` raises when a named user's home
# cannot be resolved (notably on Windows). An extracted
# home path must fail closed rather than bypass the guard.
if raw.strip().startswith("~"):
return ToolResult.error(
"Error: Command blocked by safety guard "
"(path outside working dir)"
+ _WORKSPACE_BOUNDARY_NOTE
)
continue
if self._is_benign_device_path(str(p)):
+12 -14
View File
@@ -96,7 +96,7 @@ def _normalize(text: str) -> str:
def _validate_url(url: str) -> tuple[bool, str]:
"""Validate URL scheme/domain. Does NOT check resolved IPs (use _validate_url_safe for that)."""
"""Validate URL scheme/domain. Does not resolve IPs; use the async safe helper for that."""
try:
p = urlparse(url)
if p.scheme not in ('http', 'https'):
@@ -108,18 +108,16 @@ def _validate_url(url: str) -> tuple[bool, str]:
return False, str(e)
def _validate_url_safe(url: str) -> tuple[bool, str]:
"""Validate URL with SSRF protection: scheme, domain, and resolved IP check."""
from nanobot.security.network import validate_url_target
async def _async_validate_url_safe(url: str) -> tuple[bool, str]:
from nanobot.security.network import async_validate_url_target
return validate_url_target(url)
return await async_validate_url_target(url)
def _resolve_url_safe(url: str) -> tuple[bool, str, tuple[str, ...]]:
"""Validate URL and return the resolved IPs to pin during the request."""
from nanobot.security.network import resolve_url_target
async def _async_resolve_url_safe(url: str) -> tuple[bool, str, tuple[str, ...]]:
from nanobot.security.network import async_resolve_url_target
return resolve_url_target(url)
return await async_resolve_url_target(url)
def _pinned_dns_transport() -> httpx.AsyncBaseTransport:
@@ -209,7 +207,7 @@ async def _get_with_safe_redirects(
"""GET a URL while validating every redirect target before requesting it."""
current_url = url
for _ in range(MAX_REDIRECTS + 1):
is_valid, error_msg, _ = _resolve_url_safe(current_url)
is_valid, error_msg, _ = await _async_resolve_url_safe(current_url)
if not is_valid:
return None, f"Redirect blocked: {error_msg}"
@@ -229,7 +227,7 @@ async def _get_with_safe_redirects(
return response, None
next_url = urljoin(str(response.url), location)
is_valid, error_msg = _validate_url_safe(next_url)
is_valid, error_msg = await _async_validate_url_safe(next_url)
if not is_valid:
await response.aclose()
return None, f"Redirect blocked: {error_msg}"
@@ -249,7 +247,7 @@ async def _stream_with_safe_redirects(
current_url = url
chain_carries_credentials = _url_carries_credentials(url)
for _ in range(MAX_REDIRECTS + 1):
is_valid, error_msg, _ = _resolve_url_safe(current_url)
is_valid, error_msg, _ = await _async_resolve_url_safe(current_url)
if not is_valid:
return None, None, f"Redirect blocked: {error_msg}", chain_carries_credentials
@@ -283,7 +281,7 @@ async def _stream_with_safe_redirects(
chain_carries_credentials = (
chain_carries_credentials or _url_carries_credentials(next_url)
)
is_valid, error_msg = _validate_url_safe(next_url)
is_valid, error_msg = await _async_validate_url_safe(next_url)
if not is_valid:
await stream.__aexit__(None, None, None)
return None, None, f"Redirect blocked: {error_msg}", chain_carries_credentials
@@ -1106,7 +1104,7 @@ class WebFetchTool(Tool):
url = url.strip(" \t\r\n`\"'")
extract_mode = kwargs.pop("extractMode", extract_mode)
max_chars = cast(int, kwargs.pop("maxChars", max_chars) or self.max_chars)
is_valid, error_msg = _validate_url_safe(url)
is_valid, error_msg = await _async_validate_url_safe(url)
if not is_valid:
return json.dumps({"error": f"URL validation failed: {error_msg}", "url": url}, ensure_ascii=False)
+374 -41
View File
@@ -2,15 +2,20 @@
from __future__ import annotations
import asyncio
import ctypes
import json
import os
import re
import shlex
import shutil
import signal
import subprocess
import sys
import time
from collections.abc import Iterable
from contextlib import suppress
from ctypes import wintypes
from dataclasses import dataclass
from importlib import metadata as importlib_metadata
from pathlib import Path
@@ -97,6 +102,141 @@ class CliAppsRuntimeConfig:
catalog_ttl_seconds: int = 3600
@dataclass(slots=True)
class _PreparedCliRun:
name: str
entry: str
resolved: str
args: list[str]
cwd: Path
timeout: int
env: dict[str, str]
artifact_snapshot: dict[Path, tuple[int, int]]
class _JobObjectBasicLimitInformation(ctypes.Structure):
_fields_ = [
("PerProcessUserTimeLimit", ctypes.c_int64),
("PerJobUserTimeLimit", ctypes.c_int64),
("LimitFlags", wintypes.DWORD),
("MinimumWorkingSetSize", ctypes.c_size_t),
("MaximumWorkingSetSize", ctypes.c_size_t),
("ActiveProcessLimit", wintypes.DWORD),
("Affinity", ctypes.c_size_t),
("PriorityClass", wintypes.DWORD),
("SchedulingClass", wintypes.DWORD),
]
class _IoCounters(ctypes.Structure):
_fields_ = [
("ReadOperationCount", ctypes.c_uint64),
("WriteOperationCount", ctypes.c_uint64),
("OtherOperationCount", ctypes.c_uint64),
("ReadTransferCount", ctypes.c_uint64),
("WriteTransferCount", ctypes.c_uint64),
("OtherTransferCount", ctypes.c_uint64),
]
class _JobObjectExtendedLimitInformation(ctypes.Structure):
_fields_ = [
("BasicLimitInformation", _JobObjectBasicLimitInformation),
("IoInfo", _IoCounters),
("ProcessMemoryLimit", ctypes.c_size_t),
("JobMemoryLimit", ctypes.c_size_t),
("PeakProcessMemoryUsed", ctypes.c_size_t),
("PeakJobMemoryUsed", ctypes.c_size_t),
]
class _WindowsJob:
"""Best-effort Windows process tree ownership for timeout/cancellation."""
_KILL_ON_JOB_CLOSE = 0x00002000
_EXTENDED_LIMIT_INFORMATION = 9
_PROCESS_TERMINATE = 0x0001
_PROCESS_SET_QUOTA = 0x0100
def __init__(self) -> None:
win_dll = getattr(ctypes, "WinDLL")
self._kernel32 = win_dll("kernel32", use_last_error=True)
self._kernel32.CreateJobObjectW.argtypes = [wintypes.LPVOID, wintypes.LPCWSTR]
self._kernel32.CreateJobObjectW.restype = wintypes.HANDLE
self._kernel32.OpenProcess.argtypes = [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD]
self._kernel32.OpenProcess.restype = wintypes.HANDLE
self._kernel32.AssignProcessToJobObject.argtypes = [wintypes.HANDLE, wintypes.HANDLE]
self._kernel32.AssignProcessToJobObject.restype = wintypes.BOOL
self._kernel32.SetInformationJobObject.argtypes = [
wintypes.HANDLE,
ctypes.c_int,
wintypes.LPVOID,
wintypes.DWORD,
]
self._kernel32.SetInformationJobObject.restype = wintypes.BOOL
self._kernel32.TerminateJobObject.argtypes = [wintypes.HANDLE, wintypes.UINT]
self._kernel32.TerminateJobObject.restype = wintypes.BOOL
self._kernel32.CloseHandle.argtypes = [wintypes.HANDLE]
self._kernel32.CloseHandle.restype = wintypes.BOOL
self._handle: Any = self._kernel32.CreateJobObjectW(None, None)
if not self._handle:
raise OSError(ctypes.get_last_error(), "CreateJobObjectW failed")
try:
self._set_kill_on_close(True)
except OSError:
self._kernel32.CloseHandle(self._handle)
self._handle = None
raise
@classmethod
def create(cls) -> _WindowsJob | None:
if os.name != "nt":
return None
try:
return cls()
except OSError as exc:
logger.debug("CLI Apps: Windows job object unavailable: {}", exc)
return None
def _set_kill_on_close(self, enabled: bool) -> None:
info = _JobObjectExtendedLimitInformation()
info.BasicLimitInformation.LimitFlags = self._KILL_ON_JOB_CLOSE if enabled else 0
ok = self._kernel32.SetInformationJobObject(
self._handle,
self._EXTENDED_LIMIT_INFORMATION,
ctypes.byref(info),
ctypes.sizeof(info),
)
if not ok:
raise OSError(ctypes.get_last_error(), "SetInformationJobObject failed")
def assign(self, pid: int) -> bool:
process_handle = self._kernel32.OpenProcess(
self._PROCESS_TERMINATE | self._PROCESS_SET_QUOTA,
False,
pid,
)
if not process_handle:
return False
try:
return bool(self._kernel32.AssignProcessToJobObject(self._handle, process_handle))
finally:
self._kernel32.CloseHandle(process_handle)
def terminate(self) -> None:
if self._handle and not self._kernel32.TerminateJobObject(self._handle, 1):
raise OSError(ctypes.get_last_error(), "TerminateJobObject failed")
def close(self, *, kill_descendants: bool) -> None:
if not self._handle:
return
if not kill_descendants:
with suppress(OSError):
self._set_kill_on_close(False)
self._kernel32.CloseHandle(self._handle)
self._handle = None
_BRANDS: dict[str, tuple[str, str]] = {
"1password-cli": ("1password", "#3B66BC"),
"arcgis": ("arcgis", "#2C7AC3"),
@@ -1428,6 +1568,197 @@ Use the `run_cli_app` tool with `name="{name}"` for command execution. Do not in
lines.append(f"- {rel} ({kind}, {self._format_artifact_size(path)})")
return lines
def _prepare_run(
self,
name: str,
args: list[str] | None,
*,
json_output: bool,
working_dir: str | None,
timeout: int | None,
restrict_to_workspace: bool,
) -> _PreparedCliRun:
app = self.get_app(name)
installed = self._load_installed()
app_name = str(app["name"])
if app_name not in installed:
raise CliAppError(f"CLI app '{name}' is not installed")
cwd = self._resolve_cwd(working_dir, restrict_to_workspace=restrict_to_workspace)
entry = str(installed[app_name].get("entry_point") or app.get("entry_point") or "")
resolved = shutil.which(entry)
if not entry or not resolved:
raise CliAppError(f"{entry or name} is not available on PATH")
clean_args = [str(arg) for arg in (args or [])]
if json_output and "--json" not in clean_args:
clean_args = ["--json", *clean_args]
effective_timeout = max(1, min(timeout or self.runtime.run_timeout, 600))
return _PreparedCliRun(
name=name,
entry=entry,
resolved=resolved,
args=clean_args,
cwd=cwd,
timeout=effective_timeout,
env=self._subprocess_env(),
artifact_snapshot=self._artifact_snapshot(cwd),
)
def _format_run_result(
self,
prepared: _PreparedCliRun,
*,
returncode: int,
stdout: str,
stderr: str,
) -> str:
command = " ".join([prepared.entry, *(shlex.quote(arg) for arg in prepared.args)])
output = [
f"CLI app '{prepared.name}' exited {returncode}.",
f"Command: {command}",
]
if stdout:
output.append("\nSTDOUT:\n" + stdout.rstrip())
if stderr:
output.append("\nSTDERR:\n" + stderr.rstrip())
artifacts = self._changed_artifacts(prepared.cwd, prepared.artifact_snapshot)
if artifacts:
output.append(
"\nArtifacts created or updated:\n"
+ "\n".join(self._format_artifact_lines(prepared.cwd, artifacts))
)
if any(path.suffix.lower() in _INLINE_ARTIFACT_EXTENSIONS for path in artifacts):
output.append(
"\nTo show a preview in WebUI, reference a raster artifact with Markdown "
"using its workspace-relative path, for example `![diagram](diagram.png)`."
)
return _truncate("\n".join(output))
@staticmethod
def _terminate_run_process_sync(
process: subprocess.Popen[str],
job: _WindowsJob | None,
) -> None:
if job is not None:
with suppress(OSError):
job.terminate()
job.close(kill_descendants=True)
elif os.name == "nt":
with suppress(OSError, subprocess.TimeoutExpired):
subprocess.run(
["taskkill", "/PID", str(process.pid), "/T", "/F"],
check=False,
stdout=subprocess.DEVNULL,
stderr=subprocess.DEVNULL,
timeout=5,
)
else:
with suppress(ProcessLookupError, PermissionError):
os.killpg(process.pid, signal.SIGKILL)
if process.poll() is None:
with suppress(ProcessLookupError):
process.kill()
with suppress(subprocess.TimeoutExpired):
process.wait(timeout=5)
@staticmethod
async def _terminate_run_process(
process: asyncio.subprocess.Process,
job: _WindowsJob | None,
) -> None:
if job is not None:
with suppress(OSError):
await asyncio.to_thread(job.terminate)
job.close(kill_descendants=True)
elif os.name == "nt":
with suppress(OSError, asyncio.TimeoutError):
await asyncio.wait_for(
asyncio.to_thread(
subprocess.run,
["taskkill", "/PID", str(process.pid), "/T", "/F"],
check=False,
stdout=subprocess.DEVNULL,
stderr=subprocess.DEVNULL,
timeout=5,
),
timeout=6.0,
)
else:
with suppress(ProcessLookupError, PermissionError):
os.killpg(process.pid, signal.SIGKILL)
if process.returncode is None:
with suppress(ProcessLookupError):
process.kill()
with suppress(asyncio.TimeoutError, ProcessLookupError):
await asyncio.wait_for(process.wait(), timeout=5.0)
async def run_async(
self,
name: str,
args: list[str] | None = None,
*,
json_output: bool = False,
working_dir: str | None = None,
timeout: int | None = None,
restrict_to_workspace: bool = False,
) -> str:
prepared = await asyncio.to_thread(
self._prepare_run,
name,
args,
json_output=json_output,
working_dir=working_dir,
timeout=timeout,
restrict_to_workspace=restrict_to_workspace,
)
process_kwargs: dict[str, Any] = {}
if os.name == "nt":
process_kwargs["creationflags"] = subprocess.CREATE_NEW_PROCESS_GROUP
else:
process_kwargs["start_new_session"] = True
job = _WindowsJob.create()
try:
process = await asyncio.create_subprocess_exec(
prepared.resolved,
*prepared.args,
cwd=str(prepared.cwd),
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
env=prepared.env,
**process_kwargs,
)
except BaseException:
if job is not None:
job.close(kill_descendants=False)
raise
if job is not None and not job.assign(process.pid):
job.close(kill_descendants=False)
job = None
try:
stdout_raw, stderr_raw = await asyncio.wait_for(
process.communicate(),
timeout=prepared.timeout,
)
except asyncio.TimeoutError:
await self._terminate_run_process(process, job)
return f"CLI app '{prepared.name}' timed out after {prepared.timeout}s"
except asyncio.CancelledError:
await self._terminate_run_process(process, job)
raise
except BaseException:
await self._terminate_run_process(process, job)
raise
if job is not None:
job.close(kill_descendants=False)
stdout = stdout_raw.decode("utf-8", errors="replace")
stderr = stderr_raw.decode("utf-8", errors="replace")
return await asyncio.to_thread(
self._format_run_result,
prepared,
returncode=process.returncode or 0,
stdout=stdout,
stderr=stderr,
)
def run(
self,
name: str,
@@ -1438,50 +1769,52 @@ Use the `run_cli_app` tool with `name="{name}"` for command execution. Do not in
timeout: int | None = None,
restrict_to_workspace: bool = False,
) -> str:
app = self.get_app(name)
installed = self._load_installed()
if str(app["name"]) not in installed:
raise CliAppError(f"CLI app '{name}' is not installed")
cwd = self._resolve_cwd(working_dir, restrict_to_workspace=restrict_to_workspace)
entry = str(installed[str(app["name"])].get("entry_point") or app.get("entry_point") or "")
resolved = shutil.which(entry)
if not entry or not resolved:
raise CliAppError(f"{entry or name} is not available on PATH")
clean_args = [str(arg) for arg in (args or [])]
if json_output and "--json" not in clean_args:
clean_args = ["--json", *clean_args]
effective_timeout = max(1, min(timeout or self.runtime.run_timeout, 600))
artifact_snapshot = self._artifact_snapshot(cwd)
prepared = self._prepare_run(
name,
args,
json_output=json_output,
working_dir=working_dir,
timeout=timeout,
restrict_to_workspace=restrict_to_workspace,
)
process_kwargs: dict[str, Any] = {}
if os.name == "nt":
process_kwargs["creationflags"] = subprocess.CREATE_NEW_PROCESS_GROUP
else:
process_kwargs["start_new_session"] = True
job = _WindowsJob.create()
try:
result = subprocess.run(
[resolved, *clean_args],
cwd=str(cwd),
capture_output=True,
process = subprocess.Popen(
[prepared.resolved, *prepared.args],
cwd=str(prepared.cwd),
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
text=True,
encoding="utf-8",
errors="replace",
timeout=effective_timeout,
env=self._subprocess_env(),
env=prepared.env,
**process_kwargs,
)
except BaseException:
if job is not None:
job.close(kill_descendants=False)
raise
if job is not None and not job.assign(process.pid):
job.close(kill_descendants=False)
job = None
try:
stdout, stderr = process.communicate(timeout=prepared.timeout)
except subprocess.TimeoutExpired:
return f"CLI app '{name}' timed out after {effective_timeout}s"
output = [
f"CLI app '{name}' exited {result.returncode}.",
f"Command: {entry} {' '.join(shlex.quote(arg) for arg in clean_args)}".rstrip(),
]
if result.stdout:
output.append("\nSTDOUT:\n" + result.stdout.rstrip())
if result.stderr:
output.append("\nSTDERR:\n" + result.stderr.rstrip())
artifacts = self._changed_artifacts(cwd, artifact_snapshot)
if artifacts:
output.append(
"\nArtifacts created or updated:\n"
+ "\n".join(self._format_artifact_lines(cwd, artifacts))
)
if any(path.suffix.lower() in _INLINE_ARTIFACT_EXTENSIONS for path in artifacts):
output.append(
"\nTo show a preview in WebUI, reference a raster artifact with Markdown "
"using its workspace-relative path, for example `![diagram](diagram.png)`."
)
return _truncate("\n".join(output))
self._terminate_run_process_sync(process, job)
return f"CLI app '{prepared.name}' timed out after {prepared.timeout}s"
except BaseException:
self._terminate_run_process_sync(process, job)
raise
if job is not None:
job.close(kill_descendants=False)
return self._format_run_result(
prepared,
returncode=process.returncode,
stdout=stdout,
stderr=stderr,
)
+19 -10
View File
@@ -21,7 +21,10 @@ from nanobot.bus.events import OutboundMessage
from nanobot.bus.queue import MessageBus
from nanobot.channels.base import BaseChannel
from nanobot.config.schema import Base
from nanobot.security.network import validate_resolved_url, validate_url_target
from nanobot.security.network import (
async_validate_resolved_url,
async_validate_url_target,
)
DINGTALK_MAX_REMOTE_MEDIA_BYTES = 20 * 1024 * 1024
DINGTALK_MAX_REMOTE_MEDIA_REDIRECTS = 3
@@ -417,8 +420,8 @@ class DingTalkChannel(BaseChannel):
return self._zip_bytes(filename, data)
return data, filename, content_type
def _validate_remote_media_url(self, media_ref: str) -> bool:
ok, err = validate_url_target(media_ref)
async def _validate_remote_media_url(self, media_ref: str) -> bool:
ok, err = await async_validate_url_target(media_ref)
if not ok:
self.logger.warning("remote media URL blocked ref={} reason={}", media_ref, err)
return False
@@ -434,7 +437,11 @@ class DingTalkChannel(BaseChannel):
allowed_hosts = {host.lower() for host in self.config.remote_media_redirect_allowed_hosts}
return next_host in allowed_hosts
def _next_remote_media_url(self, current_url: str, location: str | None) -> str | None:
async def _next_remote_media_url(
self,
current_url: str,
location: str | None,
) -> str | None:
if not self.config.allow_remote_media_redirects:
self.logger.warning("media download redirect refused ref={}", current_url)
return None
@@ -449,7 +456,7 @@ class DingTalkChannel(BaseChannel):
next_url,
)
return None
if not self._validate_remote_media_url(next_url):
if not await self._validate_remote_media_url(next_url):
return None
return next_url
@@ -461,7 +468,7 @@ class DingTalkChannel(BaseChannel):
if not self._http:
return None, None
if not self._validate_remote_media_url(media_ref):
if not await self._validate_remote_media_url(media_ref):
return None, None
try:
@@ -473,7 +480,7 @@ class DingTalkChannel(BaseChannel):
current_url = media_ref
for _ in range(DINGTALK_MAX_REMOTE_MEDIA_REDIRECTS + 1):
async with stream("GET", current_url, follow_redirects=False) as resp:
final_ok, final_err = validate_resolved_url(str(resp.url))
final_ok, final_err = await async_validate_resolved_url(str(resp.url))
if not final_ok:
self.logger.warning(
"remote media redirect blocked ref={} final={} reason={}",
@@ -483,7 +490,7 @@ class DingTalkChannel(BaseChannel):
)
return None, None
if 300 <= resp.status_code < 400:
next_url = self._next_remote_media_url(
next_url = await self._next_remote_media_url(
str(resp.url), resp.headers.get("location")
)
if not next_url:
@@ -516,7 +523,9 @@ class DingTalkChannel(BaseChannel):
current_url = media_ref
for _ in range(DINGTALK_MAX_REMOTE_MEDIA_REDIRECTS + 1):
resp = await self._http.get(current_url, follow_redirects=False)
final_ok, final_err = validate_resolved_url(str(getattr(resp, "url", current_url)))
final_ok, final_err = await async_validate_resolved_url(
str(getattr(resp, "url", current_url))
)
if not final_ok:
self.logger.warning(
"remote media redirect blocked ref={} final={} reason={}",
@@ -526,7 +535,7 @@ class DingTalkChannel(BaseChannel):
)
return None, None
if 300 <= resp.status_code < 400:
next_url = self._next_remote_media_url(
next_url = await self._next_remote_media_url(
str(getattr(resp, "url", current_url)), resp.headers.get("location")
)
if not next_url:
+3 -3
View File
@@ -24,7 +24,7 @@ from nanobot.bus.queue import MessageBus
from nanobot.channels.base import BaseChannel
from nanobot.config.paths import get_media_dir
from nanobot.config.schema import Base
from nanobot.security.network import validate_url_target
from nanobot.security.network import async_validate_url_target
from nanobot.utils.helpers import safe_filename
_DOWNLOAD_TIMEOUT = aiohttp.ClientTimeout(total=60)
@@ -473,7 +473,7 @@ class NapcatChannel(BaseChannel):
if not ref:
return None
if ref.startswith(("http://", "https://")):
ok, err = validate_url_target(ref)
ok, err = await async_validate_url_target(ref)
if not ok:
logger.warning("napcat: rejected remote image '{}': {}", ref, err)
return None
@@ -525,7 +525,7 @@ class NapcatChannel(BaseChannel):
# logger.debug("napcat: downloading image from {}", url)
if self._http is None:
return None
ok, err = validate_url_target(url)
ok, err = await async_validate_url_target(url)
if not ok:
logger.warning("napcat: skip image '{}': {}", url, err)
return None
@@ -149,9 +149,13 @@ async def test_download_image_rejects_redirects(tmp_path, monkeypatch) -> None:
channel = _channel()
channel._media_root = tmp_path
channel._http = _FakeHttp(_FakeResponse(status=302))
async def allow_url(_url: str) -> tuple[bool, str]:
return True, ""
monkeypatch.setattr(
"nanobot.channels.napcat.runtime.validate_url_target",
lambda _url: (True, ""),
"nanobot.channels.napcat.runtime.async_validate_url_target",
allow_url,
)
result = await channel._download_image({"url": "https://example.com/a.png", "file": "a.png"})
+2 -2
View File
@@ -40,7 +40,7 @@ from nanobot.bus.events import OutboundMessage
from nanobot.bus.queue import MessageBus
from nanobot.channels.base import BaseChannel
from nanobot.config.schema import Base
from nanobot.security.network import validate_url_target
from nanobot.security.network import async_validate_url_target
from nanobot.utils.logging_bridge import redirect_lib_logging
try:
@@ -458,7 +458,7 @@ class QQChannel(BaseChannel):
return None, None
# Remote URL
ok, err = validate_url_target(media_ref)
ok, err = await async_validate_url_target(media_ref)
if not ok:
self.logger.warning("outbound media URL validation failed url={} err={}", media_ref, err)
return None, None
+2 -2
View File
@@ -23,8 +23,8 @@ from nanobot.config.schema import Base
from nanobot.pairing import is_approved
from nanobot.security.network import (
PinnedDNSAsyncTransport,
async_validate_url_target,
httpx_env_proxy_mounts,
validate_url_target,
)
from nanobot.utils.helpers import safe_filename, split_message
@@ -95,7 +95,7 @@ _HTML_DOWNLOAD_PREFIXES = (b"<!doctype html", b"<html")
async def _validate_slack_download_request(request: httpx.Request) -> None:
"""Validate every Slack file request, including redirects, before transport."""
ok, error = validate_url_target(str(request.url))
ok, error = await async_validate_url_target(str(request.url))
if not ok:
raise httpx.RequestError(f"unsafe Slack file URL: {error}", request=request)
@@ -859,13 +859,13 @@ def _patch_download_validation(
monkeypatch: pytest.MonkeyPatch,
validated: list[str],
) -> None:
def validate(url: str) -> tuple[bool, str]:
async def validate(url: str) -> tuple[bool, str]:
validated.append(url)
if "169.254.169.254" in url:
return False, "blocked metadata address"
return True, ""
monkeypatch.setattr("nanobot.channels.slack.runtime.validate_url_target", validate)
monkeypatch.setattr("nanobot.channels.slack.runtime.async_validate_url_target", validate)
@pytest.mark.asyncio
+2 -2
View File
@@ -36,7 +36,7 @@ from nanobot.channels.base import BaseChannel
from nanobot.command.builtin import build_help_text
from nanobot.config.paths import get_media_dir
from nanobot.config.schema import Base
from nanobot.security.network import validate_url_target
from nanobot.security.network import async_validate_url_target
from nanobot.utils.helpers import split_message
from nanobot.utils.logging_bridge import redirect_lib_logging
@@ -956,7 +956,7 @@ class TelegramChannel(BaseChannel):
# Telegram Bot API accepts HTTP(S) URLs directly for media params.
if self._is_remote_media_url(media_path):
ok, error = validate_url_target(media_path)
ok, error = await async_validate_url_target(media_path)
if not ok:
raise ValueError(f"unsafe media URL: {error}")
await self._call_with_retry(
@@ -1488,7 +1488,14 @@ async def test_send_remote_media_url_after_security_validation(monkeypatch) -> N
MessageBus(),
)
_install_ready_app(channel)
monkeypatch.setattr("nanobot.channels.telegram.runtime.validate_url_target", lambda url: (True, ""))
async def allow_url(_url: str) -> tuple[bool, str]:
return True, ""
monkeypatch.setattr(
"nanobot.channels.telegram.runtime.async_validate_url_target",
allow_url,
)
await channel.send(
OutboundMessage(
@@ -1546,9 +1553,13 @@ async def test_send_blocks_unsafe_remote_media_url(monkeypatch) -> None:
MessageBus(),
)
_install_ready_app(channel)
async def deny_url(_url: str) -> tuple[bool, str]:
return False, "Blocked: example.com resolves to private/internal address 127.0.0.1"
monkeypatch.setattr(
"nanobot.channels.telegram.runtime.validate_url_target",
lambda url: (False, "Blocked: example.com resolves to private/internal address 127.0.0.1"),
"nanobot.channels.telegram.runtime.async_validate_url_target",
deny_url,
)
await channel.send(
+30 -6
View File
@@ -55,6 +55,7 @@ from nanobot.security.workspace_access import (
WORKSPACE_SCOPE_METADATA_KEY,
WorkspaceScopeError,
)
from nanobot.session.async_compat import call_session_manager
from nanobot.session.goal_state import goal_state_ws_blob
from nanobot.session.model_selection import model_preset_from_metadata
from nanobot.session.recovery import recovery_state_from_metadata
@@ -447,6 +448,26 @@ class WebSocketChannel(BaseChannel):
if sessions is None:
return {}
snapshot = sessions.read_session_metadata(f"websocket:{chat_id}")
return self._attached_model_fields_from_snapshot(chat_id, snapshot)
async def _attached_model_fields_async(self, chat_id: str) -> dict[str, Any]:
"""Build attach fields without blocking the gateway event loop."""
sessions = self.gateway.session_manager
if sessions is None:
return {}
snapshot = await call_session_manager(
sessions,
"read_session_metadata_async",
sessions.read_session_metadata,
f"websocket:{chat_id}",
)
return self._attached_model_fields_from_snapshot(chat_id, snapshot)
def _attached_model_fields_from_snapshot(
self,
chat_id: str,
snapshot: dict[str, Any] | None,
) -> dict[str, Any]:
raw_metadata = snapshot.get("metadata") if snapshot is not None else None
metadata = cast(dict[str, object], raw_metadata) if isinstance(raw_metadata, dict) else None
fields: dict[str, Any] = {}
@@ -510,13 +531,16 @@ class WebSocketChannel(BaseChannel):
fork_key: str,
) -> None:
"""Attach and hydrate a newly created WebUI chat fork."""
scope = self._workspaces.scope_for_session_key(fork_key)
scope = await asyncio.to_thread(
self._workspaces.scope_for_session_key,
fork_key,
)
self._attach(connection, fork_id)
await self._send_event(
connection,
"attached",
chat_id=fork_id,
**self._attached_model_fields(fork_id),
**await self._attached_model_fields_async(fork_id),
)
await self._send_event(
connection,
@@ -904,7 +928,7 @@ class WebSocketChannel(BaseChannel):
connection,
"attached",
chat_id=new_id,
**self._attached_model_fields(new_id),
**await self._attached_model_fields_async(new_id),
)
await self._send_event(
connection,
@@ -960,7 +984,7 @@ class WebSocketChannel(BaseChannel):
connection,
"attached",
chat_id=cid,
**self._attached_model_fields(cid),
**await self._attached_model_fields_async(cid),
)
await self._hydrate_after_subscribe(cid)
return
@@ -1263,7 +1287,7 @@ class WebSocketChannel(BaseChannel):
else False
),
)
self._workspaces.persist_scope(cid, scope)
await asyncio.to_thread(self._workspaces.persist_scope, cid, scope)
accepted = True
finally:
if not accepted and queued_owner is not None:
@@ -1553,7 +1577,7 @@ class WebSocketChannel(BaseChannel):
turn_id: str | None = None,
) -> Any | None:
try:
return resolver()
return await asyncio.to_thread(resolver)
except WorkspaceScopeError as exc:
await self._send_event(
connection,
@@ -4,6 +4,7 @@ import asyncio
import json
import random
import socket
import threading
import time
from contextlib import suppress
from pathlib import Path
@@ -2576,6 +2577,69 @@ async def test_webui_automations_route_lists_all_jobs_and_allows_user_actions(
await server_task
@pytest.mark.asyncio
async def test_webui_cron_update_rearms_started_service_on_owner_loop(
bus: MagicMock,
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
store_path = tmp_path / "cron" / "jobs.json"
cron = CronService(store_path, max_sleep_ms=60_000)
job = cron.add_job(
name="Before update",
schedule=CronSchedule(kind="every", every_ms=86_400_000),
message="Check the repo status",
session_key="websocket:abc",
origin_channel="websocket",
origin_chat_id="abc",
)
await cron.start()
owner_thread_id = threading.get_ident()
initial_timer = cron._timer_task
request_thread_ids: list[int] = []
arm_thread_ids: list[int] = []
timer_rearmed = asyncio.Event()
original_request_timer_rearm = cron._request_timer_rearm
original_arm_timer = cron._arm_timer
def tracked_request_timer_rearm() -> None:
request_thread_ids.append(threading.get_ident())
original_request_timer_rearm()
def tracked_arm_timer() -> None:
arm_thread_ids.append(threading.get_ident())
original_arm_timer()
timer_rearmed.set()
monkeypatch.setattr(cron, "_request_timer_rearm", tracked_request_timer_rearm)
monkeypatch.setattr(cron, "_arm_timer", tracked_arm_timer)
channel = _ch(bus, cron_service=cron, port=_free_port())
try:
response = await _webui_mutate(
channel,
"automation.update",
{"id": job.id, "values": {"name": "After update"}},
)
await asyncio.wait_for(timer_rearmed.wait(), timeout=1)
assert response.status_code == 200
assert request_thread_ids
assert all(thread_id != owner_thread_id for thread_id in request_thread_ids)
assert arm_thread_ids and set(arm_thread_ids) == {owner_thread_id}
assert cron._timer_task is not None
assert cron._timer_task is not initial_timer
assert not cron._timer_task.done()
stored = json.loads(store_path.read_text(encoding="utf-8"))
assert len(stored["jobs"]) == 1
assert stored["jobs"][0]["name"] == "After update"
finally:
cron.stop()
await asyncio.sleep(0)
@pytest.mark.asyncio
async def test_webui_automations_route_manages_local_triggers(
bus: MagicMock, tmp_path: Path
@@ -3769,3 +3833,77 @@ def test_bootstrap_secret_also_enforced_on_localhost(bus: MagicMock) -> None:
channel = _ch(bus, host="0.0.0.0", tokenIssueSecret="s3cret")
resp = channel.gateway.http._handle_bootstrap(_LOCAL, _NO_HEADERS)
assert resp.status_code == 401
@pytest.mark.asyncio
async def test_webui_skill_update_cancellation_waits_for_config_and_runtime_state(
bus: MagicMock,
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
from nanobot.webui import ws_http
skill_dir = tmp_path / "skills" / "cancel-safe-skill"
skill_dir.mkdir(parents=True)
(skill_dir / "SKILL.md").write_text(
"---\nname: cancel-safe-skill\ndescription: Cancellation test skill.\n---\n",
encoding="utf-8",
)
mutation_started = threading.Event()
release_mutation = threading.Event()
original_update = ws_http.set_webui_skill_enabled
update_calls = 0
def blocked_update(*args: Any, **kwargs: Any) -> dict[str, Any]:
nonlocal update_calls
update_calls += 1
mutation_started.set()
assert release_mutation.wait(timeout=1)
return original_update(*args, **kwargs)
monkeypatch.setattr(ws_http, "set_webui_skill_enabled", blocked_update)
channel = _ch(
bus,
session_manager=_seed_session(tmp_path),
workspace_path=tmp_path,
port=_free_port(),
)
runtime_states: list[set[str]] = []
channel.gateway.http.skill_state_action = runtime_states.append
task = asyncio.create_task(
_webui_mutate(
channel,
"skill.update",
{"name": "cancel-safe-skill", "enabled": False},
)
)
assert await asyncio.to_thread(mutation_started.wait, 1)
try:
task.cancel()
await asyncio.sleep(0)
assert not task.done()
assert runtime_states == []
finally:
release_mutation.set()
with pytest.raises(asyncio.CancelledError):
await asyncio.wait_for(task, timeout=1)
assert update_calls == 1
assert "cancel-safe-skill" in channel.gateway.http.disabled_skills
assert runtime_states == [{"cancel-safe-skill"}]
saved = load_config(channel.gateway.settings.config.path)
assert "cancel-safe-skill" in saved.agents.defaults.disabled_skills
settled_state = (
update_calls,
set(channel.gateway.http.disabled_skills),
list(runtime_states),
)
await asyncio.sleep(0.05)
assert (
update_calls,
set(channel.gateway.http.disabled_skills),
runtime_states,
) == settled_state
+75 -12
View File
@@ -34,6 +34,7 @@ from nanobot.config.paths import is_default_workspace
from nanobot.config.schema import Config
from nanobot.gateway.runtime import GatewayInstance
from nanobot.security.network import is_loopback_host
from nanobot.session.async_compat import call_session_manager as _call_session_manager
from nanobot.session.keys import UNIFIED_SESSION_KEY, last_channel_from_metadata
from nanobot.utils.evaluator import evaluate_response, resolve_evaluator_prompt
from nanobot.utils.helpers import sync_workspace_templates
@@ -45,6 +46,30 @@ __all__ = ["_run_gateway"]
console = Console()
_EVENT_LOOP_LAG_INTERVAL_S = 0.5
_EVENT_LOOP_LAG_WARNING_S = 0.25
async def _monitor_event_loop_lag(
*,
interval_s: float = _EVENT_LOOP_LAG_INTERVAL_S,
warning_threshold_s: float = _EVENT_LOOP_LAG_WARNING_S,
log: Any | None = None,
) -> None:
"""Log scheduler drift so gateway-wide stalls have direct evidence."""
loop = asyncio.get_running_loop()
lag_log = log or logger
while True:
expected = loop.time() + interval_s
await asyncio.sleep(interval_s)
lag_s = max(0.0, loop.time() - expected)
if lag_s >= warning_threshold_s:
lag_log.warning(
"event loop lag operation=gateway duration_ms={} interval_ms={}",
int(lag_s * 1000),
int(interval_s * 1000),
)
def _http_endpoint_responding(url: str, *, timeout_s: float = 0.25) -> bool:
"""Return whether an HTTP endpoint responds, including with an auth error."""
@@ -493,12 +518,22 @@ def _run_gateway(
and hasattr(session_manager, "save")
):
key = session_key or _channel_session_key(msg.channel, msg.chat_id)
session = session_manager.get_or_create(key)
session = await _call_session_manager(
session_manager,
"get_or_create_async",
session_manager.get_or_create,
key,
)
extra: dict[str, Any] = {"_channel_delivery": True}
if msg.media:
extra["media"] = list(msg.media)
session.add_message("assistant", msg.content, **extra)
session_manager.save(session)
await _call_session_manager(
session_manager,
"save_async",
session_manager.save,
session,
)
await bus.publish_outbound(msg)
message_tool = agent.tools.get("message")
@@ -568,14 +603,14 @@ def _run_gateway(
if sha:
logger.info("Dream commit: {}", sha)
store.compact_history()
prune_dream_sessions(agent.sessions)
await asyncio.to_thread(prune_dream_sessions, agent.sessions)
return None
# Heartbeat is a system job that checks HEARTBEAT.md for active tasks.
if job.name == "heartbeat":
heartbeat_file = config.workspace_path / "HEARTBEAT.md"
try:
content = heartbeat_file.read_text(encoding="utf-8")
content = await asyncio.to_thread(heartbeat_file.read_text, encoding="utf-8")
except OSError:
logger.debug("Heartbeat: HEARTBEAT.md missing")
return None
@@ -583,7 +618,7 @@ def _run_gateway(
logger.debug("Heartbeat: HEARTBEAT.md has no active tasks")
return None
channel, chat_id = _pick_heartbeat_target()
channel, chat_id = await _pick_heartbeat_target()
if channel == "cli":
return None
@@ -611,9 +646,19 @@ def _run_gateway(
message_tool.reset_suppress_delivery(suppress_token)
# Keep a small tail of heartbeat history so the loop stays bounded.
session = agent.sessions.get_or_create("heartbeat")
session = await _call_session_manager(
agent.sessions,
"get_or_create_async",
agent.sessions.get_or_create,
"heartbeat",
)
session.retain_recent_legal_suffix(hb_cfg.keep_recent_messages)
agent.sessions.save(session)
await _call_session_manager(
agent.sessions,
"save_async",
agent.sessions.save,
session,
)
if not resp or not resp.content:
return
@@ -690,17 +735,27 @@ def _run_gateway(
config_path=Path(config_path),
)
def _pick_heartbeat_target() -> tuple[str, str]:
async def _pick_heartbeat_target() -> tuple[str, str]:
"""Pick a routable channel/chat target for heartbeat-triggered messages."""
sidebar_state = read_webui_sidebar_state()
sidebar_state = await asyncio.to_thread(read_webui_sidebar_state)
unified_metadata = None
if config.agents.defaults.unified_session:
record = session_manager.read_session_metadata(UNIFIED_SESSION_KEY)
record = await _call_session_manager(
session_manager,
"read_session_metadata_async",
session_manager.read_session_metadata,
UNIFIED_SESSION_KEY,
)
if isinstance(record, dict) and isinstance(record.get("metadata"), dict):
unified_metadata = record["metadata"]
sessions = await _call_session_manager(
session_manager,
"list_sessions_async",
session_manager.list_sessions,
)
return _pick_heartbeat_target_from_sessions(
enabled_channels=channels.enabled_channels,
sessions=session_manager.list_sessions(),
sessions=sessions,
archived_keys=sidebar_state.get("archived_keys", []),
unified_session_metadata=unified_metadata,
)
@@ -907,6 +962,10 @@ def _run_gateway(
_monitor_local_clients(),
name="nanobot-gateway-client-monitor",
),
asyncio.create_task(
_monitor_event_loop_lag(),
name="nanobot-event-loop-lag-monitor",
),
]
if health_server_enabled:
tasks.append(asyncio.create_task(
@@ -974,7 +1033,11 @@ def _run_gateway(
# Flush all cached sessions to durable storage before exit.
# This prevents data loss on filesystems with write-back
# caching (rclone VFS, NFS, FUSE mounts, etc.).
flushed = agent.sessions.flush_all()
flushed = await _call_session_manager(
agent.sessions,
"flush_all_async",
agent.sessions.flush_all,
)
if flushed:
logger.info("Shutdown: flushed {} session(s) to disk", flushed)
finally:
+70 -17
View File
@@ -3,6 +3,7 @@
from __future__ import annotations
import asyncio
import inspect
import os
import subprocess
import sys
@@ -14,6 +15,8 @@ from typing import TYPE_CHECKING, Any, Literal, cast
from nanobot import __version__
from nanobot.bus.events import INBOUND_META_USER_SHELL, OutboundMessage
from nanobot.command.router import CommandContext, CommandRouter, normalize_command_text
from nanobot.session.async_compat import call_session_manager
from nanobot.utils.cancellation import shield_and_drain
from nanobot.utils.helpers import build_status_content
from nanobot.utils.restart import set_restart_notice_to_env
from nanobot.utils.workspace_prompts import initialize_workspace_prompt
@@ -22,6 +25,7 @@ if TYPE_CHECKING:
from nanobot.agent.loop import AgentLoop
from nanobot.session.manager import Session
from nanobot.utils.gitstore import CommitInfo
from nanobot.utils.llm_runtime import LLMRuntime
# WebUI protocol contract for how a slash command participates in turn state:
# - side_channel: returns control text without starting or ending an agent turn.
@@ -201,6 +205,52 @@ def builtin_command_starts_agent_turn(text: str) -> bool:
return spec.lifecycle == "agent_turn_with_args" and bool(args.strip())
def _has_native_coroutine_method(target: object, name: str) -> bool:
"""Check the real target class without trusting dynamic mock attributes."""
method = inspect.getattr_static(type(target), name, None)
return inspect.iscoroutinefunction(method)
async def _get_or_create_session(loop: AgentLoop, key: str) -> Session:
sessions = loop.sessions
return await call_session_manager(
sessions,
"get_or_create_async",
sessions.get_or_create,
key,
)
async def _save_session(loop: AgentLoop, session: Session) -> None:
sessions = loop.sessions
await call_session_manager(
sessions,
"save_async",
sessions.save,
session,
)
async def _runtime_for_session(loop: AgentLoop, session: Session) -> LLMRuntime:
if _has_native_coroutine_method(loop, "runtime_for_session_async"):
return await loop.runtime_for_session_async(session)
return await shield_and_drain(
asyncio.to_thread(loop.runtime_for_session, session)
)
async def _set_session_model_preset(
loop: AgentLoop,
session_key: str,
name: str,
) -> LLMRuntime:
if _has_native_coroutine_method(loop, "set_session_model_preset_async"):
return await loop.set_session_model_preset_async(session_key, name)
return await shield_and_drain(
asyncio.to_thread(loop.set_session_model_preset, session_key, name)
)
async def cmd_stop(ctx: CommandContext) -> OutboundMessage:
"""Cancel all active tasks and subagents for the session."""
loop = ctx.loop
@@ -257,8 +307,8 @@ async def cmd_restart(ctx: CommandContext) -> OutboundMessage:
async def cmd_status(ctx: CommandContext) -> OutboundMessage:
"""Build an outbound status message for a session."""
loop = ctx.loop
session = ctx.session or loop.sessions.get_or_create(ctx.key)
runtime = ctx.runtime or loop.runtime_for_session(session)
session = ctx.session or await _get_or_create_session(loop, ctx.key)
runtime = ctx.runtime or await _runtime_for_session(loop, session)
ctx_est = 0
with suppress(Exception):
ctx_est, _ = loop.consolidator.estimate_session_prompt_tokens(
@@ -306,29 +356,32 @@ async def cmd_new(ctx: CommandContext) -> OutboundMessage:
loop = ctx.loop
await loop._cancel_active_tasks(ctx.key) # pyright: ignore[reportPrivateUsage]
loop.discard_session_file_state(ctx.key)
session = ctx.session or loop.sessions.get_or_create(ctx.key)
session = ctx.session or await _get_or_create_session(loop, ctx.key)
snapshot = list(session.messages)
archive_snapshot = None
runtime = None
if session.last_consolidated < len(snapshot):
runtime = ctx.runtime or loop.runtime_for_session(session)
runtime = ctx.runtime or await _runtime_for_session(loop, session)
archive_snapshot = replace(
session,
messages=snapshot,
metadata=dict(session.metadata),
provider_state=None,
)
session.clear()
loop.sessions.save(session)
loop.sessions.invalidate(session.key)
if archive_snapshot is not None and runtime is not None:
loop.schedule_background(
loop.consolidator.archive_session( # pyright: ignore[reportUnknownMemberType]
archive_snapshot,
archive_end=len(snapshot),
runtime=runtime,
async def reset_and_schedule_archive() -> None:
session.clear()
await _save_session(loop, session)
loop.sessions.invalidate(session.key)
if archive_snapshot is not None and runtime is not None:
loop.schedule_background(
loop.consolidator.archive_session( # pyright: ignore[reportUnknownMemberType]
archive_snapshot,
archive_end=len(snapshot),
runtime=runtime,
)
)
)
await shield_and_drain(reset_and_schedule_archive())
return OutboundMessage(
channel=ctx.msg.channel, chat_id=ctx.msg.chat_id,
content="New session started.",
@@ -377,7 +430,7 @@ async def cmd_model(ctx: CommandContext) -> OutboundMessage:
metadata = {**dict(ctx.msg.metadata or {}), "render_as": "text"}
if not args:
session = ctx.session or loop.sessions.get_or_create(ctx.key)
session = ctx.session or await _get_or_create_session(loop, ctx.key)
return OutboundMessage(
channel=ctx.msg.channel,
chat_id=ctx.msg.chat_id,
@@ -387,7 +440,7 @@ async def cmd_model(ctx: CommandContext) -> OutboundMessage:
name = args
try:
runtime = loop.set_session_model_preset(ctx.key, name)
runtime = await _set_session_model_preset(loop, ctx.key, name)
except (KeyError, ValueError) as exc:
names = _model_preset_names(loop)
return OutboundMessage(
@@ -848,7 +901,7 @@ async def cmd_history(ctx: CommandContext) -> OutboundMessage:
metadata=dict(ctx.msg.metadata or {}),
)
session = ctx.session or ctx.loop.sessions.get_or_create(ctx.key)
session = ctx.session or await _get_or_create_session(ctx.loop, ctx.key)
history = session.get_history(max_messages=0, include_runtime_context=False)
visible = [_format_history_message(m) for m in history]
visible = [m for m in visible if m is not None]
+10
View File
@@ -8,6 +8,7 @@ import time
import uuid
from typing import TYPE_CHECKING, Any, Protocol
from nanobot.agent.automation_turns import AutomationTurnAcceptedCancellation
from nanobot.agent.tools.cron import CronTool
from nanobot.bus.events import InboundMessage, OutboundMessage
from nanobot.cron.session_delivery import origin_delivery_context
@@ -127,6 +128,15 @@ async def run_bound_cron_job(
session_key_override=session_key,
)
)
except AutomationTurnAcceptedCancellation:
cron.write_run_record(
run_id,
{
**run_record_base,
"status": "accepted",
},
)
raise
except (Exception, asyncio.CancelledError) as exc:
error_text = str(exc) or exc.__class__.__name__
cron.write_run_record(
+230 -87
View File
@@ -11,11 +11,12 @@ from dataclasses import asdict
from datetime import datetime
from pathlib import Path
from types import EllipsisType
from typing import Any, Callable, Coroutine, Literal
from typing import Any, Callable, Coroutine, Literal, TypeVar
from filelock import FileLock
from loguru import logger
from nanobot.agent.automation_turns import AutomationTurnAcceptedCancellation
from nanobot.cron.session_turns import is_bound_cron_job
from nanobot.cron.types import (
CronJob,
@@ -25,10 +26,14 @@ from nanobot.cron.types import (
CronSchedule,
CronStore,
)
from nanobot.utils.cancellation import shield_and_drain
from nanobot.utils.run_records import (
write_run_record as write_automation_run_record,
)
_FILE_LOCK_TIMEOUT_SECONDS = 5
_T = TypeVar("_T")
class CronJobSkippedError(Exception):
"""Raised by cron callbacks when a job was intentionally skipped."""
@@ -164,10 +169,16 @@ class CronService:
self.store_path = store_path
self._action_path = store_path.parent / "action.jsonl"
self._run_records_dir = store_path.parent / "runs"
self._lock = FileLock(str(self._action_path.parent) + ".lock")
self._lock = FileLock(
str(self._action_path.parent) + ".lock",
timeout=_FILE_LOCK_TIMEOUT_SECONDS,
)
self.on_job = on_job
self._store: CronStore | None = None
self._timer_task: asyncio.Task[None] | None = None
self._operation_lock = asyncio.Lock()
self._claimed_job_ids: set[str] = set()
self._event_loop: asyncio.AbstractEventLoop | None = None
self._running = False
self._active_executions = 0
self._store_dirty = False
@@ -451,25 +462,58 @@ class CronService:
"""Write an internal audit record for one cron execution."""
write_automation_run_record(self._run_records_dir, run_id, record)
async def start(self) -> None:
"""Start the cron service."""
self._running = True
loaded = self._load_store()
if loaded is None:
# Store file existed but was corrupt and has been preserved with
# a ``.corrupt-<ts>`` suffix. Bail out instead of starting with
# an empty store; that would call ``_save_store`` and overwrite
# the now-renamed (but still recoverable) data with [].
self._running = False
raise RuntimeError(
f"cron store at {self.store_path} is corrupt and was preserved; "
"refusing to start with an empty job list. "
"Inspect the .corrupt-<ts> backup and restore manually."
async def run_sync(
self,
operation: Callable[..., _T],
/,
*args: Any,
**kwargs: Any,
) -> _T:
"""Serialize a complete cron transaction in a worker thread.
A running thread cannot be cancelled safely. Keep the transaction lock
until it exits so cancellation is never reported while that worker can
still mutate cron state behind a later operation.
"""
async with self._operation_lock:
return await shield_and_drain(
asyncio.to_thread(operation, *args, **kwargs)
)
self._recompute_next_runs()
self._save_store()
self._arm_timer()
logger.info("Cron service started with {} jobs", len(self._store.jobs if self._store else []))
async def start(self) -> None:
"""Start the cron service and settle accepted work before cancellation."""
async def settle_start() -> None:
self._event_loop = asyncio.get_running_loop()
self._running = True
try:
async with self._operation_lock:
loaded = await asyncio.to_thread(self._load_store)
if loaded is None:
# Store file existed but was corrupt and has been preserved with
# a ``.corrupt-<ts>`` suffix. Bail out instead of starting with
# an empty store; that would call ``_save_store`` and overwrite
# the now-renamed (but still recoverable) data with [].
raise RuntimeError(
f"cron store at {self.store_path} is corrupt and was preserved; "
"refusing to start with an empty job list. "
"Inspect the .corrupt-<ts> backup and restore manually."
)
self._recompute_next_runs()
await asyncio.to_thread(self._save_store)
self._arm_timer()
logger.info(
"Cron service started with {} jobs",
len(self._store.jobs if self._store else []),
)
except BaseException:
# A failed start must not retain ownership without a timer. Caller
# cancellation is shielded until this composite either reaches the
# fully started state above or rolls back here.
self.stop()
raise
await shield_and_drain(settle_start())
def stop(self) -> None:
"""Stop the cron service."""
@@ -497,8 +541,22 @@ class CronService:
if j.enabled and j.state.next_run_at_ms]
return min(times) if times else None
def _request_timer_rearm(self) -> None:
"""Re-arm on the owning event loop, including from persistence workers."""
if not self._running:
return
loop = self._event_loop
try:
current_loop = asyncio.get_running_loop()
except RuntimeError:
current_loop = None
if current_loop is loop:
self._arm_timer()
elif loop is not None and loop.is_running():
loop.call_soon_threadsafe(self._arm_timer)
def _arm_timer(self) -> None:
"""Schedule the next timer tick."""
"""Schedule the next timer tick on the owning event loop."""
if self._timer_task:
self._timer_task.cancel()
@@ -520,7 +578,7 @@ class CronService:
self._timer_task = asyncio.create_task(tick())
async def _on_timer(self) -> None:
"""Handle timer tick - run due jobs."""
"""Run due jobs while keeping persistence transactions serialized."""
reload_store = self._active_executions == 0
self._active_executions += 1
try:
@@ -528,11 +586,17 @@ class CronService:
# to persist their advanced schedule. Persist that exact snapshot
# before reloading or executing anything else; otherwise the older
# disk state can replay the same job.
if self._store_dirty:
self._save_store()
return
async with self._operation_lock:
if self._store_dirty:
await shield_and_drain(asyncio.to_thread(self._save_store))
return
store = self._load_store(reload_during_execution=reload_store)
store = await shield_and_drain(
asyncio.to_thread(
self._load_store,
reload_during_execution=reload_store,
)
)
# If a hot reload found a corrupt store on disk, ``self._store``
# may still hold the previous, known-good in-memory snapshot.
if store is None:
@@ -547,7 +611,8 @@ class CronService:
for job in due_jobs:
await self._execute_job(job)
self._save_store()
async with self._operation_lock:
await shield_and_drain(asyncio.to_thread(self._save_store))
except Exception:
# A load/persist failure must not kill the scheduler: keep the
# in-memory store and retry on the next tick. This mirrors the
@@ -564,58 +629,124 @@ class CronService:
# single bad tick cannot silently stop all future jobs.
self._arm_timer()
async def _execute_job(self, job: CronJob) -> None:
"""Execute a single job."""
async def _claim_job(self, job_id: str) -> bool:
"""Claim one job without serializing callbacks for different jobs."""
async with self._operation_lock:
if job_id in self._claimed_job_ids:
return False
self._claimed_job_ids.add(job_id)
return True
async def _release_job_claim(self, job_id: str) -> None:
async with self._operation_lock:
self._claimed_job_ids.discard(job_id)
async def _settle_job_execution(
self,
job: CronJob,
*,
start_ms: int,
status: Literal["ok", "error", "skipped"],
error: str | None,
persist: bool = False,
) -> None:
end_ms = _now_ms()
async with self._operation_lock:
job.state.last_status = status
job.state.last_error = error
job.state.last_run_at_ms = start_ms
job.updated_at_ms = end_ms
job.state.run_history.append(CronRunRecord(
run_at_ms=start_ms,
status=status,
duration_ms=end_ms - start_ms,
error=error,
))
job.state.run_history = job.state.run_history[-self._MAX_RUN_HISTORY:]
if job.schedule.kind == "at":
if job.delete_after_run:
store = await shield_and_drain(
asyncio.to_thread(self._require_store)
)
store.jobs = [item for item in store.jobs if item.id != job.id]
else:
job.enabled = False
job.state.next_run_at_ms = None
else:
job.state.next_run_at_ms = _compute_next_run(job.schedule, _now_ms())
if persist:
await shield_and_drain(asyncio.to_thread(self._save_store))
@staticmethod
async def _drain_settlement_on_cancellation(settlement: asyncio.Task[None]) -> None:
"""Finish a short durable settlement despite repeated cancellation."""
while not settlement.done():
try:
await asyncio.shield(settlement)
except asyncio.CancelledError:
continue
settlement.result()
async def _execute_job(self, job: CronJob) -> bool:
"""Execute a claimed job and serialize its in-memory settlement."""
if not await self._claim_job(job.id):
logger.info("Cron: job '{}' ({}) is already running", job.name, job.id)
return False
start_ms = _now_ms()
logger.info("Cron: executing job '{}' ({})", job.name, job.id)
status: Literal["ok", "error", "skipped"]
error: str | None
accepted_cancellation: AutomationTurnAcceptedCancellation | None = None
try:
if self.on_job:
await self.on_job(job)
try:
if self.on_job:
await self.on_job(job)
status = "ok"
error = None
logger.info("Cron: job '{}' completed", job.name)
except AutomationTurnAcceptedCancellation as exc:
# The agent owns this turn now. Advance and persist the schedule
# before allowing shutdown cancellation to unwind the timer.
status = "ok"
error = None
accepted_cancellation = exc
logger.info("Cron: job '{}' was accepted before cancellation", job.name)
except CronJobSkippedError as exc:
status = "skipped"
error = str(exc) or None
logger.warning("Cron: job '{}' skipped: {}", job.name, error or "")
except asyncio.CancelledError as exc:
current = asyncio.current_task()
if current is not None and current.cancelling():
raise
status = "error"
error = str(exc) or exc.__class__.__name__
logger.exception("Cron: job '{}' was cancelled", job.name)
except Exception as exc:
status = "error"
error = str(exc)
logger.exception("Cron: job '{}' failed", job.name)
job.state.last_status = "ok"
job.state.last_error = None
logger.info("Cron: job '{}' completed", job.name)
except CronJobSkippedError as e:
job.state.last_status = "skipped"
job.state.last_error = str(e) or None
logger.warning("Cron: job '{}' skipped: {}", job.name, job.state.last_error or "")
except asyncio.CancelledError as e:
current = asyncio.current_task()
if current is not None and current.cancelling():
raise
job.state.last_status = "error"
job.state.last_error = str(e) or e.__class__.__name__
logger.exception("Cron: job '{}' was cancelled", job.name)
except Exception as e:
job.state.last_status = "error"
job.state.last_error = str(e)
logger.exception("Cron: job '{}' failed", job.name)
end_ms = _now_ms()
job.state.last_run_at_ms = start_ms
job.updated_at_ms = end_ms
job.state.run_history.append(CronRunRecord(
run_at_ms=start_ms,
status=job.state.last_status,
duration_ms=end_ms - start_ms,
error=job.state.last_error,
))
job.state.run_history = job.state.run_history[-self._MAX_RUN_HISTORY:]
# Handle one-shot jobs
if job.schedule.kind == "at":
if job.delete_after_run:
store = self._require_store()
store.jobs = [item for item in store.jobs if item.id != job.id]
else:
job.enabled = False
job.state.next_run_at_ms = None
else:
# Compute next run
job.state.next_run_at_ms = _compute_next_run(job.schedule, _now_ms())
settlement = asyncio.create_task(
self._settle_job_execution(
job,
start_ms=start_ms,
status=status,
error=error,
persist=accepted_cancellation is not None,
)
)
if accepted_cancellation is not None:
await self._drain_settlement_on_cancellation(settlement)
raise accepted_cancellation
await settlement
return True
finally:
await self._release_job_claim(job.id)
def _append_action(
self,
@@ -697,7 +828,7 @@ class CronService:
store = self._require_store()
store.jobs.append(job)
self._save_store()
self._arm_timer()
self._request_timer_rearm()
else:
self._append_action("add", asdict(job))
@@ -714,7 +845,7 @@ class CronService:
store.jobs = [j for j in store.jobs if j.id != job.id]
store.jobs.append(job)
self._save_store()
self._arm_timer()
self._request_timer_rearm()
logger.info("Cron: registered system job '{}' ({})", job.name, job.id)
return job
@@ -726,7 +857,7 @@ class CronService:
removed = len(store.jobs) < before
if removed:
self._save_store()
self._arm_timer()
self._request_timer_rearm()
logger.info("Cron: removed system job {}", job_id)
return removed
@@ -747,7 +878,7 @@ class CronService:
if removed:
if self._should_persist_store():
self._save_store()
self._arm_timer()
self._request_timer_rearm()
else:
self._append_action("del", {"job_id": job_id})
logger.info("Cron: removed job {}", job_id)
@@ -769,7 +900,7 @@ class CronService:
job.state.next_run_at_ms = None
if self._should_persist_store():
self._save_store()
self._arm_timer()
self._request_timer_rearm()
else:
self._append_action("update", asdict(job))
return job
@@ -825,7 +956,7 @@ class CronService:
if self._should_persist_store():
self._save_store()
self._arm_timer()
self._request_timer_rearm()
else:
self._append_action("update", asdict(job))
@@ -840,19 +971,31 @@ class CronService:
# A manual run is another side-effecting entrypoint. Do not start
# it while the result of a previous timer execution is still only
# in memory.
if self._store_dirty:
self._save_store()
store = self._require_store(reload_during_execution=reload_store)
async with self._operation_lock:
if self._store_dirty:
await shield_and_drain(asyncio.to_thread(self._save_store))
store = await shield_and_drain(
asyncio.to_thread(
self._require_store,
reload_during_execution=reload_store,
)
)
for job in store.jobs:
if job.id == job_id:
if self._is_unbound_agent_job(job):
self._enforce_agent_binding(job)
self._save_store()
async with self._operation_lock:
self._enforce_agent_binding(job)
await shield_and_drain(
asyncio.to_thread(self._save_store)
)
return False
if not force and not job.enabled:
return False
await self._execute_job(job)
self._save_store()
executed = await self._execute_job(job)
if not executed:
return False
async with self._operation_lock:
await shield_and_drain(asyncio.to_thread(self._save_store))
return True
return False
finally:
+2 -2
View File
@@ -20,7 +20,7 @@ from nanobot.providers.registry import find_by_name
from nanobot.security.network import (
PinnedDNSAsyncTransport,
UnsafeURLRequestError,
resolve_url_target,
async_resolve_url_target,
)
from nanobot.utils.helpers import detect_image_mime
@@ -174,7 +174,7 @@ async def _download_image_data_url(
current_url = url
for _ in range(_IMAGE_DOWNLOAD_MAX_REDIRECTS + 1):
if proxy:
ok, error, _ = resolve_url_target(
ok, error, _ = await async_resolve_url_target(
current_url,
trust_remote_dns=True,
)
+7 -5
View File
@@ -210,18 +210,20 @@ class RuntimeClient:
async def compact_session(self, session_key: str) -> SessionSnapshot:
"""Run token consolidation for one session."""
session = self._loop.sessions.get_or_create(session_key)
runtime = self._loop.runtime_for_session(session)
session = await self._loop.sessions.get_or_create_async(session_key)
runtime = await self._loop.runtime_for_session_async(session)
await self._loop.consolidator.maybe_consolidate_by_tokens(
session,
runtime=runtime,
)
return snapshot_from_session(self._loop.sessions.get_or_create(session_key))
return snapshot_from_session(
await self._loop.sessions.get_or_create_async(session_key)
)
async def compact_idle_session(self, session_key: str, *, max_suffix: int = 8) -> str | None:
"""Run idle-session compaction for one session and return the summary."""
session = self._loop.sessions.get_or_create(session_key)
runtime = self._loop.runtime_for_session(session)
session = await self._loop.sessions.get_or_create_async(session_key)
runtime = await self._loop.runtime_for_session_async(session)
return await self._loop.consolidator.compact_idle_session(
session_key,
runtime=runtime,
+144 -42
View File
@@ -29,6 +29,7 @@ _BLOCKED_NETWORKS = [
_URL_RE = re.compile(r"https?://[^\s\"'`;|<>]+", re.IGNORECASE)
_allowed_networks: list[ipaddress.IPv4Network | ipaddress.IPv6Network] = []
_DNS_RESOLUTION_TIMEOUT_SECONDS = 5.0
def is_loopback_host(host: str) -> bool:
@@ -75,6 +76,63 @@ def _is_private(addr: ipaddress.IPv4Address | ipaddress.IPv6Address) -> bool:
return any(normalized in net for net in _BLOCKED_NETWORKS)
def _parse_url_hostname(url: str) -> tuple[str | None, str | None]:
try:
parsed = urlparse(url)
except Exception as exc:
return None, str(exc)
if parsed.scheme not in ("http", "https"):
return None, f"Only http/https allowed, got '{parsed.scheme or 'none'}'"
if not parsed.netloc:
return None, "Missing domain"
if not parsed.hostname:
return None, "Missing hostname"
return parsed.hostname, None
def _unresolved_target_result(
hostname: str,
*,
trust_remote_dns: bool,
) -> tuple[bool, str, tuple[str, ...]]:
if not trust_remote_dns:
return False, f"Cannot resolve hostname: {hostname}", ()
normalized_hostname = hostname.rstrip(".").lower()
if normalized_hostname == "localhost" or normalized_hostname.endswith(".localhost"):
return False, f"Blocked local/internal hostname: {hostname}", ()
try:
literal_addr = ipaddress.ip_address(normalized_hostname)
except ValueError:
return True, "", ()
if _is_private(literal_addr):
return False, f"Blocked private/internal address: {literal_addr}", ()
return True, "", (str(_normalize_addr(literal_addr)),)
def _resolved_target_result(
hostname: str,
infos: list[Any],
*,
allow_loopback: bool,
) -> tuple[bool, str, tuple[str, ...]]:
addrs: list[ipaddress.IPv4Address | ipaddress.IPv6Address] = []
for info in infos:
try:
addr = ipaddress.ip_address(info[4][0])
except (IndexError, TypeError, ValueError):
continue
addrs.append(addr)
if allow_loopback and _is_allowed_loopback_target(hostname, addrs):
return True, "", tuple(dict.fromkeys(str(_normalize_addr(addr)) for addr in addrs))
for addr in addrs:
if _is_private(addr):
return False, f"Blocked: {hostname} resolves to private/internal address {addr}", ()
return True, "", tuple(dict.fromkeys(str(_normalize_addr(addr)) for addr in addrs))
def resolve_url_target(
url: str,
*,
@@ -97,52 +155,43 @@ def resolve_url_target(
resolved_ips contains the public IPs that were validated for this URL, or
is empty when an unresolved hostname is delegated to a trusted proxy.
"""
try:
p = urlparse(url)
except Exception as e:
return False, str(e), ()
if p.scheme not in ("http", "https"):
return False, f"Only http/https allowed, got '{p.scheme or 'none'}'", ()
if not p.netloc:
return False, "Missing domain", ()
hostname = p.hostname
if not hostname:
return False, "Missing hostname", ()
hostname, error = _parse_url_hostname(url)
if hostname is None:
return False, error or "Missing hostname", ()
try:
infos = socket.getaddrinfo(hostname, None, socket.AF_UNSPEC, socket.SOCK_STREAM)
except socket.gaierror:
if not trust_remote_dns:
return False, f"Cannot resolve hostname: {hostname}", ()
return _unresolved_target_result(hostname, trust_remote_dns=trust_remote_dns)
return _resolved_target_result(hostname, infos, allow_loopback=allow_loopback)
normalized_hostname = hostname.rstrip(".").lower()
if normalized_hostname == "localhost" or normalized_hostname.endswith(".localhost"):
return False, f"Blocked local/internal hostname: {hostname}", ()
try:
literal_addr = ipaddress.ip_address(normalized_hostname)
except ValueError:
return True, "", ()
if _is_private(literal_addr):
return False, f"Blocked private/internal address: {literal_addr}", ()
return True, "", (str(_normalize_addr(literal_addr)),)
addrs: list[ipaddress.IPv4Address | ipaddress.IPv6Address] = []
for info in infos:
try:
addr = ipaddress.ip_address(info[4][0])
except ValueError:
continue
addrs.append(addr)
if allow_loopback and _is_allowed_loopback_target(hostname, addrs):
return True, "", tuple(dict.fromkeys(str(_normalize_addr(addr)) for addr in addrs))
for addr in addrs:
if _is_private(addr):
return False, f"Blocked: {hostname} resolves to private/internal address {addr}", ()
return True, "", tuple(dict.fromkeys(str(_normalize_addr(addr)) for addr in addrs))
async def async_resolve_url_target(
url: str,
*,
allow_loopback: bool = False,
trust_remote_dns: bool = False,
timeout_s: float = _DNS_RESOLUTION_TIMEOUT_SECONDS,
) -> tuple[bool, str, tuple[str, ...]]:
"""Resolve and validate an HTTP target without blocking the event loop."""
hostname, error = _parse_url_hostname(url)
if hostname is None:
return False, error or "Missing hostname", ()
loop = asyncio.get_running_loop()
try:
infos = await asyncio.wait_for(
loop.getaddrinfo(
hostname,
None,
family=socket.AF_UNSPEC,
type=socket.SOCK_STREAM,
),
timeout=timeout_s,
)
except asyncio.TimeoutError:
return False, f"Timed out resolving hostname: {hostname}", ()
except socket.gaierror:
return _unresolved_target_result(hostname, trust_remote_dns=trust_remote_dns)
return _resolved_target_result(hostname, infos, allow_loopback=allow_loopback)
def validate_url_target(url: str, *, allow_loopback: bool = False) -> tuple[bool, str]:
@@ -151,6 +200,16 @@ def validate_url_target(url: str, *, allow_loopback: bool = False) -> tuple[bool
return ok, error
async def async_validate_url_target(
url: str,
*,
allow_loopback: bool = False,
) -> tuple[bool, str]:
"""Validate a URL using the event loop's asynchronous resolver."""
ok, error, _ = await async_resolve_url_target(url, allow_loopback=allow_loopback)
return ok, error
def env_proxy_applies_to_url(url: str) -> bool:
"""Return True when process proxy settings would proxy this URL."""
try:
@@ -277,7 +336,10 @@ class PinnedDNSAsyncTransport(httpx.AsyncBaseTransport):
async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
url = str(request.url)
ok, error, resolved_ips = resolve_url_target(url, allow_loopback=self._allow_loopback)
ok, error, resolved_ips = await async_resolve_url_target(
url,
allow_loopback=self._allow_loopback,
)
if not ok:
raise UnsafeURLRequestError(error, request=request)
async with self._resolver_lock:
@@ -320,6 +382,46 @@ def validate_resolved_url(url: str) -> tuple[bool, str]:
return True, ""
async def async_validate_resolved_url(url: str) -> tuple[bool, str]:
"""Validate a redirect target without blocking on domain resolution."""
try:
parsed = urlparse(url)
except Exception:
return True, ""
hostname = parsed.hostname
if not hostname:
return True, ""
try:
addr = ipaddress.ip_address(hostname)
except ValueError:
loop = asyncio.get_running_loop()
try:
infos = await asyncio.wait_for(
loop.getaddrinfo(
hostname,
None,
family=socket.AF_UNSPEC,
type=socket.SOCK_STREAM,
),
timeout=_DNS_RESOLUTION_TIMEOUT_SECONDS,
)
except asyncio.TimeoutError:
return False, f"Timed out resolving redirect hostname: {hostname}"
except socket.gaierror:
return True, ""
for info in infos:
try:
addr = ipaddress.ip_address(info[4][0])
except (IndexError, TypeError, ValueError):
continue
if _is_private(addr):
return False, f"Redirect target {hostname} resolves to private address {addr}"
return True, ""
if _is_private(addr):
return False, f"Redirect target is a private address: {addr}"
return True, ""
def contains_internal_url(command: str, *, allow_loopback: bool = False) -> bool:
"""Return True if the command string contains a URL targeting an internal/private address."""
for m in _URL_RE.finditer(command):
+29
View File
@@ -0,0 +1,29 @@
"""Compatibility bridge for asynchronous SessionManager operations."""
import asyncio
import inspect
from collections.abc import Awaitable, Callable
from typing import Any, TypeVar, cast
from nanobot.utils.cancellation import shield_and_drain
_SessionResult = TypeVar("_SessionResult")
async def call_session_manager(
manager: object,
async_method_name: str,
sync_method: Callable[..., _SessionResult],
/,
*args: Any,
**kwargs: Any,
) -> _SessionResult:
"""Prefer a class-declared coroutine, or offload the established sync contract."""
class_async_method = inspect.getattr_static(type(manager), async_method_name, None)
if inspect.iscoroutinefunction(class_async_method):
async_method = cast(
Callable[..., Awaitable[_SessionResult]],
getattr(manager, async_method_name),
)
return await async_method(*args, **kwargs)
return await shield_and_drain(asyncio.to_thread(sync_method, *args, **kwargs))
+113 -1
View File
@@ -1,5 +1,6 @@
"""Session management for conversation history."""
import asyncio
import base64
import errno
import hashlib
@@ -28,6 +29,7 @@ from nanobot.runtime_context import (
public_history_message,
)
from nanobot.session.model_selection import SESSION_MODEL_PRESET_METADATA_KEY
from nanobot.utils.cancellation import shield_and_drain
from nanobot.utils.helpers import (
content_with_media_breadcrumbs,
ensure_dir,
@@ -71,6 +73,7 @@ _WORKSPACE_STATE_DIR = ".nanobot"
_WORKSPACE_ID_FILE = "workspace-id"
_WORKSPACE_ID_RE = re.compile(r"^[0-9a-f]{32}$")
_SESSION_MIGRATION_LOCK_TIMEOUT_SECONDS = 30
_SESSION_FILES_LOCK_TIMEOUT_SECONDS = 5
_SESSION_FILES_LOCK_FILENAME = ".session-files.lock"
_COPY_CHUNK_SIZE = 1024 * 1024
@@ -560,7 +563,8 @@ class JsonlSessionStore:
self.sessions_dir = ensure_dir(root / workspace_id)
self.legacy_sessions_dir = get_legacy_sessions_dir()
self._session_files_lock = FileLock(
str(self.sessions_dir / _SESSION_FILES_LOCK_FILENAME)
str(self.sessions_dir / _SESSION_FILES_LOCK_FILENAME),
timeout=_SESSION_FILES_LOCK_TIMEOUT_SECONDS,
)
with self._session_files_lock:
self._migrate_from_workspace(canonical_workspace)
@@ -1642,6 +1646,7 @@ class SessionManager:
self._cache: OrderedDict[str, Session] = OrderedDict()
# Preserve identity for sessions held by active callers without retaining idle ones.
self._overflow_cache: WeakValueDictionary[str, Session] = WeakValueDictionary()
self._async_session_locks: WeakValueDictionary[str, asyncio.Lock] = WeakValueDictionary()
self._max_cached_sessions = SESSION_CACHE_MAX_SIZE
self._delete_observer: Callable[[str], None] | None = None
@@ -1741,6 +1746,28 @@ class SessionManager:
self._remember(session)
return session
def _async_session_lock(self, key: str) -> asyncio.Lock:
lock = self._async_session_locks.get(key)
if lock is None:
lock = asyncio.Lock()
self._async_session_locks[key] = lock
return lock
async def get_or_create_async(self, key: str) -> Session:
"""Load a session without running file I/O or lock waits on the event loop."""
cached = self.get_cached(key)
if cached is not None:
return cached
async with self._async_session_lock(key):
cached = self.get_cached(key)
if cached is not None:
return cached
session = await asyncio.to_thread(self._load, key)
if session is None:
session = Session(key=key)
self._remember(session)
return session
def get_or_create_transient(
self,
key: str,
@@ -1774,6 +1801,17 @@ class SessionManager:
self._store.save(session, fsync=fsync)
self._remember(session)
async def save_async(self, session: Session, *, fsync: bool = False) -> None:
"""Persist a session without blocking the caller's event loop."""
if not session.policy.persist:
return
async def save_and_remember() -> None:
await asyncio.to_thread(self._store.save, session, fsync=fsync)
self._remember(session)
await shield_and_drain(save_and_remember())
def save_runtime_checkpoint(self, session: Session) -> None:
"""Persist volatile recovery state without rewriting long history."""
if not session.policy.persist:
@@ -1786,6 +1824,23 @@ class SessionManager:
# they opt into a dedicated checkpoint primitive.
self.save(session)
async def save_runtime_checkpoint_async(self, session: Session) -> None:
"""Persist an in-flight checkpoint without blocking the event loop."""
if not session.policy.persist:
return
async def save_and_remember() -> None:
if self._store is self._jsonl_store:
await asyncio.to_thread(
self._jsonl_store.save_runtime_checkpoint,
session,
)
else:
await asyncio.to_thread(self._store.save, session)
self._remember(session)
await shield_and_drain(save_and_remember())
def rename_model_preset(self, old_name: str, new_name: str) -> int:
"""Rename a session-scoped model preset across durable and live sessions."""
if old_name == new_name:
@@ -1827,6 +1882,21 @@ class SessionManager:
raise
return len(changed)
async def flush_all_async(self) -> int:
"""Re-save every cached session without blocking the event loop."""
cached = dict(self._overflow_cache.items())
cached.update(self._cache)
flushed = 0
for key, session in cached.items():
try:
await shield_and_drain(
asyncio.to_thread(self._store.save, session, fsync=True)
)
flushed += 1
except Exception:
logger.warning("Failed to flush session {}", key, exc_info=True)
return flushed
def flush_all(self) -> int:
"""Re-save every cached session with fsync for durable shutdown.
@@ -1858,6 +1928,18 @@ class SessionManager:
self._delete_observer(key)
return deleted
async def delete_session_async(self, key: str) -> bool:
"""Delete a session without blocking the event loop."""
async def delete_and_notify() -> bool:
self.invalidate(key)
deleted = await asyncio.to_thread(self._store.delete, key)
if self._delete_observer is not None:
self._delete_observer(key)
return deleted
return await shield_and_drain(delete_and_notify())
def restore_sessions_to_workspace(self) -> SessionRestoreResult:
"""Restore session files to the pre-relocation path for an explicit rollback."""
return self._jsonl_store.restore_to_workspace()
@@ -1930,6 +2012,10 @@ class SessionManager:
"""Read session metadata without loading the transcript."""
return cast(dict[str, Any] | None, self._store.read_metadata(key))
async def read_session_metadata_async(self, key: str) -> dict[str, Any] | None:
"""Read session metadata without blocking the event loop."""
return await asyncio.to_thread(self.read_session_metadata, key)
def update_session_metadata(
self,
key: str,
@@ -1943,5 +2029,31 @@ class SessionManager:
session.metadata.update(deepcopy(updates))
return updated
async def update_session_metadata_async(
self,
key: str,
updates: dict[str, Any],
*,
fsync: bool = False,
) -> bool:
"""Update metadata without blocking the event loop."""
async def update_and_refresh_cache() -> bool:
updated = await asyncio.to_thread(
self._store.update_metadata,
key,
updates,
fsync=fsync,
)
if updated and (session := self.get_cached(key)) is not None:
session.metadata.update(deepcopy(updates))
return updated
return await shield_and_drain(update_and_refresh_cache())
def list_sessions(self) -> list[dict[str, Any]]:
return cast(list[dict[str, Any]], self._store.list_sessions())
async def list_sessions_async(self) -> list[dict[str, Any]]:
"""List persisted sessions without blocking the event loop."""
return await asyncio.to_thread(self.list_sessions)
+55 -23
View File
@@ -25,6 +25,7 @@ from nanobot.bus.outbound_events import (
)
from nanobot.bus.queue import MessageBus
from nanobot.session import turn_continuation
from nanobot.session.async_compat import call_session_manager
from nanobot.session.keys import UNIFIED_SESSION_KEY, last_channel_from_metadata
from nanobot.session.manager import Session, SessionManager
from nanobot.webui.metadata import WEBUI_TURN_METADATA_KEY
@@ -460,6 +461,37 @@ class RecoveryCoordinator:
repr=False,
)
async def _get_or_create_session(self, key: str) -> Session:
return await call_session_manager(
self.sessions,
"get_or_create_async",
self.sessions.get_or_create,
key,
)
async def _save_session(self, session: Session) -> None:
await call_session_manager(
self.sessions,
"save_async",
self.sessions.save,
session,
)
async def _read_session_metadata(self, key: str) -> dict[str, Any] | None:
return await call_session_manager(
self.sessions,
"read_session_metadata_async",
self.sessions.read_session_metadata,
key,
)
async def _list_sessions(self) -> list[dict[str, Any]]:
return await call_session_manager(
self.sessions,
"list_sessions_async",
self.sessions.list_sessions,
)
def register_recovery_task(self, session_key: str, task: asyncio.Task[Any]) -> None:
"""Track the task that owns an explicit recovery continuation."""
self._active_recovery_tasks[session_key] = task
@@ -482,8 +514,8 @@ class RecoveryCoordinator:
async def scan(self) -> None:
"""Recover every interrupted WebUI session once at gateway startup."""
for key in self._recovery_candidates():
metadata_payload = self.sessions.read_session_metadata(key)
for key in await self._recovery_candidates():
metadata_payload = await self._read_session_metadata(key)
raw_metadata = metadata_payload.get("metadata") if metadata_payload else None
metadata = cast(dict[str, Any], raw_metadata) if isinstance(raw_metadata, dict) else {}
route = self._websocket_route_for(key, metadata)
@@ -492,7 +524,7 @@ class RecoveryCoordinator:
unfinished = self._has_unfinished_webui_transcript(key)
if not self._needs_recovery(metadata) and not unfinished:
continue
session = self.sessions.get_or_create(key)
session = await self._get_or_create_session(key)
try:
await self._recover_session(session, route[1])
await self._requeue_pending_followups(session)
@@ -507,14 +539,14 @@ class RecoveryCoordinator:
reason="recovery_failed",
can_continue=False,
)
self.sessions.save(session)
await self._save_session(session)
await self._publish(route[1], failed)
def _recovery_candidates(self) -> list[str]:
async def _recovery_candidates(self) -> list[str]:
"""Discover canonical and transcript-only WebUI sessions cheaply."""
candidates = dict.fromkeys(
key
for item in self.sessions.list_sessions()
for item in await self._list_sessions()
if isinstance((key := item.get("key")), str)
)
try:
@@ -523,7 +555,7 @@ class RecoveryCoordinator:
# duplicating its filename and migration rules here would drift.
from nanobot.webui.session_list_index import list_webui_sessions
for item in list_webui_sessions(self.sessions):
for item in await asyncio.to_thread(list_webui_sessions, self.sessions):
key = item.get("key")
if isinstance(key, str):
candidates.setdefault(key, None)
@@ -549,7 +581,7 @@ class RecoveryCoordinator:
"""Reject stale queued recoveries and let new user input supersede them."""
recovery_id = message.metadata.get(RECOVERY_INBOUND_METADATA_KEY)
if isinstance(recovery_id, str):
session = self.sessions.get_or_create(message.session_key)
session = await self._get_or_create_session(message.session_key)
state = recovery_state_from_metadata(session.metadata)
return bool(
state
@@ -558,7 +590,7 @@ class RecoveryCoordinator:
)
if message.channel != "websocket":
return True
session = self.sessions.get_or_create(message.session_key)
session = await self._get_or_create_session(message.session_key)
state = recovery_state_from_metadata(session.metadata)
if state and state["status"] in {"resuming", "awaiting_user", "failed"}:
await self._cancel_active_recovery(message.session_key)
@@ -572,13 +604,13 @@ class RecoveryCoordinator:
attempts=cast(int, state.get("attempts", 0)),
reason="superseded",
)
self.sessions.save(session)
await self._save_session(session)
await self._publish(message.chat_id, recovered)
return True
async def turn_completed(self, session_key: str) -> None:
"""Resolve a resuming state after the recovered turn commits."""
session = self.sessions.get_or_create(session_key)
session = await self._get_or_create_session(session_key)
state = recovery_state_from_metadata(session.metadata)
if not state or state["status"] != "resuming":
return
@@ -592,7 +624,7 @@ class RecoveryCoordinator:
attempts=cast(int, state.get("attempts", 0)),
reason="continued",
)
self.sessions.save(session)
await self._save_session(session)
await self._publish(route[1], recovered)
async def handle_action(self, action: str, payload: dict[str, Any]) -> dict[str, Any]:
@@ -603,7 +635,7 @@ class RecoveryCoordinator:
raise RecoveryActionError("missing chat_id")
if not isinstance(recovery_id, str) or not recovery_id:
raise RecoveryActionError("missing recovery_id")
session = self.sessions.get_or_create(self._session_key(chat_id))
session = await self._get_or_create_session(self._session_key(chat_id))
state = recovery_state_from_metadata(session.metadata)
if not state or state["recovery_id"] != recovery_id:
raise RecoveryActionError("recovery state is stale", status=409)
@@ -618,7 +650,7 @@ class RecoveryCoordinator:
attempts=cast(int, state.get("attempts", 0)),
reason="dismissed",
)
self.sessions.save(session)
await self._save_session(session)
await self._publish(chat_id, next_state)
return next_state
if action != "continue":
@@ -635,7 +667,7 @@ class RecoveryCoordinator:
reason="user_confirmed",
resume_message_count=len(session.messages),
)
self.sessions.save(session)
await self._save_session(session)
await self._publish(chat_id, next_state)
await self._queue_continuation(session, chat_id, next_state)
return next_state
@@ -668,7 +700,7 @@ class RecoveryCoordinator:
attempts=cast(int, state.get("attempts", 1)),
reason="loop_guard",
)
self.sessions.save(session)
await self._save_session(session)
await self._publish(chat_id, next_state)
elif self._has_unfinished_webui_transcript(session.key):
# A normal last-client shutdown can materialize the checkpoint
@@ -690,7 +722,7 @@ class RecoveryCoordinator:
),
can_continue=can_continue,
)
self.sessions.save(session)
await self._save_session(session)
await self._publish(chat_id, waiting)
return
if state and state["status"] in {"awaiting_user", "failed"}:
@@ -706,7 +738,7 @@ class RecoveryCoordinator:
attempts=cast(int, state.get("attempts", 1)),
reason="loop_guard",
)
self.sessions.save(session)
await self._save_session(session)
await self._publish(chat_id, waiting)
return
@@ -724,7 +756,7 @@ class RecoveryCoordinator:
reason="checkpoint_unknown",
can_continue=False,
)
self.sessions.save(session)
await self._save_session(session)
await self._publish(chat_id, waiting)
return
if checkpoint is not None and not _runtime_checkpoint_is_well_formed(checkpoint):
@@ -738,7 +770,7 @@ class RecoveryCoordinator:
reason="checkpoint_invalid",
can_continue=False,
)
self.sessions.save(session)
await self._save_session(session)
await self._publish(chat_id, waiting)
return
if phase == "final_response":
@@ -750,7 +782,7 @@ class RecoveryCoordinator:
attempts=0,
reason="answer_restored",
)
self.sessions.save(session)
await self._save_session(session)
await self._publish(chat_id, recovered)
return
if phase in _UNCERTAIN_TOOL_PHASES or pending_calls:
@@ -762,7 +794,7 @@ class RecoveryCoordinator:
attempts=0,
reason="tool_state_unknown",
)
self.sessions.save(session)
await self._save_session(session)
await self._publish(chat_id, waiting)
return
# A gateway restart is a lifecycle boundary. Never enqueue model work
@@ -777,7 +809,7 @@ class RecoveryCoordinator:
attempts=0,
reason="restart_requires_confirmation",
)
self.sessions.save(session)
await self._save_session(session)
await self._publish(chat_id, waiting)
async def _queue_continuation(
+13 -10
View File
@@ -2,6 +2,7 @@
from __future__ import annotations
import asyncio
import re
import time
from collections.abc import Awaitable, Callable
@@ -176,7 +177,7 @@ async def maybe_generate_webui_title(
model: str,
) -> bool:
"""Generate and persist a short title for WebUI-owned sessions only."""
session = sessions.get_or_create(session_key)
session = await sessions.get_or_create_async(session_key)
if session.metadata.get(WEBUI_SESSION_METADATA_KEY) is not True:
return False
if session.metadata.get(WEBUI_TITLE_USER_EDITED_METADATA_KEY) is True:
@@ -187,7 +188,7 @@ async def maybe_generate_webui_title(
if cleaned_current_title:
if cleaned_current_title != current_title:
session.metadata[WEBUI_TITLE_METADATA_KEY] = cleaned_current_title
sessions.save(session)
await sessions.save_async(session)
return False
session.metadata.pop(WEBUI_TITLE_METADATA_KEY, None)
@@ -241,7 +242,7 @@ async def maybe_generate_webui_title(
)
return False
session.metadata[WEBUI_TITLE_METADATA_KEY] = title
sessions.save(session)
await sessions.save_async(session)
return True
@@ -437,8 +438,8 @@ class WebuiTurnRoutePolicy:
)
and route.channel == "websocket"
):
session = self.sessions.get_or_create(session_key)
if session.metadata.get(WEBUI_SESSION_METADATA_KEY) is True:
session = self.sessions.get_cached(session_key)
if session is not None and session.metadata.get(WEBUI_SESSION_METADATA_KEY) is True:
metadata = dict(route.metadata)
turn_prefix = "session-input" if internal_user_input else "subagent"
metadata.update({
@@ -580,7 +581,7 @@ class WebuiTurnCoordinator:
or not session_key.startswith("websocket:")
):
return
persisted = self.sessions.read_session_metadata(session_key)
persisted = await self.sessions.read_session_metadata_async(session_key)
metadata_value: object = persisted.get("metadata") if persisted is not None else None
metadata = (
cast(dict[str, Any], metadata_value)
@@ -591,7 +592,8 @@ class WebuiTurnCoordinator:
return
public_metadata = _session_message_public_metadata(envelope)
try:
append_session_message_input(
await asyncio.to_thread(
append_session_message_input,
session_key,
content=event.content,
created_at_ms=envelope["created_at_ms"],
@@ -616,8 +618,9 @@ class WebuiTurnCoordinator:
def _handle_session_turn_started(self, event: SessionTurnStarted) -> None:
if not self._is_websocket_event(event.context):
return
session = self.sessions.get_or_create(event.context.session_key)
mark_webui_session(session, event.context.metadata)
session = self.sessions.get_cached(event.context.session_key)
if session is not None:
mark_webui_session(session, event.context.metadata)
async def _handle_run_status_changed(self, event: TurnRunStatusChanged) -> None:
if not self._is_websocket_event(event.context):
@@ -703,7 +706,7 @@ class WebuiTurnCoordinator:
if msg.channel != "websocket":
return
session = self.sessions.get_or_create(session_key)
session = await self.sessions.get_or_create_async(session_key)
await self.bus.publish_outbound(
outbound_message_for_event(
channel=msg.channel,
+215 -59
View File
@@ -5,17 +5,24 @@ from __future__ import annotations
import asyncio
import uuid
from collections.abc import Awaitable, Callable
from typing import Any
from contextlib import suppress
from typing import Any, TypeVar
from loguru import logger
from nanobot.agent.automation_turns import AutomationTurnError
from nanobot.agent.automation_turns import (
AutomationTurnAcceptedCancellation,
AutomationTurnError,
)
from nanobot.bus.events import InboundMessage, OutboundMessage
from nanobot.triggers.local_session_turns import LOCAL_TRIGGER_META
from nanobot.triggers.local_store import LocalTriggerStore
from nanobot.triggers.local_types import LocalTrigger, TriggerDelivery
from nanobot.utils.cancellation import shield_and_drain
from nanobot.webui.metadata import WEBUI_MESSAGE_SOURCE_METADATA_KEY, WEBUI_TURN_METADATA_KEY
_T = TypeVar("_T")
async def run_local_trigger_queue(
*,
@@ -29,14 +36,16 @@ async def run_local_trigger_queue(
if submit_turn is None:
raise ValueError("run_local_trigger_queue requires submit_turn")
logger.info("Local trigger queue started")
recovered = store.recover_processing_deliveries()
recovered = await shield_and_drain(asyncio.to_thread(store.recover_processing_deliveries))
if recovered:
logger.warning(
"Trigger: recovered {} interrupted delivery file(s) from processing",
recovered,
)
while True:
deliveries = store.claim_deliveries(limit=batch_size)
deliveries = await shield_and_drain(
asyncio.to_thread(store.claim_deliveries, limit=batch_size)
)
if not deliveries:
await asyncio.sleep(poll_interval_s)
continue
@@ -49,30 +58,24 @@ async def run_local_trigger_queue(
submit_turn=submit_turn,
is_channel_enabled=is_channel_enabled,
)
store.complete_delivery(delivery)
except _DeliverySettledOnCancellation:
raise
except asyncio.CancelledError as exc:
store.retry_delivery(delivery, str(exc) or exc.__class__.__name__)
_write_delivery_run_record(
error = str(exc) or exc.__class__.__name__
await shield_and_drain(asyncio.to_thread(store.retry_delivery, delivery, error))
await _write_delivery_run_record(
store,
delivery,
status="interrupted",
error=str(exc) or exc.__class__.__name__,
error=error,
)
raise
except _TerminalDeliveryError as exc:
store.record_delivery(
delivery.trigger_id,
status="error",
error=str(exc),
run_at_ms=delivery.created_at_ms,
await _await_delivery_settlement(
_settle_failed_delivery(store, delivery, error=str(exc)),
store=store,
delivery=delivery,
)
_write_delivery_run_record(
store,
delivery,
status="error",
error=str(exc),
)
store.complete_delivery(delivery)
logger.warning(
"Trigger: dropped delivery {} for {}: {}",
delivery.id,
@@ -81,19 +84,11 @@ async def run_local_trigger_queue(
)
except AutomationTurnError as exc:
error = str(exc) or exc.__class__.__name__
store.record_delivery(
delivery.trigger_id,
status="error",
error=error,
run_at_ms=delivery.created_at_ms,
await _await_delivery_settlement(
_settle_failed_delivery(store, delivery, error=error),
store=store,
delivery=delivery,
)
_write_delivery_run_record(
store,
delivery,
status="error",
error=error,
)
store.complete_delivery(delivery)
logger.warning(
"Trigger: delivery {} for {} reached the agent but failed: {}",
delivery.id,
@@ -102,18 +97,10 @@ async def run_local_trigger_queue(
)
except Exception as exc:
error = str(exc) or exc.__class__.__name__
retried = store.retry_delivery(delivery, error)
_write_delivery_run_record(
store,
delivery,
status="retrying" if retried else "error",
error=error,
)
store.record_delivery(
delivery.trigger_id,
status="error",
error=error,
run_at_ms=delivery.created_at_ms,
retried = await _await_delivery_settlement(
_settle_retryable_delivery(store, delivery, error=error),
store=store,
delivery=delivery,
)
logger.exception(
"Trigger: failed delivery {} for {}{}",
@@ -127,6 +114,10 @@ class _TerminalDeliveryError(RuntimeError):
pass
class _DeliverySettledOnCancellation(asyncio.CancelledError):
"""Cancellation reported only after an already-submitted delivery is settled."""
async def _deliver_delivery(
store: LocalTriggerStore,
delivery: TriggerDelivery,
@@ -134,7 +125,7 @@ async def _deliver_delivery(
submit_turn: Callable[[InboundMessage], Awaitable[OutboundMessage | None]],
is_channel_enabled: Callable[[str], bool],
) -> None:
trigger = store.get(delivery.trigger_id)
trigger = await asyncio.to_thread(store.get, delivery.trigger_id)
if trigger is None:
raise _TerminalDeliveryError("trigger not found")
if not trigger.enabled:
@@ -142,7 +133,14 @@ async def _deliver_delivery(
if not is_channel_enabled(trigger.channel):
raise _TerminalDeliveryError(f"target channel is not enabled: {trigger.channel}")
store.write_delivery_run_record(delivery, trigger=trigger, status="processing")
await shield_and_drain(
asyncio.to_thread(
store.write_delivery_run_record,
delivery,
trigger=trigger,
status="processing",
)
)
msg = InboundMessage(
channel=trigger.channel,
sender_id=trigger.sender_id,
@@ -151,22 +149,177 @@ async def _deliver_delivery(
metadata=_delivery_metadata(trigger, delivery),
session_key_override=trigger.session_key,
)
response = await submit_turn(msg)
store.record_delivery(
trigger.id,
status="ok",
run_at_ms=delivery.created_at_ms,
try:
response = await submit_turn(msg)
except AutomationTurnAcceptedCancellation:
try:
await _await_delivery_settlement(
_settle_accepted_delivery(store, delivery, trigger=trigger),
store=store,
delivery=delivery,
)
except _DeliverySettledOnCancellation:
raise
except Exception:
logger.exception(
"Trigger: failed to persist accepted delivery {}; dropping retry",
delivery.id,
)
with suppress(Exception):
await shield_and_drain(asyncio.to_thread(store.complete_delivery, delivery))
raise _DeliverySettledOnCancellation from None
try:
await _await_delivery_settlement(
_settle_submitted_delivery(store, delivery, trigger=trigger, response=response),
store=store,
delivery=delivery,
)
except Exception:
logger.exception(
"Trigger: failed to persist status for submitted delivery {}; dropping retry",
delivery.id,
)
with suppress(Exception):
await shield_and_drain(asyncio.to_thread(store.complete_delivery, delivery))
async def _await_delivery_settlement(
operation: Awaitable[_T],
*,
store: LocalTriggerStore,
delivery: TriggerDelivery,
) -> _T:
settlement = asyncio.ensure_future(operation)
try:
return await asyncio.shield(settlement)
except asyncio.CancelledError:
while not settlement.done():
try:
await asyncio.shield(settlement)
except asyncio.CancelledError:
continue
try:
settlement.result()
except Exception:
logger.exception(
"Trigger: failed to settle delivery {} during cancellation",
delivery.id,
)
completion = asyncio.create_task(
shield_and_drain(asyncio.to_thread(store.complete_delivery, delivery))
)
while not completion.done():
try:
await asyncio.shield(completion)
except asyncio.CancelledError:
continue
with suppress(Exception):
completion.result()
raise _DeliverySettledOnCancellation from None
async def _settle_failed_delivery(
store: LocalTriggerStore,
delivery: TriggerDelivery,
*,
error: str,
) -> None:
await _write_delivery_run_record(
store,
delivery,
status="error",
error=error,
)
_write_delivery_run_record(
await shield_and_drain(asyncio.to_thread(store.complete_delivery, delivery))
# Publish the terminal status only after the durable delivery state is settled.
await shield_and_drain(
asyncio.to_thread(
store.record_delivery,
delivery.trigger_id,
status="error",
error=error,
run_at_ms=delivery.created_at_ms,
)
)
async def _settle_retryable_delivery(
store: LocalTriggerStore,
delivery: TriggerDelivery,
*,
error: str,
) -> bool:
retried = await shield_and_drain(asyncio.to_thread(store.retry_delivery, delivery, error))
await _write_delivery_run_record(
store,
delivery,
status="retrying" if retried else "error",
error=error,
)
await shield_and_drain(
asyncio.to_thread(
store.record_delivery,
delivery.trigger_id,
status="error",
error=error,
run_at_ms=delivery.created_at_ms,
)
)
return retried
async def _settle_accepted_delivery(
store: LocalTriggerStore,
delivery: TriggerDelivery,
*,
trigger: LocalTrigger,
) -> None:
"""Commit an accepted delivery without claiming the agent turn completed."""
await _write_delivery_run_record(
store,
delivery,
trigger=trigger,
status="accepted",
)
await shield_and_drain(asyncio.to_thread(store.complete_delivery, delivery))
await shield_and_drain(
asyncio.to_thread(
store.record_delivery,
trigger.id,
status="ok",
run_at_ms=delivery.created_at_ms,
)
)
async def _settle_submitted_delivery(
store: LocalTriggerStore,
delivery: TriggerDelivery,
*,
trigger: LocalTrigger,
response: OutboundMessage | None,
) -> None:
await _write_delivery_run_record(
store,
delivery,
trigger=trigger,
status="ok",
response=response.content if response else "",
)
await shield_and_drain(asyncio.to_thread(store.complete_delivery, delivery))
# last_status is the externally visible commit marker for a settled delivery.
await shield_and_drain(
asyncio.to_thread(
store.record_delivery,
trigger.id,
status="ok",
run_at_ms=delivery.created_at_ms,
)
)
def _write_delivery_run_record(
async def _write_delivery_run_record(
store: LocalTriggerStore,
delivery: TriggerDelivery,
*,
@@ -176,12 +329,15 @@ def _write_delivery_run_record(
response: str | None = None,
) -> None:
try:
store.write_delivery_run_record(
delivery,
trigger=trigger,
status=status,
error=error,
response=response,
await shield_and_drain(
asyncio.to_thread(
store.write_delivery_run_record,
delivery,
trigger=trigger,
status=status,
error=error,
response=response,
)
)
except Exception:
logger.exception(
+5 -1
View File
@@ -24,6 +24,7 @@ _MAX_RUN_HISTORY = 20
_MAX_DELIVERY_ATTEMPTS = 10
_RUN_RECORD_TEXT_MAX_CHARS = 4000
_PROCESSING_RECOVERY_ERROR = "delivery was recovered from interrupted processing"
_FILE_LOCK_TIMEOUT_SECONDS = 5
class TriggerStoreError(RuntimeError):
@@ -49,7 +50,10 @@ class LocalTriggerStore:
self.processing_dir = self.root / "processing"
self.failed_dir = self.root / "failed"
self.runs_dir = self.root / "runs"
self._lock = FileLock(str(self.root / ".lock"))
self._lock = FileLock(
str(self.root / ".lock"),
timeout=_FILE_LOCK_TIMEOUT_SECONDS,
)
def create(
self,
+40
View File
@@ -3,8 +3,48 @@
from __future__ import annotations
import asyncio
from collections.abc import Awaitable
from typing import TypeVar
_T = TypeVar("_T")
def task_is_cancelling() -> bool:
task = asyncio.current_task()
return task is not None and task.cancelling() > 0
async def shield_and_drain(awaitable: Awaitable[_T]) -> _T:
"""Delay caller cancellation until an accepted operation has fully settled.
``asyncio.to_thread`` cannot stop a worker that has already started. Shielding
keeps cancellation from detaching that worker, and draining also lets any
post-write in-memory settlement in ``awaitable`` finish. Cancellation is still
re-raised as soon as the accepted operation is done.
"""
settlement = asyncio.ensure_future(awaitable)
cancellation: asyncio.CancelledError | None = None
while not settlement.done():
try:
result = await asyncio.shield(settlement)
except asyncio.CancelledError as exc:
if cancellation is None:
cancellation = exc
except BaseException:
if cancellation is None:
raise
break
else:
if cancellation is not None:
raise cancellation
return result
if cancellation is not None:
try:
settlement.result()
except BaseException:
# The caller's cancellation wins once settlement has been observed.
pass
raise cancellation
return settlement.result()
+32 -19
View File
@@ -2,6 +2,7 @@
from __future__ import annotations
import asyncio
import re
import uuid
from collections.abc import Mapping
@@ -9,6 +10,7 @@ from typing import TYPE_CHECKING, Any, TypeGuard
from nanobot.session.manager import SessionManager
from nanobot.session.webui_turns import WEBUI_TITLE_METADATA_KEY, clean_generated_title
from nanobot.utils.cancellation import shield_and_drain
from nanobot.webui.transcript import (
append_fork_marker,
delete_webui_transcript,
@@ -93,24 +95,35 @@ async def handle_webui_fork_chat(
await channel.send_webui_protocol_error(connection, "session_manager_unavailable")
return
try:
forked = create_webui_chat_fork(
session_manager,
source_chat_id=source_chat_id,
before_user_index=raw_index,
title=envelope.get("title") if isinstance(envelope.get("title"), str) else None,
)
if forked is None:
await channel.send_webui_protocol_error(connection, "invalid fork source or index")
async def create_and_attach() -> None:
try:
forked = await asyncio.to_thread(
create_webui_chat_fork,
session_manager,
source_chat_id=source_chat_id,
before_user_index=raw_index,
title=(
envelope.get("title")
if isinstance(envelope.get("title"), str)
else None
),
)
if forked is None:
await channel.send_webui_protocol_error(
connection,
"invalid fork source or index",
)
return
fork_id, fork_key = forked
except Exception as exc:
channel.logger.warning("fork_chat failed: {}", exc)
await channel.send_webui_protocol_error(connection, "fork_chat_failed")
return
fork_id, fork_key = forked
except Exception as exc:
channel.logger.warning("fork_chat failed: {}", exc)
await channel.send_webui_protocol_error(connection, "fork_chat_failed")
return
await channel.attach_webui_fork(
connection,
fork_id=fork_id,
fork_key=fork_key,
)
await channel.attach_webui_fork(
connection,
fork_id=fork_id,
fork_key=fork_key,
)
await shield_and_drain(create_and_attach())
+2 -2
View File
@@ -16,7 +16,7 @@ from nanobot.agent.tools.mcp import MCPConnection, connect_mcp_servers
from nanobot.agent.tools.mcp_oauth import MCP_OAUTH_CALLBACK_PATH, MCPOAuthHandlers
from nanobot.agent.tools.registry import ToolRegistry
from nanobot.config.schema import MCPServerConfig
from nanobot.security.network import validate_url_target
from nanobot.security.network import async_validate_url_target
from nanobot.webui.http_utils import is_loopback_host
McpReload = Callable[[], Awaitable[dict[str, Any]]]
@@ -259,7 +259,7 @@ class McpOAuthManager:
):
flow.error = "The MCP server returned an unsafe authorization URL."
raise McpOAuthError(flow.error)
ok, _error = validate_url_target(authorization_url)
ok, _error = await async_validate_url_target(authorization_url)
if not ok:
flow.error = "The MCP server returned an unsafe authorization URL."
raise McpOAuthError(flow.error)
+38 -23
View File
@@ -27,6 +27,7 @@ from nanobot.providers.image_generation import (
)
from nanobot.providers.registry import find_by_name
from nanobot.security.network import is_loopback_host
from nanobot.utils.cancellation import shield_and_drain
from nanobot.webui.settings_contracts import (
QueryParams,
SettingsRequest,
@@ -640,7 +641,11 @@ class CapabilitySettingsHandler:
) -> SettingsRouteResult:
if action == "api-status":
return SettingsRouteResult.success(
api_service_payload(self.settings, operations.api_runtime())
await asyncio.to_thread(
api_service_payload,
self.settings,
operations.api_runtime(),
)
)
if action == "api-start":
return await self._start_api(request, operations)
@@ -673,17 +678,22 @@ class CapabilitySettingsHandler:
return SettingsRouteResult.failure(404, "unknown settings action")
operation, section, apply_image_reload = mutation
try:
payload = self.settings.mutate(operation, request.query)
except WebUISettingsError as exc:
return SettingsRouteResult.failure(exc.status, exc.message)
if apply_image_reload:
payload, image_restart_cleared = await self.apply_image_runtime_change(
async def mutate_and_apply() -> tuple[dict[str, Any], bool]:
payload = await self.settings.mutate_async(operation, request.query)
if not apply_image_reload:
return payload, False
return await self.apply_image_runtime_change(
payload,
operations.reload_image,
)
else:
image_restart_cleared = False
try:
payload, image_restart_cleared = await shield_and_drain(
mutate_and_apply()
)
except WebUISettingsError as exc:
return SettingsRouteResult.failure(exc.status, exc.message)
return SettingsRouteResult.success(
payload,
decorate_restart=True,
@@ -726,16 +736,17 @@ class CapabilitySettingsHandler:
400,
"API service API key must be a string",
)
try:
await asyncio.to_thread(
self.settings.mutate,
allow_install = await self._allow_feature_package_install(request)
async def mutate_and_start() -> Any:
await self.settings.mutate_async(
operations.nanobot_features_action,
"enable",
{"name": ["api"]},
allow_install=self._allow_feature_package_install(request),
allow_install=allow_install,
)
self.settings.mutate(operations.update_api, request.query)
config = self.settings.config.load()
await self.settings.mutate_async(operations.update_api, request.query)
config = await self.settings.config.load_async()
runtime = operations.api_runtime()
options = ApiStartOptions(
host=config.api.host,
@@ -744,10 +755,13 @@ class CapabilitySettingsHandler:
config_path=str(self.settings.config.path),
)
current = runtime.status()
result = await asyncio.to_thread(
return await asyncio.to_thread(
runtime.restart if current.running else runtime.start_background,
options,
)
try:
result = await shield_and_drain(mutate_and_start())
if not result.ok:
return SettingsRouteResult.failure(
500,
@@ -762,7 +776,8 @@ class CapabilitySettingsHandler:
self.logger.exception("failed to start managed API service")
return SettingsRouteResult.failure(500, str(exc))
return SettingsRouteResult.success(
api_service_payload(
await asyncio.to_thread(
api_service_payload,
self.settings,
operations.api_runtime(),
last_action="started",
@@ -775,7 +790,7 @@ class CapabilitySettingsHandler:
) -> SettingsRouteResult:
runtime = operations.api_runtime()
try:
result = await asyncio.to_thread(runtime.stop)
result = await shield_and_drain(asyncio.to_thread(runtime.stop))
except Exception as exc:
self.logger.exception("failed to stop managed API service")
return SettingsRouteResult.failure(500, str(exc))
@@ -785,20 +800,20 @@ class CapabilitySettingsHandler:
api_runtime_message(result.message),
)
return SettingsRouteResult.success(
api_service_payload(
await asyncio.to_thread(
api_service_payload,
self.settings,
operations.api_runtime(),
last_action="stopped",
)
)
def _allow_feature_package_install(self, request: SettingsRequest) -> bool:
async def _allow_feature_package_install(self, request: SettingsRequest) -> bool:
if request.local_browser:
return True
try:
return bool(
self.settings.config.load().tools.webui_allow_remote_package_install
)
config = await self.settings.config.load_async()
return bool(config.tools.webui_allow_remote_package_install)
except Exception:
self.logger.exception("failed to load remote package install policy")
return False
+68 -32
View File
@@ -29,6 +29,7 @@ from nanobot.config.schema import Config, FallbackCandidate, ModelPresetConfig,
from nanobot.providers.image_generation import get_image_gen_provider
from nanobot.providers.oauth_guidance import OAUTH_CLI_KIT_MISSING_MESSAGE
from nanobot.providers.registry import PROVIDERS, create_dynamic_spec, find_by_name
from nanobot.utils.cancellation import shield_and_drain
from nanobot.webui.settings_contracts import (
QueryParams,
SettingsRequest,
@@ -1651,6 +1652,30 @@ class ModelSettingsHandler:
if self.settings.refresh_runtime_config is not None:
self.settings.refresh_runtime_config()
async def _mutate_and_refresh(
self,
operation: SettingsOperation,
query: QueryParams,
**kwargs: Any,
) -> dict[str, Any]:
payload = await self.settings.mutate_async(operation, query, **kwargs)
self._refresh_runtime_config()
return payload
async def _update_provider_and_runtime(
self,
operation: SettingsOperation,
query: QueryParams,
apply_image_runtime_change: Callable[
[dict[str, Any]],
Awaitable[tuple[dict[str, Any], bool]],
],
) -> tuple[dict[str, Any], bool]:
payload = await self.settings.mutate_async(operation, query)
payload, image_restart_cleared = await apply_image_runtime_change(payload)
self._refresh_runtime_config()
return payload, image_restart_cleared
async def handle(
self,
action: str,
@@ -1659,8 +1684,12 @@ class ModelSettingsHandler:
) -> SettingsRouteResult:
try:
if action == "agent-update":
payload = self.settings.mutate(operations.update_agent, request.query)
self._refresh_runtime_config()
payload = await shield_and_drain(
self._mutate_and_refresh(
operations.update_agent,
request.query,
)
)
return SettingsRouteResult.success(
payload,
decorate_restart=True,
@@ -1668,12 +1697,13 @@ class ModelSettingsHandler:
)
if action == "model-update":
payload = self.settings.mutate(
operations.update_model,
request.query,
rename_model_preset=self.settings.rename_model_preset,
payload = await shield_and_drain(
self._mutate_and_refresh(
operations.update_model,
request.query,
rename_model_preset=self.settings.rename_model_preset,
)
)
self._refresh_runtime_config()
return SettingsRouteResult.success(payload, decorate_restart=True)
mutation = {
@@ -1684,19 +1714,19 @@ class ModelSettingsHandler:
"provider-create": operations.create_provider,
}.get(action)
if mutation is not None:
payload = self.settings.mutate(mutation, request.query)
self._refresh_runtime_config()
payload = await shield_and_drain(
self._mutate_and_refresh(mutation, request.query)
)
return SettingsRouteResult.success(payload, decorate_restart=True)
if action == "provider-update":
payload = self.settings.mutate(
operations.update_provider,
request.query,
payload, image_restart_cleared = await shield_and_drain(
self._update_provider_and_runtime(
operations.update_provider,
request.query,
operations.apply_image_runtime_change,
)
)
payload, image_restart_cleared = await operations.apply_image_runtime_change(
payload
)
self._refresh_runtime_config()
return SettingsRouteResult.success(
payload,
decorate_restart=True,
@@ -1724,11 +1754,13 @@ class ModelSettingsHandler:
return SettingsRouteResult.success(payload)
if action == "oauth-login":
payload = await asyncio.to_thread(
self.settings.read,
operations.oauth_login,
request.query,
oauth_flows=self.settings.oauth_flows,
payload = await shield_and_drain(
asyncio.to_thread(
self.settings.read,
operations.oauth_login,
request.query,
oauth_flows=self.settings.oauth_flows,
)
)
elif action == "oauth-complete":
raw_response = (request.payload or {}).get("authorization_response")
@@ -1736,19 +1768,23 @@ class ModelSettingsHandler:
raise WebUISettingsError(
"OAuth authorization response must be a string"
)
payload = await asyncio.to_thread(
self.settings.read,
operations.oauth_complete,
request.query,
raw_response or None,
oauth_flows=self.settings.oauth_flows,
payload = await shield_and_drain(
asyncio.to_thread(
self.settings.read,
operations.oauth_complete,
request.query,
raw_response or None,
oauth_flows=self.settings.oauth_flows,
)
)
elif action == "oauth-logout":
payload = await asyncio.to_thread(
self.settings.read,
operations.oauth_logout,
request.query,
oauth_flows=self.settings.oauth_flows,
payload = await shield_and_drain(
asyncio.to_thread(
self.settings.read,
operations.oauth_logout,
request.query,
oauth_flows=self.settings.oauth_flows,
)
)
else:
return SettingsRouteResult.failure(404, "unknown settings action")
+36 -21
View File
@@ -4,6 +4,7 @@ from __future__ import annotations
import asyncio
import html
import inspect
import json
from collections.abc import Awaitable, Callable, Mapping
from typing import Any, cast
@@ -18,6 +19,7 @@ from nanobot.bus.queue import MessageBus
from nanobot.channels.registry import load_channel_plugin
from nanobot.channels.validation import validate_channel_config
from nanobot.pairing import approve_code, deny_code, list_pending
from nanobot.utils.cancellation import shield_and_drain
from nanobot.webui import settings_capabilities as capability_domain
from nanobot.webui import settings_contracts as contracts
from nanobot.webui import settings_models as model_domain
@@ -208,6 +210,16 @@ def _payload_query(payload: dict[str, Any]) -> QueryParams:
}
async def _call_settings_handler(
handler: Callable[[], Response | Awaitable[Response]],
) -> Response:
"""Keep synchronous handlers off-loop while supporting native async handlers."""
result = await asyncio.to_thread(handler)
if inspect.isawaitable(result):
return await result
return result
class WebUISettingsRouter:
"""Authenticate and dispatch settings requests to transport-neutral domains."""
@@ -284,9 +296,9 @@ class WebUISettingsRouter:
if not self._authorized(request):
return self._unauthorized()
if route == ("root", "settings"):
return await asyncio.to_thread(self._handle_settings)
return await _call_settings_handler(self._handle_settings)
if route == ("root", "usage"):
return await asyncio.to_thread(self._handle_settings_usage)
return await _call_settings_handler(self._handle_settings_usage)
domain, action = route
domain_request = self._domain_request(
@@ -415,19 +427,18 @@ class WebUISettingsRouter:
)
return self._json_response(payload)
def _handle_settings(self) -> Response:
return self._json_response(
self._with_restart_state(
self.settings.read(
settings_payload,
surface=self._runtime_surface,
runtime_capability_overrides=self._runtime_capabilities,
)
)
async def _handle_settings(self) -> Response:
payload = await self.settings.read_async(
settings_payload,
surface=self._runtime_surface,
runtime_capability_overrides=self._runtime_capabilities,
)
return self._json_response(self._with_restart_state(payload))
def _handle_settings_usage(self) -> Response:
return self._json_response(self.settings.read(settings_usage_payload))
async def _handle_settings_usage(self) -> Response:
return self._json_response(
await self.settings.read_async(settings_usage_payload)
)
def _model_operations(self) -> model_domain.ModelSettingsOperations:
return model_domain.ModelSettingsOperations(
@@ -560,7 +571,7 @@ class WebUISettingsRouter:
allow_install=allow_install,
)
def _allow_feature_package_install(
async def _allow_feature_package_install(
self,
connection: Any,
request: WsRequest,
@@ -570,29 +581,33 @@ class WebUISettingsRouter:
request,
needs_local_browser=True,
)
return self._system.allow_feature_package_install(domain_request)
return await self._system.allow_feature_package_install(domain_request)
async def _handle_mcp_oauth_start(self, request: WsRequest) -> Response:
if not self._authorized(request):
return self._unauthorized()
if self._mcp_oauth_redirect_uri is None:
redirect_uri_for_request = self._mcp_oauth_redirect_uri
if redirect_uri_for_request is None:
return self._error_response(500, "MCP OAuth callback is not configured")
query = self._parse_mcp_settings_query(request)
try:
name, cfg = await asyncio.to_thread(
self.settings.mutate,
async def mutate_and_start() -> dict[str, Any]:
name, cfg = await self.settings.mutate_async(
ensure_mcp_oauth_server,
query,
)
redirect_uri = self._mcp_oauth_redirect_uri(request)
redirect_uri = redirect_uri_for_request(request)
reset = (_query_first(query, "reset") or "").lower() in {"1", "true", "yes"}
payload = await self._mcp_oauth.start(
return await self._mcp_oauth.start(
name,
cfg,
redirect_uri,
reload_mcp=self._reload_mcp_runtime,
reset_credentials=reset,
)
try:
payload = await shield_and_drain(mutate_and_start())
except Exception as exc:
return self._mcp_oauth_error_response(exc, action="start")
return self._json_response(payload)
+43 -1
View File
@@ -2,6 +2,7 @@
from __future__ import annotations
import asyncio
import threading
from collections.abc import Callable
from dataclasses import dataclass
@@ -12,9 +13,11 @@ from filelock import FileLock
from nanobot.config.loader import load_config, save_config
from nanobot.config.schema import Config
from nanobot.utils.cancellation import shield_and_drain
_T = TypeVar("_T")
_WEBUI_OAUTH_MAX_FLOWS = 8
_SETTINGS_FILE_LOCK_TIMEOUT_SECONDS = 5
class WebUISettingsConfig:
@@ -25,13 +28,20 @@ class WebUISettingsConfig:
self.path.parent.mkdir(parents=True, exist_ok=True)
self._lock = threading.RLock()
lock_path = self.path.with_suffix(f"{self.path.suffix}.lock")
self._file_lock = FileLock(str(lock_path))
self._file_lock = FileLock(
str(lock_path),
timeout=_SETTINGS_FILE_LOCK_TIMEOUT_SECONDS,
)
def load(self) -> Config:
"""Load this gateway's config without consulting the process-global path."""
with self._lock:
return load_config(self.path)
async def load_async(self) -> Config:
"""Load config without running file I/O or lock waits on the event loop."""
return await asyncio.to_thread(self.load)
def update(self, mutation: Callable[[Config], _T]) -> _T:
"""Apply and atomically persist one path-scoped read-modify-write operation."""
with self._lock, self._file_lock:
@@ -40,11 +50,21 @@ class WebUISettingsConfig:
save_config(config, self.path)
return result
async def update_async(self, mutation: Callable[[Config], _T]) -> _T:
"""Update config without blocking the event loop."""
return await shield_and_drain(asyncio.to_thread(self.update, mutation))
def run_serialized(self, operation: Callable[[Path], _T]) -> _T:
"""Run a path-aware read-modify-write operation under the config-file lock."""
with self._lock, self._file_lock:
return operation(self.path)
async def run_serialized_async(self, operation: Callable[[Path], _T]) -> _T:
"""Run a serialized config operation without blocking the event loop."""
return await shield_and_drain(
asyncio.to_thread(self.run_serialized, operation)
)
class WebUIOAuthFlowRegistry:
"""Bounded, thread-safe OAuth flows owned by one gateway instance."""
@@ -146,6 +166,16 @@ class WebUISettingsServices:
"""Run a settings read against this gateway's explicit config path."""
return operation(*args, config_path=self.config.path, **kwargs)
async def read_async(
self,
operation: Callable[..., _T],
/,
*args: Any,
**kwargs: Any,
) -> _T:
"""Run a settings read without blocking the event loop."""
return await asyncio.to_thread(self.read, operation, *args, **kwargs)
def mutate(
self,
operation: Callable[..., _T],
@@ -161,3 +191,15 @@ class WebUISettingsServices:
**kwargs,
)
)
async def mutate_async(
self,
operation: Callable[..., _T],
/,
*args: Any,
**kwargs: Any,
) -> _T:
"""Mutate settings without blocking the event loop."""
return await shield_and_drain(
asyncio.to_thread(self.mutate, operation, *args, **kwargs)
)
+61 -21
View File
@@ -23,6 +23,7 @@ from nanobot.config.schema import Config
from nanobot.llm_usage import llm_usage_payload
from nanobot.optional_features import OptionalFeatureError, with_channel_runtime_status
from nanobot.security.workspace_access import workspace_sandbox_status
from nanobot.utils.cancellation import shield_and_drain
from nanobot.webui.settings_capabilities import network_safety_payload
from nanobot.webui.settings_contracts import (
QueryParams,
@@ -446,12 +447,17 @@ class SystemSettingsHandler:
operations: SystemSettingsOperations,
) -> SettingsRouteResult:
try:
payload = await asyncio.to_thread(
pending = asyncio.to_thread(
operations.cli_apps_action,
action,
request.query,
config_path=self.settings.config.path,
)
payload = (
await shield_and_drain(pending)
if action in {"install", "update", "uninstall"}
else await pending
)
except WebUISettingsError as exc:
return SettingsRouteResult.failure(exc.status, exc.message)
except Exception as exc:
@@ -505,17 +511,29 @@ class SystemSettingsHandler:
action: str,
operations: SystemSettingsOperations,
) -> SettingsRouteResult:
try:
allow_install = (
action != "enable"
or await self.allow_feature_package_install(request)
)
async def mutate_and_apply() -> dict[str, Any]:
payload = await asyncio.to_thread(
self._nanobot_features_action,
action,
request.query,
operations,
allow_install=(
action != "enable"
or self.allow_feature_package_install(request)
),
allow_install=allow_install,
)
payload = await self._apply_feature_runtime_change(
action,
request.query,
payload,
operations,
)
return self._with_channel_runtime_status(payload, operations)
try:
payload = await shield_and_drain(mutate_and_apply())
except OptionalFeatureError as exc:
return SettingsRouteResult.failure(exc.status, exc.message)
except Exception as exc:
@@ -527,13 +545,6 @@ class SystemSettingsHandler:
action,
)
return SettingsRouteResult.failure(status, message)
payload = await self._apply_feature_runtime_change(
action,
request.query,
payload,
operations,
)
payload = self._with_channel_runtime_status(payload, operations)
return SettingsRouteResult.success(
payload,
decorate_restart=True,
@@ -628,6 +639,15 @@ class SystemSettingsHandler:
self,
request: SettingsRequest,
operations: SystemSettingsOperations,
) -> SettingsRouteResult:
return await shield_and_drain(
self._channel_configure_settled(request, operations)
)
async def _channel_configure_settled(
self,
request: SettingsRequest,
operations: SystemSettingsOperations,
) -> SettingsRouteResult:
name = (query_first(request.query, "name") or "").strip()
instance_id = (
@@ -682,7 +702,7 @@ class SystemSettingsHandler:
"enable",
feature_query,
operations,
allow_install=self.allow_feature_package_install(request),
allow_install=await self.allow_feature_package_install(request),
)
except OptionalFeatureError as exc:
return SettingsRouteResult.failure(
@@ -825,6 +845,22 @@ class SystemSettingsHandler:
channel_name: str,
payload: dict[str, Any],
operations: SystemSettingsOperations,
) -> dict[str, Any]:
return await shield_and_drain(
self._settle_channel_connect_success(
request,
channel_name,
payload,
operations,
)
)
async def _settle_channel_connect_success(
self,
request: SettingsRequest,
channel_name: str,
payload: dict[str, Any],
operations: SystemSettingsOperations,
) -> dict[str, Any]:
target = {"name": [channel_name]}
if payload.get("instance_id"):
@@ -835,11 +871,11 @@ class SystemSettingsHandler:
"enable",
target,
operations,
allow_install=self.allow_feature_package_install(request),
allow_install=await self.allow_feature_package_install(request),
)
except OptionalFeatureError as exc:
features = self.feature_runtime_fallback(
self._nanobot_features_payload(operations),
await asyncio.to_thread(self._nanobot_features_payload, operations),
message=(
f"{channel_name} connected, but enabling channel support failed: "
f"{exc.message}"
@@ -859,13 +895,12 @@ class SystemSettingsHandler:
)
return updated
def allow_feature_package_install(self, request: SettingsRequest) -> bool:
async def allow_feature_package_install(self, request: SettingsRequest) -> bool:
if request.local_browser:
return True
try:
return bool(
self.settings.config.load().tools.webui_allow_remote_package_install
)
config = await self.settings.config.load_async()
return bool(config.tools.webui_allow_remote_package_install)
except Exception:
self.logger.exception("failed to load remote package install policy")
return False
@@ -925,13 +960,18 @@ class SystemSettingsHandler:
operations: SystemSettingsOperations,
) -> SettingsRouteResult:
try:
payload = await operations.mcp_presets_action(
pending = operations.mcp_presets_action(
action,
request.query,
reload_mcp=operations.reload_mcp,
mcp_runtime_status=operations.mcp_runtime_status,
config=self.settings.config,
)
payload = (
await pending
if action is None
else await shield_and_drain(pending)
)
except Exception as exc:
status = getattr(exc, "status", 500)
message = getattr(exc, "message", str(exc))
+113 -24
View File
@@ -34,6 +34,7 @@ from nanobot.session.session_handles import (
SessionHandleResolver,
)
from nanobot.triggers.local_types import LocalTrigger
from nanobot.utils.cancellation import shield_and_drain
from nanobot.webui.file_preview import (
WebUIFilePreviewError,
file_preview_availability_payload,
@@ -135,6 +136,18 @@ _WEBUI_MUTATION_PAYLOAD_ATTR = "_nanobot_webui_mutation_payload"
_WEBUI_MUTATION_REQUEST_ATTR = "_nanobot_webui_mutation_request"
_NO_STORE_HEADERS = [("Cache-Control", "no-store")]
def _slow_http_operation(path: str) -> str:
"""Return a route family without logging user-controlled path/query values."""
clean_path = path.split("?", 1)[0]
if clean_path == "/webui/bootstrap":
return clean_path
parts = [part for part in clean_path.split("/") if part]
if len(parts) >= 2 and parts[0] == "api":
return f"/api/{parts[1]}"
return "/webui"
_WEBUI_MUTATION_PATHS = {
"automation.enable": "/api/webui/automations/enable",
"automation.disable": "/api/webui/automations/disable",
@@ -420,7 +433,12 @@ class GatewayHTTPHandler:
response = await self._dispatch_resolved(connection, request, got)
return response
finally:
self._log_slow_http(got, response, started)
self._log_slow_http(
got,
response,
started,
input_chars=len(request.path),
)
async def dispatch_webui_mutation(
self,
@@ -556,7 +574,14 @@ class GatewayHTTPHandler:
return connection.respond(404, "Not Found")
def _log_slow_http(self, path: str, response: Any | None, started: float) -> None:
def _log_slow_http(
self,
path: str,
response: Any | None,
started: float,
*,
input_chars: int,
) -> None:
elapsed_ms = int((time.perf_counter() - started) * 1000)
if elapsed_ms < _SLOW_WEBUI_HTTP_LOG_MS:
return
@@ -564,9 +589,10 @@ class GatewayHTTPHandler:
return
status = getattr(response, "status_code", None)
self._log.warning(
"slow webui http route path={} status={} duration_ms={}",
path,
"slow webui http operation={} status={} input_chars={} duration_ms={}",
_slow_http_operation(path),
status if status is not None else "none",
input_chars,
elapsed_ms,
)
@@ -694,7 +720,11 @@ class GatewayHTTPHandler:
async def _dispatch_session_routes(self, request: WsRequest, got: str) -> Response | None:
m = re.match(r"^/api/sessions/([^/]+)/webui-thread$", got)
if m:
return self._handle_webui_thread_get(request, m.group(1))
return await asyncio.to_thread(
self._handle_webui_thread_get,
request,
m.group(1),
)
m = re.match(r"^/api/sessions/([^/]+)/context$", got)
if m:
@@ -702,15 +732,27 @@ class GatewayHTTPHandler:
m = re.match(r"^/api/sessions/([^/]+)/file-preview$", got)
if m:
return self._handle_file_preview(request, m.group(1))
return await asyncio.to_thread(
self._handle_file_preview,
request,
m.group(1),
)
m = re.match(r"^/api/sessions/([^/]+)/automations$", got)
if m:
return self._handle_session_automations(request, m.group(1))
return await self._run_cron_transaction(
self._handle_session_automations,
request,
m.group(1),
)
m = re.match(r"^/api/sessions/([^/]+)/delete$", got)
if m:
return self._handle_session_delete(request, m.group(1))
return await self._run_cron_transaction(
self._handle_session_delete,
request,
m.group(1),
)
return None
@@ -957,13 +999,24 @@ class GatewayHTTPHandler:
# -- Automation routes --------------------------------------------------
async def _run_cron_transaction(
self,
operation: Callable[..., Any],
/,
*args: Any,
**kwargs: Any,
) -> Any:
if self.cron_service is not None:
return await self.cron_service.run_sync(operation, *args, **kwargs)
return await shield_and_drain(asyncio.to_thread(operation, *args, **kwargs))
async def _dispatch_automation_routes(
self,
request: WsRequest,
got: str,
) -> Response | None:
if got == "/api/webui/automations":
return self._handle_webui_automations(request)
return await self._run_cron_transaction(self._handle_webui_automations, request)
m = re.match(r"^/api/webui/automations/(enable|disable|delete|run|update)$", got)
if m:
return await self._handle_webui_automation_action(request, m.group(1))
@@ -1029,13 +1082,24 @@ class GatewayHTTPHandler:
job_id = (_query_first(query, "id") or _query_first(query, "job_id") or "").strip()
if not job_id:
return _http_error(400, "missing automation id")
trigger = self.local_trigger_store.get(job_id) if self.local_trigger_store else None
trigger = (
await asyncio.to_thread(self.local_trigger_store.get, job_id)
if self.local_trigger_store
else None
)
if trigger is not None:
return self._handle_local_trigger_action(request, action, trigger)
return await shield_and_drain(
asyncio.to_thread(
self._handle_local_trigger_action,
request,
action,
trigger,
)
)
if self.cron_service is None:
return _http_error(404, "automation not found")
job = self.cron_service.get_job(job_id)
job = await self.cron_service.run_sync(self.cron_service.get_job, job_id)
if job is None:
return _http_error(404, "automation not found")
if job.payload.kind == "system_event":
@@ -1044,13 +1108,23 @@ class GatewayHTTPHandler:
return _http_error(409, "automation has no linked chat")
if action == "enable":
if self.cron_service.enable_job(job_id, enabled=True) is None:
result = await self.cron_service.run_sync(
self.cron_service.enable_job,
job_id,
enabled=True,
)
if result is None:
return _http_error(404, "automation not found")
elif action == "disable":
if self.cron_service.enable_job(job_id, enabled=False) is None:
result = await self.cron_service.run_sync(
self.cron_service.enable_job,
job_id,
enabled=False,
)
if result is None:
return _http_error(404, "automation not found")
elif action == "delete":
result = self.cron_service.remove_job(job_id)
result = await self.cron_service.run_sync(self.cron_service.remove_job, job_id)
if result == "not_found":
return _http_error(404, "automation not found")
if result == "protected":
@@ -1068,7 +1142,11 @@ class GatewayHTTPHandler:
if isinstance(parsed, str):
return _http_error(400, parsed)
try:
result = self.cron_service.update_job(job_id, **parsed)
result = await self.cron_service.run_sync(
self.cron_service.update_job,
job_id,
**parsed,
)
except ValueError as exc:
return _http_error(400, str(exc))
if result == "not_found":
@@ -1078,7 +1156,7 @@ class GatewayHTTPHandler:
else:
return _http_error(404, "unknown automation action")
return self._handle_webui_automations(request)
return await self._run_cron_transaction(self._handle_webui_automations, request)
def _handle_local_trigger_action(
self,
@@ -1163,9 +1241,17 @@ class GatewayHTTPHandler:
if got == "/api/webui/skills/install":
return await self._handle_webui_skill_install(connection, request)
if got == "/api/webui/skills/update":
return self._handle_webui_skill_update(request)
return await shield_and_drain(
asyncio.to_thread(self._handle_webui_skill_update, request)
)
if got == "/api/webui/skills/delete":
return self._handle_webui_skill_delete(connection, request)
return await shield_and_drain(
asyncio.to_thread(
self._handle_webui_skill_delete,
connection,
request,
)
)
if got == "/api/webui/skills":
return self._handle_webui_skills(request)
m = re.match(r"^/api/webui/skills/([^/]+)$", got)
@@ -1276,7 +1362,7 @@ class GatewayHTTPHandler:
) -> Response:
if not self.check_api_token(request):
return _http_error(401, "Unauthorized")
if not self._allow_webui_package_install(connection, request):
if not await self._allow_webui_package_install(connection, request):
return _http_error(403, "remote skill installation is disabled")
if self._skill_install_lock.locked():
return _http_error(409, "another skill installation is already in progress")
@@ -1308,13 +1394,16 @@ class GatewayHTTPHandler:
"last_action": action,
})
def _allow_webui_package_install(self, connection: Any, request: WsRequest) -> bool:
async def _allow_webui_package_install(
self,
connection: Any,
request: WsRequest,
) -> bool:
if _is_local_browser_request(connection, request.headers):
return True
try:
return bool(
self.settings.config.load().tools.webui_allow_remote_package_install
)
config = await self.settings.config.load_async()
return bool(config.tools.webui_allow_remote_package_install)
except Exception:
self._log.exception("failed to load remote package install policy")
return False
+15 -15
View File
@@ -123,21 +123,21 @@ def allow_loopback_mcp_urls(monkeypatch: pytest.MonkeyPatch):
_resolver_lock = asyncio.Lock()
monkeypatch.setattr(mcp_module, "PinnedDNSAsyncTransport", TestPinnedDNSAsyncTransport)
monkeypatch.setattr(
mcp_module,
"validate_url_target",
lambda url, *, allow_loopback=False: (True, ""),
)
monkeypatch.setattr(
mcp_module,
"resolve_url_target",
lambda url, *, allow_loopback=False: (True, "", ("127.0.0.1",)),
)
monkeypatch.setattr(
security_network,
"resolve_url_target",
lambda url, *, allow_loopback=False: (True, "", ("127.0.0.1",)),
)
async def allow_url(url: str, *, allow_loopback: bool = False) -> tuple[bool, str]:
return True, ""
async def resolve_url(
url: str,
*,
allow_loopback: bool = False,
trust_remote_dns: bool = False,
timeout_s: float = 3.0,
) -> tuple[bool, str, tuple[str, ...]]:
return True, "", ("127.0.0.1",)
monkeypatch.setattr(mcp_module, "async_validate_url_target", allow_url)
monkeypatch.setattr(mcp_module, "async_resolve_url_target", resolve_url)
monkeypatch.setattr(security_network, "async_resolve_url_target", resolve_url)
monkeypatch.setattr(
mcp_module,
"env_proxy_applies_to_url",
+47
View File
@@ -357,3 +357,50 @@ async def test_length_finish_with_blank_content_routes_to_length_recovery():
"finish_reason='length' response with blank content"
)
assert result.final_content == "done"
@pytest.mark.asyncio
async def test_slow_tool_log_records_scale_without_argument_content(monkeypatch) -> None:
from nanobot.agent import runner as runner_module
from nanobot.agent.runner import AgentRunner
records: list[str] = []
class _Logger:
def warning(self, message: str, *args: object) -> None:
records.append(message.format(*args))
secret = "customer-token-do-not-log"
async def execute(_name, _args):
return "ok"
monkeypatch.setattr(runner_module, "_SLOW_TOOL_LOG_MS", 0)
monkeypatch.setattr(runner_module, "logger", _Logger())
runner = AgentRunner()
spec = make_run_spec(
MagicMock(spec=LLMProvider),
initial_messages=[],
tools=SimpleNamespace(execute=execute),
model="test-model",
max_iterations=1,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
)
await runner._run_tool(
spec,
ToolCallRequest(
id="call_1",
name="edit_file",
arguments={"old_text": secret, "paths": ["one", "two"]},
),
external_lookup_counts={},
workspace_violation_counts={},
)
assert len(records) == 1
assert "operation=edit_file" in records[0]
assert "input_items=4" in records[0]
assert f"input_chars={len(secret)}" in records[0]
assert "duration_ms=" in records[0]
assert secret not in records[0]
+76
View File
@@ -1233,6 +1233,32 @@ async def test_submitted_cron_turn_reports_pending_until_completed(tmp_path):
assert loop.pending_cron_job_ids_for_session(session_key) == set()
@pytest.mark.asyncio
async def test_accepted_automation_turn_preserves_wait_for_timeout(tmp_path):
"""Accepted-turn signalling must not leak through asyncio timeout semantics."""
from nanobot.bus.events import InboundMessage
from nanobot.cron.session_turns import CRON_TRIGGER_META
loop = _make_loop(tmp_path)
loop._running = True
msg = InboundMessage(
channel="websocket",
sender_id="cron",
chat_id="chat-1",
content="scheduled work",
metadata={CRON_TRIGGER_META: {"job_id": "job-1", "run_id": "run-timeout"}},
session_key_override="websocket:chat-1",
)
submit_task = asyncio.create_task(
asyncio.wait_for(loop.submit_cron_turn(msg), timeout=0.01)
)
assert await asyncio.wait_for(loop.bus.consume_inbound(), timeout=0.5) is msg
with pytest.raises(TimeoutError):
await submit_task
@pytest.mark.asyncio
async def test_submitted_local_trigger_turn_reports_pending_until_completed(tmp_path):
"""Local triggers remain marked pending while their session turn is in flight."""
@@ -1818,3 +1844,53 @@ async def test_injection_cycle_cap_on_error_path():
assert result.had_injections is True
# Should cap: _MAX_INJECTION_CYCLES drained rounds + 1 final round that breaks
assert call_count["n"] == _MAX_INJECTION_CYCLES + 1
@pytest.mark.asyncio
async def test_accepted_direct_automation_turn_keeps_single_effect_owner():
"""Cancelling a submitter after direct admission must not cancel or replay its work."""
from nanobot.agent.automation_turns import AutomationTurnAcceptedCancellation
from nanobot.agent.cron_turns import CronTurnCoordinator
from nanobot.bus.events import InboundMessage
from nanobot.cron.session_turns import CRON_TRIGGER_META
dispatch_started = asyncio.Event()
release_dispatch = asyncio.Event()
effect_completed = asyncio.Event()
effects: list[str] = []
async def dispatch(msg: InboundMessage) -> None:
dispatch_started.set()
await release_dispatch.wait()
effects.append(msg.content)
coordinator.complete(msg)
effect_completed.set()
coordinator = CronTurnCoordinator(
publish_inbound=lambda _msg: asyncio.sleep(0),
dispatch=dispatch,
is_running=lambda: False,
)
msg = InboundMessage(
channel="websocket",
sender_id="cron",
chat_id="chat-1",
content="scheduled work",
metadata={CRON_TRIGGER_META: {"job_id": "job-1", "run_id": "run-direct"}},
session_key_override="websocket:chat-1",
)
submit_task = asyncio.create_task(coordinator.submit(msg))
try:
await asyncio.wait_for(dispatch_started.wait(), timeout=0.5)
submit_task.cancel()
with pytest.raises(AutomationTurnAcceptedCancellation):
await submit_task
assert effects == []
release_dispatch.set()
await asyncio.wait_for(effect_completed.wait(), timeout=0.5)
assert effects == ["scheduled work"]
finally:
release_dispatch.set()
await asyncio.gather(submit_task, return_exceptions=True)
+14 -5
View File
@@ -3,7 +3,6 @@ import os
import subprocess
import time
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import MagicMock
import pytest
@@ -463,11 +462,21 @@ async def test_cli_app_scope_controls_working_dir(
seen: dict[str, str] = {}
def fake_run(argv, **kwargs):
seen["cwd"] = kwargs["cwd"]
return SimpleNamespace(returncode=0, stdout="ok", stderr="")
class _Process:
returncode = 0
monkeypatch.setattr("nanobot.apps.cli.service.subprocess.run", fake_run)
async def communicate(self) -> tuple[bytes, bytes]:
return b"ok", b""
async def fake_create_subprocess_exec(*_argv: str, **kwargs: object) -> _Process:
seen["cwd"] = str(kwargs["cwd"])
return _Process()
monkeypatch.setattr(
"nanobot.apps.cli.service.asyncio.create_subprocess_exec",
fake_create_subprocess_exec,
)
monkeypatch.setattr("nanobot.apps.cli.service._WindowsJob.create", lambda: None)
tool = CliAppsTool(
workspace=tmp_path,
restrict_to_workspace=True,
+60 -8
View File
@@ -3,6 +3,7 @@
from __future__ import annotations
import asyncio
import threading
from unittest.mock import AsyncMock, MagicMock
import pytest
@@ -21,9 +22,9 @@ from nanobot.agent.tools.long_task import (
from nanobot.agent.tools.registry import ToolRegistry
from nanobot.bus.outbound_events import GoalStateSyncEvent
from nanobot.bus.queue import MessageBus
from nanobot.bus.runtime_events import RuntimeEventBus
from nanobot.bus.runtime_events import GoalStateChanged, RuntimeEventBus
from nanobot.session.goal_state import GOAL_STATE_KEY, MAX_GOAL_OBJECTIVE_CHARS
from nanobot.session.manager import SessionManager
from nanobot.session.manager import Session, SessionManager
from nanobot.session.turn_continuation import should_finalize_on_max_iterations
from nanobot.session.webui_turns import WebuiTurnCoordinator
@@ -179,13 +180,13 @@ async def test_goal_state_mutations_roll_back_on_save_failure(tmp_path, monkeypa
sess = sm.get_or_create("websocket:c1")
sess.metadata["marker"] = {"keep": True}
sess.metadata["_sustained_goal_continuation_rounds"] = 12
original_save = sm.save
original_save_async = sm.save_async
create_context = _request_context()
def fail_save(_session, **_kwargs):
async def fail_save(_session, **_kwargs):
raise OSError("disk unavailable")
monkeypatch.setattr(sm, "save", fail_save)
monkeypatch.setattr(sm, "save_async", fail_save)
with pytest.raises(OSError, match="disk unavailable"):
await _execute(create, create_context, objective="Old")
@@ -195,13 +196,13 @@ async def test_goal_state_mutations_roll_back_on_save_failure(tmp_path, monkeypa
}
assert GOAL_STATE_KEY not in SessionManager(tmp_path).get_or_create("websocket:c1").metadata
monkeypatch.setattr(sm, "save", original_save)
monkeypatch.setattr(sm, "save_async", original_save_async)
assert "Goal recorded" in await _execute(create, create_context, objective="Old")
sess.metadata["_sustained_goal_continuation_rounds"] = 12
sm.save(sess)
replace_context = _request_context()
monkeypatch.setattr(sm, "save", fail_save)
monkeypatch.setattr(sm, "save_async", fail_save)
with pytest.raises(OSError, match="disk unavailable"):
await _execute(update, replace_context, action="replace", objective="New")
@@ -211,7 +212,7 @@ async def test_goal_state_mutations_roll_back_on_save_failure(tmp_path, monkeypa
assert persisted[GOAL_STATE_KEY]["objective"] == "Old"
assert persisted["_sustained_goal_continuation_rounds"] == 12
monkeypatch.setattr(sm, "save", original_save)
monkeypatch.setattr(sm, "save_async", original_save_async)
assert "Goal replaced" in await _execute(
update,
replace_context,
@@ -225,6 +226,57 @@ async def test_goal_state_mutations_roll_back_on_save_failure(tmp_path, monkeypa
)
@pytest.mark.asyncio
async def test_sync_goal_save_cancellation_settles_state_and_runtime_event() -> None:
started = threading.Event()
release = threading.Event()
session = Session(key="websocket:c1")
durable_metadata: dict[str, object] = {}
save_calls = 0
class _SyncSessionManager:
def get_or_create(self, key: str) -> Session:
assert key == session.key
return session
def save(self, target: Session) -> None:
nonlocal durable_metadata, save_calls
save_calls += 1
started.set()
assert release.wait(timeout=1)
durable_metadata = dict(target.metadata)
runtime_events = RuntimeEventBus()
published: list[GoalStateChanged] = []
runtime_events.subscribe(published.append, GoalStateChanged)
create = CreateGoalTool(
sessions=_SyncSessionManager(), # type: ignore[arg-type]
runtime_events=runtime_events,
)
task = asyncio.create_task(
_execute(create, _request_context(), objective="Persist through cancellation")
)
assert await asyncio.to_thread(started.wait, 1)
try:
task.cancel()
await asyncio.sleep(0)
assert not task.done()
finally:
release.set()
with pytest.raises(asyncio.CancelledError):
await asyncio.wait_for(task, timeout=1)
assert save_calls == 1
assert session.metadata[GOAL_STATE_KEY]["objective"] == "Persist through cancellation"
assert durable_metadata[GOAL_STATE_KEY] == session.metadata[GOAL_STATE_KEY]
assert len(published) == 1
assert published[0].session_metadata[GOAL_STATE_KEY] == session.metadata[GOAL_STATE_KEY]
await asyncio.sleep(0.05)
assert save_calls == 1
assert len(published) == 1
@pytest.mark.asyncio
async def test_goal_tools_reject_oversized_objectives(tmp_path):
sm = SessionManager(tmp_path)
+13 -9
View File
@@ -42,17 +42,21 @@ def test_run_passes_filtered_env(monkeypatch, tmp_path) -> None:
manager = CliAppManager(workspace=tmp_path, data_dir=tmp_path / "cli-apps")
captured: dict[str, object] = {}
def fake_run(*args, **kwargs):
class _Process:
returncode = 0
def communicate(self, *, timeout: int) -> tuple[str, str]:
assert timeout == 60
return "ok", ""
def fake_popen(argv: list[str], **kwargs: object) -> _Process:
assert argv == ["/bin/echo", "hi"]
assert kwargs["text"] is True
captured.update(kwargs)
return _Process()
class Result:
returncode = 0
stdout = "ok"
stderr = ""
return Result()
monkeypatch.setattr("nanobot.apps.cli.service.subprocess.run", fake_run)
monkeypatch.setattr("nanobot.apps.cli.service.subprocess.Popen", fake_popen)
monkeypatch.setattr("nanobot.apps.cli.service._WindowsJob.create", lambda: None)
monkeypatch.setattr(manager, "get_app", lambda name: {"name": name, "entry_point": "echo"})
monkeypatch.setattr(
manager,
+113 -1
View File
@@ -8,10 +8,15 @@ block the stop.
"""
import asyncio
import threading
import time
from contextlib import suppress
from nanobot.cli.gateway_runtime import _close_gateway_runtime
from nanobot.cli.gateway_runtime import (
_call_session_manager,
_close_gateway_runtime,
_monitor_event_loop_lag,
)
class _FakeAgent:
@@ -216,3 +221,110 @@ async def test_cancelled_runtime_tasks_gather_does_not_raise() -> None:
assert runtime_tasks.done() # the cancelled gather was awaited without raising
assert agent.close_calls == 1
assert provider.close_calls == 1
async def test_session_manager_call_prefers_class_declared_coroutine() -> None:
class _Manager:
async def list_sessions_async(self) -> list[str]:
return ["native-async"]
def list_sessions(self) -> list[str]:
raise AssertionError("sync fallback should not run")
manager = _Manager()
result = await _call_session_manager(
manager,
"list_sessions_async",
manager.list_sessions,
)
assert result == ["native-async"]
async def test_session_manager_call_offloads_sync_compatibility_fallback() -> None:
calling_thread = threading.get_ident()
sync_threads: list[int] = []
class _Manager:
def list_sessions(self) -> list[str]:
sync_threads.append(threading.get_ident())
return ["sync-fallback"]
manager = _Manager()
async def _fabricated_async() -> list[str]:
raise AssertionError("instance-only async stand-in should not run")
setattr(manager, "list_sessions_async", _fabricated_async)
result = await _call_session_manager(
manager,
"list_sessions_async",
manager.list_sessions,
)
assert result == ["sync-fallback"]
assert sync_threads and sync_threads[0] != calling_thread
async def test_session_manager_sync_fallback_cancellation_waits_for_worker() -> None:
started = threading.Event()
release = threading.Event()
finished = threading.Event()
mutations: list[str] = []
class _Manager:
def save(self, value: str) -> str:
started.set()
assert release.wait(timeout=1)
mutations.append(value)
finished.set()
return "saved"
manager = _Manager()
task = asyncio.create_task(
_call_session_manager(manager, "save_async", manager.save, "mutation")
)
assert await asyncio.to_thread(started.wait, 1)
try:
task.cancel()
await asyncio.sleep(0)
assert not task.done()
finally:
release.set()
with suppress(asyncio.CancelledError):
await asyncio.wait_for(task, timeout=1)
assert task.cancelled()
assert finished.is_set()
assert mutations == ["mutation"]
await asyncio.sleep(0.05)
assert mutations == ["mutation"]
async def test_event_loop_lag_monitor_logs_gateway_scheduler_drift() -> None:
records: list[str] = []
class _Logger:
def warning(self, message: str, *args: object) -> None:
records.append(message.format(*args))
task = asyncio.create_task(
_monitor_event_loop_lag(
interval_s=0.01,
warning_threshold_s=0.015,
log=_Logger(),
)
)
await asyncio.sleep(0)
time.sleep(0.04)
await asyncio.sleep(0.02)
task.cancel()
with suppress(asyncio.CancelledError):
await task
assert records
assert "operation=gateway" in records[0]
assert "duration_ms=" in records[0]
assert "interval_ms=10" in records[0]
+168 -12
View File
@@ -1,9 +1,11 @@
from __future__ import annotations
import asyncio
import json
import subprocess
import sys
import time
from contextlib import suppress
from pathlib import Path
from types import SimpleNamespace
@@ -944,19 +946,23 @@ def test_run_installed_cli_uses_argv_without_shell(
lambda entry: resolved if entry == "cli-anything-gimp" else None,
)
def fake_run(argv: list[str], **kwargs: object) -> subprocess.CompletedProcess[str]:
assert "shell" not in kwargs or kwargs["shell"] is False
class _Process:
returncode = 0
def communicate(self, *, timeout: int) -> tuple[str, str]:
assert timeout == 5
return "ARGS=['--json', 'project', 'list']", ""
def fake_popen(argv: list[str], **kwargs: object) -> _Process:
assert argv == [resolved, "--json", "project", "list"]
assert "shell" not in kwargs
assert kwargs["text"] is True
assert kwargs["encoding"] == "utf-8"
assert kwargs["errors"] == "replace"
return subprocess.CompletedProcess(
argv,
0,
stdout="ARGS=" + repr(argv[1:]),
stderr="",
)
return _Process()
monkeypatch.setattr("nanobot.apps.cli.service.subprocess.run", fake_run)
monkeypatch.setattr("nanobot.apps.cli.service.subprocess.Popen", fake_popen)
monkeypatch.setattr("nanobot.apps.cli.service._WindowsJob.create", lambda: None)
manager._save_installed(
{
"gimp": {
@@ -986,12 +992,21 @@ def test_run_reports_created_artifacts(
lambda entry: resolved if entry == "cli-anything-gimp" else None,
)
def fake_run(argv: list[str], **kwargs: object) -> subprocess.CompletedProcess[str]:
class _Process:
returncode = 0
def communicate(self, *, timeout: int) -> tuple[str, str]:
assert timeout == 5
return "done", ""
def fake_popen(argv: list[str], **kwargs: object) -> _Process:
assert argv == [resolved, "render"]
cwd = Path(str(kwargs["cwd"]))
(cwd / "diagram.png").write_bytes(b"\x89PNG\r\n\x1a\nimage")
return subprocess.CompletedProcess(argv, 0, stdout="done", stderr="")
return _Process()
monkeypatch.setattr("nanobot.apps.cli.service.subprocess.run", fake_run)
monkeypatch.setattr("nanobot.apps.cli.service.subprocess.Popen", fake_popen)
monkeypatch.setattr("nanobot.apps.cli.service._WindowsJob.create", lambda: None)
manager._save_installed({"gimp": {"entry_point": "cli-anything-gimp"}})
result = manager.run("gimp", ["render"])
@@ -1100,3 +1115,144 @@ def test_uninstall_uses_uv_pip_when_pip_unavailable(
sys.executable,
"suno-cli",
]
@pytest.mark.asyncio
async def test_run_async_keeps_event_loop_responsive(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
manager = _manager(tmp_path)
_seed_catalog(manager)
manager._save_installed({"gimp": {"entry_point": "python"}})
monkeypatch.setattr(
"nanobot.apps.cli.service.shutil.which",
lambda entry: sys.executable if entry == "python" else None,
)
task = asyncio.create_task(
manager.run_async(
"gimp",
["-c", "import time; time.sleep(0.08); print('done')"],
)
)
for _ in range(3):
await asyncio.sleep(0.01)
assert not task.done()
result = await asyncio.wait_for(task, timeout=2)
assert "CLI app 'gimp' exited 0" in result
assert "done" in result
@pytest.mark.asyncio
async def test_run_async_cancellation_terminates_process_tree(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
manager = _manager(tmp_path)
_seed_catalog(manager)
manager._save_installed({"gimp": {"entry_point": "python"}})
monkeypatch.setattr(
"nanobot.apps.cli.service.shutil.which",
lambda entry: sys.executable if entry == "python" else None,
)
started = manager.workspace / "parent-started.txt"
orphan_marker = manager.workspace / "orphan-survived.txt"
child_code = (
"import time; from pathlib import Path; time.sleep(0.5); "
f"Path({str(orphan_marker)!r}).write_text('orphan', encoding='utf-8')"
)
parent_code = (
"import subprocess, sys, time; from pathlib import Path; "
f"subprocess.Popen([sys.executable, '-c', {child_code!r}]); "
f"Path({str(started)!r}).write_text('started', encoding='utf-8'); "
"time.sleep(30)"
)
task = asyncio.create_task(manager.run_async("gimp", ["-c", parent_code]))
try:
for _ in range(200):
if started.exists():
break
await asyncio.sleep(0.01)
assert started.exists()
task.cancel()
with pytest.raises(asyncio.CancelledError):
await asyncio.wait_for(task, timeout=3)
await asyncio.sleep(0.8)
assert not orphan_marker.exists()
finally:
if not task.done():
task.cancel()
with suppress(asyncio.CancelledError):
await task
def _parent_exits_with_pipe_holding_descendant(started: Path, marker: Path) -> str:
child_code = (
"import time; from pathlib import Path; time.sleep(1.5); "
f"Path({str(marker)!r}).write_text('survived', encoding='utf-8')"
)
return (
"import subprocess, sys, time; from pathlib import Path; "
"time.sleep(0.25); "
f"subprocess.Popen([sys.executable, '-c', {child_code!r}]); "
f"Path({str(started)!r}).write_text('spawned', encoding='utf-8')"
)
@pytest.mark.asyncio
async def test_run_async_timeout_kills_pipe_holding_descendant_after_parent_exit(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
manager = _manager(tmp_path)
_seed_catalog(manager)
manager._save_installed({"gimp": {"entry_point": "python"}})
monkeypatch.setattr(
"nanobot.apps.cli.service.shutil.which",
lambda entry: sys.executable if entry == "python" else None,
)
started = manager.workspace / "descendant-started.txt"
marker = manager.workspace / "descendant-survived.txt"
parent_code = _parent_exits_with_pipe_holding_descendant(started, marker)
before = time.monotonic()
result = await asyncio.wait_for(
manager.run_async("gimp", ["-c", parent_code], timeout=1),
timeout=4,
)
elapsed = time.monotonic() - before
assert started.exists()
assert result == "CLI app 'gimp' timed out after 1s"
assert elapsed < 4
await asyncio.sleep(0.9)
assert not marker.exists()
def test_run_timeout_kills_pipe_holding_descendant_after_parent_exit(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
manager = _manager(tmp_path)
_seed_catalog(manager)
manager._save_installed({"gimp": {"entry_point": "python"}})
monkeypatch.setattr(
"nanobot.apps.cli.service.shutil.which",
lambda entry: sys.executable if entry == "python" else None,
)
started = manager.workspace / "sync-descendant-started.txt"
marker = manager.workspace / "sync-descendant-survived.txt"
parent_code = _parent_exits_with_pipe_holding_descendant(started, marker)
before = time.monotonic()
result = manager.run("gimp", ["-c", parent_code], timeout=1)
elapsed = time.monotonic() - before
assert started.exists()
assert result == "CLI app 'gimp' timed out after 1s"
assert elapsed < 4
time.sleep(0.9)
assert not marker.exists()
+106 -10
View File
@@ -2,10 +2,13 @@ from __future__ import annotations
import asyncio
import json
import subprocess
import threading
import time
from pathlib import Path
import pytest
from nanobot.agent.tools import cli_apps as cli_apps_tool
from nanobot.agent.tools.cli_apps import CliAppsTool
from nanobot.agent.tools.context import RequestContext
from nanobot.apps.cli.service import CliAppManager, CliAppsRuntimeConfig
@@ -52,16 +55,22 @@ def test_run_cli_app_uses_installed_registry_app(
lambda entry: resolved if entry == "cli-anything-gimp" else None,
)
def fake_run(argv: list[str], **kwargs: object) -> subprocess.CompletedProcess[str]:
assert "shell" not in kwargs or kwargs["shell"] is False
return subprocess.CompletedProcess(
argv,
0,
stdout="tool:" + " ".join(argv[1:]),
stderr="",
)
class _Process:
returncode = 0
monkeypatch.setattr("nanobot.apps.cli.service.subprocess.run", fake_run)
async def communicate(self) -> tuple[bytes, bytes]:
return b"tool:--json project list", b""
async def fake_create_subprocess_exec(*argv: str, **kwargs: object) -> _Process:
assert argv == (resolved, "--json", "project", "list")
assert "shell" not in kwargs
return _Process()
monkeypatch.setattr(
"nanobot.apps.cli.service.asyncio.create_subprocess_exec",
fake_create_subprocess_exec,
)
monkeypatch.setattr("nanobot.apps.cli.service._WindowsJob.create", lambda: None)
monkeypatch.setattr("nanobot.apps.cli.service.get_runtime_subdir", lambda _name: data_dir)
tool = CliAppsTool(
@@ -156,3 +165,90 @@ def test_cli_app_tool_provides_context_only_for_attachment(tmp_path: Path) -> No
assert attached is not None
assert attached.source == "cli_apps"
assert "CLI App Attachment: @drawio" in attached.content
@pytest.mark.asyncio
async def test_run_cli_app_uses_threaded_sync_fallback_for_compatible_manager(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
event_loop_thread = threading.get_ident()
seen: dict[str, object] = {}
class CompatibleManager:
def __init__(
self,
*,
workspace: Path,
runtime: CliAppsRuntimeConfig,
) -> None:
seen["workspace"] = workspace
seen["runtime"] = runtime
def run_async(self, *_args: object, **_kwargs: object) -> str:
raise AssertionError("a synchronous compatibility method must not be awaited")
def run(self, name: str, **kwargs: object) -> str:
seen["thread"] = threading.get_ident()
seen["name"] = name
seen["kwargs"] = kwargs
return "compatible result"
monkeypatch.setattr(cli_apps_tool, "CliAppManager", CompatibleManager)
runtime = CliAppsRuntimeConfig(run_timeout=7)
tool = CliAppsTool(workspace=tmp_path, restrict_to_workspace=True, runtime=runtime)
result = await tool.execute(
name="demo",
args=["show"],
json=True,
working_dir=str(tmp_path),
timeout=3,
)
assert result == "compatible result"
assert seen["workspace"] == tmp_path
assert seen["runtime"] is runtime
assert seen["thread"] != event_loop_thread
assert seen["name"] == "demo"
assert seen["kwargs"] == {
"args": ["show"],
"json_output": True,
"working_dir": str(tmp_path),
"timeout": 3,
"restrict_to_workspace": True,
}
@pytest.mark.asyncio
async def test_run_cli_app_preserves_native_manager_cancellation(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
started = asyncio.Event()
cancelled = asyncio.Event()
async def fake_run_async(self: CliAppManager, name: str, **_kwargs: object) -> str:
assert name == "demo"
started.set()
try:
await asyncio.Future()
except asyncio.CancelledError:
cancelled.set()
raise
raise AssertionError("unreachable")
def fail_sync_run(self: CliAppManager, name: str, **_kwargs: object) -> str:
raise AssertionError(f"native async manager fell back to run({name!r})")
monkeypatch.setattr(CliAppManager, "run_async", fake_run_async)
monkeypatch.setattr(CliAppManager, "run", fail_sync_run)
tool = CliAppsTool(workspace=tmp_path)
task = asyncio.create_task(tool.execute(name="demo"))
await asyncio.wait_for(started.wait(), timeout=1)
task.cancel()
with pytest.raises(asyncio.CancelledError):
await task
assert cancelled.is_set()
+53
View File
@@ -2,6 +2,8 @@
from __future__ import annotations
import asyncio
import threading
from inspect import Parameter, signature
from unittest.mock import AsyncMock, MagicMock
@@ -9,9 +11,11 @@ import pytest
from nanobot.command.builtin import (
builtin_command_starts_agent_turn,
cmd_new,
register_builtin_commands,
)
from nanobot.command.router import CommandContext, CommandRouter
from nanobot.session.manager import Session
def test_command_context_requires_loop_as_keyword_dependency() -> None:
@@ -95,6 +99,55 @@ def test_builtin_command_agent_turn_lifecycle(content: str, expected: bool) -> N
assert builtin_command_starts_agent_turn(content) is expected
@pytest.mark.asyncio
async def test_new_cancellation_waits_for_save_and_cache_invalidation() -> None:
started = threading.Event()
release = threading.Event()
session = Session(key="test:chat1")
invalidated: list[str] = []
save_calls = 0
class _SyncSessionManager:
def save(self, target: Session) -> None:
nonlocal save_calls
assert target is session
save_calls += 1
started.set()
assert release.wait(timeout=1)
def invalidate(self, key: str) -> None:
invalidated.append(key)
loop = MagicMock()
loop.sessions = _SyncSessionManager()
loop._cancel_active_tasks = AsyncMock(return_value=0)
ctx = CommandContext(
msg=MagicMock(channel="test", chat_id="chat1", metadata={}),
session=session,
key=session.key,
raw="/new",
loop=loop,
)
task = asyncio.create_task(cmd_new(ctx))
assert await asyncio.to_thread(started.wait, 1)
try:
task.cancel()
await asyncio.sleep(0)
assert not task.done()
assert invalidated == []
finally:
release.set()
with pytest.raises(asyncio.CancelledError):
await asyncio.wait_for(task, timeout=1)
assert save_calls == 1
assert invalidated == [session.key]
await asyncio.sleep(0.05)
assert save_calls == 1
assert invalidated == [session.key]
class TestMidTurnCommandDispatchedDirectly:
"""Verify that commands matching is_dispatchable_command() are dispatched
correctly when session=None (the mid-turn path)."""
+125
View File
@@ -0,0 +1,125 @@
from __future__ import annotations
import asyncio
import json
from contextlib import suppress
from pathlib import Path
import pytest
from nanobot.agent.automation_turns import AutomationTurnAcceptedCancellation
from nanobot.agent.cron_turns import CronTurnCoordinator
from nanobot.bus.events import InboundMessage, OutboundMessage
from nanobot.cron.bound_runner import run_bound_cron_job
from nanobot.cron.service import CronService
from nanobot.cron.types import CronJob, CronSchedule
class _NoCronToolRegistry:
def get(self, _name: str) -> None:
return None
class _CoordinatedAgent:
def __init__(self) -> None:
self.tools = _NoCronToolRegistry()
self.inbound: asyncio.Queue[InboundMessage] = asyncio.Queue()
self.coordinator = CronTurnCoordinator(
publish_inbound=self.inbound.put,
dispatch=lambda _msg: asyncio.sleep(0),
is_running=lambda: True,
)
async def submit_cron_turn(
self,
msg: InboundMessage,
) -> OutboundMessage | None:
return await self.coordinator.submit(msg)
@pytest.mark.asyncio
async def test_accepted_bound_turn_cancellation_advances_durable_schedule(
tmp_path: Path,
) -> None:
store_path = tmp_path / "cron" / "jobs.json"
agent = _CoordinatedAgent()
service = CronService(store_path)
job = service.add_job(
name="accepted reminder",
schedule=CronSchedule(kind="every", every_ms=60_000),
message="check accepted work",
session_key="websocket:chat-1",
origin_channel="websocket",
origin_chat_id="chat-1",
)
async def on_job(current: CronJob) -> str | None:
return await run_bound_cron_job(current, agent=agent, cron=service)
service.on_job = on_job
work_accepted = asyncio.Event()
release_work = asyncio.Event()
effects: list[str] = []
async def agent_worker() -> None:
msg = await agent.inbound.get()
work_accepted.set()
await release_work.wait()
effects.append(msg.content)
agent.coordinator.complete(msg)
worker = asyncio.create_task(agent_worker())
run = asyncio.create_task(service.run_job(job.id))
try:
await asyncio.wait_for(work_accepted.wait(), timeout=1)
run.cancel()
with pytest.raises(AutomationTurnAcceptedCancellation):
await asyncio.wait_for(run, timeout=1)
# Cancellation is prompt and the independently-owned agent turn can
# finish later, while the cron schedule is already durable.
assert effects == []
persisted = CronService(store_path).get_job(job.id)
assert persisted is not None
assert persisted.state.last_status == "ok"
assert persisted.state.last_error is None
assert persisted.state.last_run_at_ms is not None
assert persisted.state.next_run_at_ms is not None
assert persisted.state.next_run_at_ms > persisted.state.last_run_at_ms
assert len(persisted.state.run_history) == 1
assert persisted.state.run_history[0].status == "ok"
records = [
json.loads(path.read_text(encoding="utf-8"))
for path in (store_path.parent / "runs").glob("*.json")
]
assert len(records) == 1
assert records[0]["status"] == "accepted"
replayed: list[str] = []
async def replay_job(current: CronJob) -> None:
replayed.append(current.id)
restarted = CronService(store_path, on_job=replay_job)
restarted._running = True
restarted._arm_timer = lambda: None
try:
await restarted._on_timer()
finally:
restarted.stop()
assert replayed == []
release_work.set()
await asyncio.wait_for(worker, timeout=1)
assert effects == [records[0]["rendered_prompt"]]
finally:
release_work.set()
if not run.done():
run.cancel()
if not worker.done():
worker.cancel()
with suppress(asyncio.CancelledError):
await run
with suppress(asyncio.CancelledError):
await worker
+226
View File
@@ -1,5 +1,6 @@
import asyncio
import json
import threading
import time
from pathlib import Path
@@ -653,6 +654,231 @@ async def test_manual_run_persists_completion_when_callback_lists_jobs(tmp_path)
assert state["runHistory"][0]["status"] == "ok"
@pytest.mark.asyncio
async def test_same_job_manual_runs_do_not_overlap(tmp_path) -> None:
entered = asyncio.Event()
release = asyncio.Event()
calls: list[str] = []
async def on_job(job: CronJob) -> None:
calls.append(job.id)
entered.set()
await release.wait()
service = CronService(tmp_path / "cron" / "jobs.json", on_job=on_job)
job = service.add_job(
name="single-flight",
schedule=CronSchedule(kind="every", every_ms=60_000),
message="hello",
**_bound_chat(),
)
first = asyncio.create_task(service.run_job(job.id))
try:
await asyncio.wait_for(entered.wait(), timeout=1)
assert await service.run_job(job.id) is False
assert calls == [job.id]
finally:
release.set()
assert await first is True
loaded = service.get_job(job.id)
assert loaded is not None
assert len(loaded.state.run_history) == 1
@pytest.mark.asyncio
async def test_timer_and_manual_run_of_same_job_do_not_overlap(tmp_path) -> None:
entered = asyncio.Event()
release = asyncio.Event()
calls: list[str] = []
async def on_job(job: CronJob) -> None:
calls.append(job.id)
entered.set()
await release.wait()
service = CronService(tmp_path / "cron" / "jobs.json", on_job=on_job)
service._running = True
service._arm_timer = lambda: None
service._load_store()
job = service.add_job(
name="timer-single-flight",
schedule=CronSchedule(kind="every", every_ms=60_000),
message="hello",
**_bound_chat(),
)
job.state.next_run_at_ms = max(1, int(time.time() * 1000) - 1_000)
service._save_store()
timer = asyncio.create_task(service._on_timer())
try:
await asyncio.wait_for(entered.wait(), timeout=1)
assert await service.run_job(job.id, force=True) is False
assert calls == [job.id]
finally:
release.set()
await timer
service.stop()
loaded = service.get_job(job.id)
assert loaded is not None
assert len(loaded.state.run_history) == 1
@pytest.mark.asyncio
async def test_start_cancellation_waits_for_load_and_fully_starts(
tmp_path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
store_path = tmp_path / "cron" / "jobs.json"
service = CronService(store_path)
job = service.add_job(
name="load-settlement",
schedule=CronSchedule(kind="every", every_ms=60_000),
message="hello",
**_bound_chat(),
)
entered = threading.Event()
release = threading.Event()
original_load_store = service._load_store
def blocking_load_store():
entered.set()
if not release.wait(timeout=2):
raise TimeoutError("test did not release cron load worker")
return original_load_store()
monkeypatch.setattr(service, "_load_store", blocking_load_store)
start = asyncio.create_task(service.start())
assert await asyncio.to_thread(entered.wait, 1)
try:
assert start.cancel()
await asyncio.sleep(0)
assert not start.done()
assert not store_path.exists()
release.set()
with pytest.raises(asyncio.CancelledError):
await asyncio.wait_for(start, timeout=1)
assert service._running is True
timer = service._timer_task
assert timer is not None and not timer.done()
assert service._store_dirty is False
assert (tmp_path / "cron" / "action.jsonl").read_text(encoding="utf-8") == ""
persisted = store_path.read_bytes()
assert json.loads(persisted)["jobs"][0]["id"] == job.id
await asyncio.sleep(0.01)
assert store_path.read_bytes() == persisted
assert service._timer_task is timer
finally:
release.set()
await asyncio.gather(start, return_exceptions=True)
timer = service._timer_task
service.stop()
if timer is not None:
await asyncio.gather(timer, return_exceptions=True)
@pytest.mark.asyncio
async def test_start_repeated_cancellation_waits_for_persistence_and_fully_starts(
tmp_path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
store_path = tmp_path / "cron" / "jobs.json"
service = CronService(store_path)
entered = threading.Event()
release = threading.Event()
original_save_store = service._save_store
def blocking_save_store() -> None:
entered.set()
if not release.wait(timeout=2):
raise TimeoutError("test did not release cron persistence worker")
original_save_store()
monkeypatch.setattr(service, "_save_store", blocking_save_store)
start = asyncio.create_task(service.start())
assert await asyncio.to_thread(entered.wait, 1)
try:
assert start.cancel()
await asyncio.sleep(0)
assert not start.done()
assert start.cancel()
await asyncio.sleep(0)
assert not start.done()
assert not store_path.exists()
release.set()
with pytest.raises(asyncio.CancelledError):
await asyncio.wait_for(start, timeout=1)
assert service._running is True
timer = service._timer_task
assert timer is not None and not timer.done()
assert service._store_dirty is False
persisted = store_path.read_bytes()
assert json.loads(persisted) == {"version": 1, "jobs": []}
await asyncio.sleep(0.01)
assert store_path.read_bytes() == persisted
assert service._timer_task is timer
finally:
release.set()
await asyncio.gather(start, return_exceptions=True)
timer = service._timer_task
service.stop()
if timer is not None:
await asyncio.gather(timer, return_exceptions=True)
@pytest.mark.asyncio
async def test_run_sync_cancellation_drains_worker_under_transaction_lock(tmp_path) -> None:
service = CronService(tmp_path / "cron" / "jobs.json")
entered = threading.Event()
release = threading.Event()
second_entered = threading.Event()
mutations: list[str] = []
def blocking_mutation() -> str:
entered.set()
if not release.wait(timeout=2):
raise TimeoutError("test did not release cron worker")
mutations.append("first")
return "first"
def later_mutation() -> str:
second_entered.set()
mutations.append("second")
return "second"
first = asyncio.create_task(service.run_sync(blocking_mutation))
assert await asyncio.to_thread(entered.wait, 1)
first.cancel()
second = asyncio.create_task(service.run_sync(later_mutation))
try:
await asyncio.sleep(0.05)
assert not first.done()
assert not second_entered.is_set()
release.set()
with pytest.raises(asyncio.CancelledError):
await asyncio.wait_for(first, timeout=1)
assert mutations == ["first"]
assert await asyncio.wait_for(second, timeout=1) == "second"
assert mutations == ["first", "second"]
await asyncio.sleep(0.01)
assert mutations == ["first", "second"]
finally:
release.set()
await asyncio.gather(first, second, return_exceptions=True)
@pytest.mark.asyncio
async def test_overlapping_manual_runs_preserve_stopped_service_state(tmp_path) -> None:
store_path = tmp_path / "cron" / "jobs.json"
+4 -2
View File
@@ -9,6 +9,7 @@ and tightens the runtime error for ``add`` without ``message``.
from __future__ import annotations
import asyncio
from collections.abc import Iterator
import pytest
@@ -21,6 +22,9 @@ from nanobot.agent.tools.registry import ToolRegistry
class _SvcStub:
"""Minimal CronService stand-in; we only exercise schema/dispatch paths."""
async def run_sync(self, operation, /, *args, **kwargs):
return await asyncio.to_thread(operation, *args, **kwargs)
def list_jobs(self):
return []
@@ -74,8 +78,6 @@ class TestSchemaContract:
# Schema permits omitting message; the runtime must return a message
# that tells the LLM exactly what's missing and how to retry, so it
# doesn't loop like #3113 reports.
import asyncio
tool = registry._tools["cron"] # type: ignore[attr-defined]
out = asyncio.run(tool.execute(action="add", at="2030-01-01T00:00:00"))
assert "message" in out
+82
View File
@@ -0,0 +1,82 @@
"""Worker-thread regressions for live CronTool mutations."""
import asyncio
import json
import threading
import pytest
from nanobot.agent.tools.context import RequestContext, request_context
from nanobot.agent.tools.cron import CronTool
from nanobot.cron.service import CronService
@pytest.mark.asyncio
async def test_cron_tool_add_on_started_service_is_durable_and_rearms_owner_loop(
tmp_path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
store_path = tmp_path / "cron" / "jobs.json"
service = CronService(store_path, max_sleep_ms=60_000)
await service.start()
owner_thread_id = threading.get_ident()
initial_timer = service._timer_task
worker_thread_ids: list[int] = []
arm_thread_ids: list[int] = []
timer_rearmed = asyncio.Event()
tool = CronTool(service)
original_execute_sync = tool._execute_sync
def tracked_execute_sync(*args, **kwargs):
worker_thread_ids.append(threading.get_ident())
return original_execute_sync(*args, **kwargs)
original_arm_timer = service._arm_timer
def tracked_arm_timer() -> None:
arm_thread_ids.append(threading.get_ident())
original_arm_timer()
timer_rearmed.set()
monkeypatch.setattr(tool, "_execute_sync", tracked_execute_sync)
monkeypatch.setattr(service, "_arm_timer", tracked_arm_timer)
try:
with request_context(
RequestContext(
channel="websocket",
chat_id="acceptance-chat",
session_key="websocket:acceptance-chat",
)
):
result = await tool.execute(
action="add",
name="Acceptance reminder",
message="Check the acceptance result",
every_seconds=3600,
)
await asyncio.wait_for(timer_rearmed.wait(), timeout=1)
assert result.startswith("Created job 'Acceptance reminder'")
assert "no running event loop" not in result.lower()
assert worker_thread_ids
assert all(thread_id != owner_thread_id for thread_id in worker_thread_ids)
assert arm_thread_ids and set(arm_thread_ids) == {owner_thread_id}
assert service._timer_task is not None
assert service._timer_task is not initial_timer
assert not service._timer_task.done()
stored = json.loads(store_path.read_text(encoding="utf-8"))
assert len(stored["jobs"]) == 1
assert stored["jobs"][0]["name"] == "Acceptance reminder"
assert stored["jobs"][0]["payload"]["sessionKey"] == "websocket:acceptance-chat"
reloaded_jobs = CronService(store_path).list_jobs(include_disabled=True)
assert len(reloaded_jobs) == 1
assert reloaded_jobs[0].name == "Acceptance reminder"
finally:
service.stop()
await asyncio.sleep(0)
+80
View File
@@ -2,13 +2,18 @@
from __future__ import annotations
import asyncio
import ipaddress
import socket
import threading
import time
from unittest.mock import patch
import pytest
from nanobot.security.network import (
async_resolve_url_target,
async_validate_resolved_url,
configure_ssrf_whitelist,
contains_internal_url,
env_proxy_applies_to_url,
@@ -362,3 +367,78 @@ def test_whitelist_allows_ipv6_mapped_cgnat():
assert ok, f"Whitelisted IPv6-mapped CGNAT should be allowed, got: {err}"
finally:
configure_ssrf_whitelist([])
@pytest.mark.asyncio
async def test_async_dns_resolution_keeps_event_loop_responsive(
monkeypatch: pytest.MonkeyPatch,
) -> None:
started = threading.Event()
def slow_resolver(hostname, port, family=0, type=0, proto=0, flags=0): # noqa: A002
started.set()
time.sleep(0.08)
return [
(socket.AF_INET, socket.SOCK_STREAM, 0, "", ("93.184.216.34", 0)),
]
monkeypatch.setattr(socket, "getaddrinfo", slow_resolver)
task = asyncio.create_task(async_resolve_url_target("https://example.com/path"))
assert await asyncio.to_thread(started.wait, 0.5)
for _ in range(3):
await asyncio.sleep(0.01)
assert not task.done()
assert await asyncio.wait_for(task, timeout=0.5) == (
True,
"",
("93.184.216.34",),
)
@pytest.mark.asyncio
async def test_async_dns_validation_preserves_sync_security_results(
monkeypatch: pytest.MonkeyPatch,
) -> None:
def resolver(hostname, port, family=0, type=0, proto=0, flags=0): # noqa: A002
assert hostname == "evil.com"
return [
(socket.AF_INET, socket.SOCK_STREAM, 0, "", ("169.254.169.254", 0)),
]
monkeypatch.setattr(socket, "getaddrinfo", resolver)
sync_result = resolve_url_target("http://evil.com/latest")
async_result = await async_resolve_url_target("http://evil.com/latest")
redirect_result = await async_validate_resolved_url("http://evil.com/latest")
assert async_result == sync_result
assert async_result[0] is False
assert redirect_result[0] is False
@pytest.mark.asyncio
async def test_async_dns_resolution_times_out(
monkeypatch: pytest.MonkeyPatch,
) -> None:
loop = asyncio.get_running_loop()
never_resolves = asyncio.Event()
async def stalled_getaddrinfo(*args, **kwargs):
await never_resolves.wait()
monkeypatch.setattr(loop, "getaddrinfo", stalled_getaddrinfo)
result = await asyncio.wait_for(
async_resolve_url_target(
"https://example.com/path",
timeout_s=0.01,
),
timeout=0.2,
)
assert result == (
False,
"Timed out resolving hostname: example.com",
(),
)
+4 -4
View File
@@ -178,13 +178,13 @@ async def test_scan_loads_only_sessions_that_need_webui_recovery(
coordinator, _, restarted = _coordinator(tmp_path)
loaded: list[str] = []
get_or_create = restarted.get_or_create
get_or_create_async = restarted.get_or_create_async
def tracked_get_or_create(key: str) -> Session:
async def tracked_get_or_create(key: str) -> Session:
loaded.append(key)
return get_or_create(key)
return await get_or_create_async(key)
monkeypatch.setattr(restarted, "get_or_create", tracked_get_or_create)
monkeypatch.setattr(restarted, "get_or_create_async", tracked_get_or_create)
monkeypatch.setattr(
"nanobot.webui.transcript.has_unfinished_transcript_tail",
lambda _key: False,
+38
View File
@@ -1,4 +1,6 @@
import asyncio
import gc
import threading
import weakref
from nanobot.session.manager import SESSION_CACHE_MAX_SIZE, SessionManager
@@ -87,3 +89,39 @@ def test_transient_session_never_reaches_storage(tmp_path) -> None:
assert list(manager.sessions_dir.glob("*.jsonl")) == []
manager.invalidate(session.key)
assert manager.get_cached(session.key) is None
async def test_concurrent_first_async_gets_share_cached_identity(
tmp_path,
monkeypatch,
) -> None:
manager = SessionManager(tmp_path)
key = "test:concurrent-first-load"
load_started = threading.Event()
release_load = threading.Event()
load_calls = 0
def delayed_load(loaded_key: str):
nonlocal load_calls
assert loaded_key == key
load_calls += 1
load_started.set()
assert release_load.wait(timeout=1)
return None
monkeypatch.setattr(manager, "_load", delayed_load)
first_task = asyncio.create_task(manager.get_or_create_async(key))
try:
assert await asyncio.to_thread(load_started.wait, 0.5)
second_task = asyncio.create_task(manager.get_or_create_async(key))
await asyncio.sleep(0)
assert not second_task.done()
finally:
release_load.set()
first, second = await asyncio.gather(first_task, second_task)
assert load_calls == 1
assert first is second
assert manager.get_cached(key) is first
assert await manager.get_or_create_async(key) is first
+162
View File
@@ -0,0 +1,162 @@
"""Cancellation settlement guarantees for native SessionManager mutations."""
from __future__ import annotations
import asyncio
import threading
from pathlib import Path
from typing import Any
import pytest
from nanobot.session.manager import Session, SessionManager
async def _cancel_blocked_mutation(
task: asyncio.Task[Any],
*,
started: threading.Event,
release: threading.Event,
) -> None:
assert await asyncio.to_thread(started.wait, 1)
try:
task.cancel()
await asyncio.sleep(0)
assert not task.done(), "cancellation escaped before the mutation worker settled"
finally:
release.set()
with pytest.raises(asyncio.CancelledError):
await asyncio.wait_for(task, timeout=1)
async def test_save_async_cancellation_waits_for_durable_write_and_cache(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
manager = SessionManager(tmp_path)
session = Session(key="test:cancel-save")
session.add_message("user", "persist exactly once")
started = threading.Event()
release = threading.Event()
finished = threading.Event()
save_calls = 0
original_save = manager._store.save
def blocked_save(target: Session, *, fsync: bool = False) -> None:
nonlocal save_calls
save_calls += 1
started.set()
assert release.wait(timeout=1)
original_save(target, fsync=fsync)
finished.set()
monkeypatch.setattr(manager._store, "save", blocked_save)
task = asyncio.create_task(manager.save_async(session))
await _cancel_blocked_mutation(task, started=started, release=release)
assert finished.is_set()
assert save_calls == 1
assert manager.get_cached(session.key) is session
durable = manager.read_session_file(session.key)
assert durable is not None
assert durable["messages"][0]["content"] == "persist exactly once"
await asyncio.sleep(0.05)
assert save_calls == 1
assert manager.read_session_file(session.key) == durable
async def test_update_metadata_async_cancellation_waits_for_file_and_cache_refresh(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
manager = SessionManager(tmp_path)
session = manager.get_or_create("test:cancel-metadata")
session.metadata["title"] = "before"
manager.save(session)
started = threading.Event()
release = threading.Event()
finished = threading.Event()
update_calls = 0
original_update = manager._store.update_metadata
def blocked_update(
key: str,
updates: dict[str, Any],
*,
fsync: bool = False,
) -> bool:
nonlocal update_calls
update_calls += 1
started.set()
assert release.wait(timeout=1)
updated = original_update(key, updates, fsync=fsync)
finished.set()
return updated
monkeypatch.setattr(manager._store, "update_metadata", blocked_update)
task = asyncio.create_task(
manager.update_session_metadata_async(
session.key,
{"title": "after", "settled": True},
)
)
await _cancel_blocked_mutation(task, started=started, release=release)
assert finished.is_set()
assert update_calls == 1
assert session.metadata["title"] == "after"
assert session.metadata["settled"] is True
durable = manager.read_session_file(session.key)
assert durable is not None
assert durable["metadata"]["title"] == "after"
assert durable["metadata"]["settled"] is True
await asyncio.sleep(0.05)
assert update_calls == 1
assert manager.read_session_file(session.key) == durable
async def test_delete_async_cancellation_waits_for_file_cache_and_observer(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
manager = SessionManager(tmp_path)
session = manager.get_or_create("test:cancel-delete")
session.add_message("user", "delete me")
manager.save(session)
observed: list[str] = []
manager.set_delete_observer(observed.append)
started = threading.Event()
release = threading.Event()
finished = threading.Event()
delete_calls = 0
original_delete = manager._store.delete
def blocked_delete(key: str) -> bool:
nonlocal delete_calls
delete_calls += 1
started.set()
assert release.wait(timeout=1)
deleted = original_delete(key)
finished.set()
return deleted
monkeypatch.setattr(manager._store, "delete", blocked_delete)
task = asyncio.create_task(manager.delete_session_async(session.key))
await _cancel_blocked_mutation(task, started=started, release=release)
assert finished.is_set()
assert delete_calls == 1
assert manager.read_session_file(session.key) is None
assert manager.get_cached(session.key) is None
assert observed == [session.key]
await asyncio.sleep(0.05)
assert delete_calls == 1
assert observed == [session.key]
assert manager.read_session_file(session.key) is None
+127
View File
@@ -0,0 +1,127 @@
"""Cross-cutting heartbeat regressions for synchronous persistence locks."""
from __future__ import annotations
import asyncio
from contextlib import suppress
from pathlib import Path
from filelock import FileLock
from nanobot.agent.tools.context import RequestContext, request_context
from nanobot.agent.tools.cron import CronTool
from nanobot.config.loader import load_config, save_config
from nanobot.config.schema import Config
from nanobot.cron.service import CronService
from nanobot.session.manager import SessionManager
from nanobot.triggers.local_runner import run_local_trigger_queue
from nanobot.triggers.local_store import LocalTriggerStore
from nanobot.webui.settings_api import update_api_settings
from nanobot.webui.settings_services import WebUISettingsServices
async def _assert_10ms_heartbeat_runs_while_pending(task: asyncio.Task[object]) -> None:
for _ in range(3):
await asyncio.sleep(0.01)
assert not task.done()
async def test_session_lock_contention_does_not_block_event_loop(
tmp_path: Path,
monkeypatch,
) -> None:
workspace = tmp_path / "workspace"
sessions_root = tmp_path / "session-data"
legacy_root = tmp_path / "legacy-sessions"
monkeypatch.setattr(
"nanobot.session.manager.get_legacy_sessions_dir",
lambda: legacy_root,
)
manager = SessionManager(workspace, sessions_root=sessions_root)
session = manager.get_or_create("websocket:lock-test")
session.add_message("user", "hello")
lock = manager._jsonl_store._session_files_lock
assert lock.timeout == 5
blocker = FileLock(lock.lock_file, timeout=1)
blocker.acquire()
task = asyncio.create_task(manager.save_async(session))
try:
await _assert_10ms_heartbeat_runs_while_pending(task)
finally:
blocker.release()
await asyncio.wait_for(task, timeout=1)
assert manager.read_session_file(session.key) is not None
async def test_cron_lock_contention_does_not_block_event_loop(tmp_path: Path) -> None:
service = CronService(tmp_path / "cron" / "jobs.json")
tool = CronTool(service)
assert service._lock.timeout == 5
blocker = FileLock(service._lock.lock_file, timeout=1)
blocker.acquire()
with request_context(
RequestContext(
channel="websocket",
chat_id="lock-test",
session_key="websocket:lock-test",
)
):
task = asyncio.create_task(
tool.execute(action="add", message="wake up", every_seconds=60)
)
try:
await _assert_10ms_heartbeat_runs_while_pending(task)
finally:
blocker.release()
result = await asyncio.wait_for(task, timeout=1)
assert result.startswith("Created job")
async def test_local_trigger_lock_contention_does_not_block_event_loop(
tmp_path: Path,
) -> None:
store = LocalTriggerStore(tmp_path)
assert store._lock.timeout == 5
blocker = FileLock(store._lock.lock_file, timeout=1)
blocker.acquire()
async def submit_turn(_message):
return None
task = asyncio.create_task(
run_local_trigger_queue(
store=store,
submit_turn=submit_turn,
is_channel_enabled=lambda _name: True,
poll_interval_s=0.01,
)
)
try:
await _assert_10ms_heartbeat_runs_while_pending(task)
finally:
blocker.release()
await asyncio.sleep(0.05)
task.cancel()
with suppress(asyncio.CancelledError):
await asyncio.wait_for(task, timeout=1)
async def test_settings_lock_contention_does_not_block_event_loop(tmp_path: Path) -> None:
config_path = tmp_path / "config.json"
save_config(Config(), config_path)
services = WebUISettingsServices.create(config_path)
assert services.config._file_lock.timeout == 5
blocker = FileLock(services.config._file_lock.lock_file, timeout=1)
blocker.acquire()
task = asyncio.create_task(
services.mutate_async(update_api_settings, {"port": ["19001"]})
)
try:
await _assert_10ms_heartbeat_runs_while_pending(task)
finally:
blocker.release()
await asyncio.wait_for(task, timeout=1)
assert load_config(config_path).api.port == 19001
+33
View File
@@ -2,8 +2,10 @@
from __future__ import annotations
import asyncio
import socket
import sys
import threading
from unittest.mock import patch
import pytest
@@ -173,6 +175,37 @@ async def test_exec_blocks_chained_internal_url():
assert "Error" in result
@pytest.mark.asyncio
async def test_restricted_exec_url_guard_keeps_event_loop_responsive(
tmp_path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
resolver_started = threading.Event()
release_resolver = threading.Event()
def slow_private_resolver(hostname, port, family=0, type=0, proto=0, flags=0): # noqa: A002
resolver_started.set()
assert release_resolver.wait(timeout=1)
return _fake_resolve_private(hostname, port, family, type)
monkeypatch.setattr(socket, "getaddrinfo", slow_private_resolver)
tool = ExecTool(working_dir=str(tmp_path), restrict_to_workspace=True)
task = asyncio.create_task(tool.execute(command="curl https://example.com/latest"))
safety_release = threading.Timer(0.2, release_resolver.set)
safety_release.start()
try:
assert await asyncio.to_thread(resolver_started.wait, 0.5)
for _ in range(3):
await asyncio.sleep(0)
assert not task.done()
finally:
release_resolver.set()
safety_release.cancel()
result = await asyncio.wait_for(task, timeout=0.5)
assert "internal/private URL detected" in result
# --- #2989: block writes to nanobot internal state files -----------------
+190
View File
@@ -1,7 +1,12 @@
"""Tests for enhanced filesystem tools: ReadFileTool, EditFileTool, ListDirTool."""
import asyncio
import threading
import pytest
from nanobot.agent.tools import file_state
from nanobot.agent.tools import filesystem as filesystem_tools
from nanobot.agent.tools.filesystem import (
EditFileTool,
ListDirTool,
@@ -478,3 +483,188 @@ class TestWorkspaceRestriction:
)
assert "Successfully edited" in result
assert target.read_text(encoding="utf-8") == "after\n"
@pytest.mark.asyncio
async def test_edit_matching_keeps_event_loop_responsive(tmp_path, monkeypatch):
target = tmp_path / "slow.py"
target.write_text("before\n", encoding="utf-8")
tool = EditFileTool(workspace=tmp_path)
original_find_matches = filesystem_tools._find_matches
started = threading.Event()
release = threading.Event()
def blocking_find_matches(*args, **kwargs):
started.set()
if not release.wait(timeout=1):
raise TimeoutError("test did not release edit matching")
return original_find_matches(*args, **kwargs)
monkeypatch.setattr(filesystem_tools, "_find_matches", blocking_find_matches)
task = asyncio.create_task(
tool.execute(path=str(target), old_text="before", new_text="after")
)
try:
assert await asyncio.to_thread(started.wait, 0.5)
for _ in range(3):
await asyncio.sleep(0.01)
assert not task.done()
finally:
release.set()
assert "Successfully edited" in await asyncio.wait_for(task, timeout=0.5)
assert target.read_text(encoding="utf-8") == "after\n"
@pytest.mark.asyncio
async def test_edit_cancellation_before_commit_prevents_delayed_write(tmp_path, monkeypatch):
target = tmp_path / "cancel.py"
target.write_text("before\n", encoding="utf-8")
tool = EditFileTool(workspace=tmp_path)
original_find_matches = filesystem_tools._find_matches
started = threading.Event()
release = threading.Event()
def blocking_find_matches(*args, **kwargs):
started.set()
if not release.wait(timeout=1):
raise TimeoutError("test did not release edit matching")
return original_find_matches(*args, **kwargs)
monkeypatch.setattr(filesystem_tools, "_find_matches", blocking_find_matches)
task = asyncio.create_task(
tool.execute(path=str(target), old_text="before", new_text="after")
)
assert await asyncio.to_thread(started.wait, 0.5)
task.cancel()
with pytest.raises(asyncio.CancelledError):
await asyncio.wait_for(task, timeout=0.2)
release.set()
await asyncio.sleep(0.05)
assert target.read_text(encoding="utf-8") == "before\n"
@pytest.mark.asyncio
async def test_edit_repeated_cancellation_during_commit_waits_for_settlement(
tmp_path,
monkeypatch,
):
target = tmp_path / "commit.py"
target.write_text("before\n", encoding="utf-8")
tool = EditFileTool(workspace=tmp_path)
path_type = type(target)
original_write_bytes = path_type.write_bytes
original_wait_for_commit = EditFileTool._wait_for_commit
commit_started = threading.Event()
drain_started = threading.Event()
release = threading.Event()
writes: list[bytes] = []
def blocking_write_bytes(path, data):
if path == target:
commit_started.set()
if not release.wait(timeout=1):
raise TimeoutError("test did not release edit commit")
writes.append(data)
return original_write_bytes(path, data)
def observed_wait_for_commit(commit_lock):
drain_started.set()
original_wait_for_commit(commit_lock)
monkeypatch.setattr(path_type, "write_bytes", blocking_write_bytes)
monkeypatch.setattr(
EditFileTool,
"_wait_for_commit",
staticmethod(observed_wait_for_commit),
)
task = asyncio.create_task(
tool.execute(path=str(target), old_text="before", new_text="after")
)
try:
assert await asyncio.to_thread(commit_started.wait, 0.5)
assert task.cancel()
assert await asyncio.to_thread(drain_started.wait, 0.5)
# Later cancellation requests must not interrupt the in-flight commit drain.
assert task.cancel()
await asyncio.sleep(0)
assert not task.done()
assert task.cancel()
await asyncio.sleep(0)
assert not task.done()
assert target.read_text(encoding="utf-8") == "before\n"
finally:
release.set()
with pytest.raises(asyncio.CancelledError):
await asyncio.wait_for(task, timeout=0.5)
final_content = target.read_bytes()
final_state = tool._file_states.get(target)
assert target.read_text(encoding="utf-8") == "after\n"
assert writes == [final_content]
assert final_state is not None
assert tool._file_states.check_read(target) is None
await asyncio.sleep(0.05)
assert target.read_bytes() == final_content
assert tool._file_states.get(target) is final_state
assert writes == [final_content]
def test_exact_match_line_accounting_is_linear_and_occurrence_can_stop_early():
class CountingText(str):
count_span = 0
find_calls = 0
def count(self, sub, start=None, end=None):
actual_start = 0 if start is None else start
actual_end = len(self) if end is None else end
self.count_span += max(0, actual_end - actual_start)
return super().count(sub, actual_start, actual_end)
def find(self, sub, start=None, end=None):
self.find_calls += 1
actual_start = 0 if start is None else start
actual_end = len(self) if end is None else end
return super().find(sub, actual_start, actual_end)
content = CountingText("match\n" * 10_000)
matches = filesystem_tools._find_exact_matches(content, "match")
assert len(matches) == 10_000
assert matches[-1].line == 10_000
assert content.count_span <= len(content)
first_only = CountingText(content)
matches = filesystem_tools._find_matches(
first_only,
"match",
max_exact_matches=1,
)
assert len(matches) == 1
assert first_only.find_calls == 1
@pytest.mark.asyncio
async def test_edit_worker_preserves_bound_file_state_context(tmp_path):
target = tmp_path / "context.py"
target.write_text("before\n", encoding="utf-8")
states = file_state.FileStates()
states.record_read(target)
tool = EditFileTool(workspace=tmp_path)
token = file_state.bind_file_states(states)
try:
result = await tool.execute(
path=str(target),
old_text="before",
new_text="after",
)
finally:
file_state.reset_file_states(token)
assert result == f"Successfully edited {target}"
assert states.get(target) is not None
assert target.read_text(encoding="utf-8") == "after\n"
+37 -23
View File
@@ -2,7 +2,6 @@
from __future__ import annotations
import asyncio
import socket
from unittest.mock import MagicMock, patch
import httpx
@@ -56,10 +55,12 @@ async def test_probe_uses_default_port_for_http(monkeypatch: pytest.MonkeyPatch)
"""When no port is present, probe the validated address on port 80."""
attempts: list[tuple[str, int]] = []
monkeypatch.setattr(
"nanobot.agent.tools.mcp.resolve_url_target",
lambda _url: (True, "", ("93.184.216.34",)),
)
async def _resolve_url_target(
_url: str,
) -> tuple[bool, str, tuple[str, ...]]:
return True, "", ("93.184.216.34",)
monkeypatch.setattr(mcp_mod, "async_resolve_url_target", _resolve_url_target)
async def _open_connection(host: str, port: int):
attempts.append((host, port))
@@ -72,32 +73,43 @@ async def test_probe_uses_default_port_for_http(monkeypatch: pytest.MonkeyPatch)
@pytest.mark.asyncio
async def test_probe_rejects_public_name_resolving_to_loopback():
def _resolver(hostname, port, family=0, type_=0):
return [(socket.AF_INET, socket.SOCK_STREAM, 0, "", ("127.0.0.1", 0))]
async def test_probe_rejects_public_name_resolving_to_loopback(
monkeypatch: pytest.MonkeyPatch,
) -> None:
async def _resolve_url_target(
_url: str,
) -> tuple[bool, str, tuple[str, ...]]:
return False, "Blocked: example.com resolves to private/internal address 127.0.0.1", ()
with patch("nanobot.security.network.socket.getaddrinfo", _resolver):
assert await _probe_http_url("http://example.com:8765/mcp") is False
monkeypatch.setattr(mcp_mod, "async_resolve_url_target", _resolve_url_target)
assert await _probe_http_url("http://example.com:8765/mcp") is False
@pytest.mark.asyncio
async def test_probe_skips_direct_tcp_when_global_proxy_env_is_set(monkeypatch):
def _resolver(hostname, port, family=0, type_=0):
return [(socket.AF_INET, socket.SOCK_STREAM, 0, "", ("93.184.216.34", 0))]
async def test_probe_skips_direct_tcp_when_global_proxy_env_is_set(
monkeypatch: pytest.MonkeyPatch,
) -> None:
async def _resolve_url_target(
_url: str,
) -> tuple[bool, str, tuple[str, ...]]:
return True, "", ("93.184.216.34",)
async def _open_connection(*args, **kwargs):
raise AssertionError("global proxy env should skip direct TCP probe")
monkeypatch.setenv("HTTPS_PROXY", "http://proxy.example:8080")
monkeypatch.setenv("NO_PROXY", "localhost,127.0.0.1,::1")
monkeypatch.setattr(mcp_mod, "async_resolve_url_target", _resolve_url_target)
monkeypatch.setattr("nanobot.agent.tools.mcp.asyncio.open_connection", _open_connection)
with patch("nanobot.security.network.socket.getaddrinfo", _resolver):
assert await _probe_http_url("https://mcp.example.com/mcp") is True
assert await _probe_http_url("https://mcp.example.com/mcp") is True
@pytest.mark.asyncio
async def test_probe_tries_next_validated_ip_when_first_is_unreachable(monkeypatch):
async def test_probe_tries_next_validated_ip_when_first_is_unreachable(
monkeypatch: pytest.MonkeyPatch,
) -> None:
attempts: list[tuple[str, int]] = []
class FakeWriter:
@@ -107,11 +119,10 @@ async def test_probe_tries_next_validated_ip_when_first_is_unreachable(monkeypat
async def wait_closed(self):
return None
def _resolver(hostname, port, family=0, type_=0):
return [
(socket.AF_INET, socket.SOCK_STREAM, 0, "", ("93.184.216.34", 0)),
(socket.AF_INET, socket.SOCK_STREAM, 0, "", ("93.184.216.35", 0)),
]
async def _resolve_url_target(
_url: str,
) -> tuple[bool, str, tuple[str, ...]]:
return True, "", ("93.184.216.34", "93.184.216.35")
async def _open_connection(host: str, port: int):
attempts.append((host, port))
@@ -119,7 +130,7 @@ async def test_probe_tries_next_validated_ip_when_first_is_unreachable(monkeypat
raise OSError("first address unreachable")
return object(), FakeWriter()
monkeypatch.setattr("nanobot.security.network.socket.getaddrinfo", _resolver)
monkeypatch.setattr(mcp_mod, "async_resolve_url_target", _resolve_url_target)
monkeypatch.setattr("nanobot.agent.tools.mcp.asyncio.open_connection", _open_connection)
assert await _probe_http_url("http://mcp.example:8765/mcp") is True
@@ -182,10 +193,13 @@ async def test_connect_isolates_streamable_http_status_failure(
async def _reachable(_url: str) -> bool:
return True
async def _validate_url_target(_url: str) -> tuple[bool, str]:
return True, ""
def _return_http_530(request: httpx.Request) -> httpx.Response:
return httpx.Response(530, text="cloudflare error 1033", request=request)
monkeypatch.setattr(mcp_mod, "validate_url_target", lambda _url: (True, ""))
monkeypatch.setattr(mcp_mod, "async_validate_url_target", _validate_url_target)
monkeypatch.setattr(mcp_mod, "_probe_http_url", _reachable)
monkeypatch.setattr(
mcp_mod,
+19 -10
View File
@@ -213,7 +213,7 @@ async def test_saved_oauth_http_403_projects_failed_runtime_without_details(
oauth_mod.create_mcp_oauth_auth = create_auth # type: ignore[attr-defined]
oauth_mod.mcp_oauth_has_credentials = lambda _name, _url: True # type: ignore[attr-defined]
monkeypatch.setitem(sys.modules, "nanobot.agent.tools.mcp_oauth", oauth_mod)
monkeypatch.setattr(mcp_mod, "validate_url_target", lambda _url: (True, ""))
monkeypatch.setattr(mcp_mod, "async_validate_url_target", _allow_url)
monkeypatch.setattr(mcp_mod, "_probe_http_url", reachable)
monkeypatch.setattr(
sys.modules["mcp.client.streamable_http"],
@@ -1006,7 +1006,7 @@ async def test_connect_mcp_servers_env_proxy_adds_proxy_mounts_and_keeps_pinned_
async def _reachable(_url: str) -> bool:
return True
def _validate(_url: str) -> tuple[bool, str]:
async def _validate(_url: str) -> tuple[bool, str]:
return True, ""
class FakeAsyncClient:
@@ -1033,7 +1033,7 @@ async def test_connect_mcp_servers_env_proxy_adds_proxy_mounts_and_keeps_pinned_
monkeypatch.setenv("HTTPS_PROXY", "http://proxy.example:8080")
monkeypatch.setenv("NO_PROXY", "localhost,127.0.0.1,::1")
monkeypatch.setattr(mcp_mod, "validate_url_target", _validate)
monkeypatch.setattr(mcp_mod, "async_validate_url_target", _validate)
monkeypatch.setattr(mcp_mod, "_probe_http_url", _reachable)
monkeypatch.setattr(
mcp_mod,
@@ -1105,7 +1105,7 @@ async def test_connect_mcp_servers_http_clients_reject_unsafe_redirect_targets(
sent_urls: list[str] = []
used_transports: list[str] = []
def _validate(url: str, **_kwargs: object) -> tuple[bool, str]:
async def _validate(url: str, **_kwargs: object) -> tuple[bool, str]:
checked_urls.append(url)
if url == "http://127.0.0.1/private":
return False, "loopback blocked"
@@ -1145,7 +1145,7 @@ async def test_connect_mcp_servers_http_clients_reject_unsafe_redirect_targets(
await http_client.get("https://example.com/start")
yield object(), object(), object()
monkeypatch.setattr(mcp_mod, "validate_url_target", _validate)
monkeypatch.setattr(mcp_mod, "async_validate_url_target", _validate)
monkeypatch.setattr(mcp_mod, "_probe_http_url", _reachable)
monkeypatch.setattr(
mcp_mod,
@@ -1320,7 +1320,7 @@ async def test_connect_mcp_servers_streamable_http_uses_finite_timeout(
async def _reachable(_url: str) -> bool:
return True
def _validate(_url: str) -> tuple[bool, str]:
async def _validate(_url: str) -> tuple[bool, str]:
return True, ""
@asynccontextmanager
@@ -1328,7 +1328,7 @@ async def test_connect_mcp_servers_streamable_http_uses_finite_timeout(
captured["timeout"] = http_client.timeout
yield object(), object(), object()
monkeypatch.setattr(mcp_mod, "validate_url_target", _validate)
monkeypatch.setattr(mcp_mod, "async_validate_url_target", _validate)
monkeypatch.setattr(mcp_mod, "_probe_http_url", _reachable)
monkeypatch.setattr(
mcp_mod,
@@ -1371,7 +1371,7 @@ async def test_connect_mcp_servers_attaches_oauth_to_remote_http_client(
async def _reachable(_url: str) -> bool:
return True
def _validate(_url: str) -> tuple[bool, str]:
async def _validate(_url: str) -> tuple[bool, str]:
return True, ""
async def _create_auth(name: str, url: str, handlers: object) -> object:
@@ -1407,7 +1407,7 @@ async def test_connect_mcp_servers_attaches_oauth_to_remote_http_client(
assert http_client is not None
yield object(), object(), object()
monkeypatch.setattr(mcp_mod, "validate_url_target", _validate)
monkeypatch.setattr(mcp_mod, "async_validate_url_target", _validate)
monkeypatch.setattr(mcp_mod, "_probe_http_url", _reachable)
monkeypatch.setattr(mcp_mod.httpx, "AsyncClient", FakeAsyncClient)
monkeypatch.setattr(sys.modules["mcp.client.sse"], "sse_client", _capturing_sse_client)
@@ -1461,7 +1461,7 @@ async def test_connect_mcp_servers_skips_background_oauth_without_credentials(
oauth_mod.MCPAuthorizationRequiredError = AuthorizationRequiredError # type: ignore[attr-defined]
oauth_mod.create_mcp_oauth_auth = _create_auth # type: ignore[attr-defined]
monkeypatch.setitem(sys.modules, "nanobot.agent.tools.mcp_oauth", oauth_mod)
monkeypatch.setattr(mcp_mod, "validate_url_target", lambda _url: (True, ""))
monkeypatch.setattr(mcp_mod, "async_validate_url_target", _allow_url)
monkeypatch.setattr(mcp_mod, "_probe_http_url", _probe)
stacks = await connect_mcp_servers(
@@ -1959,3 +1959,12 @@ def test_long_server_name_tools_are_matched_by_server_name() -> None:
assert removed == 1
assert wrapper.name not in registry.tool_names
assert other_wrapper.name in registry.tool_names
async def _allow_url(
_url: str,
*,
allow_loopback: bool = False,
) -> tuple[bool, str]:
del allow_loopback
return True, ""
+78
View File
@@ -2,7 +2,9 @@
from __future__ import annotations
import asyncio
import os
import threading
import time
from pathlib import Path
from types import SimpleNamespace
@@ -97,6 +99,82 @@ async def test_find_files_rejects_paths_outside_workspace(tmp_path: Path) -> Non
assert result.startswith("Error:")
@pytest.mark.asyncio
async def test_find_files_scan_keeps_event_loop_responsive(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
target = tmp_path / "match.txt"
target.write_text("ok\n", encoding="utf-8")
tool = FindFilesTool(workspace=tmp_path, allowed_dir=tmp_path)
original_iter_paths = tool._iter_paths
started = threading.Event()
release = threading.Event()
def blocking_iter_paths(root: Path, *, include_dirs: bool):
started.set()
if not release.wait(timeout=1):
raise TimeoutError("test did not release find_files traversal")
yield from original_iter_paths(root, include_dirs=include_dirs)
monkeypatch.setattr(tool, "_iter_paths", blocking_iter_paths)
task = asyncio.create_task(tool.execute(path="."))
try:
assert await asyncio.to_thread(started.wait, 0.5)
for _ in range(3):
await asyncio.sleep(0.01)
assert not task.done()
finally:
release.set()
assert await asyncio.wait_for(task, timeout=0.5) == "match.txt"
@pytest.mark.asyncio
async def test_find_files_cancellation_is_prompt_and_stops_scan_before_next_path(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
(tmp_path / "match.txt").write_text("ok\n", encoding="utf-8")
tool = FindFilesTool(workspace=tmp_path, allowed_dir=tmp_path)
original_iter_paths = tool._iter_paths
started = threading.Event()
release = threading.Event()
def blocking_iter_paths(root: Path, *, include_dirs: bool):
started.set()
if not release.wait(timeout=1):
raise TimeoutError("test did not release find_files traversal")
yield from original_iter_paths(root, include_dirs=include_dirs)
monkeypatch.setattr(tool, "_iter_paths", blocking_iter_paths)
task = asyncio.create_task(tool.execute(path="."))
assert await asyncio.to_thread(started.wait, 0.5)
task.cancel()
try:
with pytest.raises(asyncio.CancelledError):
await asyncio.wait_for(task, timeout=0.2)
await asyncio.sleep(0.01)
finally:
release.set()
await asyncio.sleep(0.05)
@pytest.mark.asyncio
async def test_find_files_enforces_path_budget(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
(tmp_path / "a.txt").write_text("a\n", encoding="utf-8")
(tmp_path / "b.txt").write_text("b\n", encoding="utf-8")
monkeypatch.setattr(FindFilesTool, "_MAX_SCAN_PATHS", 1)
tool = FindFilesTool(workspace=tmp_path, allowed_dir=tmp_path)
result = await tool.execute(path=".")
assert result.startswith("Error: find_files scan exceeded 1 paths")
@pytest.mark.asyncio
async def test_grep_respects_glob_filter_and_context(tmp_path: Path) -> None:
(tmp_path / "src").mkdir()
+1 -1
View File
@@ -16,7 +16,7 @@ from nanobot.agent.tools.web import (
)
def _fake_resolve_public(hostname, port, family=0, type_=0):
def _fake_resolve_public(hostname, port, family=0, type_=0, proto=0, flags=0):
return [(socket.AF_INET, socket.SOCK_STREAM, 0, "", ("93.184.216.34", 0))]
+7 -7
View File
@@ -30,11 +30,11 @@ def _clear_proxy_env(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.delenv(name, raising=False)
def _fake_resolve_private(hostname, port, family=0, type_=0):
def _fake_resolve_private(hostname, port, family=0, type_=0, proto=0, flags=0):
return [(socket.AF_INET, socket.SOCK_STREAM, 0, "", ("169.254.169.254", 0))]
def _fake_resolve_public(hostname, port, family=0, type_=0):
def _fake_resolve_public(hostname, port, family=0, type_=0, proto=0, flags=0):
return [(socket.AF_INET, socket.SOCK_STREAM, 0, "", ("93.184.216.34", 0))]
@@ -99,7 +99,7 @@ async def test_web_fetch_blocks_private_ip():
@pytest.mark.asyncio
async def test_web_fetch_blocks_localhost():
tool = WebFetchTool()
def _resolve_localhost(hostname, port, family=0, type_=0):
def _resolve_localhost(hostname, port, family=0, type_=0, proto=0, flags=0):
return [(socket.AF_INET, socket.SOCK_STREAM, 0, "", ("127.0.0.1", 0))]
with patch("nanobot.security.network.socket.getaddrinfo", _resolve_localhost):
result = await tool.execute(url="http://localhost/admin")
@@ -112,7 +112,7 @@ async def test_web_fetch_blocks_localhost_even_in_full_workspace_scope(tmp_path)
tool = WebFetchTool()
scope = build_workspace_scope(tmp_path, "full")
def _resolve_localhost(hostname, port, family=0, type_=0):
def _resolve_localhost(hostname, port, family=0, type_=0, proto=0, flags=0):
return [(socket.AF_INET, socket.SOCK_STREAM, 0, "", ("127.0.0.1", 0))]
token = bind_workspace_scope(scope)
@@ -439,7 +439,7 @@ async def test_web_fetch_blocks_private_redirect_before_readability_request(monk
monkeypatch.setattr(web_module.httpx, "AsyncClient", FakeClient)
monkeypatch.setattr(web_module, "_pinned_dns_transport", lambda: object())
def resolve_public_start_only(hostname, port, family=0, type_=0):
def resolve_public_start_only(hostname, port, family=0, type_=0, proto=0, flags=0):
if hostname == "attacker.example":
return _fake_resolve_public(hostname, port, family, type_)
return _REAL_GETADDRINFO(hostname, port, family, type_)
@@ -485,7 +485,7 @@ async def test_web_fetch_blocks_private_redirect_before_returning_image(monkeypa
monkeypatch.setattr("nanobot.agent.tools.web.httpx.AsyncClient", TransportAsyncClient)
monkeypatch.setattr(web_module, "_pinned_dns_transport", lambda: object())
def resolve_public_start_only(hostname, port, family=0, type_=0):
def resolve_public_start_only(hostname, port, family=0, type_=0, proto=0, flags=0):
if hostname == "example.com":
return _fake_resolve_public(hostname, port, family, type_)
return _REAL_GETADDRINFO(hostname, port, family, type_)
@@ -526,7 +526,7 @@ async def test_web_fetch_does_not_request_private_redirect_target(monkeypatch):
monkeypatch.setattr(web_module.httpx, "AsyncClient", TransportAsyncClient)
monkeypatch.setattr(web_module, "_pinned_dns_transport", lambda: object())
def resolve_public_start_only(hostname, port, family=0, type_=0):
def resolve_public_start_only(hostname, port, family=0, type_=0, proto=0, flags=0):
if hostname == "attacker.example":
return _fake_resolve_public(hostname, port, family, type_)
return _REAL_GETADDRINFO(hostname, port, family, type_)
@@ -11,7 +11,7 @@ import pytest
from nanobot.agent.tools.web import WebFetchTool, _validate_url
def _fake_resolve_public(hostname, port, family=0, type_=0):
def _fake_resolve_public(hostname, port, family=0, type_=0, proto=0, flags=0):
import socket
return [(socket.AF_INET, socket.SOCK_STREAM, 0, "", ("93.184.216.34", 0))]
+304
View File
@@ -4,6 +4,7 @@ import asyncio
import errno
import json
import os
import threading
from contextlib import suppress
from pathlib import Path
@@ -13,6 +14,7 @@ from nanobot.agent.automation_turns import AutomationTurnError
from nanobot.bus.events import InboundMessage, OutboundMessage
from nanobot.triggers.local_runner import run_local_trigger_queue
from nanobot.triggers.local_store import LocalTriggerStore, TriggerDisabledError
from nanobot.triggers.local_turns import LocalTriggerTurnCoordinator
from nanobot.triggers.local_types import LocalTrigger, TriggerDelivery
from nanobot.webui.metadata import WEBUI_MESSAGE_SOURCE_METADATA_KEY, WEBUI_TURN_METADATA_KEY
@@ -276,6 +278,163 @@ def test_recover_processing_deliveries_requeues_claimed_delivery(tmp_path: Path)
assert reclaimed[0].last_error == "delivery was recovered from interrupted processing"
@pytest.mark.asyncio
async def test_local_trigger_queue_cancellation_waits_for_claim_mutation(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
store = LocalTriggerStore(tmp_path)
trigger = store.create(
name="PR review",
channel="websocket",
chat_id="chat-1",
session_key="websocket:chat-1",
)
store.enqueue(trigger.id, "Review PR #4591")
claim_started = threading.Event()
allow_claim = threading.Event()
claim_finished = threading.Event()
claim_deliveries = store.claim_deliveries
submitted: list[InboundMessage] = []
def blocked_claim_deliveries(*, limit: int = 20) -> list[TriggerDelivery]:
claim_started.set()
if not allow_claim.wait(timeout=2):
raise TimeoutError("test did not release delivery claim")
try:
return claim_deliveries(limit=limit)
finally:
claim_finished.set()
async def submit_turn(msg: InboundMessage) -> None:
submitted.append(msg)
monkeypatch.setattr(store, "claim_deliveries", blocked_claim_deliveries)
task = asyncio.create_task(
run_local_trigger_queue(
store=store,
submit_turn=submit_turn,
is_channel_enabled=_channel_is_enabled,
poll_interval_s=0.01,
)
)
try:
assert await asyncio.to_thread(claim_started.wait, 1)
task.cancel()
await asyncio.sleep(0)
assert not task.done()
assert len(list(store.inbox_dir.glob("*.json"))) == 1
assert list(store.processing_dir.glob("*.json")) == []
allow_claim.set()
with pytest.raises(asyncio.CancelledError):
await asyncio.wait_for(task, timeout=1)
assert claim_finished.is_set()
state_at_cancellation = (
sorted(path.name for path in store.inbox_dir.glob("*.json")),
sorted(path.name for path in store.processing_dir.glob("*.json")),
)
assert state_at_cancellation[0] == []
assert len(state_at_cancellation[1]) == 1
assert submitted == []
# The worker has returned before cancellation is visible, so no detached
# claim can move another file after the cancellation report.
await asyncio.sleep(0)
assert (
sorted(path.name for path in store.inbox_dir.glob("*.json")),
sorted(path.name for path in store.processing_dir.glob("*.json")),
) == state_at_cancellation
finally:
allow_claim.set()
task.cancel()
with suppress(asyncio.CancelledError):
await task
@pytest.mark.asyncio
async def test_local_trigger_queue_cancellation_waits_for_startup_recovery(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
store = LocalTriggerStore(tmp_path)
trigger = store.create(
name="PR review",
channel="websocket",
chat_id="chat-1",
session_key="websocket:chat-1",
)
store.enqueue(trigger.id, "Review PR #4591")
assert len(store.claim_deliveries()) == 1
restarted = LocalTriggerStore(tmp_path)
recovery_started = threading.Event()
allow_recovery = threading.Event()
recovery_finished = threading.Event()
recover_processing_deliveries = restarted.recover_processing_deliveries
submitted: list[InboundMessage] = []
def blocked_recovery() -> int:
recovery_started.set()
if not allow_recovery.wait(timeout=2):
raise TimeoutError("test did not release delivery recovery")
try:
return recover_processing_deliveries()
finally:
recovery_finished.set()
async def submit_turn(msg: InboundMessage) -> None:
submitted.append(msg)
monkeypatch.setattr(restarted, "recover_processing_deliveries", blocked_recovery)
task = asyncio.create_task(
run_local_trigger_queue(
store=restarted,
submit_turn=submit_turn,
is_channel_enabled=_channel_is_enabled,
poll_interval_s=0.01,
)
)
try:
assert await asyncio.to_thread(recovery_started.wait, 1)
task.cancel()
await asyncio.sleep(0)
assert not task.done()
assert list(restarted.inbox_dir.glob("*.json")) == []
assert len(list(restarted.processing_dir.glob("*.json"))) == 1
allow_recovery.set()
with pytest.raises(asyncio.CancelledError):
await asyncio.wait_for(task, timeout=1)
assert recovery_finished.is_set()
state_at_cancellation = (
sorted(path.name for path in restarted.inbox_dir.glob("*.json")),
sorted(path.name for path in restarted.processing_dir.glob("*.json")),
)
assert len(state_at_cancellation[0]) == 1
assert state_at_cancellation[1] == []
recovered_payload = json.loads(
(restarted.inbox_dir / state_at_cancellation[0][0]).read_text(encoding="utf-8")
)
assert recovered_payload["delivery"]["attempts"] == 1
assert submitted == []
await asyncio.sleep(0)
assert (
sorted(path.name for path in restarted.inbox_dir.glob("*.json")),
sorted(path.name for path in restarted.processing_dir.glob("*.json")),
) == state_at_cancellation
finally:
allow_recovery.set()
task.cancel()
with suppress(asyncio.CancelledError):
await task
@pytest.mark.asyncio
async def test_local_trigger_queue_submits_bound_inbound_message(tmp_path: Path) -> None:
store = LocalTriggerStore(tmp_path)
@@ -390,6 +549,7 @@ async def test_local_trigger_queue_rejects_unavailable_target_channel(tmp_path:
@pytest.mark.asyncio
async def test_local_trigger_queue_waits_for_submitted_turn_before_ack(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
store = LocalTriggerStore(tmp_path)
trigger = store.create(
@@ -401,6 +561,17 @@ async def test_local_trigger_queue_waits_for_submitted_turn_before_ack(
delivery = store.enqueue(trigger.id, "Review failed CI")
submitted: list[InboundMessage] = []
release = asyncio.Event()
completion_started = threading.Event()
allow_completion = threading.Event()
complete_delivery = store.complete_delivery
def _complete_delivery(claimed: TriggerDelivery) -> None:
completion_started.set()
if not allow_completion.wait(timeout=1):
raise TimeoutError("test did not release delivery completion")
complete_delivery(claimed)
monkeypatch.setattr(store, "complete_delivery", _complete_delivery)
async def _submit_turn(msg: InboundMessage):
submitted.append(msg)
@@ -430,6 +601,13 @@ async def test_local_trigger_queue_waits_for_submitted_turn_before_ack(
assert stored.last_status is None
release.set()
assert await asyncio.to_thread(completion_started.wait, 1)
stored = store.get(trigger.id)
assert stored is not None
assert stored.last_status is None
assert list(store.processing_dir.glob("*.json"))
allow_completion.set()
for _ in range(100):
stored = store.get(trigger.id)
if stored and stored.last_status == "ok":
@@ -444,6 +622,7 @@ async def test_local_trigger_queue_waits_for_submitted_turn_before_ack(
record = _read_run_record(store, delivery.id)
assert record["status"] == "ok"
finally:
allow_completion.set()
task.cancel()
with suppress(asyncio.CancelledError):
await task
@@ -495,9 +674,115 @@ async def test_local_trigger_queue_requeues_when_submitted_turn_is_interrupted(
await task
@pytest.mark.asyncio
async def test_accepted_turn_cancellation_settles_delivery_without_replay(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
store = LocalTriggerStore(tmp_path)
trigger = store.create(
name="CI review",
channel="websocket",
chat_id="chat-1",
session_key="websocket:chat-1",
)
delivery = store.enqueue(trigger.id, "Review failed CI")
inbound: asyncio.Queue[InboundMessage] = asyncio.Queue()
work_accepted = asyncio.Event()
release_work = asyncio.Event()
effects: list[str] = []
publish_count = 0
async def publish(msg: InboundMessage) -> None:
nonlocal publish_count
publish_count += 1
await inbound.put(msg)
coordinator = LocalTriggerTurnCoordinator(
publish_inbound=publish,
dispatch=lambda _msg: asyncio.sleep(0),
is_running=lambda: True,
)
async def agent_worker() -> None:
msg = await inbound.get()
work_accepted.set()
await release_work.wait()
effects.append(msg.content)
coordinator.complete(msg)
agent_task = asyncio.create_task(agent_worker())
queue_task = asyncio.create_task(
run_local_trigger_queue(
store=store,
submit_turn=coordinator.submit,
is_channel_enabled=_channel_is_enabled,
poll_interval_s=0.01,
)
)
restarted_task: asyncio.Task[None] | None = None
try:
await asyncio.wait_for(work_accepted.wait(), timeout=1)
queue_task.cancel()
with pytest.raises(asyncio.CancelledError):
await asyncio.wait_for(queue_task, timeout=1)
# The delivery is durably settled before cancellation is reported even
# though the independently-owned agent turn has not completed yet.
assert effects == []
assert not list(store.processing_dir.glob("*.json"))
assert _read_run_record(store, delivery.id)["status"] == "accepted"
restarted = LocalTriggerStore(tmp_path)
replayed: list[InboundMessage] = []
poll_completed = threading.Event()
claim_deliveries = restarted.claim_deliveries
def tracked_claim_deliveries(*, limit: int = 20) -> list[TriggerDelivery]:
try:
return claim_deliveries(limit=limit)
finally:
poll_completed.set()
async def replay_submit(msg: InboundMessage) -> None:
replayed.append(msg)
monkeypatch.setattr(restarted, "claim_deliveries", tracked_claim_deliveries)
restarted_task = asyncio.create_task(
run_local_trigger_queue(
store=restarted,
submit_turn=replay_submit,
is_channel_enabled=_channel_is_enabled,
poll_interval_s=1,
)
)
assert await asyncio.to_thread(poll_completed.wait, 1)
assert replayed == []
assert restarted.claim_deliveries() == []
release_work.set()
await asyncio.wait_for(agent_task, timeout=1)
assert effects == ["Review failed CI"]
assert publish_count == 1
finally:
release_work.set()
if restarted_task is not None:
restarted_task.cancel()
with suppress(asyncio.CancelledError):
await restarted_task
queue_task.cancel()
with suppress(asyncio.CancelledError):
await queue_task
if not agent_task.done():
agent_task.cancel()
with suppress(asyncio.CancelledError):
await agent_task
@pytest.mark.asyncio
async def test_local_trigger_queue_does_not_retry_completed_agent_failure(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
store = LocalTriggerStore(tmp_path)
trigger = store.create(
@@ -508,6 +793,17 @@ async def test_local_trigger_queue_does_not_retry_completed_agent_failure(
)
delivery = store.enqueue(trigger.id, "Review failed CI")
started = asyncio.Event()
completion_started = threading.Event()
allow_completion = threading.Event()
complete_delivery = store.complete_delivery
def _complete_delivery(claimed: TriggerDelivery) -> None:
completion_started.set()
if not allow_completion.wait(timeout=1):
raise TimeoutError("test did not release delivery completion")
complete_delivery(claimed)
monkeypatch.setattr(store, "complete_delivery", _complete_delivery)
async def _submit_turn(_msg: InboundMessage):
started.set()
@@ -523,6 +819,13 @@ async def test_local_trigger_queue_does_not_retry_completed_agent_failure(
)
try:
await asyncio.wait_for(started.wait(), timeout=1)
assert await asyncio.to_thread(completion_started.wait, 1)
stored = store.get(trigger.id)
assert stored is not None
assert stored.last_status is None
assert list(store.processing_dir.glob("*.json"))
allow_completion.set()
for _ in range(100):
stored = store.get(trigger.id)
if stored and stored.last_status == "error":
@@ -540,6 +843,7 @@ async def test_local_trigger_queue_does_not_retry_completed_agent_failure(
assert record["status"] == "error"
assert record["error"] == "model failed"
finally:
allow_completion.set()
task.cancel()
with suppress(asyncio.CancelledError):
await task
+49
View File
@@ -1,5 +1,7 @@
from __future__ import annotations
import asyncio
import threading
from types import SimpleNamespace
from unittest.mock import ANY, AsyncMock, MagicMock
@@ -160,3 +162,50 @@ async def test_fork_handler_maps_invalid_source_and_internal_failure_to_stable_e
channel.logger.warning.assert_called_once_with("fork_chat failed: {}", ANY)
channel.send_webui_protocol_error.assert_awaited_once_with(connection, "fork_chat_failed")
@pytest.mark.asyncio
async def test_fork_cancellation_waits_for_creation_and_client_attachment(
monkeypatch: pytest.MonkeyPatch,
) -> None:
started = threading.Event()
release = threading.Event()
connection = object()
def blocked_create(*_args, **_kwargs) -> tuple[str, str]:
started.set()
assert release.wait(timeout=1)
return "fork-id", "websocket:fork-id"
channel = SimpleNamespace(
send_webui_protocol_error=AsyncMock(),
attach_webui_fork=AsyncMock(),
gateway=SimpleNamespace(session_manager=object()),
logger=SimpleNamespace(warning=MagicMock()),
)
monkeypatch.setattr(forking, "create_webui_chat_fork", blocked_create)
task = asyncio.create_task(
forking.handle_webui_fork_chat(
channel,
connection,
{"source_chat_id": "source", "before_user_index": 0},
)
)
assert await asyncio.to_thread(started.wait, 1)
try:
task.cancel()
await asyncio.sleep(0)
assert not task.done()
channel.attach_webui_fork.assert_not_awaited()
finally:
release.set()
with pytest.raises(asyncio.CancelledError):
await asyncio.wait_for(task, timeout=1)
channel.attach_webui_fork.assert_awaited_once_with(
connection,
fork_id="fork-id",
fork_key="websocket:fork-id",
)
channel.send_webui_protocol_error.assert_not_awaited()
+29
View File
@@ -2,8 +2,11 @@
import gzip
import json
import time
from types import SimpleNamespace
from nanobot.webui.http_utils import http_json_response
from nanobot.webui.ws_http import GatewayHTTPHandler
def test_http_json_response_compresses_large_payload_when_gzip_is_accepted() -> None:
@@ -36,3 +39,29 @@ def test_http_json_response_does_not_compress_small_payload() -> None:
assert "Content-Encoding" not in response.headers
assert response.headers["Vary"] == "Accept-Encoding"
assert json.loads(response.body) == payload
def test_slow_http_log_records_route_scale_without_user_path_content() -> None:
records: list[str] = []
class _Logger:
def warning(self, message: str, *args: object) -> None:
records.append(message.format(*args))
handler = SimpleNamespace(_log=_Logger())
secret = "private-session-token"
GatewayHTTPHandler._log_slow_http(
handler,
f"/api/sessions/{secret}?query={secret}",
SimpleNamespace(status_code=200),
time.perf_counter() - 2,
input_chars=80,
)
assert len(records) == 1
assert "operation=/api/sessions" in records[0]
assert "status=200" in records[0]
assert "input_chars=80" in records[0]
assert "duration_ms=" in records[0]
assert secret not in records[0]
+22 -8
View File
@@ -38,9 +38,12 @@ async def test_browser_flow_retries_current_server_and_ignores_unrelated_reload_
received: dict[str, object] = {}
reload_calls = 0
async def validate_url_target(_url: str) -> tuple[bool, str]:
return True, ""
monkeypatch.setattr(
"nanobot.webui.mcp_oauth_api.validate_url_target",
lambda _url: (True, ""),
"nanobot.webui.mcp_oauth_api.async_validate_url_target",
validate_url_target,
)
async def connect(servers, _registry, *, oauth_handlers):
@@ -115,9 +118,12 @@ async def test_remote_http_flow_accepts_a_pasted_loopback_callback(
connection = _Connection()
received: dict[str, object] = {}
async def validate_url_target(_url: str) -> tuple[bool, str]:
return True, ""
monkeypatch.setattr(
"nanobot.webui.mcp_oauth_api.validate_url_target",
lambda _url: (True, ""),
"nanobot.webui.mcp_oauth_api.async_validate_url_target",
validate_url_target,
)
async def connect(_servers, _registry, *, oauth_handlers):
@@ -187,9 +193,13 @@ async def test_browser_flow_surfaces_provider_denial_without_callback_descriptio
monkeypatch: pytest.MonkeyPatch,
) -> None:
manager = McpOAuthManager()
async def validate_url_target(_url: str) -> tuple[bool, str]:
return True, ""
monkeypatch.setattr(
"nanobot.webui.mcp_oauth_api.validate_url_target",
lambda _url: (True, ""),
"nanobot.webui.mcp_oauth_api.async_validate_url_target",
validate_url_target,
)
async def connect(_servers, _registry, *, oauth_handlers):
@@ -233,9 +243,13 @@ async def test_browser_flow_blocks_unsafe_authorization_url(
state: str,
) -> None:
manager = McpOAuthManager()
async def validate_url_target(_url: str) -> tuple[bool, str]:
return url_is_safe, "private address"
monkeypatch.setattr(
"nanobot.webui.mcp_oauth_api.validate_url_target",
lambda _url: (url_is_safe, "private address"),
"nanobot.webui.mcp_oauth_api.async_validate_url_target",
validate_url_target,
)
async def connect(_servers, _registry, *, oauth_handlers):
+176
View File
@@ -526,3 +526,179 @@ async def test_version_check_route_enforces_auth_and_bounds_failures(
assert failed.status_code == 500
assert json.loads(failed.body) == {"error": "version check failed"}
assert "upstream secret body" not in failed.body.decode()
@pytest.mark.asyncio
async def test_mcp_oauth_start_cancellation_waits_for_config_and_flow_start(
monkeypatch: pytest.MonkeyPatch,
) -> None:
config_started = threading.Event()
release_config = threading.Event()
start_entered = asyncio.Event()
release_start = asyncio.Event()
effects: list[str] = []
cfg = SimpleNamespace(
type="streamableHttp",
auth="oauth",
url="https://app.xmind.com/api/mcp",
)
def ensure_server(_query, *, config_path=None):
config_started.set()
assert release_config.wait(timeout=1)
effects.append("config_saved")
return "xmind", cfg
async def start(
name,
config,
redirect_uri,
*,
reload_mcp,
reset_credentials=False,
):
assert (name, config, redirect_uri) == (
"xmind",
cfg,
"https://gateway.example/auth/mcp/callback",
)
assert callable(reload_mcp)
assert reset_credentials is False
effects.append("start_entered")
start_entered.set()
await release_start.wait()
effects.append("start_settled")
return {
"status": "authorization_required",
"flow_id": "flow-123",
"name": name,
}
monkeypatch.setattr(
"nanobot.webui.settings_routes.ensure_mcp_oauth_server",
ensure_server,
)
router = _router()
router._mcp_oauth = SimpleNamespace(start=start)
request = _mutation_request(
"/api/settings/mcp-oauth/start",
{"name": "xmind"},
)
task = asyncio.create_task(
router.dispatch(None, request, "/api/settings/mcp-oauth/start")
)
assert await asyncio.to_thread(config_started.wait, 1)
try:
task.cancel()
await asyncio.sleep(0)
assert not task.done()
release_config.set()
await asyncio.wait_for(start_entered.wait(), timeout=1)
assert effects == ["config_saved", "start_entered"]
assert not task.done()
finally:
release_config.set()
release_start.set()
with pytest.raises(asyncio.CancelledError):
await asyncio.wait_for(task, timeout=1)
assert effects == ["config_saved", "start_entered", "start_settled"]
await asyncio.sleep(0.05)
assert effects == ["config_saved", "start_entered", "start_settled"]
@pytest.mark.asyncio
async def test_model_settings_cancellation_waits_for_mutation_and_runtime_refresh(
monkeypatch: pytest.MonkeyPatch,
) -> None:
mutation_started = threading.Event()
release_mutation = threading.Event()
effects: list[str] = []
def update(_query, *, config_path=None):
mutation_started.set()
assert release_mutation.wait(timeout=1)
effects.append("config_saved")
return {"updated": True}
def refresh_runtime_config() -> None:
effects.append("runtime_refreshed")
monkeypatch.setattr(
"nanobot.webui.settings_routes.update_agent_settings",
update,
)
path = "/api/settings/update"
request = _mutation_request(path, {"model_preset": "Codex"})
task = asyncio.create_task(
_router(refresh_runtime_config=refresh_runtime_config).dispatch(
None,
request,
path,
)
)
assert await asyncio.to_thread(mutation_started.wait, 1)
try:
task.cancel()
await asyncio.sleep(0)
assert not task.done()
assert effects == []
finally:
release_mutation.set()
with pytest.raises(asyncio.CancelledError):
await asyncio.wait_for(task, timeout=1)
assert effects == ["config_saved", "runtime_refreshed"]
await asyncio.sleep(0.05)
assert effects == ["config_saved", "runtime_refreshed"]
@pytest.mark.asyncio
async def test_image_settings_cancellation_waits_for_runtime_application(
monkeypatch: pytest.MonkeyPatch,
) -> None:
reload_started = asyncio.Event()
release_reload = asyncio.Event()
effects: list[str] = []
def update(_query, *, config_path=None):
effects.append("config_saved")
return {"requires_restart": True}
async def reload_image(_bus):
effects.append("reload_started")
reload_started.set()
await release_reload.wait()
effects.append("reload_settled")
return {"ok": True, "requires_restart": False}
monkeypatch.setattr(
"nanobot.webui.settings_routes.update_image_generation_settings",
update,
)
monkeypatch.setattr(
"nanobot.webui.settings_routes.request_image_generation_reload",
reload_image,
)
path = "/api/settings/image-generation/update"
request = _mutation_request(path, {"enabled": True})
task = asyncio.create_task(_router().dispatch(None, request, path))
await asyncio.wait_for(reload_started.wait(), timeout=1)
assert effects == ["config_saved", "reload_started"]
task.cancel()
await asyncio.sleep(0)
assert not task.done()
release_reload.set()
with pytest.raises(asyncio.CancelledError):
await asyncio.wait_for(task, timeout=1)
assert effects == ["config_saved", "reload_started", "reload_settled"]
await asyncio.sleep(0.05)
assert effects == ["config_saved", "reload_started", "reload_settled"]