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

from __future__ import annotations

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

import pytest
from agent import ToolRegistry
from agent.orchestration import (
    ArtifactRef,
    ValidationEvidenceGrant,
    ValidationVerdict,
)
from agent.orchestration.evidence import EvidenceResponse
from agent.orchestration.models import ValidationEvidenceSnapshot
from agent.orchestration.validation_evidence import (
    evidence_handle,
    extend_snapshot,
    start_read_session,
)

from script_build_host.domain.errors import (
    ProtocolViolation,
    ValidationEvidenceUnauthorized,
)
from script_build_host.tools.contracts import AttemptManifest
from script_build_host.tools.gateway import LegacyScriptToolGateway
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 save_script_paragraphs(
        self,
        *,
        paragraphs: list[Mapping[str, Any]],
        expected_state_revision: str,
        context: Mapping[str, Any],
    ) -> dict[str, Any]:
        self.last_payload = {
            "paragraphs": list(paragraphs),
            "expected_state_revision": expected_state_revision,
        }
        self.last_context = dict(context)
        return {"created": 1}

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


class _TaskStore:
    def __init__(self, task: Any, ref: ArtifactRef) -> None:
        self.task = task
        snapshot = extend_snapshot(
            ValidationEvidenceSnapshot(),
            (ValidationEvidenceGrant(ref, evidence_handle(ref)),),
        )
        self.validation = SimpleNamespace(
            evidence_snapshot=snapshot,
            read_session=start_read_session(snapshot, ()),
        )

    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")},
            validations={"validation-a": self.validation},
            context_documents={},
        )


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 _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, ref)
        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_returns_only_bounded_evidence_handles() -> 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 == {
        "items": [],
        "queries_used": 1,
        "queries_remaining": 4,
        "evidence_handles": [],
    }


@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_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(
        "save_script_paragraphs",
        {
            "paragraphs": [{"client_key": "opening", "name": "opening", "content_range": {}}],
            "expected_state_revision": "revision-a",
        },
        context,
    )
    assert candidates.last_context == context
    assert candidates.last_payload == {
        "paragraphs": [{"client_key": "opening", "name": "opening", "content_range": {}}],
        "expected_state_revision": "revision-a",
    }


@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_handle_ids": [],
            }
        ],
        defects=[
            {
                "defect_code": "REALIZATION_PLACEHOLDER",
                "criterion_id": "criterion-a",
                "scope_selector": {"anchor": "mission", "path": ["body", "end"]},
                "observed_excerpt": excerpt,
                "evidence_handle_ids": [evidence_handle(ref)],
                "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_selector": {"anchor": "mission", "path": ["body"]},
        "observed_excerpt": "overlapping body",
        "evidence_handle_ids": [],
        "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_handle_ids": [evidence_handle(ref)],
                }
            ],
            defects=[defect],
            recommendation="replace",
            context=context,
        )

    defect["severity"] = "warning"
    defect["evidence_handle_ids"] = ["src_" + "f" * 20]
    with pytest.raises(
        ValidationEvidenceUnauthorized, match="unknown evidence handle"
    ):
        await gateway.submit_structured_validation(
            verdict="passed",
            criterion_results=[
                {
                    "criterion_id": "criterion-a",
                    "verdict": "passed",
                    "reason": "ok",
                    "evidence_handle_ids": [evidence_handle(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")
    image_handle = "image_" + "a" * 20
    candidates.images = [
        {
            "type": "base64",
            "media_type": "image/png",
            "data": encoded,
            "image_handle": image_handle,
            "sha256": "sha256:" + hashlib.sha256(b"safe-image-bytes").hexdigest(),
        }
    ]
    registry = ToolRegistry()
    register_script_tools(registry, gateway)
    result = await registry.execute(
        "view_frozen_images",
        {"image_handles": [image_handle]},
        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",
        {"image_handles": [image_handle]},
        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"]