monkeybot 2.1.1__py3-none-any.whl

monkeybot/gateway/sse/scheduler_routes.py

@@ -0,0 +1,200 @@
+"""REST control plane for prompt-first scheduled loops."""
+
+from __future__ import annotations
+
+import time
+import uuid
+from typing import Any, cast
+
+from fastapi import APIRouter, Depends, Request
+from pydantic import BaseModel, Field
+
+from monkeybot.core.persistence.backends import ScheduledLoopStore, StorageBackend
+from monkeybot.core.persistence.scheduled_loops import ScheduledLoopCreate, ScheduledLoopRow
+from monkeybot.core.types.content_blocks import Text
+from monkeybot.gateway.sse.loop_port import LoopPort
+from monkeybot.gateway.sse.models import APIError
+from monkeybot.gateway.sse.scheduler_wiring import GatewaySessionEnsurer
+from monkeybot.gateway.sse.session_bus import SessionRegistry
+from monkeybot.scheduler.interval import parse_interval_ms, parse_optional_duration_ms
+
+
+def _get_registry(request: Request) -> SessionRegistry:
+    return cast(SessionRegistry, request.app.state.registry)
+
+
+def _storage_backend(request: Request) -> StorageBackend:
+    backend: StorageBackend | None = getattr(request.app.state, "storage", None)
+    if backend is None:
+        raise APIError(503, "STORAGE_NOT_READY", "Storage backend not initialized", uuid.uuid4().hex)
+    return backend
+
+
+def _loop_store(request: Request) -> ScheduledLoopStore:
+    return _storage_backend(request).scheduled_loops()
+
+
+class CreateLoopRequest(BaseModel):
+    prompt: str = Field(..., min_length=1)
+    interval: str | int | float
+    session_id: str = "loop-main"
+    loop_id: str | None = None
+    max_ticks: int | None = Field(default=None, ge=1)
+    max_runtime: str | int | float | None = None
+    skip_if_busy: bool = True
+    unbounded: bool = False
+    confirmed: bool = False
+
+
+class InvokeTickRequest(BaseModel):
+    session_id: str
+    request_id: str
+    message: str
+
+
+def _row_dict(row: ScheduledLoopRow) -> dict[str, Any]:
+    return {
+        "loop_id": row.loop_id,
+        "session_id": row.session_id,
+        "status": row.status,
+        "prompt": row.prompt,
+        "interval_ms": row.interval_ms,
+        "max_ticks": row.max_ticks,
+        "max_runtime_ms": row.max_runtime_ms,
+        "skip_if_busy": row.skip_if_busy,
+        "tick_index": row.tick_index,
+        "next_tick_at_ms": row.next_tick_at_ms,
+        "started_at_ms": row.started_at_ms,
+        "last_tick_at_ms": row.last_tick_at_ms,
+        "last_error": row.last_error,
+        "stop_reason": row.stop_reason,
+        "tick_in_flight": row.tick_in_flight,
+    }
+
+
+def build_scheduler_router(*, loop_port: LoopPort, registry: SessionRegistry) -> APIRouter:
+    router = APIRouter(prefix="/scheduler", tags=["scheduler"])
+    ensurer = GatewaySessionEnsurer(registry)
+
+    @router.post("/loops", status_code=201)
+    async def create_loop(
+        body: CreateLoopRequest,
+        request: Request,
+    ) -> dict[str, Any]:
+        store = _loop_store(request)
+        if not body.confirmed:
+            raise APIError(
+                400,
+                "CONFIRMATION_REQUIRED",
+                "confirmed=true is required to register a scheduled loop via REST/CLI",
+                uuid.uuid4().hex,
+            )
+        try:
+            interval_ms = parse_interval_ms(body.interval)
+            max_runtime_ms = parse_optional_duration_ms(body.max_runtime)
+        except ValueError as exc:
+            raise APIError(400, "INVALID_INTERVAL", str(exc), uuid.uuid4().hex) from exc
+        spec = ScheduledLoopCreate(
+            prompt=body.prompt,
+            interval_ms=interval_ms,
+            session_id=body.session_id,
+            loop_id=body.loop_id,
+            max_ticks=body.max_ticks,
+            max_runtime_ms=max_runtime_ms,
+            skip_if_busy=body.skip_if_busy,
+            unbounded=body.unbounded,
+        )
+        try:
+            row = await store.create(spec)
+        except ValueError as exc:
+            msg = str(exc)
+            if "scheduled loops require" in msg:
+                raise APIError(400, "INVALID_GUARDS", msg, uuid.uuid4().hex) from exc
+            raise APIError(409, "LOOP_EXISTS", msg, uuid.uuid4().hex) from exc
+        await ensurer.ensure_session(row.session_id)
+        return {"loop": _row_dict(row)}
+
+    @router.get("/loops")
+    async def list_loops(request: Request) -> dict[str, Any]:
+        rows = await _loop_store(request).list_all()
+        return {"loops": [_row_dict(r) for r in rows]}
+
+    @router.get("/loops/{loop_id}")
+    async def get_loop(loop_id: str, request: Request) -> dict[str, Any]:
+        row = await _loop_store(request).get(loop_id)
+        if row is None:
+            raise APIError(404, "LOOP_NOT_FOUND", f"Unknown loop {loop_id}", uuid.uuid4().hex)
+        usage = await _storage_backend(request).usage().summary(
+            thread_id=row.session_id,
+            since_ms=row.started_at_ms,
+        )
+        return {
+            "loop": _row_dict(row),
+            "usage": {
+                "session_id": row.session_id,
+                "since_ms": row.started_at_ms,
+                "turns": usage.turns,
+                "input_tokens": usage.input_tokens,
+                "output_tokens": usage.output_tokens,
+                "cached_tokens": usage.cached_tokens,
+                "cache_read_tokens": usage.cache_read_tokens,
+                "cache_creation_tokens": usage.cache_creation_tokens,
+                "cost_usd": usage.cost_usd,
+            },
+        }
+
+    @router.post("/loops/{loop_id}/pause")
+    async def pause_loop(loop_id: str, request: Request) -> dict[str, Any]:
+        ok = await _loop_store(request).pause(loop_id)
+        if not ok:
+            raise APIError(404, "LOOP_NOT_FOUND", f"Unknown loop {loop_id}", uuid.uuid4().hex)
+        return {"loop_id": loop_id, "status": "paused"}
+
+    @router.post("/loops/{loop_id}/resume")
+    async def resume_loop(loop_id: str, request: Request) -> dict[str, Any]:
+        ok = await _loop_store(request).resume(loop_id)
+        if not ok:
+            raise APIError(404, "LOOP_NOT_FOUND", f"Paused loop {loop_id} not found", uuid.uuid4().hex)
+        return {"loop_id": loop_id, "status": "active"}
+
+    @router.post("/loops/{loop_id}/stop")
+    async def stop_loop(loop_id: str, request: Request) -> dict[str, Any]:
+        ok = await _loop_store(request).stop(loop_id)
+        if not ok:
+            raise APIError(404, "LOOP_NOT_FOUND", f"Unknown loop {loop_id}", uuid.uuid4().hex)
+        return {"loop_id": loop_id, "status": "completed"}
+
+    @router.post("/invoke-tick")
+    async def invoke_tick_sync(
+        body: InvokeTickRequest,
+        request: Request,
+        reg_dep: SessionRegistry = Depends(_get_registry),
+    ) -> dict[str, Any]:
+        """Run one agent turn synchronously (scheduler worker / internal automation)."""
+        if reg_dep.get(body.session_id) is None:
+            await ensurer.ensure_session(body.session_id)
+        turn_locks = _storage_backend(request).session_turns()
+        acquired = await turn_locks.try_acquire(body.session_id, body.request_id)
+        if not acquired:
+            raise APIError(
+                409,
+                "SESSION_BUSY",
+                "Session already processing a request",
+                uuid.uuid4().hex,
+            )
+        bus = reg_dep.get(body.session_id)
+        if bus is not None:
+            bus.current_request_id = body.request_id
+        try:
+            await loop_port.start_turn(
+                body.session_id,
+                body.request_id,
+                [Text(text=body.message)],
+            )
+        finally:
+            await turn_locks.release(body.session_id, body.request_id)
+            if bus is not None and bus.current_request_id == body.request_id:
+                bus.current_request_id = None
+        return {"ok": True, "session_id": body.session_id, "request_id": body.request_id}
+
+    return router

monkeybot/gateway/sse/scheduler_wiring.py

@@ -0,0 +1,96 @@
+"""Gateway adapters for the scheduled-loop engine."""
+
+from __future__ import annotations
+
+import logging
+import os
+import time
+
+from monkeybot.core.persistence.backends import SessionTurnLockStore, StorageBackend
+from monkeybot.core.types.content_blocks import ContentBlock
+from monkeybot.gateway.sse.loop_port import LoopPort
+from monkeybot.gateway.sse.session_bus import SessionAlreadyExistsError, SessionRegistry
+from monkeybot.scheduler.tick_result import TickInvokeResult
+
+logger = logging.getLogger(__name__)
+
+
+class StorageSessionBusyChecker:
+    """Check session busy state via durable storage (multi-replica safe)."""
+
+    def __init__(self, turn_locks: SessionTurnLockStore) -> None:
+        self._turn_locks = turn_locks
+
+    def is_busy(self, session_id: str) -> bool:
+        del session_id
+        return False
+
+    async def is_busy_async(self, session_id: str) -> bool:
+        return await self._turn_locks.is_busy(session_id)
+
+
+class GatewaySessionBusyChecker:
+    """Legacy in-process checker; prefer :class:`StorageSessionBusyChecker`."""
+
+    def __init__(self, registry: SessionRegistry) -> None:
+        self._registry = registry
+
+    def is_busy(self, session_id: str) -> bool:
+        bus = self._registry.get(session_id)
+        return bus is not None and bus.current_request_id is not None
+
+
+class GatewaySessionEnsurer:
+    def __init__(self, registry: SessionRegistry) -> None:
+        self._registry = registry
+
+    async def ensure_session(self, session_id: str) -> None:
+        if self._registry.get(session_id) is not None:
+            return
+        agent_md = os.environ.get("AGENT_MD")
+        try:
+            self._registry.create(
+                session_id,
+                agent_md=agent_md,
+                created_at_ms=int(time.time() * 1000),
+            )
+            logger.info("scheduler created session_id=%s", session_id)
+        except SessionAlreadyExistsError:
+            return
+
+
+class GatewayTickInvoker:
+    """Invoke ``LoopPort.start_turn`` and surface turn-level outcomes."""
+
+    def __init__(
+        self,
+        loop_port: LoopPort,
+        registry: SessionRegistry,
+        turn_locks: SessionTurnLockStore,
+    ) -> None:
+        self._loop_port = loop_port
+        self._registry = registry
+        self._turn_locks = turn_locks
+
+    async def invoke_tick(
+        self,
+        session_id: str,
+        request_id: str,
+        user_content: list[ContentBlock],
+    ) -> TickInvokeResult:
+        acquired = await self._turn_locks.try_acquire(session_id, request_id)
+        if not acquired:
+            return TickInvokeResult.session_busy()
+        bus = self._registry.get(session_id)
+        if bus is not None:
+            bus.current_request_id = request_id
+        try:
+            await self._loop_port.start_turn(session_id, request_id, user_content)
+        except Exception as exc:
+            return TickInvokeResult.fail(str(exc))
+        finally:
+            await self._turn_locks.release(session_id, request_id)
+            bus = self._registry.get(session_id)
+            if bus is not None and bus.current_request_id == request_id:
+                bus.current_request_id = None
+        return TickInvokeResult.ok()

monkeybot/gateway/sse/session_bus.py

@@ -0,0 +1,226 @@
+"""
+In-memory per-session SSE bus with replay buffer and live subscribers.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import os
+from collections import deque
+from typing import Any, Literal
+
+from monkeybot.core.attachments.catalog import SessionAttachmentCatalog
+from monkeybot.core.persistence.transcript import TranscriptWriter
+from monkeybot.core.runtime.input_admission import InputAdmission
+from monkeybot.core.tools.permission import SessionApprovals
+
+from .sse import format_data_event
+
+PENDING_RESPONSE_TIMEOUT_SEC: float = float(os.environ.get("PENDING_RESPONSE_TIMEOUT_SEC", "300"))
+
+
+def _replay_maxlen_from_env() -> int:
+    raw = os.environ.get("SSE_REPLAY_MAX", "256")
+    try:
+        n = int(raw)
+        return max(1, n)
+    except ValueError:
+        return 256
+
+
+class SessionAlreadyExistsError(Exception):
+    """Raised when POST /sessions repeats an existing client-supplied id."""
+
+
+class SessionBus:
+    """Broadcasts framed SSE events; buffers numbered data events for replay."""
+
+    def __init__(
+        self,
+        *,
+        created_at_ms: int,
+        agent_md: str | None,
+        replay_maxlen: int | None = None,
+        provider: Any | None = None,
+        model_name: str | None = None,
+    ) -> None:
+        self.created_at_ms = created_at_ms
+        self.agent_md = agent_md
+        self.provider = provider
+        self.model_name = model_name
+        self.current_request_id: str | None = None
+        self.cancel_requested_for: str | None = None
+        self._seq = 0
+        maxlen = replay_maxlen if replay_maxlen is not None else _replay_maxlen_from_env()
+        self._replay: deque[tuple[int, str]] = deque(maxlen=maxlen)
+        self._subscribers: set[asyncio.Queue[str]] = set()
+        self._lock = asyncio.Lock()
+        self.pending_responses: dict[str, asyncio.Future[Any]] = {}
+        self.terminated_pending_keys: deque[str] = deque(maxlen=256)
+        self.attachment_catalog: SessionAttachmentCatalog | None = None
+        self.transcript_writer: TranscriptWriter | None = None
+        """Lazily-created ``TranscriptWriter`` (internal debugging only); None when disabled."""
+        self.admission = InputAdmission()
+        """Process-local steer + follow-up queues (not shared across gateway replicas)."""
+        self.session_approvals = SessionApprovals()
+        """Process-local 'always allow' rememberies (not shared across gateway replicas)."""
+        self.follow_up_retry_task: asyncio.Task[None] | None = None
+        """Scheduled drain retry after a failed durable turn-lock acquire."""
+
+    def cancel_follow_up_retry(self) -> None:
+        """Cancel any pending follow-up lock-retry task."""
+        task = self.follow_up_retry_task
+        self.follow_up_retry_task = None
+        if task is not None and not task.done():
+            task.cancel()
+
+    def register_pending(self, pending_key: str) -> asyncio.Future[Any]:
+        fut = asyncio.get_running_loop().create_future()
+        self.pending_responses[pending_key] = fut
+        return fut
+
+    def resolve_pending(self, pending_key: str, payload: Any) -> bool:
+        fut = self.pending_responses.get(pending_key)
+        if fut is None or fut.done():
+            return False
+        fut.set_result(payload)
+        self.pending_responses.pop(pending_key, None)
+        self.terminated_pending_keys.append(pending_key)
+        return True
+
+    def abandon_pending_timeout(self, pending_key: str) -> None:
+        fut = self.pending_responses.pop(pending_key, None)
+        if fut is None:
+            return
+        if not fut.done():
+            fut.cancel()
+        self.terminated_pending_keys.append(pending_key)
+
+    def abandon_pending_cancel_all(self) -> None:
+        for pending_key in list(self.pending_responses.keys()):
+            fut = self.pending_responses.pop(pending_key, None)
+            if fut is not None and not fut.done():
+                fut.cancel()
+            self.terminated_pending_keys.append(pending_key)
+
+    def is_pending_or_terminal(self, pending_key: str) -> Literal["pending", "terminated", "unknown"]:
+        if pending_key in self.pending_responses:
+            return "pending"
+        if pending_key in self.terminated_pending_keys:
+            return "terminated"
+        return "unknown"
+
+    async def publish_data(self, data_json: str) -> int:
+        """Buffer and broadcast one JSON data event; returns monotonic sequence id."""
+        async with self._lock:
+            self._seq += 1
+            seq = self._seq
+            frame = format_data_event(seq, data_json)
+            self._replay.append((seq, frame))
+            subscribers = list(self._subscribers)
+        for q in subscribers:
+            await q.put(frame)
+        return seq
+
+    async def publish_comment(self, comment_line: str) -> None:
+        """Send a comment/heartbeat line to live subscribers only (no replay)."""
+        async with self._lock:
+            subscribers = list(self._subscribers)
+        for q in subscribers:
+            await q.put(comment_line)
+
+    async def subscribe(
+        self, last_event_id: int | None
+    ) -> tuple[list[str], asyncio.Queue[str]]:
+        """
+        Register a subscriber and return buffered frames after last_event_id.
+
+        If last_event_id is None, replay all buffered frames (seq > 0).
+        """
+        async with self._lock:
+            q: asyncio.Queue[str] = asyncio.Queue()
+            self._subscribers.add(q)
+            cutoff = last_event_id if last_event_id is not None else 0
+            replay_frames = [frame for seq, frame in self._replay if seq > cutoff]
+        return replay_frames, q
+
+    async def unsubscribe(self, queue: asyncio.Queue[str]) -> None:
+        """Remove a subscriber queue (call from SSE disconnect finally)."""
+        async with self._lock:
+            self._subscribers.discard(queue)
+
+
+async def _await_user_response(
+    bus: SessionBus,
+    *,
+    pending_key: str,
+    timeout_sec: float | None = None,
+) -> dict[str, Any]:
+    """Block until POST resolves *pending_key*, timeout, cancel, or disconnect policy.
+
+    Returns:
+        Normal POST payloads (structure depends on flow).
+        On timeout: ``{"_timeout": True}`` (sentinel).
+
+    Raises:
+        asyncio.CancelledError: When the backing Future was cancelled (Stop button path).
+    """
+    from monkeybot.core.runtime.loop import _await_user_response_any
+
+    fut = bus.pending_responses[pending_key]
+    return await _await_user_response_any(bus, fut, pending_key, timeout_sec=timeout_sec)
+
+
+class SessionRegistry:
+    """Process-local registry of SessionBus instances."""
+
+    def __init__(self) -> None:
+        self._sessions: dict[str, SessionBus] = {}  # ponytail: in-process registry, use Redis pub/sub for multi-instance deployments
+
+    def get(self, session_id: str) -> SessionBus | None:
+        """Return the bus for id or None."""
+        return self._sessions.get(session_id)
+
+    def create(
+        self,
+        session_id: str,
+        *,
+        agent_md: str | None,
+        created_at_ms: int,
+        provider: Any | None = None,
+        model_name: str | None = None,
+    ) -> SessionBus:
+        """Create a new session bus or raise SessionAlreadyExistsError."""
+        if session_id in self._sessions:
+            raise SessionAlreadyExistsError(session_id)
+        bus = SessionBus(
+            created_at_ms=created_at_ms,
+            agent_md=agent_md,
+            provider=provider,
+            model_name=model_name,
+        )
+        bus.attachment_catalog = SessionAttachmentCatalog(session_id=session_id)
+        self._sessions[session_id] = bus
+        return bus
+
+    def remove(self, session_id: str) -> bool:
+        """Drop a session bus and any auxiliary per-session state keyed by it.
+
+        Cancels outstanding pending-response futures so awaiting callers don't
+        hang, then evicts the memory-curation cache entry for this thread id
+        (see ``memory_prompt._curation_cache``) so both structures share the
+        same lifecycle instead of growing unbounded for the life of the process.
+
+        Returns True if a session was found and removed, False otherwise.
+        """
+        bus = self._sessions.pop(session_id, None)
+        if bus is None:
+            return False
+        bus.abandon_pending_cancel_all()
+        bus.cancel_follow_up_retry()
+        bus.admission.clear_all()
+
+        from monkeybot.core.context.memory_prompt import evict_curation_cache
+
+        evict_curation_cache(session_id)
+        return True

monkeybot/gateway/sse/sse.py

@@ -0,0 +1,46 @@
+"""
+SSE wire framing and mapping from agent events (kind) to SSE JSON (type).
+"""
+
+import json
+from dataclasses import fields as dc_fields
+from typing import Any, cast
+
+from monkeybot.core.runtime.events import AgentEvent
+
+
+def format_data_event(seq: int, data: str) -> str:
+    """Build one SSE data event with an `id` line (stored in replay buffer)."""
+    return f"id: {seq}\ndata: {data}\n\n"
+
+
+def format_ping(n: int) -> str:
+    """Heartbeat comment line; not buffered for replay."""
+    return f": ping {n}\n\n"
+
+
+def format_active_requests(request_ids: list[str]) -> str:
+    """ActiveRequests frame without an `id:` line."""
+    payload = json_dumps_wire({"type": "ActiveRequests", "request_ids": request_ids})
+    return f"data: {payload}\n\n"
+
+
+def agent_event_to_wire_dict(event: AgentEvent) -> dict[str, Any]:
+    """Map an AgentEvent to an SSE JSON dict with ``type`` + ``chat_request_id``."""
+    kind = str(getattr(event, "kind"))
+    payload = {
+        f.name: getattr(event, f.name)
+        for f in dc_fields(cast(Any, event))
+        if f.name != "kind"
+    }
+    rid = str(payload.get("request_id", ""))
+    out: dict[str, Any] = {"type": kind, "request_id": rid, "chat_request_id": rid}
+    for key, val in payload.items():
+        if key != "request_id":
+            out[key] = val
+    return out
+
+
+def json_dumps_wire(obj: dict[str, Any]) -> str:
+    """Stable JSON for SSE data lines (no extra whitespace)."""
+    return json.dumps(obj, separators=(",", ":"))

monkeybot/gateway/sse/workspace_layout.py

@@ -0,0 +1,7 @@
+"""Agent-facing workspace root (workspace file API + loop ``workspace_root``)."""
+
+from __future__ import annotations
+
+from monkeybot.core.workspace_layout import resolve_agent_workspace_root
+
+__all__ = ["resolve_agent_workspace_root"]