monkeybot 2.1.1__py3-none-any.whl

monkeybot/core/subagents/worker_pool.py

@@ -0,0 +1,403 @@
+"""Poll-and-claim worker pool for queued subagent runs."""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import logging
+import os
+import time
+import uuid
+from dataclasses import dataclass
+from pathlib import Path
+
+from monkeybot.core.config.settings import auto_schema_enabled_from_config
+from monkeybot.core.logging_utils import kv
+from monkeybot.core.persistence.backends import RunStore, StorageBackend, create_storage_backend
+from monkeybot.core.persistence.durable_runs import SubagentRunRow
+from monkeybot.core.runtime.events import Error, TurnComplete, event_to_json
+from monkeybot.core.subagents.subagent_proto import (
+    SubagentEnvelope,
+    resolve_subagent_script,
+    spawn_subagent,
+)
+
+logger = logging.getLogger(__name__)
+
+_SHUTDOWN_FAILURE_MESSAGE = "subagent run cancelled during worker shutdown"
+_STALE_CLAIM_WARN_FRACTION = 0.8
+
+
+def _env_float(name: str, default: float) -> float:
+    raw = os.environ.get(name, "").strip()
+    if not raw:
+        return default
+    try:
+        return float(raw)
+    except ValueError:
+        return default
+
+
+def _env_int(name: str, default: int) -> int:
+    raw = os.environ.get(name, "").strip()
+    if not raw:
+        return default
+    try:
+        return int(raw)
+    except ValueError:
+        return default
+
+
+@dataclass(frozen=True)
+class WorkerEnvSettings:
+    poll_interval_s: float
+    concurrency: int
+    stale_claim_ms: int
+
+
+def worker_env_settings() -> WorkerEnvSettings:
+    """Read worker-pool tuning from environment once."""
+    return WorkerEnvSettings(
+        poll_interval_s=_env_float("MONKEYBOT_WORKER_POLL_INTERVAL_S", 2.0),
+        concurrency=_env_int("MONKEYBOT_WORKER_CONCURRENCY", 1),
+        stale_claim_ms=_env_int("MONKEYBOT_WORKER_STALE_CLAIM_MS", 600_000),
+    )
+
+
+@dataclass
+class WorkerPoolHandle:
+    """In-process worker pool task plus shared shutdown state."""
+
+    task: asyncio.Task[None]
+    backend: StorageBackend
+    active_runs: set[str]
+    worker_id: str
+
+
+def resolve_worker_id() -> str:
+    explicit = os.environ.get("MONKEYBOT_WORKER_ID", "").strip()
+    if explicit:
+        return explicit
+    return f"worker-{uuid.uuid4().hex[:12]}"
+
+
+async def _still_owns_run(
+    run_store: RunStore,
+    run_id: str,
+    worker_id: str,
+    *,
+    claimed_at: int | None = None,
+) -> bool:
+    row = await run_store.get_run(run_id)
+    if row is None or row.status != "running" or row.worker_id != worker_id:
+        return False
+    if claimed_at is not None and row.claimed_at != claimed_at:
+        return False
+    return True
+
+
+async def _fail_in_flight_runs(
+    run_store: RunStore,
+    run_ids: set[str],
+    *,
+    worker_id: str,
+    message: str,
+) -> None:
+    for run_id in list(run_ids):
+        try:
+            if await _still_owns_run(run_store, run_id, worker_id):
+                await run_store.record_failed(run_id, message)
+        except Exception:
+            logger.exception("failed to mark run_id=%s as failed during shutdown", run_id)
+        finally:
+            run_ids.discard(run_id)
+
+
+async def _stale_claim_watchdog(
+    run_id: str,
+    claimed_at: int,
+    stale_claim_ms: int,
+) -> None:
+    """Warn once when a run approaches the stale-claim reclaim window."""
+    warn_after_ms = int(stale_claim_ms * _STALE_CLAIM_WARN_FRACTION)
+    deadline_s = max(0.0, (claimed_at + warn_after_ms - int(time.time() * 1000)) / 1000.0)
+    if deadline_s > 0:
+        await asyncio.sleep(deadline_s)
+    logger.warning(
+        "subagent run_id=%s has been running for >=%.0f%% of MONKEYBOT_WORKER_STALE_CLAIM_MS "
+        "(%d ms); if execution exceeds the full window another worker may reclaim and "
+        "re-execute the run",
+        run_id,
+        _STALE_CLAIM_WARN_FRACTION * 100,
+        stale_claim_ms,
+    )
+
+
+async def execute_claimed_run(
+    run_store: RunStore,
+    row: SubagentRunRow,
+    *,
+    script: Path,
+    worker_id: str,
+    stale_claim_ms: int = 600_000,
+) -> None:
+    """Run one claimed subagent row to completion and persist the outcome."""
+    envelope = SubagentEnvelope.from_json(row.envelope_json)
+    logger.debug(
+        "worker executing %s",
+        kv(run_id=row.run_id, worker_id=worker_id, parent_run_id=envelope.parent_run_id),
+    )
+    scratch = Path(row.scratch_dir)
+    errors: list[str] = []
+    last_event_json: str | None = None
+    claimed_at = row.claimed_at
+    if claimed_at is None:
+        fresh = await run_store.get_run(row.run_id)
+        if fresh is not None:
+            claimed_at = fresh.claimed_at
+    watchdog: asyncio.Task[None] | None = None
+    if claimed_at is not None and stale_claim_ms > 0:
+        watchdog = asyncio.create_task(
+            _stale_claim_watchdog(row.run_id, claimed_at, stale_claim_ms),
+            name=f"monkeybot-stale-watchdog-{row.run_id}",
+        )
+
+    async def _record_failed_if_owner(message: str) -> bool:
+        if not await _still_owns_run(
+            run_store,
+            row.run_id,
+            worker_id,
+            claimed_at=claimed_at,
+        ):
+            logger.warning(
+                "skipping failed outcome for run_id=%s worker_id=%s: claim lost",
+                row.run_id,
+                worker_id,
+            )
+            return False
+        await run_store.record_failed(row.run_id, message)
+        return True
+
+    async def _record_completed_if_owner(result_json: str) -> bool:
+        if not await _still_owns_run(
+            run_store,
+            row.run_id,
+            worker_id,
+            claimed_at=claimed_at,
+        ):
+            logger.warning(
+                "skipping completed outcome for run_id=%s worker_id=%s: claim lost",
+                row.run_id,
+                worker_id,
+            )
+            return False
+        await run_store.record_completed(row.run_id, result_json)
+        return True
+
+    extra_env: dict[str, str] | None = None
+    if envelope.agent_md:
+        extra_env = {"MONKEYBOT_SUBAGENT_AGENT_MD": envelope.agent_md}
+
+    try:
+        async for evt in spawn_subagent(
+            str(script),
+            envelope,
+            scratch_dir=scratch,
+            extra_env=extra_env,
+        ):
+            if isinstance(evt, Error):
+                errors.append(evt.error)
+            elif isinstance(evt, TurnComplete):
+                last_event_json = event_to_json(evt)
+    except asyncio.CancelledError:
+        await _record_failed_if_owner(_SHUTDOWN_FAILURE_MESSAGE)
+        raise
+    except Exception as exc:
+        logger.exception("worker failed executing run_id=%s", row.run_id)
+        errors.append(str(exc))
+    finally:
+        if watchdog is not None:
+            watchdog.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await watchdog
+
+    if errors:
+        logger.debug(
+            "worker finished %s",
+            kv(run_id=row.run_id, worker_id=worker_id, outcome="failed"),
+        )
+        await _record_failed_if_owner("; ".join(errors))
+        return
+
+    if last_event_json is None:
+        logger.debug(
+            "worker finished %s",
+            kv(run_id=row.run_id, worker_id=worker_id, outcome="failed_no_turn_complete"),
+        )
+        await _record_failed_if_owner("subagent produced no TurnComplete event")
+        return
+
+    logger.debug(
+        "worker finished %s",
+        kv(run_id=row.run_id, worker_id=worker_id, outcome="completed"),
+    )
+    await _record_completed_if_owner(last_event_json)
+
+
+async def run_worker_loop(
+    backend: StorageBackend,
+    *,
+    worker_id: str,
+    script: Path,
+    poll_interval_s: float = 2.0,
+    concurrency: int = 1,
+    stale_claim_ms: int = 600_000,
+    active_runs: set[str] | None = None,
+) -> None:
+    """Poll ``pending_runs``, claim work atomically, and execute claimed rows."""
+    run_store = backend.runs()
+    sem = asyncio.Semaphore(max(1, concurrency))
+    in_flight = active_runs if active_runs is not None else set()
+
+    async def _try_claim_and_run(row: SubagentRunRow) -> None:
+        if row.status != "pending":
+            return
+        async with sem:
+            if not await run_store.claim(row.run_id, worker_id):
+                return
+            claimed_row = await run_store.get_run(row.run_id)
+            if claimed_row is None:
+                return
+            logger.debug(
+                "worker claimed %s",
+                kv(run_id=row.run_id, worker_id=worker_id),
+            )
+            in_flight.add(row.run_id)
+            try:
+                await execute_claimed_run(
+                    run_store,
+                    claimed_row,
+                    script=script,
+                    worker_id=worker_id,
+                    stale_claim_ms=stale_claim_ms,
+                )
+            except asyncio.CancelledError:
+                await _fail_in_flight_runs(
+                    run_store,
+                    {row.run_id},
+                    worker_id=worker_id,
+                    message=_SHUTDOWN_FAILURE_MESSAGE,
+                )
+                raise
+            finally:
+                in_flight.discard(row.run_id)
+
+    try:
+        while True:
+            try:
+                reset_count = await run_store.reset_stale_claims(stale_claim_ms)
+                if reset_count:
+                    logger.info(
+                        "reset stale subagent claims %s",
+                        kv(count=reset_count, worker_id=worker_id),
+                    )
+                pending = await run_store.pending_runs()
+                tasks = [_try_claim_and_run(row) for row in pending if row.status == "pending"]
+                if tasks:
+                    await asyncio.gather(*tasks)
+            except asyncio.CancelledError:
+                raise
+            except Exception:
+                logger.exception("worker poll loop error worker_id=%s", worker_id)
+            await asyncio.sleep(poll_interval_s)
+    except asyncio.CancelledError:
+        await _fail_in_flight_runs(
+            run_store,
+            in_flight,
+            worker_id=worker_id,
+            message=_SHUTDOWN_FAILURE_MESSAGE,
+        )
+        raise
+
+
+def start_worker_pool_background(
+    backend: StorageBackend,
+    *,
+    worker_id: str | None = None,
+    script: Path | None = None,
+) -> WorkerPoolHandle:
+    """Start the poll-and-claim loop as a background asyncio task.
+
+    Development-only entry point: the loop runs on the caller's event loop (e.g. the
+    gateway process), so subagent execution competes with whatever else that loop
+    serves (SSE streams) and there is no backpressure between them. For production,
+    run standalone worker processes via ``python -m monkeybot.subagents.worker``
+    (see :func:`run_worker_main`), which scale independently of the gateway.
+    """
+    wid = worker_id or resolve_worker_id()
+    resolved_script = script or resolve_subagent_script()
+    settings = worker_env_settings()
+    active_runs: set[str] = set()
+    logger.info(
+        "starting subagent worker pool worker_id=%s script=%s concurrency=%d "
+        "stale_claim_ms=%d (MONKEYBOT_WORKER_STALE_CLAIM_MS; runs without heartbeat "
+        "may be reclaimed and re-executed after this window)",
+        wid,
+        resolved_script,
+        settings.concurrency,
+        settings.stale_claim_ms,
+    )
+    task = asyncio.create_task(
+        run_worker_loop(
+            backend,
+            worker_id=wid,
+            script=resolved_script,
+            poll_interval_s=settings.poll_interval_s,
+            concurrency=settings.concurrency,
+            stale_claim_ms=settings.stale_claim_ms,
+            active_runs=active_runs,
+        ),
+        name=f"monkeybot-worker-{wid}",
+    )
+    return WorkerPoolHandle(
+        task=task,
+        backend=backend,
+        active_runs=active_runs,
+        worker_id=wid,
+    )
+
+
+async def shutdown_worker_pool(handle: WorkerPoolHandle) -> None:
+    """Cancel an in-process worker pool and fail any still-running claimed rows."""
+    handle.task.cancel()
+    with contextlib.suppress(asyncio.CancelledError):
+        await handle.task
+    await _fail_in_flight_runs(
+        handle.backend.runs(),
+        handle.active_runs,
+        worker_id=handle.worker_id,
+        message=_SHUTDOWN_FAILURE_MESSAGE,
+    )
+
+
+async def run_worker_main() -> None:
+    """CLI entry: open storage backend and run the worker loop until cancelled."""
+    db_url = os.environ.get("DB_URL", "sqlite:///data/monkeybot.db")
+    settings = worker_env_settings()
+    logger.info(
+        "subagent worker starting stale_claim_ms=%d (MONKEYBOT_WORKER_STALE_CLAIM_MS)",
+        settings.stale_claim_ms,
+    )
+    backend = create_storage_backend(db_url)
+    await backend.open(run_schema=auto_schema_enabled_from_config())
+    try:
+        await run_worker_loop(
+            backend,
+            worker_id=resolve_worker_id(),
+            script=resolve_subagent_script(),
+            poll_interval_s=settings.poll_interval_s,
+            concurrency=settings.concurrency,
+            stale_claim_ms=settings.stale_claim_ms,
+        )
+    finally:
+        await backend.close()

monkeybot/core/testing/__init__.py

@@ -0,0 +1 @@
+"""Test doubles for core components."""

monkeybot/core/testing/mocks_provider.py

@@ -0,0 +1,86 @@
+"""Test doubles for :class:`~monkeybot.core.llm.provider.Provider` without network."""
+
+from __future__ import annotations
+
+import json
+from collections.abc import AsyncIterator, Sequence
+
+from monkeybot.core.llm.provider import Message, ProviderEvent
+from monkeybot.core.types.content_blocks import ContentBlock, Text, ToolRequest, ToolResponse
+from monkeybot.core.types.types_tools import ToolDef
+
+
+def fake_provider_prompt_tokens(messages: Sequence[Message], tools: Sequence[ToolDef]) -> int:
+    """Deterministic char÷4 tally for fake providers in unit tests (no vendor APIs)."""
+
+    def char_blocks(blocks: Sequence[ContentBlock]) -> int:
+        total = 0
+        for b in blocks:
+            if isinstance(b, Text):
+                total += len(b.text)
+            elif isinstance(b, ToolRequest):
+                total += len(b.id) + len(b.name) + len(json.dumps(b.args, sort_keys=True))
+            elif isinstance(b, ToolResponse):
+                total += len(b.id) + len(b.tool_name) + char_blocks(b.result)
+            else:
+                total += len(json.dumps(b.to_dict(), sort_keys=True))
+        return total
+
+    n = sum(char_blocks(m.content) for m in messages) // 4
+    for t in tools:
+        n += (
+            len(t.name)
+            + len(t.description)
+            + len(json.dumps(t.input_schema, sort_keys=True, default=str))
+        ) // 4
+    return max(0, n)
+
+
+class ScriptedFakeProvider:
+    """Deterministic provider yielding a fixed :class:`~monkeybot.core.llm.provider.ProviderEvent` list."""
+
+    def __init__(
+        self,
+        events: list[ProviderEvent],
+        *,
+        name: str = "fake",
+        supports_streaming: bool = True,
+    ) -> None:
+        self._events = list(events)
+        self._name = name
+        self._supports_streaming = supports_streaming
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    @property
+    def supports_streaming(self) -> bool:
+        return self._supports_streaming
+
+    async def stream(
+        self,
+        messages: Sequence[Message],
+        tools: Sequence[ToolDef],
+        *,
+        model: str,
+        thinking_budget: int | None = None,
+    ) -> AsyncIterator[ProviderEvent]:
+        del thinking_budget
+        # Copy protects against callers mutating the original list mid-stream.
+        for ev in list(self._events):
+            yield ev
+
+    async def count_input_tokens(
+        self,
+        messages: Sequence[Message],
+        tools: Sequence[ToolDef],
+        *,
+        model: str,
+        thinking_budget: int | None = None,
+    ) -> int:
+        del model, thinking_budget
+        return fake_provider_prompt_tokens(messages, tools)
+
+
+__all__ = ["ScriptedFakeProvider", "fake_provider_prompt_tokens"]

monkeybot/core/testing/mocks_realtime_provider.py

@@ -0,0 +1,137 @@
+"""Test doubles for :class:`~monkeybot.core.llm.realtime_provider.RealtimeSession`."""
+
+from __future__ import annotations
+
+from collections.abc import AsyncIterator, Sequence
+
+from monkeybot.core.llm.realtime_provider import (
+    AudioFormat,
+    RealtimeEvent,
+    RealtimeSession,
+    RealtimeSessionConfig,
+)
+
+
+class ScriptedRealtimeSession:
+    """Deterministic in-memory realtime session for unit tests."""
+
+    def __init__(
+        self,
+        events: list[RealtimeEvent],
+        *,
+        input_format: AudioFormat | None = None,
+        output_format: AudioFormat | None = None,
+    ) -> None:
+        self._events = list(events)
+        self._input_format = input_format or AudioFormat(
+            encoding="pcm_s16le", sample_rate_hz=24000, channels=1, frame_ms=200
+        )
+        self._output_format = output_format or AudioFormat(
+            encoding="pcm_s16le", sample_rate_hz=24000, channels=1, frame_ms=200
+        )
+        self._audio_sent: list[bytes] = []
+        self._text_sent: list[str] = []
+        self._context_sent: list[str] = []
+        self._tool_results_sent: list[tuple[str, str, dict[str, object], bool]] = []
+        self._audio_turn_ends = 0
+        self._interrupt_count = 0
+        self._closed = False
+        self._close_reason: str | None = None
+
+    @property
+    def input_format(self) -> AudioFormat:
+        return self._input_format
+
+    @property
+    def output_format(self) -> AudioFormat:
+        return self._output_format
+
+    async def send_audio(self, chunk: bytes) -> None:
+        self._audio_sent.append(chunk)
+
+    async def end_audio_turn(self) -> None:
+        self._audio_turn_ends += 1
+
+    async def send_text(self, text: str) -> None:
+        self._text_sent.append(text)
+
+    async def send_context(self, text: str) -> None:
+        self._context_sent.append(text)
+
+    async def send_tool_results(
+        self,
+        results: Sequence[tuple[str, str, dict[str, object], bool]],
+    ) -> None:
+        self._tool_results_sent.extend(results)
+
+    async def interrupt(self) -> None:
+        self._interrupt_count += 1
+
+    def events(self) -> AsyncIterator[RealtimeEvent]:
+        async def _gen() -> AsyncIterator[RealtimeEvent]:
+            for ev in list(self._events):
+                yield ev
+
+        return _gen()
+
+    async def close(self, *, reason: str = "session_end") -> None:
+        self._closed = True
+        self._close_reason = reason
+
+    def audio_sent(self) -> list[bytes]:
+        return list(self._audio_sent)
+
+    def audio_turn_ends(self) -> int:
+        return self._audio_turn_ends
+
+    def text_sent(self) -> list[str]:
+        return list(self._text_sent)
+
+    def context_sent(self) -> list[str]:
+        return list(self._context_sent)
+
+    def tool_results_sent(self) -> list[tuple[str, str, dict[str, object], bool]]:
+        return list(self._tool_results_sent)
+
+    def interrupt_count(self) -> int:
+        return self._interrupt_count
+
+    def is_closed(self) -> bool:
+        return self._closed
+
+    def close_reason(self) -> str | None:
+        return self._close_reason
+
+
+class ScriptedRealtimeProvider:
+    """Deterministic factory for :class:`ScriptedRealtimeSession`."""
+
+    def __init__(
+        self,
+        *,
+        name: str = "fake-realtime",
+        sessions: list[ScriptedRealtimeSession] | None = None,
+    ) -> None:
+        self._name = name
+        self._sessions = sessions or []
+        self._calls: list[tuple[str, RealtimeSessionConfig]] = []
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    async def connect(
+        self,
+        *,
+        config: RealtimeSessionConfig,
+    ) -> RealtimeSession:
+        self._calls.append((config.model, config))
+        if not self._sessions:
+            return ScriptedRealtimeSession([])
+        return self._sessions.pop(0)
+
+    def connect_calls(self) -> list[tuple[str, RealtimeSessionConfig]]:
+        return list(self._calls)
+
+
+__all__ = ["ScriptedRealtimeProvider", "ScriptedRealtimeSession"]

monkeybot/core/tools/__init__.py

@@ -0,0 +1 @@
+"""Tool execution, workspace I/O, sandbox, terminal, and inspection."""