zk-data-agent/src/bash_bg_store.py

"""Persist background bash task records.

Sibling to run_state_store.py: a SQLite table tracking bg shell tasks the
agent spawns via bash(run_in_background=true). The store survives backend
restarts so polling can be resumed; the actual remote process is owned by
nohup on the Jupyter terminal and is independent of this store.
"""

from __future__ import annotations

import sqlite3
import time
from dataclasses import dataclass
from pathlib import Path
from threading import RLock
from typing import Any


ACTIVE_BG_STATUSES = {'running'}
TERMINAL_BG_STATUSES = {'completed', 'failed', 'cancelled'}


@dataclass(frozen=True)
class BgTaskSpec:
    """Payload an agent-side handler hands to the backend on registration."""

    task_id: str
    account_key: str
    session_id: str
    run_id: str
    pid: int
    task_dir: str
    output_path: str
    pid_path: str
    exit_code_path: str
    command: str
    started_at: float
    wait_for_completion: bool


@dataclass(frozen=True)
class BgTaskStatus:
    """Snapshot the agent reads back via bash_status."""

    task_id: str
    status: str
    pid: int
    started_at: float
    finished_at: float | None
    exit_code: int | None
    output_path: str
    output_preview: str
    auto_resumed: bool


class BashBgStore:
    def __init__(self, db_path: Path) -> None:
        self.db_path = db_path
        self._lock = RLock()

    def record_start(self, spec: BgTaskSpec) -> None:
        now = time.time()
        with self._connect() as conn:
            conn.execute(
                """
                insert into bash_bg_tasks (
                    task_id, account_key, session_id, run_id, pid,
                    task_dir, output_path, pid_path, exit_code_path,
                    command, status, started_at, updated_at,
                    wait_for_completion, auto_resumed
                )
                values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 'running', ?, ?, ?, 0)
                on conflict(task_id) do nothing
                """,
                (
                    spec.task_id,
                    spec.account_key,
                    spec.session_id,
                    spec.run_id,
                    spec.pid,
                    spec.task_dir,
                    spec.output_path,
                    spec.pid_path,
                    spec.exit_code_path,
                    spec.command,
                    spec.started_at,
                    now,
                    1 if spec.wait_for_completion else 0,
                ),
            )

    def mark_completed(
        self,
        task_id: str,
        *,
        exit_code: int,
        finished_at: float | None = None,
    ) -> None:
        finished = finished_at if finished_at is not None else time.time()
        status = 'completed' if exit_code == 0 else 'failed'
        with self._connect() as conn:
            conn.execute(
                """
                update bash_bg_tasks
                set status = ?, exit_code = ?, finished_at = ?, updated_at = ?
                where task_id = ? and status = 'running'
                """,
                (status, exit_code, finished, finished, task_id),
            )

    def mark_killed(self, task_id: str) -> None:
        now = time.time()
        with self._connect() as conn:
            conn.execute(
                """
                update bash_bg_tasks
                set status = 'cancelled', finished_at = ?, updated_at = ?
                where task_id = ? and status = 'running'
                """,
                (now, now, task_id),
            )

    def mark_auto_resumed(self, task_id: str) -> bool:
        """Set auto_resumed=1 atomically; return True if we won the race."""
        with self._connect() as conn:
            cursor = conn.execute(
                """
                update bash_bg_tasks
                set auto_resumed = 1, updated_at = ?
                where task_id = ? and auto_resumed = 0
                """,
                (time.time(), task_id),
            )
            return cursor.rowcount > 0

    def list_active(self) -> list[dict[str, Any]]:
        with self._connect() as conn:
            rows = conn.execute(
                "select * from bash_bg_tasks where status = 'running'"
            ).fetchall()
            return [dict(row) for row in rows]

    def list_for_session(
        self,
        account_key: str,
        session_id: str,
    ) -> list[dict[str, Any]]:
        with self._connect() as conn:
            rows = conn.execute(
                """
                select * from bash_bg_tasks
                where account_key = ? and session_id = ?
                order by started_at desc
                """,
                (account_key, session_id),
            ).fetchall()
            return [dict(row) for row in rows]

    def list_active_for_session(
        self,
        account_key: str,
        session_id: str,
    ) -> list[dict[str, Any]]:
        with self._connect() as conn:
            rows = conn.execute(
                """
                select * from bash_bg_tasks
                where account_key = ?
                  and session_id = ?
                  and status = 'running'
                order by started_at desc
                """,
                (account_key, session_id),
            ).fetchall()
            return [dict(row) for row in rows]

    def get(self, task_id: str) -> dict[str, Any] | None:
        with self._connect() as conn:
            row = conn.execute(
                "select * from bash_bg_tasks where task_id = ?",
                (task_id,),
            ).fetchone()
            return dict(row) if row is not None else None

    def _connect(self) -> sqlite3.Connection:
        with self._lock:
            self.db_path.parent.mkdir(parents=True, exist_ok=True)
            conn = sqlite3.connect(self.db_path, timeout=10)
            conn.row_factory = sqlite3.Row
            self._init_schema(conn)
            return conn

    def _init_schema(self, conn: sqlite3.Connection) -> None:
        conn.execute('pragma journal_mode=wal')
        conn.execute(
            """
            create table if not exists bash_bg_tasks (
                task_id text primary key,
                account_key text not null,
                session_id text not null,
                run_id text not null default '',
                pid integer not null,
                task_dir text not null,
                output_path text not null,
                pid_path text not null,
                exit_code_path text not null,
                command text not null,
                status text not null,
                started_at real not null,
                finished_at real,
                exit_code integer,
                updated_at real not null,
                wait_for_completion integer not null default 1,
                auto_resumed integer not null default 0
            )
            """
        )
        conn.execute(
            """
            create index if not exists idx_bash_bg_session
            on bash_bg_tasks(account_key, session_id, started_at)
            """
        )
        conn.execute(
            """
            create index if not exists idx_bash_bg_active
            on bash_bg_tasks(status)
            """
        )