302 lines
11 KiB
Python
302 lines
11 KiB
Python
"""Durable checkpoint journal for the generic Online Agent worker.
|
|
|
|
The JSON task state remains a fast cache. This SQLite journal is the durable,
|
|
append-only evidence used for restart/resume/replay decisions.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from datetime import datetime, timezone
|
|
import hashlib
|
|
import json
|
|
from pathlib import Path
|
|
import sqlite3
|
|
from typing import Any, Mapping
|
|
|
|
PHASE_ORDER = {
|
|
"claimed": 10,
|
|
"workspace_prepared": 20,
|
|
"verification_passed": 30,
|
|
"source_capture_done": 40,
|
|
"commit_done": 40,
|
|
"reconciliation_required": 45,
|
|
"result_pending": 50,
|
|
"result_submitted": 60,
|
|
}
|
|
|
|
|
|
class CheckpointError(RuntimeError):
|
|
pass
|
|
|
|
|
|
def canonical(value: object) -> bytes:
|
|
return json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False, default=str).encode("utf-8")
|
|
|
|
|
|
def sha256_json(value: object) -> str:
|
|
return hashlib.sha256(canonical(value)).hexdigest()
|
|
|
|
|
|
def now() -> str:
|
|
return datetime.now(timezone.utc).isoformat()
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class CheckpointBinding:
|
|
task_id: str
|
|
attempt: int
|
|
execution_digest: str
|
|
job_id: str
|
|
repository: str
|
|
base_oid: str
|
|
expected_branch_oid: str
|
|
target_branch: str
|
|
generation: int
|
|
|
|
def material(self) -> dict[str, Any]:
|
|
return {
|
|
"task_id": self.task_id,
|
|
"attempt": self.attempt,
|
|
"execution_digest": self.execution_digest,
|
|
"job_id": self.job_id,
|
|
"repository": self.repository,
|
|
"base_oid": self.base_oid,
|
|
"expected_branch_oid": self.expected_branch_oid,
|
|
"target_branch": self.target_branch,
|
|
"generation": self.generation,
|
|
}
|
|
|
|
@property
|
|
def binding_sha256(self) -> str:
|
|
return sha256_json(self.material())
|
|
|
|
|
|
class CheckpointStore:
|
|
def __init__(self, path: Path) -> None:
|
|
self.path = Path(path)
|
|
self.path.parent.mkdir(parents=True, exist_ok=True)
|
|
with self._connect() as connection:
|
|
connection.executescript(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS checkpoints(
|
|
task_id TEXT NOT NULL,
|
|
attempt INTEGER NOT NULL,
|
|
sequence INTEGER NOT NULL,
|
|
phase TEXT NOT NULL,
|
|
phase_order INTEGER NOT NULL,
|
|
execution_digest TEXT NOT NULL,
|
|
job_id TEXT NOT NULL,
|
|
repository TEXT NOT NULL,
|
|
base_oid TEXT NOT NULL,
|
|
expected_branch_oid TEXT NOT NULL,
|
|
target_branch TEXT NOT NULL,
|
|
generation INTEGER NOT NULL,
|
|
binding_sha256 TEXT NOT NULL,
|
|
payload_json TEXT NOT NULL,
|
|
payload_sha256 TEXT NOT NULL,
|
|
created_at TEXT NOT NULL,
|
|
PRIMARY KEY(task_id, attempt, sequence)
|
|
);
|
|
CREATE INDEX IF NOT EXISTS idx_checkpoints_task_latest
|
|
ON checkpoints(task_id, attempt DESC, sequence DESC);
|
|
CREATE UNIQUE INDEX IF NOT EXISTS idx_checkpoints_exact_replay
|
|
ON checkpoints(task_id, attempt, phase, payload_sha256);
|
|
CREATE TABLE IF NOT EXISTS checkpoint_heads(
|
|
task_id TEXT PRIMARY KEY,
|
|
attempt INTEGER NOT NULL,
|
|
sequence INTEGER NOT NULL,
|
|
phase TEXT NOT NULL,
|
|
phase_order INTEGER NOT NULL,
|
|
execution_digest TEXT NOT NULL,
|
|
binding_sha256 TEXT NOT NULL,
|
|
payload_sha256 TEXT NOT NULL,
|
|
updated_at TEXT NOT NULL
|
|
);
|
|
"""
|
|
)
|
|
|
|
def _connect(self) -> sqlite3.Connection:
|
|
connection = sqlite3.connect(self.path, timeout=30.0)
|
|
connection.row_factory = sqlite3.Row
|
|
connection.execute("PRAGMA foreign_keys=ON")
|
|
connection.execute("PRAGMA journal_mode=WAL")
|
|
connection.execute("PRAGMA synchronous=FULL")
|
|
connection.execute("PRAGMA busy_timeout=30000")
|
|
return connection
|
|
|
|
@staticmethod
|
|
def _phase_order(phase: str) -> int:
|
|
try:
|
|
return PHASE_ORDER[phase]
|
|
except KeyError as exc:
|
|
raise CheckpointError("CHECKPOINT_PHASE_INVALID") from exc
|
|
|
|
def append(
|
|
self,
|
|
binding: CheckpointBinding,
|
|
phase: str,
|
|
payload: Mapping[str, Any],
|
|
) -> dict[str, Any]:
|
|
order = self._phase_order(phase)
|
|
payload_value = dict(payload)
|
|
payload_sha = sha256_json(payload_value)
|
|
with self._connect() as connection:
|
|
connection.execute("BEGIN IMMEDIATE")
|
|
replay = connection.execute(
|
|
"""
|
|
SELECT * FROM checkpoints
|
|
WHERE task_id=? AND attempt=? AND phase=? AND payload_sha256=?
|
|
""",
|
|
(binding.task_id, binding.attempt, phase, payload_sha),
|
|
).fetchone()
|
|
if replay is not None:
|
|
connection.commit()
|
|
return dict(replay)
|
|
|
|
head = connection.execute(
|
|
"SELECT * FROM checkpoint_heads WHERE task_id=?",
|
|
(binding.task_id,),
|
|
).fetchone()
|
|
if head is not None:
|
|
head_attempt = int(head["attempt"])
|
|
head_order = int(head["phase_order"])
|
|
if binding.attempt < head_attempt:
|
|
raise CheckpointError("CHECKPOINT_ATTEMPT_REGRESSION")
|
|
if binding.attempt == head_attempt:
|
|
if head["execution_digest"] != binding.execution_digest:
|
|
raise CheckpointError("CHECKPOINT_EXECUTION_DIGEST_CONFLICT")
|
|
if head["binding_sha256"] != binding.binding_sha256:
|
|
raise CheckpointError("CHECKPOINT_SCOPE_MISMATCH_DENY")
|
|
if order < head_order and phase != "reconciliation_required":
|
|
raise CheckpointError("CHECKPOINT_PHASE_REGRESSION")
|
|
sequence = int(head["sequence"]) + 1
|
|
else:
|
|
sequence = 1
|
|
else:
|
|
sequence = 1
|
|
|
|
created_at = now()
|
|
connection.execute(
|
|
"""
|
|
INSERT INTO checkpoints(
|
|
task_id,attempt,sequence,phase,phase_order,execution_digest,job_id,
|
|
repository,base_oid,expected_branch_oid,target_branch,generation,
|
|
binding_sha256,payload_json,payload_sha256,created_at
|
|
) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)
|
|
""",
|
|
(
|
|
binding.task_id,
|
|
binding.attempt,
|
|
sequence,
|
|
phase,
|
|
order,
|
|
binding.execution_digest,
|
|
binding.job_id,
|
|
binding.repository,
|
|
binding.base_oid,
|
|
binding.expected_branch_oid,
|
|
binding.target_branch,
|
|
binding.generation,
|
|
binding.binding_sha256,
|
|
canonical(payload_value).decode("utf-8"),
|
|
payload_sha,
|
|
created_at,
|
|
),
|
|
)
|
|
connection.execute(
|
|
"""
|
|
INSERT INTO checkpoint_heads(
|
|
task_id,attempt,sequence,phase,phase_order,execution_digest,
|
|
binding_sha256,payload_sha256,updated_at
|
|
) VALUES (?,?,?,?,?,?,?,?,?)
|
|
ON CONFLICT(task_id) DO UPDATE SET
|
|
attempt=excluded.attempt,
|
|
sequence=excluded.sequence,
|
|
phase=excluded.phase,
|
|
phase_order=excluded.phase_order,
|
|
execution_digest=excluded.execution_digest,
|
|
binding_sha256=excluded.binding_sha256,
|
|
payload_sha256=excluded.payload_sha256,
|
|
updated_at=excluded.updated_at
|
|
""",
|
|
(
|
|
binding.task_id,
|
|
binding.attempt,
|
|
sequence,
|
|
phase,
|
|
order,
|
|
binding.execution_digest,
|
|
binding.binding_sha256,
|
|
payload_sha,
|
|
created_at,
|
|
),
|
|
)
|
|
connection.commit()
|
|
row = connection.execute(
|
|
"SELECT * FROM checkpoints WHERE task_id=? AND attempt=? AND sequence=?",
|
|
(binding.task_id, binding.attempt, sequence),
|
|
).fetchone()
|
|
return dict(row)
|
|
|
|
def latest(self, task_id: str) -> dict[str, Any] | None:
|
|
with self._connect() as connection:
|
|
row = connection.execute(
|
|
"""
|
|
SELECT * FROM checkpoints
|
|
WHERE task_id=?
|
|
ORDER BY attempt DESC, sequence DESC
|
|
LIMIT 1
|
|
""",
|
|
(task_id,),
|
|
).fetchone()
|
|
if row is None:
|
|
return None
|
|
value = dict(row)
|
|
value["payload"] = json.loads(value.pop("payload_json"))
|
|
return value
|
|
|
|
def latest_for_attempt(self, task_id: str, attempt: int) -> dict[str, Any] | None:
|
|
with self._connect() as connection:
|
|
row = connection.execute(
|
|
"""
|
|
SELECT * FROM checkpoints
|
|
WHERE task_id=? AND attempt=?
|
|
ORDER BY sequence DESC
|
|
LIMIT 1
|
|
""",
|
|
(task_id, attempt),
|
|
).fetchone()
|
|
if row is None:
|
|
return None
|
|
value = dict(row)
|
|
value["payload"] = json.loads(value.pop("payload_json"))
|
|
return value
|
|
|
|
def verify_latest_binding(self, binding: CheckpointBinding) -> dict[str, Any] | None:
|
|
value = self.latest(binding.task_id)
|
|
if value is None:
|
|
return None
|
|
if int(value["attempt"]) > binding.attempt:
|
|
raise CheckpointError("CHECKPOINT_FUTURE_ATTEMPT")
|
|
# A previous attempt may be reused only when the immutable execution/scope binding matches.
|
|
if value["execution_digest"] != binding.execution_digest:
|
|
return None
|
|
expected = CheckpointBinding(
|
|
task_id=value["task_id"],
|
|
attempt=int(value["attempt"]),
|
|
execution_digest=value["execution_digest"],
|
|
job_id=value["job_id"],
|
|
repository=value["repository"],
|
|
base_oid=value["base_oid"],
|
|
expected_branch_oid=value["expected_branch_oid"],
|
|
target_branch=value["target_branch"],
|
|
generation=int(value["generation"]),
|
|
)
|
|
current_scope = binding.material()
|
|
historical_scope = expected.material()
|
|
current_scope.pop("attempt", None)
|
|
historical_scope.pop("attempt", None)
|
|
if sha256_json(current_scope) != sha256_json(historical_scope):
|
|
raise CheckpointError("CHECKPOINT_SCOPE_MISMATCH_DENY")
|
|
return value
|