Nexum-Expanded/runtime/src/nexum_runtime/tooling/drafting.py

"""Receipt-backed draft, comparison, and selection workflows."""

from __future__ import annotations

import hashlib
import hmac
import json
import os
import secrets
import time
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any, Mapping

from .artifacts import ArtifactStore
from .contracts import (
    ToolCall,
    ToolExecutionContext,
    ToolExecutionResult,
    ToolParameter,
    ToolSpec,
)
from .sandbox import control_root, ensure_control_root
from .security import SecretRedactor


_DRAFT_PREFIX = "draft_"
_HEX = frozenset("0123456789abcdef")


DRAFTING_TOOL_SPECS: tuple[ToolSpec, ...] = (
    ToolSpec(
        "DraftCreate",
        "drafting",
        "Persist model-selected candidate actions or answers before any external effect is taken.",
        "DraftCreate(objective='Repair failing test', candidates=[{'name':'minimal','content':'...'}], criteria=['passes current test'])",
        (
            ToolParameter("objective", "string", "Current objective being drafted."),
            ToolParameter(
                "candidates",
                "array",
                "Candidate actions or answers. Each row must include name and content.",
            ),
            ToolParameter(
                "criteria",
                "array",
                "Evidence criteria the model will use to compare candidates.",
            ),
        ),
        risk="workspace_write",
        parallel_safe=False,
        idempotent=False,
    ),
    ToolSpec(
        "DraftCompare",
        "drafting",
        "Attach evidence and observations to a draft without choosing for the model.",
        "DraftCompare(draft_id='draft_...', evidence_refs=['rep_...'], observations=['candidate 0 preserves API'])",
        (
            ToolParameter("draft_id", "string", "Draft receipt identifier."),
            ToolParameter(
                "evidence_refs",
                "array",
                "Optional receipt or artifact identifiers considered by the model.",
                required=False,
            ),
            ToolParameter(
                "observations",
                "array",
                "Model-observed comparison facts grounded in current evidence.",
                required=False,
            ),
        ),
        risk="workspace_write",
        parallel_safe=False,
        idempotent=False,
    ),
    ToolSpec(
        "DraftSelect",
        "drafting",
        "Persist the model's chosen draft index with cited evidence before execution or final answer.",
        "DraftSelect(draft_id='draft_...', selected_index=0, decision='smallest verified change', evidence_refs=['rep_...'])",
        (
            ToolParameter("draft_id", "string", "Draft receipt identifier."),
            ToolParameter("selected_index", "integer", "Zero-based selected candidate index."),
            ToolParameter("decision", "string", "Why this candidate was selected."),
            ToolParameter(
                "evidence_refs",
                "array",
                "Optional receipt or artifact identifiers supporting the selection.",
                required=False,
            ),
        ),
        risk="workspace_write",
        parallel_safe=False,
        idempotent=False,
    ),
    ToolSpec(
        "DraftStatus",
        "drafting",
        "Read a session-scoped draft receipt, comparisons, and final selection state.",
        "DraftStatus(draft_id='draft_...')",
        (ToolParameter("draft_id", "string", "Draft receipt identifier."),),
    ),
)
DRAFTING_TOOL_NAMES = frozenset(spec.name for spec in DRAFTING_TOOL_SPECS)
_DRAFTING_TOOL_BY_NAME = {spec.name: spec for spec in DRAFTING_TOOL_SPECS}


@dataclass(frozen=True)
class DraftCandidate:
    index: int
    name: str
    content_sha256: str
    artifact_id: str
    metadata_sha256: str


@dataclass(frozen=True)
class DraftComparison:
    comparison_id: str
    evidence_refs: tuple[str, ...]
    observations: tuple[str, ...]
    evidence_sha256: str
    created_unix_ms: int


@dataclass(frozen=True)
class DraftSelection:
    selected_index: int
    decision: str
    evidence_refs: tuple[str, ...]
    evidence_sha256: str
    created_unix_ms: int


@dataclass(frozen=True)
class DraftRecord:
    draft_id: str
    session_sha256: str
    objective: str
    criteria: tuple[str, ...]
    candidates: tuple[DraftCandidate, ...]
    comparisons: tuple[DraftComparison, ...]
    selection: DraftSelection | None
    created_unix_ms: int
    updated_unix_ms: int

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


def _session_sha256(session_id: str) -> str:
    return hashlib.sha256((session_id or "direct").encode("utf-8")).hexdigest()


def _require_session(session_sha256: str, session_id: str) -> None:
    if not hmac.compare_digest(session_sha256, _session_sha256(session_id)):
        raise PermissionError("draft receipt does not belong to this session")


def _sanitize_text(value: object, workspace: Path) -> str:
    redacted = SecretRedactor().redact(str(value or ""))
    return redacted.replace(str(workspace), "[WORKSPACE]")


def _string_rows(value: object, name: str, *, allow_empty: bool) -> tuple[str, ...]:
    if value is None and allow_empty:
        return ()
    if not isinstance(value, list) or any(not isinstance(item, str) for item in value):
        raise ValueError(f"{name} must be an array of strings")
    rows = tuple(item.strip() for item in value)
    if any(not item for item in rows) or (not rows and not allow_empty):
        raise ValueError(f"{name} must contain non-empty strings")
    return rows


def _draft_path(root: Path, draft_id: str) -> Path:
    suffix = draft_id.removeprefix(_DRAFT_PREFIX)
    if (
        not draft_id.startswith(_DRAFT_PREFIX)
        or len(suffix) != 32
        or any(char not in _HEX for char in suffix)
    ):
        raise ValueError("draft id is invalid")
    return root / f"{draft_id}.json"


def _atomic_json(path: Path, payload: Mapping[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
    os.chmod(path.parent, 0o700)
    temporary = path.with_name(f".{path.name}.{secrets.token_hex(8)}.tmp")
    temporary.write_text(
        json.dumps(payload, sort_keys=True, separators=(",", ":")),
        encoding="utf-8",
    )
    temporary.chmod(0o600)
    with temporary.open("rb") as handle:
        os.fsync(handle.fileno())
    os.replace(temporary, path)
    directory_fd = os.open(path.parent, os.O_RDONLY)
    try:
        os.fsync(directory_fd)
    finally:
        os.close(directory_fd)


def _load_record(path: Path) -> DraftRecord:
    payload = json.loads(path.read_text(encoding="utf-8"))
    if not isinstance(payload, dict):
        raise RuntimeError("draft receipt is invalid")
    candidates = tuple(DraftCandidate(**row) for row in payload.pop("candidates", []))
    comparisons = tuple(
        DraftComparison(
            **{
                **row,
                "evidence_refs": tuple(row.get("evidence_refs", [])),
                "observations": tuple(row.get("observations", [])),
            }
        )
        for row in payload.pop("comparisons", [])
    )
    raw_selection = payload.pop("selection", None)
    selection = (
        DraftSelection(
            **{
                **raw_selection,
                "evidence_refs": tuple(raw_selection.get("evidence_refs", [])),
            }
        )
        if isinstance(raw_selection, dict)
        else None
    )
    return DraftRecord(
        **{
            **payload,
            "criteria": tuple(payload.get("criteria", [])),
            "candidates": candidates,
            "comparisons": comparisons,
            "selection": selection,
        }
    )


class DraftStore:
    def __init__(self, workspace: str | Path) -> None:
        self.workspace = ensure_control_root(workspace)
        self.root = control_root(self.workspace) / "drafting"
        self.root.mkdir(parents=True, exist_ok=True, mode=0o700)
        os.chmod(self.root, 0o700)
        self.artifacts = ArtifactStore(self.workspace)

    def create(
        self,
        *,
        objective: str,
        candidates: object,
        criteria: tuple[str, ...],
        session_id: str,
    ) -> DraftRecord:
        objective_text = _sanitize_text(objective, self.workspace).strip()
        if not objective_text:
            raise ValueError("draft objective is required")
        if not isinstance(candidates, list) or not candidates:
            raise ValueError("draft candidates must be a non-empty array")
        rows: list[DraftCandidate] = []
        for index, raw_candidate in enumerate(candidates):
            if not isinstance(raw_candidate, dict):
                raise ValueError("each draft candidate must be an object")
            name = _sanitize_text(raw_candidate.get("name"), self.workspace).strip()
            content = _sanitize_text(raw_candidate.get("content"), self.workspace)
            if not name or not content.strip():
                raise ValueError("each draft candidate requires name and content")
            metadata = {
                key: _sanitize_text(value, self.workspace)
                for key, value in raw_candidate.items()
                if key not in {"name", "content"}
            }
            artifact = self.artifacts.put_text(
                content,
                media_type="text/plain; charset=utf-8",
                source="draft_candidate",
                session_id=session_id or "direct",
            )
            rows.append(
                DraftCandidate(
                    index=index,
                    name=name,
                    content_sha256=hashlib.sha256(content.encode("utf-8")).hexdigest(),
                    artifact_id=artifact.artifact_id,
                    metadata_sha256=hashlib.sha256(
                        json.dumps(
                            metadata,
                            sort_keys=True,
                            separators=(",", ":"),
                        ).encode("utf-8")
                    ).hexdigest(),
                )
            )
        if not criteria:
            raise ValueError("draft criteria must contain at least one entry")
        now = int(time.time() * 1000)
        record = DraftRecord(
            draft_id=_DRAFT_PREFIX + secrets.token_hex(16),
            session_sha256=_session_sha256(session_id),
            objective=objective_text,
            criteria=tuple(_sanitize_text(item, self.workspace) for item in criteria),
            candidates=tuple(rows),
            comparisons=(),
            selection=None,
            created_unix_ms=now,
            updated_unix_ms=now,
        )
        _atomic_json(_draft_path(self.root, record.draft_id), record.to_dict())
        return record

    def get(self, draft_id: str, *, session_id: str) -> DraftRecord:
        record = _load_record(_draft_path(self.root, draft_id))
        _require_session(record.session_sha256, session_id)
        return record

    def compare(
        self,
        draft_id: str,
        *,
        evidence_refs: tuple[str, ...],
        observations: tuple[str, ...],
        session_id: str,
    ) -> DraftRecord:
        record = self.get(draft_id, session_id=session_id)
        if not evidence_refs and not observations:
            raise ValueError("draft comparison requires evidence refs or observations")
        safe_evidence = tuple(_sanitize_text(item, self.workspace) for item in evidence_refs)
        safe_observations = tuple(_sanitize_text(item, self.workspace) for item in observations)
        evidence_sha256 = hashlib.sha256(
            json.dumps(
                {"evidence_refs": safe_evidence, "observations": safe_observations},
                sort_keys=True,
                separators=(",", ":"),
            ).encode("utf-8")
        ).hexdigest()
        comparison = DraftComparison(
            comparison_id="cmp_" + secrets.token_hex(16),
            evidence_refs=safe_evidence,
            observations=safe_observations,
            evidence_sha256=evidence_sha256,
            created_unix_ms=int(time.time() * 1000),
        )
        updated = DraftRecord(
            **{
                **record.to_dict(),
                "comparisons": (*record.comparisons, comparison),
                "updated_unix_ms": int(time.time() * 1000),
            }
        )
        _atomic_json(_draft_path(self.root, draft_id), updated.to_dict())
        return updated

    def select(
        self,
        draft_id: str,
        *,
        selected_index: int,
        decision: str,
        evidence_refs: tuple[str, ...],
        session_id: str,
    ) -> DraftRecord:
        record = self.get(draft_id, session_id=session_id)
        if selected_index < 0 or selected_index >= len(record.candidates):
            raise ValueError("selected draft index is outside the candidate set")
        decision_text = _sanitize_text(decision, self.workspace).strip()
        if not decision_text:
            raise ValueError("draft selection decision is required")
        safe_evidence = tuple(_sanitize_text(item, self.workspace) for item in evidence_refs)
        evidence_sha256 = hashlib.sha256(
            json.dumps(
                {
                    "selected_index": selected_index,
                    "decision": decision_text,
                    "evidence_refs": safe_evidence,
                },
                sort_keys=True,
                separators=(",", ":"),
            ).encode("utf-8")
        ).hexdigest()
        selection = DraftSelection(
            selected_index=selected_index,
            decision=decision_text,
            evidence_refs=safe_evidence,
            evidence_sha256=evidence_sha256,
            created_unix_ms=int(time.time() * 1000),
        )
        updated = DraftRecord(
            **{
                **record.to_dict(),
                "selection": selection,
                "updated_unix_ms": int(time.time() * 1000),
            }
        )
        _atomic_json(_draft_path(self.root, draft_id), updated.to_dict())
        return updated


def _tool_result(
    call: ToolCall,
    *,
    started: float,
    ok: bool,
    output: str = "",
    error: str = "",
) -> ToolExecutionResult:
    rendered = output or error
    return ToolExecutionResult(
        name=call.name,
        args=call.args,
        ok=ok,
        tool_call_id=call.call_id,
        output=output,
        error=error,
        executed=True,
        elapsed_s=round(time.perf_counter() - started, 4),
        source_trust="trusted_execution",
        output_sha256=hashlib.sha256(rendered.encode("utf-8")).hexdigest(),
    )


def execute_drafting_tool(
    call: ToolCall,
    context: ToolExecutionContext,
) -> ToolExecutionResult:
    started = time.perf_counter()
    spec = _DRAFTING_TOOL_BY_NAME.get(call.name)
    if spec is None:
        return _tool_result(
            call,
            started=started,
            ok=False,
            error=f"unsupported drafting tool: {call.name}",
        )
    try:
        spec.validate_arguments(call.args)
        store = DraftStore(context.workspace)
        if call.name == "DraftCreate":
            record = store.create(
                objective=str(call.args["objective"]),
                candidates=call.args["candidates"],
                criteria=_string_rows(call.args["criteria"], "criteria", allow_empty=False),
                session_id=context.session_id,
            )
            return _tool_result(
                call,
                started=started,
                ok=True,
                output=json.dumps(record.to_dict(), sort_keys=True),
            )
        if call.name == "DraftCompare":
            record = store.compare(
                str(call.args["draft_id"]),
                evidence_refs=_string_rows(
                    call.args.get("evidence_refs"), "evidence_refs", allow_empty=True
                ),
                observations=_string_rows(
                    call.args.get("observations"), "observations", allow_empty=True
                ),
                session_id=context.session_id,
            )
            return _tool_result(
                call,
                started=started,
                ok=True,
                output=json.dumps(record.to_dict(), sort_keys=True),
            )
        if call.name == "DraftSelect":
            record = store.select(
                str(call.args["draft_id"]),
                selected_index=int(call.args["selected_index"]),
                decision=str(call.args["decision"]),
                evidence_refs=_string_rows(
                    call.args.get("evidence_refs"), "evidence_refs", allow_empty=True
                ),
                session_id=context.session_id,
            )
            return _tool_result(
                call,
                started=started,
                ok=True,
                output=json.dumps(record.to_dict(), sort_keys=True),
            )
        if call.name == "DraftStatus":
            record = store.get(str(call.args["draft_id"]), session_id=context.session_id)
            return _tool_result(
                call,
                started=started,
                ok=True,
                output=json.dumps(record.to_dict(), sort_keys=True),
            )
        raise RuntimeError("drafting tool dispatch fell through")
    except (OSError, PermissionError, RuntimeError, ValueError) as exc:
        return _tool_result(
            call,
            started=started,
            ok=False,
            error=f"{type(exc).__name__}: {exc}",
        )


__all__ = [
    "DRAFTING_TOOL_NAMES",
    "DRAFTING_TOOL_SPECS",
    "DraftStore",
    "execute_drafting_tool",
]