Compare commits

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