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 dataclasses import asdict
from types import SimpleNamespace
from typing import Any

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

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", "decode://case/1"),
            summary="one immutable result",
            supports=("opening", "opening"),
        )


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_validation_query_includes_exact_current_business_artifact() -> None:
    ref = ArtifactRef(
        uri="script-build://artifact-versions/9",
        kind="paragraph",
        version="9",
        digest="sha256:" + "c" * 64,
    )
    gateway, coordinator, _candidates = _gateway(ref)

    async def query_evidence(_context: Mapping[str, Any], _request: Any) -> EvidenceResponse:
        return EvidenceResponse(items=[], evidence_refs=[], queries_used=1, queries_remaining=4)

    coordinator.query_evidence = query_evidence  # type: ignore[attr-defined]

    class Artifacts:
        async def read_by_ref(self, _ref: ArtifactRef, **_owners: Any) -> Any:
            return SimpleNamespace(
                artifact=SimpleNamespace(
                    content_payload=lambda: {
                        "schema_version": "paragraph-artifact/v1",
                        "paragraph_key": "opening",
                        "text": "grounded opening",
                    }
                )
            )

    gateway.artifacts = Artifacts()  # type: ignore[assignment]
    result = await gateway.query_validation_evidence(
        "current candidate",
        5,
        {
            "root_trace_id": "root-a",
            "task_id": "task-1",
            "attempt_id": "attempt-a",
            "snapshot_id": "snapshot-a",
            "validation_id": "validation-a",
        },
    )

    assert result["current_artifacts"] == [
        {
            "artifact_ref": {
                "uri": ref.uri,
                "kind": ref.kind,
                "version": ref.version,
                "digest": ref.digest,
                "summary": ref.summary,
                "metadata": ref.metadata,
            },
            "artifact": {
                "schema_version": "paragraph-artifact/v1",
                "paragraph_key": "opening",
                "text": "grounded opening",
            },
        }
    ]


@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)
    assert artifacts.version.artifact.source_refs == ("decode://case/1",)
    assert artifacts.version.artifact.supports == ("opening",)
    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",
                "evidence_refs": [],
            }
        ],
        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",
                    "evidence_refs": [asdict(ref)],
                }
            ],
            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="protected validation closure"):
        await gateway.submit_structured_validation(
            verdict="passed",
            criterion_results=[
                {
                    "criterion_id": "criterion-a",
                    "verdict": "passed",
                    "reason": "ok",
                    "evidence_refs": [asdict(ref)],
                }
            ],
            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, dict)
    assert rejected["_control"]["failure"] == {
        "code": "INVALID_TOOL_ARGUMENT",
        "message": "frozen images must not contain network URLs",
        "disposition": "retry_call",
        "source_tool": "view_frozen_images",
        "details": {"task_id": "task-1"},
    }
    assert "must not contain network URLs" in rejected["text"]