76 lines
2.5 KiB
Python
76 lines
2.5 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import sqlite3
|
|
from pathlib import Path
|
|
|
|
from app.core.contracts import TaskCheckpoint
|
|
|
|
|
|
class SQLiteCheckpointStore:
|
|
"""Durable checkpoint store for resumable runtime state."""
|
|
|
|
def __init__(self, db_path: str | Path) -> None:
|
|
self._db_path = Path(db_path)
|
|
self._db_path.parent.mkdir(parents=True, exist_ok=True)
|
|
self._initialize()
|
|
|
|
def save(self, checkpoint: TaskCheckpoint) -> TaskCheckpoint:
|
|
with sqlite3.connect(self._db_path) as conn:
|
|
conn.execute(
|
|
"""
|
|
INSERT OR REPLACE INTO checkpoints (
|
|
task_id, status, active_step_id, plan_snapshot_json,
|
|
context_snapshot_json, updated_at
|
|
) VALUES (?, ?, ?, ?, ?, ?)
|
|
""",
|
|
(
|
|
checkpoint.task_id,
|
|
checkpoint.status,
|
|
checkpoint.active_step_id,
|
|
json.dumps(checkpoint.plan_snapshot, default=str),
|
|
json.dumps(checkpoint.context_snapshot, default=str),
|
|
checkpoint.updated_at.isoformat(),
|
|
),
|
|
)
|
|
conn.commit()
|
|
return checkpoint
|
|
|
|
def load(self, task_id: str) -> TaskCheckpoint | None:
|
|
with sqlite3.connect(self._db_path) as conn:
|
|
row = conn.execute(
|
|
"""
|
|
SELECT task_id, status, active_step_id, plan_snapshot_json,
|
|
context_snapshot_json, updated_at
|
|
FROM checkpoints
|
|
WHERE task_id = ?
|
|
""",
|
|
(task_id,),
|
|
).fetchone()
|
|
if not row:
|
|
return None
|
|
return TaskCheckpoint(
|
|
task_id=row[0],
|
|
status=row[1],
|
|
active_step_id=row[2],
|
|
plan_snapshot=json.loads(row[3]),
|
|
context_snapshot=json.loads(row[4]),
|
|
updated_at=row[5],
|
|
)
|
|
|
|
def _initialize(self) -> None:
|
|
with sqlite3.connect(self._db_path) as conn:
|
|
conn.execute(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS checkpoints (
|
|
task_id TEXT PRIMARY KEY,
|
|
status TEXT NOT NULL,
|
|
active_step_id TEXT,
|
|
plan_snapshot_json TEXT NOT NULL,
|
|
context_snapshot_json TEXT NOT NULL,
|
|
updated_at TEXT NOT NULL
|
|
)
|
|
"""
|
|
)
|
|
conn.commit()
|