189 lines
8.4 KiB
Python
189 lines
8.4 KiB
Python
"""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),
|
|
)
|