299 lines
18 KiB
Python
299 lines
18 KiB
Python
"""SQLite authority for ``cad.single-stage.v1`` tasks.
|
|
|
|
There is no reader or migration for the deleted feature-DAG protocol. Opening
|
|
a legacy task database drops its task data before creating this schema.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from contextlib import contextmanager
|
|
import json
|
|
from pathlib import Path
|
|
import sqlite3
|
|
from threading import RLock
|
|
from typing import Any, Iterator
|
|
|
|
from app.cad_agent.domain.errors import ErrorCode
|
|
from app.cad_agent.domain.state import TaskPhase, TaskState
|
|
from app.cad_agent.ports import InvocationRecord
|
|
|
|
|
|
class SqliteTaskRepository:
|
|
PROTOCOL_VERSION = "cad.single-stage.v1"
|
|
|
|
def __init__(self, database_path: Path) -> None:
|
|
self.database_path = database_path
|
|
self.database_path.parent.mkdir(parents=True, exist_ok=True)
|
|
self._lock = RLock()
|
|
self.protocol_reset = False
|
|
self._initialize()
|
|
|
|
@contextmanager
|
|
def _connection(self) -> Iterator[sqlite3.Connection]:
|
|
connection = sqlite3.connect(self.database_path, timeout=20, isolation_level=None)
|
|
try:
|
|
connection.row_factory = sqlite3.Row
|
|
connection.execute("PRAGMA foreign_keys = ON")
|
|
connection.execute("PRAGMA journal_mode = WAL")
|
|
yield connection
|
|
finally:
|
|
connection.close()
|
|
|
|
def _initialize(self) -> None:
|
|
with self._lock, self._connection() as connection:
|
|
existing = connection.execute("SELECT sql FROM sqlite_master WHERE type='table' AND name='tasks'").fetchone()
|
|
if existing is not None and self.PROTOCOL_VERSION not in str(existing[0] or ""):
|
|
self.protocol_reset = True
|
|
# This database is dedicated to CAD task state. An old
|
|
# protocol has no safe reader or migration, so remove every
|
|
# persisted task table and let SQLite remove their indexes.
|
|
tables = connection.execute(
|
|
"SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%'"
|
|
).fetchall()
|
|
for row in tables:
|
|
name = str(row[0]).replace('"', '""')
|
|
connection.execute(f'DROP TABLE IF EXISTS "{name}"')
|
|
connection.executescript(f"""
|
|
CREATE TABLE IF NOT EXISTS tasks (
|
|
task_id TEXT PRIMARY KEY,
|
|
protocol_version TEXT NOT NULL CHECK(protocol_version = '{self.PROTOCOL_VERSION}'),
|
|
request TEXT NOT NULL, phase TEXT NOT NULL, state_version INTEGER NOT NULL,
|
|
active_revision TEXT NOT NULL DEFAULT '', repair_count INTEGER NOT NULL DEFAULT 0,
|
|
last_error TEXT, retry_from_phase TEXT NOT NULL DEFAULT '',
|
|
requirements_path TEXT NOT NULL DEFAULT '', authoring_path TEXT NOT NULL DEFAULT '',
|
|
runtime_cdsl_path TEXT NOT NULL DEFAULT '', compile_audit_path TEXT NOT NULL DEFAULT '',
|
|
diagnostics_path TEXT NOT NULL DEFAULT '', completion_path TEXT NOT NULL DEFAULT '',
|
|
clarification_path TEXT NOT NULL DEFAULT '',
|
|
created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
|
|
);
|
|
CREATE TABLE IF NOT EXISTS invocations (
|
|
invocation_id TEXT PRIMARY KEY, task_id TEXT NOT NULL REFERENCES tasks(task_id),
|
|
idempotency_key TEXT NOT NULL, status TEXT NOT NULL, result_json TEXT,
|
|
created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
|
UNIQUE(task_id,idempotency_key)
|
|
);
|
|
CREATE TABLE IF NOT EXISTS ledger (
|
|
sequence INTEGER PRIMARY KEY AUTOINCREMENT, task_id TEXT NOT NULL REFERENCES tasks(task_id),
|
|
event_json TEXT NOT NULL, created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
|
|
);
|
|
CREATE TABLE IF NOT EXISTS outbox (
|
|
event_id INTEGER PRIMARY KEY AUTOINCREMENT, task_id TEXT NOT NULL REFERENCES tasks(task_id),
|
|
event_json TEXT NOT NULL, published_at TEXT
|
|
);
|
|
CREATE TABLE IF NOT EXISTS usage_records (
|
|
usage_id INTEGER PRIMARY KEY AUTOINCREMENT, task_id TEXT NOT NULL REFERENCES tasks(task_id),
|
|
usage_json TEXT NOT NULL, created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
|
|
);
|
|
CREATE TABLE IF NOT EXISTS tool_audits (
|
|
audit_id INTEGER PRIMARY KEY AUTOINCREMENT, task_id TEXT NOT NULL REFERENCES tasks(task_id),
|
|
audit_json TEXT NOT NULL, created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
|
|
);
|
|
CREATE TABLE IF NOT EXISTS model_capabilities (
|
|
provider_id TEXT NOT NULL, model_id TEXT NOT NULL, schema_hash TEXT NOT NULL,
|
|
supported INTEGER NOT NULL, report_json TEXT NOT NULL, checked_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
|
PRIMARY KEY(provider_id,model_id,schema_hash)
|
|
);
|
|
""")
|
|
|
|
def create_task(self, task_id: str, request: str) -> TaskState:
|
|
with self._lock, self._connection() as connection:
|
|
connection.execute(
|
|
"INSERT OR IGNORE INTO tasks(task_id,protocol_version,request,phase,state_version) VALUES (?,?,?,?,0)",
|
|
(task_id, self.PROTOCOL_VERSION, request, TaskPhase.ANALYZING_REQUEST.value),
|
|
)
|
|
state = self.get_state(task_id)
|
|
if state is None:
|
|
raise RuntimeError("Task was not persisted")
|
|
return state
|
|
|
|
def get_state(self, task_id: str) -> TaskState | None:
|
|
with self._lock, self._connection() as connection:
|
|
row = connection.execute("SELECT * FROM tasks WHERE task_id=?", (task_id,)).fetchone()
|
|
return self._state(row) if row else None
|
|
|
|
def get_task_projection(self, task_id: str) -> dict[str, Any] | None:
|
|
state = self.get_state(task_id)
|
|
if state is None:
|
|
return None
|
|
events = self.ledger_events(task_id)
|
|
# A later repair can end at authoring or compilation after an earlier
|
|
# build already published an executable prefix. Its terminal event
|
|
# intentionally contains diagnostics rather than a duplicate path
|
|
# map, so keep the most recent event that actually delivered a
|
|
# revision. Otherwise the completed projection loses the existing
|
|
# STEP/GLB links exactly when best-effort publication matters most.
|
|
result: dict[str, Any] = {}
|
|
paths: dict[str, Any] = {}
|
|
for item in reversed(events):
|
|
if item.get("event") not in {
|
|
"build_completed", "published_best_effort", "build_failed",
|
|
}:
|
|
continue
|
|
candidate_paths = item.get("paths")
|
|
if isinstance(candidate_paths, dict) and candidate_paths:
|
|
result, paths = item, candidate_paths
|
|
break
|
|
return {
|
|
"schema_version": self.PROTOCOL_VERSION, "task_id": task_id, "phase": state.phase.value,
|
|
"lifecycle": self._lifecycle(state.phase), "state_version": state.version,
|
|
"active_revision": state.active_revision, "current_revision": state.active_revision,
|
|
"published_revision": state.active_revision if state.phase == TaskPhase.COMPLETED else "",
|
|
"repair_count": state.repair_count, "repair_budget": 2,
|
|
"last_error": state.last_error.value if state.last_error else "",
|
|
"requirements_path": state.requirements_path, "authoring_path": state.authoring_path,
|
|
"runtime_cdsl_path": state.runtime_cdsl_path, "compile_audit_path": state.compile_audit_path,
|
|
"diagnostics_path": state.diagnostics_path, "completion_path": state.completion_path,
|
|
"message": str(result.get("message") or ""),
|
|
"revisions": ([{
|
|
"revision_id": state.active_revision, "status": "success", "visibility": "final" if state.phase == TaskPhase.COMPLETED else "checkpoint",
|
|
"cdsl_path": paths.get("model.cdsl.json", ""), "step_path": paths.get("model.step", ""),
|
|
"glb_path": paths.get("model.glb", ""), "report_path": paths.get("rebuild-report.json", ""),
|
|
}] if state.active_revision else []),
|
|
"action_ledger_summary": events[-16:],
|
|
}
|
|
|
|
|
|
def ledger_events(self, task_id: str) -> list[dict[str, Any]]:
|
|
with self._lock, self._connection() as connection:
|
|
rows = connection.execute("SELECT sequence,event_json,created_at FROM ledger WHERE task_id=? ORDER BY sequence", (task_id,)).fetchall()
|
|
return [{"sequence": int(row["sequence"]), "at": str(row["created_at"]), **json.loads(row["event_json"])} for row in rows]
|
|
|
|
def compare_and_swap(self, state: TaskState, *, events: list[dict[str, Any]] = (), invocation_id: str | None = None, invocation_result: dict[str, Any] | None = None) -> bool:
|
|
if (invocation_id is None) != (invocation_result is None):
|
|
raise ValueError("Invocation completion requires both values")
|
|
previous = state.version - 1
|
|
with self._lock, self._connection() as connection:
|
|
connection.execute("BEGIN IMMEDIATE")
|
|
try:
|
|
cursor = connection.execute("""
|
|
UPDATE tasks SET phase=?,state_version=?,active_revision=?,repair_count=?,last_error=?,retry_from_phase=?,
|
|
requirements_path=?,authoring_path=?,runtime_cdsl_path=?,compile_audit_path=?,diagnostics_path=?,completion_path=?,clarification_path=?,updated_at=CURRENT_TIMESTAMP
|
|
WHERE task_id=? AND state_version=?
|
|
""", (
|
|
state.phase.value, state.version, state.active_revision, state.repair_count,
|
|
state.last_error.value if state.last_error else None, state.retry_from_phase.value if state.retry_from_phase else "",
|
|
state.requirements_path, state.authoring_path, state.runtime_cdsl_path, state.compile_audit_path,
|
|
state.diagnostics_path, state.completion_path, state.clarification_path, state.task_id, previous,
|
|
))
|
|
if cursor.rowcount != 1:
|
|
connection.execute("ROLLBACK")
|
|
return False
|
|
for item in events:
|
|
encoded = json.dumps(item, ensure_ascii=True, sort_keys=True)
|
|
connection.execute("INSERT INTO ledger(task_id,event_json) VALUES (?,?)", (state.task_id, encoded))
|
|
connection.execute("INSERT INTO outbox(task_id,event_json) VALUES (?,?)", (state.task_id, encoded))
|
|
if invocation_id is not None:
|
|
updated = connection.execute("UPDATE invocations SET status='finished',result_json=?,updated_at=CURRENT_TIMESTAMP WHERE invocation_id=? AND task_id=? AND status='processing'", (json.dumps(invocation_result, ensure_ascii=True), invocation_id, state.task_id))
|
|
if updated.rowcount != 1:
|
|
raise ValueError("Invocation is not active")
|
|
connection.execute("COMMIT")
|
|
return True
|
|
except Exception:
|
|
connection.execute("ROLLBACK")
|
|
raise
|
|
|
|
def begin_invocation(self, task_id: str, invocation_id: str, idempotency_key: str) -> InvocationRecord:
|
|
with self._lock, self._connection() as connection:
|
|
connection.execute("BEGIN IMMEDIATE")
|
|
try:
|
|
row = connection.execute("SELECT * FROM invocations WHERE task_id=? AND idempotency_key=?", (task_id, idempotency_key)).fetchone()
|
|
if row is None:
|
|
connection.execute("INSERT INTO invocations(invocation_id,task_id,idempotency_key,status) VALUES (?,?,?,'processing')", (invocation_id, task_id, idempotency_key))
|
|
connection.execute("COMMIT")
|
|
return InvocationRecord(invocation_id, idempotency_key, "processing")
|
|
connection.execute("COMMIT")
|
|
return self._invocation(row)
|
|
except Exception:
|
|
connection.execute("ROLLBACK")
|
|
raise
|
|
|
|
def get_invocation(self, task_id: str, invocation_id: str) -> InvocationRecord | None:
|
|
with self._lock, self._connection() as connection:
|
|
row = connection.execute("SELECT * FROM invocations WHERE task_id=? AND invocation_id=?", (task_id, invocation_id)).fetchone()
|
|
return self._invocation(row) if row else None
|
|
|
|
def finish_invocation(self, invocation_id: str, result: dict[str, Any]) -> None:
|
|
with self._lock, self._connection() as connection:
|
|
connection.execute("UPDATE invocations SET status='finished',result_json=?,updated_at=CURRENT_TIMESTAMP WHERE invocation_id=? AND status='processing'", (json.dumps(result, ensure_ascii=True), invocation_id))
|
|
|
|
def append_outbox(self, task_id: str, event: dict[str, Any]) -> None:
|
|
with self._lock, self._connection() as connection:
|
|
connection.execute("INSERT INTO outbox(task_id,event_json) VALUES (?,?)", (task_id, json.dumps(event, ensure_ascii=True, sort_keys=True)))
|
|
|
|
def pending_outbox(self, limit: int = 100, *, task_id: str | None = None) -> list[dict[str, Any]]:
|
|
query = "SELECT event_id,task_id,event_json FROM outbox WHERE published_at IS NULL"
|
|
args: tuple[Any, ...] = ()
|
|
if task_id is not None:
|
|
query += " AND task_id=?"
|
|
args += (task_id,)
|
|
query += " ORDER BY event_id LIMIT ?"
|
|
args += (limit,)
|
|
with self._lock, self._connection() as connection:
|
|
rows = connection.execute(query, args).fetchall()
|
|
return [{"event_id": int(row["event_id"]), "task_id": str(row["task_id"]), **json.loads(row["event_json"])} for row in rows]
|
|
|
|
def mark_outbox_published(self, event_id: int) -> None:
|
|
with self._lock, self._connection() as connection:
|
|
connection.execute("UPDATE outbox SET published_at=CURRENT_TIMESTAMP WHERE event_id=? AND published_at IS NULL", (event_id,))
|
|
|
|
def record_usage(self, task_id: str, payload: dict[str, Any]) -> None:
|
|
with self._lock, self._connection() as connection:
|
|
connection.execute("INSERT INTO usage_records(task_id,usage_json) VALUES (?,?)", (task_id, json.dumps(payload, ensure_ascii=True, sort_keys=True)))
|
|
|
|
def record_tool_audit(self, task_id: str, payload: dict[str, Any]) -> None:
|
|
with self._lock, self._connection() as connection:
|
|
connection.execute("INSERT INTO tool_audits(task_id,audit_json) VALUES (?,?)", (task_id, json.dumps(payload, ensure_ascii=True, sort_keys=True)))
|
|
|
|
def tool_audits(self, task_id: str) -> list[dict[str, Any]]:
|
|
with self._lock, self._connection() as connection:
|
|
rows = connection.execute("SELECT audit_id,audit_json,created_at FROM tool_audits WHERE task_id=? ORDER BY audit_id", (task_id,)).fetchall()
|
|
return [{"audit_id": int(row["audit_id"]), "at": str(row["created_at"]), **json.loads(row["audit_json"])} for row in rows]
|
|
|
|
def usage_summary(self, task_id: str) -> dict[str, Any]:
|
|
with self._lock, self._connection() as connection:
|
|
rows = connection.execute("SELECT usage_json FROM usage_records WHERE task_id=? ORDER BY usage_id", (task_id,)).fetchall()
|
|
return {"calls": len(rows), "records": [json.loads(row["usage_json"]) for row in rows]}
|
|
|
|
def running_task_ids(self) -> list[str]:
|
|
terminal = tuple(item.value for item in (TaskPhase.COMPLETED, TaskPhase.FAILED, TaskPhase.CANCELLED, TaskPhase.WAITING_FOR_USER))
|
|
with self._lock, self._connection() as connection:
|
|
rows = connection.execute("SELECT task_id FROM tasks WHERE phase NOT IN (?,?,?,?) ORDER BY created_at", terminal).fetchall()
|
|
return [str(row["task_id"]) for row in rows]
|
|
|
|
def model_capability(self, provider_id: str, model_id: str, schema_hash: str) -> dict[str, Any] | None:
|
|
with self._lock, self._connection() as connection:
|
|
row = connection.execute("SELECT supported,report_json,checked_at FROM model_capabilities WHERE provider_id=? AND model_id=? AND schema_hash=?", (provider_id, model_id, schema_hash)).fetchone()
|
|
return {"supported": bool(row["supported"]), "report": json.loads(row["report_json"]), "checked_at": str(row["checked_at"])} if row else None
|
|
|
|
def record_model_capability(self, provider_id: str, model_id: str, schema_hash: str, report: dict[str, Any]) -> None:
|
|
with self._lock, self._connection() as connection:
|
|
connection.execute("""INSERT INTO model_capabilities(provider_id,model_id,schema_hash,supported,report_json) VALUES (?,?,?,?,?)
|
|
ON CONFLICT(provider_id,model_id,schema_hash) DO UPDATE SET supported=excluded.supported,report_json=excluded.report_json,checked_at=CURRENT_TIMESTAMP""", (provider_id, model_id, schema_hash, int(bool(report.get("supported"))), json.dumps(report, ensure_ascii=True, sort_keys=True)))
|
|
|
|
@staticmethod
|
|
def _state(row: sqlite3.Row) -> TaskState:
|
|
return TaskState(
|
|
task_id=str(row["task_id"]), phase=TaskPhase(str(row["phase"])), version=int(row["state_version"]),
|
|
active_revision=str(row["active_revision"]), repair_count=int(row["repair_count"]),
|
|
last_error=ErrorCode(str(row["last_error"])) if row["last_error"] else None,
|
|
retry_from_phase=TaskPhase(str(row["retry_from_phase"])) if row["retry_from_phase"] else None,
|
|
requirements_path=str(row["requirements_path"]), authoring_path=str(row["authoring_path"]),
|
|
runtime_cdsl_path=str(row["runtime_cdsl_path"]), compile_audit_path=str(row["compile_audit_path"]),
|
|
diagnostics_path=str(row["diagnostics_path"]), completion_path=str(row["completion_path"]), clarification_path=str(row["clarification_path"]),
|
|
)
|
|
|
|
@staticmethod
|
|
def _invocation(row: sqlite3.Row) -> InvocationRecord:
|
|
return InvocationRecord(str(row["invocation_id"]), str(row["idempotency_key"]), str(row["status"]), json.loads(row["result_json"]) if row["result_json"] else None)
|
|
|
|
@staticmethod
|
|
def _lifecycle(phase: TaskPhase) -> str:
|
|
if phase == TaskPhase.COMPLETED:
|
|
return "completed"
|
|
if phase == TaskPhase.FAILED:
|
|
return "failed"
|
|
if phase == TaskPhase.CANCELLED:
|
|
return "cancelled"
|
|
if phase == TaskPhase.WAITING_FOR_USER:
|
|
return "waiting_for_user"
|
|
return "running"
|