script-build-and-framework-tmp/script_build_host/tests/test_phase_two_agent_tools.py

from __future__ import annotations

import base64
from collections.abc import Mapping
from types import SimpleNamespace
from typing import Any

import pytest
from agent import ToolRegistry
from agent.orchestration import ArtifactRef, ValidationVerdict

from script_build_host.domain.artifacts import ArtifactState
from script_build_host.domain.errors import ArtifactNotFound, ProtocolViolation
from script_build_host.tools.contracts import AttemptManifest
from script_build_host.tools.gateway import LegacyScriptToolGateway, RetrievalResult
from script_build_host.tools.registry import register_script_tools


class _CandidateTools:
    def __init__(self, artifact_ref: ArtifactRef) -> None:
        self.artifact_ref = artifact_ref
        self.last_payload: dict[str, Any] | None = None
        self.last_context: dict[str, Any] | None = None
        self.images: list[dict[str, Any]] = []

    async def resolve_attempt_manifest(self, *, context: Mapping[str, Any]) -> AttemptManifest:
        self.last_context = dict(context)
        return AttemptManifest(
            artifact_ref=self.artifact_ref,
            scope_ref="script-build://scopes/body",
        )

    async def create_script_paragraph(
        self, *, payload: Mapping[str, Any], context: Mapping[str, Any]
    ) -> dict[str, int]:
        self.last_payload = dict(payload)
        self.last_context = dict(context)
        return {"paragraph_id": 7}

    async def view_frozen_images(
        self,
        *,
        raw_artifact_refs: list[str],
        context: Mapping[str, Any],
    ) -> list[dict[str, Any]]:
        self.last_payload = {"raw_artifact_refs": raw_artifact_refs}
        self.last_context = dict(context)
        return self.images


class _TaskStore:
    def __init__(self, task: Any) -> None:
        self.task = task

    async def load(self, _root_trace_id: str) -> Any:
        return SimpleNamespace(
            tasks={"task-1": self.task},
            attempts={"attempt-a": SimpleNamespace(attempt_id="attempt-a", task_id="task-1")},
        )


class _ArtifactStore:
    def __init__(self, ref: ArtifactRef) -> None:
        self.ref = ref

    async def get(self, _root_trace_id: str, _snapshot_id: str) -> Any:
        return SimpleNamespace(attempt_id="attempt-a", artifact_refs=[self.ref], evidence_refs=[])


class _Bindings:
    async def get_by_root(self, _root_trace_id: str) -> Any:
        return SimpleNamespace(script_build_id=9, input_snapshot_id=3)


class _BudgetGuard:
    def __init__(self, error: Exception | None = None) -> None:
        self.error = error
        self.calls: list[tuple[str, str, str]] = []

    async def ensure_tool_budget(self, *, root_trace_id: str, task_id: str, entry: str) -> None:
        self.calls.append((root_trace_id, task_id, entry))
        if self.error is not None:
            raise self.error


class _BusinessArtifacts:
    def __init__(self, ref: ArtifactRef) -> None:
        self.ref = ref

    async def read_by_ref(
        self,
        ref: ArtifactRef,
        *,
        script_build_id: int,
        task_id: str,
        attempt_id: str,
    ) -> Any:
        assert (ref, script_build_id, task_id, attempt_id) == (
            self.ref,
            9,
            "task-1",
            "attempt-a",
        )
        return SimpleNamespace(
            artifact_type=SimpleNamespace(value=ref.kind),
            canonical_sha256=ref.digest,
            state=SimpleNamespace(value="frozen"),
        )


class _RetrievalArtifacts:
    def __init__(self) -> None:
        self.version: Any | None = None

    async def get_by_attempt(self, **_owners: Any) -> Any:
        if self.version is None:
            raise ArtifactNotFound()
        return self.version

    async def freeze(self, *, artifact: Any, **owners: Any) -> tuple[Any, ArtifactRef]:
        ref = ArtifactRef(
            "script-build://artifact-versions/8",
            "evidence",
            "8",
            "sha256:" + "b" * 64,
        )
        self.version = SimpleNamespace(
            artifact_version_id=8,
            state=ArtifactState.FROZEN,
            artifact=artifact,
        )
        return self.version, ref


class _CountingRetrievalAdapter:
    def __init__(self) -> None:
        self.calls = 0

    async def retrieve(self, **_kwargs: Any) -> RetrievalResult:
        self.calls += 1
        return RetrievalResult(
            source_refs=("decode://case/1",),
            summary="one immutable result",
        )


class _Coordinator:
    def __init__(self, ref: ArtifactRef) -> None:
        criterion = SimpleNamespace(criterion_id="criterion-a", hard=True)
        task = SimpleNamespace(
            task_id="task-1",
            validation_ids=(),
            current_spec=SimpleNamespace(
                acceptance_criteria=(criterion,),
                context_refs=("script-build://task-kinds/paragraph",),
            ),
        )
        self.task_store = _TaskStore(task)
        self.artifact_store = _ArtifactStore(ref)
        self.attempt_submission: tuple[dict[str, Any], Any] | None = None
        self.validation_submission: tuple[dict[str, Any], tuple[Any, ...]] | None = None

    async def submit_attempt(self, context: Mapping[str, Any], submission: Any) -> dict[str, Any]:
        self.attempt_submission = (dict(context), submission)
        return {
            "attempt_id": context["attempt_id"],
            "snapshot_id": "snapshot-a",
            "status": "awaiting_validation",
        }

    async def submit_validation(self, context: Mapping[str, Any], *values: Any) -> dict[str, Any]:
        self.validation_submission = (dict(context), values)
        return {
            "validation_id": context["validation_id"],
            "verdict": values[0].value,
            "status": "awaiting_decision",
        }


def _gateway(ref: ArtifactRef) -> tuple[LegacyScriptToolGateway, _Coordinator, _CandidateTools]:
    coordinator = _Coordinator(ref)
    candidates = _CandidateTools(ref)
    gateway = LegacyScriptToolGateway(
        bindings=_Bindings(),  # type: ignore[arg-type]
        snapshots=SimpleNamespace(),
        artifacts=_BusinessArtifacts(ref),  # type: ignore[arg-type]
        coordinator=coordinator,
        candidate_tools=candidates,  # type: ignore[arg-type]
    )
    return gateway, coordinator, candidates


@pytest.mark.asyncio
async def test_attempt_submission_is_derived_from_protected_attempt_artifact() -> None:
    ref = ArtifactRef(
        uri="script-build://artifact-versions/7",
        kind="paragraph",
        version="7",
        digest="sha256:" + "a" * 64,
    )
    gateway, coordinator, candidates = _gateway(ref)
    budget = _BudgetGuard()
    gateway.budget_guard = budget
    context = {
        "root_trace_id": "root-a",
        "task_id": "task-1",
        "attempt_id": "attempt-a",
        "spec_version": 1,
    }
    await gateway.submit_current_attempt(context)

    assert coordinator.attempt_submission is not None
    submitted_context, submission = coordinator.attempt_submission
    assert submitted_context == context
    assert submission.artifact_refs == [ref]
    assert submission.evidence_refs == []
    assert "paragraph" in submission.summary
    assert "sha256:" in submission.summary
    assert candidates.last_context == context
    assert budget.calls == [("root-a", "task-1", "submit")]


@pytest.mark.asyncio
async def test_retrieval_budget_is_checked_before_adapter_or_snapshot_access() -> None:
    ref = ArtifactRef(
        uri="script-build://artifact-versions/7",
        kind="evidence",
        version="7",
        digest="sha256:" + "a" * 64,
    )
    gateway, _, _ = _gateway(ref)
    budget = _BudgetGuard(ProtocolViolation("retrieval budget exhausted"))
    gateway.budget_guard = budget
    with pytest.raises(ProtocolViolation, match="budget exhausted"):
        await gateway.retrieve(
            "decode",
            "search_script_decode_case",
            {"query": "opening"},
            {
                "root_trace_id": "root-a",
                "task_id": "task-1",
                "attempt_id": "attempt-a",
            },
        )
    assert budget.calls == [("root-a", "task-1", "retrieval")]


@pytest.mark.asyncio
async def test_one_retrieval_attempt_never_calls_adapter_twice() -> None:
    artifacts = _RetrievalArtifacts()
    adapter = _CountingRetrievalAdapter()

    class Snapshots:
        async def get(self, snapshot_id: str, *, script_build_id: int) -> Any:
            assert (snapshot_id, script_build_id) == ("3", 9)
            return SimpleNamespace(account={})

    gateway = LegacyScriptToolGateway(
        bindings=_Bindings(),  # type: ignore[arg-type]
        snapshots=Snapshots(),  # type: ignore[arg-type]
        artifacts=artifacts,  # type: ignore[arg-type]
        coordinator=SimpleNamespace(),
        retrieval_adapters={"decode": adapter},
    )
    context = {
        "root_trace_id": "root-a",
        "task_id": "task-1",
        "attempt_id": "attempt-a",
        "spec_version": 1,
    }
    await gateway.retrieve("decode", "search_script_decode_case", {"query": "opening"}, context)
    with pytest.raises(ProtocolViolation, match="only one Evidence"):
        await gateway.retrieve(
            "decode", "search_script_decode_case", {"query": "different"}, context
        )
    assert adapter.calls == 1


@pytest.mark.asyncio
async def test_candidate_write_receives_identity_only_from_protected_context() -> None:
    ref = ArtifactRef(
        uri="script-build://artifact-versions/7",
        kind="paragraph",
        version="7",
        digest="sha256:" + "a" * 64,
    )
    gateway, _, candidates = _gateway(ref)
    context = {
        "root_trace_id": "root-a",
        "task_id": "task-1",
        "attempt_id": "attempt-a",
        "spec_version": 1,
    }
    await gateway.candidate_command(
        "create_script_paragraph",
        {"paragraph_index": 1, "name": "opening", "content_range": {}},
        context,
    )
    assert candidates.last_context == context
    assert candidates.last_payload == {
        "paragraph_index": 1,
        "name": "opening",
        "content_range": {},
    }


@pytest.mark.asyncio
async def test_structured_validation_keeps_excerpt_out_of_framework_summary() -> None:
    ref = ArtifactRef(
        uri="script-build://artifact-versions/7",
        kind="paragraph",
        version="7",
        digest="sha256:" + "a" * 64,
    )
    gateway, coordinator, _ = _gateway(ref)
    context = {
        "root_trace_id": "root-a",
        "task_id": "task-1",
        "attempt_id": "attempt-a",
        "snapshot_id": "snapshot-a",
        "validation_id": "validation-a",
    }
    excerpt = "this exact candidate passage must not enter the bounded summary"
    await gateway.submit_structured_validation(
        verdict="failed",
        criterion_results=[
            {"criterion_id": "criterion-a", "verdict": "failed", "reason": "not realized"}
        ],
        defects=[
            {
                "defect_code": "REALIZATION_PLACEHOLDER",
                "criterion_id": "criterion-a",
                "scope_ref": "script-build://scopes/body/end",
                "observed_excerpt": excerpt,
                "evidence_refs": [
                    {
                        "uri": ref.uri,
                        "kind": ref.kind,
                        "version": ref.version,
                        "digest": ref.digest,
                    }
                ],
                "severity": "hard",
                "invalidated_inputs": ["decision-a"],
                "recommended_action_class": "split",
            }
        ],
        recommendation="split",
        context=context,
    )

    assert coordinator.validation_submission is not None
    _, values = coordinator.validation_submission
    verdict, results, summary, evidence_refs, unverified, risks, recommendation = values
    assert verdict is ValidationVerdict.FAILED
    assert results[0].reason.endswith("defect_codes=REALIZATION_PLACEHOLDER")
    assert excerpt not in summary
    assert len(summary) <= 500
    assert evidence_refs == (ref,)
    assert unverified == ()
    assert risks == ("REALIZATION_PLACEHOLDER",)
    assert recommendation == "split"


@pytest.mark.asyncio
async def test_passed_validation_rejects_blocking_defect_and_unfrozen_evidence() -> None:
    ref = ArtifactRef(
        uri="script-build://artifact-versions/7",
        kind="paragraph",
        version="7",
        digest="sha256:" + "a" * 64,
    )
    gateway, _, _ = _gateway(ref)
    context = {
        "root_trace_id": "root-a",
        "task_id": "task-1",
        "attempt_id": "attempt-a",
        "snapshot_id": "snapshot-a",
        "validation_id": "validation-a",
    }
    defect: dict[str, Any] = {
        "defect_code": "WRITE_SCOPE_CONFLICT",
        "criterion_id": "criterion-a",
        "scope_ref": "script-build://scopes/body",
        "observed_excerpt": "overlapping body",
        "evidence_refs": [],
        "severity": "hard",
        "invalidated_inputs": [],
        "recommended_action_class": "replace",
    }
    with pytest.raises(ProtocolViolation, match="cannot contain hard"):
        await gateway.submit_structured_validation(
            verdict="passed",
            criterion_results=[
                {"criterion_id": "criterion-a", "verdict": "passed", "reason": "ok"}
            ],
            defects=[defect],
            recommendation="replace",
            context=context,
        )

    defect["severity"] = "warning"
    defect["evidence_refs"] = [
        {
            "uri": "script-build://artifact-versions/8",
            "kind": "paragraph",
            "version": "8",
            "digest": "sha256:" + "b" * 64,
        }
    ]
    with pytest.raises(ProtocolViolation, match="fixed validation snapshot"):
        await gateway.submit_structured_validation(
            verdict="passed",
            criterion_results=[
                {"criterion_id": "criterion-a", "verdict": "passed", "reason": "ok"}
            ],
            defects=[defect],
            recommendation="accept",
            context=context,
        )


@pytest.mark.asyncio
async def test_view_frozen_images_returns_multimodal_without_url_or_data_in_text() -> None:
    ref = ArtifactRef(
        uri="script-build://artifact-versions/7",
        kind="paragraph",
        version="7",
        digest="sha256:" + "a" * 64,
    )
    gateway, _, candidates = _gateway(ref)
    encoded = base64.b64encode(b"safe-image-bytes").decode("ascii")
    candidates.images = [
        {
            "type": "base64",
            "media_type": "image/png",
            "data": encoded,
            "raw_artifact_ref": "script-build://raw-artifacts/sha256/abc",
        }
    ]
    registry = ToolRegistry()
    register_script_tools(registry, gateway)
    result = await registry.execute(
        "view_frozen_images",
        {"raw_artifact_refs": ["script-build://raw-artifacts/sha256/abc"]},
        context={"root_trace_id": "root-a", "task_id": "task-1"},
    )
    assert isinstance(result, dict)
    assert result["images"] == [{"type": "base64", "media_type": "image/png", "data": encoded}]
    assert encoded not in result["text"]
    assert "http" not in result["text"]

    candidates.images = [
        {
            "type": "base64",
            "media_type": "image/png",
            "data": encoded,
            "url": "https://forbidden.invalid/image.png",
        }
    ]
    rejected = await registry.execute(
        "view_frozen_images",
        {"raw_artifact_refs": ["script-build://raw-artifacts/sha256/abc"]},
        context={"root_trace_id": "root-a", "task_id": "task-1"},
    )
    assert isinstance(rejected, str)
    assert "must not contain network URLs" in rejected