monkeybot 2.1.1__py3-none-any.whl

monkeybot/core/persistence/firestore_scheduled_loops.py

@@ -0,0 +1,336 @@
+"""Firestore persistence for prompt-first scheduled agent loops."""
+
+from __future__ import annotations
+
+import time
+from collections.abc import Awaitable, Callable
+from typing import cast
+
+from google.cloud import firestore
+from google.cloud.firestore import AsyncClient
+from google.cloud.firestore_v1.base_query import FieldFilter
+
+from monkeybot.core.persistence.scheduled_loops import (
+    ScheduledLoopCreate,
+    ScheduledLoopRow,
+    _loop_id_from_create,
+    doc_to_scheduled_loop_row,
+    validate_loop_guards,
+)
+
+_LOOP_STATUSES = frozenset({"active", "paused", "completed", "failed"})
+
+
+def _collection_name(prefix: str, base: str) -> str:
+    if not prefix:
+        return base
+    return f"{prefix}_{base}"
+
+
+class FirestoreScheduledLoopStore:
+    """Firestore-backed scheduled-loop store."""
+
+    def __init__(self, client: AsyncClient, prefix: str) -> None:
+        self._client = client
+        self._collection = _collection_name(prefix, "scheduled_loops")
+
+    def _doc(self, loop_id: str) -> firestore.AsyncDocumentReference:
+        return self._client.collection(self._collection).document(loop_id)
+
+    async def create(self, spec: ScheduledLoopCreate) -> ScheduledLoopRow:
+        loop_id = _loop_id_from_create(spec)
+        validate_loop_guards(
+            max_ticks=spec.max_ticks,
+            max_runtime_ms=spec.max_runtime_ms,
+            unbounded=spec.unbounded,
+        )
+        now_ms = int(time.time() * 1000)
+        payload = {
+            "session_id": spec.session_id.strip() or "loop-main",
+            "status": "active",
+            "prompt": spec.prompt.strip(),
+            "interval_ms": spec.interval_ms,
+            "max_ticks": spec.max_ticks,
+            "max_runtime_ms": spec.max_runtime_ms,
+            "skip_if_busy": 1 if spec.skip_if_busy else 0,
+            "tick_index": 0,
+            "next_tick_at_ms": now_ms,
+            "started_at_ms": now_ms,
+            "last_tick_at_ms": None,
+            "last_error": None,
+            "stop_reason": None,
+            "tick_in_flight": 0,
+            "worker_id": None,
+            "claimed_at_ms": None,
+        }
+        doc_ref = self._doc(loop_id)
+        transaction = self._client.transaction()
+
+        async def _create_body(
+            txn: firestore.AsyncTransaction,
+            ref: firestore.AsyncDocumentReference,
+        ) -> None:
+            snapshot = await ref.get(transaction=txn)
+            if snapshot.exists:
+                raise ValueError(f"scheduled loop already exists: {loop_id}")
+            txn.set(ref, payload)
+
+        create_txn = cast(
+            Callable[[firestore.AsyncTransaction, firestore.AsyncDocumentReference], Awaitable[None]],
+            firestore.async_transactional(_create_body),
+        )
+        await create_txn(transaction, doc_ref)
+        row = await self.get(loop_id)
+        if row is None:
+            raise RuntimeError("failed to read scheduled loop after insert")
+        return row
+
+    async def get(self, loop_id: str) -> ScheduledLoopRow | None:
+        snapshot = await self._doc(loop_id).get()
+        if not snapshot.exists:
+            return None
+        return doc_to_scheduled_loop_row(snapshot.id, snapshot.to_dict() or {})
+
+    async def list_all(self) -> list[ScheduledLoopRow]:
+        rows: list[ScheduledLoopRow] = []
+        query = self._client.collection(self._collection).order_by(
+            "started_at_ms", direction=firestore.Query.DESCENDING
+        )
+        async for doc in query.stream():
+            rows.append(doc_to_scheduled_loop_row(doc.id, doc.to_dict() or {}))
+        return rows
+
+    async def list_due(self, now_ms: int) -> list[ScheduledLoopRow]:
+        rows: list[ScheduledLoopRow] = []
+        query = (
+            self._client.collection(self._collection)
+            .where(filter=FieldFilter("status", "==", "active"))
+            .where(filter=FieldFilter("tick_in_flight", "==", 0))
+            .where(filter=FieldFilter("next_tick_at_ms", "<=", now_ms))
+            .order_by("next_tick_at_ms")
+        )
+        async for doc in query.stream():
+            rows.append(doc_to_scheduled_loop_row(doc.id, doc.to_dict() or {}))
+        return rows
+
+    async def claim_tick(self, loop_id: str, worker_id: str) -> ScheduledLoopRow | None:
+        doc_ref = self._doc(loop_id)
+        transaction = self._client.transaction()
+        now_ms = int(time.time() * 1000)
+
+        async def _claim_body(
+            txn: firestore.AsyncTransaction,
+            ref: firestore.AsyncDocumentReference,
+        ) -> bool:
+            snapshot = await ref.get(transaction=txn)
+            if not snapshot.exists:
+                return False
+            data = snapshot.to_dict() or {}
+            if data.get("status") != "active":
+                return False
+            if int(cast(int, data.get("tick_in_flight", 0))) != 0:
+                return False
+            if int(cast(int, data.get("next_tick_at_ms", 0))) > now_ms:
+                return False
+            txn.update(
+                ref,
+                {
+                    "tick_in_flight": 1,
+                    "worker_id": worker_id,
+                    "claimed_at_ms": now_ms,
+                },
+            )
+            return True
+
+        claim_txn = cast(
+            Callable[[firestore.AsyncTransaction, firestore.AsyncDocumentReference], Awaitable[bool]],
+            firestore.async_transactional(_claim_body),
+        )
+        if not await claim_txn(transaction, doc_ref):
+            return None
+        return await self.get(loop_id)
+
+    async def release_stale_claims(self, stale_after_ms: int) -> int:
+        cutoff = int(time.time() * 1000) - stale_after_ms
+        query = (
+            self._client.collection(self._collection)
+            .where(filter=FieldFilter("tick_in_flight", "==", 1))
+            .where(filter=FieldFilter("claimed_at_ms", "<", cutoff))
+        )
+        reset = 0
+        batch = self._client.batch()
+        count = 0
+        async for doc in query.stream():
+            data = doc.to_dict() or {}
+            if data.get("claimed_at_ms") is None:
+                continue
+            batch.update(
+                doc.reference,
+                {
+                    "tick_in_flight": 0,
+                    "worker_id": None,
+                    "claimed_at_ms": None,
+                    "last_error": data.get("last_error") or "stale tick claim released",
+                },
+            )
+            reset += 1
+            count += 1
+            if count >= 400:
+                await batch.commit()
+                batch = self._client.batch()
+                count = 0
+        if count:
+            await batch.commit()
+        return reset
+
+    async def complete_tick(
+        self,
+        loop_id: str,
+        *,
+        worker_id: str,
+        error: str | None = None,
+    ) -> ScheduledLoopRow | None:
+        doc_ref = self._doc(loop_id)
+        transaction = self._client.transaction()
+        now_ms = int(time.time() * 1000)
+
+        async def _complete_body(
+            txn: firestore.AsyncTransaction,
+            ref: firestore.AsyncDocumentReference,
+        ) -> bool:
+            snapshot = await ref.get(transaction=txn)
+            if not snapshot.exists:
+                return False
+            data = snapshot.to_dict() or {}
+            if data.get("worker_id") != worker_id or int(data.get("tick_in_flight", 0)) != 1:
+                return False
+            tick_index = int(cast(int, data.get("tick_index", 0))) + 1
+            status = str(data.get("status", "active"))
+            stop_reason: str | None = None
+            interval_ms = int(cast(int, data.get("interval_ms", 0)))
+            started_at_ms = int(cast(int, data.get("started_at_ms", now_ms)))
+            if error:
+                status = "failed"
+                stop_reason = "tick_error"
+            else:
+                max_ticks = data.get("max_ticks")
+                max_runtime_ms = data.get("max_runtime_ms")
+                if max_ticks is not None and tick_index >= int(cast(int, max_ticks)):
+                    status = "completed"
+                    stop_reason = "max_ticks"
+                elif (
+                    max_runtime_ms is not None
+                    and (now_ms - started_at_ms) >= int(cast(int, max_runtime_ms))
+                ):
+                    status = "completed"
+                    stop_reason = "max_runtime"
+            next_tick_at = (
+                now_ms + interval_ms
+                if status == "active"
+                else int(cast(int, data.get("next_tick_at_ms", now_ms)))
+            )
+            txn.update(
+                ref,
+                {
+                    "tick_index": tick_index,
+                    "last_tick_at_ms": now_ms,
+                    "last_error": error,
+                    "status": status,
+                    "stop_reason": stop_reason,
+                    "next_tick_at_ms": next_tick_at,
+                    "tick_in_flight": 0,
+                    "worker_id": None,
+                    "claimed_at_ms": None,
+                },
+            )
+            return True
+
+        complete_txn = cast(
+            Callable[[firestore.AsyncTransaction, firestore.AsyncDocumentReference], Awaitable[bool]],
+            firestore.async_transactional(_complete_body),
+        )
+        if not await complete_txn(transaction, doc_ref):
+            return None
+        return await self.get(loop_id)
+
+    async def defer_tick(self, loop_id: str, *, worker_id: str, reason: str) -> None:
+        row = await self.get(loop_id)
+        if row is None or row.worker_id != worker_id:
+            return
+        now_ms = int(time.time() * 1000)
+        await self._doc(loop_id).update(
+            {
+                "tick_in_flight": 0,
+                "worker_id": None,
+                "claimed_at_ms": None,
+                "next_tick_at_ms": now_ms + row.interval_ms,
+                "last_error": reason,
+            }
+        )
+
+    async def renew_tick_claim(self, loop_id: str, worker_id: str) -> bool:
+        doc_ref = self._doc(loop_id)
+        transaction = self._client.transaction()
+        now_ms = int(time.time() * 1000)
+
+        async def _renew_body(
+            txn: firestore.AsyncTransaction,
+            ref: firestore.AsyncDocumentReference,
+        ) -> bool:
+            snapshot = await ref.get(transaction=txn)
+            if not snapshot.exists:
+                return False
+            data = snapshot.to_dict() or {}
+            if data.get("worker_id") != worker_id or int(data.get("tick_in_flight", 0)) != 1:
+                return False
+            txn.update(ref, {"claimed_at_ms": now_ms})
+            return True
+
+        renew_txn = cast(
+            Callable[[firestore.AsyncTransaction, firestore.AsyncDocumentReference], Awaitable[bool]],
+            firestore.async_transactional(_renew_body),
+        )
+        return bool(await renew_txn(transaction, doc_ref))
+
+    async def set_status(self, loop_id: str, status: str, *, stop_reason: str | None = None) -> bool:
+        if status not in _LOOP_STATUSES:
+            raise ValueError(f"invalid loop status: {status}")
+        snapshot = await self._doc(loop_id).get()
+        if not snapshot.exists:
+            return False
+        payload: dict[str, object] = {
+            "status": status,
+            "tick_in_flight": 0,
+            "worker_id": None,
+            "claimed_at_ms": None,
+        }
+        if stop_reason is not None:
+            payload["stop_reason"] = stop_reason
+        await self._doc(loop_id).update(payload)
+        return True
+
+    async def pause(self, loop_id: str) -> bool:
+        return await self.set_status(loop_id, "paused")
+
+    async def resume(self, loop_id: str) -> bool:
+        snapshot = await self._doc(loop_id).get()
+        if not snapshot.exists:
+            return False
+        data = snapshot.to_dict() or {}
+        if data.get("status") != "paused":
+            return False
+        now_ms = int(time.time() * 1000)
+        await self._doc(loop_id).update(
+            {
+                "status": "active",
+                "stop_reason": None,
+                "next_tick_at_ms": now_ms,
+                "tick_in_flight": 0,
+                "worker_id": None,
+                "claimed_at_ms": None,
+            }
+        )
+        return True
+
+    async def stop(self, loop_id: str, *, stop_reason: str = "manual") -> bool:
+        return await self.set_status(loop_id, "completed", stop_reason=stop_reason)

monkeybot/core/persistence/history.py

@@ -0,0 +1,156 @@
+"""SQLite-backed conversation history store.
+
+Durable conversation facts live here as ``Message`` rows with typed
+``ContentBlock`` JSON (not as ``AgentEvent`` rows). Tool settlement is the
+``ToolResponse`` block(s) on a user row after a tool batch; the live SSE
+``ToolCallResult`` event is the progress mirror (see ``events.is_durable_event``).
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import time
+from typing import cast
+
+import aiosqlite
+
+from monkeybot.core.types.content_blocks import ContentBlock
+from monkeybot.core.llm.provider import Message, Role
+from monkeybot.core.persistence.thread_summary import ChatThreadSummary, preview_from_content_blob
+
+logger = logging.getLogger("monkeybot.core.persistence.history")
+
+_VALID_ROLES: tuple[str, ...] = ("user", "assistant", "system")
+
+
+def _validate_message(message: Message) -> None:
+    """Reject persisted shapes that SQLite/provider contracts disallow."""
+    role = message.role
+    if role not in _VALID_ROLES:
+        raise ValueError(f"invalid role: {role!r}")
+
+
+class SQLiteHistoryStore:
+    """Append/read conversation rows keyed by ``thread_id``."""
+
+    def __init__(self, conn: aiosqlite.Connection) -> None:
+        self._conn = conn
+
+    async def append(self, thread_id: str, message: Message) -> None:
+        """Validate ``message``, JSON-encode content blocks, insert one row."""
+        _validate_message(message)
+        payload = json.dumps(
+            [b.to_dict() for b in message.content],
+            separators=(",", ":"),
+            ensure_ascii=False,
+        )
+        created_at = int(time.time() * 1000)
+        await self._conn.execute(
+            """
+            INSERT INTO conversation_history(thread_id, role, content, created_at)
+            VALUES (?, ?, ?, ?)
+            """,
+            (thread_id, message.role, payload, created_at),
+        )
+        await self._conn.commit()
+
+    async def load(self, thread_id: str, limit: int = 100) -> list[Message]:
+        """Return up to ``limit`` messages for ``thread_id``, oldest first."""
+        cursor = await self._conn.execute(
+            """
+            SELECT id, role, content, created_at
+            FROM conversation_history
+            WHERE thread_id = ?
+            ORDER BY created_at DESC, id DESC
+            LIMIT ?
+            """,
+            (thread_id, limit),
+        )
+        rows = await cursor.fetchall()
+        await cursor.close()
+        rows_chrono = list(reversed(list(rows)))
+        out: list[Message] = []
+        for row in rows_chrono:
+            row_id = int(row[0])
+            role = row[1]
+            content_blob = row[2]
+            try:
+                raw = json.loads(content_blob)
+                if not isinstance(raw, list):
+                    raise ValueError("stored content must be a JSON array")
+                blocks = [ContentBlock.from_dict(b) for b in raw]
+            except (json.JSONDecodeError, ValueError, TypeError) as exc:
+                logger.error(
+                    "Unparseable history row id=%s thread_id=%s",
+                    row_id,
+                    thread_id,
+                    exc_info=True,
+                )
+                raise ValueError(f"history row {row_id} unparseable: {exc}") from exc
+            if role not in _VALID_ROLES:
+                raise ValueError(f"history row {row_id} has invalid role: {role!r}")
+            out.append(Message(role=cast(Role, role), content=blocks))
+        return out
+
+    async def clear(self, thread_id: str) -> None:
+        """Delete every stored message for ``thread_id``."""
+        await self._conn.execute(
+            "DELETE FROM conversation_history WHERE thread_id = ?", (thread_id,)
+        )
+        await self._conn.commit()
+
+    async def reset(self, thread_id: str, messages: list[Message]) -> None:
+        """Replace the thread transcript with ``messages`` (validated like ``append``)."""
+        await self._conn.execute(
+            "DELETE FROM conversation_history WHERE thread_id = ?", (thread_id,)
+        )
+        await self._conn.commit()
+        for msg in messages:
+            await self.append(thread_id, msg)
+
+    async def list_threads(self, limit: int = 50) -> list[ChatThreadSummary]:
+        """Return recent threads ordered by last activity (newest first).
+
+        The correlated subquery on ``last_content`` is O(threads × messages) per call.
+        For production SQLite load, add a composite index on
+        ``conversation_history(thread_id, created_at DESC, id DESC)``.
+        """
+        cap = max(1, min(limit, 200))
+        cursor = await self._conn.execute(
+            """
+            SELECT
+                h.thread_id,
+                MAX(h.created_at) AS last_message_at,
+                COUNT(*) AS message_count,
+                (
+                    SELECT h2.content
+                    FROM conversation_history h2
+                    WHERE h2.thread_id = h.thread_id
+                    ORDER BY h2.created_at DESC, h2.id DESC
+                    LIMIT 1
+                ) AS last_content
+            FROM conversation_history h
+            GROUP BY h.thread_id
+            ORDER BY last_message_at DESC
+            LIMIT ?
+            """,
+            (cap,),
+        )
+        rows = await cursor.fetchall()
+        await cursor.close()
+        out: list[ChatThreadSummary] = []
+        for row in rows:
+            thread_id = str(row[0])
+            last_at = int(row[1])
+            count = int(row[2])
+            preview = preview_from_content_blob(str(row[3] or ""))
+            out.append(
+                ChatThreadSummary(
+                    thread_id=thread_id,
+                    last_message_at=last_at,
+                    message_count=count,
+                    preview=preview or "(empty)",
+                )
+            )
+        return out