"""Minimal Agent-local database bootstrap; no MultiClaw database is used.""" from __future__ import annotations import sqlite3 from pathlib import Path class AgentDatabase: def __init__(self, database_path: Path) -> None: self._database_path = database_path def initialize(self) -> None: self._database_path.parent.mkdir(parents=True, exist_ok=True) with sqlite3.connect(self._database_path) as connection: connection.execute( """ CREATE TABLE IF NOT EXISTS application_meta ( key TEXT PRIMARY KEY, value TEXT NOT NULL ) """ ) connection.execute( """ CREATE TABLE IF NOT EXISTS event_outbox ( event_id INTEGER PRIMARY KEY AUTOINCREMENT, payload TEXT NOT NULL, created_at TEXT NOT NULL ) """ ) connection.execute( "INSERT OR REPLACE INTO application_meta(key, value) VALUES (?, ?)", ("schema_version", "0"), ) connection.execute( """ CREATE TABLE IF NOT EXISTS task_runs ( task_id TEXT PRIMARY KEY, instruction TEXT NOT NULL, status TEXT NOT NULL, received_at TEXT NOT NULL, started_at TEXT, finished_at TEXT, output TEXT, error TEXT, execution_source TEXT ) """ ) try: connection.execute("ALTER TABLE task_runs ADD COLUMN execution_source TEXT") except sqlite3.OperationalError as error: if "duplicate column name" not in str(error).lower(): raise connection.execute( """ CREATE TABLE IF NOT EXISTS task_events ( event_id INTEGER PRIMARY KEY AUTOINCREMENT, task_id TEXT NOT NULL, status TEXT NOT NULL, message TEXT NOT NULL, created_at TEXT NOT NULL ) """ ) def task_status(self, task_id: str) -> str | None: with sqlite3.connect(self._database_path) as connection: row = connection.execute("SELECT status FROM task_runs WHERE task_id = ?", (task_id,)).fetchone() return str(row[0]) if row is not None else None def active_task_ids(self) -> list[str]: """Return only work which can still be stopped by this Agent.""" with sqlite3.connect(self._database_path) as connection: rows = connection.execute( "SELECT task_id FROM task_runs WHERE status IN ('accepted', 'running')" ).fetchall() return [str(row[0]) for row in rows] def create_task(self, task_id: str, instruction: str, received_at: str) -> None: with sqlite3.connect(self._database_path) as connection: connection.execute( "INSERT INTO task_runs(task_id, instruction, status, received_at) VALUES (?, ?, 'accepted', ?)", (task_id, instruction, received_at), ) self._record_event(connection, task_id, "accepted", "任务已接收。", received_at) def cancel_unreceived_task(self, task_id: str, reason: str, occurred_at: str) -> None: """Store a cancellation tombstone when cancel overtakes dispatch.""" with sqlite3.connect(self._database_path) as connection: connection.execute( "INSERT OR IGNORE INTO task_runs(task_id, instruction, status, received_at, finished_at, error) VALUES (?, '', 'cancelled', ?, ?, ?)", (task_id, occurred_at, occurred_at, reason), ) self._record_event(connection, task_id, "cancelled", reason, occurred_at) def get_meta(self, key: str) -> str: with sqlite3.connect(self._database_path) as connection: row = connection.execute("SELECT value FROM application_meta WHERE key = ?", (key,)).fetchone() return str(row[0]) if row is not None else "" def set_meta(self, key: str, value: str) -> None: with sqlite3.connect(self._database_path) as connection: connection.execute( "INSERT INTO application_meta(key, value) VALUES (?, ?) ON CONFLICT(key) DO UPDATE SET value = excluded.value", (key, value), ) def delete_meta_prefix(self, prefix: str) -> None: """Remove local state only for controller-owned entries in one namespace.""" with sqlite3.connect(self._database_path) as connection: connection.execute("DELETE FROM application_meta WHERE key LIKE ?", (prefix + "%",)) def delete_meta(self, key: str) -> None: with sqlite3.connect(self._database_path) as connection: connection.execute("DELETE FROM application_meta WHERE key = ?", (key,)) def meta_entries(self, prefix: str) -> list[tuple[str, str]]: with sqlite3.connect(self._database_path) as connection: rows = connection.execute( "SELECT key, value FROM application_meta WHERE key LIKE ? ORDER BY key", (prefix + "%",) ).fetchall() return [(str(row[0]), str(row[1])) for row in rows] def enqueue_event(self, payload: str, created_at: str) -> None: with sqlite3.connect(self._database_path) as connection: connection.execute("INSERT INTO event_outbox(payload, created_at) VALUES (?, ?)", (payload, created_at)) def pending_events(self) -> list[tuple[int, str]]: with sqlite3.connect(self._database_path) as connection: rows = connection.execute("SELECT event_id, payload FROM event_outbox ORDER BY event_id").fetchall() return [(int(row[0]), str(row[1])) for row in rows] def remove_event(self, event_id: int) -> None: with sqlite3.connect(self._database_path) as connection: connection.execute("DELETE FROM event_outbox WHERE event_id = ?", (event_id,)) def recover_interrupted_tasks(self, occurred_at: str) -> list[str]: with sqlite3.connect(self._database_path) as connection: rows = connection.execute("SELECT task_id FROM task_runs WHERE status IN ('accepted', 'running')").fetchall() task_ids = [str(row[0]) for row in rows] for task_id in task_ids: connection.execute("UPDATE task_runs SET status = 'failed', finished_at = ?, error = ? WHERE task_id = ?", (occurred_at, "Agent 重启前的任务无法安全恢复。", task_id)) self._record_event(connection, task_id, "failed", "Agent 重启后已安全结束未完成任务。", occurred_at) return task_ids def update_task( self, task_id: str, status: str, message: str, occurred_at: str, output: str | None = None, error: str | None = None, execution_source: str | None = None, ) -> None: with sqlite3.connect(self._database_path) as connection: fields = ["status = ?"] values: list[object] = [status] if status == "running": fields.append("started_at = COALESCE(started_at, ?)") values.append(occurred_at) if status in {"completed", "failed", "cancelled", "timed_out"}: fields.append("finished_at = ?") values.append(occurred_at) if output is not None: fields.append("output = ?") values.append(output) if error is not None: fields.append("error = ?") values.append(error) if execution_source is not None: fields.append("execution_source = ?") values.append(execution_source) values.append(task_id) connection.execute(f"UPDATE task_runs SET {', '.join(fields)} WHERE task_id = ?", values) self._record_event(connection, task_id, status, message, occurred_at) @staticmethod def _record_event( connection: sqlite3.Connection, task_id: str, status: str, message: str, created_at: str ) -> None: connection.execute( "INSERT INTO task_events(task_id, status, message, created_at) VALUES (?, ?, ?, ?)", (task_id, status, message, created_at), )