From 31538812377b4cf9c2e89ee6b50ad00c410b8015 Mon Sep 17 00:00:00 2001 From: binrogithub Date: Sun, 8 Feb 2026 17:59:51 +0300 Subject: [PATCH] Add run store with atomic writes --- hc_agent/core/run_store.py | 216 +++++++++++++++++++++++++++++++++++++ 1 file changed, 216 insertions(+) create mode 100644 hc_agent/core/run_store.py diff --git a/hc_agent/core/run_store.py b/hc_agent/core/run_store.py new file mode 100644 index 0000000000..707dd93d74 --- /dev/null +++ b/hc_agent/core/run_store.py @@ -0,0 +1,216 @@ +"""Run storage utilities with atomic writes and index maintenance.""" + +from __future__ import annotations + +import json +import os +import shutil +import uuid +from contextlib import contextmanager +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Dict, Iterable, List, Optional + + +@dataclass +class RunIndexEntry: + run_id: str + tenant_id: Optional[str] + created_at: str + stage: Optional[str] + ok: Optional[bool] + + +class RunStore: + def __init__(self, root_dir: str | Path) -> None: + self.root_dir = Path(root_dir) + self.runs_dir = self.root_dir / "runs" + self.index_path = self.runs_dir / "index.json" + self.lock_path = self.runs_dir / ".index.lock" + + def write_json( + self, + run_id: str, + relpath: str, + payload: Any, + *, + tenant_id: Optional[str] = None, + stage: Optional[str] = None, + ok: Optional[bool] = None, + created_at: Optional[str] = None, + ) -> Path: + target = self._run_path(run_id, relpath) + serialized = json.dumps(payload, indent=2, sort_keys=True) + self._write_text_atomic(target, serialized) + self._update_index( + run_id, + tenant_id=tenant_id, + stage=stage, + ok=ok, + created_at=created_at, + ) + return target + + def write_text( + self, + run_id: str, + relpath: str, + content: str, + *, + tenant_id: Optional[str] = None, + stage: Optional[str] = None, + ok: Optional[bool] = None, + created_at: Optional[str] = None, + ) -> Path: + target = self._run_path(run_id, relpath) + self._write_text_atomic(target, content) + self._update_index( + run_id, + tenant_id=tenant_id, + stage=stage, + ok=ok, + created_at=created_at, + ) + return target + + def read_json(self, run_id: str, relpath: str) -> Any: + target = self._run_path(run_id, relpath) + with target.open("r", encoding="utf-8") as handle: + return json.load(handle) + + def read_text(self, run_id: str, relpath: str) -> str: + target = self._run_path(run_id, relpath) + return target.read_text(encoding="utf-8") + + def delete_run(self, run_id: str) -> None: + run_dir = self.runs_dir / run_id + with self._index_lock(): + if run_dir.exists(): + shutil.rmtree(run_dir) + self._remove_from_index({run_id}) + + def gc(self, max_age_seconds: float) -> List[str]: + now = datetime.now(timezone.utc) + threshold = now.timestamp() - max_age_seconds + removed: List[str] = [] + with self._index_lock(): + entries = self._load_index_entries() + for entry in entries: + created_at = self._parse_time(entry.get("created_at")) + if created_at is None: + continue + if created_at.timestamp() < threshold: + run_id = entry.get("run_id") + if run_id: + run_dir = self.runs_dir / run_id + if run_dir.exists(): + shutil.rmtree(run_dir) + removed.append(run_id) + if removed: + self._remove_from_index(set(removed)) + return removed + + def _run_path(self, run_id: str, relpath: str) -> Path: + target = self.runs_dir / run_id / relpath + target.parent.mkdir(parents=True, exist_ok=True) + return target + + def _write_text_atomic(self, target: Path, content: str) -> None: + data = content.encode("utf-8") + self._write_bytes_atomic(target, data) + + def _write_bytes_atomic(self, target: Path, data: bytes) -> None: + target.parent.mkdir(parents=True, exist_ok=True) + temp_name = f".{target.name}.tmp.{uuid.uuid4().hex}" + temp_path = target.with_name(temp_name) + with temp_path.open("wb") as handle: + handle.write(data) + handle.flush() + os.fsync(handle.fileno()) + os.replace(temp_path, target) + + def _update_index( + self, + run_id: str, + *, + tenant_id: Optional[str], + stage: Optional[str], + ok: Optional[bool], + created_at: Optional[str], + ) -> None: + with self._index_lock(): + entries = self._load_index_entries() + entry_map = {entry["run_id"]: entry for entry in entries if "run_id" in entry} + entry = entry_map.get(run_id, {}) + entry["run_id"] = run_id + if tenant_id is not None: + entry["tenant_id"] = tenant_id + if stage is not None: + entry["stage"] = stage + if ok is not None: + entry["ok"] = ok + if created_at is not None: + entry["created_at"] = created_at + else: + entry.setdefault("created_at", self._now_iso()) + entry.setdefault("tenant_id", None) + entry.setdefault("stage", None) + entry.setdefault("ok", None) + entry_map[run_id] = entry + self._write_index_entries(entry_map.values()) + + def _remove_from_index(self, run_ids: set[str]) -> None: + entries = self._load_index_entries() + remaining = [entry for entry in entries if entry.get("run_id") not in run_ids] + self._write_index_entries(remaining) + + def _write_index_entries(self, entries: Iterable[Dict[str, Any]]) -> None: + self.runs_dir.mkdir(parents=True, exist_ok=True) + normalized = sorted(entries, key=lambda entry: entry.get("created_at", "")) + data = json.dumps(normalized, indent=2, sort_keys=True) + self._write_text_atomic(self.index_path, data) + + def _load_index_entries(self) -> List[Dict[str, Any]]: + if not self.index_path.exists(): + return [] + try: + with self.index_path.open("r", encoding="utf-8") as handle: + data = json.load(handle) + except json.JSONDecodeError: + return [] + if isinstance(data, list): + return [entry for entry in data if isinstance(entry, dict)] + return [] + + def _now_iso(self) -> str: + return datetime.now(timezone.utc).isoformat() + + def _parse_time(self, value: Optional[str]) -> Optional[datetime]: + if not value: + return None + try: + parsed = datetime.fromisoformat(value) + except ValueError: + return None + if parsed.tzinfo is None: + return parsed.replace(tzinfo=timezone.utc) + return parsed.astimezone(timezone.utc) + + @contextmanager + def _index_lock(self) -> Iterable[None]: + self.runs_dir.mkdir(parents=True, exist_ok=True) + fd: Optional[int] = None + try: + fd = os.open(self.lock_path, os.O_CREAT | os.O_EXCL | os.O_WRONLY) + except FileExistsError: + fd = None + try: + yield + finally: + if fd is not None: + os.close(fd) + try: + os.unlink(self.lock_path) + except FileNotFoundError: + pass