Managed Agent の長時間走行がサンドボックス再生成で消える前に — チェックポイントと冪等リジュームの設計

import sqlite3
import json
import time
from contextlib import contextmanager
 
class Checkpoint:
    """サンドボックスの外に置く、進捗と確定済み結果の永続ストア。"""
 
    def __init__(self, path: str, run_id: str):
        self.run_id = run_id
        self.conn = sqlite3.connect(path, isolation_level=None)  # 自前でBEGIN/COMMIT
        self.conn.execute("PRAGMA journal_mode=WAL;")
        self._init_schema()
 
    def _init_schema(self):
        self.conn.executescript(
            """
            CREATE TABLE IF NOT EXISTS progress (
                run_id     TEXT PRIMARY KEY,
                step       INTEGER NOT NULL,
                state_json TEXT NOT NULL,
                updated_at REAL NOT NULL
            );
            CREATE TABLE IF NOT EXISTS committed (
                idem_key   TEXT PRIMARY KEY,
                run_id     TEXT NOT NULL,
                result_json TEXT NOT NULL,
                committed_at REAL NOT NULL
            );
            CREATE TABLE IF NOT EXISTS claims (
                idem_key   TEXT PRIMARY KEY,
                run_id     TEXT NOT NULL,
                claimed_at REAL NOT NULL
            );
            """
        )
 
    @contextmanager
    def _tx(self):
        self.conn.execute("BEGIN IMMEDIATE;")
        try:
            yield
            self.conn.execute("COMMIT;")
        except Exception:
            self.conn.execute("ROLLBACK;")
            raise
 
    def load(self):
        """再開時に呼ぶ。到達済みステップと再開状態を返す。"""
        row = self.conn.execute(
            "SELECT step, state_json FROM progress WHERE run_id=?",
            (self.run_id,),
        ).fetchone()
        if row is None:
            return 0, {}
        return row[0], json.loads(row[1])
 
    def save_progress(self, step: int, state: dict):
        """地図だけを更新する。荷物は committed 側に既に入っている。"""
        with self._tx():
            self.conn.execute(
                "INSERT INTO progress(run_id, step, state_json, updated_at) "
                "VALUES(?,?,?,?) "
                "ON CONFLICT(run_id) DO UPDATE SET "
                "step=excluded.step, state_json=excluded.state_json, "
                "updated_at=excluded.updated_at",
                (self.run_id, step, json.dumps(state), time.time()),
            )

    def run_once(self, idem_key: str, side_effect):
        """claim → execute → commit。同じ idem_key では副作用を一度だけ走らせる。"""
        # --- claim 相: 既に確定済みなら結果を返して終了 ---
        done = self.conn.execute(
            "SELECT result_json FROM committed WHERE idem_key=?",
            (idem_key,),
        ).fetchone()
        if done is not None:
            return json.loads(done[0]), "skipped"
 
        # 予約を立てる。競合したら（別走行が実行中なら）確定を待つ。
        try:
            with self._tx():
                self.conn.execute(
                    "INSERT INTO claims(idem_key, run_id, claimed_at) VALUES(?,?,?)",
                    (idem_key, self.run_id, time.time()),
                )
        except sqlite3.IntegrityError:
            return self._await_commit(idem_key), "awaited"
 
        # --- execute 相: ここに来た走行だけが副作用を実行する ---
        result = side_effect()
 
        # --- commit 相: 結果と確定印を同一トランザクションで書く ---
        with self._tx():
            self.conn.execute(
                "INSERT INTO committed(idem_key, run_id, result_json, committed_at) "
                "VALUES(?,?,?,?)",
                (idem_key, self.run_id, json.dumps(result), time.time()),
            )
            self.conn.execute("DELETE FROM claims WHERE idem_key=?", (idem_key,))
        return result, "executed"
 
    def _await_commit(self, idem_key: str, timeout: float = 30.0):
        deadline = time.time() + timeout
        while time.time() < deadline:
            row = self.conn.execute(
                "SELECT result_json FROM committed WHERE idem_key=?",
                (idem_key,),
            ).fetchone()
            if row is not None:
                return json.loads(row[0])
            time.sleep(0.5)
        raise TimeoutError(f"commit を待てませんでした: {idem_key}")

def run_agent(inputs, ckpt: Checkpoint, call_model):
    start_step, state = ckpt.load()   # 再開点をストアから得る
    processed = state.get("processed", 0)
 
    for step in range(start_step, len(inputs)):
        item = inputs[step]
        idem_key = f"{ckpt.run_id}:step:{step}"
 
        # 副作用（モデル呼び出し＋外部への確定書き込み）を一度だけ
        result, status = ckpt.run_once(
            idem_key,
            lambda: call_model(item),
        )
        processed += 1 if status == "executed" else 0
 
        # 地図だけを更新する。結果は committed 側に既にある。
        ckpt.save_progress(step + 1, {"processed": processed})
 
    return processed