mirror of
https://github.com/HKUDS/nanobot.git
synced 2026-08-31 08:13:11 +03:00
fix(gateway): keep event loop responsive (NAN-33)
Move blocking filesystem, persistence, subprocess, media, and DNS work off the gateway event loop while preserving existing contracts. Add bounded cancellation and responsiveness regression coverage.
This commit is contained in:
@@ -2,7 +2,9 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Collection
|
||||
import asyncio
|
||||
import inspect
|
||||
from collections.abc import Awaitable, Collection
|
||||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Any, Callable, Coroutine
|
||||
|
||||
@@ -47,9 +49,27 @@ class AutoCompact:
|
||||
return idle_seconds >= self._ttl * 60
|
||||
|
||||
def _has_unarchived_messages(self, key: str) -> bool:
|
||||
session = self.sessions.get_or_create(key)
|
||||
return self._session_has_unarchived_messages(self.sessions.get_or_create(key))
|
||||
|
||||
@staticmethod
|
||||
def _session_has_unarchived_messages(session: Session) -> bool:
|
||||
return session.last_consolidated < len(session.messages)
|
||||
|
||||
def _has_native_async_session_method(self, name: str) -> bool:
|
||||
"""Check the manager's real class, not mock-generated instance attributes."""
|
||||
method = inspect.getattr_static(type(self.sessions), name, None)
|
||||
return inspect.iscoroutinefunction(method)
|
||||
|
||||
async def _list_sessions_nonblocking(self) -> list[dict[str, Any]]:
|
||||
if self._has_native_async_session_method("list_sessions_async"):
|
||||
return await self.sessions.list_sessions_async()
|
||||
return await asyncio.to_thread(self.sessions.list_sessions)
|
||||
|
||||
async def _get_or_create_nonblocking(self, key: str) -> Session:
|
||||
if self._has_native_async_session_method("get_or_create_async"):
|
||||
return await self.sessions.get_or_create_async(key)
|
||||
return await asyncio.to_thread(self.sessions.get_or_create, key)
|
||||
|
||||
@classmethod
|
||||
def _is_internal_session(cls, key: str) -> bool:
|
||||
return key.startswith(cls._INTERNAL_SESSION_PREFIXES)
|
||||
@@ -79,6 +99,31 @@ class AutoCompact:
|
||||
self._archiving.add(key)
|
||||
schedule_background(self._archive(key, runtime=runtime))
|
||||
|
||||
async def check_expired_async(
|
||||
self,
|
||||
schedule_background: Callable[[Coroutine[Any, Any, None]], None],
|
||||
resolve_runtime: Callable[[Session], Awaitable[LLMRuntime]],
|
||||
active_session_keys: Collection[str] = (),
|
||||
) -> None:
|
||||
"""Schedule idle archival without blocking the event loop."""
|
||||
now = datetime.now()
|
||||
active_keys = set(active_session_keys)
|
||||
for info in await self._list_sessions_nonblocking():
|
||||
key = info.get("key", "")
|
||||
if not key or self._is_internal_session(key) or key in self._archiving:
|
||||
continue
|
||||
if key in active_keys or not self._is_expired(info.get("updated_at"), now):
|
||||
continue
|
||||
session = await self._get_or_create_nonblocking(key)
|
||||
if not self._session_has_unarchived_messages(session):
|
||||
continue
|
||||
try:
|
||||
runtime = await resolve_runtime(session)
|
||||
except (KeyError, ValueError):
|
||||
continue
|
||||
self._archiving.add(key)
|
||||
schedule_background(self._archive_async(key, runtime=runtime))
|
||||
|
||||
async def _archive(self, key: str, *, runtime: LLMRuntime) -> None:
|
||||
if self._is_internal_session(key):
|
||||
self._archiving.discard(key)
|
||||
@@ -90,18 +135,38 @@ class AutoCompact:
|
||||
max_suffix=self._RECENT_SUFFIX_MESSAGES,
|
||||
)
|
||||
if summary and summary != "(nothing)":
|
||||
session = self.sessions.get_or_create(key)
|
||||
stored = session_summary_from_metadata(
|
||||
session.metadata,
|
||||
fallback_last_active=session.updated_at,
|
||||
)
|
||||
if stored is not None:
|
||||
self._summaries[key] = stored
|
||||
self._record_stored_summary(key, self.sessions.get_or_create(key))
|
||||
except Exception:
|
||||
logger.exception("Auto-compact: failed for {}", key)
|
||||
finally:
|
||||
self._archiving.discard(key)
|
||||
|
||||
async def _archive_async(self, key: str, *, runtime: LLMRuntime) -> None:
|
||||
if self._is_internal_session(key):
|
||||
self._archiving.discard(key)
|
||||
return
|
||||
try:
|
||||
summary = await self.consolidator.compact_idle_session(
|
||||
key,
|
||||
runtime=runtime,
|
||||
max_suffix=self._RECENT_SUFFIX_MESSAGES,
|
||||
)
|
||||
if summary and summary != "(nothing)":
|
||||
session = await self._get_or_create_nonblocking(key)
|
||||
self._record_stored_summary(key, session)
|
||||
except Exception:
|
||||
logger.exception("Auto-compact: failed for {}", key)
|
||||
finally:
|
||||
self._archiving.discard(key)
|
||||
|
||||
def _record_stored_summary(self, key: str, session: Session) -> None:
|
||||
stored = session_summary_from_metadata(
|
||||
session.metadata,
|
||||
fallback_last_active=session.updated_at,
|
||||
)
|
||||
if stored is not None:
|
||||
self._summaries[key] = stored
|
||||
|
||||
def prepare_session(self, session: Session, key: str) -> tuple[Session, SessionSummary | None]:
|
||||
if self._is_internal_session(key):
|
||||
self._archiving.discard(key)
|
||||
@@ -110,6 +175,28 @@ class AutoCompact:
|
||||
if key in self._archiving or self._is_expired(session.updated_at):
|
||||
logger.info("Auto-compact: reloading session {} (archiving={})", key, key in self._archiving)
|
||||
session = self.sessions.get_or_create(key)
|
||||
return self._prepared_summary(session, key)
|
||||
|
||||
async def prepare_session_async(
|
||||
self,
|
||||
session: Session,
|
||||
key: str,
|
||||
) -> tuple[Session, SessionSummary | None]:
|
||||
"""Prepare a session without blocking on a reload."""
|
||||
if self._is_internal_session(key):
|
||||
self._archiving.discard(key)
|
||||
self._summaries.pop(key, None)
|
||||
return session, None
|
||||
if key in self._archiving or self._is_expired(session.updated_at):
|
||||
logger.info("Auto-compact: reloading session {} (archiving={})", key, key in self._archiving)
|
||||
session = await self._get_or_create_nonblocking(key)
|
||||
return self._prepared_summary(session, key)
|
||||
|
||||
def _prepared_summary(
|
||||
self,
|
||||
session: Session,
|
||||
key: str,
|
||||
) -> tuple[Session, SessionSummary | None]:
|
||||
# Hot path: summary from in-memory dict (process hasn't restarted).
|
||||
entry = self._summaries.pop(key, None)
|
||||
if entry:
|
||||
|
||||
@@ -13,6 +13,19 @@ class AutomationTurnError(RuntimeError):
|
||||
"""Raised when an automation turn reaches the agent and finishes with an error."""
|
||||
|
||||
|
||||
class AutomationTurnAcceptedCancellation(asyncio.CancelledError):
|
||||
"""Cancellation raised after an automation turn was accepted for processing.
|
||||
|
||||
Callers must not replay the turn: the accepted agent work now has independent
|
||||
ownership and may continue after the submitting task is cancelled.
|
||||
"""
|
||||
|
||||
|
||||
def _consume_future_exception(future: asyncio.Future[object]) -> None:
|
||||
if not future.cancelled():
|
||||
future.exception()
|
||||
|
||||
|
||||
async def publish_next_deferred_turn(
|
||||
*,
|
||||
deferred_queues: dict[str, list[InboundMessage]],
|
||||
@@ -70,19 +83,36 @@ class AutomationTurnCoordinator:
|
||||
future: asyncio.Future[OutboundMessage | None] = loop.create_future()
|
||||
self._waiters[turn_id] = future
|
||||
self._pending_messages_by_turn_id[turn_id] = msg
|
||||
accepted = False
|
||||
try:
|
||||
if self._is_running():
|
||||
await self._publish_inbound(msg)
|
||||
accepted = True
|
||||
else:
|
||||
await self._dispatch(msg)
|
||||
# Direct dispatch is given independent task ownership for the
|
||||
# same reason as publishing to the inbound queue: once admitted,
|
||||
# cancelling this submitter must not cancel and then replay the
|
||||
# already-running agent turn.
|
||||
dispatch_future: asyncio.Future[object] = asyncio.ensure_future(
|
||||
self._dispatch(msg)
|
||||
)
|
||||
dispatch_future.add_done_callback(_consume_future_exception)
|
||||
accepted = True
|
||||
await asyncio.shield(dispatch_future)
|
||||
try:
|
||||
return await future
|
||||
except asyncio.CancelledError:
|
||||
except asyncio.CancelledError as exc:
|
||||
if accepted:
|
||||
raise AutomationTurnAcceptedCancellation(*exc.args) from None
|
||||
raise
|
||||
except AutomationTurnError:
|
||||
raise
|
||||
except Exception as exc:
|
||||
raise AutomationTurnError(str(exc) or exc.__class__.__name__) from exc
|
||||
except asyncio.CancelledError as exc:
|
||||
if accepted and not isinstance(exc, AutomationTurnAcceptedCancellation):
|
||||
raise AutomationTurnAcceptedCancellation(*exc.args) from None
|
||||
raise
|
||||
finally:
|
||||
self._waiters.pop(turn_id, None)
|
||||
self._pending_messages_by_turn_id.pop(turn_id, None)
|
||||
|
||||
+155
-32
@@ -67,6 +67,7 @@ from nanobot.security.workspace_access import (
|
||||
reset_workspace_scope,
|
||||
)
|
||||
from nanobot.session import turn_continuation
|
||||
from nanobot.session.async_compat import call_session_manager
|
||||
from nanobot.session.automation_turns import automation_history_overrides
|
||||
from nanobot.session.goal_state import (
|
||||
goal_state_runtime_lines,
|
||||
@@ -540,6 +541,33 @@ class AgentLoop:
|
||||
**extra,
|
||||
)
|
||||
|
||||
async def _get_or_create_session(self, key: str) -> Session:
|
||||
"""Use native async session loading, with a compatibility fallback."""
|
||||
return await call_session_manager(
|
||||
self.sessions,
|
||||
"get_or_create_async",
|
||||
self.sessions.get_or_create,
|
||||
key,
|
||||
)
|
||||
|
||||
async def _save_session(self, session: Session) -> None:
|
||||
"""Use native async session saving, with a compatibility fallback."""
|
||||
await call_session_manager(
|
||||
self.sessions,
|
||||
"save_async",
|
||||
self.sessions.save,
|
||||
session,
|
||||
)
|
||||
|
||||
async def _save_runtime_checkpoint(self, session: Session) -> None:
|
||||
"""Use native async checkpoint saving, with a compatibility fallback."""
|
||||
await call_session_manager(
|
||||
self.sessions,
|
||||
"save_runtime_checkpoint_async",
|
||||
self.sessions.save_runtime_checkpoint,
|
||||
session,
|
||||
)
|
||||
|
||||
def _sync_subagent_runtime_limits(self) -> None:
|
||||
"""Keep subagent runtime limits aligned with mutable loop settings."""
|
||||
self.subagents.max_iterations = self.max_iterations
|
||||
@@ -579,6 +607,30 @@ class AgentLoop:
|
||||
self.sessions.save(session)
|
||||
return self.llm_runtime()
|
||||
|
||||
async def runtime_for_session_async(
|
||||
self,
|
||||
session: Session,
|
||||
*,
|
||||
recover_removed: bool = True,
|
||||
) -> LLMRuntime:
|
||||
"""Resolve a session runtime without blocking on recovery persistence."""
|
||||
name = model_preset_from_metadata(session.metadata)
|
||||
if name is None:
|
||||
return self.llm_runtime()
|
||||
try:
|
||||
return self.runtime_resolver.resolve_preset(name)
|
||||
except KeyError:
|
||||
if not recover_removed or name in self.runtime_resolver.model_presets:
|
||||
raise
|
||||
logger.warning(
|
||||
"Session '{}' references removed model preset '{}'; falling back to default",
|
||||
session.key,
|
||||
name,
|
||||
)
|
||||
session.metadata.pop(SESSION_MODEL_PRESET_METADATA_KEY, None)
|
||||
await self._save_session(session)
|
||||
return self.llm_runtime()
|
||||
|
||||
def set_session_model_preset(
|
||||
self,
|
||||
session_key: str,
|
||||
@@ -591,6 +643,18 @@ class AgentLoop:
|
||||
self.sessions.save(session)
|
||||
return runtime
|
||||
|
||||
async def set_session_model_preset_async(
|
||||
self,
|
||||
session_key: str,
|
||||
name: str,
|
||||
) -> LLMRuntime:
|
||||
"""Validate and persist one session's preset selection without blocking."""
|
||||
runtime = self.runtime_resolver.resolve_preset(name)
|
||||
session = await self._get_or_create_session(session_key)
|
||||
session.metadata[SESSION_MODEL_PRESET_METADATA_KEY] = runtime.model_preset
|
||||
await self._save_session(session)
|
||||
return runtime
|
||||
|
||||
def _publish_runtime_selection(
|
||||
self,
|
||||
runtime: LLMRuntime,
|
||||
@@ -703,17 +767,14 @@ class AgentLoop:
|
||||
session_key=session_key,
|
||||
)
|
||||
|
||||
def _persist_user_message_early(
|
||||
def _stage_user_message_early(
|
||||
self,
|
||||
msg: InboundMessage,
|
||||
session: Session,
|
||||
runtime_context_blocks: list[RuntimeContextBlock] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> bool:
|
||||
"""Persist the triggering user message before the turn starts.
|
||||
|
||||
Returns True if the message was persisted.
|
||||
"""
|
||||
"""Add the triggering user message and recovery markers in memory."""
|
||||
if not turn_continuation.should_persist_user_message(msg.metadata):
|
||||
return False
|
||||
media_paths = [
|
||||
@@ -742,10 +803,45 @@ class AgentLoop:
|
||||
followup_id = msg.metadata.get(PENDING_FOLLOWUP_ID_KEY)
|
||||
if isinstance(followup_id, str) and followup_id:
|
||||
acknowledge_pending_followups(session, [followup_id])
|
||||
self.sessions.save(session)
|
||||
return True
|
||||
return False
|
||||
|
||||
def _persist_user_message_early(
|
||||
self,
|
||||
msg: InboundMessage,
|
||||
session: Session,
|
||||
runtime_context_blocks: list[RuntimeContextBlock] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> bool:
|
||||
"""Synchronously persist the user message for compatibility callers."""
|
||||
persisted = self._stage_user_message_early(
|
||||
msg,
|
||||
session,
|
||||
runtime_context_blocks,
|
||||
**kwargs,
|
||||
)
|
||||
if persisted:
|
||||
self.sessions.save(session)
|
||||
return persisted
|
||||
|
||||
async def _persist_user_message_early_async(
|
||||
self,
|
||||
msg: InboundMessage,
|
||||
session: Session,
|
||||
runtime_context_blocks: list[RuntimeContextBlock] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> bool:
|
||||
"""Persist the user message without blocking the event loop."""
|
||||
persisted = self._stage_user_message_early(
|
||||
msg,
|
||||
session,
|
||||
runtime_context_blocks,
|
||||
**kwargs,
|
||||
)
|
||||
if persisted:
|
||||
await self._save_session(session)
|
||||
return persisted
|
||||
|
||||
def _build_initial_messages(self, ctx: TurnContext) -> list[dict[str, Any]]:
|
||||
"""Build the initial message list for the LLM turn."""
|
||||
assert ctx.session is not None
|
||||
@@ -843,7 +939,7 @@ class AgentLoop:
|
||||
if tool is None:
|
||||
content = "Shell execution is disabled in this nanobot configuration."
|
||||
else:
|
||||
session = ctx.session or self.sessions.get_or_create(ctx.key)
|
||||
session = ctx.session or await AgentLoop._get_or_create_session(self, ctx.key)
|
||||
scope = self.workspace_scopes.for_turn(
|
||||
channel=ctx.msg.channel,
|
||||
message_metadata=metadata,
|
||||
@@ -1001,7 +1097,7 @@ class AgentLoop:
|
||||
public_payload[self._PROVIDER_STATE_CHECKPOINT_VERSION_KEY] = (
|
||||
self._PROVIDER_STATE_CHECKPOINT_VERSION
|
||||
)
|
||||
self._set_runtime_checkpoint(session, public_payload)
|
||||
await self._set_runtime_checkpoint_async(session, public_payload)
|
||||
|
||||
async def _drain_pending(*, limit: int = _MAX_INJECTIONS_PER_TURN) -> list[dict[str, Any]]:
|
||||
"""Drain follow-up messages from the pending queue.
|
||||
@@ -1239,18 +1335,33 @@ class AgentLoop:
|
||||
logger.error("LLM returned error: {}", (result.final_content or "")[:200])
|
||||
return result.final_content, result.tools_used, result.messages, result.stop_reason, result.had_injections
|
||||
|
||||
def _check_expired_sessions_if_due(self) -> None:
|
||||
"""Scan idle sessions no more often than the configured interval."""
|
||||
def _idle_compact_scan_due(self) -> bool:
|
||||
now = time.monotonic()
|
||||
if now < self._next_idle_compact_check_at:
|
||||
return
|
||||
return False
|
||||
self._next_idle_compact_check_at = now + self._idle_compact_check_interval_s
|
||||
return True
|
||||
|
||||
def _check_expired_sessions_if_due(self) -> None:
|
||||
"""Synchronously scan idle sessions for compatibility with direct callers."""
|
||||
if not self._idle_compact_scan_due():
|
||||
return
|
||||
self.auto_compact.check_expired(
|
||||
self.schedule_background,
|
||||
self.runtime_for_session,
|
||||
active_session_keys=self._pending_queues.keys(),
|
||||
)
|
||||
|
||||
async def _check_expired_sessions_if_due_async(self) -> None:
|
||||
"""Scan idle sessions without blocking the event loop."""
|
||||
if not self._idle_compact_scan_due():
|
||||
return
|
||||
await self.auto_compact.check_expired_async(
|
||||
self.schedule_background,
|
||||
self.runtime_for_session_async,
|
||||
active_session_keys=self._pending_queues.keys(),
|
||||
)
|
||||
|
||||
async def run(self) -> None:
|
||||
"""Run the agent loop, dispatching messages as tasks to stay responsive to /stop."""
|
||||
self._running = True
|
||||
@@ -1261,7 +1372,7 @@ class AgentLoop:
|
||||
try:
|
||||
msg = await asyncio.wait_for(self.bus.consume_inbound(), timeout=1.0)
|
||||
except asyncio.TimeoutError:
|
||||
self._check_expired_sessions_if_due()
|
||||
await self._check_expired_sessions_if_due_async()
|
||||
continue
|
||||
except asyncio.CancelledError:
|
||||
# Preserve real task cancellation so shutdown can complete cleanly.
|
||||
@@ -1339,7 +1450,7 @@ class AgentLoop:
|
||||
)
|
||||
continue
|
||||
pending_msg = routed_msg
|
||||
session = self.sessions.get_or_create(effective_key)
|
||||
session = await self._get_or_create_session(effective_key)
|
||||
followup_id = record_pending_followup(session, pending_msg)
|
||||
if followup_id is not None:
|
||||
pending_msg = dataclasses.replace(
|
||||
@@ -1349,7 +1460,7 @@ class AgentLoop:
|
||||
PENDING_FOLLOWUP_ID_KEY: followup_id,
|
||||
},
|
||||
)
|
||||
self.sessions.save(session)
|
||||
await self._save_session(session)
|
||||
try:
|
||||
self._pending_queues[effective_key].put_nowait(pending_msg)
|
||||
except asyncio.QueueFull:
|
||||
@@ -1460,10 +1571,10 @@ class AgentLoop:
|
||||
raise
|
||||
try:
|
||||
key = self._effective_session_key(msg)
|
||||
session = self.sessions.get_or_create(key)
|
||||
session = await self._get_or_create_session(key)
|
||||
if self._restore_runtime_checkpoint(session):
|
||||
self._clear_pending_user_turn(session)
|
||||
self.sessions.save(session)
|
||||
await self._save_session(session)
|
||||
logger.info(
|
||||
"Restored partial context for cancelled session {}",
|
||||
key,
|
||||
@@ -1788,7 +1899,7 @@ class AgentLoop:
|
||||
if ctx.session is None:
|
||||
raise RuntimeError("required session is not active")
|
||||
else:
|
||||
ctx.session = self.sessions.get_or_create(ctx.session_key)
|
||||
ctx.session = await self._get_or_create_session(ctx.session_key)
|
||||
session = ctx.session
|
||||
ctx.ephemeral = ctx.ephemeral or not session.policy.persist
|
||||
tools = ctx.tools or self.tools
|
||||
@@ -1819,16 +1930,16 @@ class AgentLoop:
|
||||
self.workspace_scopes.persist_message_scope(session, msg)
|
||||
|
||||
if self._restore_runtime_checkpoint(session):
|
||||
self.sessions.save(session)
|
||||
await self._save_session(session)
|
||||
if (
|
||||
RECOVERY_INBOUND_METADATA_KEY not in msg.metadata
|
||||
and restore_pending_interruption(session)
|
||||
):
|
||||
self.sessions.save(session)
|
||||
await self._save_session(session)
|
||||
|
||||
async def _compact_session(self, ctx: TurnContext) -> None:
|
||||
session = ctx.require_session()
|
||||
ctx.session, pending = self.auto_compact.prepare_session(
|
||||
ctx.session, pending = await self.auto_compact.prepare_session_async(
|
||||
session,
|
||||
ctx.session_key,
|
||||
)
|
||||
@@ -1865,14 +1976,14 @@ class AgentLoop:
|
||||
# them out of LLM context. /new is excluded because it
|
||||
# intentionally clears the session.
|
||||
if cmd_ctx.raw.lower() != "/new":
|
||||
ctx.input_persisted_early = self._persist_user_message_early(
|
||||
ctx.input_persisted_early = await self._persist_user_message_early_async(
|
||||
ctx.msg, session, _command=True
|
||||
)
|
||||
session.add_message(
|
||||
"assistant", result.content, _command=True
|
||||
)
|
||||
self._clear_pending_user_turn(session)
|
||||
self.sessions.save(session)
|
||||
await self._save_session(session)
|
||||
if not ctx.ephemeral:
|
||||
await self.runtime_event_publisher.session_turn_persisted(
|
||||
ctx.msg,
|
||||
@@ -1887,7 +1998,7 @@ class AgentLoop:
|
||||
session = ctx.require_session()
|
||||
runtime = ctx.runtime
|
||||
if runtime is None:
|
||||
runtime = self.runtime_for_session(session)
|
||||
runtime = await self.runtime_for_session_async(session)
|
||||
ctx.runtime = runtime
|
||||
if ctx.session_key.startswith("dream:"):
|
||||
logger.info(
|
||||
@@ -1931,7 +2042,7 @@ class AgentLoop:
|
||||
# provider compatibility or prompt assembly work. A compatible
|
||||
# staged state replaces this in a second atomic save below.
|
||||
session.provider_state = None
|
||||
self.sessions.save(session)
|
||||
await self._save_session(session)
|
||||
ctx.input_persisted_early = True
|
||||
await ctx.delivery.runtime_admitted(runtime)
|
||||
|
||||
@@ -1985,7 +2096,7 @@ class AgentLoop:
|
||||
elif stored_state is not None:
|
||||
session.provider_state = None
|
||||
if ctx.kind is TurnKind.USER:
|
||||
ctx.input_persisted_early = self._persist_user_message_early(
|
||||
ctx.input_persisted_early = await self._persist_user_message_early_async(
|
||||
ctx.msg,
|
||||
session,
|
||||
runtime_context_blocks=ctx.runtime_context_blocks,
|
||||
@@ -1995,7 +2106,7 @@ class AgentLoop:
|
||||
elif subagent_followup_persisted and staged_provider_state:
|
||||
# Upgrade the replay-safe baseline to the resumable state before
|
||||
# prompt assembly and the first model checkpoint.
|
||||
self.sessions.save(session)
|
||||
await self._save_session(session)
|
||||
ctx.initial_messages = self._build_initial_messages(ctx)
|
||||
|
||||
if ctx.on_progress is None:
|
||||
@@ -2071,6 +2182,9 @@ class AgentLoop:
|
||||
turn_latency_ms=ctx.turn_latency_ms,
|
||||
)
|
||||
ctx.delivery.record_latency(ctx.turn_latency_ms)
|
||||
self._clear_pending_user_turn(session)
|
||||
self._clear_runtime_checkpoint(session)
|
||||
await self._save_session(session)
|
||||
if not ctx.ephemeral:
|
||||
self.schedule_background(
|
||||
self.consolidator.maybe_consolidate_by_tokens(
|
||||
@@ -2078,10 +2192,6 @@ class AgentLoop:
|
||||
runtime=runtime,
|
||||
)
|
||||
)
|
||||
self._clear_pending_user_turn(session)
|
||||
self._clear_runtime_checkpoint(session)
|
||||
self.sessions.save(session)
|
||||
if not ctx.ephemeral:
|
||||
await self.runtime_event_publisher.session_turn_persisted(
|
||||
ctx.msg,
|
||||
ctx.session_key,
|
||||
@@ -2294,11 +2404,24 @@ class AgentLoop:
|
||||
)
|
||||
return True
|
||||
|
||||
def _set_runtime_checkpoint(self, session: Session, payload: dict[str, Any]) -> None:
|
||||
"""Persist the latest in-flight turn state into session metadata."""
|
||||
def _set_runtime_checkpoint(
|
||||
self,
|
||||
session: Session,
|
||||
payload: dict[str, Any],
|
||||
) -> None:
|
||||
"""Synchronously persist a checkpoint for compatibility callers."""
|
||||
session.metadata[self._RUNTIME_CHECKPOINT_KEY] = payload
|
||||
self.sessions.save_runtime_checkpoint(session)
|
||||
|
||||
async def _set_runtime_checkpoint_async(
|
||||
self,
|
||||
session: Session,
|
||||
payload: dict[str, Any],
|
||||
) -> None:
|
||||
"""Persist the latest in-flight turn state without blocking the event loop."""
|
||||
session.metadata[self._RUNTIME_CHECKPOINT_KEY] = payload
|
||||
await self._save_runtime_checkpoint(session)
|
||||
|
||||
def _mark_pending_user_turn(self, session: Session) -> None:
|
||||
session.metadata[self._PENDING_USER_TURN_KEY] = True
|
||||
|
||||
|
||||
+26
-9
@@ -22,6 +22,7 @@ from loguru import logger
|
||||
|
||||
from nanobot.llm_usage.context import llm_usage_source
|
||||
from nanobot.runtime_context import public_history_messages
|
||||
from nanobot.session.async_compat import call_session_manager
|
||||
from nanobot.session.manager import (
|
||||
MIN_COMPACTED_REPLAY_MESSAGES,
|
||||
Session,
|
||||
@@ -824,6 +825,22 @@ class Consolidator:
|
||||
weakref.WeakValueDictionary()
|
||||
)
|
||||
|
||||
async def _get_or_create_session(self, key: str) -> Session:
|
||||
return await call_session_manager(
|
||||
self.sessions,
|
||||
"get_or_create_async",
|
||||
self.sessions.get_or_create,
|
||||
key,
|
||||
)
|
||||
|
||||
async def _save_session(self, session: Session) -> None:
|
||||
await call_session_manager(
|
||||
self.sessions,
|
||||
"save_async",
|
||||
self.sessions.save,
|
||||
session,
|
||||
)
|
||||
|
||||
def get_lock(self, session_key: str) -> asyncio.Lock:
|
||||
"""Return the shared consolidation lock for one session."""
|
||||
return self._locks.setdefault(session_key, asyncio.Lock())
|
||||
@@ -859,13 +876,13 @@ class Consolidator:
|
||||
return []
|
||||
return session.get_history()
|
||||
|
||||
def _persist_last_summary(self, session: Session, summary: str | None) -> None:
|
||||
async def _persist_last_summary(self, session: Session, summary: str | None) -> None:
|
||||
if summary and summary != "(nothing)":
|
||||
session.metadata["_last_summary"] = {
|
||||
"text": summary,
|
||||
"last_active": session.updated_at.isoformat(),
|
||||
}
|
||||
self.sessions.save(session)
|
||||
await self._save_session(session)
|
||||
|
||||
def estimate_session_prompt_tokens(
|
||||
self,
|
||||
@@ -1057,7 +1074,7 @@ class Consolidator:
|
||||
lock = self.get_lock(session.key)
|
||||
async with lock:
|
||||
# Refresh session reference: AutoCompact may have replaced it.
|
||||
fresh = self.sessions.get_or_create(session.key)
|
||||
fresh = await self._get_or_create_session(session.key)
|
||||
if fresh is not session:
|
||||
session = fresh
|
||||
if not session.messages:
|
||||
@@ -1071,7 +1088,7 @@ class Consolidator:
|
||||
runtime=runtime,
|
||||
)
|
||||
if estimated <= 0:
|
||||
self._persist_last_summary(session, last_summary)
|
||||
await self._persist_last_summary(session, last_summary)
|
||||
return
|
||||
if estimated < budget:
|
||||
unconsolidated_count = len(session.messages) - session.last_consolidated
|
||||
@@ -1083,7 +1100,7 @@ class Consolidator:
|
||||
source,
|
||||
unconsolidated_count,
|
||||
)
|
||||
self._persist_last_summary(session, last_summary)
|
||||
await self._persist_last_summary(session, last_summary)
|
||||
return
|
||||
|
||||
for round_num in range(self._MAX_CONSOLIDATION_ROUNDS):
|
||||
@@ -1127,7 +1144,7 @@ class Consolidator:
|
||||
last_summary = summary
|
||||
session.last_consolidated = end_idx
|
||||
session.provider_state = None
|
||||
self.sessions.save(session)
|
||||
await self._save_session(session)
|
||||
if not summary:
|
||||
# LLM is degraded — stop hammering it this call;
|
||||
# the next invocation can retry a fresh chunk.
|
||||
@@ -1143,7 +1160,7 @@ class Consolidator:
|
||||
# Persist the last summary to session metadata so it can be injected
|
||||
# into the runtime context on the next prepare_session() call, aligning
|
||||
# the summary injection strategy with AutoCompact._archive().
|
||||
self._persist_last_summary(session, last_summary)
|
||||
await self._persist_last_summary(session, last_summary)
|
||||
|
||||
async def compact_idle_session(
|
||||
self,
|
||||
@@ -1168,7 +1185,7 @@ class Consolidator:
|
||||
lock = self.get_lock(session_key)
|
||||
async with lock:
|
||||
self.sessions.invalidate(session_key)
|
||||
session = self.sessions.get_or_create(session_key)
|
||||
session = await self._get_or_create_session(session_key)
|
||||
|
||||
archive_start = session.last_consolidated
|
||||
messages_to_archive = list(session.messages[archive_start:])
|
||||
@@ -1193,7 +1210,7 @@ class Consolidator:
|
||||
# through the captured batch so new messages remain eligible next time.
|
||||
session.last_consolidated = archive_end
|
||||
session.provider_state = None
|
||||
self.sessions.save(session)
|
||||
await self._save_session(session)
|
||||
|
||||
visible = session.get_history(
|
||||
max_messages=MIN_COMPACTED_REPLAY_MESSAGES,
|
||||
|
||||
+29
-1
@@ -6,7 +6,7 @@ import asyncio
|
||||
import inspect
|
||||
import os
|
||||
import time
|
||||
from collections.abc import Awaitable, Callable, Iterable
|
||||
from collections.abc import Awaitable, Callable, Iterable, Sized
|
||||
from copy import deepcopy
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
@@ -83,6 +83,22 @@ _MAX_EMPTY_RETRIES = 2
|
||||
_MAX_LENGTH_RECOVERIES = 3
|
||||
_MAX_INJECTIONS_PER_TURN = 3
|
||||
_MAX_INJECTION_CYCLES = 5
|
||||
_SLOW_TOOL_LOG_MS = 1_000
|
||||
|
||||
|
||||
def _tool_input_scale(params: object) -> tuple[int, int]:
|
||||
"""Return bounded structural counts without logging argument content."""
|
||||
if not isinstance(params, dict):
|
||||
return 0, len(params) if isinstance(params, str | bytes) else 0
|
||||
params_dict = cast(dict[object, object], params)
|
||||
items = len(params_dict)
|
||||
chars = 0
|
||||
for value in params_dict.values():
|
||||
if isinstance(value, str | bytes):
|
||||
chars += len(value)
|
||||
elif isinstance(value, list | tuple | set | dict):
|
||||
items += len(cast(Sized, value))
|
||||
return items, chars
|
||||
|
||||
|
||||
def _restore_outer_whitespace(content: str, original: str | None) -> str:
|
||||
@@ -1521,6 +1537,7 @@ class AgentRunner:
|
||||
RuntimeError(prep_error) if spec.fail_on_tool_error else None
|
||||
)
|
||||
await hook.before_execute_tool(context, tool_call, tool, params)
|
||||
tool_started_at = time.perf_counter()
|
||||
try:
|
||||
if tool is not None:
|
||||
result = await tool.execute(**params)
|
||||
@@ -1549,6 +1566,17 @@ class AgentRunner:
|
||||
if spec.fail_on_tool_error:
|
||||
return payload, event, exc
|
||||
return payload, event, None
|
||||
finally:
|
||||
duration_ms = int((time.perf_counter() - tool_started_at) * 1000)
|
||||
if duration_ms >= _SLOW_TOOL_LOG_MS:
|
||||
input_items, input_chars = _tool_input_scale(params)
|
||||
logger.warning(
|
||||
"slow tool operation={} input_items={} input_chars={} duration_ms={}",
|
||||
tool_call.name,
|
||||
input_items,
|
||||
input_chars,
|
||||
duration_ms,
|
||||
)
|
||||
|
||||
if is_tool_error_result(result):
|
||||
await hook.on_execute_tool_error(context, tool_call, tool, params, result)
|
||||
|
||||
@@ -4,7 +4,10 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import inspect
|
||||
from pathlib import Path
|
||||
from typing import TypedDict
|
||||
|
||||
from pydantic import Field
|
||||
|
||||
@@ -24,6 +27,14 @@ from nanobot.runtime_context import RuntimeContextBlock, wrap_runtime_context_li
|
||||
from nanobot.security.workspace_access import current_tool_workspace
|
||||
|
||||
|
||||
class _CliAppRunKwargs(TypedDict):
|
||||
args: list[str]
|
||||
json_output: bool
|
||||
working_dir: str | None
|
||||
timeout: int | None
|
||||
restrict_to_workspace: bool
|
||||
|
||||
|
||||
class CliAppsToolConfig(Base):
|
||||
"""CLI Apps tool configuration."""
|
||||
|
||||
@@ -147,14 +158,17 @@ class CliAppsTool(Tool):
|
||||
)
|
||||
workspace = access.project_path or self.workspace
|
||||
manager = CliAppManager(workspace=workspace, runtime=self.runtime)
|
||||
run_kwargs: _CliAppRunKwargs = {
|
||||
"args": args or [],
|
||||
"json_output": bool(json),
|
||||
"working_dir": working_dir,
|
||||
"timeout": timeout,
|
||||
"restrict_to_workspace": access.restrict_to_workspace,
|
||||
}
|
||||
try:
|
||||
return manager.run(
|
||||
name,
|
||||
args=args or [],
|
||||
json_output=bool(json),
|
||||
working_dir=working_dir,
|
||||
timeout=timeout,
|
||||
restrict_to_workspace=access.restrict_to_workspace,
|
||||
)
|
||||
run_async = inspect.getattr_static(type(manager), "run_async", None)
|
||||
if inspect.iscoroutinefunction(run_async):
|
||||
return await manager.run_async(name, **run_kwargs)
|
||||
return await asyncio.to_thread(manager.run, name, **run_kwargs)
|
||||
except CliAppError as exc:
|
||||
return ToolResult.error(f"Error: {exc.message}")
|
||||
|
||||
@@ -143,14 +143,41 @@ class CronTool(Tool):
|
||||
tz: str | None = None,
|
||||
at: str | None = None,
|
||||
job_id: str | None = None,
|
||||
) -> str:
|
||||
if action == "add" and self._in_cron_context.get():
|
||||
return ToolResult.error(
|
||||
"Error: cannot schedule new jobs from within a cron job execution"
|
||||
)
|
||||
return await self._cron.run_sync(
|
||||
self._execute_sync,
|
||||
action,
|
||||
name,
|
||||
message,
|
||||
every_seconds,
|
||||
cron_expr,
|
||||
tz,
|
||||
at,
|
||||
job_id,
|
||||
)
|
||||
|
||||
def _execute_sync(
|
||||
self,
|
||||
action: str,
|
||||
name: str | None,
|
||||
message: str,
|
||||
every_seconds: int | None,
|
||||
cron_expr: str | None,
|
||||
tz: str | None,
|
||||
at: str | None,
|
||||
job_id: str | None,
|
||||
) -> str:
|
||||
if action == "add":
|
||||
if self._in_cron_context.get():
|
||||
return ToolResult.error("Error: cannot schedule new jobs from within a cron job execution")
|
||||
return self._add_job(name, message, every_seconds, cron_expr, tz, at)
|
||||
elif action == "list":
|
||||
if action == "list":
|
||||
return self._list_jobs()
|
||||
elif action == "remove":
|
||||
if action == "remove":
|
||||
return self._remove_job(job_id)
|
||||
return f"Unknown action: {action}"
|
||||
|
||||
|
||||
@@ -2,9 +2,11 @@
|
||||
|
||||
# pyright: reportPrivateUsage=false, reportUnusedFunction=false
|
||||
|
||||
import asyncio
|
||||
import difflib
|
||||
import mimetypes
|
||||
import os
|
||||
import threading
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
@@ -21,6 +23,7 @@ from nanobot.agent.tools.schema import (
|
||||
)
|
||||
from nanobot.config_base import Base
|
||||
from nanobot.security.workspace_access import current_tool_workspace
|
||||
from nanobot.utils.cancellation import shield_and_drain
|
||||
from nanobot.utils.helpers import build_image_content_blocks, detect_image_mime
|
||||
|
||||
|
||||
@@ -664,22 +667,31 @@ def _match_covers_line(match: _MatchSpan, line: int) -> bool:
|
||||
return match.line <= line <= _match_end_line(match)
|
||||
|
||||
|
||||
def _find_exact_matches(content: str, old_text: str) -> list[_MatchSpan]:
|
||||
def _find_exact_matches(
|
||||
content: str,
|
||||
old_text: str,
|
||||
*,
|
||||
max_matches: int | None = None,
|
||||
) -> list[_MatchSpan]:
|
||||
matches: list[_MatchSpan] = []
|
||||
start = 0
|
||||
while True:
|
||||
idx = content.find(old_text, start)
|
||||
search_start = 0
|
||||
line_start = 0
|
||||
line = 1
|
||||
while max_matches is None or len(matches) < max_matches:
|
||||
idx = content.find(old_text, search_start)
|
||||
if idx == -1:
|
||||
break
|
||||
line += content.count("\n", line_start, idx)
|
||||
matches.append(
|
||||
_MatchSpan(
|
||||
start=idx,
|
||||
end=idx + len(old_text),
|
||||
text=content[idx : idx + len(old_text)],
|
||||
line=content.count("\n", 0, idx) + 1,
|
||||
line=line,
|
||||
)
|
||||
)
|
||||
start = idx + max(1, len(old_text))
|
||||
line_start = idx
|
||||
search_start = idx + max(1, len(old_text))
|
||||
return matches
|
||||
|
||||
|
||||
@@ -735,27 +747,36 @@ def _find_quote_matches(content: str, old_text: str) -> list[_MatchSpan]:
|
||||
norm_content = _normalize_quotes(content)
|
||||
norm_old = _normalize_quotes(old_text)
|
||||
matches: list[_MatchSpan] = []
|
||||
start = 0
|
||||
search_start = 0
|
||||
line_start = 0
|
||||
line = 1
|
||||
while True:
|
||||
idx = norm_content.find(norm_old, start)
|
||||
idx = norm_content.find(norm_old, search_start)
|
||||
if idx == -1:
|
||||
break
|
||||
line += content.count("\n", line_start, idx)
|
||||
matches.append(
|
||||
_MatchSpan(
|
||||
start=idx,
|
||||
end=idx + len(old_text),
|
||||
text=content[idx : idx + len(old_text)],
|
||||
line=content.count("\n", 0, idx) + 1,
|
||||
line=line,
|
||||
)
|
||||
)
|
||||
start = idx + max(1, len(norm_old))
|
||||
line_start = idx
|
||||
search_start = idx + max(1, len(norm_old))
|
||||
return matches
|
||||
|
||||
|
||||
def _find_matches(content: str, old_text: str) -> list[_MatchSpan]:
|
||||
"""Locate all matches using progressively looser strategies."""
|
||||
def _find_matches(
|
||||
content: str,
|
||||
old_text: str,
|
||||
*,
|
||||
max_exact_matches: int | None = None,
|
||||
) -> list[_MatchSpan]:
|
||||
"""Locate matches using progressively looser strategies."""
|
||||
for matcher in (
|
||||
lambda: _find_exact_matches(content, old_text),
|
||||
lambda: _find_exact_matches(content, old_text, max_matches=max_exact_matches),
|
||||
lambda: _find_trim_matches(content, old_text),
|
||||
lambda: _find_trim_matches(content, old_text, normalize_quotes=True),
|
||||
lambda: _find_quote_matches(content, old_text),
|
||||
@@ -869,6 +890,43 @@ class EditFileTool(_FsTool):
|
||||
new_text: str | None = None,
|
||||
replace_all: bool = False, occurrence: int | None = None,
|
||||
line_hint: int | None = None, expected_replacements: int | None = None, **kwargs: Any,
|
||||
) -> str:
|
||||
cancelled = threading.Event()
|
||||
commit_lock = threading.Lock()
|
||||
try:
|
||||
return await asyncio.to_thread(
|
||||
self._execute_sync,
|
||||
path=path,
|
||||
old_text=old_text,
|
||||
new_text=new_text,
|
||||
replace_all=replace_all,
|
||||
occurrence=occurrence,
|
||||
line_hint=line_hint,
|
||||
expected_replacements=expected_replacements,
|
||||
cancelled=cancelled,
|
||||
commit_lock=commit_lock,
|
||||
)
|
||||
except asyncio.CancelledError:
|
||||
cancelled.set()
|
||||
# If a commit already started, do not report cancellation until the
|
||||
# file bytes and FileStates record are settled. Otherwise, taking
|
||||
# the lock first guarantees the worker observes ``cancelled`` before
|
||||
# it can mutate the target.
|
||||
await shield_and_drain(
|
||||
asyncio.to_thread(self._wait_for_commit, commit_lock)
|
||||
)
|
||||
raise
|
||||
|
||||
@staticmethod
|
||||
def _wait_for_commit(commit_lock: threading.Lock) -> None:
|
||||
with commit_lock:
|
||||
pass
|
||||
|
||||
def _execute_sync(
|
||||
self, *, path: str | None, old_text: str | None,
|
||||
new_text: str | None, replace_all: bool, occurrence: int | None,
|
||||
line_hint: int | None, expected_replacements: int | None,
|
||||
cancelled: threading.Event, commit_lock: threading.Lock,
|
||||
) -> str:
|
||||
try:
|
||||
if not path:
|
||||
@@ -892,9 +950,12 @@ class EditFileTool(_FsTool):
|
||||
# Create-file semantics: old_text='' + file doesn't exist → create
|
||||
if not file_exists:
|
||||
if old_text == "":
|
||||
fp.parent.mkdir(parents=True, exist_ok=True)
|
||||
fp.write_text(new_text, encoding="utf-8")
|
||||
self._file_states.record_write(fp)
|
||||
with commit_lock:
|
||||
if cancelled.is_set():
|
||||
return ToolResult.error("Error: edit_file cancelled.")
|
||||
fp.parent.mkdir(parents=True, exist_ok=True)
|
||||
fp.write_text(new_text, encoding="utf-8")
|
||||
self._file_states.record_write(fp)
|
||||
return f"Successfully created {fp}"
|
||||
return self._file_not_found_msg(path, fp)
|
||||
|
||||
@@ -912,8 +973,11 @@ class EditFileTool(_FsTool):
|
||||
content = raw.decode("utf-8")
|
||||
if content.strip():
|
||||
return ToolResult.error(f"Error: Cannot create file — {path} already exists and is not empty.")
|
||||
fp.write_text(new_text, encoding="utf-8")
|
||||
self._file_states.record_write(fp)
|
||||
with commit_lock:
|
||||
if cancelled.is_set():
|
||||
return ToolResult.error("Error: edit_file cancelled.")
|
||||
fp.write_text(new_text, encoding="utf-8")
|
||||
self._file_states.record_write(fp)
|
||||
return f"Successfully edited {fp}"
|
||||
|
||||
# Read-before-edit check
|
||||
@@ -923,7 +987,11 @@ class EditFileTool(_FsTool):
|
||||
uses_crlf = b"\r\n" in raw
|
||||
content = raw.decode("utf-8").replace("\r\n", "\n")
|
||||
norm_old = old_text.replace("\r\n", "\n")
|
||||
matches = _find_matches(content, norm_old)
|
||||
matches = _find_matches(
|
||||
content,
|
||||
norm_old,
|
||||
max_exact_matches=occurrence,
|
||||
)
|
||||
|
||||
if not matches:
|
||||
return self._not_found_msg(old_text, content, path)
|
||||
@@ -1000,8 +1068,11 @@ class EditFileTool(_FsTool):
|
||||
if uses_crlf:
|
||||
new_content = new_content.replace("\n", "\r\n")
|
||||
|
||||
fp.write_bytes(new_content.encode("utf-8"))
|
||||
self._file_states.record_write(fp)
|
||||
with commit_lock:
|
||||
if cancelled.is_set():
|
||||
return ToolResult.error("Error: edit_file cancelled.")
|
||||
fp.write_bytes(new_content.encode("utf-8"))
|
||||
self._file_states.record_write(fp)
|
||||
msg = f"Successfully edited {fp}"
|
||||
if warning:
|
||||
msg = f"{warning}\n{msg}"
|
||||
|
||||
@@ -17,6 +17,7 @@ from nanobot.agent.tools.context import RequestContext, ToolContext, current_req
|
||||
from nanobot.agent.tools.schema import StringSchema, tool_parameters_schema
|
||||
from nanobot.bus.runtime_events import GoalStateChanged, RuntimeEventBus, RuntimeEventContext
|
||||
from nanobot.runtime_context import RuntimeContextBlock, wrap_runtime_context_lines
|
||||
from nanobot.session.async_compat import call_session_manager
|
||||
from nanobot.session.goal_state import (
|
||||
GOAL_STATE_KEY,
|
||||
MAX_GOAL_OBJECTIVE_CHARS,
|
||||
@@ -28,6 +29,7 @@ from nanobot.session.goal_state import (
|
||||
sustained_goal_active,
|
||||
)
|
||||
from nanobot.session.turn_continuation import reset_goal_continuation_rounds
|
||||
from nanobot.utils.cancellation import shield_and_drain
|
||||
from nanobot.utils.prompt_templates import render_template
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -60,36 +62,68 @@ class _GoalToolsMixin:
|
||||
self._sessions = sessions
|
||||
self._runtime_events = runtime_events
|
||||
|
||||
def _session(self):
|
||||
async def _get_or_create_session(self, key: str):
|
||||
return await call_session_manager(
|
||||
self._sessions,
|
||||
"get_or_create_async",
|
||||
self._sessions.get_or_create,
|
||||
key,
|
||||
)
|
||||
|
||||
async def _save_session(self, session: Any) -> None:
|
||||
await call_session_manager(
|
||||
self._sessions,
|
||||
"save_async",
|
||||
self._sessions.save,
|
||||
session,
|
||||
)
|
||||
|
||||
async def _session(self):
|
||||
request_ctx = current_request_context()
|
||||
if request_ctx is None:
|
||||
return None
|
||||
key = request_ctx.session_key
|
||||
if not key:
|
||||
return None
|
||||
return self._sessions.get_or_create(key)
|
||||
return await self._get_or_create_session(key)
|
||||
|
||||
def _goal_mutation_allowed(self) -> bool:
|
||||
return current_request_context() is not None and goal_mutation_allowed()
|
||||
|
||||
def _save_goal_state(
|
||||
async def _save_goal_state(
|
||||
self,
|
||||
sess: Any,
|
||||
blob: dict[str, Any],
|
||||
*,
|
||||
reset_continuation: bool = False,
|
||||
revoke_permission: bool = False,
|
||||
) -> None:
|
||||
previous_metadata = deepcopy(sess.metadata)
|
||||
sess.metadata[GOAL_STATE_KEY] = blob
|
||||
discard_legacy_goal_state_key(sess.metadata)
|
||||
if reset_continuation:
|
||||
reset_goal_continuation_rounds(sess.metadata)
|
||||
saved = False
|
||||
|
||||
async def save_and_publish() -> None:
|
||||
nonlocal saved
|
||||
sess.metadata[GOAL_STATE_KEY] = blob
|
||||
discard_legacy_goal_state_key(sess.metadata)
|
||||
if reset_continuation:
|
||||
reset_goal_continuation_rounds(sess.metadata)
|
||||
try:
|
||||
await self._save_session(sess)
|
||||
except BaseException:
|
||||
sess.metadata.clear()
|
||||
sess.metadata.update(previous_metadata)
|
||||
raise
|
||||
saved = True
|
||||
await self._publish_goal_state_changed(sess.metadata)
|
||||
|
||||
try:
|
||||
self._sessions.save(sess)
|
||||
except BaseException:
|
||||
sess.metadata.clear()
|
||||
sess.metadata.update(previous_metadata)
|
||||
raise
|
||||
await shield_and_drain(save_and_publish())
|
||||
finally:
|
||||
# This ContextVar belongs to the caller task, not the settlement task.
|
||||
# Apply the post-save permission effect here even when cancellation was
|
||||
# delayed until the durable save and runtime notification completed.
|
||||
if revoke_permission and saved:
|
||||
revoke_goal_mutation_permission()
|
||||
|
||||
async def _publish_goal_state_changed(self, metadata: dict[str, Any]) -> None:
|
||||
runtime_events = self._runtime_events
|
||||
@@ -175,7 +209,7 @@ class CreateGoalTool(Tool, _GoalToolsMixin):
|
||||
) -> RuntimeContextBlock | None:
|
||||
if not request.session_key:
|
||||
return None
|
||||
session = self._sessions.get_or_create(request.session_key)
|
||||
session = await self._get_or_create_session(request.session_key)
|
||||
goal_start_requested = explicit_goal_requested(request.metadata)
|
||||
goal_active = sustained_goal_active(session.metadata)
|
||||
if not goal_start_requested and not goal_active:
|
||||
@@ -197,7 +231,7 @@ class CreateGoalTool(Tool, _GoalToolsMixin):
|
||||
ui_summary: str | None = None,
|
||||
**kwargs: Any,
|
||||
) -> str:
|
||||
sess = self._session()
|
||||
sess = await self._session()
|
||||
if sess is None:
|
||||
return ToolResult.error(
|
||||
"Error: create_goal requires an active chat session (missing routing context)."
|
||||
@@ -225,8 +259,7 @@ class CreateGoalTool(Tool, _GoalToolsMixin):
|
||||
"ui_summary": summary,
|
||||
"started_at": _iso_now(),
|
||||
}
|
||||
self._save_goal_state(sess, blob, reset_continuation=True)
|
||||
await self._publish_goal_state_changed(sess.metadata)
|
||||
await self._save_goal_state(sess, blob, reset_continuation=True)
|
||||
extra = f"\nSummary line: {summary}" if summary else ""
|
||||
return (
|
||||
"Goal recorded. Keep working toward the objective using ordinary tools. "
|
||||
@@ -305,7 +338,7 @@ class UpdateGoalTool(Tool, _GoalToolsMixin):
|
||||
ui_summary: str | None = None,
|
||||
**kwargs: Any,
|
||||
) -> str:
|
||||
sess = self._session()
|
||||
sess = await self._session()
|
||||
if sess is None:
|
||||
return ToolResult.error("Error: update_goal requires an active chat session.")
|
||||
prior = parse_goal_state(goal_state_raw(sess.metadata))
|
||||
@@ -340,8 +373,7 @@ class UpdateGoalTool(Tool, _GoalToolsMixin):
|
||||
"previous_objective": str(prior.get("objective") or ""),
|
||||
"recap": (recap or "").strip(),
|
||||
}
|
||||
self._save_goal_state(sess, blob, reset_continuation=True)
|
||||
await self._publish_goal_state_changed(sess.metadata)
|
||||
await self._save_goal_state(sess, blob, reset_continuation=True)
|
||||
extra = f"\nSummary line: {summary}" if summary else ""
|
||||
return "Goal replaced. Continue toward the new objective using ordinary tools." + extra
|
||||
|
||||
@@ -359,9 +391,7 @@ class UpdateGoalTool(Tool, _GoalToolsMixin):
|
||||
}
|
||||
if normalized == "complete":
|
||||
blob["completed_at"] = ended
|
||||
self._save_goal_state(sess, blob)
|
||||
revoke_goal_mutation_permission()
|
||||
await self._publish_goal_state_changed(sess.metadata)
|
||||
await self._save_goal_state(sess, blob, revoke_permission=True)
|
||||
|
||||
tail = (recap or "").strip()
|
||||
label = {
|
||||
|
||||
@@ -20,10 +20,10 @@ from nanobot.agent.tools.base import Tool, ToolResult
|
||||
from nanobot.agent.tools.registry import ToolRegistry
|
||||
from nanobot.security.network import (
|
||||
PinnedDNSAsyncTransport,
|
||||
async_resolve_url_target,
|
||||
async_validate_url_target,
|
||||
env_proxy_applies_to_url,
|
||||
httpx_env_proxy_mounts,
|
||||
resolve_url_target,
|
||||
validate_url_target,
|
||||
)
|
||||
from nanobot.utils.cancellation import task_is_cancelling
|
||||
|
||||
@@ -249,7 +249,7 @@ async def _probe_http_url(url: str, timeout: float = 3.0) -> bool:
|
||||
port = parsed.port
|
||||
if not port:
|
||||
port = 443 if parsed.scheme == "https" else 80
|
||||
ok, _, resolved_ips = resolve_url_target(url)
|
||||
ok, _, resolved_ips = await async_resolve_url_target(url)
|
||||
if not ok:
|
||||
return False
|
||||
if env_proxy_applies_to_url(url):
|
||||
@@ -298,7 +298,7 @@ def _pinned_transport_kwargs() -> dict[str, Any]:
|
||||
|
||||
async def _validate_mcp_request_url(request: httpx.Request) -> None:
|
||||
"""Validate each outgoing MCP HTTP request, including redirect targets."""
|
||||
ok, error = validate_url_target(str(request.url))
|
||||
ok, error = await async_validate_url_target(str(request.url))
|
||||
if not ok:
|
||||
raise httpx.RequestError(
|
||||
f"Blocked unsafe MCP URL {_redact_url(str(request.url))} ({error})",
|
||||
@@ -1031,7 +1031,7 @@ async def connect_mcp_servers(
|
||||
return False
|
||||
|
||||
if transport_type in {"sse", "streamableHttp"}:
|
||||
ok, error = validate_url_target(cfg.url)
|
||||
ok, error = await async_validate_url_target(cfg.url)
|
||||
if not ok:
|
||||
logger.warning(
|
||||
"MCP server '{}': blocked unsafe URL {} ({})",
|
||||
|
||||
@@ -107,6 +107,13 @@ class RuntimeControl(Protocol):
|
||||
session_key: str | None,
|
||||
) -> LLMRuntime: ...
|
||||
|
||||
async def set_model_preset_async(
|
||||
self,
|
||||
name: str,
|
||||
*,
|
||||
session_key: str | None,
|
||||
) -> LLMRuntime: ...
|
||||
|
||||
def set_max_iterations(self, value: int) -> None: ...
|
||||
|
||||
def set_context_window_tokens(self, value: int) -> LLMRuntime: ...
|
||||
@@ -162,6 +169,12 @@ class _RuntimeControlTarget(Protocol):
|
||||
|
||||
def set_session_model_preset(self, session_key: str, name: str) -> LLMRuntime: ...
|
||||
|
||||
async def set_session_model_preset_async(
|
||||
self,
|
||||
session_key: str,
|
||||
name: str,
|
||||
) -> LLMRuntime: ...
|
||||
|
||||
|
||||
class AgentRuntimeControl:
|
||||
"""Allowlisted adapter from agent-loop state to ``RuntimeControl``."""
|
||||
@@ -208,6 +221,16 @@ class AgentRuntimeControl:
|
||||
return self.__target.set_session_model_preset(session_key, name)
|
||||
return self.__target.set_model_preset(name)
|
||||
|
||||
async def set_model_preset_async(
|
||||
self,
|
||||
name: str,
|
||||
*,
|
||||
session_key: str | None,
|
||||
) -> LLMRuntime:
|
||||
if session_key is not None:
|
||||
return await self.__target.set_session_model_preset_async(session_key, name)
|
||||
return self.__target.set_model_preset(name)
|
||||
|
||||
def set_max_iterations(self, value: int) -> None:
|
||||
self.__target.max_iterations = value
|
||||
self.__target.subagents.max_iterations = value
|
||||
|
||||
+103
-53
@@ -4,9 +4,12 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import fnmatch
|
||||
import os
|
||||
import re
|
||||
import threading
|
||||
import time
|
||||
from contextlib import suppress
|
||||
from pathlib import Path, PurePosixPath
|
||||
from typing import Any, Iterable, TypeVar
|
||||
@@ -125,6 +128,8 @@ class _SearchTool(_FsTool):
|
||||
class FindFilesTool(_SearchTool):
|
||||
"""Find files by path fragment, glob, or type."""
|
||||
_scopes = {"core", "subagent"}
|
||||
_MAX_SCAN_PATHS = 500_000
|
||||
_MAX_SCAN_SECONDS = 30.0
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
@@ -218,66 +223,111 @@ class FindFilesTool(_SearchTool):
|
||||
offset: int = 0,
|
||||
**kwargs: Any,
|
||||
) -> str:
|
||||
cancelled = threading.Event()
|
||||
try:
|
||||
target = self._resolve(path or ".")
|
||||
if not target.exists():
|
||||
return ToolResult.error(f"Error: Path not found: {path}")
|
||||
if not (target.is_dir() or target.is_file()):
|
||||
return ToolResult.error(f"Error: Unsupported path: {path}")
|
||||
|
||||
if sort not in {"path", "modified"}:
|
||||
return ToolResult.error("Error: sort must be 'path' or 'modified'")
|
||||
|
||||
limit = (
|
||||
_DEFAULT_FILE_HEAD_LIMIT
|
||||
if head_limit is None
|
||||
else None if head_limit == 0 else head_limit
|
||||
return await asyncio.to_thread(
|
||||
self._execute_sync,
|
||||
path=path,
|
||||
query=query,
|
||||
glob=glob,
|
||||
file_type=type,
|
||||
include_dirs=include_dirs,
|
||||
sort=sort,
|
||||
head_limit=head_limit,
|
||||
offset=offset,
|
||||
cancelled=cancelled,
|
||||
)
|
||||
root = target if target.is_dir() else target.parent
|
||||
matches: list[tuple[str, float]] = []
|
||||
|
||||
for candidate in self._iter_paths(target, include_dirs=include_dirs):
|
||||
if candidate.is_dir() and not include_dirs:
|
||||
continue
|
||||
rel_path = candidate.relative_to(root).as_posix()
|
||||
display_path = self._display_path(candidate, root)
|
||||
name = candidate.name
|
||||
|
||||
if glob and not _match_glob(rel_path, name, glob):
|
||||
continue
|
||||
if candidate.is_file() and not _matches_type(name, type):
|
||||
continue
|
||||
if candidate.is_dir() and type:
|
||||
continue
|
||||
if not _matches_query(display_path, query):
|
||||
continue
|
||||
try:
|
||||
mtime = candidate.stat().st_mtime
|
||||
except OSError:
|
||||
mtime = 0.0
|
||||
suffix = "/" if candidate.is_dir() else ""
|
||||
matches.append((display_path + suffix, mtime))
|
||||
|
||||
if sort == "modified":
|
||||
matches.sort(key=lambda item: (-item[1], item[0]))
|
||||
else:
|
||||
matches.sort(key=lambda item: item[0])
|
||||
|
||||
paths = [item[0] for item in matches]
|
||||
paged, truncated = _paginate(paths, limit, offset)
|
||||
if not paged:
|
||||
return "No files found"
|
||||
|
||||
result = "\n".join(paged)
|
||||
note = _pagination_note(limit, offset, truncated)
|
||||
if note:
|
||||
result += "\n\n" + note
|
||||
return result
|
||||
except asyncio.CancelledError:
|
||||
cancelled.set()
|
||||
raise
|
||||
except PermissionError as e:
|
||||
return ToolResult.error(f"Error: {e}")
|
||||
except Exception as e:
|
||||
return ToolResult.error(f"Error finding files: {e}")
|
||||
|
||||
def _execute_sync(
|
||||
self,
|
||||
*,
|
||||
path: str,
|
||||
query: str | None,
|
||||
glob: str | None,
|
||||
file_type: str | None,
|
||||
include_dirs: bool,
|
||||
sort: str,
|
||||
head_limit: int | None,
|
||||
offset: int,
|
||||
cancelled: threading.Event,
|
||||
) -> str:
|
||||
target = self._resolve(path or ".")
|
||||
if not target.exists():
|
||||
return ToolResult.error(f"Error: Path not found: {path}")
|
||||
if not (target.is_dir() or target.is_file()):
|
||||
return ToolResult.error(f"Error: Unsupported path: {path}")
|
||||
|
||||
if sort not in {"path", "modified"}:
|
||||
return ToolResult.error("Error: sort must be 'path' or 'modified'")
|
||||
|
||||
limit = (
|
||||
_DEFAULT_FILE_HEAD_LIMIT
|
||||
if head_limit is None
|
||||
else None if head_limit == 0 else head_limit
|
||||
)
|
||||
root = target if target.is_dir() else target.parent
|
||||
matches: list[tuple[str, float]] = []
|
||||
deadline = time.monotonic() + self._MAX_SCAN_SECONDS
|
||||
scanned = 0
|
||||
|
||||
for candidate in self._iter_paths(target, include_dirs=include_dirs):
|
||||
if cancelled.is_set():
|
||||
raise RuntimeError("find_files scan cancelled")
|
||||
scanned += 1
|
||||
if scanned > self._MAX_SCAN_PATHS:
|
||||
return ToolResult.error(
|
||||
f"Error: find_files scan exceeded {self._MAX_SCAN_PATHS} paths; "
|
||||
"narrow path, query, glob, or type and retry."
|
||||
)
|
||||
if time.monotonic() > deadline:
|
||||
return ToolResult.error(
|
||||
f"Error: find_files scan exceeded {self._MAX_SCAN_SECONDS:g} seconds; "
|
||||
"narrow path, query, glob, or type and retry."
|
||||
)
|
||||
if candidate.is_dir() and not include_dirs:
|
||||
continue
|
||||
rel_path = candidate.relative_to(root).as_posix()
|
||||
display_path = self._display_path(candidate, root)
|
||||
name = candidate.name
|
||||
|
||||
if glob and not _match_glob(rel_path, name, glob):
|
||||
continue
|
||||
if candidate.is_file() and not _matches_type(name, file_type):
|
||||
continue
|
||||
if candidate.is_dir() and file_type:
|
||||
continue
|
||||
if not _matches_query(display_path, query):
|
||||
continue
|
||||
try:
|
||||
mtime = candidate.stat().st_mtime
|
||||
except OSError:
|
||||
mtime = 0.0
|
||||
suffix = "/" if candidate.is_dir() else ""
|
||||
matches.append((display_path + suffix, mtime))
|
||||
|
||||
if sort == "modified":
|
||||
matches.sort(key=lambda item: (-item[1], item[0]))
|
||||
else:
|
||||
matches.sort(key=lambda item: item[0])
|
||||
|
||||
paths = [item[0] for item in matches]
|
||||
paged, truncated = _paginate(paths, limit, offset)
|
||||
if not paged:
|
||||
return "No files found"
|
||||
|
||||
result = "\n".join(paged)
|
||||
note = _pagination_note(limit, offset, truncated)
|
||||
if note:
|
||||
result += "\n\n" + note
|
||||
return result
|
||||
|
||||
|
||||
class GrepTool(_SearchTool):
|
||||
"""Search file contents using a regex-like pattern."""
|
||||
|
||||
@@ -370,7 +370,7 @@ class MyTool(Tool):
|
||||
if not self._modify_allowed:
|
||||
return ToolResult.error("Error: set is disabled (tools.my.allow_set is false)")
|
||||
if action in ("modify", "set"):
|
||||
return self._modify(key, value)
|
||||
return await self._modify_async(key, value)
|
||||
return f"Unknown action: {action}"
|
||||
|
||||
# -- inspect --
|
||||
@@ -492,6 +492,11 @@ class MyTool(Tool):
|
||||
return ToolResult.error(f"Error: '{key}' is read-only and cannot be modified")
|
||||
return self._modify_scratchpad(key, value)
|
||||
|
||||
async def _modify_async(self, key: str | None, value: Any) -> str:
|
||||
if key == "model_preset":
|
||||
return await self._modify_model_preset_async(value)
|
||||
return self._modify(key, value)
|
||||
|
||||
def _modify_model_preset(self, value: Any) -> str:
|
||||
if not isinstance(value, str) or not value.strip():
|
||||
return ToolResult.error("Error: 'model_preset' must be a non-empty string")
|
||||
@@ -520,6 +525,34 @@ class MyTool(Tool):
|
||||
f"context_window_tokens is now {runtime.context_window_tokens!r}"
|
||||
)
|
||||
|
||||
async def _modify_model_preset_async(self, value: Any) -> str:
|
||||
if not isinstance(value, str) or not value.strip():
|
||||
return ToolResult.error("Error: 'model_preset' must be a non-empty string")
|
||||
name = value.strip()
|
||||
session_key = current_request_session_key()
|
||||
old = self._runtime_control.snapshot().model_preset
|
||||
try:
|
||||
runtime = await self._runtime_control.set_model_preset_async(
|
||||
name,
|
||||
session_key=session_key,
|
||||
)
|
||||
except (KeyError, ValueError) as exc:
|
||||
message = str(exc.args[0]) if exc.args else str(exc)
|
||||
punctuation = "" if message.endswith((".", "!", "?")) else "."
|
||||
return ToolResult.error(f"Error: {message}{punctuation}")
|
||||
if session_key:
|
||||
self._audit("modify", f"model_preset = {name!r}")
|
||||
return (
|
||||
f"Set model_preset = {name!r} for the next turn; "
|
||||
f"model will be {runtime.model!r}; "
|
||||
f"context_window_tokens will be {runtime.context_window_tokens!r}"
|
||||
)
|
||||
self._audit("modify", f"model_preset: {old!r} -> {name!r}")
|
||||
return (
|
||||
f"Set model_preset = {name!r} (was {old!r}); model is now {runtime.model!r}; "
|
||||
f"context_window_tokens is now {runtime.context_window_tokens!r}"
|
||||
)
|
||||
|
||||
def _modify_restricted(self, key: str, value: Any) -> str:
|
||||
spec = self.RESTRICTED[key]
|
||||
expected = cast(type[Any], spec["type"])
|
||||
|
||||
@@ -267,6 +267,7 @@ class ExecTool(Tool):
|
||||
|
||||
_MAX_TIMEOUT = 600
|
||||
_MAX_OUTPUT = 10_000
|
||||
_PREPARE_TIMEOUT_SECONDS = 6.0
|
||||
|
||||
# Kernel device files safe as stdio redirect targets (#3599).
|
||||
_BENIGN_DEVICE_PATHS: frozenset[str] = frozenset({
|
||||
@@ -324,7 +325,20 @@ class ExecTool(Tool):
|
||||
if max_output_chars is None:
|
||||
max_output_chars = max_output_tokens
|
||||
|
||||
prepared = self._prepare_command(command, working_dir, timeout, shell, login)
|
||||
try:
|
||||
prepared = await asyncio.wait_for(
|
||||
asyncio.to_thread(
|
||||
self._prepare_command,
|
||||
command,
|
||||
working_dir,
|
||||
timeout,
|
||||
shell,
|
||||
login,
|
||||
),
|
||||
timeout=self._PREPARE_TIMEOUT_SECONDS,
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
return ToolResult.error("Error: command validation timed out")
|
||||
if isinstance(prepared, str):
|
||||
return prepared
|
||||
|
||||
@@ -916,6 +930,15 @@ class ExecTool(Tool):
|
||||
if self._is_benign_device_path(expanded):
|
||||
continue
|
||||
except Exception:
|
||||
# ``Path.expanduser()`` raises when a named user's home
|
||||
# cannot be resolved (notably on Windows). An extracted
|
||||
# home path must fail closed rather than bypass the guard.
|
||||
if raw.strip().startswith("~"):
|
||||
return ToolResult.error(
|
||||
"Error: Command blocked by safety guard "
|
||||
"(path outside working dir)"
|
||||
+ _WORKSPACE_BOUNDARY_NOTE
|
||||
)
|
||||
continue
|
||||
|
||||
if self._is_benign_device_path(str(p)):
|
||||
|
||||
+12
-14
@@ -96,7 +96,7 @@ def _normalize(text: str) -> str:
|
||||
|
||||
|
||||
def _validate_url(url: str) -> tuple[bool, str]:
|
||||
"""Validate URL scheme/domain. Does NOT check resolved IPs (use _validate_url_safe for that)."""
|
||||
"""Validate URL scheme/domain. Does not resolve IPs; use the async safe helper for that."""
|
||||
try:
|
||||
p = urlparse(url)
|
||||
if p.scheme not in ('http', 'https'):
|
||||
@@ -108,18 +108,16 @@ def _validate_url(url: str) -> tuple[bool, str]:
|
||||
return False, str(e)
|
||||
|
||||
|
||||
def _validate_url_safe(url: str) -> tuple[bool, str]:
|
||||
"""Validate URL with SSRF protection: scheme, domain, and resolved IP check."""
|
||||
from nanobot.security.network import validate_url_target
|
||||
async def _async_validate_url_safe(url: str) -> tuple[bool, str]:
|
||||
from nanobot.security.network import async_validate_url_target
|
||||
|
||||
return validate_url_target(url)
|
||||
return await async_validate_url_target(url)
|
||||
|
||||
|
||||
def _resolve_url_safe(url: str) -> tuple[bool, str, tuple[str, ...]]:
|
||||
"""Validate URL and return the resolved IPs to pin during the request."""
|
||||
from nanobot.security.network import resolve_url_target
|
||||
async def _async_resolve_url_safe(url: str) -> tuple[bool, str, tuple[str, ...]]:
|
||||
from nanobot.security.network import async_resolve_url_target
|
||||
|
||||
return resolve_url_target(url)
|
||||
return await async_resolve_url_target(url)
|
||||
|
||||
|
||||
def _pinned_dns_transport() -> httpx.AsyncBaseTransport:
|
||||
@@ -209,7 +207,7 @@ async def _get_with_safe_redirects(
|
||||
"""GET a URL while validating every redirect target before requesting it."""
|
||||
current_url = url
|
||||
for _ in range(MAX_REDIRECTS + 1):
|
||||
is_valid, error_msg, _ = _resolve_url_safe(current_url)
|
||||
is_valid, error_msg, _ = await _async_resolve_url_safe(current_url)
|
||||
if not is_valid:
|
||||
return None, f"Redirect blocked: {error_msg}"
|
||||
|
||||
@@ -229,7 +227,7 @@ async def _get_with_safe_redirects(
|
||||
return response, None
|
||||
|
||||
next_url = urljoin(str(response.url), location)
|
||||
is_valid, error_msg = _validate_url_safe(next_url)
|
||||
is_valid, error_msg = await _async_validate_url_safe(next_url)
|
||||
if not is_valid:
|
||||
await response.aclose()
|
||||
return None, f"Redirect blocked: {error_msg}"
|
||||
@@ -249,7 +247,7 @@ async def _stream_with_safe_redirects(
|
||||
current_url = url
|
||||
chain_carries_credentials = _url_carries_credentials(url)
|
||||
for _ in range(MAX_REDIRECTS + 1):
|
||||
is_valid, error_msg, _ = _resolve_url_safe(current_url)
|
||||
is_valid, error_msg, _ = await _async_resolve_url_safe(current_url)
|
||||
if not is_valid:
|
||||
return None, None, f"Redirect blocked: {error_msg}", chain_carries_credentials
|
||||
|
||||
@@ -283,7 +281,7 @@ async def _stream_with_safe_redirects(
|
||||
chain_carries_credentials = (
|
||||
chain_carries_credentials or _url_carries_credentials(next_url)
|
||||
)
|
||||
is_valid, error_msg = _validate_url_safe(next_url)
|
||||
is_valid, error_msg = await _async_validate_url_safe(next_url)
|
||||
if not is_valid:
|
||||
await stream.__aexit__(None, None, None)
|
||||
return None, None, f"Redirect blocked: {error_msg}", chain_carries_credentials
|
||||
@@ -1106,7 +1104,7 @@ class WebFetchTool(Tool):
|
||||
url = url.strip(" \t\r\n`\"'")
|
||||
extract_mode = kwargs.pop("extractMode", extract_mode)
|
||||
max_chars = cast(int, kwargs.pop("maxChars", max_chars) or self.max_chars)
|
||||
is_valid, error_msg = _validate_url_safe(url)
|
||||
is_valid, error_msg = await _async_validate_url_safe(url)
|
||||
if not is_valid:
|
||||
return json.dumps({"error": f"URL validation failed: {error_msg}", "url": url}, ensure_ascii=False)
|
||||
|
||||
|
||||
+374
-41
@@ -2,15 +2,20 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import ctypes
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import shlex
|
||||
import shutil
|
||||
import signal
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from collections.abc import Iterable
|
||||
from contextlib import suppress
|
||||
from ctypes import wintypes
|
||||
from dataclasses import dataclass
|
||||
from importlib import metadata as importlib_metadata
|
||||
from pathlib import Path
|
||||
@@ -97,6 +102,141 @@ class CliAppsRuntimeConfig:
|
||||
catalog_ttl_seconds: int = 3600
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class _PreparedCliRun:
|
||||
name: str
|
||||
entry: str
|
||||
resolved: str
|
||||
args: list[str]
|
||||
cwd: Path
|
||||
timeout: int
|
||||
env: dict[str, str]
|
||||
artifact_snapshot: dict[Path, tuple[int, int]]
|
||||
|
||||
|
||||
class _JobObjectBasicLimitInformation(ctypes.Structure):
|
||||
_fields_ = [
|
||||
("PerProcessUserTimeLimit", ctypes.c_int64),
|
||||
("PerJobUserTimeLimit", ctypes.c_int64),
|
||||
("LimitFlags", wintypes.DWORD),
|
||||
("MinimumWorkingSetSize", ctypes.c_size_t),
|
||||
("MaximumWorkingSetSize", ctypes.c_size_t),
|
||||
("ActiveProcessLimit", wintypes.DWORD),
|
||||
("Affinity", ctypes.c_size_t),
|
||||
("PriorityClass", wintypes.DWORD),
|
||||
("SchedulingClass", wintypes.DWORD),
|
||||
]
|
||||
|
||||
|
||||
class _IoCounters(ctypes.Structure):
|
||||
_fields_ = [
|
||||
("ReadOperationCount", ctypes.c_uint64),
|
||||
("WriteOperationCount", ctypes.c_uint64),
|
||||
("OtherOperationCount", ctypes.c_uint64),
|
||||
("ReadTransferCount", ctypes.c_uint64),
|
||||
("WriteTransferCount", ctypes.c_uint64),
|
||||
("OtherTransferCount", ctypes.c_uint64),
|
||||
]
|
||||
|
||||
|
||||
class _JobObjectExtendedLimitInformation(ctypes.Structure):
|
||||
_fields_ = [
|
||||
("BasicLimitInformation", _JobObjectBasicLimitInformation),
|
||||
("IoInfo", _IoCounters),
|
||||
("ProcessMemoryLimit", ctypes.c_size_t),
|
||||
("JobMemoryLimit", ctypes.c_size_t),
|
||||
("PeakProcessMemoryUsed", ctypes.c_size_t),
|
||||
("PeakJobMemoryUsed", ctypes.c_size_t),
|
||||
]
|
||||
|
||||
|
||||
class _WindowsJob:
|
||||
"""Best-effort Windows process tree ownership for timeout/cancellation."""
|
||||
|
||||
_KILL_ON_JOB_CLOSE = 0x00002000
|
||||
_EXTENDED_LIMIT_INFORMATION = 9
|
||||
_PROCESS_TERMINATE = 0x0001
|
||||
_PROCESS_SET_QUOTA = 0x0100
|
||||
|
||||
def __init__(self) -> None:
|
||||
win_dll = getattr(ctypes, "WinDLL")
|
||||
self._kernel32 = win_dll("kernel32", use_last_error=True)
|
||||
self._kernel32.CreateJobObjectW.argtypes = [wintypes.LPVOID, wintypes.LPCWSTR]
|
||||
self._kernel32.CreateJobObjectW.restype = wintypes.HANDLE
|
||||
self._kernel32.OpenProcess.argtypes = [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD]
|
||||
self._kernel32.OpenProcess.restype = wintypes.HANDLE
|
||||
self._kernel32.AssignProcessToJobObject.argtypes = [wintypes.HANDLE, wintypes.HANDLE]
|
||||
self._kernel32.AssignProcessToJobObject.restype = wintypes.BOOL
|
||||
self._kernel32.SetInformationJobObject.argtypes = [
|
||||
wintypes.HANDLE,
|
||||
ctypes.c_int,
|
||||
wintypes.LPVOID,
|
||||
wintypes.DWORD,
|
||||
]
|
||||
self._kernel32.SetInformationJobObject.restype = wintypes.BOOL
|
||||
self._kernel32.TerminateJobObject.argtypes = [wintypes.HANDLE, wintypes.UINT]
|
||||
self._kernel32.TerminateJobObject.restype = wintypes.BOOL
|
||||
self._kernel32.CloseHandle.argtypes = [wintypes.HANDLE]
|
||||
self._kernel32.CloseHandle.restype = wintypes.BOOL
|
||||
self._handle: Any = self._kernel32.CreateJobObjectW(None, None)
|
||||
if not self._handle:
|
||||
raise OSError(ctypes.get_last_error(), "CreateJobObjectW failed")
|
||||
try:
|
||||
self._set_kill_on_close(True)
|
||||
except OSError:
|
||||
self._kernel32.CloseHandle(self._handle)
|
||||
self._handle = None
|
||||
raise
|
||||
|
||||
@classmethod
|
||||
def create(cls) -> _WindowsJob | None:
|
||||
if os.name != "nt":
|
||||
return None
|
||||
try:
|
||||
return cls()
|
||||
except OSError as exc:
|
||||
logger.debug("CLI Apps: Windows job object unavailable: {}", exc)
|
||||
return None
|
||||
|
||||
def _set_kill_on_close(self, enabled: bool) -> None:
|
||||
info = _JobObjectExtendedLimitInformation()
|
||||
info.BasicLimitInformation.LimitFlags = self._KILL_ON_JOB_CLOSE if enabled else 0
|
||||
ok = self._kernel32.SetInformationJobObject(
|
||||
self._handle,
|
||||
self._EXTENDED_LIMIT_INFORMATION,
|
||||
ctypes.byref(info),
|
||||
ctypes.sizeof(info),
|
||||
)
|
||||
if not ok:
|
||||
raise OSError(ctypes.get_last_error(), "SetInformationJobObject failed")
|
||||
|
||||
def assign(self, pid: int) -> bool:
|
||||
process_handle = self._kernel32.OpenProcess(
|
||||
self._PROCESS_TERMINATE | self._PROCESS_SET_QUOTA,
|
||||
False,
|
||||
pid,
|
||||
)
|
||||
if not process_handle:
|
||||
return False
|
||||
try:
|
||||
return bool(self._kernel32.AssignProcessToJobObject(self._handle, process_handle))
|
||||
finally:
|
||||
self._kernel32.CloseHandle(process_handle)
|
||||
|
||||
def terminate(self) -> None:
|
||||
if self._handle and not self._kernel32.TerminateJobObject(self._handle, 1):
|
||||
raise OSError(ctypes.get_last_error(), "TerminateJobObject failed")
|
||||
|
||||
def close(self, *, kill_descendants: bool) -> None:
|
||||
if not self._handle:
|
||||
return
|
||||
if not kill_descendants:
|
||||
with suppress(OSError):
|
||||
self._set_kill_on_close(False)
|
||||
self._kernel32.CloseHandle(self._handle)
|
||||
self._handle = None
|
||||
|
||||
|
||||
_BRANDS: dict[str, tuple[str, str]] = {
|
||||
"1password-cli": ("1password", "#3B66BC"),
|
||||
"arcgis": ("arcgis", "#2C7AC3"),
|
||||
@@ -1428,6 +1568,197 @@ Use the `run_cli_app` tool with `name="{name}"` for command execution. Do not in
|
||||
lines.append(f"- {rel} ({kind}, {self._format_artifact_size(path)})")
|
||||
return lines
|
||||
|
||||
def _prepare_run(
|
||||
self,
|
||||
name: str,
|
||||
args: list[str] | None,
|
||||
*,
|
||||
json_output: bool,
|
||||
working_dir: str | None,
|
||||
timeout: int | None,
|
||||
restrict_to_workspace: bool,
|
||||
) -> _PreparedCliRun:
|
||||
app = self.get_app(name)
|
||||
installed = self._load_installed()
|
||||
app_name = str(app["name"])
|
||||
if app_name not in installed:
|
||||
raise CliAppError(f"CLI app '{name}' is not installed")
|
||||
cwd = self._resolve_cwd(working_dir, restrict_to_workspace=restrict_to_workspace)
|
||||
entry = str(installed[app_name].get("entry_point") or app.get("entry_point") or "")
|
||||
resolved = shutil.which(entry)
|
||||
if not entry or not resolved:
|
||||
raise CliAppError(f"{entry or name} is not available on PATH")
|
||||
clean_args = [str(arg) for arg in (args or [])]
|
||||
if json_output and "--json" not in clean_args:
|
||||
clean_args = ["--json", *clean_args]
|
||||
effective_timeout = max(1, min(timeout or self.runtime.run_timeout, 600))
|
||||
return _PreparedCliRun(
|
||||
name=name,
|
||||
entry=entry,
|
||||
resolved=resolved,
|
||||
args=clean_args,
|
||||
cwd=cwd,
|
||||
timeout=effective_timeout,
|
||||
env=self._subprocess_env(),
|
||||
artifact_snapshot=self._artifact_snapshot(cwd),
|
||||
)
|
||||
|
||||
def _format_run_result(
|
||||
self,
|
||||
prepared: _PreparedCliRun,
|
||||
*,
|
||||
returncode: int,
|
||||
stdout: str,
|
||||
stderr: str,
|
||||
) -> str:
|
||||
command = " ".join([prepared.entry, *(shlex.quote(arg) for arg in prepared.args)])
|
||||
output = [
|
||||
f"CLI app '{prepared.name}' exited {returncode}.",
|
||||
f"Command: {command}",
|
||||
]
|
||||
if stdout:
|
||||
output.append("\nSTDOUT:\n" + stdout.rstrip())
|
||||
if stderr:
|
||||
output.append("\nSTDERR:\n" + stderr.rstrip())
|
||||
artifacts = self._changed_artifacts(prepared.cwd, prepared.artifact_snapshot)
|
||||
if artifacts:
|
||||
output.append(
|
||||
"\nArtifacts created or updated:\n"
|
||||
+ "\n".join(self._format_artifact_lines(prepared.cwd, artifacts))
|
||||
)
|
||||
if any(path.suffix.lower() in _INLINE_ARTIFACT_EXTENSIONS for path in artifacts):
|
||||
output.append(
|
||||
"\nTo show a preview in WebUI, reference a raster artifact with Markdown "
|
||||
"using its workspace-relative path, for example ``."
|
||||
)
|
||||
return _truncate("\n".join(output))
|
||||
|
||||
@staticmethod
|
||||
def _terminate_run_process_sync(
|
||||
process: subprocess.Popen[str],
|
||||
job: _WindowsJob | None,
|
||||
) -> None:
|
||||
if job is not None:
|
||||
with suppress(OSError):
|
||||
job.terminate()
|
||||
job.close(kill_descendants=True)
|
||||
elif os.name == "nt":
|
||||
with suppress(OSError, subprocess.TimeoutExpired):
|
||||
subprocess.run(
|
||||
["taskkill", "/PID", str(process.pid), "/T", "/F"],
|
||||
check=False,
|
||||
stdout=subprocess.DEVNULL,
|
||||
stderr=subprocess.DEVNULL,
|
||||
timeout=5,
|
||||
)
|
||||
else:
|
||||
with suppress(ProcessLookupError, PermissionError):
|
||||
os.killpg(process.pid, signal.SIGKILL)
|
||||
if process.poll() is None:
|
||||
with suppress(ProcessLookupError):
|
||||
process.kill()
|
||||
with suppress(subprocess.TimeoutExpired):
|
||||
process.wait(timeout=5)
|
||||
|
||||
@staticmethod
|
||||
async def _terminate_run_process(
|
||||
process: asyncio.subprocess.Process,
|
||||
job: _WindowsJob | None,
|
||||
) -> None:
|
||||
if job is not None:
|
||||
with suppress(OSError):
|
||||
await asyncio.to_thread(job.terminate)
|
||||
job.close(kill_descendants=True)
|
||||
elif os.name == "nt":
|
||||
with suppress(OSError, asyncio.TimeoutError):
|
||||
await asyncio.wait_for(
|
||||
asyncio.to_thread(
|
||||
subprocess.run,
|
||||
["taskkill", "/PID", str(process.pid), "/T", "/F"],
|
||||
check=False,
|
||||
stdout=subprocess.DEVNULL,
|
||||
stderr=subprocess.DEVNULL,
|
||||
timeout=5,
|
||||
),
|
||||
timeout=6.0,
|
||||
)
|
||||
else:
|
||||
with suppress(ProcessLookupError, PermissionError):
|
||||
os.killpg(process.pid, signal.SIGKILL)
|
||||
if process.returncode is None:
|
||||
with suppress(ProcessLookupError):
|
||||
process.kill()
|
||||
with suppress(asyncio.TimeoutError, ProcessLookupError):
|
||||
await asyncio.wait_for(process.wait(), timeout=5.0)
|
||||
|
||||
async def run_async(
|
||||
self,
|
||||
name: str,
|
||||
args: list[str] | None = None,
|
||||
*,
|
||||
json_output: bool = False,
|
||||
working_dir: str | None = None,
|
||||
timeout: int | None = None,
|
||||
restrict_to_workspace: bool = False,
|
||||
) -> str:
|
||||
prepared = await asyncio.to_thread(
|
||||
self._prepare_run,
|
||||
name,
|
||||
args,
|
||||
json_output=json_output,
|
||||
working_dir=working_dir,
|
||||
timeout=timeout,
|
||||
restrict_to_workspace=restrict_to_workspace,
|
||||
)
|
||||
process_kwargs: dict[str, Any] = {}
|
||||
if os.name == "nt":
|
||||
process_kwargs["creationflags"] = subprocess.CREATE_NEW_PROCESS_GROUP
|
||||
else:
|
||||
process_kwargs["start_new_session"] = True
|
||||
job = _WindowsJob.create()
|
||||
try:
|
||||
process = await asyncio.create_subprocess_exec(
|
||||
prepared.resolved,
|
||||
*prepared.args,
|
||||
cwd=str(prepared.cwd),
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
env=prepared.env,
|
||||
**process_kwargs,
|
||||
)
|
||||
except BaseException:
|
||||
if job is not None:
|
||||
job.close(kill_descendants=False)
|
||||
raise
|
||||
if job is not None and not job.assign(process.pid):
|
||||
job.close(kill_descendants=False)
|
||||
job = None
|
||||
try:
|
||||
stdout_raw, stderr_raw = await asyncio.wait_for(
|
||||
process.communicate(),
|
||||
timeout=prepared.timeout,
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
await self._terminate_run_process(process, job)
|
||||
return f"CLI app '{prepared.name}' timed out after {prepared.timeout}s"
|
||||
except asyncio.CancelledError:
|
||||
await self._terminate_run_process(process, job)
|
||||
raise
|
||||
except BaseException:
|
||||
await self._terminate_run_process(process, job)
|
||||
raise
|
||||
if job is not None:
|
||||
job.close(kill_descendants=False)
|
||||
stdout = stdout_raw.decode("utf-8", errors="replace")
|
||||
stderr = stderr_raw.decode("utf-8", errors="replace")
|
||||
return await asyncio.to_thread(
|
||||
self._format_run_result,
|
||||
prepared,
|
||||
returncode=process.returncode or 0,
|
||||
stdout=stdout,
|
||||
stderr=stderr,
|
||||
)
|
||||
|
||||
def run(
|
||||
self,
|
||||
name: str,
|
||||
@@ -1438,50 +1769,52 @@ Use the `run_cli_app` tool with `name="{name}"` for command execution. Do not in
|
||||
timeout: int | None = None,
|
||||
restrict_to_workspace: bool = False,
|
||||
) -> str:
|
||||
app = self.get_app(name)
|
||||
installed = self._load_installed()
|
||||
if str(app["name"]) not in installed:
|
||||
raise CliAppError(f"CLI app '{name}' is not installed")
|
||||
cwd = self._resolve_cwd(working_dir, restrict_to_workspace=restrict_to_workspace)
|
||||
entry = str(installed[str(app["name"])].get("entry_point") or app.get("entry_point") or "")
|
||||
resolved = shutil.which(entry)
|
||||
if not entry or not resolved:
|
||||
raise CliAppError(f"{entry or name} is not available on PATH")
|
||||
clean_args = [str(arg) for arg in (args or [])]
|
||||
if json_output and "--json" not in clean_args:
|
||||
clean_args = ["--json", *clean_args]
|
||||
effective_timeout = max(1, min(timeout or self.runtime.run_timeout, 600))
|
||||
artifact_snapshot = self._artifact_snapshot(cwd)
|
||||
prepared = self._prepare_run(
|
||||
name,
|
||||
args,
|
||||
json_output=json_output,
|
||||
working_dir=working_dir,
|
||||
timeout=timeout,
|
||||
restrict_to_workspace=restrict_to_workspace,
|
||||
)
|
||||
process_kwargs: dict[str, Any] = {}
|
||||
if os.name == "nt":
|
||||
process_kwargs["creationflags"] = subprocess.CREATE_NEW_PROCESS_GROUP
|
||||
else:
|
||||
process_kwargs["start_new_session"] = True
|
||||
job = _WindowsJob.create()
|
||||
try:
|
||||
result = subprocess.run(
|
||||
[resolved, *clean_args],
|
||||
cwd=str(cwd),
|
||||
capture_output=True,
|
||||
process = subprocess.Popen(
|
||||
[prepared.resolved, *prepared.args],
|
||||
cwd=str(prepared.cwd),
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
text=True,
|
||||
encoding="utf-8",
|
||||
errors="replace",
|
||||
timeout=effective_timeout,
|
||||
env=self._subprocess_env(),
|
||||
env=prepared.env,
|
||||
**process_kwargs,
|
||||
)
|
||||
except BaseException:
|
||||
if job is not None:
|
||||
job.close(kill_descendants=False)
|
||||
raise
|
||||
if job is not None and not job.assign(process.pid):
|
||||
job.close(kill_descendants=False)
|
||||
job = None
|
||||
try:
|
||||
stdout, stderr = process.communicate(timeout=prepared.timeout)
|
||||
except subprocess.TimeoutExpired:
|
||||
return f"CLI app '{name}' timed out after {effective_timeout}s"
|
||||
output = [
|
||||
f"CLI app '{name}' exited {result.returncode}.",
|
||||
f"Command: {entry} {' '.join(shlex.quote(arg) for arg in clean_args)}".rstrip(),
|
||||
]
|
||||
if result.stdout:
|
||||
output.append("\nSTDOUT:\n" + result.stdout.rstrip())
|
||||
if result.stderr:
|
||||
output.append("\nSTDERR:\n" + result.stderr.rstrip())
|
||||
artifacts = self._changed_artifacts(cwd, artifact_snapshot)
|
||||
if artifacts:
|
||||
output.append(
|
||||
"\nArtifacts created or updated:\n"
|
||||
+ "\n".join(self._format_artifact_lines(cwd, artifacts))
|
||||
)
|
||||
if any(path.suffix.lower() in _INLINE_ARTIFACT_EXTENSIONS for path in artifacts):
|
||||
output.append(
|
||||
"\nTo show a preview in WebUI, reference a raster artifact with Markdown "
|
||||
"using its workspace-relative path, for example ``."
|
||||
)
|
||||
return _truncate("\n".join(output))
|
||||
self._terminate_run_process_sync(process, job)
|
||||
return f"CLI app '{prepared.name}' timed out after {prepared.timeout}s"
|
||||
except BaseException:
|
||||
self._terminate_run_process_sync(process, job)
|
||||
raise
|
||||
if job is not None:
|
||||
job.close(kill_descendants=False)
|
||||
return self._format_run_result(
|
||||
prepared,
|
||||
returncode=process.returncode,
|
||||
stdout=stdout,
|
||||
stderr=stderr,
|
||||
)
|
||||
|
||||
@@ -21,7 +21,10 @@ from nanobot.bus.events import OutboundMessage
|
||||
from nanobot.bus.queue import MessageBus
|
||||
from nanobot.channels.base import BaseChannel
|
||||
from nanobot.config.schema import Base
|
||||
from nanobot.security.network import validate_resolved_url, validate_url_target
|
||||
from nanobot.security.network import (
|
||||
async_validate_resolved_url,
|
||||
async_validate_url_target,
|
||||
)
|
||||
|
||||
DINGTALK_MAX_REMOTE_MEDIA_BYTES = 20 * 1024 * 1024
|
||||
DINGTALK_MAX_REMOTE_MEDIA_REDIRECTS = 3
|
||||
@@ -417,8 +420,8 @@ class DingTalkChannel(BaseChannel):
|
||||
return self._zip_bytes(filename, data)
|
||||
return data, filename, content_type
|
||||
|
||||
def _validate_remote_media_url(self, media_ref: str) -> bool:
|
||||
ok, err = validate_url_target(media_ref)
|
||||
async def _validate_remote_media_url(self, media_ref: str) -> bool:
|
||||
ok, err = await async_validate_url_target(media_ref)
|
||||
if not ok:
|
||||
self.logger.warning("remote media URL blocked ref={} reason={}", media_ref, err)
|
||||
return False
|
||||
@@ -434,7 +437,11 @@ class DingTalkChannel(BaseChannel):
|
||||
allowed_hosts = {host.lower() for host in self.config.remote_media_redirect_allowed_hosts}
|
||||
return next_host in allowed_hosts
|
||||
|
||||
def _next_remote_media_url(self, current_url: str, location: str | None) -> str | None:
|
||||
async def _next_remote_media_url(
|
||||
self,
|
||||
current_url: str,
|
||||
location: str | None,
|
||||
) -> str | None:
|
||||
if not self.config.allow_remote_media_redirects:
|
||||
self.logger.warning("media download redirect refused ref={}", current_url)
|
||||
return None
|
||||
@@ -449,7 +456,7 @@ class DingTalkChannel(BaseChannel):
|
||||
next_url,
|
||||
)
|
||||
return None
|
||||
if not self._validate_remote_media_url(next_url):
|
||||
if not await self._validate_remote_media_url(next_url):
|
||||
return None
|
||||
return next_url
|
||||
|
||||
@@ -461,7 +468,7 @@ class DingTalkChannel(BaseChannel):
|
||||
if not self._http:
|
||||
return None, None
|
||||
|
||||
if not self._validate_remote_media_url(media_ref):
|
||||
if not await self._validate_remote_media_url(media_ref):
|
||||
return None, None
|
||||
|
||||
try:
|
||||
@@ -473,7 +480,7 @@ class DingTalkChannel(BaseChannel):
|
||||
current_url = media_ref
|
||||
for _ in range(DINGTALK_MAX_REMOTE_MEDIA_REDIRECTS + 1):
|
||||
async with stream("GET", current_url, follow_redirects=False) as resp:
|
||||
final_ok, final_err = validate_resolved_url(str(resp.url))
|
||||
final_ok, final_err = await async_validate_resolved_url(str(resp.url))
|
||||
if not final_ok:
|
||||
self.logger.warning(
|
||||
"remote media redirect blocked ref={} final={} reason={}",
|
||||
@@ -483,7 +490,7 @@ class DingTalkChannel(BaseChannel):
|
||||
)
|
||||
return None, None
|
||||
if 300 <= resp.status_code < 400:
|
||||
next_url = self._next_remote_media_url(
|
||||
next_url = await self._next_remote_media_url(
|
||||
str(resp.url), resp.headers.get("location")
|
||||
)
|
||||
if not next_url:
|
||||
@@ -516,7 +523,9 @@ class DingTalkChannel(BaseChannel):
|
||||
current_url = media_ref
|
||||
for _ in range(DINGTALK_MAX_REMOTE_MEDIA_REDIRECTS + 1):
|
||||
resp = await self._http.get(current_url, follow_redirects=False)
|
||||
final_ok, final_err = validate_resolved_url(str(getattr(resp, "url", current_url)))
|
||||
final_ok, final_err = await async_validate_resolved_url(
|
||||
str(getattr(resp, "url", current_url))
|
||||
)
|
||||
if not final_ok:
|
||||
self.logger.warning(
|
||||
"remote media redirect blocked ref={} final={} reason={}",
|
||||
@@ -526,7 +535,7 @@ class DingTalkChannel(BaseChannel):
|
||||
)
|
||||
return None, None
|
||||
if 300 <= resp.status_code < 400:
|
||||
next_url = self._next_remote_media_url(
|
||||
next_url = await self._next_remote_media_url(
|
||||
str(getattr(resp, "url", current_url)), resp.headers.get("location")
|
||||
)
|
||||
if not next_url:
|
||||
|
||||
@@ -24,7 +24,7 @@ from nanobot.bus.queue import MessageBus
|
||||
from nanobot.channels.base import BaseChannel
|
||||
from nanobot.config.paths import get_media_dir
|
||||
from nanobot.config.schema import Base
|
||||
from nanobot.security.network import validate_url_target
|
||||
from nanobot.security.network import async_validate_url_target
|
||||
from nanobot.utils.helpers import safe_filename
|
||||
|
||||
_DOWNLOAD_TIMEOUT = aiohttp.ClientTimeout(total=60)
|
||||
@@ -473,7 +473,7 @@ class NapcatChannel(BaseChannel):
|
||||
if not ref:
|
||||
return None
|
||||
if ref.startswith(("http://", "https://")):
|
||||
ok, err = validate_url_target(ref)
|
||||
ok, err = await async_validate_url_target(ref)
|
||||
if not ok:
|
||||
logger.warning("napcat: rejected remote image '{}': {}", ref, err)
|
||||
return None
|
||||
@@ -525,7 +525,7 @@ class NapcatChannel(BaseChannel):
|
||||
# logger.debug("napcat: downloading image from {}", url)
|
||||
if self._http is None:
|
||||
return None
|
||||
ok, err = validate_url_target(url)
|
||||
ok, err = await async_validate_url_target(url)
|
||||
if not ok:
|
||||
logger.warning("napcat: skip image '{}': {}", url, err)
|
||||
return None
|
||||
|
||||
@@ -149,9 +149,13 @@ async def test_download_image_rejects_redirects(tmp_path, monkeypatch) -> None:
|
||||
channel = _channel()
|
||||
channel._media_root = tmp_path
|
||||
channel._http = _FakeHttp(_FakeResponse(status=302))
|
||||
|
||||
async def allow_url(_url: str) -> tuple[bool, str]:
|
||||
return True, ""
|
||||
|
||||
monkeypatch.setattr(
|
||||
"nanobot.channels.napcat.runtime.validate_url_target",
|
||||
lambda _url: (True, ""),
|
||||
"nanobot.channels.napcat.runtime.async_validate_url_target",
|
||||
allow_url,
|
||||
)
|
||||
|
||||
result = await channel._download_image({"url": "https://example.com/a.png", "file": "a.png"})
|
||||
|
||||
@@ -40,7 +40,7 @@ from nanobot.bus.events import OutboundMessage
|
||||
from nanobot.bus.queue import MessageBus
|
||||
from nanobot.channels.base import BaseChannel
|
||||
from nanobot.config.schema import Base
|
||||
from nanobot.security.network import validate_url_target
|
||||
from nanobot.security.network import async_validate_url_target
|
||||
from nanobot.utils.logging_bridge import redirect_lib_logging
|
||||
|
||||
try:
|
||||
@@ -458,7 +458,7 @@ class QQChannel(BaseChannel):
|
||||
return None, None
|
||||
|
||||
# Remote URL
|
||||
ok, err = validate_url_target(media_ref)
|
||||
ok, err = await async_validate_url_target(media_ref)
|
||||
if not ok:
|
||||
self.logger.warning("outbound media URL validation failed url={} err={}", media_ref, err)
|
||||
return None, None
|
||||
|
||||
@@ -23,8 +23,8 @@ from nanobot.config.schema import Base
|
||||
from nanobot.pairing import is_approved
|
||||
from nanobot.security.network import (
|
||||
PinnedDNSAsyncTransport,
|
||||
async_validate_url_target,
|
||||
httpx_env_proxy_mounts,
|
||||
validate_url_target,
|
||||
)
|
||||
from nanobot.utils.helpers import safe_filename, split_message
|
||||
|
||||
@@ -95,7 +95,7 @@ _HTML_DOWNLOAD_PREFIXES = (b"<!doctype html", b"<html")
|
||||
|
||||
async def _validate_slack_download_request(request: httpx.Request) -> None:
|
||||
"""Validate every Slack file request, including redirects, before transport."""
|
||||
ok, error = validate_url_target(str(request.url))
|
||||
ok, error = await async_validate_url_target(str(request.url))
|
||||
if not ok:
|
||||
raise httpx.RequestError(f"unsafe Slack file URL: {error}", request=request)
|
||||
|
||||
|
||||
@@ -859,13 +859,13 @@ def _patch_download_validation(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
validated: list[str],
|
||||
) -> None:
|
||||
def validate(url: str) -> tuple[bool, str]:
|
||||
async def validate(url: str) -> tuple[bool, str]:
|
||||
validated.append(url)
|
||||
if "169.254.169.254" in url:
|
||||
return False, "blocked metadata address"
|
||||
return True, ""
|
||||
|
||||
monkeypatch.setattr("nanobot.channels.slack.runtime.validate_url_target", validate)
|
||||
monkeypatch.setattr("nanobot.channels.slack.runtime.async_validate_url_target", validate)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
||||
@@ -36,7 +36,7 @@ from nanobot.channels.base import BaseChannel
|
||||
from nanobot.command.builtin import build_help_text
|
||||
from nanobot.config.paths import get_media_dir
|
||||
from nanobot.config.schema import Base
|
||||
from nanobot.security.network import validate_url_target
|
||||
from nanobot.security.network import async_validate_url_target
|
||||
from nanobot.utils.helpers import split_message
|
||||
from nanobot.utils.logging_bridge import redirect_lib_logging
|
||||
|
||||
@@ -956,7 +956,7 @@ class TelegramChannel(BaseChannel):
|
||||
|
||||
# Telegram Bot API accepts HTTP(S) URLs directly for media params.
|
||||
if self._is_remote_media_url(media_path):
|
||||
ok, error = validate_url_target(media_path)
|
||||
ok, error = await async_validate_url_target(media_path)
|
||||
if not ok:
|
||||
raise ValueError(f"unsafe media URL: {error}")
|
||||
await self._call_with_retry(
|
||||
|
||||
@@ -1488,7 +1488,14 @@ async def test_send_remote_media_url_after_security_validation(monkeypatch) -> N
|
||||
MessageBus(),
|
||||
)
|
||||
_install_ready_app(channel)
|
||||
monkeypatch.setattr("nanobot.channels.telegram.runtime.validate_url_target", lambda url: (True, ""))
|
||||
|
||||
async def allow_url(_url: str) -> tuple[bool, str]:
|
||||
return True, ""
|
||||
|
||||
monkeypatch.setattr(
|
||||
"nanobot.channels.telegram.runtime.async_validate_url_target",
|
||||
allow_url,
|
||||
)
|
||||
|
||||
await channel.send(
|
||||
OutboundMessage(
|
||||
@@ -1546,9 +1553,13 @@ async def test_send_blocks_unsafe_remote_media_url(monkeypatch) -> None:
|
||||
MessageBus(),
|
||||
)
|
||||
_install_ready_app(channel)
|
||||
|
||||
async def deny_url(_url: str) -> tuple[bool, str]:
|
||||
return False, "Blocked: example.com resolves to private/internal address 127.0.0.1"
|
||||
|
||||
monkeypatch.setattr(
|
||||
"nanobot.channels.telegram.runtime.validate_url_target",
|
||||
lambda url: (False, "Blocked: example.com resolves to private/internal address 127.0.0.1"),
|
||||
"nanobot.channels.telegram.runtime.async_validate_url_target",
|
||||
deny_url,
|
||||
)
|
||||
|
||||
await channel.send(
|
||||
|
||||
@@ -55,6 +55,7 @@ from nanobot.security.workspace_access import (
|
||||
WORKSPACE_SCOPE_METADATA_KEY,
|
||||
WorkspaceScopeError,
|
||||
)
|
||||
from nanobot.session.async_compat import call_session_manager
|
||||
from nanobot.session.goal_state import goal_state_ws_blob
|
||||
from nanobot.session.model_selection import model_preset_from_metadata
|
||||
from nanobot.session.recovery import recovery_state_from_metadata
|
||||
@@ -447,6 +448,26 @@ class WebSocketChannel(BaseChannel):
|
||||
if sessions is None:
|
||||
return {}
|
||||
snapshot = sessions.read_session_metadata(f"websocket:{chat_id}")
|
||||
return self._attached_model_fields_from_snapshot(chat_id, snapshot)
|
||||
|
||||
async def _attached_model_fields_async(self, chat_id: str) -> dict[str, Any]:
|
||||
"""Build attach fields without blocking the gateway event loop."""
|
||||
sessions = self.gateway.session_manager
|
||||
if sessions is None:
|
||||
return {}
|
||||
snapshot = await call_session_manager(
|
||||
sessions,
|
||||
"read_session_metadata_async",
|
||||
sessions.read_session_metadata,
|
||||
f"websocket:{chat_id}",
|
||||
)
|
||||
return self._attached_model_fields_from_snapshot(chat_id, snapshot)
|
||||
|
||||
def _attached_model_fields_from_snapshot(
|
||||
self,
|
||||
chat_id: str,
|
||||
snapshot: dict[str, Any] | None,
|
||||
) -> dict[str, Any]:
|
||||
raw_metadata = snapshot.get("metadata") if snapshot is not None else None
|
||||
metadata = cast(dict[str, object], raw_metadata) if isinstance(raw_metadata, dict) else None
|
||||
fields: dict[str, Any] = {}
|
||||
@@ -510,13 +531,16 @@ class WebSocketChannel(BaseChannel):
|
||||
fork_key: str,
|
||||
) -> None:
|
||||
"""Attach and hydrate a newly created WebUI chat fork."""
|
||||
scope = self._workspaces.scope_for_session_key(fork_key)
|
||||
scope = await asyncio.to_thread(
|
||||
self._workspaces.scope_for_session_key,
|
||||
fork_key,
|
||||
)
|
||||
self._attach(connection, fork_id)
|
||||
await self._send_event(
|
||||
connection,
|
||||
"attached",
|
||||
chat_id=fork_id,
|
||||
**self._attached_model_fields(fork_id),
|
||||
**await self._attached_model_fields_async(fork_id),
|
||||
)
|
||||
await self._send_event(
|
||||
connection,
|
||||
@@ -904,7 +928,7 @@ class WebSocketChannel(BaseChannel):
|
||||
connection,
|
||||
"attached",
|
||||
chat_id=new_id,
|
||||
**self._attached_model_fields(new_id),
|
||||
**await self._attached_model_fields_async(new_id),
|
||||
)
|
||||
await self._send_event(
|
||||
connection,
|
||||
@@ -960,7 +984,7 @@ class WebSocketChannel(BaseChannel):
|
||||
connection,
|
||||
"attached",
|
||||
chat_id=cid,
|
||||
**self._attached_model_fields(cid),
|
||||
**await self._attached_model_fields_async(cid),
|
||||
)
|
||||
await self._hydrate_after_subscribe(cid)
|
||||
return
|
||||
@@ -1263,7 +1287,7 @@ class WebSocketChannel(BaseChannel):
|
||||
else False
|
||||
),
|
||||
)
|
||||
self._workspaces.persist_scope(cid, scope)
|
||||
await asyncio.to_thread(self._workspaces.persist_scope, cid, scope)
|
||||
accepted = True
|
||||
finally:
|
||||
if not accepted and queued_owner is not None:
|
||||
@@ -1553,7 +1577,7 @@ class WebSocketChannel(BaseChannel):
|
||||
turn_id: str | None = None,
|
||||
) -> Any | None:
|
||||
try:
|
||||
return resolver()
|
||||
return await asyncio.to_thread(resolver)
|
||||
except WorkspaceScopeError as exc:
|
||||
await self._send_event(
|
||||
connection,
|
||||
|
||||
@@ -4,6 +4,7 @@ import asyncio
|
||||
import json
|
||||
import random
|
||||
import socket
|
||||
import threading
|
||||
import time
|
||||
from contextlib import suppress
|
||||
from pathlib import Path
|
||||
@@ -2576,6 +2577,69 @@ async def test_webui_automations_route_lists_all_jobs_and_allows_user_actions(
|
||||
await server_task
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_webui_cron_update_rearms_started_service_on_owner_loop(
|
||||
bus: MagicMock,
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
store_path = tmp_path / "cron" / "jobs.json"
|
||||
cron = CronService(store_path, max_sleep_ms=60_000)
|
||||
job = cron.add_job(
|
||||
name="Before update",
|
||||
schedule=CronSchedule(kind="every", every_ms=86_400_000),
|
||||
message="Check the repo status",
|
||||
session_key="websocket:abc",
|
||||
origin_channel="websocket",
|
||||
origin_chat_id="abc",
|
||||
)
|
||||
await cron.start()
|
||||
|
||||
owner_thread_id = threading.get_ident()
|
||||
initial_timer = cron._timer_task
|
||||
request_thread_ids: list[int] = []
|
||||
arm_thread_ids: list[int] = []
|
||||
timer_rearmed = asyncio.Event()
|
||||
original_request_timer_rearm = cron._request_timer_rearm
|
||||
original_arm_timer = cron._arm_timer
|
||||
|
||||
def tracked_request_timer_rearm() -> None:
|
||||
request_thread_ids.append(threading.get_ident())
|
||||
original_request_timer_rearm()
|
||||
|
||||
def tracked_arm_timer() -> None:
|
||||
arm_thread_ids.append(threading.get_ident())
|
||||
original_arm_timer()
|
||||
timer_rearmed.set()
|
||||
|
||||
monkeypatch.setattr(cron, "_request_timer_rearm", tracked_request_timer_rearm)
|
||||
monkeypatch.setattr(cron, "_arm_timer", tracked_arm_timer)
|
||||
channel = _ch(bus, cron_service=cron, port=_free_port())
|
||||
|
||||
try:
|
||||
response = await _webui_mutate(
|
||||
channel,
|
||||
"automation.update",
|
||||
{"id": job.id, "values": {"name": "After update"}},
|
||||
)
|
||||
await asyncio.wait_for(timer_rearmed.wait(), timeout=1)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert request_thread_ids
|
||||
assert all(thread_id != owner_thread_id for thread_id in request_thread_ids)
|
||||
assert arm_thread_ids and set(arm_thread_ids) == {owner_thread_id}
|
||||
assert cron._timer_task is not None
|
||||
assert cron._timer_task is not initial_timer
|
||||
assert not cron._timer_task.done()
|
||||
|
||||
stored = json.loads(store_path.read_text(encoding="utf-8"))
|
||||
assert len(stored["jobs"]) == 1
|
||||
assert stored["jobs"][0]["name"] == "After update"
|
||||
finally:
|
||||
cron.stop()
|
||||
await asyncio.sleep(0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_webui_automations_route_manages_local_triggers(
|
||||
bus: MagicMock, tmp_path: Path
|
||||
@@ -3769,3 +3833,77 @@ def test_bootstrap_secret_also_enforced_on_localhost(bus: MagicMock) -> None:
|
||||
channel = _ch(bus, host="0.0.0.0", tokenIssueSecret="s3cret")
|
||||
resp = channel.gateway.http._handle_bootstrap(_LOCAL, _NO_HEADERS)
|
||||
assert resp.status_code == 401
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_webui_skill_update_cancellation_waits_for_config_and_runtime_state(
|
||||
bus: MagicMock,
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
from nanobot.webui import ws_http
|
||||
|
||||
skill_dir = tmp_path / "skills" / "cancel-safe-skill"
|
||||
skill_dir.mkdir(parents=True)
|
||||
(skill_dir / "SKILL.md").write_text(
|
||||
"---\nname: cancel-safe-skill\ndescription: Cancellation test skill.\n---\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
mutation_started = threading.Event()
|
||||
release_mutation = threading.Event()
|
||||
original_update = ws_http.set_webui_skill_enabled
|
||||
update_calls = 0
|
||||
|
||||
def blocked_update(*args: Any, **kwargs: Any) -> dict[str, Any]:
|
||||
nonlocal update_calls
|
||||
update_calls += 1
|
||||
mutation_started.set()
|
||||
assert release_mutation.wait(timeout=1)
|
||||
return original_update(*args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(ws_http, "set_webui_skill_enabled", blocked_update)
|
||||
channel = _ch(
|
||||
bus,
|
||||
session_manager=_seed_session(tmp_path),
|
||||
workspace_path=tmp_path,
|
||||
port=_free_port(),
|
||||
)
|
||||
runtime_states: list[set[str]] = []
|
||||
channel.gateway.http.skill_state_action = runtime_states.append
|
||||
task = asyncio.create_task(
|
||||
_webui_mutate(
|
||||
channel,
|
||||
"skill.update",
|
||||
{"name": "cancel-safe-skill", "enabled": False},
|
||||
)
|
||||
)
|
||||
|
||||
assert await asyncio.to_thread(mutation_started.wait, 1)
|
||||
try:
|
||||
task.cancel()
|
||||
await asyncio.sleep(0)
|
||||
assert not task.done()
|
||||
assert runtime_states == []
|
||||
finally:
|
||||
release_mutation.set()
|
||||
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await asyncio.wait_for(task, timeout=1)
|
||||
|
||||
assert update_calls == 1
|
||||
assert "cancel-safe-skill" in channel.gateway.http.disabled_skills
|
||||
assert runtime_states == [{"cancel-safe-skill"}]
|
||||
saved = load_config(channel.gateway.settings.config.path)
|
||||
assert "cancel-safe-skill" in saved.agents.defaults.disabled_skills
|
||||
|
||||
settled_state = (
|
||||
update_calls,
|
||||
set(channel.gateway.http.disabled_skills),
|
||||
list(runtime_states),
|
||||
)
|
||||
await asyncio.sleep(0.05)
|
||||
assert (
|
||||
update_calls,
|
||||
set(channel.gateway.http.disabled_skills),
|
||||
runtime_states,
|
||||
) == settled_state
|
||||
|
||||
@@ -34,6 +34,7 @@ from nanobot.config.paths import is_default_workspace
|
||||
from nanobot.config.schema import Config
|
||||
from nanobot.gateway.runtime import GatewayInstance
|
||||
from nanobot.security.network import is_loopback_host
|
||||
from nanobot.session.async_compat import call_session_manager as _call_session_manager
|
||||
from nanobot.session.keys import UNIFIED_SESSION_KEY, last_channel_from_metadata
|
||||
from nanobot.utils.evaluator import evaluate_response, resolve_evaluator_prompt
|
||||
from nanobot.utils.helpers import sync_workspace_templates
|
||||
@@ -45,6 +46,30 @@ __all__ = ["_run_gateway"]
|
||||
|
||||
console = Console()
|
||||
|
||||
_EVENT_LOOP_LAG_INTERVAL_S = 0.5
|
||||
_EVENT_LOOP_LAG_WARNING_S = 0.25
|
||||
|
||||
|
||||
async def _monitor_event_loop_lag(
|
||||
*,
|
||||
interval_s: float = _EVENT_LOOP_LAG_INTERVAL_S,
|
||||
warning_threshold_s: float = _EVENT_LOOP_LAG_WARNING_S,
|
||||
log: Any | None = None,
|
||||
) -> None:
|
||||
"""Log scheduler drift so gateway-wide stalls have direct evidence."""
|
||||
loop = asyncio.get_running_loop()
|
||||
lag_log = log or logger
|
||||
while True:
|
||||
expected = loop.time() + interval_s
|
||||
await asyncio.sleep(interval_s)
|
||||
lag_s = max(0.0, loop.time() - expected)
|
||||
if lag_s >= warning_threshold_s:
|
||||
lag_log.warning(
|
||||
"event loop lag operation=gateway duration_ms={} interval_ms={}",
|
||||
int(lag_s * 1000),
|
||||
int(interval_s * 1000),
|
||||
)
|
||||
|
||||
|
||||
def _http_endpoint_responding(url: str, *, timeout_s: float = 0.25) -> bool:
|
||||
"""Return whether an HTTP endpoint responds, including with an auth error."""
|
||||
@@ -493,12 +518,22 @@ def _run_gateway(
|
||||
and hasattr(session_manager, "save")
|
||||
):
|
||||
key = session_key or _channel_session_key(msg.channel, msg.chat_id)
|
||||
session = session_manager.get_or_create(key)
|
||||
session = await _call_session_manager(
|
||||
session_manager,
|
||||
"get_or_create_async",
|
||||
session_manager.get_or_create,
|
||||
key,
|
||||
)
|
||||
extra: dict[str, Any] = {"_channel_delivery": True}
|
||||
if msg.media:
|
||||
extra["media"] = list(msg.media)
|
||||
session.add_message("assistant", msg.content, **extra)
|
||||
session_manager.save(session)
|
||||
await _call_session_manager(
|
||||
session_manager,
|
||||
"save_async",
|
||||
session_manager.save,
|
||||
session,
|
||||
)
|
||||
await bus.publish_outbound(msg)
|
||||
|
||||
message_tool = agent.tools.get("message")
|
||||
@@ -568,14 +603,14 @@ def _run_gateway(
|
||||
if sha:
|
||||
logger.info("Dream commit: {}", sha)
|
||||
store.compact_history()
|
||||
prune_dream_sessions(agent.sessions)
|
||||
await asyncio.to_thread(prune_dream_sessions, agent.sessions)
|
||||
return None
|
||||
|
||||
# Heartbeat is a system job that checks HEARTBEAT.md for active tasks.
|
||||
if job.name == "heartbeat":
|
||||
heartbeat_file = config.workspace_path / "HEARTBEAT.md"
|
||||
try:
|
||||
content = heartbeat_file.read_text(encoding="utf-8")
|
||||
content = await asyncio.to_thread(heartbeat_file.read_text, encoding="utf-8")
|
||||
except OSError:
|
||||
logger.debug("Heartbeat: HEARTBEAT.md missing")
|
||||
return None
|
||||
@@ -583,7 +618,7 @@ def _run_gateway(
|
||||
logger.debug("Heartbeat: HEARTBEAT.md has no active tasks")
|
||||
return None
|
||||
|
||||
channel, chat_id = _pick_heartbeat_target()
|
||||
channel, chat_id = await _pick_heartbeat_target()
|
||||
if channel == "cli":
|
||||
return None
|
||||
|
||||
@@ -611,9 +646,19 @@ def _run_gateway(
|
||||
message_tool.reset_suppress_delivery(suppress_token)
|
||||
|
||||
# Keep a small tail of heartbeat history so the loop stays bounded.
|
||||
session = agent.sessions.get_or_create("heartbeat")
|
||||
session = await _call_session_manager(
|
||||
agent.sessions,
|
||||
"get_or_create_async",
|
||||
agent.sessions.get_or_create,
|
||||
"heartbeat",
|
||||
)
|
||||
session.retain_recent_legal_suffix(hb_cfg.keep_recent_messages)
|
||||
agent.sessions.save(session)
|
||||
await _call_session_manager(
|
||||
agent.sessions,
|
||||
"save_async",
|
||||
agent.sessions.save,
|
||||
session,
|
||||
)
|
||||
|
||||
if not resp or not resp.content:
|
||||
return
|
||||
@@ -690,17 +735,27 @@ def _run_gateway(
|
||||
config_path=Path(config_path),
|
||||
)
|
||||
|
||||
def _pick_heartbeat_target() -> tuple[str, str]:
|
||||
async def _pick_heartbeat_target() -> tuple[str, str]:
|
||||
"""Pick a routable channel/chat target for heartbeat-triggered messages."""
|
||||
sidebar_state = read_webui_sidebar_state()
|
||||
sidebar_state = await asyncio.to_thread(read_webui_sidebar_state)
|
||||
unified_metadata = None
|
||||
if config.agents.defaults.unified_session:
|
||||
record = session_manager.read_session_metadata(UNIFIED_SESSION_KEY)
|
||||
record = await _call_session_manager(
|
||||
session_manager,
|
||||
"read_session_metadata_async",
|
||||
session_manager.read_session_metadata,
|
||||
UNIFIED_SESSION_KEY,
|
||||
)
|
||||
if isinstance(record, dict) and isinstance(record.get("metadata"), dict):
|
||||
unified_metadata = record["metadata"]
|
||||
sessions = await _call_session_manager(
|
||||
session_manager,
|
||||
"list_sessions_async",
|
||||
session_manager.list_sessions,
|
||||
)
|
||||
return _pick_heartbeat_target_from_sessions(
|
||||
enabled_channels=channels.enabled_channels,
|
||||
sessions=session_manager.list_sessions(),
|
||||
sessions=sessions,
|
||||
archived_keys=sidebar_state.get("archived_keys", []),
|
||||
unified_session_metadata=unified_metadata,
|
||||
)
|
||||
@@ -907,6 +962,10 @@ def _run_gateway(
|
||||
_monitor_local_clients(),
|
||||
name="nanobot-gateway-client-monitor",
|
||||
),
|
||||
asyncio.create_task(
|
||||
_monitor_event_loop_lag(),
|
||||
name="nanobot-event-loop-lag-monitor",
|
||||
),
|
||||
]
|
||||
if health_server_enabled:
|
||||
tasks.append(asyncio.create_task(
|
||||
@@ -974,7 +1033,11 @@ def _run_gateway(
|
||||
# Flush all cached sessions to durable storage before exit.
|
||||
# This prevents data loss on filesystems with write-back
|
||||
# caching (rclone VFS, NFS, FUSE mounts, etc.).
|
||||
flushed = agent.sessions.flush_all()
|
||||
flushed = await _call_session_manager(
|
||||
agent.sessions,
|
||||
"flush_all_async",
|
||||
agent.sessions.flush_all,
|
||||
)
|
||||
if flushed:
|
||||
logger.info("Shutdown: flushed {} session(s) to disk", flushed)
|
||||
finally:
|
||||
|
||||
+70
-17
@@ -3,6 +3,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import inspect
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
@@ -14,6 +15,8 @@ from typing import TYPE_CHECKING, Any, Literal, cast
|
||||
from nanobot import __version__
|
||||
from nanobot.bus.events import INBOUND_META_USER_SHELL, OutboundMessage
|
||||
from nanobot.command.router import CommandContext, CommandRouter, normalize_command_text
|
||||
from nanobot.session.async_compat import call_session_manager
|
||||
from nanobot.utils.cancellation import shield_and_drain
|
||||
from nanobot.utils.helpers import build_status_content
|
||||
from nanobot.utils.restart import set_restart_notice_to_env
|
||||
from nanobot.utils.workspace_prompts import initialize_workspace_prompt
|
||||
@@ -22,6 +25,7 @@ if TYPE_CHECKING:
|
||||
from nanobot.agent.loop import AgentLoop
|
||||
from nanobot.session.manager import Session
|
||||
from nanobot.utils.gitstore import CommitInfo
|
||||
from nanobot.utils.llm_runtime import LLMRuntime
|
||||
|
||||
# WebUI protocol contract for how a slash command participates in turn state:
|
||||
# - side_channel: returns control text without starting or ending an agent turn.
|
||||
@@ -201,6 +205,52 @@ def builtin_command_starts_agent_turn(text: str) -> bool:
|
||||
return spec.lifecycle == "agent_turn_with_args" and bool(args.strip())
|
||||
|
||||
|
||||
def _has_native_coroutine_method(target: object, name: str) -> bool:
|
||||
"""Check the real target class without trusting dynamic mock attributes."""
|
||||
method = inspect.getattr_static(type(target), name, None)
|
||||
return inspect.iscoroutinefunction(method)
|
||||
|
||||
|
||||
async def _get_or_create_session(loop: AgentLoop, key: str) -> Session:
|
||||
sessions = loop.sessions
|
||||
return await call_session_manager(
|
||||
sessions,
|
||||
"get_or_create_async",
|
||||
sessions.get_or_create,
|
||||
key,
|
||||
)
|
||||
|
||||
|
||||
async def _save_session(loop: AgentLoop, session: Session) -> None:
|
||||
sessions = loop.sessions
|
||||
await call_session_manager(
|
||||
sessions,
|
||||
"save_async",
|
||||
sessions.save,
|
||||
session,
|
||||
)
|
||||
|
||||
|
||||
async def _runtime_for_session(loop: AgentLoop, session: Session) -> LLMRuntime:
|
||||
if _has_native_coroutine_method(loop, "runtime_for_session_async"):
|
||||
return await loop.runtime_for_session_async(session)
|
||||
return await shield_and_drain(
|
||||
asyncio.to_thread(loop.runtime_for_session, session)
|
||||
)
|
||||
|
||||
|
||||
async def _set_session_model_preset(
|
||||
loop: AgentLoop,
|
||||
session_key: str,
|
||||
name: str,
|
||||
) -> LLMRuntime:
|
||||
if _has_native_coroutine_method(loop, "set_session_model_preset_async"):
|
||||
return await loop.set_session_model_preset_async(session_key, name)
|
||||
return await shield_and_drain(
|
||||
asyncio.to_thread(loop.set_session_model_preset, session_key, name)
|
||||
)
|
||||
|
||||
|
||||
async def cmd_stop(ctx: CommandContext) -> OutboundMessage:
|
||||
"""Cancel all active tasks and subagents for the session."""
|
||||
loop = ctx.loop
|
||||
@@ -257,8 +307,8 @@ async def cmd_restart(ctx: CommandContext) -> OutboundMessage:
|
||||
async def cmd_status(ctx: CommandContext) -> OutboundMessage:
|
||||
"""Build an outbound status message for a session."""
|
||||
loop = ctx.loop
|
||||
session = ctx.session or loop.sessions.get_or_create(ctx.key)
|
||||
runtime = ctx.runtime or loop.runtime_for_session(session)
|
||||
session = ctx.session or await _get_or_create_session(loop, ctx.key)
|
||||
runtime = ctx.runtime or await _runtime_for_session(loop, session)
|
||||
ctx_est = 0
|
||||
with suppress(Exception):
|
||||
ctx_est, _ = loop.consolidator.estimate_session_prompt_tokens(
|
||||
@@ -306,29 +356,32 @@ async def cmd_new(ctx: CommandContext) -> OutboundMessage:
|
||||
loop = ctx.loop
|
||||
await loop._cancel_active_tasks(ctx.key) # pyright: ignore[reportPrivateUsage]
|
||||
loop.discard_session_file_state(ctx.key)
|
||||
session = ctx.session or loop.sessions.get_or_create(ctx.key)
|
||||
session = ctx.session or await _get_or_create_session(loop, ctx.key)
|
||||
snapshot = list(session.messages)
|
||||
archive_snapshot = None
|
||||
runtime = None
|
||||
if session.last_consolidated < len(snapshot):
|
||||
runtime = ctx.runtime or loop.runtime_for_session(session)
|
||||
runtime = ctx.runtime or await _runtime_for_session(loop, session)
|
||||
archive_snapshot = replace(
|
||||
session,
|
||||
messages=snapshot,
|
||||
metadata=dict(session.metadata),
|
||||
provider_state=None,
|
||||
)
|
||||
session.clear()
|
||||
loop.sessions.save(session)
|
||||
loop.sessions.invalidate(session.key)
|
||||
if archive_snapshot is not None and runtime is not None:
|
||||
loop.schedule_background(
|
||||
loop.consolidator.archive_session( # pyright: ignore[reportUnknownMemberType]
|
||||
archive_snapshot,
|
||||
archive_end=len(snapshot),
|
||||
runtime=runtime,
|
||||
async def reset_and_schedule_archive() -> None:
|
||||
session.clear()
|
||||
await _save_session(loop, session)
|
||||
loop.sessions.invalidate(session.key)
|
||||
if archive_snapshot is not None and runtime is not None:
|
||||
loop.schedule_background(
|
||||
loop.consolidator.archive_session( # pyright: ignore[reportUnknownMemberType]
|
||||
archive_snapshot,
|
||||
archive_end=len(snapshot),
|
||||
runtime=runtime,
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
await shield_and_drain(reset_and_schedule_archive())
|
||||
return OutboundMessage(
|
||||
channel=ctx.msg.channel, chat_id=ctx.msg.chat_id,
|
||||
content="New session started.",
|
||||
@@ -377,7 +430,7 @@ async def cmd_model(ctx: CommandContext) -> OutboundMessage:
|
||||
metadata = {**dict(ctx.msg.metadata or {}), "render_as": "text"}
|
||||
|
||||
if not args:
|
||||
session = ctx.session or loop.sessions.get_or_create(ctx.key)
|
||||
session = ctx.session or await _get_or_create_session(loop, ctx.key)
|
||||
return OutboundMessage(
|
||||
channel=ctx.msg.channel,
|
||||
chat_id=ctx.msg.chat_id,
|
||||
@@ -387,7 +440,7 @@ async def cmd_model(ctx: CommandContext) -> OutboundMessage:
|
||||
|
||||
name = args
|
||||
try:
|
||||
runtime = loop.set_session_model_preset(ctx.key, name)
|
||||
runtime = await _set_session_model_preset(loop, ctx.key, name)
|
||||
except (KeyError, ValueError) as exc:
|
||||
names = _model_preset_names(loop)
|
||||
return OutboundMessage(
|
||||
@@ -848,7 +901,7 @@ async def cmd_history(ctx: CommandContext) -> OutboundMessage:
|
||||
metadata=dict(ctx.msg.metadata or {}),
|
||||
)
|
||||
|
||||
session = ctx.session or ctx.loop.sessions.get_or_create(ctx.key)
|
||||
session = ctx.session or await _get_or_create_session(ctx.loop, ctx.key)
|
||||
history = session.get_history(max_messages=0, include_runtime_context=False)
|
||||
visible = [_format_history_message(m) for m in history]
|
||||
visible = [m for m in visible if m is not None]
|
||||
|
||||
@@ -8,6 +8,7 @@ import time
|
||||
import uuid
|
||||
from typing import TYPE_CHECKING, Any, Protocol
|
||||
|
||||
from nanobot.agent.automation_turns import AutomationTurnAcceptedCancellation
|
||||
from nanobot.agent.tools.cron import CronTool
|
||||
from nanobot.bus.events import InboundMessage, OutboundMessage
|
||||
from nanobot.cron.session_delivery import origin_delivery_context
|
||||
@@ -127,6 +128,15 @@ async def run_bound_cron_job(
|
||||
session_key_override=session_key,
|
||||
)
|
||||
)
|
||||
except AutomationTurnAcceptedCancellation:
|
||||
cron.write_run_record(
|
||||
run_id,
|
||||
{
|
||||
**run_record_base,
|
||||
"status": "accepted",
|
||||
},
|
||||
)
|
||||
raise
|
||||
except (Exception, asyncio.CancelledError) as exc:
|
||||
error_text = str(exc) or exc.__class__.__name__
|
||||
cron.write_run_record(
|
||||
|
||||
+230
-87
@@ -11,11 +11,12 @@ from dataclasses import asdict
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from types import EllipsisType
|
||||
from typing import Any, Callable, Coroutine, Literal
|
||||
from typing import Any, Callable, Coroutine, Literal, TypeVar
|
||||
|
||||
from filelock import FileLock
|
||||
from loguru import logger
|
||||
|
||||
from nanobot.agent.automation_turns import AutomationTurnAcceptedCancellation
|
||||
from nanobot.cron.session_turns import is_bound_cron_job
|
||||
from nanobot.cron.types import (
|
||||
CronJob,
|
||||
@@ -25,10 +26,14 @@ from nanobot.cron.types import (
|
||||
CronSchedule,
|
||||
CronStore,
|
||||
)
|
||||
from nanobot.utils.cancellation import shield_and_drain
|
||||
from nanobot.utils.run_records import (
|
||||
write_run_record as write_automation_run_record,
|
||||
)
|
||||
|
||||
_FILE_LOCK_TIMEOUT_SECONDS = 5
|
||||
_T = TypeVar("_T")
|
||||
|
||||
|
||||
class CronJobSkippedError(Exception):
|
||||
"""Raised by cron callbacks when a job was intentionally skipped."""
|
||||
@@ -164,10 +169,16 @@ class CronService:
|
||||
self.store_path = store_path
|
||||
self._action_path = store_path.parent / "action.jsonl"
|
||||
self._run_records_dir = store_path.parent / "runs"
|
||||
self._lock = FileLock(str(self._action_path.parent) + ".lock")
|
||||
self._lock = FileLock(
|
||||
str(self._action_path.parent) + ".lock",
|
||||
timeout=_FILE_LOCK_TIMEOUT_SECONDS,
|
||||
)
|
||||
self.on_job = on_job
|
||||
self._store: CronStore | None = None
|
||||
self._timer_task: asyncio.Task[None] | None = None
|
||||
self._operation_lock = asyncio.Lock()
|
||||
self._claimed_job_ids: set[str] = set()
|
||||
self._event_loop: asyncio.AbstractEventLoop | None = None
|
||||
self._running = False
|
||||
self._active_executions = 0
|
||||
self._store_dirty = False
|
||||
@@ -451,25 +462,58 @@ class CronService:
|
||||
"""Write an internal audit record for one cron execution."""
|
||||
write_automation_run_record(self._run_records_dir, run_id, record)
|
||||
|
||||
async def start(self) -> None:
|
||||
"""Start the cron service."""
|
||||
self._running = True
|
||||
loaded = self._load_store()
|
||||
if loaded is None:
|
||||
# Store file existed but was corrupt and has been preserved with
|
||||
# a ``.corrupt-<ts>`` suffix. Bail out instead of starting with
|
||||
# an empty store; that would call ``_save_store`` and overwrite
|
||||
# the now-renamed (but still recoverable) data with [].
|
||||
self._running = False
|
||||
raise RuntimeError(
|
||||
f"cron store at {self.store_path} is corrupt and was preserved; "
|
||||
"refusing to start with an empty job list. "
|
||||
"Inspect the .corrupt-<ts> backup and restore manually."
|
||||
async def run_sync(
|
||||
self,
|
||||
operation: Callable[..., _T],
|
||||
/,
|
||||
*args: Any,
|
||||
**kwargs: Any,
|
||||
) -> _T:
|
||||
"""Serialize a complete cron transaction in a worker thread.
|
||||
|
||||
A running thread cannot be cancelled safely. Keep the transaction lock
|
||||
until it exits so cancellation is never reported while that worker can
|
||||
still mutate cron state behind a later operation.
|
||||
"""
|
||||
async with self._operation_lock:
|
||||
return await shield_and_drain(
|
||||
asyncio.to_thread(operation, *args, **kwargs)
|
||||
)
|
||||
self._recompute_next_runs()
|
||||
self._save_store()
|
||||
self._arm_timer()
|
||||
logger.info("Cron service started with {} jobs", len(self._store.jobs if self._store else []))
|
||||
|
||||
async def start(self) -> None:
|
||||
"""Start the cron service and settle accepted work before cancellation."""
|
||||
|
||||
async def settle_start() -> None:
|
||||
self._event_loop = asyncio.get_running_loop()
|
||||
self._running = True
|
||||
try:
|
||||
async with self._operation_lock:
|
||||
loaded = await asyncio.to_thread(self._load_store)
|
||||
if loaded is None:
|
||||
# Store file existed but was corrupt and has been preserved with
|
||||
# a ``.corrupt-<ts>`` suffix. Bail out instead of starting with
|
||||
# an empty store; that would call ``_save_store`` and overwrite
|
||||
# the now-renamed (but still recoverable) data with [].
|
||||
raise RuntimeError(
|
||||
f"cron store at {self.store_path} is corrupt and was preserved; "
|
||||
"refusing to start with an empty job list. "
|
||||
"Inspect the .corrupt-<ts> backup and restore manually."
|
||||
)
|
||||
self._recompute_next_runs()
|
||||
await asyncio.to_thread(self._save_store)
|
||||
self._arm_timer()
|
||||
logger.info(
|
||||
"Cron service started with {} jobs",
|
||||
len(self._store.jobs if self._store else []),
|
||||
)
|
||||
except BaseException:
|
||||
# A failed start must not retain ownership without a timer. Caller
|
||||
# cancellation is shielded until this composite either reaches the
|
||||
# fully started state above or rolls back here.
|
||||
self.stop()
|
||||
raise
|
||||
|
||||
await shield_and_drain(settle_start())
|
||||
|
||||
def stop(self) -> None:
|
||||
"""Stop the cron service."""
|
||||
@@ -497,8 +541,22 @@ class CronService:
|
||||
if j.enabled and j.state.next_run_at_ms]
|
||||
return min(times) if times else None
|
||||
|
||||
def _request_timer_rearm(self) -> None:
|
||||
"""Re-arm on the owning event loop, including from persistence workers."""
|
||||
if not self._running:
|
||||
return
|
||||
loop = self._event_loop
|
||||
try:
|
||||
current_loop = asyncio.get_running_loop()
|
||||
except RuntimeError:
|
||||
current_loop = None
|
||||
if current_loop is loop:
|
||||
self._arm_timer()
|
||||
elif loop is not None and loop.is_running():
|
||||
loop.call_soon_threadsafe(self._arm_timer)
|
||||
|
||||
def _arm_timer(self) -> None:
|
||||
"""Schedule the next timer tick."""
|
||||
"""Schedule the next timer tick on the owning event loop."""
|
||||
if self._timer_task:
|
||||
self._timer_task.cancel()
|
||||
|
||||
@@ -520,7 +578,7 @@ class CronService:
|
||||
self._timer_task = asyncio.create_task(tick())
|
||||
|
||||
async def _on_timer(self) -> None:
|
||||
"""Handle timer tick - run due jobs."""
|
||||
"""Run due jobs while keeping persistence transactions serialized."""
|
||||
reload_store = self._active_executions == 0
|
||||
self._active_executions += 1
|
||||
try:
|
||||
@@ -528,11 +586,17 @@ class CronService:
|
||||
# to persist their advanced schedule. Persist that exact snapshot
|
||||
# before reloading or executing anything else; otherwise the older
|
||||
# disk state can replay the same job.
|
||||
if self._store_dirty:
|
||||
self._save_store()
|
||||
return
|
||||
async with self._operation_lock:
|
||||
if self._store_dirty:
|
||||
await shield_and_drain(asyncio.to_thread(self._save_store))
|
||||
return
|
||||
|
||||
store = self._load_store(reload_during_execution=reload_store)
|
||||
store = await shield_and_drain(
|
||||
asyncio.to_thread(
|
||||
self._load_store,
|
||||
reload_during_execution=reload_store,
|
||||
)
|
||||
)
|
||||
# If a hot reload found a corrupt store on disk, ``self._store``
|
||||
# may still hold the previous, known-good in-memory snapshot.
|
||||
if store is None:
|
||||
@@ -547,7 +611,8 @@ class CronService:
|
||||
for job in due_jobs:
|
||||
await self._execute_job(job)
|
||||
|
||||
self._save_store()
|
||||
async with self._operation_lock:
|
||||
await shield_and_drain(asyncio.to_thread(self._save_store))
|
||||
except Exception:
|
||||
# A load/persist failure must not kill the scheduler: keep the
|
||||
# in-memory store and retry on the next tick. This mirrors the
|
||||
@@ -564,58 +629,124 @@ class CronService:
|
||||
# single bad tick cannot silently stop all future jobs.
|
||||
self._arm_timer()
|
||||
|
||||
async def _execute_job(self, job: CronJob) -> None:
|
||||
"""Execute a single job."""
|
||||
async def _claim_job(self, job_id: str) -> bool:
|
||||
"""Claim one job without serializing callbacks for different jobs."""
|
||||
async with self._operation_lock:
|
||||
if job_id in self._claimed_job_ids:
|
||||
return False
|
||||
self._claimed_job_ids.add(job_id)
|
||||
return True
|
||||
|
||||
async def _release_job_claim(self, job_id: str) -> None:
|
||||
async with self._operation_lock:
|
||||
self._claimed_job_ids.discard(job_id)
|
||||
|
||||
async def _settle_job_execution(
|
||||
self,
|
||||
job: CronJob,
|
||||
*,
|
||||
start_ms: int,
|
||||
status: Literal["ok", "error", "skipped"],
|
||||
error: str | None,
|
||||
persist: bool = False,
|
||||
) -> None:
|
||||
end_ms = _now_ms()
|
||||
async with self._operation_lock:
|
||||
job.state.last_status = status
|
||||
job.state.last_error = error
|
||||
job.state.last_run_at_ms = start_ms
|
||||
job.updated_at_ms = end_ms
|
||||
job.state.run_history.append(CronRunRecord(
|
||||
run_at_ms=start_ms,
|
||||
status=status,
|
||||
duration_ms=end_ms - start_ms,
|
||||
error=error,
|
||||
))
|
||||
job.state.run_history = job.state.run_history[-self._MAX_RUN_HISTORY:]
|
||||
|
||||
if job.schedule.kind == "at":
|
||||
if job.delete_after_run:
|
||||
store = await shield_and_drain(
|
||||
asyncio.to_thread(self._require_store)
|
||||
)
|
||||
store.jobs = [item for item in store.jobs if item.id != job.id]
|
||||
else:
|
||||
job.enabled = False
|
||||
job.state.next_run_at_ms = None
|
||||
else:
|
||||
job.state.next_run_at_ms = _compute_next_run(job.schedule, _now_ms())
|
||||
|
||||
if persist:
|
||||
await shield_and_drain(asyncio.to_thread(self._save_store))
|
||||
|
||||
@staticmethod
|
||||
async def _drain_settlement_on_cancellation(settlement: asyncio.Task[None]) -> None:
|
||||
"""Finish a short durable settlement despite repeated cancellation."""
|
||||
while not settlement.done():
|
||||
try:
|
||||
await asyncio.shield(settlement)
|
||||
except asyncio.CancelledError:
|
||||
continue
|
||||
settlement.result()
|
||||
|
||||
async def _execute_job(self, job: CronJob) -> bool:
|
||||
"""Execute a claimed job and serialize its in-memory settlement."""
|
||||
if not await self._claim_job(job.id):
|
||||
logger.info("Cron: job '{}' ({}) is already running", job.name, job.id)
|
||||
return False
|
||||
|
||||
start_ms = _now_ms()
|
||||
logger.info("Cron: executing job '{}' ({})", job.name, job.id)
|
||||
status: Literal["ok", "error", "skipped"]
|
||||
error: str | None
|
||||
accepted_cancellation: AutomationTurnAcceptedCancellation | None = None
|
||||
|
||||
try:
|
||||
if self.on_job:
|
||||
await self.on_job(job)
|
||||
try:
|
||||
if self.on_job:
|
||||
await self.on_job(job)
|
||||
status = "ok"
|
||||
error = None
|
||||
logger.info("Cron: job '{}' completed", job.name)
|
||||
except AutomationTurnAcceptedCancellation as exc:
|
||||
# The agent owns this turn now. Advance and persist the schedule
|
||||
# before allowing shutdown cancellation to unwind the timer.
|
||||
status = "ok"
|
||||
error = None
|
||||
accepted_cancellation = exc
|
||||
logger.info("Cron: job '{}' was accepted before cancellation", job.name)
|
||||
except CronJobSkippedError as exc:
|
||||
status = "skipped"
|
||||
error = str(exc) or None
|
||||
logger.warning("Cron: job '{}' skipped: {}", job.name, error or "")
|
||||
except asyncio.CancelledError as exc:
|
||||
current = asyncio.current_task()
|
||||
if current is not None and current.cancelling():
|
||||
raise
|
||||
status = "error"
|
||||
error = str(exc) or exc.__class__.__name__
|
||||
logger.exception("Cron: job '{}' was cancelled", job.name)
|
||||
except Exception as exc:
|
||||
status = "error"
|
||||
error = str(exc)
|
||||
logger.exception("Cron: job '{}' failed", job.name)
|
||||
|
||||
job.state.last_status = "ok"
|
||||
job.state.last_error = None
|
||||
logger.info("Cron: job '{}' completed", job.name)
|
||||
|
||||
except CronJobSkippedError as e:
|
||||
job.state.last_status = "skipped"
|
||||
job.state.last_error = str(e) or None
|
||||
logger.warning("Cron: job '{}' skipped: {}", job.name, job.state.last_error or "")
|
||||
except asyncio.CancelledError as e:
|
||||
current = asyncio.current_task()
|
||||
if current is not None and current.cancelling():
|
||||
raise
|
||||
job.state.last_status = "error"
|
||||
job.state.last_error = str(e) or e.__class__.__name__
|
||||
logger.exception("Cron: job '{}' was cancelled", job.name)
|
||||
except Exception as e:
|
||||
job.state.last_status = "error"
|
||||
job.state.last_error = str(e)
|
||||
logger.exception("Cron: job '{}' failed", job.name)
|
||||
|
||||
end_ms = _now_ms()
|
||||
job.state.last_run_at_ms = start_ms
|
||||
job.updated_at_ms = end_ms
|
||||
|
||||
job.state.run_history.append(CronRunRecord(
|
||||
run_at_ms=start_ms,
|
||||
status=job.state.last_status,
|
||||
duration_ms=end_ms - start_ms,
|
||||
error=job.state.last_error,
|
||||
))
|
||||
job.state.run_history = job.state.run_history[-self._MAX_RUN_HISTORY:]
|
||||
|
||||
# Handle one-shot jobs
|
||||
if job.schedule.kind == "at":
|
||||
if job.delete_after_run:
|
||||
store = self._require_store()
|
||||
store.jobs = [item for item in store.jobs if item.id != job.id]
|
||||
else:
|
||||
job.enabled = False
|
||||
job.state.next_run_at_ms = None
|
||||
else:
|
||||
# Compute next run
|
||||
job.state.next_run_at_ms = _compute_next_run(job.schedule, _now_ms())
|
||||
settlement = asyncio.create_task(
|
||||
self._settle_job_execution(
|
||||
job,
|
||||
start_ms=start_ms,
|
||||
status=status,
|
||||
error=error,
|
||||
persist=accepted_cancellation is not None,
|
||||
)
|
||||
)
|
||||
if accepted_cancellation is not None:
|
||||
await self._drain_settlement_on_cancellation(settlement)
|
||||
raise accepted_cancellation
|
||||
await settlement
|
||||
return True
|
||||
finally:
|
||||
await self._release_job_claim(job.id)
|
||||
|
||||
def _append_action(
|
||||
self,
|
||||
@@ -697,7 +828,7 @@ class CronService:
|
||||
store = self._require_store()
|
||||
store.jobs.append(job)
|
||||
self._save_store()
|
||||
self._arm_timer()
|
||||
self._request_timer_rearm()
|
||||
else:
|
||||
self._append_action("add", asdict(job))
|
||||
|
||||
@@ -714,7 +845,7 @@ class CronService:
|
||||
store.jobs = [j for j in store.jobs if j.id != job.id]
|
||||
store.jobs.append(job)
|
||||
self._save_store()
|
||||
self._arm_timer()
|
||||
self._request_timer_rearm()
|
||||
logger.info("Cron: registered system job '{}' ({})", job.name, job.id)
|
||||
return job
|
||||
|
||||
@@ -726,7 +857,7 @@ class CronService:
|
||||
removed = len(store.jobs) < before
|
||||
if removed:
|
||||
self._save_store()
|
||||
self._arm_timer()
|
||||
self._request_timer_rearm()
|
||||
logger.info("Cron: removed system job {}", job_id)
|
||||
return removed
|
||||
|
||||
@@ -747,7 +878,7 @@ class CronService:
|
||||
if removed:
|
||||
if self._should_persist_store():
|
||||
self._save_store()
|
||||
self._arm_timer()
|
||||
self._request_timer_rearm()
|
||||
else:
|
||||
self._append_action("del", {"job_id": job_id})
|
||||
logger.info("Cron: removed job {}", job_id)
|
||||
@@ -769,7 +900,7 @@ class CronService:
|
||||
job.state.next_run_at_ms = None
|
||||
if self._should_persist_store():
|
||||
self._save_store()
|
||||
self._arm_timer()
|
||||
self._request_timer_rearm()
|
||||
else:
|
||||
self._append_action("update", asdict(job))
|
||||
return job
|
||||
@@ -825,7 +956,7 @@ class CronService:
|
||||
|
||||
if self._should_persist_store():
|
||||
self._save_store()
|
||||
self._arm_timer()
|
||||
self._request_timer_rearm()
|
||||
else:
|
||||
self._append_action("update", asdict(job))
|
||||
|
||||
@@ -840,19 +971,31 @@ class CronService:
|
||||
# A manual run is another side-effecting entrypoint. Do not start
|
||||
# it while the result of a previous timer execution is still only
|
||||
# in memory.
|
||||
if self._store_dirty:
|
||||
self._save_store()
|
||||
store = self._require_store(reload_during_execution=reload_store)
|
||||
async with self._operation_lock:
|
||||
if self._store_dirty:
|
||||
await shield_and_drain(asyncio.to_thread(self._save_store))
|
||||
store = await shield_and_drain(
|
||||
asyncio.to_thread(
|
||||
self._require_store,
|
||||
reload_during_execution=reload_store,
|
||||
)
|
||||
)
|
||||
for job in store.jobs:
|
||||
if job.id == job_id:
|
||||
if self._is_unbound_agent_job(job):
|
||||
self._enforce_agent_binding(job)
|
||||
self._save_store()
|
||||
async with self._operation_lock:
|
||||
self._enforce_agent_binding(job)
|
||||
await shield_and_drain(
|
||||
asyncio.to_thread(self._save_store)
|
||||
)
|
||||
return False
|
||||
if not force and not job.enabled:
|
||||
return False
|
||||
await self._execute_job(job)
|
||||
self._save_store()
|
||||
executed = await self._execute_job(job)
|
||||
if not executed:
|
||||
return False
|
||||
async with self._operation_lock:
|
||||
await shield_and_drain(asyncio.to_thread(self._save_store))
|
||||
return True
|
||||
return False
|
||||
finally:
|
||||
|
||||
@@ -20,7 +20,7 @@ from nanobot.providers.registry import find_by_name
|
||||
from nanobot.security.network import (
|
||||
PinnedDNSAsyncTransport,
|
||||
UnsafeURLRequestError,
|
||||
resolve_url_target,
|
||||
async_resolve_url_target,
|
||||
)
|
||||
from nanobot.utils.helpers import detect_image_mime
|
||||
|
||||
@@ -174,7 +174,7 @@ async def _download_image_data_url(
|
||||
current_url = url
|
||||
for _ in range(_IMAGE_DOWNLOAD_MAX_REDIRECTS + 1):
|
||||
if proxy:
|
||||
ok, error, _ = resolve_url_target(
|
||||
ok, error, _ = await async_resolve_url_target(
|
||||
current_url,
|
||||
trust_remote_dns=True,
|
||||
)
|
||||
|
||||
@@ -210,18 +210,20 @@ class RuntimeClient:
|
||||
|
||||
async def compact_session(self, session_key: str) -> SessionSnapshot:
|
||||
"""Run token consolidation for one session."""
|
||||
session = self._loop.sessions.get_or_create(session_key)
|
||||
runtime = self._loop.runtime_for_session(session)
|
||||
session = await self._loop.sessions.get_or_create_async(session_key)
|
||||
runtime = await self._loop.runtime_for_session_async(session)
|
||||
await self._loop.consolidator.maybe_consolidate_by_tokens(
|
||||
session,
|
||||
runtime=runtime,
|
||||
)
|
||||
return snapshot_from_session(self._loop.sessions.get_or_create(session_key))
|
||||
return snapshot_from_session(
|
||||
await self._loop.sessions.get_or_create_async(session_key)
|
||||
)
|
||||
|
||||
async def compact_idle_session(self, session_key: str, *, max_suffix: int = 8) -> str | None:
|
||||
"""Run idle-session compaction for one session and return the summary."""
|
||||
session = self._loop.sessions.get_or_create(session_key)
|
||||
runtime = self._loop.runtime_for_session(session)
|
||||
session = await self._loop.sessions.get_or_create_async(session_key)
|
||||
runtime = await self._loop.runtime_for_session_async(session)
|
||||
return await self._loop.consolidator.compact_idle_session(
|
||||
session_key,
|
||||
runtime=runtime,
|
||||
|
||||
+144
-42
@@ -29,6 +29,7 @@ _BLOCKED_NETWORKS = [
|
||||
|
||||
_URL_RE = re.compile(r"https?://[^\s\"'`;|<>]+", re.IGNORECASE)
|
||||
_allowed_networks: list[ipaddress.IPv4Network | ipaddress.IPv6Network] = []
|
||||
_DNS_RESOLUTION_TIMEOUT_SECONDS = 5.0
|
||||
|
||||
|
||||
def is_loopback_host(host: str) -> bool:
|
||||
@@ -75,6 +76,63 @@ def _is_private(addr: ipaddress.IPv4Address | ipaddress.IPv6Address) -> bool:
|
||||
return any(normalized in net for net in _BLOCKED_NETWORKS)
|
||||
|
||||
|
||||
def _parse_url_hostname(url: str) -> tuple[str | None, str | None]:
|
||||
try:
|
||||
parsed = urlparse(url)
|
||||
except Exception as exc:
|
||||
return None, str(exc)
|
||||
if parsed.scheme not in ("http", "https"):
|
||||
return None, f"Only http/https allowed, got '{parsed.scheme or 'none'}'"
|
||||
if not parsed.netloc:
|
||||
return None, "Missing domain"
|
||||
if not parsed.hostname:
|
||||
return None, "Missing hostname"
|
||||
return parsed.hostname, None
|
||||
|
||||
|
||||
def _unresolved_target_result(
|
||||
hostname: str,
|
||||
*,
|
||||
trust_remote_dns: bool,
|
||||
) -> tuple[bool, str, tuple[str, ...]]:
|
||||
if not trust_remote_dns:
|
||||
return False, f"Cannot resolve hostname: {hostname}", ()
|
||||
|
||||
normalized_hostname = hostname.rstrip(".").lower()
|
||||
if normalized_hostname == "localhost" or normalized_hostname.endswith(".localhost"):
|
||||
return False, f"Blocked local/internal hostname: {hostname}", ()
|
||||
|
||||
try:
|
||||
literal_addr = ipaddress.ip_address(normalized_hostname)
|
||||
except ValueError:
|
||||
return True, "", ()
|
||||
if _is_private(literal_addr):
|
||||
return False, f"Blocked private/internal address: {literal_addr}", ()
|
||||
return True, "", (str(_normalize_addr(literal_addr)),)
|
||||
|
||||
|
||||
def _resolved_target_result(
|
||||
hostname: str,
|
||||
infos: list[Any],
|
||||
*,
|
||||
allow_loopback: bool,
|
||||
) -> tuple[bool, str, tuple[str, ...]]:
|
||||
addrs: list[ipaddress.IPv4Address | ipaddress.IPv6Address] = []
|
||||
for info in infos:
|
||||
try:
|
||||
addr = ipaddress.ip_address(info[4][0])
|
||||
except (IndexError, TypeError, ValueError):
|
||||
continue
|
||||
addrs.append(addr)
|
||||
if allow_loopback and _is_allowed_loopback_target(hostname, addrs):
|
||||
return True, "", tuple(dict.fromkeys(str(_normalize_addr(addr)) for addr in addrs))
|
||||
for addr in addrs:
|
||||
if _is_private(addr):
|
||||
return False, f"Blocked: {hostname} resolves to private/internal address {addr}", ()
|
||||
|
||||
return True, "", tuple(dict.fromkeys(str(_normalize_addr(addr)) for addr in addrs))
|
||||
|
||||
|
||||
def resolve_url_target(
|
||||
url: str,
|
||||
*,
|
||||
@@ -97,52 +155,43 @@ def resolve_url_target(
|
||||
resolved_ips contains the public IPs that were validated for this URL, or
|
||||
is empty when an unresolved hostname is delegated to a trusted proxy.
|
||||
"""
|
||||
try:
|
||||
p = urlparse(url)
|
||||
except Exception as e:
|
||||
return False, str(e), ()
|
||||
|
||||
if p.scheme not in ("http", "https"):
|
||||
return False, f"Only http/https allowed, got '{p.scheme or 'none'}'", ()
|
||||
if not p.netloc:
|
||||
return False, "Missing domain", ()
|
||||
|
||||
hostname = p.hostname
|
||||
if not hostname:
|
||||
return False, "Missing hostname", ()
|
||||
|
||||
hostname, error = _parse_url_hostname(url)
|
||||
if hostname is None:
|
||||
return False, error or "Missing hostname", ()
|
||||
try:
|
||||
infos = socket.getaddrinfo(hostname, None, socket.AF_UNSPEC, socket.SOCK_STREAM)
|
||||
except socket.gaierror:
|
||||
if not trust_remote_dns:
|
||||
return False, f"Cannot resolve hostname: {hostname}", ()
|
||||
return _unresolved_target_result(hostname, trust_remote_dns=trust_remote_dns)
|
||||
return _resolved_target_result(hostname, infos, allow_loopback=allow_loopback)
|
||||
|
||||
normalized_hostname = hostname.rstrip(".").lower()
|
||||
if normalized_hostname == "localhost" or normalized_hostname.endswith(".localhost"):
|
||||
return False, f"Blocked local/internal hostname: {hostname}", ()
|
||||
|
||||
try:
|
||||
literal_addr = ipaddress.ip_address(normalized_hostname)
|
||||
except ValueError:
|
||||
return True, "", ()
|
||||
if _is_private(literal_addr):
|
||||
return False, f"Blocked private/internal address: {literal_addr}", ()
|
||||
return True, "", (str(_normalize_addr(literal_addr)),)
|
||||
|
||||
addrs: list[ipaddress.IPv4Address | ipaddress.IPv6Address] = []
|
||||
for info in infos:
|
||||
try:
|
||||
addr = ipaddress.ip_address(info[4][0])
|
||||
except ValueError:
|
||||
continue
|
||||
addrs.append(addr)
|
||||
if allow_loopback and _is_allowed_loopback_target(hostname, addrs):
|
||||
return True, "", tuple(dict.fromkeys(str(_normalize_addr(addr)) for addr in addrs))
|
||||
for addr in addrs:
|
||||
if _is_private(addr):
|
||||
return False, f"Blocked: {hostname} resolves to private/internal address {addr}", ()
|
||||
|
||||
return True, "", tuple(dict.fromkeys(str(_normalize_addr(addr)) for addr in addrs))
|
||||
async def async_resolve_url_target(
|
||||
url: str,
|
||||
*,
|
||||
allow_loopback: bool = False,
|
||||
trust_remote_dns: bool = False,
|
||||
timeout_s: float = _DNS_RESOLUTION_TIMEOUT_SECONDS,
|
||||
) -> tuple[bool, str, tuple[str, ...]]:
|
||||
"""Resolve and validate an HTTP target without blocking the event loop."""
|
||||
hostname, error = _parse_url_hostname(url)
|
||||
if hostname is None:
|
||||
return False, error or "Missing hostname", ()
|
||||
loop = asyncio.get_running_loop()
|
||||
try:
|
||||
infos = await asyncio.wait_for(
|
||||
loop.getaddrinfo(
|
||||
hostname,
|
||||
None,
|
||||
family=socket.AF_UNSPEC,
|
||||
type=socket.SOCK_STREAM,
|
||||
),
|
||||
timeout=timeout_s,
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
return False, f"Timed out resolving hostname: {hostname}", ()
|
||||
except socket.gaierror:
|
||||
return _unresolved_target_result(hostname, trust_remote_dns=trust_remote_dns)
|
||||
return _resolved_target_result(hostname, infos, allow_loopback=allow_loopback)
|
||||
|
||||
|
||||
def validate_url_target(url: str, *, allow_loopback: bool = False) -> tuple[bool, str]:
|
||||
@@ -151,6 +200,16 @@ def validate_url_target(url: str, *, allow_loopback: bool = False) -> tuple[bool
|
||||
return ok, error
|
||||
|
||||
|
||||
async def async_validate_url_target(
|
||||
url: str,
|
||||
*,
|
||||
allow_loopback: bool = False,
|
||||
) -> tuple[bool, str]:
|
||||
"""Validate a URL using the event loop's asynchronous resolver."""
|
||||
ok, error, _ = await async_resolve_url_target(url, allow_loopback=allow_loopback)
|
||||
return ok, error
|
||||
|
||||
|
||||
def env_proxy_applies_to_url(url: str) -> bool:
|
||||
"""Return True when process proxy settings would proxy this URL."""
|
||||
try:
|
||||
@@ -277,7 +336,10 @@ class PinnedDNSAsyncTransport(httpx.AsyncBaseTransport):
|
||||
|
||||
async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
|
||||
url = str(request.url)
|
||||
ok, error, resolved_ips = resolve_url_target(url, allow_loopback=self._allow_loopback)
|
||||
ok, error, resolved_ips = await async_resolve_url_target(
|
||||
url,
|
||||
allow_loopback=self._allow_loopback,
|
||||
)
|
||||
if not ok:
|
||||
raise UnsafeURLRequestError(error, request=request)
|
||||
async with self._resolver_lock:
|
||||
@@ -320,6 +382,46 @@ def validate_resolved_url(url: str) -> tuple[bool, str]:
|
||||
return True, ""
|
||||
|
||||
|
||||
async def async_validate_resolved_url(url: str) -> tuple[bool, str]:
|
||||
"""Validate a redirect target without blocking on domain resolution."""
|
||||
try:
|
||||
parsed = urlparse(url)
|
||||
except Exception:
|
||||
return True, ""
|
||||
hostname = parsed.hostname
|
||||
if not hostname:
|
||||
return True, ""
|
||||
try:
|
||||
addr = ipaddress.ip_address(hostname)
|
||||
except ValueError:
|
||||
loop = asyncio.get_running_loop()
|
||||
try:
|
||||
infos = await asyncio.wait_for(
|
||||
loop.getaddrinfo(
|
||||
hostname,
|
||||
None,
|
||||
family=socket.AF_UNSPEC,
|
||||
type=socket.SOCK_STREAM,
|
||||
),
|
||||
timeout=_DNS_RESOLUTION_TIMEOUT_SECONDS,
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
return False, f"Timed out resolving redirect hostname: {hostname}"
|
||||
except socket.gaierror:
|
||||
return True, ""
|
||||
for info in infos:
|
||||
try:
|
||||
addr = ipaddress.ip_address(info[4][0])
|
||||
except (IndexError, TypeError, ValueError):
|
||||
continue
|
||||
if _is_private(addr):
|
||||
return False, f"Redirect target {hostname} resolves to private address {addr}"
|
||||
return True, ""
|
||||
if _is_private(addr):
|
||||
return False, f"Redirect target is a private address: {addr}"
|
||||
return True, ""
|
||||
|
||||
|
||||
def contains_internal_url(command: str, *, allow_loopback: bool = False) -> bool:
|
||||
"""Return True if the command string contains a URL targeting an internal/private address."""
|
||||
for m in _URL_RE.finditer(command):
|
||||
|
||||
@@ -0,0 +1,29 @@
|
||||
"""Compatibility bridge for asynchronous SessionManager operations."""
|
||||
|
||||
import asyncio
|
||||
import inspect
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import Any, TypeVar, cast
|
||||
|
||||
from nanobot.utils.cancellation import shield_and_drain
|
||||
|
||||
_SessionResult = TypeVar("_SessionResult")
|
||||
|
||||
|
||||
async def call_session_manager(
|
||||
manager: object,
|
||||
async_method_name: str,
|
||||
sync_method: Callable[..., _SessionResult],
|
||||
/,
|
||||
*args: Any,
|
||||
**kwargs: Any,
|
||||
) -> _SessionResult:
|
||||
"""Prefer a class-declared coroutine, or offload the established sync contract."""
|
||||
class_async_method = inspect.getattr_static(type(manager), async_method_name, None)
|
||||
if inspect.iscoroutinefunction(class_async_method):
|
||||
async_method = cast(
|
||||
Callable[..., Awaitable[_SessionResult]],
|
||||
getattr(manager, async_method_name),
|
||||
)
|
||||
return await async_method(*args, **kwargs)
|
||||
return await shield_and_drain(asyncio.to_thread(sync_method, *args, **kwargs))
|
||||
+113
-1
@@ -1,5 +1,6 @@
|
||||
"""Session management for conversation history."""
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import errno
|
||||
import hashlib
|
||||
@@ -28,6 +29,7 @@ from nanobot.runtime_context import (
|
||||
public_history_message,
|
||||
)
|
||||
from nanobot.session.model_selection import SESSION_MODEL_PRESET_METADATA_KEY
|
||||
from nanobot.utils.cancellation import shield_and_drain
|
||||
from nanobot.utils.helpers import (
|
||||
content_with_media_breadcrumbs,
|
||||
ensure_dir,
|
||||
@@ -71,6 +73,7 @@ _WORKSPACE_STATE_DIR = ".nanobot"
|
||||
_WORKSPACE_ID_FILE = "workspace-id"
|
||||
_WORKSPACE_ID_RE = re.compile(r"^[0-9a-f]{32}$")
|
||||
_SESSION_MIGRATION_LOCK_TIMEOUT_SECONDS = 30
|
||||
_SESSION_FILES_LOCK_TIMEOUT_SECONDS = 5
|
||||
_SESSION_FILES_LOCK_FILENAME = ".session-files.lock"
|
||||
_COPY_CHUNK_SIZE = 1024 * 1024
|
||||
|
||||
@@ -560,7 +563,8 @@ class JsonlSessionStore:
|
||||
self.sessions_dir = ensure_dir(root / workspace_id)
|
||||
self.legacy_sessions_dir = get_legacy_sessions_dir()
|
||||
self._session_files_lock = FileLock(
|
||||
str(self.sessions_dir / _SESSION_FILES_LOCK_FILENAME)
|
||||
str(self.sessions_dir / _SESSION_FILES_LOCK_FILENAME),
|
||||
timeout=_SESSION_FILES_LOCK_TIMEOUT_SECONDS,
|
||||
)
|
||||
with self._session_files_lock:
|
||||
self._migrate_from_workspace(canonical_workspace)
|
||||
@@ -1642,6 +1646,7 @@ class SessionManager:
|
||||
self._cache: OrderedDict[str, Session] = OrderedDict()
|
||||
# Preserve identity for sessions held by active callers without retaining idle ones.
|
||||
self._overflow_cache: WeakValueDictionary[str, Session] = WeakValueDictionary()
|
||||
self._async_session_locks: WeakValueDictionary[str, asyncio.Lock] = WeakValueDictionary()
|
||||
self._max_cached_sessions = SESSION_CACHE_MAX_SIZE
|
||||
self._delete_observer: Callable[[str], None] | None = None
|
||||
|
||||
@@ -1741,6 +1746,28 @@ class SessionManager:
|
||||
self._remember(session)
|
||||
return session
|
||||
|
||||
def _async_session_lock(self, key: str) -> asyncio.Lock:
|
||||
lock = self._async_session_locks.get(key)
|
||||
if lock is None:
|
||||
lock = asyncio.Lock()
|
||||
self._async_session_locks[key] = lock
|
||||
return lock
|
||||
|
||||
async def get_or_create_async(self, key: str) -> Session:
|
||||
"""Load a session without running file I/O or lock waits on the event loop."""
|
||||
cached = self.get_cached(key)
|
||||
if cached is not None:
|
||||
return cached
|
||||
async with self._async_session_lock(key):
|
||||
cached = self.get_cached(key)
|
||||
if cached is not None:
|
||||
return cached
|
||||
session = await asyncio.to_thread(self._load, key)
|
||||
if session is None:
|
||||
session = Session(key=key)
|
||||
self._remember(session)
|
||||
return session
|
||||
|
||||
def get_or_create_transient(
|
||||
self,
|
||||
key: str,
|
||||
@@ -1774,6 +1801,17 @@ class SessionManager:
|
||||
self._store.save(session, fsync=fsync)
|
||||
self._remember(session)
|
||||
|
||||
async def save_async(self, session: Session, *, fsync: bool = False) -> None:
|
||||
"""Persist a session without blocking the caller's event loop."""
|
||||
if not session.policy.persist:
|
||||
return
|
||||
|
||||
async def save_and_remember() -> None:
|
||||
await asyncio.to_thread(self._store.save, session, fsync=fsync)
|
||||
self._remember(session)
|
||||
|
||||
await shield_and_drain(save_and_remember())
|
||||
|
||||
def save_runtime_checkpoint(self, session: Session) -> None:
|
||||
"""Persist volatile recovery state without rewriting long history."""
|
||||
if not session.policy.persist:
|
||||
@@ -1786,6 +1824,23 @@ class SessionManager:
|
||||
# they opt into a dedicated checkpoint primitive.
|
||||
self.save(session)
|
||||
|
||||
async def save_runtime_checkpoint_async(self, session: Session) -> None:
|
||||
"""Persist an in-flight checkpoint without blocking the event loop."""
|
||||
if not session.policy.persist:
|
||||
return
|
||||
|
||||
async def save_and_remember() -> None:
|
||||
if self._store is self._jsonl_store:
|
||||
await asyncio.to_thread(
|
||||
self._jsonl_store.save_runtime_checkpoint,
|
||||
session,
|
||||
)
|
||||
else:
|
||||
await asyncio.to_thread(self._store.save, session)
|
||||
self._remember(session)
|
||||
|
||||
await shield_and_drain(save_and_remember())
|
||||
|
||||
def rename_model_preset(self, old_name: str, new_name: str) -> int:
|
||||
"""Rename a session-scoped model preset across durable and live sessions."""
|
||||
if old_name == new_name:
|
||||
@@ -1827,6 +1882,21 @@ class SessionManager:
|
||||
raise
|
||||
return len(changed)
|
||||
|
||||
async def flush_all_async(self) -> int:
|
||||
"""Re-save every cached session without blocking the event loop."""
|
||||
cached = dict(self._overflow_cache.items())
|
||||
cached.update(self._cache)
|
||||
flushed = 0
|
||||
for key, session in cached.items():
|
||||
try:
|
||||
await shield_and_drain(
|
||||
asyncio.to_thread(self._store.save, session, fsync=True)
|
||||
)
|
||||
flushed += 1
|
||||
except Exception:
|
||||
logger.warning("Failed to flush session {}", key, exc_info=True)
|
||||
return flushed
|
||||
|
||||
def flush_all(self) -> int:
|
||||
"""Re-save every cached session with fsync for durable shutdown.
|
||||
|
||||
@@ -1858,6 +1928,18 @@ class SessionManager:
|
||||
self._delete_observer(key)
|
||||
return deleted
|
||||
|
||||
async def delete_session_async(self, key: str) -> bool:
|
||||
"""Delete a session without blocking the event loop."""
|
||||
|
||||
async def delete_and_notify() -> bool:
|
||||
self.invalidate(key)
|
||||
deleted = await asyncio.to_thread(self._store.delete, key)
|
||||
if self._delete_observer is not None:
|
||||
self._delete_observer(key)
|
||||
return deleted
|
||||
|
||||
return await shield_and_drain(delete_and_notify())
|
||||
|
||||
def restore_sessions_to_workspace(self) -> SessionRestoreResult:
|
||||
"""Restore session files to the pre-relocation path for an explicit rollback."""
|
||||
return self._jsonl_store.restore_to_workspace()
|
||||
@@ -1930,6 +2012,10 @@ class SessionManager:
|
||||
"""Read session metadata without loading the transcript."""
|
||||
return cast(dict[str, Any] | None, self._store.read_metadata(key))
|
||||
|
||||
async def read_session_metadata_async(self, key: str) -> dict[str, Any] | None:
|
||||
"""Read session metadata without blocking the event loop."""
|
||||
return await asyncio.to_thread(self.read_session_metadata, key)
|
||||
|
||||
def update_session_metadata(
|
||||
self,
|
||||
key: str,
|
||||
@@ -1943,5 +2029,31 @@ class SessionManager:
|
||||
session.metadata.update(deepcopy(updates))
|
||||
return updated
|
||||
|
||||
async def update_session_metadata_async(
|
||||
self,
|
||||
key: str,
|
||||
updates: dict[str, Any],
|
||||
*,
|
||||
fsync: bool = False,
|
||||
) -> bool:
|
||||
"""Update metadata without blocking the event loop."""
|
||||
|
||||
async def update_and_refresh_cache() -> bool:
|
||||
updated = await asyncio.to_thread(
|
||||
self._store.update_metadata,
|
||||
key,
|
||||
updates,
|
||||
fsync=fsync,
|
||||
)
|
||||
if updated and (session := self.get_cached(key)) is not None:
|
||||
session.metadata.update(deepcopy(updates))
|
||||
return updated
|
||||
|
||||
return await shield_and_drain(update_and_refresh_cache())
|
||||
|
||||
def list_sessions(self) -> list[dict[str, Any]]:
|
||||
return cast(list[dict[str, Any]], self._store.list_sessions())
|
||||
|
||||
async def list_sessions_async(self) -> list[dict[str, Any]]:
|
||||
"""List persisted sessions without blocking the event loop."""
|
||||
return await asyncio.to_thread(self.list_sessions)
|
||||
|
||||
+55
-23
@@ -25,6 +25,7 @@ from nanobot.bus.outbound_events import (
|
||||
)
|
||||
from nanobot.bus.queue import MessageBus
|
||||
from nanobot.session import turn_continuation
|
||||
from nanobot.session.async_compat import call_session_manager
|
||||
from nanobot.session.keys import UNIFIED_SESSION_KEY, last_channel_from_metadata
|
||||
from nanobot.session.manager import Session, SessionManager
|
||||
from nanobot.webui.metadata import WEBUI_TURN_METADATA_KEY
|
||||
@@ -460,6 +461,37 @@ class RecoveryCoordinator:
|
||||
repr=False,
|
||||
)
|
||||
|
||||
async def _get_or_create_session(self, key: str) -> Session:
|
||||
return await call_session_manager(
|
||||
self.sessions,
|
||||
"get_or_create_async",
|
||||
self.sessions.get_or_create,
|
||||
key,
|
||||
)
|
||||
|
||||
async def _save_session(self, session: Session) -> None:
|
||||
await call_session_manager(
|
||||
self.sessions,
|
||||
"save_async",
|
||||
self.sessions.save,
|
||||
session,
|
||||
)
|
||||
|
||||
async def _read_session_metadata(self, key: str) -> dict[str, Any] | None:
|
||||
return await call_session_manager(
|
||||
self.sessions,
|
||||
"read_session_metadata_async",
|
||||
self.sessions.read_session_metadata,
|
||||
key,
|
||||
)
|
||||
|
||||
async def _list_sessions(self) -> list[dict[str, Any]]:
|
||||
return await call_session_manager(
|
||||
self.sessions,
|
||||
"list_sessions_async",
|
||||
self.sessions.list_sessions,
|
||||
)
|
||||
|
||||
def register_recovery_task(self, session_key: str, task: asyncio.Task[Any]) -> None:
|
||||
"""Track the task that owns an explicit recovery continuation."""
|
||||
self._active_recovery_tasks[session_key] = task
|
||||
@@ -482,8 +514,8 @@ class RecoveryCoordinator:
|
||||
|
||||
async def scan(self) -> None:
|
||||
"""Recover every interrupted WebUI session once at gateway startup."""
|
||||
for key in self._recovery_candidates():
|
||||
metadata_payload = self.sessions.read_session_metadata(key)
|
||||
for key in await self._recovery_candidates():
|
||||
metadata_payload = await self._read_session_metadata(key)
|
||||
raw_metadata = metadata_payload.get("metadata") if metadata_payload else None
|
||||
metadata = cast(dict[str, Any], raw_metadata) if isinstance(raw_metadata, dict) else {}
|
||||
route = self._websocket_route_for(key, metadata)
|
||||
@@ -492,7 +524,7 @@ class RecoveryCoordinator:
|
||||
unfinished = self._has_unfinished_webui_transcript(key)
|
||||
if not self._needs_recovery(metadata) and not unfinished:
|
||||
continue
|
||||
session = self.sessions.get_or_create(key)
|
||||
session = await self._get_or_create_session(key)
|
||||
try:
|
||||
await self._recover_session(session, route[1])
|
||||
await self._requeue_pending_followups(session)
|
||||
@@ -507,14 +539,14 @@ class RecoveryCoordinator:
|
||||
reason="recovery_failed",
|
||||
can_continue=False,
|
||||
)
|
||||
self.sessions.save(session)
|
||||
await self._save_session(session)
|
||||
await self._publish(route[1], failed)
|
||||
|
||||
def _recovery_candidates(self) -> list[str]:
|
||||
async def _recovery_candidates(self) -> list[str]:
|
||||
"""Discover canonical and transcript-only WebUI sessions cheaply."""
|
||||
candidates = dict.fromkeys(
|
||||
key
|
||||
for item in self.sessions.list_sessions()
|
||||
for item in await self._list_sessions()
|
||||
if isinstance((key := item.get("key")), str)
|
||||
)
|
||||
try:
|
||||
@@ -523,7 +555,7 @@ class RecoveryCoordinator:
|
||||
# duplicating its filename and migration rules here would drift.
|
||||
from nanobot.webui.session_list_index import list_webui_sessions
|
||||
|
||||
for item in list_webui_sessions(self.sessions):
|
||||
for item in await asyncio.to_thread(list_webui_sessions, self.sessions):
|
||||
key = item.get("key")
|
||||
if isinstance(key, str):
|
||||
candidates.setdefault(key, None)
|
||||
@@ -549,7 +581,7 @@ class RecoveryCoordinator:
|
||||
"""Reject stale queued recoveries and let new user input supersede them."""
|
||||
recovery_id = message.metadata.get(RECOVERY_INBOUND_METADATA_KEY)
|
||||
if isinstance(recovery_id, str):
|
||||
session = self.sessions.get_or_create(message.session_key)
|
||||
session = await self._get_or_create_session(message.session_key)
|
||||
state = recovery_state_from_metadata(session.metadata)
|
||||
return bool(
|
||||
state
|
||||
@@ -558,7 +590,7 @@ class RecoveryCoordinator:
|
||||
)
|
||||
if message.channel != "websocket":
|
||||
return True
|
||||
session = self.sessions.get_or_create(message.session_key)
|
||||
session = await self._get_or_create_session(message.session_key)
|
||||
state = recovery_state_from_metadata(session.metadata)
|
||||
if state and state["status"] in {"resuming", "awaiting_user", "failed"}:
|
||||
await self._cancel_active_recovery(message.session_key)
|
||||
@@ -572,13 +604,13 @@ class RecoveryCoordinator:
|
||||
attempts=cast(int, state.get("attempts", 0)),
|
||||
reason="superseded",
|
||||
)
|
||||
self.sessions.save(session)
|
||||
await self._save_session(session)
|
||||
await self._publish(message.chat_id, recovered)
|
||||
return True
|
||||
|
||||
async def turn_completed(self, session_key: str) -> None:
|
||||
"""Resolve a resuming state after the recovered turn commits."""
|
||||
session = self.sessions.get_or_create(session_key)
|
||||
session = await self._get_or_create_session(session_key)
|
||||
state = recovery_state_from_metadata(session.metadata)
|
||||
if not state or state["status"] != "resuming":
|
||||
return
|
||||
@@ -592,7 +624,7 @@ class RecoveryCoordinator:
|
||||
attempts=cast(int, state.get("attempts", 0)),
|
||||
reason="continued",
|
||||
)
|
||||
self.sessions.save(session)
|
||||
await self._save_session(session)
|
||||
await self._publish(route[1], recovered)
|
||||
|
||||
async def handle_action(self, action: str, payload: dict[str, Any]) -> dict[str, Any]:
|
||||
@@ -603,7 +635,7 @@ class RecoveryCoordinator:
|
||||
raise RecoveryActionError("missing chat_id")
|
||||
if not isinstance(recovery_id, str) or not recovery_id:
|
||||
raise RecoveryActionError("missing recovery_id")
|
||||
session = self.sessions.get_or_create(self._session_key(chat_id))
|
||||
session = await self._get_or_create_session(self._session_key(chat_id))
|
||||
state = recovery_state_from_metadata(session.metadata)
|
||||
if not state or state["recovery_id"] != recovery_id:
|
||||
raise RecoveryActionError("recovery state is stale", status=409)
|
||||
@@ -618,7 +650,7 @@ class RecoveryCoordinator:
|
||||
attempts=cast(int, state.get("attempts", 0)),
|
||||
reason="dismissed",
|
||||
)
|
||||
self.sessions.save(session)
|
||||
await self._save_session(session)
|
||||
await self._publish(chat_id, next_state)
|
||||
return next_state
|
||||
if action != "continue":
|
||||
@@ -635,7 +667,7 @@ class RecoveryCoordinator:
|
||||
reason="user_confirmed",
|
||||
resume_message_count=len(session.messages),
|
||||
)
|
||||
self.sessions.save(session)
|
||||
await self._save_session(session)
|
||||
await self._publish(chat_id, next_state)
|
||||
await self._queue_continuation(session, chat_id, next_state)
|
||||
return next_state
|
||||
@@ -668,7 +700,7 @@ class RecoveryCoordinator:
|
||||
attempts=cast(int, state.get("attempts", 1)),
|
||||
reason="loop_guard",
|
||||
)
|
||||
self.sessions.save(session)
|
||||
await self._save_session(session)
|
||||
await self._publish(chat_id, next_state)
|
||||
elif self._has_unfinished_webui_transcript(session.key):
|
||||
# A normal last-client shutdown can materialize the checkpoint
|
||||
@@ -690,7 +722,7 @@ class RecoveryCoordinator:
|
||||
),
|
||||
can_continue=can_continue,
|
||||
)
|
||||
self.sessions.save(session)
|
||||
await self._save_session(session)
|
||||
await self._publish(chat_id, waiting)
|
||||
return
|
||||
if state and state["status"] in {"awaiting_user", "failed"}:
|
||||
@@ -706,7 +738,7 @@ class RecoveryCoordinator:
|
||||
attempts=cast(int, state.get("attempts", 1)),
|
||||
reason="loop_guard",
|
||||
)
|
||||
self.sessions.save(session)
|
||||
await self._save_session(session)
|
||||
await self._publish(chat_id, waiting)
|
||||
return
|
||||
|
||||
@@ -724,7 +756,7 @@ class RecoveryCoordinator:
|
||||
reason="checkpoint_unknown",
|
||||
can_continue=False,
|
||||
)
|
||||
self.sessions.save(session)
|
||||
await self._save_session(session)
|
||||
await self._publish(chat_id, waiting)
|
||||
return
|
||||
if checkpoint is not None and not _runtime_checkpoint_is_well_formed(checkpoint):
|
||||
@@ -738,7 +770,7 @@ class RecoveryCoordinator:
|
||||
reason="checkpoint_invalid",
|
||||
can_continue=False,
|
||||
)
|
||||
self.sessions.save(session)
|
||||
await self._save_session(session)
|
||||
await self._publish(chat_id, waiting)
|
||||
return
|
||||
if phase == "final_response":
|
||||
@@ -750,7 +782,7 @@ class RecoveryCoordinator:
|
||||
attempts=0,
|
||||
reason="answer_restored",
|
||||
)
|
||||
self.sessions.save(session)
|
||||
await self._save_session(session)
|
||||
await self._publish(chat_id, recovered)
|
||||
return
|
||||
if phase in _UNCERTAIN_TOOL_PHASES or pending_calls:
|
||||
@@ -762,7 +794,7 @@ class RecoveryCoordinator:
|
||||
attempts=0,
|
||||
reason="tool_state_unknown",
|
||||
)
|
||||
self.sessions.save(session)
|
||||
await self._save_session(session)
|
||||
await self._publish(chat_id, waiting)
|
||||
return
|
||||
# A gateway restart is a lifecycle boundary. Never enqueue model work
|
||||
@@ -777,7 +809,7 @@ class RecoveryCoordinator:
|
||||
attempts=0,
|
||||
reason="restart_requires_confirmation",
|
||||
)
|
||||
self.sessions.save(session)
|
||||
await self._save_session(session)
|
||||
await self._publish(chat_id, waiting)
|
||||
|
||||
async def _queue_continuation(
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import re
|
||||
import time
|
||||
from collections.abc import Awaitable, Callable
|
||||
@@ -176,7 +177,7 @@ async def maybe_generate_webui_title(
|
||||
model: str,
|
||||
) -> bool:
|
||||
"""Generate and persist a short title for WebUI-owned sessions only."""
|
||||
session = sessions.get_or_create(session_key)
|
||||
session = await sessions.get_or_create_async(session_key)
|
||||
if session.metadata.get(WEBUI_SESSION_METADATA_KEY) is not True:
|
||||
return False
|
||||
if session.metadata.get(WEBUI_TITLE_USER_EDITED_METADATA_KEY) is True:
|
||||
@@ -187,7 +188,7 @@ async def maybe_generate_webui_title(
|
||||
if cleaned_current_title:
|
||||
if cleaned_current_title != current_title:
|
||||
session.metadata[WEBUI_TITLE_METADATA_KEY] = cleaned_current_title
|
||||
sessions.save(session)
|
||||
await sessions.save_async(session)
|
||||
return False
|
||||
session.metadata.pop(WEBUI_TITLE_METADATA_KEY, None)
|
||||
|
||||
@@ -241,7 +242,7 @@ async def maybe_generate_webui_title(
|
||||
)
|
||||
return False
|
||||
session.metadata[WEBUI_TITLE_METADATA_KEY] = title
|
||||
sessions.save(session)
|
||||
await sessions.save_async(session)
|
||||
return True
|
||||
|
||||
|
||||
@@ -437,8 +438,8 @@ class WebuiTurnRoutePolicy:
|
||||
)
|
||||
and route.channel == "websocket"
|
||||
):
|
||||
session = self.sessions.get_or_create(session_key)
|
||||
if session.metadata.get(WEBUI_SESSION_METADATA_KEY) is True:
|
||||
session = self.sessions.get_cached(session_key)
|
||||
if session is not None and session.metadata.get(WEBUI_SESSION_METADATA_KEY) is True:
|
||||
metadata = dict(route.metadata)
|
||||
turn_prefix = "session-input" if internal_user_input else "subagent"
|
||||
metadata.update({
|
||||
@@ -580,7 +581,7 @@ class WebuiTurnCoordinator:
|
||||
or not session_key.startswith("websocket:")
|
||||
):
|
||||
return
|
||||
persisted = self.sessions.read_session_metadata(session_key)
|
||||
persisted = await self.sessions.read_session_metadata_async(session_key)
|
||||
metadata_value: object = persisted.get("metadata") if persisted is not None else None
|
||||
metadata = (
|
||||
cast(dict[str, Any], metadata_value)
|
||||
@@ -591,7 +592,8 @@ class WebuiTurnCoordinator:
|
||||
return
|
||||
public_metadata = _session_message_public_metadata(envelope)
|
||||
try:
|
||||
append_session_message_input(
|
||||
await asyncio.to_thread(
|
||||
append_session_message_input,
|
||||
session_key,
|
||||
content=event.content,
|
||||
created_at_ms=envelope["created_at_ms"],
|
||||
@@ -616,8 +618,9 @@ class WebuiTurnCoordinator:
|
||||
def _handle_session_turn_started(self, event: SessionTurnStarted) -> None:
|
||||
if not self._is_websocket_event(event.context):
|
||||
return
|
||||
session = self.sessions.get_or_create(event.context.session_key)
|
||||
mark_webui_session(session, event.context.metadata)
|
||||
session = self.sessions.get_cached(event.context.session_key)
|
||||
if session is not None:
|
||||
mark_webui_session(session, event.context.metadata)
|
||||
|
||||
async def _handle_run_status_changed(self, event: TurnRunStatusChanged) -> None:
|
||||
if not self._is_websocket_event(event.context):
|
||||
@@ -703,7 +706,7 @@ class WebuiTurnCoordinator:
|
||||
if msg.channel != "websocket":
|
||||
return
|
||||
|
||||
session = self.sessions.get_or_create(session_key)
|
||||
session = await self.sessions.get_or_create_async(session_key)
|
||||
await self.bus.publish_outbound(
|
||||
outbound_message_for_event(
|
||||
channel=msg.channel,
|
||||
|
||||
@@ -5,17 +5,24 @@ from __future__ import annotations
|
||||
import asyncio
|
||||
import uuid
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import Any
|
||||
from contextlib import suppress
|
||||
from typing import Any, TypeVar
|
||||
|
||||
from loguru import logger
|
||||
|
||||
from nanobot.agent.automation_turns import AutomationTurnError
|
||||
from nanobot.agent.automation_turns import (
|
||||
AutomationTurnAcceptedCancellation,
|
||||
AutomationTurnError,
|
||||
)
|
||||
from nanobot.bus.events import InboundMessage, OutboundMessage
|
||||
from nanobot.triggers.local_session_turns import LOCAL_TRIGGER_META
|
||||
from nanobot.triggers.local_store import LocalTriggerStore
|
||||
from nanobot.triggers.local_types import LocalTrigger, TriggerDelivery
|
||||
from nanobot.utils.cancellation import shield_and_drain
|
||||
from nanobot.webui.metadata import WEBUI_MESSAGE_SOURCE_METADATA_KEY, WEBUI_TURN_METADATA_KEY
|
||||
|
||||
_T = TypeVar("_T")
|
||||
|
||||
|
||||
async def run_local_trigger_queue(
|
||||
*,
|
||||
@@ -29,14 +36,16 @@ async def run_local_trigger_queue(
|
||||
if submit_turn is None:
|
||||
raise ValueError("run_local_trigger_queue requires submit_turn")
|
||||
logger.info("Local trigger queue started")
|
||||
recovered = store.recover_processing_deliveries()
|
||||
recovered = await shield_and_drain(asyncio.to_thread(store.recover_processing_deliveries))
|
||||
if recovered:
|
||||
logger.warning(
|
||||
"Trigger: recovered {} interrupted delivery file(s) from processing",
|
||||
recovered,
|
||||
)
|
||||
while True:
|
||||
deliveries = store.claim_deliveries(limit=batch_size)
|
||||
deliveries = await shield_and_drain(
|
||||
asyncio.to_thread(store.claim_deliveries, limit=batch_size)
|
||||
)
|
||||
if not deliveries:
|
||||
await asyncio.sleep(poll_interval_s)
|
||||
continue
|
||||
@@ -49,30 +58,24 @@ async def run_local_trigger_queue(
|
||||
submit_turn=submit_turn,
|
||||
is_channel_enabled=is_channel_enabled,
|
||||
)
|
||||
store.complete_delivery(delivery)
|
||||
except _DeliverySettledOnCancellation:
|
||||
raise
|
||||
except asyncio.CancelledError as exc:
|
||||
store.retry_delivery(delivery, str(exc) or exc.__class__.__name__)
|
||||
_write_delivery_run_record(
|
||||
error = str(exc) or exc.__class__.__name__
|
||||
await shield_and_drain(asyncio.to_thread(store.retry_delivery, delivery, error))
|
||||
await _write_delivery_run_record(
|
||||
store,
|
||||
delivery,
|
||||
status="interrupted",
|
||||
error=str(exc) or exc.__class__.__name__,
|
||||
error=error,
|
||||
)
|
||||
raise
|
||||
except _TerminalDeliveryError as exc:
|
||||
store.record_delivery(
|
||||
delivery.trigger_id,
|
||||
status="error",
|
||||
error=str(exc),
|
||||
run_at_ms=delivery.created_at_ms,
|
||||
await _await_delivery_settlement(
|
||||
_settle_failed_delivery(store, delivery, error=str(exc)),
|
||||
store=store,
|
||||
delivery=delivery,
|
||||
)
|
||||
_write_delivery_run_record(
|
||||
store,
|
||||
delivery,
|
||||
status="error",
|
||||
error=str(exc),
|
||||
)
|
||||
store.complete_delivery(delivery)
|
||||
logger.warning(
|
||||
"Trigger: dropped delivery {} for {}: {}",
|
||||
delivery.id,
|
||||
@@ -81,19 +84,11 @@ async def run_local_trigger_queue(
|
||||
)
|
||||
except AutomationTurnError as exc:
|
||||
error = str(exc) or exc.__class__.__name__
|
||||
store.record_delivery(
|
||||
delivery.trigger_id,
|
||||
status="error",
|
||||
error=error,
|
||||
run_at_ms=delivery.created_at_ms,
|
||||
await _await_delivery_settlement(
|
||||
_settle_failed_delivery(store, delivery, error=error),
|
||||
store=store,
|
||||
delivery=delivery,
|
||||
)
|
||||
_write_delivery_run_record(
|
||||
store,
|
||||
delivery,
|
||||
status="error",
|
||||
error=error,
|
||||
)
|
||||
store.complete_delivery(delivery)
|
||||
logger.warning(
|
||||
"Trigger: delivery {} for {} reached the agent but failed: {}",
|
||||
delivery.id,
|
||||
@@ -102,18 +97,10 @@ async def run_local_trigger_queue(
|
||||
)
|
||||
except Exception as exc:
|
||||
error = str(exc) or exc.__class__.__name__
|
||||
retried = store.retry_delivery(delivery, error)
|
||||
_write_delivery_run_record(
|
||||
store,
|
||||
delivery,
|
||||
status="retrying" if retried else "error",
|
||||
error=error,
|
||||
)
|
||||
store.record_delivery(
|
||||
delivery.trigger_id,
|
||||
status="error",
|
||||
error=error,
|
||||
run_at_ms=delivery.created_at_ms,
|
||||
retried = await _await_delivery_settlement(
|
||||
_settle_retryable_delivery(store, delivery, error=error),
|
||||
store=store,
|
||||
delivery=delivery,
|
||||
)
|
||||
logger.exception(
|
||||
"Trigger: failed delivery {} for {}{}",
|
||||
@@ -127,6 +114,10 @@ class _TerminalDeliveryError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
class _DeliverySettledOnCancellation(asyncio.CancelledError):
|
||||
"""Cancellation reported only after an already-submitted delivery is settled."""
|
||||
|
||||
|
||||
async def _deliver_delivery(
|
||||
store: LocalTriggerStore,
|
||||
delivery: TriggerDelivery,
|
||||
@@ -134,7 +125,7 @@ async def _deliver_delivery(
|
||||
submit_turn: Callable[[InboundMessage], Awaitable[OutboundMessage | None]],
|
||||
is_channel_enabled: Callable[[str], bool],
|
||||
) -> None:
|
||||
trigger = store.get(delivery.trigger_id)
|
||||
trigger = await asyncio.to_thread(store.get, delivery.trigger_id)
|
||||
if trigger is None:
|
||||
raise _TerminalDeliveryError("trigger not found")
|
||||
if not trigger.enabled:
|
||||
@@ -142,7 +133,14 @@ async def _deliver_delivery(
|
||||
if not is_channel_enabled(trigger.channel):
|
||||
raise _TerminalDeliveryError(f"target channel is not enabled: {trigger.channel}")
|
||||
|
||||
store.write_delivery_run_record(delivery, trigger=trigger, status="processing")
|
||||
await shield_and_drain(
|
||||
asyncio.to_thread(
|
||||
store.write_delivery_run_record,
|
||||
delivery,
|
||||
trigger=trigger,
|
||||
status="processing",
|
||||
)
|
||||
)
|
||||
msg = InboundMessage(
|
||||
channel=trigger.channel,
|
||||
sender_id=trigger.sender_id,
|
||||
@@ -151,22 +149,177 @@ async def _deliver_delivery(
|
||||
metadata=_delivery_metadata(trigger, delivery),
|
||||
session_key_override=trigger.session_key,
|
||||
)
|
||||
response = await submit_turn(msg)
|
||||
store.record_delivery(
|
||||
trigger.id,
|
||||
status="ok",
|
||||
run_at_ms=delivery.created_at_ms,
|
||||
try:
|
||||
response = await submit_turn(msg)
|
||||
except AutomationTurnAcceptedCancellation:
|
||||
try:
|
||||
await _await_delivery_settlement(
|
||||
_settle_accepted_delivery(store, delivery, trigger=trigger),
|
||||
store=store,
|
||||
delivery=delivery,
|
||||
)
|
||||
except _DeliverySettledOnCancellation:
|
||||
raise
|
||||
except Exception:
|
||||
logger.exception(
|
||||
"Trigger: failed to persist accepted delivery {}; dropping retry",
|
||||
delivery.id,
|
||||
)
|
||||
with suppress(Exception):
|
||||
await shield_and_drain(asyncio.to_thread(store.complete_delivery, delivery))
|
||||
raise _DeliverySettledOnCancellation from None
|
||||
|
||||
try:
|
||||
await _await_delivery_settlement(
|
||||
_settle_submitted_delivery(store, delivery, trigger=trigger, response=response),
|
||||
store=store,
|
||||
delivery=delivery,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception(
|
||||
"Trigger: failed to persist status for submitted delivery {}; dropping retry",
|
||||
delivery.id,
|
||||
)
|
||||
with suppress(Exception):
|
||||
await shield_and_drain(asyncio.to_thread(store.complete_delivery, delivery))
|
||||
|
||||
|
||||
async def _await_delivery_settlement(
|
||||
operation: Awaitable[_T],
|
||||
*,
|
||||
store: LocalTriggerStore,
|
||||
delivery: TriggerDelivery,
|
||||
) -> _T:
|
||||
settlement = asyncio.ensure_future(operation)
|
||||
try:
|
||||
return await asyncio.shield(settlement)
|
||||
except asyncio.CancelledError:
|
||||
while not settlement.done():
|
||||
try:
|
||||
await asyncio.shield(settlement)
|
||||
except asyncio.CancelledError:
|
||||
continue
|
||||
try:
|
||||
settlement.result()
|
||||
except Exception:
|
||||
logger.exception(
|
||||
"Trigger: failed to settle delivery {} during cancellation",
|
||||
delivery.id,
|
||||
)
|
||||
completion = asyncio.create_task(
|
||||
shield_and_drain(asyncio.to_thread(store.complete_delivery, delivery))
|
||||
)
|
||||
while not completion.done():
|
||||
try:
|
||||
await asyncio.shield(completion)
|
||||
except asyncio.CancelledError:
|
||||
continue
|
||||
with suppress(Exception):
|
||||
completion.result()
|
||||
raise _DeliverySettledOnCancellation from None
|
||||
|
||||
|
||||
async def _settle_failed_delivery(
|
||||
store: LocalTriggerStore,
|
||||
delivery: TriggerDelivery,
|
||||
*,
|
||||
error: str,
|
||||
) -> None:
|
||||
await _write_delivery_run_record(
|
||||
store,
|
||||
delivery,
|
||||
status="error",
|
||||
error=error,
|
||||
)
|
||||
_write_delivery_run_record(
|
||||
await shield_and_drain(asyncio.to_thread(store.complete_delivery, delivery))
|
||||
# Publish the terminal status only after the durable delivery state is settled.
|
||||
await shield_and_drain(
|
||||
asyncio.to_thread(
|
||||
store.record_delivery,
|
||||
delivery.trigger_id,
|
||||
status="error",
|
||||
error=error,
|
||||
run_at_ms=delivery.created_at_ms,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
async def _settle_retryable_delivery(
|
||||
store: LocalTriggerStore,
|
||||
delivery: TriggerDelivery,
|
||||
*,
|
||||
error: str,
|
||||
) -> bool:
|
||||
retried = await shield_and_drain(asyncio.to_thread(store.retry_delivery, delivery, error))
|
||||
await _write_delivery_run_record(
|
||||
store,
|
||||
delivery,
|
||||
status="retrying" if retried else "error",
|
||||
error=error,
|
||||
)
|
||||
await shield_and_drain(
|
||||
asyncio.to_thread(
|
||||
store.record_delivery,
|
||||
delivery.trigger_id,
|
||||
status="error",
|
||||
error=error,
|
||||
run_at_ms=delivery.created_at_ms,
|
||||
)
|
||||
)
|
||||
return retried
|
||||
|
||||
|
||||
async def _settle_accepted_delivery(
|
||||
store: LocalTriggerStore,
|
||||
delivery: TriggerDelivery,
|
||||
*,
|
||||
trigger: LocalTrigger,
|
||||
) -> None:
|
||||
"""Commit an accepted delivery without claiming the agent turn completed."""
|
||||
await _write_delivery_run_record(
|
||||
store,
|
||||
delivery,
|
||||
trigger=trigger,
|
||||
status="accepted",
|
||||
)
|
||||
await shield_and_drain(asyncio.to_thread(store.complete_delivery, delivery))
|
||||
await shield_and_drain(
|
||||
asyncio.to_thread(
|
||||
store.record_delivery,
|
||||
trigger.id,
|
||||
status="ok",
|
||||
run_at_ms=delivery.created_at_ms,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
async def _settle_submitted_delivery(
|
||||
store: LocalTriggerStore,
|
||||
delivery: TriggerDelivery,
|
||||
*,
|
||||
trigger: LocalTrigger,
|
||||
response: OutboundMessage | None,
|
||||
) -> None:
|
||||
await _write_delivery_run_record(
|
||||
store,
|
||||
delivery,
|
||||
trigger=trigger,
|
||||
status="ok",
|
||||
response=response.content if response else "",
|
||||
)
|
||||
await shield_and_drain(asyncio.to_thread(store.complete_delivery, delivery))
|
||||
# last_status is the externally visible commit marker for a settled delivery.
|
||||
await shield_and_drain(
|
||||
asyncio.to_thread(
|
||||
store.record_delivery,
|
||||
trigger.id,
|
||||
status="ok",
|
||||
run_at_ms=delivery.created_at_ms,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _write_delivery_run_record(
|
||||
async def _write_delivery_run_record(
|
||||
store: LocalTriggerStore,
|
||||
delivery: TriggerDelivery,
|
||||
*,
|
||||
@@ -176,12 +329,15 @@ def _write_delivery_run_record(
|
||||
response: str | None = None,
|
||||
) -> None:
|
||||
try:
|
||||
store.write_delivery_run_record(
|
||||
delivery,
|
||||
trigger=trigger,
|
||||
status=status,
|
||||
error=error,
|
||||
response=response,
|
||||
await shield_and_drain(
|
||||
asyncio.to_thread(
|
||||
store.write_delivery_run_record,
|
||||
delivery,
|
||||
trigger=trigger,
|
||||
status=status,
|
||||
error=error,
|
||||
response=response,
|
||||
)
|
||||
)
|
||||
except Exception:
|
||||
logger.exception(
|
||||
|
||||
@@ -24,6 +24,7 @@ _MAX_RUN_HISTORY = 20
|
||||
_MAX_DELIVERY_ATTEMPTS = 10
|
||||
_RUN_RECORD_TEXT_MAX_CHARS = 4000
|
||||
_PROCESSING_RECOVERY_ERROR = "delivery was recovered from interrupted processing"
|
||||
_FILE_LOCK_TIMEOUT_SECONDS = 5
|
||||
|
||||
|
||||
class TriggerStoreError(RuntimeError):
|
||||
@@ -49,7 +50,10 @@ class LocalTriggerStore:
|
||||
self.processing_dir = self.root / "processing"
|
||||
self.failed_dir = self.root / "failed"
|
||||
self.runs_dir = self.root / "runs"
|
||||
self._lock = FileLock(str(self.root / ".lock"))
|
||||
self._lock = FileLock(
|
||||
str(self.root / ".lock"),
|
||||
timeout=_FILE_LOCK_TIMEOUT_SECONDS,
|
||||
)
|
||||
|
||||
def create(
|
||||
self,
|
||||
|
||||
@@ -3,8 +3,48 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Awaitable
|
||||
from typing import TypeVar
|
||||
|
||||
_T = TypeVar("_T")
|
||||
|
||||
|
||||
def task_is_cancelling() -> bool:
|
||||
task = asyncio.current_task()
|
||||
return task is not None and task.cancelling() > 0
|
||||
|
||||
|
||||
async def shield_and_drain(awaitable: Awaitable[_T]) -> _T:
|
||||
"""Delay caller cancellation until an accepted operation has fully settled.
|
||||
|
||||
``asyncio.to_thread`` cannot stop a worker that has already started. Shielding
|
||||
keeps cancellation from detaching that worker, and draining also lets any
|
||||
post-write in-memory settlement in ``awaitable`` finish. Cancellation is still
|
||||
re-raised as soon as the accepted operation is done.
|
||||
"""
|
||||
settlement = asyncio.ensure_future(awaitable)
|
||||
cancellation: asyncio.CancelledError | None = None
|
||||
|
||||
while not settlement.done():
|
||||
try:
|
||||
result = await asyncio.shield(settlement)
|
||||
except asyncio.CancelledError as exc:
|
||||
if cancellation is None:
|
||||
cancellation = exc
|
||||
except BaseException:
|
||||
if cancellation is None:
|
||||
raise
|
||||
break
|
||||
else:
|
||||
if cancellation is not None:
|
||||
raise cancellation
|
||||
return result
|
||||
|
||||
if cancellation is not None:
|
||||
try:
|
||||
settlement.result()
|
||||
except BaseException:
|
||||
# The caller's cancellation wins once settlement has been observed.
|
||||
pass
|
||||
raise cancellation
|
||||
return settlement.result()
|
||||
|
||||
+32
-19
@@ -2,6 +2,7 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import re
|
||||
import uuid
|
||||
from collections.abc import Mapping
|
||||
@@ -9,6 +10,7 @@ from typing import TYPE_CHECKING, Any, TypeGuard
|
||||
|
||||
from nanobot.session.manager import SessionManager
|
||||
from nanobot.session.webui_turns import WEBUI_TITLE_METADATA_KEY, clean_generated_title
|
||||
from nanobot.utils.cancellation import shield_and_drain
|
||||
from nanobot.webui.transcript import (
|
||||
append_fork_marker,
|
||||
delete_webui_transcript,
|
||||
@@ -93,24 +95,35 @@ async def handle_webui_fork_chat(
|
||||
await channel.send_webui_protocol_error(connection, "session_manager_unavailable")
|
||||
return
|
||||
|
||||
try:
|
||||
forked = create_webui_chat_fork(
|
||||
session_manager,
|
||||
source_chat_id=source_chat_id,
|
||||
before_user_index=raw_index,
|
||||
title=envelope.get("title") if isinstance(envelope.get("title"), str) else None,
|
||||
)
|
||||
if forked is None:
|
||||
await channel.send_webui_protocol_error(connection, "invalid fork source or index")
|
||||
async def create_and_attach() -> None:
|
||||
try:
|
||||
forked = await asyncio.to_thread(
|
||||
create_webui_chat_fork,
|
||||
session_manager,
|
||||
source_chat_id=source_chat_id,
|
||||
before_user_index=raw_index,
|
||||
title=(
|
||||
envelope.get("title")
|
||||
if isinstance(envelope.get("title"), str)
|
||||
else None
|
||||
),
|
||||
)
|
||||
if forked is None:
|
||||
await channel.send_webui_protocol_error(
|
||||
connection,
|
||||
"invalid fork source or index",
|
||||
)
|
||||
return
|
||||
fork_id, fork_key = forked
|
||||
except Exception as exc:
|
||||
channel.logger.warning("fork_chat failed: {}", exc)
|
||||
await channel.send_webui_protocol_error(connection, "fork_chat_failed")
|
||||
return
|
||||
fork_id, fork_key = forked
|
||||
except Exception as exc:
|
||||
channel.logger.warning("fork_chat failed: {}", exc)
|
||||
await channel.send_webui_protocol_error(connection, "fork_chat_failed")
|
||||
return
|
||||
|
||||
await channel.attach_webui_fork(
|
||||
connection,
|
||||
fork_id=fork_id,
|
||||
fork_key=fork_key,
|
||||
)
|
||||
await channel.attach_webui_fork(
|
||||
connection,
|
||||
fork_id=fork_id,
|
||||
fork_key=fork_key,
|
||||
)
|
||||
|
||||
await shield_and_drain(create_and_attach())
|
||||
|
||||
@@ -16,7 +16,7 @@ from nanobot.agent.tools.mcp import MCPConnection, connect_mcp_servers
|
||||
from nanobot.agent.tools.mcp_oauth import MCP_OAUTH_CALLBACK_PATH, MCPOAuthHandlers
|
||||
from nanobot.agent.tools.registry import ToolRegistry
|
||||
from nanobot.config.schema import MCPServerConfig
|
||||
from nanobot.security.network import validate_url_target
|
||||
from nanobot.security.network import async_validate_url_target
|
||||
from nanobot.webui.http_utils import is_loopback_host
|
||||
|
||||
McpReload = Callable[[], Awaitable[dict[str, Any]]]
|
||||
@@ -259,7 +259,7 @@ class McpOAuthManager:
|
||||
):
|
||||
flow.error = "The MCP server returned an unsafe authorization URL."
|
||||
raise McpOAuthError(flow.error)
|
||||
ok, _error = validate_url_target(authorization_url)
|
||||
ok, _error = await async_validate_url_target(authorization_url)
|
||||
if not ok:
|
||||
flow.error = "The MCP server returned an unsafe authorization URL."
|
||||
raise McpOAuthError(flow.error)
|
||||
|
||||
@@ -27,6 +27,7 @@ from nanobot.providers.image_generation import (
|
||||
)
|
||||
from nanobot.providers.registry import find_by_name
|
||||
from nanobot.security.network import is_loopback_host
|
||||
from nanobot.utils.cancellation import shield_and_drain
|
||||
from nanobot.webui.settings_contracts import (
|
||||
QueryParams,
|
||||
SettingsRequest,
|
||||
@@ -640,7 +641,11 @@ class CapabilitySettingsHandler:
|
||||
) -> SettingsRouteResult:
|
||||
if action == "api-status":
|
||||
return SettingsRouteResult.success(
|
||||
api_service_payload(self.settings, operations.api_runtime())
|
||||
await asyncio.to_thread(
|
||||
api_service_payload,
|
||||
self.settings,
|
||||
operations.api_runtime(),
|
||||
)
|
||||
)
|
||||
if action == "api-start":
|
||||
return await self._start_api(request, operations)
|
||||
@@ -673,17 +678,22 @@ class CapabilitySettingsHandler:
|
||||
return SettingsRouteResult.failure(404, "unknown settings action")
|
||||
|
||||
operation, section, apply_image_reload = mutation
|
||||
try:
|
||||
payload = self.settings.mutate(operation, request.query)
|
||||
except WebUISettingsError as exc:
|
||||
return SettingsRouteResult.failure(exc.status, exc.message)
|
||||
if apply_image_reload:
|
||||
payload, image_restart_cleared = await self.apply_image_runtime_change(
|
||||
|
||||
async def mutate_and_apply() -> tuple[dict[str, Any], bool]:
|
||||
payload = await self.settings.mutate_async(operation, request.query)
|
||||
if not apply_image_reload:
|
||||
return payload, False
|
||||
return await self.apply_image_runtime_change(
|
||||
payload,
|
||||
operations.reload_image,
|
||||
)
|
||||
else:
|
||||
image_restart_cleared = False
|
||||
|
||||
try:
|
||||
payload, image_restart_cleared = await shield_and_drain(
|
||||
mutate_and_apply()
|
||||
)
|
||||
except WebUISettingsError as exc:
|
||||
return SettingsRouteResult.failure(exc.status, exc.message)
|
||||
return SettingsRouteResult.success(
|
||||
payload,
|
||||
decorate_restart=True,
|
||||
@@ -726,16 +736,17 @@ class CapabilitySettingsHandler:
|
||||
400,
|
||||
"API service API key must be a string",
|
||||
)
|
||||
try:
|
||||
await asyncio.to_thread(
|
||||
self.settings.mutate,
|
||||
allow_install = await self._allow_feature_package_install(request)
|
||||
|
||||
async def mutate_and_start() -> Any:
|
||||
await self.settings.mutate_async(
|
||||
operations.nanobot_features_action,
|
||||
"enable",
|
||||
{"name": ["api"]},
|
||||
allow_install=self._allow_feature_package_install(request),
|
||||
allow_install=allow_install,
|
||||
)
|
||||
self.settings.mutate(operations.update_api, request.query)
|
||||
config = self.settings.config.load()
|
||||
await self.settings.mutate_async(operations.update_api, request.query)
|
||||
config = await self.settings.config.load_async()
|
||||
runtime = operations.api_runtime()
|
||||
options = ApiStartOptions(
|
||||
host=config.api.host,
|
||||
@@ -744,10 +755,13 @@ class CapabilitySettingsHandler:
|
||||
config_path=str(self.settings.config.path),
|
||||
)
|
||||
current = runtime.status()
|
||||
result = await asyncio.to_thread(
|
||||
return await asyncio.to_thread(
|
||||
runtime.restart if current.running else runtime.start_background,
|
||||
options,
|
||||
)
|
||||
|
||||
try:
|
||||
result = await shield_and_drain(mutate_and_start())
|
||||
if not result.ok:
|
||||
return SettingsRouteResult.failure(
|
||||
500,
|
||||
@@ -762,7 +776,8 @@ class CapabilitySettingsHandler:
|
||||
self.logger.exception("failed to start managed API service")
|
||||
return SettingsRouteResult.failure(500, str(exc))
|
||||
return SettingsRouteResult.success(
|
||||
api_service_payload(
|
||||
await asyncio.to_thread(
|
||||
api_service_payload,
|
||||
self.settings,
|
||||
operations.api_runtime(),
|
||||
last_action="started",
|
||||
@@ -775,7 +790,7 @@ class CapabilitySettingsHandler:
|
||||
) -> SettingsRouteResult:
|
||||
runtime = operations.api_runtime()
|
||||
try:
|
||||
result = await asyncio.to_thread(runtime.stop)
|
||||
result = await shield_and_drain(asyncio.to_thread(runtime.stop))
|
||||
except Exception as exc:
|
||||
self.logger.exception("failed to stop managed API service")
|
||||
return SettingsRouteResult.failure(500, str(exc))
|
||||
@@ -785,20 +800,20 @@ class CapabilitySettingsHandler:
|
||||
api_runtime_message(result.message),
|
||||
)
|
||||
return SettingsRouteResult.success(
|
||||
api_service_payload(
|
||||
await asyncio.to_thread(
|
||||
api_service_payload,
|
||||
self.settings,
|
||||
operations.api_runtime(),
|
||||
last_action="stopped",
|
||||
)
|
||||
)
|
||||
|
||||
def _allow_feature_package_install(self, request: SettingsRequest) -> bool:
|
||||
async def _allow_feature_package_install(self, request: SettingsRequest) -> bool:
|
||||
if request.local_browser:
|
||||
return True
|
||||
try:
|
||||
return bool(
|
||||
self.settings.config.load().tools.webui_allow_remote_package_install
|
||||
)
|
||||
config = await self.settings.config.load_async()
|
||||
return bool(config.tools.webui_allow_remote_package_install)
|
||||
except Exception:
|
||||
self.logger.exception("failed to load remote package install policy")
|
||||
return False
|
||||
|
||||
@@ -29,6 +29,7 @@ from nanobot.config.schema import Config, FallbackCandidate, ModelPresetConfig,
|
||||
from nanobot.providers.image_generation import get_image_gen_provider
|
||||
from nanobot.providers.oauth_guidance import OAUTH_CLI_KIT_MISSING_MESSAGE
|
||||
from nanobot.providers.registry import PROVIDERS, create_dynamic_spec, find_by_name
|
||||
from nanobot.utils.cancellation import shield_and_drain
|
||||
from nanobot.webui.settings_contracts import (
|
||||
QueryParams,
|
||||
SettingsRequest,
|
||||
@@ -1651,6 +1652,30 @@ class ModelSettingsHandler:
|
||||
if self.settings.refresh_runtime_config is not None:
|
||||
self.settings.refresh_runtime_config()
|
||||
|
||||
async def _mutate_and_refresh(
|
||||
self,
|
||||
operation: SettingsOperation,
|
||||
query: QueryParams,
|
||||
**kwargs: Any,
|
||||
) -> dict[str, Any]:
|
||||
payload = await self.settings.mutate_async(operation, query, **kwargs)
|
||||
self._refresh_runtime_config()
|
||||
return payload
|
||||
|
||||
async def _update_provider_and_runtime(
|
||||
self,
|
||||
operation: SettingsOperation,
|
||||
query: QueryParams,
|
||||
apply_image_runtime_change: Callable[
|
||||
[dict[str, Any]],
|
||||
Awaitable[tuple[dict[str, Any], bool]],
|
||||
],
|
||||
) -> tuple[dict[str, Any], bool]:
|
||||
payload = await self.settings.mutate_async(operation, query)
|
||||
payload, image_restart_cleared = await apply_image_runtime_change(payload)
|
||||
self._refresh_runtime_config()
|
||||
return payload, image_restart_cleared
|
||||
|
||||
async def handle(
|
||||
self,
|
||||
action: str,
|
||||
@@ -1659,8 +1684,12 @@ class ModelSettingsHandler:
|
||||
) -> SettingsRouteResult:
|
||||
try:
|
||||
if action == "agent-update":
|
||||
payload = self.settings.mutate(operations.update_agent, request.query)
|
||||
self._refresh_runtime_config()
|
||||
payload = await shield_and_drain(
|
||||
self._mutate_and_refresh(
|
||||
operations.update_agent,
|
||||
request.query,
|
||||
)
|
||||
)
|
||||
return SettingsRouteResult.success(
|
||||
payload,
|
||||
decorate_restart=True,
|
||||
@@ -1668,12 +1697,13 @@ class ModelSettingsHandler:
|
||||
)
|
||||
|
||||
if action == "model-update":
|
||||
payload = self.settings.mutate(
|
||||
operations.update_model,
|
||||
request.query,
|
||||
rename_model_preset=self.settings.rename_model_preset,
|
||||
payload = await shield_and_drain(
|
||||
self._mutate_and_refresh(
|
||||
operations.update_model,
|
||||
request.query,
|
||||
rename_model_preset=self.settings.rename_model_preset,
|
||||
)
|
||||
)
|
||||
self._refresh_runtime_config()
|
||||
return SettingsRouteResult.success(payload, decorate_restart=True)
|
||||
|
||||
mutation = {
|
||||
@@ -1684,19 +1714,19 @@ class ModelSettingsHandler:
|
||||
"provider-create": operations.create_provider,
|
||||
}.get(action)
|
||||
if mutation is not None:
|
||||
payload = self.settings.mutate(mutation, request.query)
|
||||
self._refresh_runtime_config()
|
||||
payload = await shield_and_drain(
|
||||
self._mutate_and_refresh(mutation, request.query)
|
||||
)
|
||||
return SettingsRouteResult.success(payload, decorate_restart=True)
|
||||
|
||||
if action == "provider-update":
|
||||
payload = self.settings.mutate(
|
||||
operations.update_provider,
|
||||
request.query,
|
||||
payload, image_restart_cleared = await shield_and_drain(
|
||||
self._update_provider_and_runtime(
|
||||
operations.update_provider,
|
||||
request.query,
|
||||
operations.apply_image_runtime_change,
|
||||
)
|
||||
)
|
||||
payload, image_restart_cleared = await operations.apply_image_runtime_change(
|
||||
payload
|
||||
)
|
||||
self._refresh_runtime_config()
|
||||
return SettingsRouteResult.success(
|
||||
payload,
|
||||
decorate_restart=True,
|
||||
@@ -1724,11 +1754,13 @@ class ModelSettingsHandler:
|
||||
return SettingsRouteResult.success(payload)
|
||||
|
||||
if action == "oauth-login":
|
||||
payload = await asyncio.to_thread(
|
||||
self.settings.read,
|
||||
operations.oauth_login,
|
||||
request.query,
|
||||
oauth_flows=self.settings.oauth_flows,
|
||||
payload = await shield_and_drain(
|
||||
asyncio.to_thread(
|
||||
self.settings.read,
|
||||
operations.oauth_login,
|
||||
request.query,
|
||||
oauth_flows=self.settings.oauth_flows,
|
||||
)
|
||||
)
|
||||
elif action == "oauth-complete":
|
||||
raw_response = (request.payload or {}).get("authorization_response")
|
||||
@@ -1736,19 +1768,23 @@ class ModelSettingsHandler:
|
||||
raise WebUISettingsError(
|
||||
"OAuth authorization response must be a string"
|
||||
)
|
||||
payload = await asyncio.to_thread(
|
||||
self.settings.read,
|
||||
operations.oauth_complete,
|
||||
request.query,
|
||||
raw_response or None,
|
||||
oauth_flows=self.settings.oauth_flows,
|
||||
payload = await shield_and_drain(
|
||||
asyncio.to_thread(
|
||||
self.settings.read,
|
||||
operations.oauth_complete,
|
||||
request.query,
|
||||
raw_response or None,
|
||||
oauth_flows=self.settings.oauth_flows,
|
||||
)
|
||||
)
|
||||
elif action == "oauth-logout":
|
||||
payload = await asyncio.to_thread(
|
||||
self.settings.read,
|
||||
operations.oauth_logout,
|
||||
request.query,
|
||||
oauth_flows=self.settings.oauth_flows,
|
||||
payload = await shield_and_drain(
|
||||
asyncio.to_thread(
|
||||
self.settings.read,
|
||||
operations.oauth_logout,
|
||||
request.query,
|
||||
oauth_flows=self.settings.oauth_flows,
|
||||
)
|
||||
)
|
||||
else:
|
||||
return SettingsRouteResult.failure(404, "unknown settings action")
|
||||
|
||||
@@ -4,6 +4,7 @@ from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import html
|
||||
import inspect
|
||||
import json
|
||||
from collections.abc import Awaitable, Callable, Mapping
|
||||
from typing import Any, cast
|
||||
@@ -18,6 +19,7 @@ from nanobot.bus.queue import MessageBus
|
||||
from nanobot.channels.registry import load_channel_plugin
|
||||
from nanobot.channels.validation import validate_channel_config
|
||||
from nanobot.pairing import approve_code, deny_code, list_pending
|
||||
from nanobot.utils.cancellation import shield_and_drain
|
||||
from nanobot.webui import settings_capabilities as capability_domain
|
||||
from nanobot.webui import settings_contracts as contracts
|
||||
from nanobot.webui import settings_models as model_domain
|
||||
@@ -208,6 +210,16 @@ def _payload_query(payload: dict[str, Any]) -> QueryParams:
|
||||
}
|
||||
|
||||
|
||||
async def _call_settings_handler(
|
||||
handler: Callable[[], Response | Awaitable[Response]],
|
||||
) -> Response:
|
||||
"""Keep synchronous handlers off-loop while supporting native async handlers."""
|
||||
result = await asyncio.to_thread(handler)
|
||||
if inspect.isawaitable(result):
|
||||
return await result
|
||||
return result
|
||||
|
||||
|
||||
class WebUISettingsRouter:
|
||||
"""Authenticate and dispatch settings requests to transport-neutral domains."""
|
||||
|
||||
@@ -284,9 +296,9 @@ class WebUISettingsRouter:
|
||||
if not self._authorized(request):
|
||||
return self._unauthorized()
|
||||
if route == ("root", "settings"):
|
||||
return await asyncio.to_thread(self._handle_settings)
|
||||
return await _call_settings_handler(self._handle_settings)
|
||||
if route == ("root", "usage"):
|
||||
return await asyncio.to_thread(self._handle_settings_usage)
|
||||
return await _call_settings_handler(self._handle_settings_usage)
|
||||
|
||||
domain, action = route
|
||||
domain_request = self._domain_request(
|
||||
@@ -415,19 +427,18 @@ class WebUISettingsRouter:
|
||||
)
|
||||
return self._json_response(payload)
|
||||
|
||||
def _handle_settings(self) -> Response:
|
||||
return self._json_response(
|
||||
self._with_restart_state(
|
||||
self.settings.read(
|
||||
settings_payload,
|
||||
surface=self._runtime_surface,
|
||||
runtime_capability_overrides=self._runtime_capabilities,
|
||||
)
|
||||
)
|
||||
async def _handle_settings(self) -> Response:
|
||||
payload = await self.settings.read_async(
|
||||
settings_payload,
|
||||
surface=self._runtime_surface,
|
||||
runtime_capability_overrides=self._runtime_capabilities,
|
||||
)
|
||||
return self._json_response(self._with_restart_state(payload))
|
||||
|
||||
def _handle_settings_usage(self) -> Response:
|
||||
return self._json_response(self.settings.read(settings_usage_payload))
|
||||
async def _handle_settings_usage(self) -> Response:
|
||||
return self._json_response(
|
||||
await self.settings.read_async(settings_usage_payload)
|
||||
)
|
||||
|
||||
def _model_operations(self) -> model_domain.ModelSettingsOperations:
|
||||
return model_domain.ModelSettingsOperations(
|
||||
@@ -560,7 +571,7 @@ class WebUISettingsRouter:
|
||||
allow_install=allow_install,
|
||||
)
|
||||
|
||||
def _allow_feature_package_install(
|
||||
async def _allow_feature_package_install(
|
||||
self,
|
||||
connection: Any,
|
||||
request: WsRequest,
|
||||
@@ -570,29 +581,33 @@ class WebUISettingsRouter:
|
||||
request,
|
||||
needs_local_browser=True,
|
||||
)
|
||||
return self._system.allow_feature_package_install(domain_request)
|
||||
return await self._system.allow_feature_package_install(domain_request)
|
||||
|
||||
async def _handle_mcp_oauth_start(self, request: WsRequest) -> Response:
|
||||
if not self._authorized(request):
|
||||
return self._unauthorized()
|
||||
if self._mcp_oauth_redirect_uri is None:
|
||||
redirect_uri_for_request = self._mcp_oauth_redirect_uri
|
||||
if redirect_uri_for_request is None:
|
||||
return self._error_response(500, "MCP OAuth callback is not configured")
|
||||
query = self._parse_mcp_settings_query(request)
|
||||
try:
|
||||
name, cfg = await asyncio.to_thread(
|
||||
self.settings.mutate,
|
||||
|
||||
async def mutate_and_start() -> dict[str, Any]:
|
||||
name, cfg = await self.settings.mutate_async(
|
||||
ensure_mcp_oauth_server,
|
||||
query,
|
||||
)
|
||||
redirect_uri = self._mcp_oauth_redirect_uri(request)
|
||||
redirect_uri = redirect_uri_for_request(request)
|
||||
reset = (_query_first(query, "reset") or "").lower() in {"1", "true", "yes"}
|
||||
payload = await self._mcp_oauth.start(
|
||||
return await self._mcp_oauth.start(
|
||||
name,
|
||||
cfg,
|
||||
redirect_uri,
|
||||
reload_mcp=self._reload_mcp_runtime,
|
||||
reset_credentials=reset,
|
||||
)
|
||||
|
||||
try:
|
||||
payload = await shield_and_drain(mutate_and_start())
|
||||
except Exception as exc:
|
||||
return self._mcp_oauth_error_response(exc, action="start")
|
||||
return self._json_response(payload)
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import threading
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
@@ -12,9 +13,11 @@ from filelock import FileLock
|
||||
|
||||
from nanobot.config.loader import load_config, save_config
|
||||
from nanobot.config.schema import Config
|
||||
from nanobot.utils.cancellation import shield_and_drain
|
||||
|
||||
_T = TypeVar("_T")
|
||||
_WEBUI_OAUTH_MAX_FLOWS = 8
|
||||
_SETTINGS_FILE_LOCK_TIMEOUT_SECONDS = 5
|
||||
|
||||
|
||||
class WebUISettingsConfig:
|
||||
@@ -25,13 +28,20 @@ class WebUISettingsConfig:
|
||||
self.path.parent.mkdir(parents=True, exist_ok=True)
|
||||
self._lock = threading.RLock()
|
||||
lock_path = self.path.with_suffix(f"{self.path.suffix}.lock")
|
||||
self._file_lock = FileLock(str(lock_path))
|
||||
self._file_lock = FileLock(
|
||||
str(lock_path),
|
||||
timeout=_SETTINGS_FILE_LOCK_TIMEOUT_SECONDS,
|
||||
)
|
||||
|
||||
def load(self) -> Config:
|
||||
"""Load this gateway's config without consulting the process-global path."""
|
||||
with self._lock:
|
||||
return load_config(self.path)
|
||||
|
||||
async def load_async(self) -> Config:
|
||||
"""Load config without running file I/O or lock waits on the event loop."""
|
||||
return await asyncio.to_thread(self.load)
|
||||
|
||||
def update(self, mutation: Callable[[Config], _T]) -> _T:
|
||||
"""Apply and atomically persist one path-scoped read-modify-write operation."""
|
||||
with self._lock, self._file_lock:
|
||||
@@ -40,11 +50,21 @@ class WebUISettingsConfig:
|
||||
save_config(config, self.path)
|
||||
return result
|
||||
|
||||
async def update_async(self, mutation: Callable[[Config], _T]) -> _T:
|
||||
"""Update config without blocking the event loop."""
|
||||
return await shield_and_drain(asyncio.to_thread(self.update, mutation))
|
||||
|
||||
def run_serialized(self, operation: Callable[[Path], _T]) -> _T:
|
||||
"""Run a path-aware read-modify-write operation under the config-file lock."""
|
||||
with self._lock, self._file_lock:
|
||||
return operation(self.path)
|
||||
|
||||
async def run_serialized_async(self, operation: Callable[[Path], _T]) -> _T:
|
||||
"""Run a serialized config operation without blocking the event loop."""
|
||||
return await shield_and_drain(
|
||||
asyncio.to_thread(self.run_serialized, operation)
|
||||
)
|
||||
|
||||
|
||||
class WebUIOAuthFlowRegistry:
|
||||
"""Bounded, thread-safe OAuth flows owned by one gateway instance."""
|
||||
@@ -146,6 +166,16 @@ class WebUISettingsServices:
|
||||
"""Run a settings read against this gateway's explicit config path."""
|
||||
return operation(*args, config_path=self.config.path, **kwargs)
|
||||
|
||||
async def read_async(
|
||||
self,
|
||||
operation: Callable[..., _T],
|
||||
/,
|
||||
*args: Any,
|
||||
**kwargs: Any,
|
||||
) -> _T:
|
||||
"""Run a settings read without blocking the event loop."""
|
||||
return await asyncio.to_thread(self.read, operation, *args, **kwargs)
|
||||
|
||||
def mutate(
|
||||
self,
|
||||
operation: Callable[..., _T],
|
||||
@@ -161,3 +191,15 @@ class WebUISettingsServices:
|
||||
**kwargs,
|
||||
)
|
||||
)
|
||||
|
||||
async def mutate_async(
|
||||
self,
|
||||
operation: Callable[..., _T],
|
||||
/,
|
||||
*args: Any,
|
||||
**kwargs: Any,
|
||||
) -> _T:
|
||||
"""Mutate settings without blocking the event loop."""
|
||||
return await shield_and_drain(
|
||||
asyncio.to_thread(self.mutate, operation, *args, **kwargs)
|
||||
)
|
||||
|
||||
@@ -23,6 +23,7 @@ from nanobot.config.schema import Config
|
||||
from nanobot.llm_usage import llm_usage_payload
|
||||
from nanobot.optional_features import OptionalFeatureError, with_channel_runtime_status
|
||||
from nanobot.security.workspace_access import workspace_sandbox_status
|
||||
from nanobot.utils.cancellation import shield_and_drain
|
||||
from nanobot.webui.settings_capabilities import network_safety_payload
|
||||
from nanobot.webui.settings_contracts import (
|
||||
QueryParams,
|
||||
@@ -446,12 +447,17 @@ class SystemSettingsHandler:
|
||||
operations: SystemSettingsOperations,
|
||||
) -> SettingsRouteResult:
|
||||
try:
|
||||
payload = await asyncio.to_thread(
|
||||
pending = asyncio.to_thread(
|
||||
operations.cli_apps_action,
|
||||
action,
|
||||
request.query,
|
||||
config_path=self.settings.config.path,
|
||||
)
|
||||
payload = (
|
||||
await shield_and_drain(pending)
|
||||
if action in {"install", "update", "uninstall"}
|
||||
else await pending
|
||||
)
|
||||
except WebUISettingsError as exc:
|
||||
return SettingsRouteResult.failure(exc.status, exc.message)
|
||||
except Exception as exc:
|
||||
@@ -505,17 +511,29 @@ class SystemSettingsHandler:
|
||||
action: str,
|
||||
operations: SystemSettingsOperations,
|
||||
) -> SettingsRouteResult:
|
||||
try:
|
||||
allow_install = (
|
||||
action != "enable"
|
||||
or await self.allow_feature_package_install(request)
|
||||
)
|
||||
|
||||
async def mutate_and_apply() -> dict[str, Any]:
|
||||
payload = await asyncio.to_thread(
|
||||
self._nanobot_features_action,
|
||||
action,
|
||||
request.query,
|
||||
operations,
|
||||
allow_install=(
|
||||
action != "enable"
|
||||
or self.allow_feature_package_install(request)
|
||||
),
|
||||
allow_install=allow_install,
|
||||
)
|
||||
payload = await self._apply_feature_runtime_change(
|
||||
action,
|
||||
request.query,
|
||||
payload,
|
||||
operations,
|
||||
)
|
||||
return self._with_channel_runtime_status(payload, operations)
|
||||
|
||||
try:
|
||||
payload = await shield_and_drain(mutate_and_apply())
|
||||
except OptionalFeatureError as exc:
|
||||
return SettingsRouteResult.failure(exc.status, exc.message)
|
||||
except Exception as exc:
|
||||
@@ -527,13 +545,6 @@ class SystemSettingsHandler:
|
||||
action,
|
||||
)
|
||||
return SettingsRouteResult.failure(status, message)
|
||||
payload = await self._apply_feature_runtime_change(
|
||||
action,
|
||||
request.query,
|
||||
payload,
|
||||
operations,
|
||||
)
|
||||
payload = self._with_channel_runtime_status(payload, operations)
|
||||
return SettingsRouteResult.success(
|
||||
payload,
|
||||
decorate_restart=True,
|
||||
@@ -628,6 +639,15 @@ class SystemSettingsHandler:
|
||||
self,
|
||||
request: SettingsRequest,
|
||||
operations: SystemSettingsOperations,
|
||||
) -> SettingsRouteResult:
|
||||
return await shield_and_drain(
|
||||
self._channel_configure_settled(request, operations)
|
||||
)
|
||||
|
||||
async def _channel_configure_settled(
|
||||
self,
|
||||
request: SettingsRequest,
|
||||
operations: SystemSettingsOperations,
|
||||
) -> SettingsRouteResult:
|
||||
name = (query_first(request.query, "name") or "").strip()
|
||||
instance_id = (
|
||||
@@ -682,7 +702,7 @@ class SystemSettingsHandler:
|
||||
"enable",
|
||||
feature_query,
|
||||
operations,
|
||||
allow_install=self.allow_feature_package_install(request),
|
||||
allow_install=await self.allow_feature_package_install(request),
|
||||
)
|
||||
except OptionalFeatureError as exc:
|
||||
return SettingsRouteResult.failure(
|
||||
@@ -825,6 +845,22 @@ class SystemSettingsHandler:
|
||||
channel_name: str,
|
||||
payload: dict[str, Any],
|
||||
operations: SystemSettingsOperations,
|
||||
) -> dict[str, Any]:
|
||||
return await shield_and_drain(
|
||||
self._settle_channel_connect_success(
|
||||
request,
|
||||
channel_name,
|
||||
payload,
|
||||
operations,
|
||||
)
|
||||
)
|
||||
|
||||
async def _settle_channel_connect_success(
|
||||
self,
|
||||
request: SettingsRequest,
|
||||
channel_name: str,
|
||||
payload: dict[str, Any],
|
||||
operations: SystemSettingsOperations,
|
||||
) -> dict[str, Any]:
|
||||
target = {"name": [channel_name]}
|
||||
if payload.get("instance_id"):
|
||||
@@ -835,11 +871,11 @@ class SystemSettingsHandler:
|
||||
"enable",
|
||||
target,
|
||||
operations,
|
||||
allow_install=self.allow_feature_package_install(request),
|
||||
allow_install=await self.allow_feature_package_install(request),
|
||||
)
|
||||
except OptionalFeatureError as exc:
|
||||
features = self.feature_runtime_fallback(
|
||||
self._nanobot_features_payload(operations),
|
||||
await asyncio.to_thread(self._nanobot_features_payload, operations),
|
||||
message=(
|
||||
f"{channel_name} connected, but enabling channel support failed: "
|
||||
f"{exc.message}"
|
||||
@@ -859,13 +895,12 @@ class SystemSettingsHandler:
|
||||
)
|
||||
return updated
|
||||
|
||||
def allow_feature_package_install(self, request: SettingsRequest) -> bool:
|
||||
async def allow_feature_package_install(self, request: SettingsRequest) -> bool:
|
||||
if request.local_browser:
|
||||
return True
|
||||
try:
|
||||
return bool(
|
||||
self.settings.config.load().tools.webui_allow_remote_package_install
|
||||
)
|
||||
config = await self.settings.config.load_async()
|
||||
return bool(config.tools.webui_allow_remote_package_install)
|
||||
except Exception:
|
||||
self.logger.exception("failed to load remote package install policy")
|
||||
return False
|
||||
@@ -925,13 +960,18 @@ class SystemSettingsHandler:
|
||||
operations: SystemSettingsOperations,
|
||||
) -> SettingsRouteResult:
|
||||
try:
|
||||
payload = await operations.mcp_presets_action(
|
||||
pending = operations.mcp_presets_action(
|
||||
action,
|
||||
request.query,
|
||||
reload_mcp=operations.reload_mcp,
|
||||
mcp_runtime_status=operations.mcp_runtime_status,
|
||||
config=self.settings.config,
|
||||
)
|
||||
payload = (
|
||||
await pending
|
||||
if action is None
|
||||
else await shield_and_drain(pending)
|
||||
)
|
||||
except Exception as exc:
|
||||
status = getattr(exc, "status", 500)
|
||||
message = getattr(exc, "message", str(exc))
|
||||
|
||||
+113
-24
@@ -34,6 +34,7 @@ from nanobot.session.session_handles import (
|
||||
SessionHandleResolver,
|
||||
)
|
||||
from nanobot.triggers.local_types import LocalTrigger
|
||||
from nanobot.utils.cancellation import shield_and_drain
|
||||
from nanobot.webui.file_preview import (
|
||||
WebUIFilePreviewError,
|
||||
file_preview_availability_payload,
|
||||
@@ -135,6 +136,18 @@ _WEBUI_MUTATION_PAYLOAD_ATTR = "_nanobot_webui_mutation_payload"
|
||||
_WEBUI_MUTATION_REQUEST_ATTR = "_nanobot_webui_mutation_request"
|
||||
_NO_STORE_HEADERS = [("Cache-Control", "no-store")]
|
||||
|
||||
|
||||
def _slow_http_operation(path: str) -> str:
|
||||
"""Return a route family without logging user-controlled path/query values."""
|
||||
clean_path = path.split("?", 1)[0]
|
||||
if clean_path == "/webui/bootstrap":
|
||||
return clean_path
|
||||
parts = [part for part in clean_path.split("/") if part]
|
||||
if len(parts) >= 2 and parts[0] == "api":
|
||||
return f"/api/{parts[1]}"
|
||||
return "/webui"
|
||||
|
||||
|
||||
_WEBUI_MUTATION_PATHS = {
|
||||
"automation.enable": "/api/webui/automations/enable",
|
||||
"automation.disable": "/api/webui/automations/disable",
|
||||
@@ -420,7 +433,12 @@ class GatewayHTTPHandler:
|
||||
response = await self._dispatch_resolved(connection, request, got)
|
||||
return response
|
||||
finally:
|
||||
self._log_slow_http(got, response, started)
|
||||
self._log_slow_http(
|
||||
got,
|
||||
response,
|
||||
started,
|
||||
input_chars=len(request.path),
|
||||
)
|
||||
|
||||
async def dispatch_webui_mutation(
|
||||
self,
|
||||
@@ -556,7 +574,14 @@ class GatewayHTTPHandler:
|
||||
|
||||
return connection.respond(404, "Not Found")
|
||||
|
||||
def _log_slow_http(self, path: str, response: Any | None, started: float) -> None:
|
||||
def _log_slow_http(
|
||||
self,
|
||||
path: str,
|
||||
response: Any | None,
|
||||
started: float,
|
||||
*,
|
||||
input_chars: int,
|
||||
) -> None:
|
||||
elapsed_ms = int((time.perf_counter() - started) * 1000)
|
||||
if elapsed_ms < _SLOW_WEBUI_HTTP_LOG_MS:
|
||||
return
|
||||
@@ -564,9 +589,10 @@ class GatewayHTTPHandler:
|
||||
return
|
||||
status = getattr(response, "status_code", None)
|
||||
self._log.warning(
|
||||
"slow webui http route path={} status={} duration_ms={}",
|
||||
path,
|
||||
"slow webui http operation={} status={} input_chars={} duration_ms={}",
|
||||
_slow_http_operation(path),
|
||||
status if status is not None else "none",
|
||||
input_chars,
|
||||
elapsed_ms,
|
||||
)
|
||||
|
||||
@@ -694,7 +720,11 @@ class GatewayHTTPHandler:
|
||||
async def _dispatch_session_routes(self, request: WsRequest, got: str) -> Response | None:
|
||||
m = re.match(r"^/api/sessions/([^/]+)/webui-thread$", got)
|
||||
if m:
|
||||
return self._handle_webui_thread_get(request, m.group(1))
|
||||
return await asyncio.to_thread(
|
||||
self._handle_webui_thread_get,
|
||||
request,
|
||||
m.group(1),
|
||||
)
|
||||
|
||||
m = re.match(r"^/api/sessions/([^/]+)/context$", got)
|
||||
if m:
|
||||
@@ -702,15 +732,27 @@ class GatewayHTTPHandler:
|
||||
|
||||
m = re.match(r"^/api/sessions/([^/]+)/file-preview$", got)
|
||||
if m:
|
||||
return self._handle_file_preview(request, m.group(1))
|
||||
return await asyncio.to_thread(
|
||||
self._handle_file_preview,
|
||||
request,
|
||||
m.group(1),
|
||||
)
|
||||
|
||||
m = re.match(r"^/api/sessions/([^/]+)/automations$", got)
|
||||
if m:
|
||||
return self._handle_session_automations(request, m.group(1))
|
||||
return await self._run_cron_transaction(
|
||||
self._handle_session_automations,
|
||||
request,
|
||||
m.group(1),
|
||||
)
|
||||
|
||||
m = re.match(r"^/api/sessions/([^/]+)/delete$", got)
|
||||
if m:
|
||||
return self._handle_session_delete(request, m.group(1))
|
||||
return await self._run_cron_transaction(
|
||||
self._handle_session_delete,
|
||||
request,
|
||||
m.group(1),
|
||||
)
|
||||
|
||||
return None
|
||||
|
||||
@@ -957,13 +999,24 @@ class GatewayHTTPHandler:
|
||||
|
||||
# -- Automation routes --------------------------------------------------
|
||||
|
||||
async def _run_cron_transaction(
|
||||
self,
|
||||
operation: Callable[..., Any],
|
||||
/,
|
||||
*args: Any,
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
if self.cron_service is not None:
|
||||
return await self.cron_service.run_sync(operation, *args, **kwargs)
|
||||
return await shield_and_drain(asyncio.to_thread(operation, *args, **kwargs))
|
||||
|
||||
async def _dispatch_automation_routes(
|
||||
self,
|
||||
request: WsRequest,
|
||||
got: str,
|
||||
) -> Response | None:
|
||||
if got == "/api/webui/automations":
|
||||
return self._handle_webui_automations(request)
|
||||
return await self._run_cron_transaction(self._handle_webui_automations, request)
|
||||
m = re.match(r"^/api/webui/automations/(enable|disable|delete|run|update)$", got)
|
||||
if m:
|
||||
return await self._handle_webui_automation_action(request, m.group(1))
|
||||
@@ -1029,13 +1082,24 @@ class GatewayHTTPHandler:
|
||||
job_id = (_query_first(query, "id") or _query_first(query, "job_id") or "").strip()
|
||||
if not job_id:
|
||||
return _http_error(400, "missing automation id")
|
||||
trigger = self.local_trigger_store.get(job_id) if self.local_trigger_store else None
|
||||
trigger = (
|
||||
await asyncio.to_thread(self.local_trigger_store.get, job_id)
|
||||
if self.local_trigger_store
|
||||
else None
|
||||
)
|
||||
if trigger is not None:
|
||||
return self._handle_local_trigger_action(request, action, trigger)
|
||||
return await shield_and_drain(
|
||||
asyncio.to_thread(
|
||||
self._handle_local_trigger_action,
|
||||
request,
|
||||
action,
|
||||
trigger,
|
||||
)
|
||||
)
|
||||
|
||||
if self.cron_service is None:
|
||||
return _http_error(404, "automation not found")
|
||||
job = self.cron_service.get_job(job_id)
|
||||
job = await self.cron_service.run_sync(self.cron_service.get_job, job_id)
|
||||
if job is None:
|
||||
return _http_error(404, "automation not found")
|
||||
if job.payload.kind == "system_event":
|
||||
@@ -1044,13 +1108,23 @@ class GatewayHTTPHandler:
|
||||
return _http_error(409, "automation has no linked chat")
|
||||
|
||||
if action == "enable":
|
||||
if self.cron_service.enable_job(job_id, enabled=True) is None:
|
||||
result = await self.cron_service.run_sync(
|
||||
self.cron_service.enable_job,
|
||||
job_id,
|
||||
enabled=True,
|
||||
)
|
||||
if result is None:
|
||||
return _http_error(404, "automation not found")
|
||||
elif action == "disable":
|
||||
if self.cron_service.enable_job(job_id, enabled=False) is None:
|
||||
result = await self.cron_service.run_sync(
|
||||
self.cron_service.enable_job,
|
||||
job_id,
|
||||
enabled=False,
|
||||
)
|
||||
if result is None:
|
||||
return _http_error(404, "automation not found")
|
||||
elif action == "delete":
|
||||
result = self.cron_service.remove_job(job_id)
|
||||
result = await self.cron_service.run_sync(self.cron_service.remove_job, job_id)
|
||||
if result == "not_found":
|
||||
return _http_error(404, "automation not found")
|
||||
if result == "protected":
|
||||
@@ -1068,7 +1142,11 @@ class GatewayHTTPHandler:
|
||||
if isinstance(parsed, str):
|
||||
return _http_error(400, parsed)
|
||||
try:
|
||||
result = self.cron_service.update_job(job_id, **parsed)
|
||||
result = await self.cron_service.run_sync(
|
||||
self.cron_service.update_job,
|
||||
job_id,
|
||||
**parsed,
|
||||
)
|
||||
except ValueError as exc:
|
||||
return _http_error(400, str(exc))
|
||||
if result == "not_found":
|
||||
@@ -1078,7 +1156,7 @@ class GatewayHTTPHandler:
|
||||
else:
|
||||
return _http_error(404, "unknown automation action")
|
||||
|
||||
return self._handle_webui_automations(request)
|
||||
return await self._run_cron_transaction(self._handle_webui_automations, request)
|
||||
|
||||
def _handle_local_trigger_action(
|
||||
self,
|
||||
@@ -1163,9 +1241,17 @@ class GatewayHTTPHandler:
|
||||
if got == "/api/webui/skills/install":
|
||||
return await self._handle_webui_skill_install(connection, request)
|
||||
if got == "/api/webui/skills/update":
|
||||
return self._handle_webui_skill_update(request)
|
||||
return await shield_and_drain(
|
||||
asyncio.to_thread(self._handle_webui_skill_update, request)
|
||||
)
|
||||
if got == "/api/webui/skills/delete":
|
||||
return self._handle_webui_skill_delete(connection, request)
|
||||
return await shield_and_drain(
|
||||
asyncio.to_thread(
|
||||
self._handle_webui_skill_delete,
|
||||
connection,
|
||||
request,
|
||||
)
|
||||
)
|
||||
if got == "/api/webui/skills":
|
||||
return self._handle_webui_skills(request)
|
||||
m = re.match(r"^/api/webui/skills/([^/]+)$", got)
|
||||
@@ -1276,7 +1362,7 @@ class GatewayHTTPHandler:
|
||||
) -> Response:
|
||||
if not self.check_api_token(request):
|
||||
return _http_error(401, "Unauthorized")
|
||||
if not self._allow_webui_package_install(connection, request):
|
||||
if not await self._allow_webui_package_install(connection, request):
|
||||
return _http_error(403, "remote skill installation is disabled")
|
||||
if self._skill_install_lock.locked():
|
||||
return _http_error(409, "another skill installation is already in progress")
|
||||
@@ -1308,13 +1394,16 @@ class GatewayHTTPHandler:
|
||||
"last_action": action,
|
||||
})
|
||||
|
||||
def _allow_webui_package_install(self, connection: Any, request: WsRequest) -> bool:
|
||||
async def _allow_webui_package_install(
|
||||
self,
|
||||
connection: Any,
|
||||
request: WsRequest,
|
||||
) -> bool:
|
||||
if _is_local_browser_request(connection, request.headers):
|
||||
return True
|
||||
try:
|
||||
return bool(
|
||||
self.settings.config.load().tools.webui_allow_remote_package_install
|
||||
)
|
||||
config = await self.settings.config.load_async()
|
||||
return bool(config.tools.webui_allow_remote_package_install)
|
||||
except Exception:
|
||||
self._log.exception("failed to load remote package install policy")
|
||||
return False
|
||||
|
||||
@@ -123,21 +123,21 @@ def allow_loopback_mcp_urls(monkeypatch: pytest.MonkeyPatch):
|
||||
_resolver_lock = asyncio.Lock()
|
||||
|
||||
monkeypatch.setattr(mcp_module, "PinnedDNSAsyncTransport", TestPinnedDNSAsyncTransport)
|
||||
monkeypatch.setattr(
|
||||
mcp_module,
|
||||
"validate_url_target",
|
||||
lambda url, *, allow_loopback=False: (True, ""),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
mcp_module,
|
||||
"resolve_url_target",
|
||||
lambda url, *, allow_loopback=False: (True, "", ("127.0.0.1",)),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
security_network,
|
||||
"resolve_url_target",
|
||||
lambda url, *, allow_loopback=False: (True, "", ("127.0.0.1",)),
|
||||
)
|
||||
async def allow_url(url: str, *, allow_loopback: bool = False) -> tuple[bool, str]:
|
||||
return True, ""
|
||||
|
||||
async def resolve_url(
|
||||
url: str,
|
||||
*,
|
||||
allow_loopback: bool = False,
|
||||
trust_remote_dns: bool = False,
|
||||
timeout_s: float = 3.0,
|
||||
) -> tuple[bool, str, tuple[str, ...]]:
|
||||
return True, "", ("127.0.0.1",)
|
||||
|
||||
monkeypatch.setattr(mcp_module, "async_validate_url_target", allow_url)
|
||||
monkeypatch.setattr(mcp_module, "async_resolve_url_target", resolve_url)
|
||||
monkeypatch.setattr(security_network, "async_resolve_url_target", resolve_url)
|
||||
monkeypatch.setattr(
|
||||
mcp_module,
|
||||
"env_proxy_applies_to_url",
|
||||
|
||||
@@ -357,3 +357,50 @@ async def test_length_finish_with_blank_content_routes_to_length_recovery():
|
||||
"finish_reason='length' response with blank content"
|
||||
)
|
||||
assert result.final_content == "done"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_slow_tool_log_records_scale_without_argument_content(monkeypatch) -> None:
|
||||
from nanobot.agent import runner as runner_module
|
||||
from nanobot.agent.runner import AgentRunner
|
||||
|
||||
records: list[str] = []
|
||||
|
||||
class _Logger:
|
||||
def warning(self, message: str, *args: object) -> None:
|
||||
records.append(message.format(*args))
|
||||
|
||||
secret = "customer-token-do-not-log"
|
||||
|
||||
async def execute(_name, _args):
|
||||
return "ok"
|
||||
|
||||
monkeypatch.setattr(runner_module, "_SLOW_TOOL_LOG_MS", 0)
|
||||
monkeypatch.setattr(runner_module, "logger", _Logger())
|
||||
runner = AgentRunner()
|
||||
spec = make_run_spec(
|
||||
MagicMock(spec=LLMProvider),
|
||||
initial_messages=[],
|
||||
tools=SimpleNamespace(execute=execute),
|
||||
model="test-model",
|
||||
max_iterations=1,
|
||||
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
|
||||
)
|
||||
|
||||
await runner._run_tool(
|
||||
spec,
|
||||
ToolCallRequest(
|
||||
id="call_1",
|
||||
name="edit_file",
|
||||
arguments={"old_text": secret, "paths": ["one", "two"]},
|
||||
),
|
||||
external_lookup_counts={},
|
||||
workspace_violation_counts={},
|
||||
)
|
||||
|
||||
assert len(records) == 1
|
||||
assert "operation=edit_file" in records[0]
|
||||
assert "input_items=4" in records[0]
|
||||
assert f"input_chars={len(secret)}" in records[0]
|
||||
assert "duration_ms=" in records[0]
|
||||
assert secret not in records[0]
|
||||
|
||||
@@ -1233,6 +1233,32 @@ async def test_submitted_cron_turn_reports_pending_until_completed(tmp_path):
|
||||
assert loop.pending_cron_job_ids_for_session(session_key) == set()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_accepted_automation_turn_preserves_wait_for_timeout(tmp_path):
|
||||
"""Accepted-turn signalling must not leak through asyncio timeout semantics."""
|
||||
from nanobot.bus.events import InboundMessage
|
||||
from nanobot.cron.session_turns import CRON_TRIGGER_META
|
||||
|
||||
loop = _make_loop(tmp_path)
|
||||
loop._running = True
|
||||
msg = InboundMessage(
|
||||
channel="websocket",
|
||||
sender_id="cron",
|
||||
chat_id="chat-1",
|
||||
content="scheduled work",
|
||||
metadata={CRON_TRIGGER_META: {"job_id": "job-1", "run_id": "run-timeout"}},
|
||||
session_key_override="websocket:chat-1",
|
||||
)
|
||||
|
||||
submit_task = asyncio.create_task(
|
||||
asyncio.wait_for(loop.submit_cron_turn(msg), timeout=0.01)
|
||||
)
|
||||
assert await asyncio.wait_for(loop.bus.consume_inbound(), timeout=0.5) is msg
|
||||
|
||||
with pytest.raises(TimeoutError):
|
||||
await submit_task
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_submitted_local_trigger_turn_reports_pending_until_completed(tmp_path):
|
||||
"""Local triggers remain marked pending while their session turn is in flight."""
|
||||
@@ -1818,3 +1844,53 @@ async def test_injection_cycle_cap_on_error_path():
|
||||
assert result.had_injections is True
|
||||
# Should cap: _MAX_INJECTION_CYCLES drained rounds + 1 final round that breaks
|
||||
assert call_count["n"] == _MAX_INJECTION_CYCLES + 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_accepted_direct_automation_turn_keeps_single_effect_owner():
|
||||
"""Cancelling a submitter after direct admission must not cancel or replay its work."""
|
||||
from nanobot.agent.automation_turns import AutomationTurnAcceptedCancellation
|
||||
from nanobot.agent.cron_turns import CronTurnCoordinator
|
||||
from nanobot.bus.events import InboundMessage
|
||||
from nanobot.cron.session_turns import CRON_TRIGGER_META
|
||||
|
||||
dispatch_started = asyncio.Event()
|
||||
release_dispatch = asyncio.Event()
|
||||
effect_completed = asyncio.Event()
|
||||
effects: list[str] = []
|
||||
|
||||
async def dispatch(msg: InboundMessage) -> None:
|
||||
dispatch_started.set()
|
||||
await release_dispatch.wait()
|
||||
effects.append(msg.content)
|
||||
coordinator.complete(msg)
|
||||
effect_completed.set()
|
||||
|
||||
coordinator = CronTurnCoordinator(
|
||||
publish_inbound=lambda _msg: asyncio.sleep(0),
|
||||
dispatch=dispatch,
|
||||
is_running=lambda: False,
|
||||
)
|
||||
msg = InboundMessage(
|
||||
channel="websocket",
|
||||
sender_id="cron",
|
||||
chat_id="chat-1",
|
||||
content="scheduled work",
|
||||
metadata={CRON_TRIGGER_META: {"job_id": "job-1", "run_id": "run-direct"}},
|
||||
session_key_override="websocket:chat-1",
|
||||
)
|
||||
|
||||
submit_task = asyncio.create_task(coordinator.submit(msg))
|
||||
try:
|
||||
await asyncio.wait_for(dispatch_started.wait(), timeout=0.5)
|
||||
submit_task.cancel()
|
||||
with pytest.raises(AutomationTurnAcceptedCancellation):
|
||||
await submit_task
|
||||
assert effects == []
|
||||
|
||||
release_dispatch.set()
|
||||
await asyncio.wait_for(effect_completed.wait(), timeout=0.5)
|
||||
assert effects == ["scheduled work"]
|
||||
finally:
|
||||
release_dispatch.set()
|
||||
await asyncio.gather(submit_task, return_exceptions=True)
|
||||
|
||||
@@ -3,7 +3,6 @@ import os
|
||||
import subprocess
|
||||
import time
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
@@ -463,11 +462,21 @@ async def test_cli_app_scope_controls_working_dir(
|
||||
|
||||
seen: dict[str, str] = {}
|
||||
|
||||
def fake_run(argv, **kwargs):
|
||||
seen["cwd"] = kwargs["cwd"]
|
||||
return SimpleNamespace(returncode=0, stdout="ok", stderr="")
|
||||
class _Process:
|
||||
returncode = 0
|
||||
|
||||
monkeypatch.setattr("nanobot.apps.cli.service.subprocess.run", fake_run)
|
||||
async def communicate(self) -> tuple[bytes, bytes]:
|
||||
return b"ok", b""
|
||||
|
||||
async def fake_create_subprocess_exec(*_argv: str, **kwargs: object) -> _Process:
|
||||
seen["cwd"] = str(kwargs["cwd"])
|
||||
return _Process()
|
||||
|
||||
monkeypatch.setattr(
|
||||
"nanobot.apps.cli.service.asyncio.create_subprocess_exec",
|
||||
fake_create_subprocess_exec,
|
||||
)
|
||||
monkeypatch.setattr("nanobot.apps.cli.service._WindowsJob.create", lambda: None)
|
||||
tool = CliAppsTool(
|
||||
workspace=tmp_path,
|
||||
restrict_to_workspace=True,
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import threading
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
@@ -21,9 +22,9 @@ from nanobot.agent.tools.long_task import (
|
||||
from nanobot.agent.tools.registry import ToolRegistry
|
||||
from nanobot.bus.outbound_events import GoalStateSyncEvent
|
||||
from nanobot.bus.queue import MessageBus
|
||||
from nanobot.bus.runtime_events import RuntimeEventBus
|
||||
from nanobot.bus.runtime_events import GoalStateChanged, RuntimeEventBus
|
||||
from nanobot.session.goal_state import GOAL_STATE_KEY, MAX_GOAL_OBJECTIVE_CHARS
|
||||
from nanobot.session.manager import SessionManager
|
||||
from nanobot.session.manager import Session, SessionManager
|
||||
from nanobot.session.turn_continuation import should_finalize_on_max_iterations
|
||||
from nanobot.session.webui_turns import WebuiTurnCoordinator
|
||||
|
||||
@@ -179,13 +180,13 @@ async def test_goal_state_mutations_roll_back_on_save_failure(tmp_path, monkeypa
|
||||
sess = sm.get_or_create("websocket:c1")
|
||||
sess.metadata["marker"] = {"keep": True}
|
||||
sess.metadata["_sustained_goal_continuation_rounds"] = 12
|
||||
original_save = sm.save
|
||||
original_save_async = sm.save_async
|
||||
create_context = _request_context()
|
||||
|
||||
def fail_save(_session, **_kwargs):
|
||||
async def fail_save(_session, **_kwargs):
|
||||
raise OSError("disk unavailable")
|
||||
|
||||
monkeypatch.setattr(sm, "save", fail_save)
|
||||
monkeypatch.setattr(sm, "save_async", fail_save)
|
||||
with pytest.raises(OSError, match="disk unavailable"):
|
||||
await _execute(create, create_context, objective="Old")
|
||||
|
||||
@@ -195,13 +196,13 @@ async def test_goal_state_mutations_roll_back_on_save_failure(tmp_path, monkeypa
|
||||
}
|
||||
assert GOAL_STATE_KEY not in SessionManager(tmp_path).get_or_create("websocket:c1").metadata
|
||||
|
||||
monkeypatch.setattr(sm, "save", original_save)
|
||||
monkeypatch.setattr(sm, "save_async", original_save_async)
|
||||
assert "Goal recorded" in await _execute(create, create_context, objective="Old")
|
||||
sess.metadata["_sustained_goal_continuation_rounds"] = 12
|
||||
sm.save(sess)
|
||||
replace_context = _request_context()
|
||||
|
||||
monkeypatch.setattr(sm, "save", fail_save)
|
||||
monkeypatch.setattr(sm, "save_async", fail_save)
|
||||
with pytest.raises(OSError, match="disk unavailable"):
|
||||
await _execute(update, replace_context, action="replace", objective="New")
|
||||
|
||||
@@ -211,7 +212,7 @@ async def test_goal_state_mutations_roll_back_on_save_failure(tmp_path, monkeypa
|
||||
assert persisted[GOAL_STATE_KEY]["objective"] == "Old"
|
||||
assert persisted["_sustained_goal_continuation_rounds"] == 12
|
||||
|
||||
monkeypatch.setattr(sm, "save", original_save)
|
||||
monkeypatch.setattr(sm, "save_async", original_save_async)
|
||||
assert "Goal replaced" in await _execute(
|
||||
update,
|
||||
replace_context,
|
||||
@@ -225,6 +226,57 @@ async def test_goal_state_mutations_roll_back_on_save_failure(tmp_path, monkeypa
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sync_goal_save_cancellation_settles_state_and_runtime_event() -> None:
|
||||
started = threading.Event()
|
||||
release = threading.Event()
|
||||
session = Session(key="websocket:c1")
|
||||
durable_metadata: dict[str, object] = {}
|
||||
save_calls = 0
|
||||
|
||||
class _SyncSessionManager:
|
||||
def get_or_create(self, key: str) -> Session:
|
||||
assert key == session.key
|
||||
return session
|
||||
|
||||
def save(self, target: Session) -> None:
|
||||
nonlocal durable_metadata, save_calls
|
||||
save_calls += 1
|
||||
started.set()
|
||||
assert release.wait(timeout=1)
|
||||
durable_metadata = dict(target.metadata)
|
||||
|
||||
runtime_events = RuntimeEventBus()
|
||||
published: list[GoalStateChanged] = []
|
||||
runtime_events.subscribe(published.append, GoalStateChanged)
|
||||
create = CreateGoalTool(
|
||||
sessions=_SyncSessionManager(), # type: ignore[arg-type]
|
||||
runtime_events=runtime_events,
|
||||
)
|
||||
task = asyncio.create_task(
|
||||
_execute(create, _request_context(), objective="Persist through cancellation")
|
||||
)
|
||||
assert await asyncio.to_thread(started.wait, 1)
|
||||
try:
|
||||
task.cancel()
|
||||
await asyncio.sleep(0)
|
||||
assert not task.done()
|
||||
finally:
|
||||
release.set()
|
||||
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await asyncio.wait_for(task, timeout=1)
|
||||
|
||||
assert save_calls == 1
|
||||
assert session.metadata[GOAL_STATE_KEY]["objective"] == "Persist through cancellation"
|
||||
assert durable_metadata[GOAL_STATE_KEY] == session.metadata[GOAL_STATE_KEY]
|
||||
assert len(published) == 1
|
||||
assert published[0].session_metadata[GOAL_STATE_KEY] == session.metadata[GOAL_STATE_KEY]
|
||||
await asyncio.sleep(0.05)
|
||||
assert save_calls == 1
|
||||
assert len(published) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_goal_tools_reject_oversized_objectives(tmp_path):
|
||||
sm = SessionManager(tmp_path)
|
||||
|
||||
@@ -42,17 +42,21 @@ def test_run_passes_filtered_env(monkeypatch, tmp_path) -> None:
|
||||
manager = CliAppManager(workspace=tmp_path, data_dir=tmp_path / "cli-apps")
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
def fake_run(*args, **kwargs):
|
||||
class _Process:
|
||||
returncode = 0
|
||||
|
||||
def communicate(self, *, timeout: int) -> tuple[str, str]:
|
||||
assert timeout == 60
|
||||
return "ok", ""
|
||||
|
||||
def fake_popen(argv: list[str], **kwargs: object) -> _Process:
|
||||
assert argv == ["/bin/echo", "hi"]
|
||||
assert kwargs["text"] is True
|
||||
captured.update(kwargs)
|
||||
return _Process()
|
||||
|
||||
class Result:
|
||||
returncode = 0
|
||||
stdout = "ok"
|
||||
stderr = ""
|
||||
|
||||
return Result()
|
||||
|
||||
monkeypatch.setattr("nanobot.apps.cli.service.subprocess.run", fake_run)
|
||||
monkeypatch.setattr("nanobot.apps.cli.service.subprocess.Popen", fake_popen)
|
||||
monkeypatch.setattr("nanobot.apps.cli.service._WindowsJob.create", lambda: None)
|
||||
monkeypatch.setattr(manager, "get_app", lambda name: {"name": name, "entry_point": "echo"})
|
||||
monkeypatch.setattr(
|
||||
manager,
|
||||
|
||||
@@ -8,10 +8,15 @@ block the stop.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import threading
|
||||
import time
|
||||
from contextlib import suppress
|
||||
|
||||
from nanobot.cli.gateway_runtime import _close_gateway_runtime
|
||||
from nanobot.cli.gateway_runtime import (
|
||||
_call_session_manager,
|
||||
_close_gateway_runtime,
|
||||
_monitor_event_loop_lag,
|
||||
)
|
||||
|
||||
|
||||
class _FakeAgent:
|
||||
@@ -216,3 +221,110 @@ async def test_cancelled_runtime_tasks_gather_does_not_raise() -> None:
|
||||
assert runtime_tasks.done() # the cancelled gather was awaited without raising
|
||||
assert agent.close_calls == 1
|
||||
assert provider.close_calls == 1
|
||||
|
||||
|
||||
async def test_session_manager_call_prefers_class_declared_coroutine() -> None:
|
||||
class _Manager:
|
||||
async def list_sessions_async(self) -> list[str]:
|
||||
return ["native-async"]
|
||||
|
||||
def list_sessions(self) -> list[str]:
|
||||
raise AssertionError("sync fallback should not run")
|
||||
|
||||
manager = _Manager()
|
||||
|
||||
result = await _call_session_manager(
|
||||
manager,
|
||||
"list_sessions_async",
|
||||
manager.list_sessions,
|
||||
)
|
||||
|
||||
assert result == ["native-async"]
|
||||
|
||||
|
||||
async def test_session_manager_call_offloads_sync_compatibility_fallback() -> None:
|
||||
calling_thread = threading.get_ident()
|
||||
sync_threads: list[int] = []
|
||||
|
||||
class _Manager:
|
||||
def list_sessions(self) -> list[str]:
|
||||
sync_threads.append(threading.get_ident())
|
||||
return ["sync-fallback"]
|
||||
|
||||
manager = _Manager()
|
||||
|
||||
async def _fabricated_async() -> list[str]:
|
||||
raise AssertionError("instance-only async stand-in should not run")
|
||||
|
||||
setattr(manager, "list_sessions_async", _fabricated_async)
|
||||
result = await _call_session_manager(
|
||||
manager,
|
||||
"list_sessions_async",
|
||||
manager.list_sessions,
|
||||
)
|
||||
|
||||
assert result == ["sync-fallback"]
|
||||
assert sync_threads and sync_threads[0] != calling_thread
|
||||
|
||||
|
||||
async def test_session_manager_sync_fallback_cancellation_waits_for_worker() -> None:
|
||||
started = threading.Event()
|
||||
release = threading.Event()
|
||||
finished = threading.Event()
|
||||
mutations: list[str] = []
|
||||
|
||||
class _Manager:
|
||||
def save(self, value: str) -> str:
|
||||
started.set()
|
||||
assert release.wait(timeout=1)
|
||||
mutations.append(value)
|
||||
finished.set()
|
||||
return "saved"
|
||||
|
||||
manager = _Manager()
|
||||
task = asyncio.create_task(
|
||||
_call_session_manager(manager, "save_async", manager.save, "mutation")
|
||||
)
|
||||
assert await asyncio.to_thread(started.wait, 1)
|
||||
try:
|
||||
task.cancel()
|
||||
await asyncio.sleep(0)
|
||||
assert not task.done()
|
||||
finally:
|
||||
release.set()
|
||||
|
||||
with suppress(asyncio.CancelledError):
|
||||
await asyncio.wait_for(task, timeout=1)
|
||||
|
||||
assert task.cancelled()
|
||||
assert finished.is_set()
|
||||
assert mutations == ["mutation"]
|
||||
await asyncio.sleep(0.05)
|
||||
assert mutations == ["mutation"]
|
||||
|
||||
|
||||
async def test_event_loop_lag_monitor_logs_gateway_scheduler_drift() -> None:
|
||||
records: list[str] = []
|
||||
|
||||
class _Logger:
|
||||
def warning(self, message: str, *args: object) -> None:
|
||||
records.append(message.format(*args))
|
||||
|
||||
task = asyncio.create_task(
|
||||
_monitor_event_loop_lag(
|
||||
interval_s=0.01,
|
||||
warning_threshold_s=0.015,
|
||||
log=_Logger(),
|
||||
)
|
||||
)
|
||||
await asyncio.sleep(0)
|
||||
time.sleep(0.04)
|
||||
await asyncio.sleep(0.02)
|
||||
task.cancel()
|
||||
with suppress(asyncio.CancelledError):
|
||||
await task
|
||||
|
||||
assert records
|
||||
assert "operation=gateway" in records[0]
|
||||
assert "duration_ms=" in records[0]
|
||||
assert "interval_ms=10" in records[0]
|
||||
|
||||
+168
-12
@@ -1,9 +1,11 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from contextlib import suppress
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
@@ -944,19 +946,23 @@ def test_run_installed_cli_uses_argv_without_shell(
|
||||
lambda entry: resolved if entry == "cli-anything-gimp" else None,
|
||||
)
|
||||
|
||||
def fake_run(argv: list[str], **kwargs: object) -> subprocess.CompletedProcess[str]:
|
||||
assert "shell" not in kwargs or kwargs["shell"] is False
|
||||
class _Process:
|
||||
returncode = 0
|
||||
|
||||
def communicate(self, *, timeout: int) -> tuple[str, str]:
|
||||
assert timeout == 5
|
||||
return "ARGS=['--json', 'project', 'list']", ""
|
||||
|
||||
def fake_popen(argv: list[str], **kwargs: object) -> _Process:
|
||||
assert argv == [resolved, "--json", "project", "list"]
|
||||
assert "shell" not in kwargs
|
||||
assert kwargs["text"] is True
|
||||
assert kwargs["encoding"] == "utf-8"
|
||||
assert kwargs["errors"] == "replace"
|
||||
return subprocess.CompletedProcess(
|
||||
argv,
|
||||
0,
|
||||
stdout="ARGS=" + repr(argv[1:]),
|
||||
stderr="",
|
||||
)
|
||||
return _Process()
|
||||
|
||||
monkeypatch.setattr("nanobot.apps.cli.service.subprocess.run", fake_run)
|
||||
monkeypatch.setattr("nanobot.apps.cli.service.subprocess.Popen", fake_popen)
|
||||
monkeypatch.setattr("nanobot.apps.cli.service._WindowsJob.create", lambda: None)
|
||||
manager._save_installed(
|
||||
{
|
||||
"gimp": {
|
||||
@@ -986,12 +992,21 @@ def test_run_reports_created_artifacts(
|
||||
lambda entry: resolved if entry == "cli-anything-gimp" else None,
|
||||
)
|
||||
|
||||
def fake_run(argv: list[str], **kwargs: object) -> subprocess.CompletedProcess[str]:
|
||||
class _Process:
|
||||
returncode = 0
|
||||
|
||||
def communicate(self, *, timeout: int) -> tuple[str, str]:
|
||||
assert timeout == 5
|
||||
return "done", ""
|
||||
|
||||
def fake_popen(argv: list[str], **kwargs: object) -> _Process:
|
||||
assert argv == [resolved, "render"]
|
||||
cwd = Path(str(kwargs["cwd"]))
|
||||
(cwd / "diagram.png").write_bytes(b"\x89PNG\r\n\x1a\nimage")
|
||||
return subprocess.CompletedProcess(argv, 0, stdout="done", stderr="")
|
||||
return _Process()
|
||||
|
||||
monkeypatch.setattr("nanobot.apps.cli.service.subprocess.run", fake_run)
|
||||
monkeypatch.setattr("nanobot.apps.cli.service.subprocess.Popen", fake_popen)
|
||||
monkeypatch.setattr("nanobot.apps.cli.service._WindowsJob.create", lambda: None)
|
||||
manager._save_installed({"gimp": {"entry_point": "cli-anything-gimp"}})
|
||||
|
||||
result = manager.run("gimp", ["render"])
|
||||
@@ -1100,3 +1115,144 @@ def test_uninstall_uses_uv_pip_when_pip_unavailable(
|
||||
sys.executable,
|
||||
"suno-cli",
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_async_keeps_event_loop_responsive(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
manager = _manager(tmp_path)
|
||||
_seed_catalog(manager)
|
||||
manager._save_installed({"gimp": {"entry_point": "python"}})
|
||||
monkeypatch.setattr(
|
||||
"nanobot.apps.cli.service.shutil.which",
|
||||
lambda entry: sys.executable if entry == "python" else None,
|
||||
)
|
||||
|
||||
task = asyncio.create_task(
|
||||
manager.run_async(
|
||||
"gimp",
|
||||
["-c", "import time; time.sleep(0.08); print('done')"],
|
||||
)
|
||||
)
|
||||
for _ in range(3):
|
||||
await asyncio.sleep(0.01)
|
||||
assert not task.done()
|
||||
|
||||
result = await asyncio.wait_for(task, timeout=2)
|
||||
assert "CLI app 'gimp' exited 0" in result
|
||||
assert "done" in result
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_async_cancellation_terminates_process_tree(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
manager = _manager(tmp_path)
|
||||
_seed_catalog(manager)
|
||||
manager._save_installed({"gimp": {"entry_point": "python"}})
|
||||
monkeypatch.setattr(
|
||||
"nanobot.apps.cli.service.shutil.which",
|
||||
lambda entry: sys.executable if entry == "python" else None,
|
||||
)
|
||||
started = manager.workspace / "parent-started.txt"
|
||||
orphan_marker = manager.workspace / "orphan-survived.txt"
|
||||
child_code = (
|
||||
"import time; from pathlib import Path; time.sleep(0.5); "
|
||||
f"Path({str(orphan_marker)!r}).write_text('orphan', encoding='utf-8')"
|
||||
)
|
||||
parent_code = (
|
||||
"import subprocess, sys, time; from pathlib import Path; "
|
||||
f"subprocess.Popen([sys.executable, '-c', {child_code!r}]); "
|
||||
f"Path({str(started)!r}).write_text('started', encoding='utf-8'); "
|
||||
"time.sleep(30)"
|
||||
)
|
||||
task = asyncio.create_task(manager.run_async("gimp", ["-c", parent_code]))
|
||||
try:
|
||||
for _ in range(200):
|
||||
if started.exists():
|
||||
break
|
||||
await asyncio.sleep(0.01)
|
||||
assert started.exists()
|
||||
|
||||
task.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await asyncio.wait_for(task, timeout=3)
|
||||
await asyncio.sleep(0.8)
|
||||
assert not orphan_marker.exists()
|
||||
finally:
|
||||
if not task.done():
|
||||
task.cancel()
|
||||
with suppress(asyncio.CancelledError):
|
||||
await task
|
||||
|
||||
|
||||
def _parent_exits_with_pipe_holding_descendant(started: Path, marker: Path) -> str:
|
||||
child_code = (
|
||||
"import time; from pathlib import Path; time.sleep(1.5); "
|
||||
f"Path({str(marker)!r}).write_text('survived', encoding='utf-8')"
|
||||
)
|
||||
return (
|
||||
"import subprocess, sys, time; from pathlib import Path; "
|
||||
"time.sleep(0.25); "
|
||||
f"subprocess.Popen([sys.executable, '-c', {child_code!r}]); "
|
||||
f"Path({str(started)!r}).write_text('spawned', encoding='utf-8')"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_async_timeout_kills_pipe_holding_descendant_after_parent_exit(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
manager = _manager(tmp_path)
|
||||
_seed_catalog(manager)
|
||||
manager._save_installed({"gimp": {"entry_point": "python"}})
|
||||
monkeypatch.setattr(
|
||||
"nanobot.apps.cli.service.shutil.which",
|
||||
lambda entry: sys.executable if entry == "python" else None,
|
||||
)
|
||||
started = manager.workspace / "descendant-started.txt"
|
||||
marker = manager.workspace / "descendant-survived.txt"
|
||||
parent_code = _parent_exits_with_pipe_holding_descendant(started, marker)
|
||||
|
||||
before = time.monotonic()
|
||||
result = await asyncio.wait_for(
|
||||
manager.run_async("gimp", ["-c", parent_code], timeout=1),
|
||||
timeout=4,
|
||||
)
|
||||
elapsed = time.monotonic() - before
|
||||
|
||||
assert started.exists()
|
||||
assert result == "CLI app 'gimp' timed out after 1s"
|
||||
assert elapsed < 4
|
||||
await asyncio.sleep(0.9)
|
||||
assert not marker.exists()
|
||||
|
||||
|
||||
def test_run_timeout_kills_pipe_holding_descendant_after_parent_exit(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
manager = _manager(tmp_path)
|
||||
_seed_catalog(manager)
|
||||
manager._save_installed({"gimp": {"entry_point": "python"}})
|
||||
monkeypatch.setattr(
|
||||
"nanobot.apps.cli.service.shutil.which",
|
||||
lambda entry: sys.executable if entry == "python" else None,
|
||||
)
|
||||
started = manager.workspace / "sync-descendant-started.txt"
|
||||
marker = manager.workspace / "sync-descendant-survived.txt"
|
||||
parent_code = _parent_exits_with_pipe_holding_descendant(started, marker)
|
||||
|
||||
before = time.monotonic()
|
||||
result = manager.run("gimp", ["-c", parent_code], timeout=1)
|
||||
elapsed = time.monotonic() - before
|
||||
|
||||
assert started.exists()
|
||||
assert result == "CLI app 'gimp' timed out after 1s"
|
||||
assert elapsed < 4
|
||||
time.sleep(0.9)
|
||||
assert not marker.exists()
|
||||
|
||||
+106
-10
@@ -2,10 +2,13 @@ from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import subprocess
|
||||
import threading
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from nanobot.agent.tools import cli_apps as cli_apps_tool
|
||||
from nanobot.agent.tools.cli_apps import CliAppsTool
|
||||
from nanobot.agent.tools.context import RequestContext
|
||||
from nanobot.apps.cli.service import CliAppManager, CliAppsRuntimeConfig
|
||||
@@ -52,16 +55,22 @@ def test_run_cli_app_uses_installed_registry_app(
|
||||
lambda entry: resolved if entry == "cli-anything-gimp" else None,
|
||||
)
|
||||
|
||||
def fake_run(argv: list[str], **kwargs: object) -> subprocess.CompletedProcess[str]:
|
||||
assert "shell" not in kwargs or kwargs["shell"] is False
|
||||
return subprocess.CompletedProcess(
|
||||
argv,
|
||||
0,
|
||||
stdout="tool:" + " ".join(argv[1:]),
|
||||
stderr="",
|
||||
)
|
||||
class _Process:
|
||||
returncode = 0
|
||||
|
||||
monkeypatch.setattr("nanobot.apps.cli.service.subprocess.run", fake_run)
|
||||
async def communicate(self) -> tuple[bytes, bytes]:
|
||||
return b"tool:--json project list", b""
|
||||
|
||||
async def fake_create_subprocess_exec(*argv: str, **kwargs: object) -> _Process:
|
||||
assert argv == (resolved, "--json", "project", "list")
|
||||
assert "shell" not in kwargs
|
||||
return _Process()
|
||||
|
||||
monkeypatch.setattr(
|
||||
"nanobot.apps.cli.service.asyncio.create_subprocess_exec",
|
||||
fake_create_subprocess_exec,
|
||||
)
|
||||
monkeypatch.setattr("nanobot.apps.cli.service._WindowsJob.create", lambda: None)
|
||||
monkeypatch.setattr("nanobot.apps.cli.service.get_runtime_subdir", lambda _name: data_dir)
|
||||
|
||||
tool = CliAppsTool(
|
||||
@@ -156,3 +165,90 @@ def test_cli_app_tool_provides_context_only_for_attachment(tmp_path: Path) -> No
|
||||
assert attached is not None
|
||||
assert attached.source == "cli_apps"
|
||||
assert "CLI App Attachment: @drawio" in attached.content
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_cli_app_uses_threaded_sync_fallback_for_compatible_manager(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
event_loop_thread = threading.get_ident()
|
||||
seen: dict[str, object] = {}
|
||||
|
||||
class CompatibleManager:
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
workspace: Path,
|
||||
runtime: CliAppsRuntimeConfig,
|
||||
) -> None:
|
||||
seen["workspace"] = workspace
|
||||
seen["runtime"] = runtime
|
||||
|
||||
def run_async(self, *_args: object, **_kwargs: object) -> str:
|
||||
raise AssertionError("a synchronous compatibility method must not be awaited")
|
||||
|
||||
def run(self, name: str, **kwargs: object) -> str:
|
||||
seen["thread"] = threading.get_ident()
|
||||
seen["name"] = name
|
||||
seen["kwargs"] = kwargs
|
||||
return "compatible result"
|
||||
|
||||
monkeypatch.setattr(cli_apps_tool, "CliAppManager", CompatibleManager)
|
||||
runtime = CliAppsRuntimeConfig(run_timeout=7)
|
||||
tool = CliAppsTool(workspace=tmp_path, restrict_to_workspace=True, runtime=runtime)
|
||||
|
||||
result = await tool.execute(
|
||||
name="demo",
|
||||
args=["show"],
|
||||
json=True,
|
||||
working_dir=str(tmp_path),
|
||||
timeout=3,
|
||||
)
|
||||
|
||||
assert result == "compatible result"
|
||||
assert seen["workspace"] == tmp_path
|
||||
assert seen["runtime"] is runtime
|
||||
assert seen["thread"] != event_loop_thread
|
||||
assert seen["name"] == "demo"
|
||||
assert seen["kwargs"] == {
|
||||
"args": ["show"],
|
||||
"json_output": True,
|
||||
"working_dir": str(tmp_path),
|
||||
"timeout": 3,
|
||||
"restrict_to_workspace": True,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_cli_app_preserves_native_manager_cancellation(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
started = asyncio.Event()
|
||||
cancelled = asyncio.Event()
|
||||
|
||||
async def fake_run_async(self: CliAppManager, name: str, **_kwargs: object) -> str:
|
||||
assert name == "demo"
|
||||
started.set()
|
||||
try:
|
||||
await asyncio.Future()
|
||||
except asyncio.CancelledError:
|
||||
cancelled.set()
|
||||
raise
|
||||
raise AssertionError("unreachable")
|
||||
|
||||
def fail_sync_run(self: CliAppManager, name: str, **_kwargs: object) -> str:
|
||||
raise AssertionError(f"native async manager fell back to run({name!r})")
|
||||
|
||||
monkeypatch.setattr(CliAppManager, "run_async", fake_run_async)
|
||||
monkeypatch.setattr(CliAppManager, "run", fail_sync_run)
|
||||
tool = CliAppsTool(workspace=tmp_path)
|
||||
|
||||
task = asyncio.create_task(tool.execute(name="demo"))
|
||||
await asyncio.wait_for(started.wait(), timeout=1)
|
||||
task.cancel()
|
||||
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
assert cancelled.is_set()
|
||||
|
||||
@@ -2,6 +2,8 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import threading
|
||||
from inspect import Parameter, signature
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
@@ -9,9 +11,11 @@ import pytest
|
||||
|
||||
from nanobot.command.builtin import (
|
||||
builtin_command_starts_agent_turn,
|
||||
cmd_new,
|
||||
register_builtin_commands,
|
||||
)
|
||||
from nanobot.command.router import CommandContext, CommandRouter
|
||||
from nanobot.session.manager import Session
|
||||
|
||||
|
||||
def test_command_context_requires_loop_as_keyword_dependency() -> None:
|
||||
@@ -95,6 +99,55 @@ def test_builtin_command_agent_turn_lifecycle(content: str, expected: bool) -> N
|
||||
assert builtin_command_starts_agent_turn(content) is expected
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_new_cancellation_waits_for_save_and_cache_invalidation() -> None:
|
||||
started = threading.Event()
|
||||
release = threading.Event()
|
||||
session = Session(key="test:chat1")
|
||||
invalidated: list[str] = []
|
||||
save_calls = 0
|
||||
|
||||
class _SyncSessionManager:
|
||||
def save(self, target: Session) -> None:
|
||||
nonlocal save_calls
|
||||
assert target is session
|
||||
save_calls += 1
|
||||
started.set()
|
||||
assert release.wait(timeout=1)
|
||||
|
||||
def invalidate(self, key: str) -> None:
|
||||
invalidated.append(key)
|
||||
|
||||
loop = MagicMock()
|
||||
loop.sessions = _SyncSessionManager()
|
||||
loop._cancel_active_tasks = AsyncMock(return_value=0)
|
||||
ctx = CommandContext(
|
||||
msg=MagicMock(channel="test", chat_id="chat1", metadata={}),
|
||||
session=session,
|
||||
key=session.key,
|
||||
raw="/new",
|
||||
loop=loop,
|
||||
)
|
||||
task = asyncio.create_task(cmd_new(ctx))
|
||||
assert await asyncio.to_thread(started.wait, 1)
|
||||
try:
|
||||
task.cancel()
|
||||
await asyncio.sleep(0)
|
||||
assert not task.done()
|
||||
assert invalidated == []
|
||||
finally:
|
||||
release.set()
|
||||
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await asyncio.wait_for(task, timeout=1)
|
||||
|
||||
assert save_calls == 1
|
||||
assert invalidated == [session.key]
|
||||
await asyncio.sleep(0.05)
|
||||
assert save_calls == 1
|
||||
assert invalidated == [session.key]
|
||||
|
||||
|
||||
class TestMidTurnCommandDispatchedDirectly:
|
||||
"""Verify that commands matching is_dispatchable_command() are dispatched
|
||||
correctly when session=None (the mid-turn path)."""
|
||||
|
||||
@@ -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
|
||||
@@ -1,5 +1,6 @@
|
||||
import asyncio
|
||||
import json
|
||||
import threading
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
@@ -653,6 +654,231 @@ async def test_manual_run_persists_completion_when_callback_lists_jobs(tmp_path)
|
||||
assert state["runHistory"][0]["status"] == "ok"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_same_job_manual_runs_do_not_overlap(tmp_path) -> None:
|
||||
entered = asyncio.Event()
|
||||
release = asyncio.Event()
|
||||
calls: list[str] = []
|
||||
|
||||
async def on_job(job: CronJob) -> None:
|
||||
calls.append(job.id)
|
||||
entered.set()
|
||||
await release.wait()
|
||||
|
||||
service = CronService(tmp_path / "cron" / "jobs.json", on_job=on_job)
|
||||
job = service.add_job(
|
||||
name="single-flight",
|
||||
schedule=CronSchedule(kind="every", every_ms=60_000),
|
||||
message="hello",
|
||||
**_bound_chat(),
|
||||
)
|
||||
|
||||
first = asyncio.create_task(service.run_job(job.id))
|
||||
try:
|
||||
await asyncio.wait_for(entered.wait(), timeout=1)
|
||||
assert await service.run_job(job.id) is False
|
||||
assert calls == [job.id]
|
||||
finally:
|
||||
release.set()
|
||||
assert await first is True
|
||||
|
||||
loaded = service.get_job(job.id)
|
||||
assert loaded is not None
|
||||
assert len(loaded.state.run_history) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_timer_and_manual_run_of_same_job_do_not_overlap(tmp_path) -> None:
|
||||
entered = asyncio.Event()
|
||||
release = asyncio.Event()
|
||||
calls: list[str] = []
|
||||
|
||||
async def on_job(job: CronJob) -> None:
|
||||
calls.append(job.id)
|
||||
entered.set()
|
||||
await release.wait()
|
||||
|
||||
service = CronService(tmp_path / "cron" / "jobs.json", on_job=on_job)
|
||||
service._running = True
|
||||
service._arm_timer = lambda: None
|
||||
service._load_store()
|
||||
job = service.add_job(
|
||||
name="timer-single-flight",
|
||||
schedule=CronSchedule(kind="every", every_ms=60_000),
|
||||
message="hello",
|
||||
**_bound_chat(),
|
||||
)
|
||||
job.state.next_run_at_ms = max(1, int(time.time() * 1000) - 1_000)
|
||||
service._save_store()
|
||||
|
||||
timer = asyncio.create_task(service._on_timer())
|
||||
try:
|
||||
await asyncio.wait_for(entered.wait(), timeout=1)
|
||||
assert await service.run_job(job.id, force=True) is False
|
||||
assert calls == [job.id]
|
||||
finally:
|
||||
release.set()
|
||||
await timer
|
||||
service.stop()
|
||||
|
||||
loaded = service.get_job(job.id)
|
||||
assert loaded is not None
|
||||
assert len(loaded.state.run_history) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_start_cancellation_waits_for_load_and_fully_starts(
|
||||
tmp_path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
store_path = tmp_path / "cron" / "jobs.json"
|
||||
service = CronService(store_path)
|
||||
job = service.add_job(
|
||||
name="load-settlement",
|
||||
schedule=CronSchedule(kind="every", every_ms=60_000),
|
||||
message="hello",
|
||||
**_bound_chat(),
|
||||
)
|
||||
entered = threading.Event()
|
||||
release = threading.Event()
|
||||
original_load_store = service._load_store
|
||||
|
||||
def blocking_load_store():
|
||||
entered.set()
|
||||
if not release.wait(timeout=2):
|
||||
raise TimeoutError("test did not release cron load worker")
|
||||
return original_load_store()
|
||||
|
||||
monkeypatch.setattr(service, "_load_store", blocking_load_store)
|
||||
start = asyncio.create_task(service.start())
|
||||
assert await asyncio.to_thread(entered.wait, 1)
|
||||
|
||||
try:
|
||||
assert start.cancel()
|
||||
await asyncio.sleep(0)
|
||||
assert not start.done()
|
||||
assert not store_path.exists()
|
||||
|
||||
release.set()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await asyncio.wait_for(start, timeout=1)
|
||||
|
||||
assert service._running is True
|
||||
timer = service._timer_task
|
||||
assert timer is not None and not timer.done()
|
||||
assert service._store_dirty is False
|
||||
assert (tmp_path / "cron" / "action.jsonl").read_text(encoding="utf-8") == ""
|
||||
persisted = store_path.read_bytes()
|
||||
assert json.loads(persisted)["jobs"][0]["id"] == job.id
|
||||
|
||||
await asyncio.sleep(0.01)
|
||||
assert store_path.read_bytes() == persisted
|
||||
assert service._timer_task is timer
|
||||
finally:
|
||||
release.set()
|
||||
await asyncio.gather(start, return_exceptions=True)
|
||||
timer = service._timer_task
|
||||
service.stop()
|
||||
if timer is not None:
|
||||
await asyncio.gather(timer, return_exceptions=True)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_start_repeated_cancellation_waits_for_persistence_and_fully_starts(
|
||||
tmp_path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
store_path = tmp_path / "cron" / "jobs.json"
|
||||
service = CronService(store_path)
|
||||
entered = threading.Event()
|
||||
release = threading.Event()
|
||||
original_save_store = service._save_store
|
||||
|
||||
def blocking_save_store() -> None:
|
||||
entered.set()
|
||||
if not release.wait(timeout=2):
|
||||
raise TimeoutError("test did not release cron persistence worker")
|
||||
original_save_store()
|
||||
|
||||
monkeypatch.setattr(service, "_save_store", blocking_save_store)
|
||||
start = asyncio.create_task(service.start())
|
||||
assert await asyncio.to_thread(entered.wait, 1)
|
||||
|
||||
try:
|
||||
assert start.cancel()
|
||||
await asyncio.sleep(0)
|
||||
assert not start.done()
|
||||
assert start.cancel()
|
||||
await asyncio.sleep(0)
|
||||
assert not start.done()
|
||||
assert not store_path.exists()
|
||||
|
||||
release.set()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await asyncio.wait_for(start, timeout=1)
|
||||
|
||||
assert service._running is True
|
||||
timer = service._timer_task
|
||||
assert timer is not None and not timer.done()
|
||||
assert service._store_dirty is False
|
||||
persisted = store_path.read_bytes()
|
||||
assert json.loads(persisted) == {"version": 1, "jobs": []}
|
||||
|
||||
await asyncio.sleep(0.01)
|
||||
assert store_path.read_bytes() == persisted
|
||||
assert service._timer_task is timer
|
||||
finally:
|
||||
release.set()
|
||||
await asyncio.gather(start, return_exceptions=True)
|
||||
timer = service._timer_task
|
||||
service.stop()
|
||||
if timer is not None:
|
||||
await asyncio.gather(timer, return_exceptions=True)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_sync_cancellation_drains_worker_under_transaction_lock(tmp_path) -> None:
|
||||
service = CronService(tmp_path / "cron" / "jobs.json")
|
||||
entered = threading.Event()
|
||||
release = threading.Event()
|
||||
second_entered = threading.Event()
|
||||
mutations: list[str] = []
|
||||
|
||||
def blocking_mutation() -> str:
|
||||
entered.set()
|
||||
if not release.wait(timeout=2):
|
||||
raise TimeoutError("test did not release cron worker")
|
||||
mutations.append("first")
|
||||
return "first"
|
||||
|
||||
def later_mutation() -> str:
|
||||
second_entered.set()
|
||||
mutations.append("second")
|
||||
return "second"
|
||||
|
||||
first = asyncio.create_task(service.run_sync(blocking_mutation))
|
||||
assert await asyncio.to_thread(entered.wait, 1)
|
||||
first.cancel()
|
||||
second = asyncio.create_task(service.run_sync(later_mutation))
|
||||
try:
|
||||
await asyncio.sleep(0.05)
|
||||
assert not first.done()
|
||||
assert not second_entered.is_set()
|
||||
|
||||
release.set()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await asyncio.wait_for(first, timeout=1)
|
||||
assert mutations == ["first"]
|
||||
|
||||
assert await asyncio.wait_for(second, timeout=1) == "second"
|
||||
assert mutations == ["first", "second"]
|
||||
await asyncio.sleep(0.01)
|
||||
assert mutations == ["first", "second"]
|
||||
finally:
|
||||
release.set()
|
||||
await asyncio.gather(first, second, return_exceptions=True)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_overlapping_manual_runs_preserve_stopped_service_state(tmp_path) -> None:
|
||||
store_path = tmp_path / "cron" / "jobs.json"
|
||||
|
||||
@@ -9,6 +9,7 @@ and tightens the runtime error for ``add`` without ``message``.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Iterator
|
||||
|
||||
import pytest
|
||||
@@ -21,6 +22,9 @@ from nanobot.agent.tools.registry import ToolRegistry
|
||||
class _SvcStub:
|
||||
"""Minimal CronService stand-in; we only exercise schema/dispatch paths."""
|
||||
|
||||
async def run_sync(self, operation, /, *args, **kwargs):
|
||||
return await asyncio.to_thread(operation, *args, **kwargs)
|
||||
|
||||
def list_jobs(self):
|
||||
return []
|
||||
|
||||
@@ -74,8 +78,6 @@ class TestSchemaContract:
|
||||
# Schema permits omitting message; the runtime must return a message
|
||||
# that tells the LLM exactly what's missing and how to retry, so it
|
||||
# doesn't loop like #3113 reports.
|
||||
import asyncio
|
||||
|
||||
tool = registry._tools["cron"] # type: ignore[attr-defined]
|
||||
out = asyncio.run(tool.execute(action="add", at="2030-01-01T00:00:00"))
|
||||
assert "message" in out
|
||||
|
||||
@@ -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)
|
||||
@@ -2,13 +2,18 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import ipaddress
|
||||
import socket
|
||||
import threading
|
||||
import time
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from nanobot.security.network import (
|
||||
async_resolve_url_target,
|
||||
async_validate_resolved_url,
|
||||
configure_ssrf_whitelist,
|
||||
contains_internal_url,
|
||||
env_proxy_applies_to_url,
|
||||
@@ -362,3 +367,78 @@ def test_whitelist_allows_ipv6_mapped_cgnat():
|
||||
assert ok, f"Whitelisted IPv6-mapped CGNAT should be allowed, got: {err}"
|
||||
finally:
|
||||
configure_ssrf_whitelist([])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_dns_resolution_keeps_event_loop_responsive(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
started = threading.Event()
|
||||
|
||||
def slow_resolver(hostname, port, family=0, type=0, proto=0, flags=0): # noqa: A002
|
||||
started.set()
|
||||
time.sleep(0.08)
|
||||
return [
|
||||
(socket.AF_INET, socket.SOCK_STREAM, 0, "", ("93.184.216.34", 0)),
|
||||
]
|
||||
|
||||
monkeypatch.setattr(socket, "getaddrinfo", slow_resolver)
|
||||
task = asyncio.create_task(async_resolve_url_target("https://example.com/path"))
|
||||
assert await asyncio.to_thread(started.wait, 0.5)
|
||||
for _ in range(3):
|
||||
await asyncio.sleep(0.01)
|
||||
assert not task.done()
|
||||
|
||||
assert await asyncio.wait_for(task, timeout=0.5) == (
|
||||
True,
|
||||
"",
|
||||
("93.184.216.34",),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_dns_validation_preserves_sync_security_results(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
def resolver(hostname, port, family=0, type=0, proto=0, flags=0): # noqa: A002
|
||||
assert hostname == "evil.com"
|
||||
return [
|
||||
(socket.AF_INET, socket.SOCK_STREAM, 0, "", ("169.254.169.254", 0)),
|
||||
]
|
||||
|
||||
monkeypatch.setattr(socket, "getaddrinfo", resolver)
|
||||
|
||||
sync_result = resolve_url_target("http://evil.com/latest")
|
||||
async_result = await async_resolve_url_target("http://evil.com/latest")
|
||||
redirect_result = await async_validate_resolved_url("http://evil.com/latest")
|
||||
|
||||
assert async_result == sync_result
|
||||
assert async_result[0] is False
|
||||
assert redirect_result[0] is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_dns_resolution_times_out(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
loop = asyncio.get_running_loop()
|
||||
never_resolves = asyncio.Event()
|
||||
|
||||
async def stalled_getaddrinfo(*args, **kwargs):
|
||||
await never_resolves.wait()
|
||||
|
||||
monkeypatch.setattr(loop, "getaddrinfo", stalled_getaddrinfo)
|
||||
|
||||
result = await asyncio.wait_for(
|
||||
async_resolve_url_target(
|
||||
"https://example.com/path",
|
||||
timeout_s=0.01,
|
||||
),
|
||||
timeout=0.2,
|
||||
)
|
||||
|
||||
assert result == (
|
||||
False,
|
||||
"Timed out resolving hostname: example.com",
|
||||
(),
|
||||
)
|
||||
|
||||
@@ -178,13 +178,13 @@ async def test_scan_loads_only_sessions_that_need_webui_recovery(
|
||||
|
||||
coordinator, _, restarted = _coordinator(tmp_path)
|
||||
loaded: list[str] = []
|
||||
get_or_create = restarted.get_or_create
|
||||
get_or_create_async = restarted.get_or_create_async
|
||||
|
||||
def tracked_get_or_create(key: str) -> Session:
|
||||
async def tracked_get_or_create(key: str) -> Session:
|
||||
loaded.append(key)
|
||||
return get_or_create(key)
|
||||
return await get_or_create_async(key)
|
||||
|
||||
monkeypatch.setattr(restarted, "get_or_create", tracked_get_or_create)
|
||||
monkeypatch.setattr(restarted, "get_or_create_async", tracked_get_or_create)
|
||||
monkeypatch.setattr(
|
||||
"nanobot.webui.transcript.has_unfinished_transcript_tail",
|
||||
lambda _key: False,
|
||||
|
||||
@@ -1,4 +1,6 @@
|
||||
import asyncio
|
||||
import gc
|
||||
import threading
|
||||
import weakref
|
||||
|
||||
from nanobot.session.manager import SESSION_CACHE_MAX_SIZE, SessionManager
|
||||
@@ -87,3 +89,39 @@ def test_transient_session_never_reaches_storage(tmp_path) -> None:
|
||||
assert list(manager.sessions_dir.glob("*.jsonl")) == []
|
||||
manager.invalidate(session.key)
|
||||
assert manager.get_cached(session.key) is None
|
||||
|
||||
|
||||
async def test_concurrent_first_async_gets_share_cached_identity(
|
||||
tmp_path,
|
||||
monkeypatch,
|
||||
) -> None:
|
||||
manager = SessionManager(tmp_path)
|
||||
key = "test:concurrent-first-load"
|
||||
load_started = threading.Event()
|
||||
release_load = threading.Event()
|
||||
load_calls = 0
|
||||
|
||||
def delayed_load(loaded_key: str):
|
||||
nonlocal load_calls
|
||||
assert loaded_key == key
|
||||
load_calls += 1
|
||||
load_started.set()
|
||||
assert release_load.wait(timeout=1)
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(manager, "_load", delayed_load)
|
||||
first_task = asyncio.create_task(manager.get_or_create_async(key))
|
||||
try:
|
||||
assert await asyncio.to_thread(load_started.wait, 0.5)
|
||||
second_task = asyncio.create_task(manager.get_or_create_async(key))
|
||||
await asyncio.sleep(0)
|
||||
assert not second_task.done()
|
||||
finally:
|
||||
release_load.set()
|
||||
|
||||
first, second = await asyncio.gather(first_task, second_task)
|
||||
|
||||
assert load_calls == 1
|
||||
assert first is second
|
||||
assert manager.get_cached(key) is first
|
||||
assert await manager.get_or_create_async(key) is first
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -2,8 +2,10 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import socket
|
||||
import sys
|
||||
import threading
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
@@ -173,6 +175,37 @@ async def test_exec_blocks_chained_internal_url():
|
||||
assert "Error" in result
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_restricted_exec_url_guard_keeps_event_loop_responsive(
|
||||
tmp_path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
resolver_started = threading.Event()
|
||||
release_resolver = threading.Event()
|
||||
|
||||
def slow_private_resolver(hostname, port, family=0, type=0, proto=0, flags=0): # noqa: A002
|
||||
resolver_started.set()
|
||||
assert release_resolver.wait(timeout=1)
|
||||
return _fake_resolve_private(hostname, port, family, type)
|
||||
|
||||
monkeypatch.setattr(socket, "getaddrinfo", slow_private_resolver)
|
||||
tool = ExecTool(working_dir=str(tmp_path), restrict_to_workspace=True)
|
||||
task = asyncio.create_task(tool.execute(command="curl https://example.com/latest"))
|
||||
safety_release = threading.Timer(0.2, release_resolver.set)
|
||||
safety_release.start()
|
||||
try:
|
||||
assert await asyncio.to_thread(resolver_started.wait, 0.5)
|
||||
for _ in range(3):
|
||||
await asyncio.sleep(0)
|
||||
assert not task.done()
|
||||
finally:
|
||||
release_resolver.set()
|
||||
safety_release.cancel()
|
||||
|
||||
result = await asyncio.wait_for(task, timeout=0.5)
|
||||
assert "internal/private URL detected" in result
|
||||
|
||||
|
||||
# --- #2989: block writes to nanobot internal state files -----------------
|
||||
|
||||
|
||||
|
||||
@@ -1,7 +1,12 @@
|
||||
"""Tests for enhanced filesystem tools: ReadFileTool, EditFileTool, ListDirTool."""
|
||||
|
||||
import asyncio
|
||||
import threading
|
||||
|
||||
import pytest
|
||||
|
||||
from nanobot.agent.tools import file_state
|
||||
from nanobot.agent.tools import filesystem as filesystem_tools
|
||||
from nanobot.agent.tools.filesystem import (
|
||||
EditFileTool,
|
||||
ListDirTool,
|
||||
@@ -478,3 +483,188 @@ class TestWorkspaceRestriction:
|
||||
)
|
||||
assert "Successfully edited" in result
|
||||
assert target.read_text(encoding="utf-8") == "after\n"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_edit_matching_keeps_event_loop_responsive(tmp_path, monkeypatch):
|
||||
target = tmp_path / "slow.py"
|
||||
target.write_text("before\n", encoding="utf-8")
|
||||
tool = EditFileTool(workspace=tmp_path)
|
||||
original_find_matches = filesystem_tools._find_matches
|
||||
started = threading.Event()
|
||||
release = threading.Event()
|
||||
|
||||
def blocking_find_matches(*args, **kwargs):
|
||||
started.set()
|
||||
if not release.wait(timeout=1):
|
||||
raise TimeoutError("test did not release edit matching")
|
||||
return original_find_matches(*args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(filesystem_tools, "_find_matches", blocking_find_matches)
|
||||
task = asyncio.create_task(
|
||||
tool.execute(path=str(target), old_text="before", new_text="after")
|
||||
)
|
||||
try:
|
||||
assert await asyncio.to_thread(started.wait, 0.5)
|
||||
for _ in range(3):
|
||||
await asyncio.sleep(0.01)
|
||||
assert not task.done()
|
||||
finally:
|
||||
release.set()
|
||||
|
||||
assert "Successfully edited" in await asyncio.wait_for(task, timeout=0.5)
|
||||
assert target.read_text(encoding="utf-8") == "after\n"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_edit_cancellation_before_commit_prevents_delayed_write(tmp_path, monkeypatch):
|
||||
target = tmp_path / "cancel.py"
|
||||
target.write_text("before\n", encoding="utf-8")
|
||||
tool = EditFileTool(workspace=tmp_path)
|
||||
original_find_matches = filesystem_tools._find_matches
|
||||
started = threading.Event()
|
||||
release = threading.Event()
|
||||
|
||||
def blocking_find_matches(*args, **kwargs):
|
||||
started.set()
|
||||
if not release.wait(timeout=1):
|
||||
raise TimeoutError("test did not release edit matching")
|
||||
return original_find_matches(*args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(filesystem_tools, "_find_matches", blocking_find_matches)
|
||||
task = asyncio.create_task(
|
||||
tool.execute(path=str(target), old_text="before", new_text="after")
|
||||
)
|
||||
assert await asyncio.to_thread(started.wait, 0.5)
|
||||
task.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await asyncio.wait_for(task, timeout=0.2)
|
||||
release.set()
|
||||
await asyncio.sleep(0.05)
|
||||
|
||||
assert target.read_text(encoding="utf-8") == "before\n"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_edit_repeated_cancellation_during_commit_waits_for_settlement(
|
||||
tmp_path,
|
||||
monkeypatch,
|
||||
):
|
||||
target = tmp_path / "commit.py"
|
||||
target.write_text("before\n", encoding="utf-8")
|
||||
tool = EditFileTool(workspace=tmp_path)
|
||||
path_type = type(target)
|
||||
original_write_bytes = path_type.write_bytes
|
||||
original_wait_for_commit = EditFileTool._wait_for_commit
|
||||
commit_started = threading.Event()
|
||||
drain_started = threading.Event()
|
||||
release = threading.Event()
|
||||
writes: list[bytes] = []
|
||||
|
||||
def blocking_write_bytes(path, data):
|
||||
if path == target:
|
||||
commit_started.set()
|
||||
if not release.wait(timeout=1):
|
||||
raise TimeoutError("test did not release edit commit")
|
||||
writes.append(data)
|
||||
return original_write_bytes(path, data)
|
||||
|
||||
def observed_wait_for_commit(commit_lock):
|
||||
drain_started.set()
|
||||
original_wait_for_commit(commit_lock)
|
||||
|
||||
monkeypatch.setattr(path_type, "write_bytes", blocking_write_bytes)
|
||||
monkeypatch.setattr(
|
||||
EditFileTool,
|
||||
"_wait_for_commit",
|
||||
staticmethod(observed_wait_for_commit),
|
||||
)
|
||||
task = asyncio.create_task(
|
||||
tool.execute(path=str(target), old_text="before", new_text="after")
|
||||
)
|
||||
try:
|
||||
assert await asyncio.to_thread(commit_started.wait, 0.5)
|
||||
assert task.cancel()
|
||||
assert await asyncio.to_thread(drain_started.wait, 0.5)
|
||||
|
||||
# Later cancellation requests must not interrupt the in-flight commit drain.
|
||||
assert task.cancel()
|
||||
await asyncio.sleep(0)
|
||||
assert not task.done()
|
||||
assert task.cancel()
|
||||
await asyncio.sleep(0)
|
||||
assert not task.done()
|
||||
assert target.read_text(encoding="utf-8") == "before\n"
|
||||
finally:
|
||||
release.set()
|
||||
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await asyncio.wait_for(task, timeout=0.5)
|
||||
|
||||
final_content = target.read_bytes()
|
||||
final_state = tool._file_states.get(target)
|
||||
assert target.read_text(encoding="utf-8") == "after\n"
|
||||
assert writes == [final_content]
|
||||
assert final_state is not None
|
||||
assert tool._file_states.check_read(target) is None
|
||||
|
||||
await asyncio.sleep(0.05)
|
||||
assert target.read_bytes() == final_content
|
||||
assert tool._file_states.get(target) is final_state
|
||||
assert writes == [final_content]
|
||||
|
||||
|
||||
def test_exact_match_line_accounting_is_linear_and_occurrence_can_stop_early():
|
||||
class CountingText(str):
|
||||
count_span = 0
|
||||
find_calls = 0
|
||||
|
||||
def count(self, sub, start=None, end=None):
|
||||
actual_start = 0 if start is None else start
|
||||
actual_end = len(self) if end is None else end
|
||||
self.count_span += max(0, actual_end - actual_start)
|
||||
return super().count(sub, actual_start, actual_end)
|
||||
|
||||
def find(self, sub, start=None, end=None):
|
||||
self.find_calls += 1
|
||||
actual_start = 0 if start is None else start
|
||||
actual_end = len(self) if end is None else end
|
||||
return super().find(sub, actual_start, actual_end)
|
||||
|
||||
content = CountingText("match\n" * 10_000)
|
||||
matches = filesystem_tools._find_exact_matches(content, "match")
|
||||
|
||||
assert len(matches) == 10_000
|
||||
assert matches[-1].line == 10_000
|
||||
assert content.count_span <= len(content)
|
||||
|
||||
first_only = CountingText(content)
|
||||
matches = filesystem_tools._find_matches(
|
||||
first_only,
|
||||
"match",
|
||||
max_exact_matches=1,
|
||||
)
|
||||
assert len(matches) == 1
|
||||
assert first_only.find_calls == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_edit_worker_preserves_bound_file_state_context(tmp_path):
|
||||
target = tmp_path / "context.py"
|
||||
target.write_text("before\n", encoding="utf-8")
|
||||
states = file_state.FileStates()
|
||||
states.record_read(target)
|
||||
tool = EditFileTool(workspace=tmp_path)
|
||||
token = file_state.bind_file_states(states)
|
||||
try:
|
||||
result = await tool.execute(
|
||||
path=str(target),
|
||||
old_text="before",
|
||||
new_text="after",
|
||||
)
|
||||
finally:
|
||||
file_state.reset_file_states(token)
|
||||
|
||||
assert result == f"Successfully edited {target}"
|
||||
assert states.get(target) is not None
|
||||
assert target.read_text(encoding="utf-8") == "after\n"
|
||||
|
||||
@@ -2,7 +2,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import socket
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import httpx
|
||||
@@ -56,10 +55,12 @@ async def test_probe_uses_default_port_for_http(monkeypatch: pytest.MonkeyPatch)
|
||||
"""When no port is present, probe the validated address on port 80."""
|
||||
attempts: list[tuple[str, int]] = []
|
||||
|
||||
monkeypatch.setattr(
|
||||
"nanobot.agent.tools.mcp.resolve_url_target",
|
||||
lambda _url: (True, "", ("93.184.216.34",)),
|
||||
)
|
||||
async def _resolve_url_target(
|
||||
_url: str,
|
||||
) -> tuple[bool, str, tuple[str, ...]]:
|
||||
return True, "", ("93.184.216.34",)
|
||||
|
||||
monkeypatch.setattr(mcp_mod, "async_resolve_url_target", _resolve_url_target)
|
||||
|
||||
async def _open_connection(host: str, port: int):
|
||||
attempts.append((host, port))
|
||||
@@ -72,32 +73,43 @@ async def test_probe_uses_default_port_for_http(monkeypatch: pytest.MonkeyPatch)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_probe_rejects_public_name_resolving_to_loopback():
|
||||
def _resolver(hostname, port, family=0, type_=0):
|
||||
return [(socket.AF_INET, socket.SOCK_STREAM, 0, "", ("127.0.0.1", 0))]
|
||||
async def test_probe_rejects_public_name_resolving_to_loopback(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
async def _resolve_url_target(
|
||||
_url: str,
|
||||
) -> tuple[bool, str, tuple[str, ...]]:
|
||||
return False, "Blocked: example.com resolves to private/internal address 127.0.0.1", ()
|
||||
|
||||
with patch("nanobot.security.network.socket.getaddrinfo", _resolver):
|
||||
assert await _probe_http_url("http://example.com:8765/mcp") is False
|
||||
monkeypatch.setattr(mcp_mod, "async_resolve_url_target", _resolve_url_target)
|
||||
|
||||
assert await _probe_http_url("http://example.com:8765/mcp") is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_probe_skips_direct_tcp_when_global_proxy_env_is_set(monkeypatch):
|
||||
def _resolver(hostname, port, family=0, type_=0):
|
||||
return [(socket.AF_INET, socket.SOCK_STREAM, 0, "", ("93.184.216.34", 0))]
|
||||
async def test_probe_skips_direct_tcp_when_global_proxy_env_is_set(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
async def _resolve_url_target(
|
||||
_url: str,
|
||||
) -> tuple[bool, str, tuple[str, ...]]:
|
||||
return True, "", ("93.184.216.34",)
|
||||
|
||||
async def _open_connection(*args, **kwargs):
|
||||
raise AssertionError("global proxy env should skip direct TCP probe")
|
||||
|
||||
monkeypatch.setenv("HTTPS_PROXY", "http://proxy.example:8080")
|
||||
monkeypatch.setenv("NO_PROXY", "localhost,127.0.0.1,::1")
|
||||
monkeypatch.setattr(mcp_mod, "async_resolve_url_target", _resolve_url_target)
|
||||
monkeypatch.setattr("nanobot.agent.tools.mcp.asyncio.open_connection", _open_connection)
|
||||
|
||||
with patch("nanobot.security.network.socket.getaddrinfo", _resolver):
|
||||
assert await _probe_http_url("https://mcp.example.com/mcp") is True
|
||||
assert await _probe_http_url("https://mcp.example.com/mcp") is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_probe_tries_next_validated_ip_when_first_is_unreachable(monkeypatch):
|
||||
async def test_probe_tries_next_validated_ip_when_first_is_unreachable(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
attempts: list[tuple[str, int]] = []
|
||||
|
||||
class FakeWriter:
|
||||
@@ -107,11 +119,10 @@ async def test_probe_tries_next_validated_ip_when_first_is_unreachable(monkeypat
|
||||
async def wait_closed(self):
|
||||
return None
|
||||
|
||||
def _resolver(hostname, port, family=0, type_=0):
|
||||
return [
|
||||
(socket.AF_INET, socket.SOCK_STREAM, 0, "", ("93.184.216.34", 0)),
|
||||
(socket.AF_INET, socket.SOCK_STREAM, 0, "", ("93.184.216.35", 0)),
|
||||
]
|
||||
async def _resolve_url_target(
|
||||
_url: str,
|
||||
) -> tuple[bool, str, tuple[str, ...]]:
|
||||
return True, "", ("93.184.216.34", "93.184.216.35")
|
||||
|
||||
async def _open_connection(host: str, port: int):
|
||||
attempts.append((host, port))
|
||||
@@ -119,7 +130,7 @@ async def test_probe_tries_next_validated_ip_when_first_is_unreachable(monkeypat
|
||||
raise OSError("first address unreachable")
|
||||
return object(), FakeWriter()
|
||||
|
||||
monkeypatch.setattr("nanobot.security.network.socket.getaddrinfo", _resolver)
|
||||
monkeypatch.setattr(mcp_mod, "async_resolve_url_target", _resolve_url_target)
|
||||
monkeypatch.setattr("nanobot.agent.tools.mcp.asyncio.open_connection", _open_connection)
|
||||
|
||||
assert await _probe_http_url("http://mcp.example:8765/mcp") is True
|
||||
@@ -182,10 +193,13 @@ async def test_connect_isolates_streamable_http_status_failure(
|
||||
async def _reachable(_url: str) -> bool:
|
||||
return True
|
||||
|
||||
async def _validate_url_target(_url: str) -> tuple[bool, str]:
|
||||
return True, ""
|
||||
|
||||
def _return_http_530(request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(530, text="cloudflare error 1033", request=request)
|
||||
|
||||
monkeypatch.setattr(mcp_mod, "validate_url_target", lambda _url: (True, ""))
|
||||
monkeypatch.setattr(mcp_mod, "async_validate_url_target", _validate_url_target)
|
||||
monkeypatch.setattr(mcp_mod, "_probe_http_url", _reachable)
|
||||
monkeypatch.setattr(
|
||||
mcp_mod,
|
||||
|
||||
@@ -213,7 +213,7 @@ async def test_saved_oauth_http_403_projects_failed_runtime_without_details(
|
||||
oauth_mod.create_mcp_oauth_auth = create_auth # type: ignore[attr-defined]
|
||||
oauth_mod.mcp_oauth_has_credentials = lambda _name, _url: True # type: ignore[attr-defined]
|
||||
monkeypatch.setitem(sys.modules, "nanobot.agent.tools.mcp_oauth", oauth_mod)
|
||||
monkeypatch.setattr(mcp_mod, "validate_url_target", lambda _url: (True, ""))
|
||||
monkeypatch.setattr(mcp_mod, "async_validate_url_target", _allow_url)
|
||||
monkeypatch.setattr(mcp_mod, "_probe_http_url", reachable)
|
||||
monkeypatch.setattr(
|
||||
sys.modules["mcp.client.streamable_http"],
|
||||
@@ -1006,7 +1006,7 @@ async def test_connect_mcp_servers_env_proxy_adds_proxy_mounts_and_keeps_pinned_
|
||||
async def _reachable(_url: str) -> bool:
|
||||
return True
|
||||
|
||||
def _validate(_url: str) -> tuple[bool, str]:
|
||||
async def _validate(_url: str) -> tuple[bool, str]:
|
||||
return True, ""
|
||||
|
||||
class FakeAsyncClient:
|
||||
@@ -1033,7 +1033,7 @@ async def test_connect_mcp_servers_env_proxy_adds_proxy_mounts_and_keeps_pinned_
|
||||
|
||||
monkeypatch.setenv("HTTPS_PROXY", "http://proxy.example:8080")
|
||||
monkeypatch.setenv("NO_PROXY", "localhost,127.0.0.1,::1")
|
||||
monkeypatch.setattr(mcp_mod, "validate_url_target", _validate)
|
||||
monkeypatch.setattr(mcp_mod, "async_validate_url_target", _validate)
|
||||
monkeypatch.setattr(mcp_mod, "_probe_http_url", _reachable)
|
||||
monkeypatch.setattr(
|
||||
mcp_mod,
|
||||
@@ -1105,7 +1105,7 @@ async def test_connect_mcp_servers_http_clients_reject_unsafe_redirect_targets(
|
||||
sent_urls: list[str] = []
|
||||
used_transports: list[str] = []
|
||||
|
||||
def _validate(url: str, **_kwargs: object) -> tuple[bool, str]:
|
||||
async def _validate(url: str, **_kwargs: object) -> tuple[bool, str]:
|
||||
checked_urls.append(url)
|
||||
if url == "http://127.0.0.1/private":
|
||||
return False, "loopback blocked"
|
||||
@@ -1145,7 +1145,7 @@ async def test_connect_mcp_servers_http_clients_reject_unsafe_redirect_targets(
|
||||
await http_client.get("https://example.com/start")
|
||||
yield object(), object(), object()
|
||||
|
||||
monkeypatch.setattr(mcp_mod, "validate_url_target", _validate)
|
||||
monkeypatch.setattr(mcp_mod, "async_validate_url_target", _validate)
|
||||
monkeypatch.setattr(mcp_mod, "_probe_http_url", _reachable)
|
||||
monkeypatch.setattr(
|
||||
mcp_mod,
|
||||
@@ -1320,7 +1320,7 @@ async def test_connect_mcp_servers_streamable_http_uses_finite_timeout(
|
||||
async def _reachable(_url: str) -> bool:
|
||||
return True
|
||||
|
||||
def _validate(_url: str) -> tuple[bool, str]:
|
||||
async def _validate(_url: str) -> tuple[bool, str]:
|
||||
return True, ""
|
||||
|
||||
@asynccontextmanager
|
||||
@@ -1328,7 +1328,7 @@ async def test_connect_mcp_servers_streamable_http_uses_finite_timeout(
|
||||
captured["timeout"] = http_client.timeout
|
||||
yield object(), object(), object()
|
||||
|
||||
monkeypatch.setattr(mcp_mod, "validate_url_target", _validate)
|
||||
monkeypatch.setattr(mcp_mod, "async_validate_url_target", _validate)
|
||||
monkeypatch.setattr(mcp_mod, "_probe_http_url", _reachable)
|
||||
monkeypatch.setattr(
|
||||
mcp_mod,
|
||||
@@ -1371,7 +1371,7 @@ async def test_connect_mcp_servers_attaches_oauth_to_remote_http_client(
|
||||
async def _reachable(_url: str) -> bool:
|
||||
return True
|
||||
|
||||
def _validate(_url: str) -> tuple[bool, str]:
|
||||
async def _validate(_url: str) -> tuple[bool, str]:
|
||||
return True, ""
|
||||
|
||||
async def _create_auth(name: str, url: str, handlers: object) -> object:
|
||||
@@ -1407,7 +1407,7 @@ async def test_connect_mcp_servers_attaches_oauth_to_remote_http_client(
|
||||
assert http_client is not None
|
||||
yield object(), object(), object()
|
||||
|
||||
monkeypatch.setattr(mcp_mod, "validate_url_target", _validate)
|
||||
monkeypatch.setattr(mcp_mod, "async_validate_url_target", _validate)
|
||||
monkeypatch.setattr(mcp_mod, "_probe_http_url", _reachable)
|
||||
monkeypatch.setattr(mcp_mod.httpx, "AsyncClient", FakeAsyncClient)
|
||||
monkeypatch.setattr(sys.modules["mcp.client.sse"], "sse_client", _capturing_sse_client)
|
||||
@@ -1461,7 +1461,7 @@ async def test_connect_mcp_servers_skips_background_oauth_without_credentials(
|
||||
oauth_mod.MCPAuthorizationRequiredError = AuthorizationRequiredError # type: ignore[attr-defined]
|
||||
oauth_mod.create_mcp_oauth_auth = _create_auth # type: ignore[attr-defined]
|
||||
monkeypatch.setitem(sys.modules, "nanobot.agent.tools.mcp_oauth", oauth_mod)
|
||||
monkeypatch.setattr(mcp_mod, "validate_url_target", lambda _url: (True, ""))
|
||||
monkeypatch.setattr(mcp_mod, "async_validate_url_target", _allow_url)
|
||||
monkeypatch.setattr(mcp_mod, "_probe_http_url", _probe)
|
||||
|
||||
stacks = await connect_mcp_servers(
|
||||
@@ -1959,3 +1959,12 @@ def test_long_server_name_tools_are_matched_by_server_name() -> None:
|
||||
assert removed == 1
|
||||
assert wrapper.name not in registry.tool_names
|
||||
assert other_wrapper.name in registry.tool_names
|
||||
|
||||
|
||||
async def _allow_url(
|
||||
_url: str,
|
||||
*,
|
||||
allow_loopback: bool = False,
|
||||
) -> tuple[bool, str]:
|
||||
del allow_loopback
|
||||
return True, ""
|
||||
|
||||
@@ -2,7 +2,9 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
@@ -97,6 +99,82 @@ async def test_find_files_rejects_paths_outside_workspace(tmp_path: Path) -> Non
|
||||
assert result.startswith("Error:")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_find_files_scan_keeps_event_loop_responsive(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
target = tmp_path / "match.txt"
|
||||
target.write_text("ok\n", encoding="utf-8")
|
||||
tool = FindFilesTool(workspace=tmp_path, allowed_dir=tmp_path)
|
||||
original_iter_paths = tool._iter_paths
|
||||
started = threading.Event()
|
||||
release = threading.Event()
|
||||
|
||||
def blocking_iter_paths(root: Path, *, include_dirs: bool):
|
||||
started.set()
|
||||
if not release.wait(timeout=1):
|
||||
raise TimeoutError("test did not release find_files traversal")
|
||||
yield from original_iter_paths(root, include_dirs=include_dirs)
|
||||
|
||||
monkeypatch.setattr(tool, "_iter_paths", blocking_iter_paths)
|
||||
task = asyncio.create_task(tool.execute(path="."))
|
||||
try:
|
||||
assert await asyncio.to_thread(started.wait, 0.5)
|
||||
for _ in range(3):
|
||||
await asyncio.sleep(0.01)
|
||||
assert not task.done()
|
||||
finally:
|
||||
release.set()
|
||||
|
||||
assert await asyncio.wait_for(task, timeout=0.5) == "match.txt"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_find_files_cancellation_is_prompt_and_stops_scan_before_next_path(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
(tmp_path / "match.txt").write_text("ok\n", encoding="utf-8")
|
||||
tool = FindFilesTool(workspace=tmp_path, allowed_dir=tmp_path)
|
||||
original_iter_paths = tool._iter_paths
|
||||
started = threading.Event()
|
||||
release = threading.Event()
|
||||
|
||||
def blocking_iter_paths(root: Path, *, include_dirs: bool):
|
||||
started.set()
|
||||
if not release.wait(timeout=1):
|
||||
raise TimeoutError("test did not release find_files traversal")
|
||||
yield from original_iter_paths(root, include_dirs=include_dirs)
|
||||
|
||||
monkeypatch.setattr(tool, "_iter_paths", blocking_iter_paths)
|
||||
task = asyncio.create_task(tool.execute(path="."))
|
||||
assert await asyncio.to_thread(started.wait, 0.5)
|
||||
task.cancel()
|
||||
try:
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await asyncio.wait_for(task, timeout=0.2)
|
||||
await asyncio.sleep(0.01)
|
||||
finally:
|
||||
release.set()
|
||||
await asyncio.sleep(0.05)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_find_files_enforces_path_budget(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
(tmp_path / "a.txt").write_text("a\n", encoding="utf-8")
|
||||
(tmp_path / "b.txt").write_text("b\n", encoding="utf-8")
|
||||
monkeypatch.setattr(FindFilesTool, "_MAX_SCAN_PATHS", 1)
|
||||
tool = FindFilesTool(workspace=tmp_path, allowed_dir=tmp_path)
|
||||
|
||||
result = await tool.execute(path=".")
|
||||
|
||||
assert result.startswith("Error: find_files scan exceeded 1 paths")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_grep_respects_glob_filter_and_context(tmp_path: Path) -> None:
|
||||
(tmp_path / "src").mkdir()
|
||||
|
||||
@@ -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))]
|
||||
|
||||
|
||||
|
||||
@@ -30,11 +30,11 @@ def _clear_proxy_env(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.delenv(name, raising=False)
|
||||
|
||||
|
||||
def _fake_resolve_private(hostname, port, family=0, type_=0):
|
||||
def _fake_resolve_private(hostname, port, family=0, type_=0, proto=0, flags=0):
|
||||
return [(socket.AF_INET, socket.SOCK_STREAM, 0, "", ("169.254.169.254", 0))]
|
||||
|
||||
|
||||
def _fake_resolve_public(hostname, port, family=0, type_=0):
|
||||
def _fake_resolve_public(hostname, port, family=0, type_=0, proto=0, flags=0):
|
||||
return [(socket.AF_INET, socket.SOCK_STREAM, 0, "", ("93.184.216.34", 0))]
|
||||
|
||||
|
||||
@@ -99,7 +99,7 @@ async def test_web_fetch_blocks_private_ip():
|
||||
@pytest.mark.asyncio
|
||||
async def test_web_fetch_blocks_localhost():
|
||||
tool = WebFetchTool()
|
||||
def _resolve_localhost(hostname, port, family=0, type_=0):
|
||||
def _resolve_localhost(hostname, port, family=0, type_=0, proto=0, flags=0):
|
||||
return [(socket.AF_INET, socket.SOCK_STREAM, 0, "", ("127.0.0.1", 0))]
|
||||
with patch("nanobot.security.network.socket.getaddrinfo", _resolve_localhost):
|
||||
result = await tool.execute(url="http://localhost/admin")
|
||||
@@ -112,7 +112,7 @@ async def test_web_fetch_blocks_localhost_even_in_full_workspace_scope(tmp_path)
|
||||
tool = WebFetchTool()
|
||||
scope = build_workspace_scope(tmp_path, "full")
|
||||
|
||||
def _resolve_localhost(hostname, port, family=0, type_=0):
|
||||
def _resolve_localhost(hostname, port, family=0, type_=0, proto=0, flags=0):
|
||||
return [(socket.AF_INET, socket.SOCK_STREAM, 0, "", ("127.0.0.1", 0))]
|
||||
|
||||
token = bind_workspace_scope(scope)
|
||||
@@ -439,7 +439,7 @@ async def test_web_fetch_blocks_private_redirect_before_readability_request(monk
|
||||
monkeypatch.setattr(web_module.httpx, "AsyncClient", FakeClient)
|
||||
monkeypatch.setattr(web_module, "_pinned_dns_transport", lambda: object())
|
||||
|
||||
def resolve_public_start_only(hostname, port, family=0, type_=0):
|
||||
def resolve_public_start_only(hostname, port, family=0, type_=0, proto=0, flags=0):
|
||||
if hostname == "attacker.example":
|
||||
return _fake_resolve_public(hostname, port, family, type_)
|
||||
return _REAL_GETADDRINFO(hostname, port, family, type_)
|
||||
@@ -485,7 +485,7 @@ async def test_web_fetch_blocks_private_redirect_before_returning_image(monkeypa
|
||||
monkeypatch.setattr("nanobot.agent.tools.web.httpx.AsyncClient", TransportAsyncClient)
|
||||
monkeypatch.setattr(web_module, "_pinned_dns_transport", lambda: object())
|
||||
|
||||
def resolve_public_start_only(hostname, port, family=0, type_=0):
|
||||
def resolve_public_start_only(hostname, port, family=0, type_=0, proto=0, flags=0):
|
||||
if hostname == "example.com":
|
||||
return _fake_resolve_public(hostname, port, family, type_)
|
||||
return _REAL_GETADDRINFO(hostname, port, family, type_)
|
||||
@@ -526,7 +526,7 @@ async def test_web_fetch_does_not_request_private_redirect_target(monkeypatch):
|
||||
monkeypatch.setattr(web_module.httpx, "AsyncClient", TransportAsyncClient)
|
||||
monkeypatch.setattr(web_module, "_pinned_dns_transport", lambda: object())
|
||||
|
||||
def resolve_public_start_only(hostname, port, family=0, type_=0):
|
||||
def resolve_public_start_only(hostname, port, family=0, type_=0, proto=0, flags=0):
|
||||
if hostname == "attacker.example":
|
||||
return _fake_resolve_public(hostname, port, family, type_)
|
||||
return _REAL_GETADDRINFO(hostname, port, family, type_)
|
||||
|
||||
@@ -11,7 +11,7 @@ import pytest
|
||||
from nanobot.agent.tools.web import WebFetchTool, _validate_url
|
||||
|
||||
|
||||
def _fake_resolve_public(hostname, port, family=0, type_=0):
|
||||
def _fake_resolve_public(hostname, port, family=0, type_=0, proto=0, flags=0):
|
||||
import socket
|
||||
return [(socket.AF_INET, socket.SOCK_STREAM, 0, "", ("93.184.216.34", 0))]
|
||||
|
||||
|
||||
@@ -4,6 +4,7 @@ import asyncio
|
||||
import errno
|
||||
import json
|
||||
import os
|
||||
import threading
|
||||
from contextlib import suppress
|
||||
from pathlib import Path
|
||||
|
||||
@@ -13,6 +14,7 @@ from nanobot.agent.automation_turns import AutomationTurnError
|
||||
from nanobot.bus.events import InboundMessage, OutboundMessage
|
||||
from nanobot.triggers.local_runner import run_local_trigger_queue
|
||||
from nanobot.triggers.local_store import LocalTriggerStore, TriggerDisabledError
|
||||
from nanobot.triggers.local_turns import LocalTriggerTurnCoordinator
|
||||
from nanobot.triggers.local_types import LocalTrigger, TriggerDelivery
|
||||
from nanobot.webui.metadata import WEBUI_MESSAGE_SOURCE_METADATA_KEY, WEBUI_TURN_METADATA_KEY
|
||||
|
||||
@@ -276,6 +278,163 @@ def test_recover_processing_deliveries_requeues_claimed_delivery(tmp_path: Path)
|
||||
assert reclaimed[0].last_error == "delivery was recovered from interrupted processing"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_local_trigger_queue_cancellation_waits_for_claim_mutation(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
store = LocalTriggerStore(tmp_path)
|
||||
trigger = store.create(
|
||||
name="PR review",
|
||||
channel="websocket",
|
||||
chat_id="chat-1",
|
||||
session_key="websocket:chat-1",
|
||||
)
|
||||
store.enqueue(trigger.id, "Review PR #4591")
|
||||
claim_started = threading.Event()
|
||||
allow_claim = threading.Event()
|
||||
claim_finished = threading.Event()
|
||||
claim_deliveries = store.claim_deliveries
|
||||
submitted: list[InboundMessage] = []
|
||||
|
||||
def blocked_claim_deliveries(*, limit: int = 20) -> list[TriggerDelivery]:
|
||||
claim_started.set()
|
||||
if not allow_claim.wait(timeout=2):
|
||||
raise TimeoutError("test did not release delivery claim")
|
||||
try:
|
||||
return claim_deliveries(limit=limit)
|
||||
finally:
|
||||
claim_finished.set()
|
||||
|
||||
async def submit_turn(msg: InboundMessage) -> None:
|
||||
submitted.append(msg)
|
||||
|
||||
monkeypatch.setattr(store, "claim_deliveries", blocked_claim_deliveries)
|
||||
task = asyncio.create_task(
|
||||
run_local_trigger_queue(
|
||||
store=store,
|
||||
submit_turn=submit_turn,
|
||||
is_channel_enabled=_channel_is_enabled,
|
||||
poll_interval_s=0.01,
|
||||
)
|
||||
)
|
||||
try:
|
||||
assert await asyncio.to_thread(claim_started.wait, 1)
|
||||
task.cancel()
|
||||
await asyncio.sleep(0)
|
||||
|
||||
assert not task.done()
|
||||
assert len(list(store.inbox_dir.glob("*.json"))) == 1
|
||||
assert list(store.processing_dir.glob("*.json")) == []
|
||||
|
||||
allow_claim.set()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await asyncio.wait_for(task, timeout=1)
|
||||
|
||||
assert claim_finished.is_set()
|
||||
state_at_cancellation = (
|
||||
sorted(path.name for path in store.inbox_dir.glob("*.json")),
|
||||
sorted(path.name for path in store.processing_dir.glob("*.json")),
|
||||
)
|
||||
assert state_at_cancellation[0] == []
|
||||
assert len(state_at_cancellation[1]) == 1
|
||||
assert submitted == []
|
||||
|
||||
# The worker has returned before cancellation is visible, so no detached
|
||||
# claim can move another file after the cancellation report.
|
||||
await asyncio.sleep(0)
|
||||
assert (
|
||||
sorted(path.name for path in store.inbox_dir.glob("*.json")),
|
||||
sorted(path.name for path in store.processing_dir.glob("*.json")),
|
||||
) == state_at_cancellation
|
||||
finally:
|
||||
allow_claim.set()
|
||||
task.cancel()
|
||||
with suppress(asyncio.CancelledError):
|
||||
await task
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_local_trigger_queue_cancellation_waits_for_startup_recovery(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
store = LocalTriggerStore(tmp_path)
|
||||
trigger = store.create(
|
||||
name="PR review",
|
||||
channel="websocket",
|
||||
chat_id="chat-1",
|
||||
session_key="websocket:chat-1",
|
||||
)
|
||||
store.enqueue(trigger.id, "Review PR #4591")
|
||||
assert len(store.claim_deliveries()) == 1
|
||||
|
||||
restarted = LocalTriggerStore(tmp_path)
|
||||
recovery_started = threading.Event()
|
||||
allow_recovery = threading.Event()
|
||||
recovery_finished = threading.Event()
|
||||
recover_processing_deliveries = restarted.recover_processing_deliveries
|
||||
submitted: list[InboundMessage] = []
|
||||
|
||||
def blocked_recovery() -> int:
|
||||
recovery_started.set()
|
||||
if not allow_recovery.wait(timeout=2):
|
||||
raise TimeoutError("test did not release delivery recovery")
|
||||
try:
|
||||
return recover_processing_deliveries()
|
||||
finally:
|
||||
recovery_finished.set()
|
||||
|
||||
async def submit_turn(msg: InboundMessage) -> None:
|
||||
submitted.append(msg)
|
||||
|
||||
monkeypatch.setattr(restarted, "recover_processing_deliveries", blocked_recovery)
|
||||
task = asyncio.create_task(
|
||||
run_local_trigger_queue(
|
||||
store=restarted,
|
||||
submit_turn=submit_turn,
|
||||
is_channel_enabled=_channel_is_enabled,
|
||||
poll_interval_s=0.01,
|
||||
)
|
||||
)
|
||||
try:
|
||||
assert await asyncio.to_thread(recovery_started.wait, 1)
|
||||
task.cancel()
|
||||
await asyncio.sleep(0)
|
||||
|
||||
assert not task.done()
|
||||
assert list(restarted.inbox_dir.glob("*.json")) == []
|
||||
assert len(list(restarted.processing_dir.glob("*.json"))) == 1
|
||||
|
||||
allow_recovery.set()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await asyncio.wait_for(task, timeout=1)
|
||||
|
||||
assert recovery_finished.is_set()
|
||||
state_at_cancellation = (
|
||||
sorted(path.name for path in restarted.inbox_dir.glob("*.json")),
|
||||
sorted(path.name for path in restarted.processing_dir.glob("*.json")),
|
||||
)
|
||||
assert len(state_at_cancellation[0]) == 1
|
||||
assert state_at_cancellation[1] == []
|
||||
recovered_payload = json.loads(
|
||||
(restarted.inbox_dir / state_at_cancellation[0][0]).read_text(encoding="utf-8")
|
||||
)
|
||||
assert recovered_payload["delivery"]["attempts"] == 1
|
||||
assert submitted == []
|
||||
|
||||
await asyncio.sleep(0)
|
||||
assert (
|
||||
sorted(path.name for path in restarted.inbox_dir.glob("*.json")),
|
||||
sorted(path.name for path in restarted.processing_dir.glob("*.json")),
|
||||
) == state_at_cancellation
|
||||
finally:
|
||||
allow_recovery.set()
|
||||
task.cancel()
|
||||
with suppress(asyncio.CancelledError):
|
||||
await task
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_local_trigger_queue_submits_bound_inbound_message(tmp_path: Path) -> None:
|
||||
store = LocalTriggerStore(tmp_path)
|
||||
@@ -390,6 +549,7 @@ async def test_local_trigger_queue_rejects_unavailable_target_channel(tmp_path:
|
||||
@pytest.mark.asyncio
|
||||
async def test_local_trigger_queue_waits_for_submitted_turn_before_ack(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
store = LocalTriggerStore(tmp_path)
|
||||
trigger = store.create(
|
||||
@@ -401,6 +561,17 @@ async def test_local_trigger_queue_waits_for_submitted_turn_before_ack(
|
||||
delivery = store.enqueue(trigger.id, "Review failed CI")
|
||||
submitted: list[InboundMessage] = []
|
||||
release = asyncio.Event()
|
||||
completion_started = threading.Event()
|
||||
allow_completion = threading.Event()
|
||||
complete_delivery = store.complete_delivery
|
||||
|
||||
def _complete_delivery(claimed: TriggerDelivery) -> None:
|
||||
completion_started.set()
|
||||
if not allow_completion.wait(timeout=1):
|
||||
raise TimeoutError("test did not release delivery completion")
|
||||
complete_delivery(claimed)
|
||||
|
||||
monkeypatch.setattr(store, "complete_delivery", _complete_delivery)
|
||||
|
||||
async def _submit_turn(msg: InboundMessage):
|
||||
submitted.append(msg)
|
||||
@@ -430,6 +601,13 @@ async def test_local_trigger_queue_waits_for_submitted_turn_before_ack(
|
||||
assert stored.last_status is None
|
||||
|
||||
release.set()
|
||||
assert await asyncio.to_thread(completion_started.wait, 1)
|
||||
stored = store.get(trigger.id)
|
||||
assert stored is not None
|
||||
assert stored.last_status is None
|
||||
assert list(store.processing_dir.glob("*.json"))
|
||||
|
||||
allow_completion.set()
|
||||
for _ in range(100):
|
||||
stored = store.get(trigger.id)
|
||||
if stored and stored.last_status == "ok":
|
||||
@@ -444,6 +622,7 @@ async def test_local_trigger_queue_waits_for_submitted_turn_before_ack(
|
||||
record = _read_run_record(store, delivery.id)
|
||||
assert record["status"] == "ok"
|
||||
finally:
|
||||
allow_completion.set()
|
||||
task.cancel()
|
||||
with suppress(asyncio.CancelledError):
|
||||
await task
|
||||
@@ -495,9 +674,115 @@ async def test_local_trigger_queue_requeues_when_submitted_turn_is_interrupted(
|
||||
await task
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_accepted_turn_cancellation_settles_delivery_without_replay(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
store = LocalTriggerStore(tmp_path)
|
||||
trigger = store.create(
|
||||
name="CI review",
|
||||
channel="websocket",
|
||||
chat_id="chat-1",
|
||||
session_key="websocket:chat-1",
|
||||
)
|
||||
delivery = store.enqueue(trigger.id, "Review failed CI")
|
||||
inbound: asyncio.Queue[InboundMessage] = asyncio.Queue()
|
||||
work_accepted = asyncio.Event()
|
||||
release_work = asyncio.Event()
|
||||
effects: list[str] = []
|
||||
publish_count = 0
|
||||
|
||||
async def publish(msg: InboundMessage) -> None:
|
||||
nonlocal publish_count
|
||||
publish_count += 1
|
||||
await inbound.put(msg)
|
||||
|
||||
coordinator = LocalTriggerTurnCoordinator(
|
||||
publish_inbound=publish,
|
||||
dispatch=lambda _msg: asyncio.sleep(0),
|
||||
is_running=lambda: True,
|
||||
)
|
||||
|
||||
async def agent_worker() -> None:
|
||||
msg = await inbound.get()
|
||||
work_accepted.set()
|
||||
await release_work.wait()
|
||||
effects.append(msg.content)
|
||||
coordinator.complete(msg)
|
||||
|
||||
agent_task = asyncio.create_task(agent_worker())
|
||||
queue_task = asyncio.create_task(
|
||||
run_local_trigger_queue(
|
||||
store=store,
|
||||
submit_turn=coordinator.submit,
|
||||
is_channel_enabled=_channel_is_enabled,
|
||||
poll_interval_s=0.01,
|
||||
)
|
||||
)
|
||||
restarted_task: asyncio.Task[None] | None = None
|
||||
try:
|
||||
await asyncio.wait_for(work_accepted.wait(), timeout=1)
|
||||
queue_task.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await asyncio.wait_for(queue_task, timeout=1)
|
||||
|
||||
# The delivery is durably settled before cancellation is reported even
|
||||
# though the independently-owned agent turn has not completed yet.
|
||||
assert effects == []
|
||||
assert not list(store.processing_dir.glob("*.json"))
|
||||
assert _read_run_record(store, delivery.id)["status"] == "accepted"
|
||||
|
||||
restarted = LocalTriggerStore(tmp_path)
|
||||
replayed: list[InboundMessage] = []
|
||||
poll_completed = threading.Event()
|
||||
claim_deliveries = restarted.claim_deliveries
|
||||
|
||||
def tracked_claim_deliveries(*, limit: int = 20) -> list[TriggerDelivery]:
|
||||
try:
|
||||
return claim_deliveries(limit=limit)
|
||||
finally:
|
||||
poll_completed.set()
|
||||
|
||||
async def replay_submit(msg: InboundMessage) -> None:
|
||||
replayed.append(msg)
|
||||
|
||||
monkeypatch.setattr(restarted, "claim_deliveries", tracked_claim_deliveries)
|
||||
restarted_task = asyncio.create_task(
|
||||
run_local_trigger_queue(
|
||||
store=restarted,
|
||||
submit_turn=replay_submit,
|
||||
is_channel_enabled=_channel_is_enabled,
|
||||
poll_interval_s=1,
|
||||
)
|
||||
)
|
||||
assert await asyncio.to_thread(poll_completed.wait, 1)
|
||||
assert replayed == []
|
||||
assert restarted.claim_deliveries() == []
|
||||
|
||||
release_work.set()
|
||||
await asyncio.wait_for(agent_task, timeout=1)
|
||||
assert effects == ["Review failed CI"]
|
||||
assert publish_count == 1
|
||||
finally:
|
||||
release_work.set()
|
||||
if restarted_task is not None:
|
||||
restarted_task.cancel()
|
||||
with suppress(asyncio.CancelledError):
|
||||
await restarted_task
|
||||
queue_task.cancel()
|
||||
with suppress(asyncio.CancelledError):
|
||||
await queue_task
|
||||
if not agent_task.done():
|
||||
agent_task.cancel()
|
||||
with suppress(asyncio.CancelledError):
|
||||
await agent_task
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_local_trigger_queue_does_not_retry_completed_agent_failure(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
store = LocalTriggerStore(tmp_path)
|
||||
trigger = store.create(
|
||||
@@ -508,6 +793,17 @@ async def test_local_trigger_queue_does_not_retry_completed_agent_failure(
|
||||
)
|
||||
delivery = store.enqueue(trigger.id, "Review failed CI")
|
||||
started = asyncio.Event()
|
||||
completion_started = threading.Event()
|
||||
allow_completion = threading.Event()
|
||||
complete_delivery = store.complete_delivery
|
||||
|
||||
def _complete_delivery(claimed: TriggerDelivery) -> None:
|
||||
completion_started.set()
|
||||
if not allow_completion.wait(timeout=1):
|
||||
raise TimeoutError("test did not release delivery completion")
|
||||
complete_delivery(claimed)
|
||||
|
||||
monkeypatch.setattr(store, "complete_delivery", _complete_delivery)
|
||||
|
||||
async def _submit_turn(_msg: InboundMessage):
|
||||
started.set()
|
||||
@@ -523,6 +819,13 @@ async def test_local_trigger_queue_does_not_retry_completed_agent_failure(
|
||||
)
|
||||
try:
|
||||
await asyncio.wait_for(started.wait(), timeout=1)
|
||||
assert await asyncio.to_thread(completion_started.wait, 1)
|
||||
stored = store.get(trigger.id)
|
||||
assert stored is not None
|
||||
assert stored.last_status is None
|
||||
assert list(store.processing_dir.glob("*.json"))
|
||||
|
||||
allow_completion.set()
|
||||
for _ in range(100):
|
||||
stored = store.get(trigger.id)
|
||||
if stored and stored.last_status == "error":
|
||||
@@ -540,6 +843,7 @@ async def test_local_trigger_queue_does_not_retry_completed_agent_failure(
|
||||
assert record["status"] == "error"
|
||||
assert record["error"] == "model failed"
|
||||
finally:
|
||||
allow_completion.set()
|
||||
task.cancel()
|
||||
with suppress(asyncio.CancelledError):
|
||||
await task
|
||||
|
||||
@@ -1,5 +1,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import threading
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import ANY, AsyncMock, MagicMock
|
||||
|
||||
@@ -160,3 +162,50 @@ async def test_fork_handler_maps_invalid_source_and_internal_failure_to_stable_e
|
||||
|
||||
channel.logger.warning.assert_called_once_with("fork_chat failed: {}", ANY)
|
||||
channel.send_webui_protocol_error.assert_awaited_once_with(connection, "fork_chat_failed")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fork_cancellation_waits_for_creation_and_client_attachment(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
started = threading.Event()
|
||||
release = threading.Event()
|
||||
connection = object()
|
||||
|
||||
def blocked_create(*_args, **_kwargs) -> tuple[str, str]:
|
||||
started.set()
|
||||
assert release.wait(timeout=1)
|
||||
return "fork-id", "websocket:fork-id"
|
||||
|
||||
channel = SimpleNamespace(
|
||||
send_webui_protocol_error=AsyncMock(),
|
||||
attach_webui_fork=AsyncMock(),
|
||||
gateway=SimpleNamespace(session_manager=object()),
|
||||
logger=SimpleNamespace(warning=MagicMock()),
|
||||
)
|
||||
monkeypatch.setattr(forking, "create_webui_chat_fork", blocked_create)
|
||||
task = asyncio.create_task(
|
||||
forking.handle_webui_fork_chat(
|
||||
channel,
|
||||
connection,
|
||||
{"source_chat_id": "source", "before_user_index": 0},
|
||||
)
|
||||
)
|
||||
assert await asyncio.to_thread(started.wait, 1)
|
||||
try:
|
||||
task.cancel()
|
||||
await asyncio.sleep(0)
|
||||
assert not task.done()
|
||||
channel.attach_webui_fork.assert_not_awaited()
|
||||
finally:
|
||||
release.set()
|
||||
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await asyncio.wait_for(task, timeout=1)
|
||||
|
||||
channel.attach_webui_fork.assert_awaited_once_with(
|
||||
connection,
|
||||
fork_id="fork-id",
|
||||
fork_key="websocket:fork-id",
|
||||
)
|
||||
channel.send_webui_protocol_error.assert_not_awaited()
|
||||
|
||||
@@ -2,8 +2,11 @@
|
||||
|
||||
import gzip
|
||||
import json
|
||||
import time
|
||||
from types import SimpleNamespace
|
||||
|
||||
from nanobot.webui.http_utils import http_json_response
|
||||
from nanobot.webui.ws_http import GatewayHTTPHandler
|
||||
|
||||
|
||||
def test_http_json_response_compresses_large_payload_when_gzip_is_accepted() -> None:
|
||||
@@ -36,3 +39,29 @@ def test_http_json_response_does_not_compress_small_payload() -> None:
|
||||
assert "Content-Encoding" not in response.headers
|
||||
assert response.headers["Vary"] == "Accept-Encoding"
|
||||
assert json.loads(response.body) == payload
|
||||
|
||||
|
||||
def test_slow_http_log_records_route_scale_without_user_path_content() -> None:
|
||||
records: list[str] = []
|
||||
|
||||
class _Logger:
|
||||
def warning(self, message: str, *args: object) -> None:
|
||||
records.append(message.format(*args))
|
||||
|
||||
handler = SimpleNamespace(_log=_Logger())
|
||||
secret = "private-session-token"
|
||||
|
||||
GatewayHTTPHandler._log_slow_http(
|
||||
handler,
|
||||
f"/api/sessions/{secret}?query={secret}",
|
||||
SimpleNamespace(status_code=200),
|
||||
time.perf_counter() - 2,
|
||||
input_chars=80,
|
||||
)
|
||||
|
||||
assert len(records) == 1
|
||||
assert "operation=/api/sessions" in records[0]
|
||||
assert "status=200" in records[0]
|
||||
assert "input_chars=80" in records[0]
|
||||
assert "duration_ms=" in records[0]
|
||||
assert secret not in records[0]
|
||||
|
||||
@@ -38,9 +38,12 @@ async def test_browser_flow_retries_current_server_and_ignores_unrelated_reload_
|
||||
received: dict[str, object] = {}
|
||||
reload_calls = 0
|
||||
|
||||
async def validate_url_target(_url: str) -> tuple[bool, str]:
|
||||
return True, ""
|
||||
|
||||
monkeypatch.setattr(
|
||||
"nanobot.webui.mcp_oauth_api.validate_url_target",
|
||||
lambda _url: (True, ""),
|
||||
"nanobot.webui.mcp_oauth_api.async_validate_url_target",
|
||||
validate_url_target,
|
||||
)
|
||||
|
||||
async def connect(servers, _registry, *, oauth_handlers):
|
||||
@@ -115,9 +118,12 @@ async def test_remote_http_flow_accepts_a_pasted_loopback_callback(
|
||||
connection = _Connection()
|
||||
received: dict[str, object] = {}
|
||||
|
||||
async def validate_url_target(_url: str) -> tuple[bool, str]:
|
||||
return True, ""
|
||||
|
||||
monkeypatch.setattr(
|
||||
"nanobot.webui.mcp_oauth_api.validate_url_target",
|
||||
lambda _url: (True, ""),
|
||||
"nanobot.webui.mcp_oauth_api.async_validate_url_target",
|
||||
validate_url_target,
|
||||
)
|
||||
|
||||
async def connect(_servers, _registry, *, oauth_handlers):
|
||||
@@ -187,9 +193,13 @@ async def test_browser_flow_surfaces_provider_denial_without_callback_descriptio
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
manager = McpOAuthManager()
|
||||
|
||||
async def validate_url_target(_url: str) -> tuple[bool, str]:
|
||||
return True, ""
|
||||
|
||||
monkeypatch.setattr(
|
||||
"nanobot.webui.mcp_oauth_api.validate_url_target",
|
||||
lambda _url: (True, ""),
|
||||
"nanobot.webui.mcp_oauth_api.async_validate_url_target",
|
||||
validate_url_target,
|
||||
)
|
||||
|
||||
async def connect(_servers, _registry, *, oauth_handlers):
|
||||
@@ -233,9 +243,13 @@ async def test_browser_flow_blocks_unsafe_authorization_url(
|
||||
state: str,
|
||||
) -> None:
|
||||
manager = McpOAuthManager()
|
||||
|
||||
async def validate_url_target(_url: str) -> tuple[bool, str]:
|
||||
return url_is_safe, "private address"
|
||||
|
||||
monkeypatch.setattr(
|
||||
"nanobot.webui.mcp_oauth_api.validate_url_target",
|
||||
lambda _url: (url_is_safe, "private address"),
|
||||
"nanobot.webui.mcp_oauth_api.async_validate_url_target",
|
||||
validate_url_target,
|
||||
)
|
||||
|
||||
async def connect(_servers, _registry, *, oauth_handlers):
|
||||
|
||||
@@ -526,3 +526,179 @@ async def test_version_check_route_enforces_auth_and_bounds_failures(
|
||||
assert failed.status_code == 500
|
||||
assert json.loads(failed.body) == {"error": "version check failed"}
|
||||
assert "upstream secret body" not in failed.body.decode()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mcp_oauth_start_cancellation_waits_for_config_and_flow_start(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
config_started = threading.Event()
|
||||
release_config = threading.Event()
|
||||
start_entered = asyncio.Event()
|
||||
release_start = asyncio.Event()
|
||||
effects: list[str] = []
|
||||
cfg = SimpleNamespace(
|
||||
type="streamableHttp",
|
||||
auth="oauth",
|
||||
url="https://app.xmind.com/api/mcp",
|
||||
)
|
||||
|
||||
def ensure_server(_query, *, config_path=None):
|
||||
config_started.set()
|
||||
assert release_config.wait(timeout=1)
|
||||
effects.append("config_saved")
|
||||
return "xmind", cfg
|
||||
|
||||
async def start(
|
||||
name,
|
||||
config,
|
||||
redirect_uri,
|
||||
*,
|
||||
reload_mcp,
|
||||
reset_credentials=False,
|
||||
):
|
||||
assert (name, config, redirect_uri) == (
|
||||
"xmind",
|
||||
cfg,
|
||||
"https://gateway.example/auth/mcp/callback",
|
||||
)
|
||||
assert callable(reload_mcp)
|
||||
assert reset_credentials is False
|
||||
effects.append("start_entered")
|
||||
start_entered.set()
|
||||
await release_start.wait()
|
||||
effects.append("start_settled")
|
||||
return {
|
||||
"status": "authorization_required",
|
||||
"flow_id": "flow-123",
|
||||
"name": name,
|
||||
}
|
||||
|
||||
monkeypatch.setattr(
|
||||
"nanobot.webui.settings_routes.ensure_mcp_oauth_server",
|
||||
ensure_server,
|
||||
)
|
||||
router = _router()
|
||||
router._mcp_oauth = SimpleNamespace(start=start)
|
||||
request = _mutation_request(
|
||||
"/api/settings/mcp-oauth/start",
|
||||
{"name": "xmind"},
|
||||
)
|
||||
task = asyncio.create_task(
|
||||
router.dispatch(None, request, "/api/settings/mcp-oauth/start")
|
||||
)
|
||||
|
||||
assert await asyncio.to_thread(config_started.wait, 1)
|
||||
try:
|
||||
task.cancel()
|
||||
await asyncio.sleep(0)
|
||||
assert not task.done()
|
||||
|
||||
release_config.set()
|
||||
await asyncio.wait_for(start_entered.wait(), timeout=1)
|
||||
assert effects == ["config_saved", "start_entered"]
|
||||
assert not task.done()
|
||||
finally:
|
||||
release_config.set()
|
||||
release_start.set()
|
||||
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await asyncio.wait_for(task, timeout=1)
|
||||
|
||||
assert effects == ["config_saved", "start_entered", "start_settled"]
|
||||
await asyncio.sleep(0.05)
|
||||
assert effects == ["config_saved", "start_entered", "start_settled"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_model_settings_cancellation_waits_for_mutation_and_runtime_refresh(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
mutation_started = threading.Event()
|
||||
release_mutation = threading.Event()
|
||||
effects: list[str] = []
|
||||
|
||||
def update(_query, *, config_path=None):
|
||||
mutation_started.set()
|
||||
assert release_mutation.wait(timeout=1)
|
||||
effects.append("config_saved")
|
||||
return {"updated": True}
|
||||
|
||||
def refresh_runtime_config() -> None:
|
||||
effects.append("runtime_refreshed")
|
||||
|
||||
monkeypatch.setattr(
|
||||
"nanobot.webui.settings_routes.update_agent_settings",
|
||||
update,
|
||||
)
|
||||
path = "/api/settings/update"
|
||||
request = _mutation_request(path, {"model_preset": "Codex"})
|
||||
task = asyncio.create_task(
|
||||
_router(refresh_runtime_config=refresh_runtime_config).dispatch(
|
||||
None,
|
||||
request,
|
||||
path,
|
||||
)
|
||||
)
|
||||
|
||||
assert await asyncio.to_thread(mutation_started.wait, 1)
|
||||
try:
|
||||
task.cancel()
|
||||
await asyncio.sleep(0)
|
||||
assert not task.done()
|
||||
assert effects == []
|
||||
finally:
|
||||
release_mutation.set()
|
||||
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await asyncio.wait_for(task, timeout=1)
|
||||
|
||||
assert effects == ["config_saved", "runtime_refreshed"]
|
||||
await asyncio.sleep(0.05)
|
||||
assert effects == ["config_saved", "runtime_refreshed"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_image_settings_cancellation_waits_for_runtime_application(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
reload_started = asyncio.Event()
|
||||
release_reload = asyncio.Event()
|
||||
effects: list[str] = []
|
||||
|
||||
def update(_query, *, config_path=None):
|
||||
effects.append("config_saved")
|
||||
return {"requires_restart": True}
|
||||
|
||||
async def reload_image(_bus):
|
||||
effects.append("reload_started")
|
||||
reload_started.set()
|
||||
await release_reload.wait()
|
||||
effects.append("reload_settled")
|
||||
return {"ok": True, "requires_restart": False}
|
||||
|
||||
monkeypatch.setattr(
|
||||
"nanobot.webui.settings_routes.update_image_generation_settings",
|
||||
update,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"nanobot.webui.settings_routes.request_image_generation_reload",
|
||||
reload_image,
|
||||
)
|
||||
path = "/api/settings/image-generation/update"
|
||||
request = _mutation_request(path, {"enabled": True})
|
||||
task = asyncio.create_task(_router().dispatch(None, request, path))
|
||||
|
||||
await asyncio.wait_for(reload_started.wait(), timeout=1)
|
||||
assert effects == ["config_saved", "reload_started"]
|
||||
task.cancel()
|
||||
await asyncio.sleep(0)
|
||||
assert not task.done()
|
||||
|
||||
release_reload.set()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await asyncio.wait_for(task, timeout=1)
|
||||
|
||||
assert effects == ["config_saved", "reload_started", "reload_settled"]
|
||||
await asyncio.sleep(0.05)
|
||||
assert effects == ["config_saved", "reload_started", "reload_settled"]
|
||||
|
||||
Reference in New Issue
Block a user