store.py

"""
Persistent store: tasks, agents, events, budgets.
Append-only event log is the source of truth — the dashboard and any
post-hoc analysis read from this. Agents and tasks are summary tables
maintained for fast lookup.
"""
import sqlite3
import json
import time
import uuid
from contextlib import contextmanager
from pathlib import Path
from typing import Any, Optional

DB_PATH = Path(__file__).parent / "logs" / "orchestrator.db"

SCHEMA = """
CREATE TABLE IF NOT EXISTS tasks (
    id TEXT PRIMARY KEY,
    parent_task_id TEXT,
    description TEXT NOT NULL,
    status TEXT NOT NULL,           -- pending | dispatching | running | done | failed | killed
    complexity_score REAL,
    initial_complexity REAL,
    topology TEXT,                  -- JSON: dispatcher's plan (roles, lead, fast_path)
    dispatch_decision TEXT,         -- "fast_path" | "single_agent" | "team"
    token_budget INTEGER NOT NULL,
    tokens_used INTEGER DEFAULT 0,
    agent_slot_budget INTEGER NOT NULL,
    agent_slots_used INTEGER DEFAULT 0,
    max_depth INTEGER NOT NULL,
    created_at REAL NOT NULL,
    updated_at REAL NOT NULL,
    result TEXT
);

CREATE TABLE IF NOT EXISTS agents (
    id TEXT PRIMARY KEY,
    task_id TEXT NOT NULL,
    parent_agent_id TEXT,
    role TEXT NOT NULL,
    depth INTEGER NOT NULL,
    status TEXT NOT NULL,           -- spawning | running | waiting | done | failed | killed
    pid INTEGER,
    workdir TEXT,
    tokens_used INTEGER DEFAULT 0,
    last_heartbeat REAL,
    created_at REAL NOT NULL,
    finished_at REAL,
    FOREIGN KEY (task_id) REFERENCES tasks(id)
);

CREATE TABLE IF NOT EXISTS events (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    ts REAL NOT NULL,
    task_id TEXT,
    agent_id TEXT,
    kind TEXT NOT NULL,             -- spawn_request | spawn_approved | spawn_denied |
                                    -- agent_started | agent_output | agent_done |
                                    -- complexity_check | budget_warning | error | task_done
    payload TEXT NOT NULL           -- JSON
);

CREATE INDEX IF NOT EXISTS idx_events_task ON events(task_id, ts);
CREATE INDEX IF NOT EXISTS idx_events_agent ON events(agent_id, ts);
CREATE INDEX IF NOT EXISTS idx_agents_task ON agents(task_id);
"""


@contextmanager
def conn():
    DB_PATH.parent.mkdir(parents=True, exist_ok=True)
    c = sqlite3.connect(DB_PATH, timeout=10.0)
    c.row_factory = sqlite3.Row
    c.execute("PRAGMA journal_mode=WAL")
    c.execute("PRAGMA synchronous=NORMAL")
    try:
        yield c
        c.commit()
    finally:
        c.close()


def init_db() -> None:
    with conn() as c:
        c.executescript(SCHEMA)


def new_id(prefix: str) -> str:
    return f"{prefix}_{uuid.uuid4().hex[:10]}"


def create_task(
    description: str,
    token_budget: int,
    agent_slot_budget: int,
    max_depth: int,
    parent_task_id: Optional[str] = None,
) -> str:
    tid = new_id("task")
    now = time.time()
    with conn() as c:
        c.execute(
            """INSERT INTO tasks (id, parent_task_id, description, status,
               token_budget, agent_slot_budget, max_depth, created_at, updated_at)
               VALUES (?, ?, ?, 'pending', ?, ?, ?, ?, ?)""",
            (tid, parent_task_id, description, token_budget,
             agent_slot_budget, max_depth, now, now),
        )
    log_event("task_created", task_id=tid,
              payload={"description": description,
                       "token_budget": token_budget,
                       "agent_slot_budget": agent_slot_budget})
    return tid


def update_task(task_id: str, **fields: Any) -> None:
    if not fields:
        return
    fields["updated_at"] = time.time()
    cols = ", ".join(f"{k} = ?" for k in fields)
    with conn() as c:
        c.execute(f"UPDATE tasks SET {cols} WHERE id = ?",
                  (*fields.values(), task_id))


def get_task(task_id: str) -> Optional[dict]:
    with conn() as c:
        row = c.execute("SELECT * FROM tasks WHERE id = ?", (task_id,)).fetchone()
        return dict(row) if row else None


def create_agent(
    task_id: str,
    role: str,
    depth: int,
    parent_agent_id: Optional[str] = None,
    workdir: Optional[str] = None,
) -> str:
    aid = new_id("agent")
    now = time.time()
    with conn() as c:
        c.execute(
            """INSERT INTO agents (id, task_id, parent_agent_id, role, depth,
               status, workdir, last_heartbeat, created_at)
               VALUES (?, ?, ?, ?, ?, 'spawning', ?, ?, ?)""",
            (aid, task_id, parent_agent_id, role, depth, workdir, now, now),
        )
        c.execute(
            "UPDATE tasks SET agent_slots_used = agent_slots_used + 1 WHERE id = ?",
            (task_id,),
        )
    return aid


def update_agent(agent_id: str, **fields: Any) -> None:
    if not fields:
        return
    cols = ", ".join(f"{k} = ?" for k in fields)
    with conn() as c:
        c.execute(f"UPDATE agents SET {cols} WHERE id = ?",
                  (*fields.values(), agent_id))


def get_agent(agent_id: str) -> Optional[dict]:
    with conn() as c:
        row = c.execute("SELECT * FROM agents WHERE id = ?", (agent_id,)).fetchone()
        return dict(row) if row else None


def list_agents_for_task(task_id: str) -> list[dict]:
    with conn() as c:
        rows = c.execute(
            "SELECT * FROM agents WHERE task_id = ? ORDER BY created_at",
            (task_id,),
        ).fetchall()
        return [dict(r) for r in rows]


def log_event(
    kind: str,
    *,
    task_id: Optional[str] = None,
    agent_id: Optional[str] = None,
    payload: Optional[dict] = None,
) -> None:
    with conn() as c:
        c.execute(
            "INSERT INTO events (ts, task_id, agent_id, kind, payload) VALUES (?, ?, ?, ?, ?)",
            (time.time(), task_id, agent_id, kind, json.dumps(payload or {})),
        )


def recent_events(task_id: str, since: float = 0.0, limit: int = 200) -> list[dict]:
    with conn() as c:
        rows = c.execute(
            """SELECT * FROM events WHERE task_id = ? AND ts > ?
               ORDER BY ts DESC LIMIT ?""",
            (task_id, since, limit),
        ).fetchall()
        return [dict(r) for r in rows]