script-build-and-framework-tmp/agent/agent/orchestration/store.py

"""Filesystem reference adapters for orchestration ports."""

from __future__ import annotations

import asyncio
import base64
import hashlib
import json
import os
import tempfile
from dataclasses import asdict
from pathlib import Path
from typing import Any, Dict, List, Optional, Sequence, Tuple

from filelock import FileLock

from .models import (
    ArtifactSnapshot,
    AttemptSubmission,
    BackgroundOperation,
    EventDraft,
    EventPage,
    OperationStatus,
    OrchestrationEvent,
    TaskLedger,
    json_values,
    utc_now,
)
from .protocols import CommitResult


class TaskStoreError(RuntimeError):
    pass


class TaskStoreNotFound(TaskStoreError):
    pass


class RevisionConflict(TaskStoreError):
    pass


class ArtifactConflict(TaskStoreError):
    pass


class FileSystemTaskStore:
    """Atomic JSON store with cross-process locking and optimistic revisions."""

    _EVENTS_KEY = "_events"
    _NEXT_EVENT_SEQUENCE_KEY = "_next_event_sequence"

    def __init__(self, base_path: str = ".trace") -> None:
        self.base_path = Path(base_path).resolve()

    def orchestration_dir(self, root_trace_id: str) -> Path:
        return _safe_child(self.base_path, root_trace_id, "orchestration")

    def ledger_path(self, root_trace_id: str) -> Path:
        return self.orchestration_dir(root_trace_id) / "ledger.json"

    def _lock_path(self, root_trace_id: str) -> Path:
        return self.orchestration_dir(root_trace_id) / ".ledger.lock"

    async def load(self, root_trace_id: str) -> TaskLedger:
        return await asyncio.to_thread(self._load_locked, root_trace_id)

    def _load_locked(self, root_trace_id: str) -> TaskLedger:
        lock_path = self._lock_path(root_trace_id)
        lock_path.parent.mkdir(parents=True, exist_ok=True)
        with FileLock(str(lock_path)):
            raw = self._read_raw(root_trace_id)
            try:
                return TaskLedger.from_dict(raw)
            except (AttributeError, KeyError, TypeError, ValueError) as exc:
                raise TaskStoreError(f"Cannot parse task ledger {self.ledger_path(root_trace_id)}: {exc}") from exc

    def _read_raw(self, root_trace_id: str) -> Dict[str, Any]:
        path = self.ledger_path(root_trace_id)
        if not path.exists():
            raise TaskStoreNotFound(f"No task ledger for root trace {root_trace_id}")
        try:
            data = json.loads(path.read_text(encoding="utf-8"))
            if not isinstance(data, dict):
                raise TypeError("ledger root must be an object")
            return data
        except (OSError, ValueError, TypeError, json.JSONDecodeError) as exc:
            raise TaskStoreError(f"Cannot load task ledger {path}: {exc}") from exc

    async def commit(
        self,
        ledger: TaskLedger,
        expected_revision: int,
        idempotency_key: Optional[str] = None,
        event: Optional[EventDraft] = None,
    ) -> CommitResult:
        # idempotency_key remains accepted for V1 callers. Durable command
        # replay is represented by TaskLedger.command_records in V2.
        del idempotency_key
        return await asyncio.to_thread(self._commit_locked, ledger, expected_revision, event)

    def _commit_locked(
        self,
        ledger: TaskLedger,
        expected_revision: int,
        event: Optional[EventDraft],
    ) -> CommitResult:
        root_trace_id = ledger.root_trace_id
        lock_path = self._lock_path(root_trace_id)
        lock_path.parent.mkdir(parents=True, exist_ok=True)
        with FileLock(str(lock_path)):
            raw: Dict[str, Any] = {}
            if self.ledger_path(root_trace_id).exists():
                raw = self._read_raw(root_trace_id)
            current_revision = int(raw.get("revision", -1))
            if current_revision != expected_revision:
                raise RevisionConflict(
                    f"Task ledger {root_trace_id} revision conflict: "
                    f"expected {expected_revision}, actual {current_revision}"
                )

            new_revision = expected_revision + 1
            updated_at = utc_now()
            data = ledger.to_dict()
            data["revision"] = new_revision
            data["updated_at"] = updated_at
            events = self._read_events(raw)
            next_sequence = self._next_sequence(raw, events)
            if event and not self._event_already_recorded(events, event):
                recorded = OrchestrationEvent(
                    schema_version=1,
                    event_id=_event_id(root_trace_id, next_sequence),
                    sequence=next_sequence,
                    root_trace_id=root_trace_id,
                    ledger_revision=new_revision,
                    event_type=event.event_type,
                    occurred_at=utc_now(),
                    payload=json_values(event.payload),
                    command_id=_public_command_id(event.command_id),
                    operation_id=event.operation_id,
                )
                events.append(json_values(asdict(recorded)))
                next_sequence += 1
            if events:
                data[self._EVENTS_KEY] = events
            data[self._NEXT_EVENT_SEQUENCE_KEY] = next_sequence

            path = self.ledger_path(root_trace_id)
            path.parent.mkdir(parents=True, exist_ok=True)
            self._atomic_json_write(path, data)
            ledger.revision = new_revision
            ledger.updated_at = updated_at
            return CommitResult(revision=ledger.revision, ledger=ledger)

    async def list_events(
        self,
        root_trace_id: str,
        cursor: Optional[str] = None,
        limit: int = 100,
    ) -> EventPage:
        if not 1 <= limit <= 1000:
            raise ValueError("Event page limit must be between 1 and 1000")
        after = _decode_cursor(cursor, root_trace_id) if cursor else 0
        return await asyncio.to_thread(self._list_events_locked, root_trace_id, after, limit, cursor)

    def _list_events_locked(
        self,
        root_trace_id: str,
        after: int,
        limit: int,
        prior_cursor: Optional[str],
    ) -> EventPage:
        lock_path = self._lock_path(root_trace_id)
        lock_path.parent.mkdir(parents=True, exist_ok=True)
        with FileLock(str(lock_path)):
            raw = self._read_raw(root_trace_id)
            try:
                available = [
                    OrchestrationEvent.from_dict(item)
                    for item in self._read_events(raw)
                    if int(item.get("sequence", 0)) > after
                ]
            except (KeyError, TypeError, ValueError) as exc:
                raise TaskStoreError("Ledger durable event journal contains an invalid event") from exc
        page = available[:limit]
        next_cursor = _encode_cursor(root_trace_id, page[-1].sequence) if page else prior_cursor
        return EventPage(events=page, next_cursor=next_cursor, has_more=len(available) > limit)

    async def list_recoverable(self, root_trace_id: str) -> List[BackgroundOperation]:
        ledger = await self.load(root_trace_id)
        recoverable = {
            OperationStatus.PENDING,
            OperationStatus.RUNNING,
            OperationStatus.STOP_REQUESTED,
        }
        return [operation for operation in ledger.operations.values() if operation.status in recoverable]

    @classmethod
    def _read_events(cls, raw: Dict[str, Any]) -> List[Dict[str, Any]]:
        events = raw.get(cls._EVENTS_KEY, [])
        if not isinstance(events, list) or any(not isinstance(item, dict) for item in events):
            raise TaskStoreError("Ledger durable event journal is invalid")
        return list(events)

    @classmethod
    def _next_sequence(cls, raw: Dict[str, Any], events: Sequence[Dict[str, Any]]) -> int:
        fallback = max((int(item.get("sequence", 0)) for item in events), default=0) + 1
        value = int(raw.get(cls._NEXT_EVENT_SEQUENCE_KEY, fallback))
        return max(value, fallback, 1)

    @staticmethod
    def _event_already_recorded(events: Sequence[Dict[str, Any]], event: EventDraft) -> bool:
        command_id = _public_command_id(event.command_id)
        return bool(command_id) and any(item.get("command_id") == command_id for item in events)

    @staticmethod
    def _atomic_json_write(path: Path, data: Dict[str, Any]) -> None:
        fd, tmp_name = tempfile.mkstemp(prefix=f".{path.name}.", suffix=".tmp", dir=str(path.parent))
        try:
            with os.fdopen(fd, "w", encoding="utf-8") as handle:
                json.dump(data, handle, ensure_ascii=False, indent=2, sort_keys=True)
                handle.flush()
                os.fsync(handle.fileno())
            os.replace(tmp_name, path)
            _fsync_directory(path.parent)
        except Exception:
            try:
                os.unlink(tmp_name)
            except OSError:
                pass
            raise


class FileSystemArtifactStore:
    """Immutable snapshots with same-attempt first-write-wins semantics."""

    def __init__(self, base_path: str = ".trace", root_trace_id: Optional[str] = None) -> None:
        self.base_path = Path(base_path).resolve()
        self.root_trace_id = root_trace_id

    def for_root(self, root_trace_id: str) -> "FileSystemArtifactStore":
        """V1 compatibility adapter; new callers should pass root_trace_id."""
        return FileSystemArtifactStore(str(self.base_path), root_trace_id=root_trace_id)

    def _artifacts_dir(self, root_trace_id: str) -> Path:
        return _safe_child(self.base_path, root_trace_id, "orchestration", "artifacts")

    def _bound_root(self) -> str:
        if not self.root_trace_id:
            raise ValueError("root_trace_id is required")
        return self.root_trace_id

    async def freeze(
        self,
        root_trace_id: str,
        attempt_id: Any,
        submission: Optional[AttemptSubmission] = None,
    ) -> ArtifactSnapshot:
        if submission is None:
            # V1 bound form: bound.freeze(attempt_id, submission)
            submission = attempt_id
            attempt_id = root_trace_id
            root_trace_id = self._bound_root()
        if not isinstance(attempt_id, str) or not isinstance(submission, AttemptSubmission):
            raise TypeError("freeze requires root_trace_id, attempt_id and AttemptSubmission")
        return await asyncio.to_thread(self._freeze_locked, root_trace_id, attempt_id, submission)

    def _freeze_locked(
        self,
        root_trace_id: str,
        attempt_id: str,
        submission: AttemptSubmission,
    ) -> ArtifactSnapshot:
        artifacts_dir = self._artifacts_dir(root_trace_id)
        artifacts_dir.mkdir(parents=True, exist_ok=True)
        with FileLock(str(artifacts_dir / ".artifacts.lock")):
            existing = self._find_for_attempt(artifacts_dir, attempt_id)
            normalized, digest = _normalize_submission(submission)
            if existing:
                if existing.sha256 == digest:
                    return existing
                raise ArtifactConflict(f"Attempt {attempt_id} already has a different artifact snapshot")

            snapshot = ArtifactSnapshot(
                snapshot_id=_snapshot_id(attempt_id),
                attempt_id=attempt_id,
                normalized_content=normalized,
                sha256=digest,
                artifact_refs=list(submission.artifact_refs),
                evidence_refs=list(submission.evidence_refs),
            )
            FileSystemTaskStore._atomic_json_write(
                artifacts_dir / f"{snapshot.snapshot_id}.json",
                json_values(asdict(snapshot)),
            )
            return snapshot

    async def get(self, root_trace_id: str, snapshot_id: Optional[str] = None) -> ArtifactSnapshot:
        if snapshot_id is None:
            snapshot_id = root_trace_id
            root_trace_id = self._bound_root()
        return await asyncio.to_thread(self._get_locked, root_trace_id, snapshot_id)

    def _get_locked(self, root_trace_id: str, snapshot_id: str) -> ArtifactSnapshot:
        artifacts_dir = self._artifacts_dir(root_trace_id)
        path = _safe_child(artifacts_dir, f"{snapshot_id}.json")
        artifacts_dir.mkdir(parents=True, exist_ok=True)
        with FileLock(str(artifacts_dir / ".artifacts.lock")):
            if not path.exists():
                raise FileNotFoundError(f"Artifact snapshot not found: {snapshot_id}")
            return self._read_snapshot(path)

    async def get_for_attempt(
        self,
        root_trace_id: str,
        attempt_id: Optional[str] = None,
    ) -> Optional[ArtifactSnapshot]:
        if attempt_id is None:
            attempt_id = root_trace_id
            root_trace_id = self._bound_root()
        return await asyncio.to_thread(self._get_for_attempt_locked, root_trace_id, attempt_id)

    def _get_for_attempt_locked(self, root_trace_id: str, attempt_id: str) -> Optional[ArtifactSnapshot]:
        artifacts_dir = self._artifacts_dir(root_trace_id)
        artifacts_dir.mkdir(parents=True, exist_ok=True)
        with FileLock(str(artifacts_dir / ".artifacts.lock")):
            return self._find_for_attempt(artifacts_dir, attempt_id)

    async def list_orphans(
        self,
        root_trace_id: str,
        known_attempt_ids: Sequence[str],
    ) -> List[ArtifactSnapshot]:
        artifacts_dir = self._artifacts_dir(root_trace_id)
        return await asyncio.to_thread(self._list_orphans_locked, artifacts_dir, set(known_attempt_ids))

    def _list_orphans_locked(self, artifacts_dir: Path, known: set[str]) -> List[ArtifactSnapshot]:
        artifacts_dir.mkdir(parents=True, exist_ok=True)
        with FileLock(str(artifacts_dir / ".artifacts.lock")):
            return [snapshot for snapshot in self._read_all(artifacts_dir) if snapshot.attempt_id not in known]

    async def cleanup_orphans(
        self,
        root_trace_id: str,
        known_attempt_ids: Sequence[str],
    ) -> List[str]:
        artifacts_dir = self._artifacts_dir(root_trace_id)
        return await asyncio.to_thread(self._cleanup_orphans_locked, artifacts_dir, set(known_attempt_ids))

    def _cleanup_orphans_locked(self, artifacts_dir: Path, known: set[str]) -> List[str]:
        artifacts_dir.mkdir(parents=True, exist_ok=True)
        with FileLock(str(artifacts_dir / ".artifacts.lock")):
            orphans = [
                (path, self._read_snapshot(path))
                for path in sorted(artifacts_dir.glob("*.json"))
            ]
            orphans = [(path, snapshot) for path, snapshot in orphans if snapshot.attempt_id not in known]
            for path, _snapshot in orphans:
                # Delete only the directory entry that was actually scanned;
                # snapshot JSON is untrusted and must not select a path.
                path.unlink()
            if orphans:
                _fsync_directory(artifacts_dir)
            return [snapshot.snapshot_id for _path, snapshot in orphans]

    def _find_for_attempt(self, artifacts_dir: Path, attempt_id: str) -> Optional[ArtifactSnapshot]:
        return next((item for item in self._read_all(artifacts_dir) if item.attempt_id == attempt_id), None)

    def _read_all(self, artifacts_dir: Path) -> List[ArtifactSnapshot]:
        return [self._read_snapshot(path) for path in sorted(artifacts_dir.glob("*.json"))]

    @staticmethod
    def _read_snapshot(path: Path) -> ArtifactSnapshot:
        try:
            return ArtifactSnapshot.from_dict(json.loads(path.read_text(encoding="utf-8")))
        except (OSError, ValueError, TypeError, json.JSONDecodeError) as exc:
            raise TaskStoreError(f"Cannot load artifact snapshot {path}: {exc}") from exc


class TraceEventSink:
    """Best-effort legacy event mirror; durable events live in TaskStore."""

    def __init__(self, base_path: str = ".trace") -> None:
        self.base_path = Path(base_path).resolve()

    async def emit(self, root_trace_id: str, event_type: str, payload: Dict[str, Any]) -> None:
        await asyncio.to_thread(self._emit_locked, root_trace_id, event_type, payload)

    def _emit_locked(self, root_trace_id: str, event_type: str, payload: Dict[str, Any]) -> None:
        directory = _safe_child(self.base_path, root_trace_id, "orchestration")
        directory.mkdir(parents=True, exist_ok=True)
        path = directory / "events.jsonl"
        event = {
            "event_id": _event_id(root_trace_id, os.urandom(8).hex()),
            "event_type": event_type,
            "root_trace_id": root_trace_id,
            "created_at": utc_now(),
            "payload": json_values(payload),
        }
        with FileLock(str(directory / ".events.lock")):
            with path.open("a", encoding="utf-8") as handle:
                handle.write(json.dumps(event, ensure_ascii=False, sort_keys=True) + "\n")
                handle.flush()


def _safe_child(base: Path, *parts: str) -> Path:
    if any(not isinstance(part, str) or not part for part in parts):
        raise ValueError("Path identifiers must be non-empty strings")
    candidate = base.joinpath(*parts).resolve()
    try:
        candidate.relative_to(base.resolve())
    except ValueError as exc:
        raise ValueError("Path identifier escapes the configured base path") from exc
    return candidate


def _normalize_submission(submission: AttemptSubmission) -> Tuple[Dict[str, Any], str]:
    normalized = {
        "summary": submission.summary.strip(),
        "artifact_refs": [json_values(asdict(item)) for item in submission.artifact_refs],
        "evidence_refs": [json_values(asdict(item)) for item in submission.evidence_refs],
    }
    canonical = json.dumps(normalized, ensure_ascii=False, sort_keys=True, separators=(",", ":"))
    return normalized, hashlib.sha256(canonical.encode("utf-8")).hexdigest()


def _snapshot_id(attempt_id: str) -> str:
    return "attempt-" + hashlib.sha256(attempt_id.encode("utf-8")).hexdigest()


def _event_id(root_trace_id: str, sequence: Any) -> str:
    source = f"{root_trace_id}:{sequence}".encode("utf-8")
    return hashlib.sha256(source).hexdigest()


def _public_command_id(command_id: Optional[str]) -> Optional[str]:
    if command_id is None:
        return None
    return "cmd_" + hashlib.sha256(command_id.encode("utf-8")).hexdigest()


def _encode_cursor(root_trace_id: str, sequence: int) -> str:
    root_digest = hashlib.sha256(root_trace_id.encode("utf-8")).hexdigest()[:16]
    raw = json.dumps({"v": 1, "root": root_digest, "after": sequence}, separators=(",", ":"))
    return base64.urlsafe_b64encode(raw.encode("utf-8")).decode("ascii").rstrip("=")


def _decode_cursor(cursor: str, root_trace_id: str) -> int:
    try:
        padding = "=" * (-len(cursor) % 4)
        data = json.loads(base64.urlsafe_b64decode(cursor + padding).decode("utf-8"))
        expected = hashlib.sha256(root_trace_id.encode("utf-8")).hexdigest()[:16]
        if data.get("v") != 1 or data.get("root") != expected:
            raise ValueError
        sequence = int(data["after"])
        if sequence < 0:
            raise ValueError
        return sequence
    except (KeyError, TypeError, ValueError, UnicodeDecodeError, json.JSONDecodeError) as exc:
        raise ValueError("Invalid event cursor") from exc


def _fsync_directory(path: Path) -> None:
    try:
        descriptor = os.open(path, os.O_RDONLY)
    except OSError:
        return
    try:
        os.fsync(descriptor)
    finally:
        os.close(descriptor)


__all__ = [
    "TaskStoreError", "TaskStoreNotFound", "RevisionConflict", "ArtifactConflict",
    "FileSystemTaskStore", "FileSystemArtifactStore", "TraceEventSink",
]