ThothII/harness/tht/corpus/store.py

"""Durable immutable corpus generations and an atomic ACTIVE pointer."""

from __future__ import annotations

import json
import fcntl
import os
import re
import stat
import shutil
import uuid
from pathlib import Path
from contextlib import contextmanager

from tht.corpus.models import CorpusManifest


_GENERATION = re.compile(r"^gen:[0-9a-f]{32}$")


class UnsafeCorpusPath(RuntimeError):
    pass


def _atomic_write(path: Path, payload: bytes) -> None:
    temporary = path.with_name(f".{path.name}.{uuid.uuid4().hex}.tmp")
    fd = os.open(temporary, os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW, 0o600)
    try:
        with os.fdopen(fd, "wb") as stream:
            stream.write(payload)
            stream.flush()
            os.fsync(stream.fileno())
        os.replace(temporary, path)
        directory = os.open(path.parent, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW)
        try:
            os.fsync(directory)
        finally:
            os.close(directory)
    except BaseException:
        temporary.unlink(missing_ok=True)
        raise


class CorpusStore:
    def __init__(self, root: Path) -> None:
        self.root = Path(root)
        self.active_path = self.root / "ACTIVE"
        self._replace = os.replace
        self._ensure_root()

    def _ensure_root(self) -> None:
        if self.root.is_symlink():
            raise UnsafeCorpusPath("corpus root must not be a symlink")
        self.root.mkdir(parents=True, exist_ok=True, mode=0o700)
        info = self.root.lstat()
        if not stat.S_ISDIR(info.st_mode) or info.st_uid != os.getuid():
            raise UnsafeCorpusPath("corpus root is unsafe")

    @contextmanager
    def writer_lock(self):
        lock_path = self.root / ".writer.lock"
        fd = os.open(lock_path, os.O_RDWR | os.O_CREAT | os.O_NOFOLLOW | os.O_CLOEXEC, 0o600)
        try:
            info = os.fstat(fd)
            if not stat.S_ISREG(info.st_mode) or info.st_uid != os.getuid() or info.st_nlink != 1:
                raise UnsafeCorpusPath("corpus writer lock is unsafe")
            fcntl.flock(fd, fcntl.LOCK_EX)
            yield
        finally:
            fcntl.flock(fd, fcntl.LOCK_UN)
            os.close(fd)

    def generation_path(self, generation: str) -> Path:
        if not _GENERATION.fullmatch(generation):
            raise UnsafeCorpusPath("invalid corpus generation")
        path = self.root / generation.replace(":", "-")
        if path.is_symlink():
            raise UnsafeCorpusPath("generation must not be a symlink")
        return path

    def stage(
        self,
        manifest: CorpusManifest,
        materialized: dict[str, str],
        *,
        generation: str | None = None,
    ) -> str:
        generation = generation or f"gen:{uuid.uuid4().hex}"
        path = self.generation_path(generation)
        try:
            path.mkdir(mode=0o700)
        except FileExistsError:
            raise UnsafeCorpusPath("generation already exists") from None
        documents = path / "documents"
        documents.mkdir(mode=0o700)
        files: dict[str, str] = {}
        for document in manifest.documents:
            relative = f"documents/{document.document_id.removeprefix('doc:')}.md"
            _atomic_write(path / relative, materialized[document.document_id].encode("utf-8"))
            files[document.document_id] = relative
        payload = json.loads(manifest.model_dump_json())
        metadata = payload["metadata"]
        metadata["files"] = files
        payload.update({"manifest_id": generation, "metadata": metadata})
        staged = CorpusManifest.model_validate(payload)
        _atomic_write(path / "manifest.json", (staged.model_dump_json(indent=2) + "\n").encode())
        return generation

    def publish(self, generation: str) -> str:
        manifest = self.manifest(generation)
        if manifest.manifest_id != generation:
            raise UnsafeCorpusPath("manifest generation mismatch")
        temporary = self.active_path.with_name(f".ACTIVE.{uuid.uuid4().hex}.tmp")
        _atomic_write(temporary, (generation + "\n").encode())
        self._replace(temporary, self.active_path)
        directory = os.open(self.root, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW)
        try:
            os.fsync(directory)
        finally:
            os.close(directory)
        return generation

    def active_generation(self) -> str | None:
        try:
            if self.active_path.is_symlink():
                raise UnsafeCorpusPath("ACTIVE must not be a symlink")
            value = self.active_path.read_text(encoding="ascii").strip()
        except FileNotFoundError:
            return None
        if not _GENERATION.fullmatch(value):
            raise UnsafeCorpusPath("ACTIVE contains an invalid generation")
        return value

    def manifest(self, generation: str) -> CorpusManifest:
        path = self.generation_path(generation)
        manifest_path = path / "manifest.json"
        if manifest_path.is_symlink():
            raise UnsafeCorpusPath("manifest must not be a symlink")
        return CorpusManifest.model_validate_json(manifest_path.read_text(encoding="utf-8"))

    def discard(self, generation: str) -> None:
        path = self.generation_path(generation)
        if path.exists():
            if path.is_symlink() or not stat.S_ISDIR(path.lstat().st_mode):
                raise UnsafeCorpusPath("generation cleanup target is unsafe")
            shutil.rmtree(path)

    def active_manifest(self) -> CorpusManifest | None:
        generation = self.active_generation()
        return self.manifest(generation) if generation else None

    def resolve_document(self, document_id: str, generation: str | None = None) -> Path | None:
        generation = generation or self.active_generation()
        if generation is None:
            return None
        manifest = self.manifest(generation)
        relative = manifest.metadata.get("files", {}).get(document_id)
        if not isinstance(relative, str):
            return None
        base = self.generation_path(generation).resolve()
        path = (base / relative).resolve()
        if not path.is_relative_to(base) or path.is_symlink():
            raise UnsafeCorpusPath("materialized document escapes its generation")
        return path