hermes-mac13
93758f3720
- local_rig/ 完整包（7 源文件 + ui.html + CHANGELOG） - tests/test_local_rig.py + test_rig_real_agent.py（纯 local_rig 测试） - scripts/15-criterion4-live.sh（判据④活环境门） - README.md（定位/快速开始/与 tutti-selfhost 的关系） - requirements.txt（唯一外部依赖 pyyaml） - .gitignore 来源: figmar/tutti-selfhost @ 3b23853 (2026-10-06) 原 PR #43 已合入 tutti-selfhost main
248 lines
9.1 KiB
Python
248 lines
9.1 KiB
Python
"""SQLite 持久化任务队列（重启不丢，借鉴 OpenRig durable queue）。
|
||
|
||
设计：
|
||
- 单文件 SQLite（WAL 模式）
|
||
- 任务状态机: pending → running → done | failed | timeout
|
||
- 幂等：同一 task_id 重复入队只保留一条
|
||
- 原子领取：BEGIN IMMEDIATE + UPDATE ... WHERE status='pending'
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import sqlite3
|
||
import time
|
||
import uuid
|
||
from dataclasses import dataclass, field
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
# 幂等窗口：契约 R-1 固化为 **24h**（同 idem_key 24h 内只执行一次）
|
||
IDEM_TTL_S = 24 * 3600
|
||
|
||
SCHEMA = """
|
||
CREATE TABLE IF NOT EXISTS tasks (
|
||
id TEXT PRIMARY KEY,
|
||
assignee TEXT NOT NULL,
|
||
prompt TEXT NOT NULL,
|
||
context TEXT DEFAULT '{}',
|
||
status TEXT NOT NULL DEFAULT 'pending',
|
||
result TEXT,
|
||
error TEXT,
|
||
created_at REAL NOT NULL,
|
||
claimed_at REAL,
|
||
finished_at REAL,
|
||
deadline_s REAL DEFAULT 600,
|
||
idem_key TEXT UNIQUE
|
||
);
|
||
CREATE INDEX IF NOT EXISTS idx_tasks_status ON tasks(status);
|
||
CREATE INDEX IF NOT EXISTS idx_tasks_assignee ON tasks(assignee, status);
|
||
|
||
CREATE TABLE IF NOT EXISTS rig_meta (
|
||
key TEXT PRIMARY KEY,
|
||
value TEXT
|
||
);
|
||
"""
|
||
|
||
|
||
@dataclass
|
||
class Task:
|
||
id: str
|
||
assignee: str
|
||
prompt: str
|
||
context: dict[str, Any] = field(default_factory=dict)
|
||
status: str = "pending"
|
||
result: str | None = None
|
||
error: str | None = None
|
||
created_at: float = field(default_factory=time.time)
|
||
claimed_at: float | None = None
|
||
finished_at: float | None = None
|
||
deadline_s: float = 600.0
|
||
idem_key: str | None = None
|
||
|
||
def to_row(self) -> tuple:
|
||
return (
|
||
self.id, self.assignee, self.prompt, json.dumps(self.context, ensure_ascii=False),
|
||
self.status, self.result, self.error,
|
||
self.created_at, self.claimed_at, self.finished_at, self.deadline_s,
|
||
)
|
||
|
||
@classmethod
|
||
def from_row(cls, row: tuple) -> Task:
|
||
return cls(
|
||
id=row[0], assignee=row[1], prompt=row[2],
|
||
context=json.loads(row[3] or "{}"), status=row[4],
|
||
result=row[5], error=row[6], created_at=row[7],
|
||
claimed_at=row[8], finished_at=row[9], deadline_s=row[10],
|
||
idem_key=(row[11] if len(row) > 11 else None), # G1：SCHEMA 第 12 列；老库迁移后同样存在
|
||
)
|
||
|
||
|
||
def ensure_idem_column(conn) -> None:
|
||
"""G1：老库补列。SQLite `ADD COLUMN` 不能带 UNIQUE ⇒ 用**唯一索引**实现同等约束；
|
||
历史行 idem_key 为 NULL（NULL 不违反唯一索引），不影响既有数据。"""
|
||
cols = [r[1] for r in conn.execute("PRAGMA table_info(tasks)").fetchall()]
|
||
if "idem_key" not in cols:
|
||
conn.execute("ALTER TABLE tasks ADD COLUMN idem_key TEXT")
|
||
conn.execute("CREATE UNIQUE INDEX IF NOT EXISTS idx_tasks_idem ON tasks(idem_key)")
|
||
conn.commit()
|
||
|
||
|
||
class TaskQueue:
|
||
"""SQLite 持久化任务队列。线程安全（sqlite3 check_same_thread=False + 锁）。"""
|
||
|
||
def __init__(self, db_path: str | Path = ":memory:"):
|
||
self.db_path = str(db_path)
|
||
self.conn = sqlite3.connect(self.db_path, check_same_thread=False)
|
||
self.conn.execute("PRAGMA journal_mode=WAL")
|
||
self.conn.execute("PRAGMA busy_timeout=5000")
|
||
self.conn.executescript(SCHEMA)
|
||
ensure_idem_column(self.conn) # G1：老库补列（新库已在 SCHEMA 里）
|
||
self.conn.commit()
|
||
|
||
def close(self):
|
||
self.conn.close()
|
||
|
||
# ── 入队 ──
|
||
|
||
def enqueue(self, assignee: str, prompt: str,
|
||
context: dict | None = None,
|
||
task_id: str | None = None,
|
||
deadline_s: float = 600.0,
|
||
idem_key: str | None = None) -> Task:
|
||
"""入队。**G1 幂等**：给了 `idem_key` 且 24h 内已有同键任务 ⇒ 直接返回那条，**不重复执行**
|
||
（对齐契约 R-1 / PRD A4「重发三次只执行一次」）。"""
|
||
if idem_key:
|
||
row = self.conn.execute(
|
||
"SELECT id FROM tasks WHERE idem_key=? AND created_at > ?",
|
||
(idem_key, time.time() - IDEM_TTL_S)).fetchone()
|
||
if row:
|
||
got = self.get(row[0])
|
||
if got is not None:
|
||
return got
|
||
tid = task_id or f"task-{uuid.uuid4().hex[:12]}"
|
||
task = Task(id=tid, assignee=assignee, prompt=prompt,
|
||
context=context or {}, deadline_s=deadline_s, idem_key=idem_key)
|
||
# 幂等：已存在则跳过
|
||
cur = self.conn.execute("SELECT id FROM tasks WHERE id=?", (tid,))
|
||
if cur.fetchone():
|
||
return self.get(tid) # type: ignore
|
||
self.conn.execute(
|
||
"INSERT INTO tasks (id,assignee,prompt,context,status,created_at,deadline_s,idem_key) "
|
||
"VALUES (?,?,?,?,?,?,?,?)",
|
||
(tid, assignee, prompt, json.dumps(context or {}, ensure_ascii=False),
|
||
"pending", task.created_at, deadline_s, idem_key),
|
||
)
|
||
self.conn.commit()
|
||
return task
|
||
|
||
# ── 领取（原子） ──
|
||
|
||
def claim(self, assignee: str | None = None, task_id: str | None = None) -> Task | None:
|
||
"""原子领取一个 pending 任务。assignee=None 时不限制 assignee。"""
|
||
self.conn.execute("BEGIN IMMEDIATE")
|
||
try:
|
||
if task_id:
|
||
cur = self.conn.execute(
|
||
"SELECT * FROM tasks WHERE id=? AND status='pending'", (task_id,))
|
||
else:
|
||
where = "status='pending'"
|
||
params: list[Any] = []
|
||
if assignee:
|
||
where += " AND assignee=?"
|
||
params.append(assignee)
|
||
cur = self.conn.execute(
|
||
f"SELECT * FROM tasks WHERE {where} ORDER BY created_at LIMIT 1", params)
|
||
row = cur.fetchone()
|
||
if not row:
|
||
self.conn.execute("COMMIT")
|
||
return None
|
||
task = Task.from_row(row)
|
||
self.conn.execute(
|
||
"UPDATE tasks SET status='running', claimed_at=? WHERE id=?",
|
||
(time.time(), task.id))
|
||
self.conn.execute("COMMIT")
|
||
task.status = "running"
|
||
task.claimed_at = time.time()
|
||
return task
|
||
except Exception:
|
||
self.conn.execute("ROLLBACK")
|
||
raise
|
||
|
||
# ── 完成 ──
|
||
|
||
def complete(self, task_id: str, result: str | None = None):
|
||
self.conn.execute(
|
||
"UPDATE tasks SET status='done', result=?, finished_at=? WHERE id=?",
|
||
(result, time.time(), task_id))
|
||
self.conn.commit()
|
||
|
||
def fail(self, task_id: str, error: str):
|
||
self.conn.execute(
|
||
"UPDATE tasks SET status='failed', error=?, finished_at=? WHERE id=?",
|
||
(error, time.time(), task_id))
|
||
self.conn.commit()
|
||
|
||
def timeout(self, task_id: str):
|
||
self.conn.execute(
|
||
"UPDATE tasks SET status='timeout', finished_at=? WHERE id=?",
|
||
(time.time(), task_id))
|
||
self.conn.commit()
|
||
|
||
# ── 查询 ──
|
||
|
||
def get(self, task_id: str) -> Task | None:
|
||
cur = self.conn.execute("SELECT * FROM tasks WHERE id=?", (task_id,))
|
||
row = cur.fetchone()
|
||
return Task.from_row(row) if row else None
|
||
|
||
def list(self, status: str | None = None, assignee: str | None = None,
|
||
limit: int = 50) -> list[Task]:
|
||
where, params = ["1=1"], []
|
||
if status:
|
||
where.append("status=?")
|
||
params.append(status)
|
||
if assignee:
|
||
where.append("assignee=?")
|
||
params.append(assignee)
|
||
cur = self.conn.execute(
|
||
f"SELECT * FROM tasks WHERE {' AND '.join(where)} ORDER BY created_at DESC LIMIT ?",
|
||
params + [limit])
|
||
return [Task.from_row(r) for r in cur.fetchall()]
|
||
|
||
def count(self, status: str | None = None) -> int:
|
||
if status:
|
||
cur = self.conn.execute("SELECT COUNT(*) FROM tasks WHERE status=?", (status,))
|
||
else:
|
||
cur = self.conn.execute("SELECT COUNT(*) FROM tasks")
|
||
return cur.fetchone()[0]
|
||
|
||
def pending_for(self, assignee: str) -> int:
|
||
cur = self.conn.execute(
|
||
"SELECT COUNT(*) FROM tasks WHERE assignee=? AND status='pending'", (assignee,))
|
||
return cur.fetchone()[0]
|
||
|
||
# ── 超时扫描 ──
|
||
|
||
def sweep_timeouts(self) -> list[str]:
|
||
"""把超时的 running 任务标为 timeout。返回被标记的 task_id 列表。"""
|
||
now = time.time()
|
||
cur = self.conn.execute(
|
||
"SELECT id FROM tasks WHERE status='running' AND claimed_at IS NOT NULL "
|
||
"AND (claimed_at + deadline_s) < ?", (now,))
|
||
ids = [r[0] for r in cur.fetchall()]
|
||
for tid in ids:
|
||
self.timeout(tid)
|
||
return ids
|
||
|
||
# ── meta ──
|
||
|
||
def set_meta(self, key: str, value: str):
|
||
self.conn.execute(
|
||
"INSERT OR REPLACE INTO rig_meta (key,value) VALUES (?,?)", (key, value))
|
||
self.conn.commit()
|
||
|
||
def get_meta(self, key: str) -> str | None:
|
||
cur = self.conn.execute("SELECT value FROM rig_meta WHERE key=?", (key,))
|
||
row = cur.fetchone()
|
||
return row[0] if row else None
|