codex-autorunner 0.1.2__py3-none-any.whl → 1.1.0__py3-none-any.whl

codex_autorunner/core/flows/store.py

@@ -0,0 +1,568 @@
+import json
+import logging
+import sqlite3
+import threading
+from contextlib import contextmanager
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any, Dict, Generator, List, Optional, cast
+
+from ..sqlite_utils import SQLITE_PRAGMAS
+from .models import (
+    FlowArtifact,
+    FlowEvent,
+    FlowEventType,
+    FlowRunRecord,
+    FlowRunStatus,
+)
+
+_logger = logging.getLogger(__name__)
+
+SCHEMA_VERSION = 2
+UNSET = object()
+
+
+def now_iso() -> str:
+    return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
+
+
+class FlowStore:
+    def __init__(self, db_path: Path):
+        self.db_path = db_path
+        self._local: threading.local = threading.local()
+
+    def _get_conn(self) -> sqlite3.Connection:
+        if not hasattr(self._local, "conn"):
+            # Ensure parent directory exists so sqlite can create/open the file.
+            try:
+                self.db_path.parent.mkdir(parents=True, exist_ok=True)
+            except Exception:
+                # Let sqlite raise a clearer error below if directory creation failed.
+                pass
+            self._local.conn = sqlite3.connect(
+                self.db_path, check_same_thread=False, isolation_level=None
+            )
+            self._local.conn.row_factory = sqlite3.Row
+            for pragma in SQLITE_PRAGMAS:
+                self._local.conn.execute(pragma)
+        return cast(sqlite3.Connection, self._local.conn)
+
+    @contextmanager
+    def transaction(self) -> Generator[sqlite3.Connection, None, None]:
+        conn = self._get_conn()
+        try:
+            conn.execute("BEGIN IMMEDIATE")
+            yield conn
+            conn.commit()
+        except Exception:
+            conn.rollback()
+            raise
+
+    def initialize(self) -> None:
+        with self.transaction() as conn:
+            self._create_schema(conn)
+            self._ensure_schema_version(conn)
+
+    def _create_schema(self, conn: sqlite3.Connection) -> None:
+        conn.execute(
+            """
+            CREATE TABLE IF NOT EXISTS schema_info (
+                version INTEGER NOT NULL PRIMARY KEY
+            )
+        """
+        )
+
+        conn.execute(
+            """
+            CREATE TABLE IF NOT EXISTS flow_runs (
+                id TEXT PRIMARY KEY,
+                flow_type TEXT NOT NULL,
+                status TEXT NOT NULL,
+                input_data TEXT NOT NULL,
+                state TEXT NOT NULL,
+                current_step TEXT,
+                stop_requested INTEGER NOT NULL DEFAULT 0,
+                created_at TEXT NOT NULL,
+                started_at TEXT,
+                finished_at TEXT,
+                error_message TEXT,
+                metadata TEXT NOT NULL DEFAULT '{}'
+            )
+        """
+        )
+
+        conn.execute(
+            """
+            CREATE TABLE IF NOT EXISTS flow_events (
+                seq INTEGER PRIMARY KEY AUTOINCREMENT,
+                id TEXT NOT NULL UNIQUE,
+                run_id TEXT NOT NULL,
+                event_type TEXT NOT NULL,
+                timestamp TEXT NOT NULL,
+                data TEXT NOT NULL,
+                step_id TEXT,
+                FOREIGN KEY (run_id) REFERENCES flow_runs(id) ON DELETE CASCADE
+            )
+        """
+        )
+
+        conn.execute(
+            """
+            CREATE TABLE IF NOT EXISTS flow_artifacts (
+                id TEXT PRIMARY KEY,
+                run_id TEXT NOT NULL,
+                kind TEXT NOT NULL,
+                path TEXT NOT NULL,
+                created_at TEXT NOT NULL,
+                metadata TEXT NOT NULL DEFAULT '{}',
+                FOREIGN KEY (run_id) REFERENCES flow_runs(id) ON DELETE CASCADE
+            )
+        """
+        )
+
+        conn.execute(
+            "CREATE INDEX IF NOT EXISTS idx_flow_runs_status ON flow_runs(status)"
+        )
+        conn.execute(
+            "CREATE INDEX IF NOT EXISTS idx_flow_events_run_id ON flow_events(run_id, seq)"
+        )
+        conn.execute(
+            "CREATE INDEX IF NOT EXISTS idx_flow_artifacts_run_id ON flow_artifacts(run_id)"
+        )
+
+    def _ensure_schema_version(self, conn: sqlite3.Connection) -> None:
+        result = conn.execute("SELECT version FROM schema_info").fetchone()
+        if result is None:
+            conn.execute(
+                "INSERT INTO schema_info (version) VALUES (?)", (SCHEMA_VERSION,)
+            )
+        else:
+            current_version = result[0]
+            if current_version < SCHEMA_VERSION:
+                self._migrate_schema(conn, current_version, SCHEMA_VERSION)
+
+    def _migrate_schema(
+        self, conn: sqlite3.Connection, from_version: int, to_version: int
+    ) -> None:
+        _logger.info("Migrating schema from version %d to %d", from_version, to_version)
+        for version in range(from_version, to_version):
+            self._apply_migration(conn, version + 1)
+        conn.execute("UPDATE schema_info SET version = ?", (to_version,))
+
+    def _apply_migration(self, conn: sqlite3.Connection, version: int) -> None:
+        if version == 1:
+            pass
+        elif version == 2:
+            conn.execute("ALTER TABLE flow_events RENAME TO flow_events_old")
+            conn.execute(
+                """
+                CREATE TABLE flow_events (
+                    seq INTEGER PRIMARY KEY AUTOINCREMENT,
+                    id TEXT NOT NULL UNIQUE,
+                    run_id TEXT NOT NULL,
+                    event_type TEXT NOT NULL,
+                    timestamp TEXT NOT NULL,
+                    data TEXT NOT NULL,
+                    step_id TEXT,
+                    FOREIGN KEY (run_id) REFERENCES flow_runs(id) ON DELETE CASCADE
+                )
+                """
+            )
+            conn.execute(
+                """
+                INSERT INTO flow_events (id, run_id, event_type, timestamp, data, step_id)
+                SELECT id, run_id, event_type, timestamp, data, step_id
+                FROM flow_events_old
+                ORDER BY timestamp ASC
+                """
+            )
+            conn.execute("DROP TABLE flow_events_old")
+            conn.execute(
+                "CREATE INDEX IF NOT EXISTS idx_flow_events_run_id ON flow_events(run_id, seq)"
+            )
+
+    def create_flow_run(
+        self,
+        run_id: str,
+        flow_type: str,
+        input_data: Dict[str, Any],
+        metadata: Optional[Dict[str, Any]] = None,
+        state: Optional[Dict[str, Any]] = None,
+        current_step: Optional[str] = None,
+    ) -> FlowRunRecord:
+        now = now_iso()
+        record = FlowRunRecord(
+            id=run_id,
+            flow_type=flow_type,
+            status=FlowRunStatus.PENDING,
+            input_data=input_data,
+            state=state or {},
+            current_step=current_step,
+            stop_requested=False,
+            created_at=now,
+            metadata=metadata or {},
+        )
+
+        with self.transaction() as conn:
+            conn.execute(
+                """
+                INSERT INTO flow_runs (
+                    id, flow_type, status, input_data, state, current_step,
+                    stop_requested, created_at, metadata
+                ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
+                """,
+                (
+                    record.id,
+                    record.flow_type,
+                    record.status.value,
+                    json.dumps(record.input_data),
+                    json.dumps(record.state),
+                    record.current_step,
+                    1 if record.stop_requested else 0,
+                    record.created_at,
+                    json.dumps(record.metadata),
+                ),
+            )
+
+        return record
+
+    def get_flow_run(self, run_id: str) -> Optional[FlowRunRecord]:
+        conn = self._get_conn()
+        row = conn.execute("SELECT * FROM flow_runs WHERE id = ?", (run_id,)).fetchone()
+        if row is None:
+            return None
+        return self._row_to_flow_run(row)
+
+    def update_flow_run_status(
+        self,
+        run_id: str,
+        status: FlowRunStatus,
+        current_step: Any = UNSET,
+        state: Any = UNSET,
+        started_at: Any = UNSET,
+        finished_at: Any = UNSET,
+        error_message: Any = UNSET,
+    ) -> Optional[FlowRunRecord]:
+        updates = ["status = ?"]
+        params: List[Any] = [status.value]
+
+        if current_step is not UNSET:
+            updates.append("current_step = ?")
+            params.append(current_step)
+
+        if state is not UNSET:
+            updates.append("state = ?")
+            params.append(json.dumps(state))
+
+        if started_at is not UNSET:
+            updates.append("started_at = ?")
+            params.append(started_at)
+
+        if finished_at is not UNSET:
+            updates.append("finished_at = ?")
+            params.append(finished_at)
+
+        if error_message is not UNSET:
+            updates.append("error_message = ?")
+            params.append(error_message)
+
+        params.append(run_id)
+
+        with self.transaction() as conn:
+            conn.execute(
+                f"UPDATE flow_runs SET {', '.join(updates)} WHERE id = ?",
+                params,
+            )
+            row = conn.execute(
+                "SELECT * FROM flow_runs WHERE id = ?", (run_id,)
+            ).fetchone()
+            if row is None:
+                return None
+            return self._row_to_flow_run(row)
+
+    def set_stop_requested(
+        self, run_id: str, stop_requested: bool
+    ) -> Optional[FlowRunRecord]:
+        with self.transaction() as conn:
+            conn.execute(
+                "UPDATE flow_runs SET stop_requested = ? WHERE id = ?",
+                (1 if stop_requested else 0, run_id),
+            )
+            row = conn.execute(
+                "SELECT * FROM flow_runs WHERE id = ?", (run_id,)
+            ).fetchone()
+            if row is None:
+                return None
+            return self._row_to_flow_run(row)
+
+    def update_current_step(
+        self, run_id: str, current_step: str
+    ) -> Optional[FlowRunRecord]:
+        with self.transaction() as conn:
+            conn.execute(
+                "UPDATE flow_runs SET current_step = ? WHERE id = ?",
+                (current_step, run_id),
+            )
+            row = conn.execute(
+                "SELECT * FROM flow_runs WHERE id = ?", (run_id,)
+            ).fetchone()
+            if row is None:
+                return None
+            return self._row_to_flow_run(row)
+
+    def list_flow_runs(
+        self, flow_type: Optional[str] = None, status: Optional[FlowRunStatus] = None
+    ) -> List[FlowRunRecord]:
+        conn = self._get_conn()
+        query = "SELECT * FROM flow_runs WHERE 1=1"
+        params: List[Any] = []
+
+        if flow_type is not None:
+            query += " AND flow_type = ?"
+            params.append(flow_type)
+
+        if status is not None:
+            query += " AND status = ?"
+            params.append(status.value)
+
+        query += " ORDER BY created_at DESC"
+
+        rows = conn.execute(query, params).fetchall()
+        return [self._row_to_flow_run(row) for row in rows]
+
+    def create_event(
+        self,
+        event_id: str,
+        run_id: str,
+        event_type: FlowEventType,
+        data: Optional[Dict[str, Any]] = None,
+        step_id: Optional[str] = None,
+    ) -> FlowEvent:
+        timestamp = now_iso()
+
+        with self.transaction() as conn:
+            conn.execute(
+                """
+                INSERT INTO flow_events (id, run_id, event_type, timestamp, data, step_id)
+                VALUES (?, ?, ?, ?, ?, ?)
+                """,
+                (
+                    event_id,
+                    run_id,
+                    event_type.value,
+                    timestamp,
+                    json.dumps(data or {}),
+                    step_id,
+                ),
+            )
+            row = conn.execute(
+                "SELECT * FROM flow_events WHERE id = ?", (event_id,)
+            ).fetchone()
+
+        if row is None:
+            raise RuntimeError("Failed to persist flow event")
+
+        return self._row_to_flow_event(row)
+
+    def get_events(
+        self,
+        run_id: str,
+        after_seq: Optional[int] = None,
+        limit: Optional[int] = None,
+    ) -> List[FlowEvent]:
+        conn = self._get_conn()
+        query = "SELECT * FROM flow_events WHERE run_id = ?"
+        params: List[Any] = [run_id]
+
+        if after_seq is not None:
+            query += " AND seq > ?"
+            params.append(after_seq)
+
+        query += " ORDER BY seq ASC"
+
+        if limit is not None:
+            query += " LIMIT ?"
+            params.append(limit)
+
+        rows = conn.execute(query, params).fetchall()
+        return [self._row_to_flow_event(row) for row in rows]
+
+    def get_last_event_meta(self, run_id: str) -> tuple[Optional[int], Optional[str]]:
+        conn = self._get_conn()
+        row = conn.execute(
+            "SELECT seq, timestamp FROM flow_events WHERE run_id = ? ORDER BY seq DESC LIMIT 1",
+            (run_id,),
+        ).fetchone()
+        if row is None:
+            return None, None
+        return row["seq"], row["timestamp"]
+
+    def get_last_event_seq_by_types(
+        self, run_id: str, event_types: list[FlowEventType]
+    ) -> Optional[int]:
+        if not event_types:
+            return None
+        conn = self._get_conn()
+        placeholders = ", ".join("?" for _ in event_types)
+        params = [run_id, *[t.value for t in event_types]]
+        row = conn.execute(
+            f"""
+            SELECT seq
+            FROM flow_events
+            WHERE run_id = ? AND event_type IN ({placeholders})
+            ORDER BY seq DESC
+            LIMIT 1
+            """,
+            params,
+        ).fetchone()
+        if row is None:
+            return None
+        return cast(int, row["seq"])
+
+    def get_last_event_by_type(
+        self, run_id: str, event_type: FlowEventType
+    ) -> Optional[FlowEvent]:
+        conn = self._get_conn()
+        row = conn.execute(
+            """
+            SELECT *
+            FROM flow_events
+            WHERE run_id = ? AND event_type = ?
+            ORDER BY seq DESC
+            LIMIT 1
+            """,
+            (run_id, event_type.value),
+        ).fetchone()
+        if row is None:
+            return None
+        return self._row_to_flow_event(row)
+
+    def get_latest_step_progress_current_ticket(
+        self, run_id: str, *, after_seq: Optional[int] = None, limit: int = 50
+    ) -> Optional[str]:
+        """Return the most recent step_progress.data.current_ticket for a run.
+
+        This is intentionally lightweight to support UI polling endpoints.
+        """
+        conn = self._get_conn()
+        query = """
+            SELECT seq, data
+            FROM flow_events
+            WHERE run_id = ? AND event_type = ?
+        """
+        params: List[Any] = [run_id, FlowEventType.STEP_PROGRESS.value]
+        if after_seq is not None:
+            query += " AND seq > ?"
+            params.append(after_seq)
+        query += " ORDER BY seq DESC LIMIT ?"
+        params.append(limit)
+        rows = conn.execute(query, params).fetchall()
+        for row in rows:
+            try:
+                data = json.loads(row["data"] or "{}")
+            except Exception:
+                data = {}
+            current_ticket = data.get("current_ticket")
+            if isinstance(current_ticket, str) and current_ticket.strip():
+                return current_ticket.strip()
+        return None
+
+    def create_artifact(
+        self,
+        artifact_id: str,
+        run_id: str,
+        kind: str,
+        path: str,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> FlowArtifact:
+        artifact = FlowArtifact(
+            id=artifact_id,
+            run_id=run_id,
+            kind=kind,
+            path=path,
+            created_at=now_iso(),
+            metadata=metadata or {},
+        )
+
+        with self.transaction() as conn:
+            conn.execute(
+                """
+                INSERT INTO flow_artifacts (id, run_id, kind, path, created_at, metadata)
+                VALUES (?, ?, ?, ?, ?, ?)
+                """,
+                (
+                    artifact.id,
+                    artifact.run_id,
+                    artifact.kind,
+                    artifact.path,
+                    artifact.created_at,
+                    json.dumps(artifact.metadata),
+                ),
+            )
+
+        return artifact
+
+    def get_artifacts(self, run_id: str) -> List[FlowArtifact]:
+        conn = self._get_conn()
+        rows = conn.execute(
+            "SELECT * FROM flow_artifacts WHERE run_id = ? ORDER BY created_at ASC",
+            (run_id,),
+        ).fetchall()
+        return [self._row_to_flow_artifact(row) for row in rows]
+
+    def get_artifact(self, artifact_id: str) -> Optional[FlowArtifact]:
+        conn = self._get_conn()
+        row = conn.execute(
+            "SELECT * FROM flow_artifacts WHERE id = ?", (artifact_id,)
+        ).fetchone()
+        if row is None:
+            return None
+        return self._row_to_flow_artifact(row)
+
+    def delete_flow_run(self, run_id: str) -> bool:
+        """Delete a flow run and its events/artifacts (cascading)."""
+        with self.transaction() as conn:
+            cursor = conn.execute("DELETE FROM flow_runs WHERE id = ?", (run_id,))
+            return cursor.rowcount > 0
+
+    def _row_to_flow_run(self, row: sqlite3.Row) -> FlowRunRecord:
+        return FlowRunRecord(
+            id=row["id"],
+            flow_type=row["flow_type"],
+            status=FlowRunStatus(row["status"]),
+            input_data=json.loads(row["input_data"]),
+            state=json.loads(row["state"]),
+            current_step=row["current_step"],
+            stop_requested=bool(row["stop_requested"]),
+            created_at=row["created_at"],
+            started_at=row["started_at"],
+            finished_at=row["finished_at"],
+            error_message=row["error_message"],
+            metadata=json.loads(row["metadata"]),
+        )
+
+    def _row_to_flow_event(self, row: sqlite3.Row) -> FlowEvent:
+        return FlowEvent(
+            seq=row["seq"],
+            id=row["id"],
+            run_id=row["run_id"],
+            event_type=FlowEventType(row["event_type"]),
+            timestamp=row["timestamp"],
+            data=json.loads(row["data"]),
+            step_id=row["step_id"],
+        )
+
+    def _row_to_flow_artifact(self, row: sqlite3.Row) -> FlowArtifact:
+        return FlowArtifact(
+            id=row["id"],
+            run_id=row["run_id"],
+            kind=row["kind"],
+            path=row["path"],
+            created_at=row["created_at"],
+            metadata=json.loads(row["metadata"]),
+        )
+
+    def close(self) -> None:
+        if hasattr(self._local, "conn"):
+            self._local.conn.close()
+            del self._local.conn

codex_autorunner/core/flows/transition.py

@@ -0,0 +1,138 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Optional
+
+from codex_autorunner.core.flows.models import FlowRunRecord, FlowRunStatus
+from codex_autorunner.core.flows.reasons import ensure_reason_summary
+from codex_autorunner.core.flows.store import now_iso
+
+
+@dataclass(frozen=True)
+class TransitionDecision:
+    """Result of resolving a flow's next status.
+
+    Attributes
+    ----------
+    status: FlowRunStatus
+        The resolved outer status.
+    finished_at: Optional[str]
+        Completion timestamp when the flow reaches a terminal state.
+    state: dict[str, Any]
+        Updated state payload (ticket_engine etc.).
+    note: Optional[str]
+        Reason for the transition (useful for tests/logging).
+    """
+
+    status: FlowRunStatus
+    finished_at: Optional[str]
+    state: dict[str, Any]
+    note: Optional[str] = None
+
+
+def resolve_flow_transition(
+    record: FlowRunRecord, health: Any, now: Optional[str] = None
+) -> TransitionDecision:
+    """Derive the flow status from worker liveness and inner ticket_engine state.
+
+    This is intentionally pure and side-effect free to keep recovery/test logic simple.
+    """
+
+    now = now or now_iso()
+    state: dict[str, Any] = record.state or {}
+    engine_raw = state.get("ticket_engine") if isinstance(state, dict) else {}
+    engine: dict[str, Any] = engine_raw if isinstance(engine_raw, dict) else {}
+    inner_status = engine.get("status")
+    reason_code = engine.get("reason_code")
+
+    # 1) Worker liveness overrides for active flows.
+    if (
+        record.status in (FlowRunStatus.RUNNING, FlowRunStatus.STOPPING)
+        and not health.is_alive
+    ):
+        new_status = (
+            FlowRunStatus.STOPPED
+            if record.status == FlowRunStatus.STOPPING
+            else FlowRunStatus.FAILED
+        )
+        state = ensure_reason_summary(state, status=new_status, default="Worker died")
+        return TransitionDecision(
+            status=new_status, finished_at=now, state=state, note="worker-dead"
+        )
+
+    # 2) Inner engine reconciliation (worker is alive or not required).
+    if record.status == FlowRunStatus.RUNNING:
+        if inner_status == "paused":
+            state = ensure_reason_summary(state, status=FlowRunStatus.PAUSED)
+            return TransitionDecision(
+                status=FlowRunStatus.PAUSED,
+                finished_at=None,
+                state=state,
+                note="engine-paused",
+            )
+
+        if inner_status == "completed":
+            return TransitionDecision(
+                status=FlowRunStatus.COMPLETED,
+                finished_at=now,
+                state=state,
+                note="engine-completed",
+            )
+
+        return TransitionDecision(
+            status=FlowRunStatus.RUNNING, finished_at=None, state=state, note="running"
+        )
+
+    if record.status == FlowRunStatus.PAUSED:
+        if inner_status == "completed":
+            return TransitionDecision(
+                status=FlowRunStatus.COMPLETED,
+                finished_at=now,
+                state=state,
+                note="paused-engine-completed",
+            )
+
+        if (
+            inner_status in (None, "running")
+            and reason_code != "user_pause"
+            and health.is_alive
+        ):
+            # Treat as stale pause; resume and clear pause metadata.
+            engine.pop("reason", None)
+            engine.pop("reason_details", None)
+            engine.pop("reason_code", None)
+            state.pop("reason_summary", None)
+            engine["status"] = "running"
+            state["ticket_engine"] = engine
+            return TransitionDecision(
+                status=FlowRunStatus.RUNNING,
+                finished_at=None,
+                state=state,
+                note="stale-pause-resumed",
+            )
+
+        if not health.is_alive:
+            return TransitionDecision(
+                status=FlowRunStatus.PAUSED,
+                finished_at=None,
+                state=state,
+                note="paused-worker-dead",
+            )
+
+        state = ensure_reason_summary(state, status=FlowRunStatus.PAUSED)
+        return TransitionDecision(
+            status=FlowRunStatus.PAUSED, finished_at=None, state=state, note="paused"
+        )
+
+    # STOPPING/STOPPED/COMPLETED/FAILED: leave unchanged.
+    if record.status.is_terminal() or record.status == FlowRunStatus.STOPPED:
+        return TransitionDecision(
+            status=record.status,
+            finished_at=record.finished_at,
+            state=state,
+            note="terminal",
+        )
+
+    return TransitionDecision(
+        status=record.status, finished_at=None, state=state, note="unchanged"
+    )