"""MemoryTaskQueue - in-memory async task queue with lifecycle status tracking.""" from __future__ import annotations import asyncio import functools import inspect import json import logging import time from collections import deque from collections.abc import Awaitable, Callable from dataclasses import dataclass, field from enum import Enum from typing import Any from uuid import uuid4 from cyberbrain26.data_plane.sqlite_store import SQLiteDatabase logger = logging.getLogger(__name__) TaskFn = Callable[..., Awaitable[Any]] DurableTaskHandler = Callable[[dict[str, Any]], Awaitable[Any]] # Durable-result key that structurally marks a job as refused before the code # it carries was ever handed to an executor. It is the only signal # ``SQLiteTaskQueue.retry_blocked`` trusts: a result without this key, or with # any other value, is treated as "may have executed" and is never re-dispatched. EXECUTION_ATTEMPTED_KEY = "execution_attempted" class TaskStatus(str, Enum): PENDING = "pending" RUNNING = "running" COMPLETED = "completed" FAILED = "failed" @dataclass class TaskRecord: """A single queued unit of asynchronous work.""" id: str func: Callable[[], Awaitable[Any]] status: TaskStatus = TaskStatus.PENDING result: Any = None error: Exception | None = None attempts: int = 0 created_at: float = field(default_factory=time.time) started_at: float | None = None completed_at: float | None = None class MemoryTaskQueue: """Async FIFO task queue executing callables with status tracking. Tasks are enqueued as callables returning coroutines. `run` drains the pending queue (including tasks enqueued during dispatch), dispatching each through the optional concurrency semaphore and recording PENDING -> RUNNING -> COMPLETED|FAILED transitions. """ def __init__(self, max_concurrency: int = 1, max_retries: int = 0) -> None: self._queue: deque[TaskRecord] = deque() self._registry: dict[str, TaskRecord] = {} self._semaphore = asyncio.Semaphore(max(1, max_concurrency)) self._lock = asyncio.Lock() self._counter = 0 self.max_retries = max_retries @property def size(self) -> int: return len(self._queue) @property def pending(self) -> int: return sum(1 for t in self._registry.values() if t.status == TaskStatus.PENDING) def enqueue(self, func: TaskFn, *args: Any, **kwargs: Any) -> str: """Schedule `func(*args, **kwargs)` for execution and return its task id.""" self._counter += 1 task_id = uuid4().hex record = TaskRecord(id=task_id, func=functools.partial(func, *args, **kwargs)) self._queue.append(record) self._registry[task_id] = record logger.debug( "enqueued task %s", task_id, extra={"task_id": task_id, "queue_size": len(self._queue)}, ) return task_id def dequeue(self) -> TaskRecord | None: """Pop and return the next pending TaskRecord without executing it.""" if not self._queue: return None return self._queue.popleft() def get(self, task_id: str) -> TaskRecord | None: return self._registry.get(task_id) def status(self, task_id: str) -> TaskStatus | None: rec = self._registry.get(task_id) return rec.status if rec else None async def dispatch(self, record: TaskRecord) -> Any: """Execute a single TaskRecord, updating its status. Returns its result. On failure, retries up to ``self.max_retries`` times with exponential backoff (base delay 0.1s, doubling each attempt). """ record.status = TaskStatus.RUNNING record.started_at = time.time() logger.info( "dispatched task %s", record.id, extra={"task_id": record.id, "status": TaskStatus.RUNNING.value}, ) attempt = 0 while True: try: async with self._semaphore: record.attempts += 1 record.result = await record.func() record.status = TaskStatus.COMPLETED record.completed_at = time.time() logger.info( "task %s completed", record.id, extra={"task_id": record.id, "status": TaskStatus.COMPLETED.value}, ) return record.result except Exception as exc: record.error = exc record.status = TaskStatus.FAILED logger.exception( "task %s failed (attempt %d/%d)", record.id, attempt + 1, self.max_retries + 1, extra={ "task_id": record.id, "status": TaskStatus.FAILED.value, "error": str(exc), }, ) if attempt >= self.max_retries: record.completed_at = time.time() raise attempt += 1 backoff = 0.1 * (2 ** (attempt - 1)) logger.debug("retrying task %s after %.1fs", record.id, backoff) await asyncio.sleep(backoff) record.status = TaskStatus.RUNNING async def run(self) -> list[TaskRecord]: """Drain the queue, dispatching pending tasks (including those enqueued during dispatch) until no tasks remain.""" executed: list[TaskRecord] = [] async with self._lock: while self._queue: record = self._queue.popleft() executed.append(record) try: await self.dispatch(record) except Exception as exc: # noqa: BLE001 - isolation boundary, drain continues logger.debug( "past failure on draining task %s: %s", record.id, exc, extra={"task_id": record.id, "error": str(exc)}, ) return executed class PersistentTaskStatus(str, Enum): """Lifecycle state for JSON-backed jobs stored in SQLite.""" PENDING = "pending" RUNNING = "running" COMPLETED = "completed" FAILED = "failed" INTERRUPTED = "interrupted" @dataclass class PersistentTaskRecord: """Serializable job state that can be inspected after process restart.""" id: str handler: str payload: dict[str, Any] status: PersistentTaskStatus result: Any = None error: str | None = None attempts: int = 0 created_at: float = 0.0 started_at: float | None = None completed_at: float | None = None class SQLiteTaskQueue(MemoryTaskQueue): """Memory-compatible queue with an additional durable, named-job API. Existing agents may continue to enqueue Python callables in memory. Jobs that must survive restart use :meth:`enqueue_job` with a stable handler key and JSON payload. A process crash changes a running job to ``interrupted``; callers must explicitly retry it to avoid silently repeating side effects. """ def __init__( self, database: SQLiteDatabase, *, handlers: dict[str, DurableTaskHandler] | None = None, max_concurrency: int = 1, max_retries: int = 0, ) -> None: super().__init__(max_concurrency=max_concurrency, max_retries=max_retries) self.database = database self._handlers: dict[str, DurableTaskHandler] = dict(handlers or {}) with self.database.transaction() as connection: connection.execute( "UPDATE tasks SET status = ?, error = ? WHERE status = ?", ( PersistentTaskStatus.INTERRUPTED.value, "process stopped while task was running; explicit retry required", PersistentTaskStatus.RUNNING.value, ), ) @property def size(self) -> int: with self.database.connection() as connection: durable_pending = connection.execute( "SELECT COUNT(*) FROM tasks WHERE status = ?", (PersistentTaskStatus.PENDING.value,), ).fetchone()[0] return super().size + int(durable_pending) @property def pending(self) -> int: with self.database.connection() as connection: durable_pending = connection.execute( "SELECT COUNT(*) FROM tasks WHERE status = ?", (PersistentTaskStatus.PENDING.value,), ).fetchone()[0] return super().pending + int(durable_pending) def register_handler( self, name: str, handler: DurableTaskHandler, *, replace: bool = False, ) -> None: """Register a stable handler key used by persisted jobs.""" if not name or name.strip() != name: raise ValueError("handler name must be non-empty and trimmed") if name in self._handlers and not replace: raise ValueError(f"task handler already registered: {name}") self._handlers[name] = handler def enqueue_job( self, handler: str, payload: dict[str, Any], *, task_id: str | None = None, ) -> str: """Persist a named job and a JSON-serializable payload.""" if not handler or handler.strip() != handler: raise ValueError("handler name must be non-empty and trimmed") if not isinstance(payload, dict): raise TypeError("durable task payload must be a dictionary") try: payload_json = json.dumps(payload, separators=(",", ":"), allow_nan=False) except (TypeError, ValueError) as exc: raise TypeError("durable task payload must contain JSON values") from exc identifier = task_id or uuid4().hex created_at = time.time() with self.database.transaction() as connection: connection.execute( "INSERT INTO tasks(task_id, handler, payload_json, status, attempts, created_at) " "VALUES (?, ?, ?, ?, 0, ?)", ( identifier, handler, payload_json, PersistentTaskStatus.PENDING.value, created_at, ), ) return identifier def get_job(self, task_id: str) -> PersistentTaskRecord | None: with self.database.connection() as connection: row = connection.execute( "SELECT * FROM tasks WHERE task_id = ?", (task_id,) ).fetchone() return self._row_to_record(row) if row is not None else None def delete_job(self, task_id: str) -> bool: """Delete a durable job by task ID. Used by ``cyberbrain plan`` to discard a stale durable result for a planner step before re-planning with the same ``--id``, so the new plan is always freshly dispatched rather than read back from a prior run. Only terminal tasks (completed/failed) are removable; a running task is left untouched so an in-flight dispatch is never orphaned mid-flight. """ with self.database.transaction() as connection: deleted = connection.execute( "DELETE FROM tasks WHERE task_id = ? AND status != ?", ( task_id, PersistentTaskStatus.RUNNING.value, ), ).rowcount return deleted == 1 def list_jobs( self, status: PersistentTaskStatus | None = None, *, limit: int = 100, ) -> list[PersistentTaskRecord]: if limit < 0: raise ValueError("limit must not be negative") with self.database.connection() as connection: if status is None: rows = connection.execute( "SELECT * FROM tasks ORDER BY created_at DESC LIMIT ?", (limit,) ).fetchall() else: rows = connection.execute( "SELECT * FROM tasks WHERE status = ? ORDER BY created_at DESC LIMIT ?", (status.value, limit), ).fetchall() return [self._row_to_record(row) for row in rows] def retry_interrupted(self, task_id: str) -> bool: """Explicitly make an interrupted job eligible for another run.""" with self.database.transaction() as connection: updated = connection.execute( "UPDATE tasks SET status = ?, error = NULL WHERE task_id = ? AND status = ?", ( PersistentTaskStatus.PENDING.value, task_id, PersistentTaskStatus.INTERRUPTED.value, ), ).rowcount return updated == 1 def retry_failed(self, task_id: str) -> bool: """Explicitly make a failed job eligible for another dispatch. ``failed`` is only reached when the handler raised, so the job never produced a successful result. Re-dispatching it re-runs the handler from the top, which for Ptah means the full authorization chain again. This is deliberately different from ``retry_interrupted``: an interrupted job may have applied side effects before the process died, whereas a failed job has no approved result to protect. """ with self.database.transaction() as connection: updated = connection.execute( "UPDATE tasks SET status = ?, error = NULL WHERE task_id = ? AND status = ?", ( PersistentTaskStatus.PENDING.value, task_id, PersistentTaskStatus.FAILED.value, ), ).rowcount return updated == 1 def retry_blocked(self, task_id: str) -> bool: """Re-open a completed job that was refused before any code ran. A refusal is not a queue failure: the agent catches it and reports a result, so the job is stored as COMPLETED with ``success=False`` and ``retry_failed`` never applies. This method is the only way to make such a job dispatchable again, and it accepts nothing but a record that structurally states ``execution_attempted`` is false. Everything else is refused on purpose: a completed job whose result may represent executed code, and a result with no such field at all (a legacy row, written before the marker existed). Guessing there could duplicate side effects of code that really ran, so the check is a structural equality against a boolean, never a heuristic on messages or error text. """ with self.database.transaction() as connection: row = connection.execute( "SELECT status, result_json FROM tasks WHERE task_id = ?", (task_id,) ).fetchone() if row is None or row["status"] != PersistentTaskStatus.COMPLETED.value: return False result = json.loads(row["result_json"]) if row["result_json"] else None if not isinstance(result, dict): return False if result.get(EXECUTION_ATTEMPTED_KEY) is not False: return False if result.get("success"): return False updated = connection.execute( "UPDATE tasks SET status = ?, error = NULL, result_json = NULL, " "completed_at = NULL WHERE task_id = ? AND status = ?", ( PersistentTaskStatus.PENDING.value, task_id, PersistentTaskStatus.COMPLETED.value, ), ).rowcount return updated == 1 async def run_persisted(self) -> list[PersistentTaskRecord]: """Run pending named jobs in order and return their durable records.""" results: list[PersistentTaskRecord] = [] for job in reversed(self.list_jobs(PersistentTaskStatus.PENDING, limit=100_000)): result = await self.dispatch_job(job.id) results.append(result) return results async def dispatch_job(self, task_id: str) -> PersistentTaskRecord: """Dispatch one pending durable job by ID.""" job = self.get_job(task_id) if job is None: raise KeyError(f"durable task not found: {task_id}") handler = self._handlers.get(job.handler) if handler is None: raise LookupError( f"no registered handler for durable task {job.id}: {job.handler}" ) if job.status is not PersistentTaskStatus.PENDING: return job return await self._dispatch_job(job, handler) async def run(self) -> list[TaskRecord]: """Run legacy callable work, then durable named jobs.""" records = await super().run() await self.run_persisted() return records async def _dispatch_job( self, job: PersistentTaskRecord, handler: DurableTaskHandler, ) -> PersistentTaskRecord: attempt = 0 while True: started_at = time.time() with self.database.transaction() as connection: claimed = connection.execute( "UPDATE tasks SET status = ?, started_at = ?, attempts = attempts + 1, " "error = NULL WHERE task_id = ? AND status = ?", ( PersistentTaskStatus.RUNNING.value, started_at, job.id, PersistentTaskStatus.PENDING.value, ), ).rowcount if claimed != 1: current = self.get_job(job.id) return current if current is not None else job try: value = handler(job.payload) if not inspect.isawaitable(value): raise TypeError("durable task handlers must return an awaitable") result = await value result_json = json.dumps(result, separators=(",", ":"), allow_nan=False) except asyncio.CancelledError: with self.database.transaction() as connection: connection.execute( "UPDATE tasks SET status = ?, error = ?, completed_at = NULL " "WHERE task_id = ? AND status = ?", ( PersistentTaskStatus.INTERRUPTED.value, "task cancelled; explicit retry required", job.id, PersistentTaskStatus.RUNNING.value, ), ) raise except Exception as exc: attempt += 1 status = ( PersistentTaskStatus.PENDING if attempt <= self.max_retries else PersistentTaskStatus.FAILED ) completed_at = None if status is PersistentTaskStatus.PENDING else time.time() with self.database.transaction() as connection: connection.execute( "UPDATE tasks SET status = ?, error = ?, completed_at = ? " "WHERE task_id = ? AND status = ?", (status.value, f"{type(exc).__name__}: {exc}", completed_at, job.id, PersistentTaskStatus.RUNNING.value), ) if status is PersistentTaskStatus.FAILED: logger.exception("durable task %s failed", job.id) break await asyncio.sleep(0.1 * (2 ** (attempt - 1))) job = self.get_job(job.id) or job continue with self.database.transaction() as connection: connection.execute( "UPDATE tasks SET status = ?, result_json = ?, error = NULL, completed_at = ? " "WHERE task_id = ? AND status = ?", ( PersistentTaskStatus.COMPLETED.value, result_json, time.time(), job.id, PersistentTaskStatus.RUNNING.value, ), ) break latest = self.get_job(job.id) return latest if latest is not None else job @staticmethod def _row_to_record(row: Any) -> PersistentTaskRecord: return PersistentTaskRecord( id=row["task_id"], handler=row["handler"], payload=json.loads(row["payload_json"]), status=PersistentTaskStatus(row["status"]), result=json.loads(row["result_json"]) if row["result_json"] is not None else None, error=row["error"], attempts=row["attempts"], created_at=row["created_at"], started_at=row["started_at"], completed_at=row["completed_at"], ) __all__ = [ "EXECUTION_ATTEMPTED_KEY", "DurableTaskHandler", "MemoryTaskQueue", "PersistentTaskRecord", "PersistentTaskStatus", "SQLiteTaskQueue", "TaskFn", "TaskRecord", "TaskStatus", ]