diff --git a/nanobot/agent/autocompact.py b/nanobot/agent/autocompact.py index a8ed9c5cc..11db45f42 100644 --- a/nanobot/agent/autocompact.py +++ b/nanobot/agent/autocompact.py @@ -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: diff --git a/nanobot/agent/automation_turns.py b/nanobot/agent/automation_turns.py index 072223b7e..55d10a5f7 100644 --- a/nanobot/agent/automation_turns.py +++ b/nanobot/agent/automation_turns.py @@ -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) diff --git a/nanobot/agent/loop.py b/nanobot/agent/loop.py index 458345ed4..acecefb39 100644 --- a/nanobot/agent/loop.py +++ b/nanobot/agent/loop.py @@ -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 diff --git a/nanobot/agent/memory.py b/nanobot/agent/memory.py index 98e89c03d..fe14c3103 100644 --- a/nanobot/agent/memory.py +++ b/nanobot/agent/memory.py @@ -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, diff --git a/nanobot/agent/runner.py b/nanobot/agent/runner.py index feffd7cc7..4877d5a95 100644 --- a/nanobot/agent/runner.py +++ b/nanobot/agent/runner.py @@ -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) diff --git a/nanobot/agent/tools/cli_apps.py b/nanobot/agent/tools/cli_apps.py index 76758a386..b46228790 100644 --- a/nanobot/agent/tools/cli_apps.py +++ b/nanobot/agent/tools/cli_apps.py @@ -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}") diff --git a/nanobot/agent/tools/cron.py b/nanobot/agent/tools/cron.py index 6a4fd653a..f610cc1da 100644 --- a/nanobot/agent/tools/cron.py +++ b/nanobot/agent/tools/cron.py @@ -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}" diff --git a/nanobot/agent/tools/filesystem.py b/nanobot/agent/tools/filesystem.py index 8b8de45b3..d4c6113d7 100644 --- a/nanobot/agent/tools/filesystem.py +++ b/nanobot/agent/tools/filesystem.py @@ -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}" diff --git a/nanobot/agent/tools/long_task.py b/nanobot/agent/tools/long_task.py index 359b60c6c..356ef55ea 100644 --- a/nanobot/agent/tools/long_task.py +++ b/nanobot/agent/tools/long_task.py @@ -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 = { diff --git a/nanobot/agent/tools/mcp.py b/nanobot/agent/tools/mcp.py index f0b8bbe5c..0c3136e36 100644 --- a/nanobot/agent/tools/mcp.py +++ b/nanobot/agent/tools/mcp.py @@ -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 {} ({})", diff --git a/nanobot/agent/tools/runtime_control.py b/nanobot/agent/tools/runtime_control.py index 85227b4de..00d68fdd1 100644 --- a/nanobot/agent/tools/runtime_control.py +++ b/nanobot/agent/tools/runtime_control.py @@ -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 diff --git a/nanobot/agent/tools/search.py b/nanobot/agent/tools/search.py index feba01a56..c5fb8f58b 100644 --- a/nanobot/agent/tools/search.py +++ b/nanobot/agent/tools/search.py @@ -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.""" diff --git a/nanobot/agent/tools/self.py b/nanobot/agent/tools/self.py index 1bc5000d0..1846bfa81 100644 --- a/nanobot/agent/tools/self.py +++ b/nanobot/agent/tools/self.py @@ -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"]) diff --git a/nanobot/agent/tools/shell.py b/nanobot/agent/tools/shell.py index 54881fe23..57638a8b9 100644 --- a/nanobot/agent/tools/shell.py +++ b/nanobot/agent/tools/shell.py @@ -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)): diff --git a/nanobot/agent/tools/web.py b/nanobot/agent/tools/web.py index e5a4125bb..7c89605d5 100644 --- a/nanobot/agent/tools/web.py +++ b/nanobot/agent/tools/web.py @@ -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) diff --git a/nanobot/apps/cli/service.py b/nanobot/apps/cli/service.py index 2ccf3305b..bda9465e8 100644 --- a/nanobot/apps/cli/service.py +++ b/nanobot/apps/cli/service.py @@ -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, + ) diff --git a/nanobot/channels/dingtalk/runtime.py b/nanobot/channels/dingtalk/runtime.py index f00d75e49..edcedacfa 100644 --- a/nanobot/channels/dingtalk/runtime.py +++ b/nanobot/channels/dingtalk/runtime.py @@ -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: diff --git a/nanobot/channels/napcat/runtime.py b/nanobot/channels/napcat/runtime.py index b431f2359..e1a0d9725 100644 --- a/nanobot/channels/napcat/runtime.py +++ b/nanobot/channels/napcat/runtime.py @@ -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 diff --git a/nanobot/channels/napcat/tests/test_napcat_channel.py b/nanobot/channels/napcat/tests/test_napcat_channel.py index 0b89e3a7a..8cb495d0e 100644 --- a/nanobot/channels/napcat/tests/test_napcat_channel.py +++ b/nanobot/channels/napcat/tests/test_napcat_channel.py @@ -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"}) diff --git a/nanobot/channels/qq/runtime.py b/nanobot/channels/qq/runtime.py index b744fec14..b39552aa8 100644 --- a/nanobot/channels/qq/runtime.py +++ b/nanobot/channels/qq/runtime.py @@ -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 diff --git a/nanobot/channels/slack/runtime.py b/nanobot/channels/slack/runtime.py index eb81a9d6b..a632983a1 100644 --- a/nanobot/channels/slack/runtime.py +++ b/nanobot/channels/slack/runtime.py @@ -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" 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) diff --git a/nanobot/channels/slack/tests/test_slack_channel.py b/nanobot/channels/slack/tests/test_slack_channel.py index 69e858e89..32f2fef8a 100644 --- a/nanobot/channels/slack/tests/test_slack_channel.py +++ b/nanobot/channels/slack/tests/test_slack_channel.py @@ -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 diff --git a/nanobot/channels/telegram/runtime.py b/nanobot/channels/telegram/runtime.py index e475b1afe..b234f7b68 100644 --- a/nanobot/channels/telegram/runtime.py +++ b/nanobot/channels/telegram/runtime.py @@ -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( diff --git a/nanobot/channels/telegram/tests/test_telegram_channel.py b/nanobot/channels/telegram/tests/test_telegram_channel.py index bf5ded907..9b5f6fcde 100644 --- a/nanobot/channels/telegram/tests/test_telegram_channel.py +++ b/nanobot/channels/telegram/tests/test_telegram_channel.py @@ -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( diff --git a/nanobot/channels/websocket/runtime.py b/nanobot/channels/websocket/runtime.py index fc8e8bf20..572fc9579 100644 --- a/nanobot/channels/websocket/runtime.py +++ b/nanobot/channels/websocket/runtime.py @@ -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, diff --git a/nanobot/channels/websocket/tests/test_websocket_http_routes.py b/nanobot/channels/websocket/tests/test_websocket_http_routes.py index 7db009867..d64009856 100644 --- a/nanobot/channels/websocket/tests/test_websocket_http_routes.py +++ b/nanobot/channels/websocket/tests/test_websocket_http_routes.py @@ -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 diff --git a/nanobot/cli/gateway_runtime.py b/nanobot/cli/gateway_runtime.py index 8d5d0f3a0..50b6965db 100644 --- a/nanobot/cli/gateway_runtime.py +++ b/nanobot/cli/gateway_runtime.py @@ -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: diff --git a/nanobot/command/builtin.py b/nanobot/command/builtin.py index ac8f8421a..1604f4d4c 100644 --- a/nanobot/command/builtin.py +++ b/nanobot/command/builtin.py @@ -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] diff --git a/nanobot/cron/bound_runner.py b/nanobot/cron/bound_runner.py index eff121a13..f4d1bc1bb 100644 --- a/nanobot/cron/bound_runner.py +++ b/nanobot/cron/bound_runner.py @@ -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( diff --git a/nanobot/cron/service.py b/nanobot/cron/service.py index 538eeb801..7673c509d 100644 --- a/nanobot/cron/service.py +++ b/nanobot/cron/service.py @@ -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-`` 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- 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-`` 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- 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: diff --git a/nanobot/providers/image_generation.py b/nanobot/providers/image_generation.py index 996e63b9b..0daf914f7 100644 --- a/nanobot/providers/image_generation.py +++ b/nanobot/providers/image_generation.py @@ -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, ) diff --git a/nanobot/sdk/clients.py b/nanobot/sdk/clients.py index 08ff1b9db..f0ce6024f 100644 --- a/nanobot/sdk/clients.py +++ b/nanobot/sdk/clients.py @@ -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, diff --git a/nanobot/security/network.py b/nanobot/security/network.py index a8ffb9fa6..23b5f94d1 100644 --- a/nanobot/security/network.py +++ b/nanobot/security/network.py @@ -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): diff --git a/nanobot/session/async_compat.py b/nanobot/session/async_compat.py new file mode 100644 index 000000000..0d264e608 --- /dev/null +++ b/nanobot/session/async_compat.py @@ -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)) diff --git a/nanobot/session/manager.py b/nanobot/session/manager.py index 8c919c3a5..2fbee9865 100644 --- a/nanobot/session/manager.py +++ b/nanobot/session/manager.py @@ -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) diff --git a/nanobot/session/recovery.py b/nanobot/session/recovery.py index 8d8e248b8..ce40868ea 100644 --- a/nanobot/session/recovery.py +++ b/nanobot/session/recovery.py @@ -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( diff --git a/nanobot/session/webui_turns.py b/nanobot/session/webui_turns.py index a403b69b0..45f434ad4 100644 --- a/nanobot/session/webui_turns.py +++ b/nanobot/session/webui_turns.py @@ -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, diff --git a/nanobot/triggers/local_runner.py b/nanobot/triggers/local_runner.py index 39f82c1c0..b5a3e71b8 100644 --- a/nanobot/triggers/local_runner.py +++ b/nanobot/triggers/local_runner.py @@ -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( diff --git a/nanobot/triggers/local_store.py b/nanobot/triggers/local_store.py index 3e5e93d8b..3d4df115d 100644 --- a/nanobot/triggers/local_store.py +++ b/nanobot/triggers/local_store.py @@ -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, diff --git a/nanobot/utils/cancellation.py b/nanobot/utils/cancellation.py index 29f68948e..c00b4d3f8 100644 --- a/nanobot/utils/cancellation.py +++ b/nanobot/utils/cancellation.py @@ -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() diff --git a/nanobot/webui/forking.py b/nanobot/webui/forking.py index 169777556..40aba7cdb 100644 --- a/nanobot/webui/forking.py +++ b/nanobot/webui/forking.py @@ -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()) diff --git a/nanobot/webui/mcp_oauth_api.py b/nanobot/webui/mcp_oauth_api.py index 9a64023a1..9b7640a62 100644 --- a/nanobot/webui/mcp_oauth_api.py +++ b/nanobot/webui/mcp_oauth_api.py @@ -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) diff --git a/nanobot/webui/settings_capabilities.py b/nanobot/webui/settings_capabilities.py index c4bc5cdd7..cd909f5e1 100644 --- a/nanobot/webui/settings_capabilities.py +++ b/nanobot/webui/settings_capabilities.py @@ -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 diff --git a/nanobot/webui/settings_models.py b/nanobot/webui/settings_models.py index a5cb4f9b2..2aa319090 100644 --- a/nanobot/webui/settings_models.py +++ b/nanobot/webui/settings_models.py @@ -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") diff --git a/nanobot/webui/settings_routes.py b/nanobot/webui/settings_routes.py index cbb5b4df0..4f35dd425 100644 --- a/nanobot/webui/settings_routes.py +++ b/nanobot/webui/settings_routes.py @@ -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) diff --git a/nanobot/webui/settings_services.py b/nanobot/webui/settings_services.py index 5f635b1d4..c696df878 100644 --- a/nanobot/webui/settings_services.py +++ b/nanobot/webui/settings_services.py @@ -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) + ) diff --git a/nanobot/webui/settings_system.py b/nanobot/webui/settings_system.py index fa4a5f7b2..88eaf2566 100644 --- a/nanobot/webui/settings_system.py +++ b/nanobot/webui/settings_system.py @@ -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)) diff --git a/nanobot/webui/ws_http.py b/nanobot/webui/ws_http.py index eb64e778b..dbb29ee82 100644 --- a/nanobot/webui/ws_http.py +++ b/nanobot/webui/ws_http.py @@ -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 diff --git a/tests/agent/test_mcp_reconnect_crash.py b/tests/agent/test_mcp_reconnect_crash.py index 86df12611..6ea004821 100644 --- a/tests/agent/test_mcp_reconnect_crash.py +++ b/tests/agent/test_mcp_reconnect_crash.py @@ -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", diff --git a/tests/agent/test_runner_errors.py b/tests/agent/test_runner_errors.py index 11332c093..37ea47bc3 100644 --- a/tests/agent/test_runner_errors.py +++ b/tests/agent/test_runner_errors.py @@ -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] diff --git a/tests/agent/test_runner_injections.py b/tests/agent/test_runner_injections.py index 0d17c2e5a..204759495 100644 --- a/tests/agent/test_runner_injections.py +++ b/tests/agent/test_runner_injections.py @@ -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) diff --git a/tests/agent/test_workspace_scope.py b/tests/agent/test_workspace_scope.py index 515246045..47da2855a 100644 --- a/tests/agent/test_workspace_scope.py +++ b/tests/agent/test_workspace_scope.py @@ -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, diff --git a/tests/agent/tools/test_long_task.py b/tests/agent/tools/test_long_task.py index 0889b061b..521156606 100644 --- a/tests/agent/tools/test_long_task.py +++ b/tests/agent/tools/test_long_task.py @@ -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) diff --git a/tests/apps/test_cli_subprocess_env.py b/tests/apps/test_cli_subprocess_env.py index 7f0ae1f56..243fbddcf 100644 --- a/tests/apps/test_cli_subprocess_env.py +++ b/tests/apps/test_cli_subprocess_env.py @@ -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, diff --git a/tests/cli/test_gateway_runtime.py b/tests/cli/test_gateway_runtime.py index d5fb0c512..347c0e5a9 100644 --- a/tests/cli/test_gateway_runtime.py +++ b/tests/cli/test_gateway_runtime.py @@ -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] diff --git a/tests/cli_apps/test_service.py b/tests/cli_apps/test_service.py index ebd0038aa..ed2cb0319 100644 --- a/tests/cli_apps/test_service.py +++ b/tests/cli_apps/test_service.py @@ -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() diff --git a/tests/cli_apps/test_tool.py b/tests/cli_apps/test_tool.py index 81738b9bc..a46104a0f 100644 --- a/tests/cli_apps/test_tool.py +++ b/tests/cli_apps/test_tool.py @@ -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() diff --git a/tests/command/test_router_dispatchable.py b/tests/command/test_router_dispatchable.py index ef745f979..ad2ade220 100644 --- a/tests/command/test_router_dispatchable.py +++ b/tests/command/test_router_dispatchable.py @@ -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).""" diff --git a/tests/cron/test_bound_runner.py b/tests/cron/test_bound_runner.py new file mode 100644 index 000000000..26f2bea91 --- /dev/null +++ b/tests/cron/test_bound_runner.py @@ -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 diff --git a/tests/cron/test_cron_service.py b/tests/cron/test_cron_service.py index 26d53dc62..4841378e1 100644 --- a/tests/cron/test_cron_service.py +++ b/tests/cron/test_cron_service.py @@ -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" diff --git a/tests/cron/test_cron_tool_schema_contract.py b/tests/cron/test_cron_tool_schema_contract.py index 9fadd1e2f..51674143c 100644 --- a/tests/cron/test_cron_tool_schema_contract.py +++ b/tests/cron/test_cron_tool_schema_contract.py @@ -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 diff --git a/tests/cron/test_cron_tool_worker.py b/tests/cron/test_cron_tool_worker.py new file mode 100644 index 000000000..5855bff0c --- /dev/null +++ b/tests/cron/test_cron_tool_worker.py @@ -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) diff --git a/tests/security/test_security_network.py b/tests/security/test_security_network.py index 25b64631b..b59038cbe 100644 --- a/tests/security/test_security_network.py +++ b/tests/security/test_security_network.py @@ -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", + (), + ) diff --git a/tests/session/test_recovery.py b/tests/session/test_recovery.py index 99829bd8e..fe83ce2c6 100644 --- a/tests/session/test_recovery.py +++ b/tests/session/test_recovery.py @@ -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, diff --git a/tests/session/test_session_cache.py b/tests/session/test_session_cache.py index d9fca02a3..9998bd7d9 100644 --- a/tests/session/test_session_cache.py +++ b/tests/session/test_session_cache.py @@ -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 diff --git a/tests/session/test_session_cancellation.py b/tests/session/test_session_cancellation.py new file mode 100644 index 000000000..8cdbe02bb --- /dev/null +++ b/tests/session/test_session_cancellation.py @@ -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 diff --git a/tests/test_event_loop_responsiveness.py b/tests/test_event_loop_responsiveness.py new file mode 100644 index 000000000..ce8078385 --- /dev/null +++ b/tests/test_event_loop_responsiveness.py @@ -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 diff --git a/tests/tools/test_exec_security.py b/tests/tools/test_exec_security.py index 6a6b1337c..d294b349b 100644 --- a/tests/tools/test_exec_security.py +++ b/tests/tools/test_exec_security.py @@ -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 ----------------- diff --git a/tests/tools/test_filesystem_tools.py b/tests/tools/test_filesystem_tools.py index a1e152c59..9c263d93e 100644 --- a/tests/tools/test_filesystem_tools.py +++ b/tests/tools/test_filesystem_tools.py @@ -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" diff --git a/tests/tools/test_mcp_probe.py b/tests/tools/test_mcp_probe.py index 0cad030ac..0721139e6 100644 --- a/tests/tools/test_mcp_probe.py +++ b/tests/tools/test_mcp_probe.py @@ -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, diff --git a/tests/tools/test_mcp_tool.py b/tests/tools/test_mcp_tool.py index 8ed4944f4..b931f76af 100644 --- a/tests/tools/test_mcp_tool.py +++ b/tests/tools/test_mcp_tool.py @@ -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, "" diff --git a/tests/tools/test_search_tools.py b/tests/tools/test_search_tools.py index 10d9f9b5b..679ddf892 100644 --- a/tests/tools/test_search_tools.py +++ b/tests/tools/test_search_tools.py @@ -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() diff --git a/tests/tools/test_web_fetch_jina_privacy.py b/tests/tools/test_web_fetch_jina_privacy.py index fca6632c3..d879ffdf8 100644 --- a/tests/tools/test_web_fetch_jina_privacy.py +++ b/tests/tools/test_web_fetch_jina_privacy.py @@ -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))] diff --git a/tests/tools/test_web_fetch_security.py b/tests/tools/test_web_fetch_security.py index 560522b7f..bba840aa9 100644 --- a/tests/tools/test_web_fetch_security.py +++ b/tests/tools/test_web_fetch_security.py @@ -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_) diff --git a/tests/tools/test_web_fetch_url_sanitization.py b/tests/tools/test_web_fetch_url_sanitization.py index bba12194f..67d7040b3 100644 --- a/tests/tools/test_web_fetch_url_sanitization.py +++ b/tests/tools/test_web_fetch_url_sanitization.py @@ -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))] diff --git a/tests/triggers/test_local_triggers.py b/tests/triggers/test_local_triggers.py index 890c1b78c..07a113dab 100644 --- a/tests/triggers/test_local_triggers.py +++ b/tests/triggers/test_local_triggers.py @@ -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 diff --git a/tests/webui/test_forking.py b/tests/webui/test_forking.py index b99a4a40e..cbfb414a2 100644 --- a/tests/webui/test_forking.py +++ b/tests/webui/test_forking.py @@ -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() diff --git a/tests/webui/test_http_utils.py b/tests/webui/test_http_utils.py index e23710332..29b5e42b3 100644 --- a/tests/webui/test_http_utils.py +++ b/tests/webui/test_http_utils.py @@ -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] diff --git a/tests/webui/test_mcp_oauth_api.py b/tests/webui/test_mcp_oauth_api.py index 1bdeb5d7f..88cb94792 100644 --- a/tests/webui/test_mcp_oauth_api.py +++ b/tests/webui/test_mcp_oauth_api.py @@ -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): diff --git a/tests/webui/test_settings_routes.py b/tests/webui/test_settings_routes.py index 3c4aafafb..60832e69c 100644 --- a/tests/webui/test_settings_routes.py +++ b/tests/webui/test_settings_routes.py @@ -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"]