| """
|
| Checkpoint Management for Agent State Persistence
|
| Part of SOVEREIGN PYTHON LLM ENGINE
|
|
|
| Backed by binary struct storage (CheckpointFile from core.storage).
|
| No JSONL, no text encoding, no injection surface.
|
|
|
| Features:
|
| - Save/restore agent state to disk (binary struct format)
|
| - Cryptographic signing (Ed25519)
|
| - WORM-compatible append-only log
|
| - Support for pause/resume of long-running agents
|
| - Time-travel debugging (replay from any checkpoint)
|
| """
|
|
|
| from dataclasses import dataclass, asdict
|
| from datetime import datetime, timezone
|
| from pathlib import Path
|
| from typing import Any, AsyncIterator
|
| import json
|
| import uuid
|
|
|
| from ..core.crypto import ContentHash, Signature, SigningKey, hash_content, sign_artifact
|
| from ..core.storage import CheckpointFile, CheckpointRecord
|
|
|
|
|
|
|
|
|
|
|
|
|
| @dataclass
|
| class AgentState:
|
| """
|
| Complete agent state snapshot.
|
|
|
| Fields:
|
| agent_id: Unique agent identifier
|
| task: Current task description
|
| step_number: Current step in execution
|
| messages: Conversation history
|
| tool_results: Results from tool executions
|
| reasoning: Agent's reasoning trace
|
| metadata: Additional state data
|
| """
|
| agent_id: str
|
| task: str
|
| step_number: int
|
| messages: list[dict[str, str]]
|
| tool_results: list[dict[str, Any]]
|
| reasoning: list[str]
|
| metadata: dict[str, Any]
|
|
|
| def to_dict(self) -> dict[str, Any]:
|
| return asdict(self)
|
|
|
| @classmethod
|
| def from_dict(cls, data: dict[str, Any]) -> 'AgentState':
|
| return cls(**data)
|
|
|
| def to_bytes(self) -> bytes:
|
| return json.dumps(self.to_dict(), sort_keys=True, separators=(',', ':')).encode('utf-8')
|
|
|
| @classmethod
|
| def from_bytes(cls, raw: bytes) -> 'AgentState':
|
| return cls.from_dict(json.loads(raw.decode('utf-8')))
|
|
|
|
|
|
|
|
|
|
|
|
|
| @dataclass(frozen=True)
|
| class Checkpoint:
|
| """
|
| Immutable checkpoint record.
|
|
|
| Fields:
|
| checkpoint_id: Unique checkpoint ID (UUID)
|
| agent_id: Agent identifier
|
| step_number: Step number when checkpoint was created
|
| state: Full agent state snapshot
|
| timestamp: UTC timestamp (ISO 8601)
|
| previous_checkpoint: ID of previous checkpoint (for chain)
|
| hash: Blake2b hash of state
|
| signature: Ed25519 signature
|
| """
|
| checkpoint_id: str
|
| agent_id: str
|
| step_number: int
|
| state: AgentState
|
| timestamp: str
|
| previous_checkpoint: str | None
|
| hash: ContentHash
|
| signature: Signature
|
|
|
| def to_dict(self) -> dict[str, Any]:
|
| return {
|
| 'checkpoint_id': self.checkpoint_id,
|
| 'agent_id': self.agent_id,
|
| 'step_number': self.step_number,
|
| 'state': self.state.to_dict(),
|
| 'timestamp': self.timestamp,
|
| 'previous_checkpoint': self.previous_checkpoint,
|
| 'hash': self.hash,
|
| 'signature': self.signature
|
| }
|
|
|
| @classmethod
|
| def from_dict(cls, data: dict[str, Any]) -> 'Checkpoint':
|
| state_data = data['state']
|
| state = AgentState.from_dict(state_data)
|
|
|
| return cls(
|
| checkpoint_id=data['checkpoint_id'],
|
| agent_id=data['agent_id'],
|
| step_number=data['step_number'],
|
| state=state,
|
| timestamp=data['timestamp'],
|
| previous_checkpoint=data.get('previous_checkpoint'),
|
| hash=ContentHash(data['hash']),
|
| signature=Signature(data['signature'])
|
| )
|
|
|
| @classmethod
|
| def from_record(cls, rec: CheckpointRecord) -> 'Checkpoint':
|
| state = AgentState.from_bytes(rec.state)
|
| return cls(
|
| checkpoint_id=rec.checkpoint_id,
|
| agent_id=rec.agent_id,
|
| step_number=rec.step_number,
|
| state=state,
|
| timestamp=datetime.fromtimestamp(
|
| rec.timestamp_ns / 1_000_000_000, tz=timezone.utc
|
| ).isoformat(),
|
| previous_checkpoint=rec.prev_id or None,
|
| hash=rec.content_hash,
|
| signature=rec.signature
|
| )
|
|
|
|
|
|
|
|
|
|
|
|
|
| class CheckpointManager:
|
| """
|
| Manage agent checkpoints with binary struct storage.
|
|
|
| Uses CheckpointFile from core.storage (128-byte fixed headers,
|
| Blake2b hash chain, Ed25519 signatures). No JSONL.
|
| """
|
|
|
| def __init__(self, checkpoint_dir: Path, signing_key: SigningKey):
|
| self.checkpoint_dir = checkpoint_dir
|
| self.signing_key = signing_key
|
| self.checkpoint_dir.mkdir(parents=True, exist_ok=True)
|
| self._last_checkpoint: dict[str, str] = {}
|
|
|
| def _get_checkpoint_file(self, agent_id: str) -> CheckpointFile:
|
| path = self.checkpoint_dir / f"{agent_id}.ckpt"
|
| return CheckpointFile(path, signing_key=self.signing_key)
|
|
|
| async def save_checkpoint(
|
| self,
|
| agent_id: str,
|
| state: AgentState,
|
| step_number: int | None = None
|
| ) -> Checkpoint:
|
| """
|
| Save agent state checkpoint to binary storage.
|
| """
|
| if step_number is None:
|
| step_number = state.step_number
|
|
|
| checkpoint_id = str(uuid.uuid4())
|
| previous_checkpoint = self._last_checkpoint.get(agent_id, "")
|
|
|
| state_bytes = state.to_bytes()
|
| state_hash = hash_content(state_bytes)
|
|
|
| signing_payload = f"{checkpoint_id}|{agent_id}|{step_number}|{state_hash}".encode('utf-8')
|
| signature = sign_artifact(self.signing_key, signing_payload)
|
|
|
| ckpt_file = self._get_checkpoint_file(agent_id)
|
| rec = ckpt_file.append(
|
| checkpoint_id=checkpoint_id,
|
| agent_id=agent_id,
|
| state=state_bytes,
|
| step=step_number,
|
| prev_id=previous_checkpoint
|
| )
|
|
|
| self._last_checkpoint[agent_id] = checkpoint_id
|
|
|
| return Checkpoint.from_record(rec)
|
|
|
| async def load_checkpoint(
|
| self,
|
| agent_id: str,
|
| checkpoint_id: str | None = None,
|
| step_number: int | None = None
|
| ) -> Checkpoint | None:
|
| """
|
| Load checkpoint by ID or step number.
|
| """
|
| ckpt_file = self._get_checkpoint_file(agent_id)
|
|
|
| if checkpoint_id:
|
| rec = ckpt_file.get_by_id(checkpoint_id)
|
| if rec:
|
| return Checkpoint.from_record(rec)
|
| return None
|
|
|
| if step_number is not None:
|
| rec = ckpt_file.get_by_step(step_number)
|
| if rec:
|
| return Checkpoint.from_record(rec)
|
| return None
|
|
|
| return None
|
|
|
| async def load_latest_checkpoint(self, agent_id: str) -> Checkpoint | None:
|
| """Load most recent checkpoint for agent."""
|
| ckpt_file = self._get_checkpoint_file(agent_id)
|
| rec = ckpt_file.get_latest()
|
| if rec:
|
| return Checkpoint.from_record(rec)
|
| return None
|
|
|
| async def list_checkpoints(
|
| self,
|
| agent_id: str,
|
| after_step: int | None = None
|
| ) -> list[Checkpoint]:
|
| """List all checkpoints for agent."""
|
| ckpt_file = self._get_checkpoint_file(agent_id)
|
| checkpoints = []
|
| for rec in ckpt_file.scan():
|
| if after_step is not None and rec.step_number <= after_step:
|
| continue
|
| checkpoints.append(Checkpoint.from_record(rec))
|
| return checkpoints
|
|
|
| async def verify_checkpoint_chain(self, agent_id: str) -> bool:
|
| """Verify cryptographic integrity of checkpoint chain."""
|
| checkpoints = await self.list_checkpoints(agent_id)
|
| if not checkpoints:
|
| return True
|
|
|
| expected_prev = ""
|
| for checkpoint in checkpoints:
|
| if (checkpoint.previous_checkpoint or "") != expected_prev:
|
| return False
|
| expected_prev = checkpoint.checkpoint_id
|
| return True
|
|
|
| async def get_checkpoint_stats(self, agent_id: str) -> dict[str, Any]:
|
| """Get checkpoint statistics for agent."""
|
| checkpoints = await self.list_checkpoints(agent_id)
|
| if not checkpoints:
|
| return {
|
| 'total_checkpoints': 0,
|
| 'first_checkpoint': None,
|
| 'last_checkpoint': None,
|
| 'chain_valid': True
|
| }
|
| return {
|
| 'total_checkpoints': len(checkpoints),
|
| 'first_checkpoint': checkpoints[0].checkpoint_id,
|
| 'last_checkpoint': checkpoints[-1].checkpoint_id,
|
| 'first_timestamp': checkpoints[0].timestamp,
|
| 'last_timestamp': checkpoints[-1].timestamp,
|
| 'step_range': (checkpoints[0].step_number, checkpoints[-1].step_number),
|
| 'chain_valid': await self.verify_checkpoint_chain(agent_id)
|
| }
|
|
|
|
|
|
|
|
|
|
|
|
|
| class CheckpointIterator:
|
| """Iterate through checkpoints in order."""
|
|
|
| def __init__(self, manager: CheckpointManager, agent_id: str):
|
| self.manager = manager
|
| self.agent_id = agent_id
|
| self._checkpoints: list[Checkpoint] | None = None
|
| self._index = 0
|
|
|
| async def __aiter__(self) -> AsyncIterator[Checkpoint]:
|
| if self._checkpoints is None:
|
| self._checkpoints = await self.manager.list_checkpoints(self.agent_id)
|
| for checkpoint in self._checkpoints:
|
| yield checkpoint
|
|
|
| async def next_checkpoint(self) -> Checkpoint | None:
|
| if self._checkpoints is None:
|
| self._checkpoints = await self.manager.list_checkpoints(self.agent_id)
|
| if self._index >= len(self._checkpoints):
|
| return None
|
| checkpoint = self._checkpoints[self._index]
|
| self._index += 1
|
| return checkpoint
|
|
|
| async def prev_checkpoint(self) -> Checkpoint | None:
|
| if self._checkpoints is None:
|
| self._checkpoints = await self.manager.list_checkpoints(self.agent_id)
|
| if self._index <= 0:
|
| return None
|
| self._index -= 1
|
| return self._checkpoints[self._index]
|
|
|
| def reset(self) -> None:
|
| self._index = 0
|
|
|
