tasks.py

"""
Persistent task memory (Phase 2).

A lightweight SQLite store that keeps the full lifecycle, plan and event
history of every agent run. The DB lives in $TASK_DB_PATH (default
``/home/user/app/data/tasks.db``); the directory is created on demand so the
container starts cleanly even on a fresh volume.

This module is INTENTIONALLY additive — none of the Phase-1 endpoints
import it, so existing behaviour cannot regress if SQLite is unavailable.

States (must match TaskState below):

    queued → planning → thinking → executing → retrying → completed | failed
                                          ↘────────────↗

Concurrency model: a single ``sqlite3.connect`` per call, ``WAL`` journal mode,
all writes guarded by a process-wide ``asyncio.Lock``.  This is plenty for the
single-worker uvicorn deployment we run on HF Spaces.
"""

from __future__ import annotations

import asyncio
import json
import logging
import os
import sqlite3
import time
import uuid
from contextlib import contextmanager
from dataclasses import dataclass, field
from typing import Any, Dict, Iterable, List, Optional

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Lifecycle states
# ---------------------------------------------------------------------------

class TaskState:
    QUEUED = "queued"
    PLANNING = "planning"
    THINKING = "thinking"
    EXECUTING = "executing"
    RETRYING = "retrying"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"

    ALL = {QUEUED, PLANNING, THINKING, EXECUTING, RETRYING, COMPLETED, FAILED, CANCELLED}
    TERMINAL = {COMPLETED, FAILED, CANCELLED}


# ---------------------------------------------------------------------------
# DB location & schema
# ---------------------------------------------------------------------------

def _db_path() -> str:
    raw = os.environ.get("TASK_DB_PATH", "/home/user/app/data/tasks.db")
    os.makedirs(os.path.dirname(raw) or ".", exist_ok=True)
    return raw


_SCHEMA = """
CREATE TABLE IF NOT EXISTS tasks (
    id           TEXT PRIMARY KEY,
    created_at   REAL NOT NULL,
    updated_at   REAL NOT NULL,
    state        TEXT NOT NULL,
    user_message TEXT NOT NULL,
    metadata     TEXT NOT NULL DEFAULT '{}',
    sandbox_id   TEXT,
    final_reply  TEXT,
    error        TEXT
);
CREATE INDEX IF NOT EXISTS idx_tasks_created_at ON tasks(created_at DESC);
CREATE INDEX IF NOT EXISTS idx_tasks_state      ON tasks(state);

CREATE TABLE IF NOT EXISTS steps (
    id          INTEGER PRIMARY KEY AUTOINCREMENT,
    task_id     TEXT NOT NULL REFERENCES tasks(id) ON DELETE CASCADE,
    idx         INTEGER NOT NULL,
    title       TEXT NOT NULL,
    description TEXT NOT NULL DEFAULT '',
    state       TEXT NOT NULL DEFAULT 'queued',
    attempts    INTEGER NOT NULL DEFAULT 0,
    started_at  REAL,
    finished_at REAL,
    result      TEXT,
    error       TEXT
);
CREATE INDEX IF NOT EXISTS idx_steps_task ON steps(task_id, idx);

CREATE TABLE IF NOT EXISTS events (
    id        INTEGER PRIMARY KEY AUTOINCREMENT,
    task_id   TEXT NOT NULL REFERENCES tasks(id) ON DELETE CASCADE,
    step_idx  INTEGER,
    ts        REAL NOT NULL,
    kind      TEXT NOT NULL,
    payload   TEXT NOT NULL DEFAULT '{}'
);
CREATE INDEX IF NOT EXISTS idx_events_task ON events(task_id, id);
"""


_init_lock = asyncio.Lock()
_initialised = False
_write_lock = asyncio.Lock()


@contextmanager
def _conn():
    path = _db_path()
    c = sqlite3.connect(path, timeout=10, isolation_level=None)
    c.row_factory = sqlite3.Row
    try:
        c.execute("PRAGMA journal_mode=WAL;")
        c.execute("PRAGMA synchronous=NORMAL;")
        c.execute("PRAGMA foreign_keys=ON;")
        yield c
    finally:
        c.close()


async def init() -> None:
    """Create tables if necessary. Safe to call many times."""
    global _initialised
    if _initialised:
        return
    async with _init_lock:
        if _initialised:
            return
        def _do():
            with _conn() as c:
                c.executescript(_SCHEMA)
        await asyncio.to_thread(_do)
        _initialised = True
        logger.info("task DB ready at %s", _db_path())


# ---------------------------------------------------------------------------
# Dataclasses
# ---------------------------------------------------------------------------

@dataclass
class Step:
    idx: int
    title: str
    description: str = ""
    state: str = TaskState.QUEUED
    attempts: int = 0
    started_at: Optional[float] = None
    finished_at: Optional[float] = None
    result: Optional[str] = None
    error: Optional[str] = None

    def to_dict(self) -> Dict[str, Any]:
        return {
            "idx": self.idx,
            "title": self.title,
            "description": self.description,
            "state": self.state,
            "attempts": self.attempts,
            "started_at": self.started_at,
            "finished_at": self.finished_at,
            "result": self.result,
            "error": self.error,
        }


@dataclass
class Task:
    id: str
    created_at: float
    updated_at: float
    state: str
    user_message: str
    metadata: Dict[str, Any] = field(default_factory=dict)
    sandbox_id: Optional[str] = None
    final_reply: Optional[str] = None
    error: Optional[str] = None
    steps: List[Step] = field(default_factory=list)

    def to_dict(self, include_steps: bool = True) -> Dict[str, Any]:
        out: Dict[str, Any] = {
            "id": self.id,
            "created_at": self.created_at,
            "updated_at": self.updated_at,
            "state": self.state,
            "user_message": self.user_message,
            "metadata": self.metadata,
            "sandbox_id": self.sandbox_id,
            "final_reply": self.final_reply,
            "error": self.error,
        }
        if include_steps:
            out["steps"] = [s.to_dict() for s in self.steps]
        return out


# ---------------------------------------------------------------------------
# CRUD helpers
# ---------------------------------------------------------------------------

async def create_task(user_message: str, metadata: Optional[Dict[str, Any]] = None) -> Task:
    await init()
    now = time.time()
    task = Task(
        id=uuid.uuid4().hex,
        created_at=now,
        updated_at=now,
        state=TaskState.QUEUED,
        user_message=user_message,
        metadata=metadata or {},
    )
    def _do():
        with _conn() as c:
            c.execute(
                "INSERT INTO tasks (id, created_at, updated_at, state, user_message, metadata) "
                "VALUES (?,?,?,?,?,?)",
                (task.id, task.created_at, task.updated_at, task.state,
                 task.user_message, json.dumps(task.metadata)),
            )
    async with _write_lock:
        await asyncio.to_thread(_do)
    return task


async def update_state(task_id: str, state: str, *, error: Optional[str] = None,
                       sandbox_id: Optional[str] = None,
                       final_reply: Optional[str] = None) -> None:
    await init()
    if state not in TaskState.ALL:
        raise ValueError(f"invalid state: {state}")
    now = time.time()
    def _do():
        fields = ["state = ?", "updated_at = ?"]
        params: List[Any] = [state, now]
        if error is not None:
            fields.append("error = ?")
            params.append(error)
        if sandbox_id is not None:
            fields.append("sandbox_id = ?")
            params.append(sandbox_id)
        if final_reply is not None:
            fields.append("final_reply = ?")
            params.append(final_reply)
        params.append(task_id)
        with _conn() as c:
            c.execute(f"UPDATE tasks SET {', '.join(fields)} WHERE id = ?", params)
    async with _write_lock:
        await asyncio.to_thread(_do)


async def set_steps(task_id: str, steps: Iterable[Dict[str, str]]) -> List[Step]:
    """Replace the plan for a task. Each input dict needs ``title`` and
    optional ``description``."""
    await init()
    rows: List[Step] = []
    for i, raw in enumerate(steps):
        rows.append(Step(idx=i, title=str(raw.get("title", f"Step {i+1}")),
                         description=str(raw.get("description", ""))))
    def _do():
        with _conn() as c:
            c.execute("DELETE FROM steps WHERE task_id = ?", (task_id,))
            c.executemany(
                "INSERT INTO steps (task_id, idx, title, description, state) VALUES (?,?,?,?,?)",
                [(task_id, s.idx, s.title, s.description, s.state) for s in rows],
            )
            c.execute("UPDATE tasks SET updated_at = ? WHERE id = ?", (time.time(), task_id))
    async with _write_lock:
        await asyncio.to_thread(_do)
    return rows


async def update_step(task_id: str, idx: int, *, state: Optional[str] = None,
                      attempts_delta: int = 0, result: Optional[str] = None,
                      error: Optional[str] = None) -> None:
    await init()
    now = time.time()
    def _do():
        fields: List[str] = []
        params: List[Any] = []
        if state is not None:
            fields.append("state = ?")
            params.append(state)
            if state == TaskState.EXECUTING:
                fields.append("started_at = COALESCE(started_at, ?)")
                params.append(now)
            elif state in TaskState.TERMINAL or state in (TaskState.COMPLETED, TaskState.FAILED):
                fields.append("finished_at = ?")
                params.append(now)
        if attempts_delta:
            fields.append("attempts = attempts + ?")
            params.append(attempts_delta)
        if result is not None:
            fields.append("result = ?")
            params.append(result[:4000])
        if error is not None:
            fields.append("error = ?")
            params.append(error[:4000])
        if not fields:
            return
        params.extend([task_id, idx])
        with _conn() as c:
            c.execute(
                f"UPDATE steps SET {', '.join(fields)} WHERE task_id = ? AND idx = ?",
                params,
            )
            c.execute("UPDATE tasks SET updated_at = ? WHERE id = ?", (now, task_id))
    async with _write_lock:
        await asyncio.to_thread(_do)


async def append_event(task_id: str, kind: str, payload: Any, step_idx: Optional[int] = None) -> None:
    await init()
    def _do():
        with _conn() as c:
            c.execute(
                "INSERT INTO events (task_id, step_idx, ts, kind, payload) VALUES (?,?,?,?,?)",
                (task_id, step_idx, time.time(), kind,
                 json.dumps(payload, ensure_ascii=False, default=str)[:8000]),
            )
    async with _write_lock:
        try:
            await asyncio.to_thread(_do)
        except Exception as e:
            # Logging only — events are diagnostics; never break the stream.
            logger.warning("append_event failed: %s", e)


# ---------------------------------------------------------------------------
# Read helpers
# ---------------------------------------------------------------------------

async def get_task(task_id: str) -> Optional[Task]:
    await init()
    def _do():
        with _conn() as c:
            row = c.execute("SELECT * FROM tasks WHERE id = ?", (task_id,)).fetchone()
            if not row:
                return None
            steps_rows = c.execute(
                "SELECT * FROM steps WHERE task_id = ? ORDER BY idx", (task_id,)
            ).fetchall()
            return row, steps_rows
    result = await asyncio.to_thread(_do)
    if not result:
        return None
    row, steps_rows = result
    return Task(
        id=row["id"],
        created_at=row["created_at"],
        updated_at=row["updated_at"],
        state=row["state"],
        user_message=row["user_message"],
        metadata=json.loads(row["metadata"] or "{}"),
        sandbox_id=row["sandbox_id"],
        final_reply=row["final_reply"],
        error=row["error"],
        steps=[
            Step(
                idx=s["idx"], title=s["title"], description=s["description"] or "",
                state=s["state"], attempts=s["attempts"],
                started_at=s["started_at"], finished_at=s["finished_at"],
                result=s["result"], error=s["error"],
            )
            for s in steps_rows
        ],
    )


async def list_tasks(limit: int = 50, state: Optional[str] = None) -> List[Dict[str, Any]]:
    await init()
    def _do():
        with _conn() as c:
            if state:
                rows = c.execute(
                    "SELECT id, created_at, updated_at, state, user_message, sandbox_id "
                    "FROM tasks WHERE state = ? ORDER BY created_at DESC LIMIT ?",
                    (state, limit),
                ).fetchall()
            else:
                rows = c.execute(
                    "SELECT id, created_at, updated_at, state, user_message, sandbox_id "
                    "FROM tasks ORDER BY created_at DESC LIMIT ?",
                    (limit,),
                ).fetchall()
            return [dict(r) for r in rows]
    return await asyncio.to_thread(_do)


async def get_events(task_id: str, after_id: int = 0, limit: int = 500) -> List[Dict[str, Any]]:
    await init()
    def _do():
        with _conn() as c:
            rows = c.execute(
                "SELECT id, step_idx, ts, kind, payload FROM events "
                "WHERE task_id = ? AND id > ? ORDER BY id LIMIT ?",
                (task_id, after_id, limit),
            ).fetchall()
            out: List[Dict[str, Any]] = []
            for r in rows:
                try:
                    payload = json.loads(r["payload"])
                except Exception:
                    payload = {"raw": r["payload"]}
                out.append({
                    "id": r["id"],
                    "step_idx": r["step_idx"],
                    "ts": r["ts"],
                    "kind": r["kind"],
                    "payload": payload,
                })
            return out
    return await asyncio.to_thread(_do)


async def delete_task(task_id: str) -> bool:
    await init()
    def _do():
        with _conn() as c:
            cur = c.execute("DELETE FROM tasks WHERE id = ?", (task_id,))
            return cur.rowcount > 0
    async with _write_lock:
        return await asyncio.to_thread(_do)