Agent/tests/test_recursive_validation_artifacts.py

import json
import unittest

from pydantic import ValidationError

from cyber_agent.core.artifacts import (
    ArtifactAccessDenied,
    ArtifactNotFound,
    ArtifactRef,
    ArtifactTemporarilyUnavailable,
    ValidationMaterial,
    extract_artifact_refs,
    material_content_hash,
    resolve_artifact_refs,
)
from cyber_agent.tools.models import ToolResult
from cyber_agent.tools.registry import ToolRegistry


CONTENT = {"script": "official value is 12%", "revision": 2}
CONTENT_HASH = material_content_hash(CONTENT)
REF = ArtifactRef(
    artifact_id="script:42",
    version="v2",
    content_hash=CONTENT_HASH,
    kind="script",
    mime_type="application/json",
)


def material(**overrides):
    values = {
        **REF.model_dump(mode="json"),
        "root_trace_id": "root-1",
        "uid": "user-1",
        "content": CONTENT,
    }
    values.update(overrides)
    return ValidationMaterial.model_validate(values)


class FakeResolver:
    def __init__(self, value):
        self.value = value
        self.calls = []

    async def resolve(self, ref, root_trace_id, uid):
        self.calls.append((ref, root_trace_id, uid))
        if isinstance(self.value, Exception):
            raise self.value
        return self.value


class ArtifactProtocolTest(unittest.IsolatedAsyncioTestCase):
    def test_ref_is_strict_and_hash_is_lowercase_sha256(self):
        with self.assertRaises(ValidationError):
            ArtifactRef.model_validate({
                **REF.model_dump(),
                "content_hash": "A" * 64,
            })
        with self.assertRaises(ValidationError):
            ArtifactRef.model_validate({**REF.model_dump(), "extra": True})

    def test_hash_is_stable_for_json_key_order(self):
        self.assertEqual(
            material_content_hash({"b": 2, "a": 1}),
            material_content_hash({"a": 1, "b": 2}),
        )

    def test_extract_only_reads_top_level_standard_metadata(self):
        self.assertEqual([], extract_artifact_refs({
            "result": {"artifact_refs": [REF.model_dump()]},
        }))
        refs = extract_artifact_refs({
            "artifact_refs": [REF.model_dump(), REF.model_dump()],
            "untrusted": {"artifact_refs": []},
        })
        self.assertEqual([REF], refs)

    def test_malformed_artifact_metadata_fails_closed(self):
        for payload in (
            {"artifact_refs": "not-an-array"},
            {"artifact_refs": [{"artifact_id": "missing-fields"}]},
        ):
            with self.subTest(payload=payload), self.assertRaises(
                (ValueError, ValidationError)
            ):
                extract_artifact_refs(payload)

    async def test_valid_resolver_material_is_accepted(self):
        resolver = FakeResolver(material())
        materials, issues = await resolve_artifact_refs(
            [REF],
            resolver=resolver,
            root_trace_id="root-1",
            uid="user-1",
        )
        self.assertEqual([], issues)
        self.assertEqual([material()], materials)
        self.assertEqual((REF, "root-1", "user-1"), resolver.calls[0])

    async def test_missing_resolver_is_unknown(self):
        materials, issues = await resolve_artifact_refs(
            [REF],
            resolver=None,
            root_trace_id="root-1",
            uid="user-1",
        )
        self.assertEqual([], materials)
        self.assertEqual("unknown", issues[0].outcome)

    async def test_explicit_not_found_is_failed(self):
        _, issues = await resolve_artifact_refs(
            [REF],
            resolver=FakeResolver(ArtifactNotFound("gone")),
            root_trace_id="root-1",
            uid="user-1",
        )
        self.assertEqual("failed", issues[0].outcome)

    async def test_temporary_failure_is_unknown(self):
        _, issues = await resolve_artifact_refs(
            [REF],
            resolver=FakeResolver(ArtifactTemporarilyUnavailable("retry later")),
            root_trace_id="root-1",
            uid="user-1",
        )
        self.assertEqual("unknown", issues[0].outcome)

    async def test_access_denied_and_cross_tree_are_errors(self):
        for value in (
            ArtifactAccessDenied("denied"),
            material(root_trace_id="other-root"),
            material(uid="other-user"),
        ):
            with self.subTest(value=value):
                _, issues = await resolve_artifact_refs(
                    [REF],
                    resolver=FakeResolver(value),
                    root_trace_id="root-1",
                    uid="user-1",
                )
                self.assertEqual("error", issues[0].outcome)

    async def test_version_or_metadata_mismatch_is_error(self):
        changed = material(version="v3")
        _, issues = await resolve_artifact_refs(
            [REF],
            resolver=FakeResolver(changed),
            root_trace_id="root-1",
            uid="user-1",
        )
        self.assertEqual("error", issues[0].outcome)

    async def test_content_hash_mismatch_is_failed(self):
        forged = ValidationMaterial.model_validate({
            **REF.model_dump(mode="json"),
            "root_trace_id": "root-1",
            "uid": "user-1",
            "content": {"script": "forged 30%"},
        })
        _, issues = await resolve_artifact_refs(
            [REF],
            resolver=FakeResolver(forged),
            root_trace_id="root-1",
            uid="user-1",
        )
        self.assertEqual("failed", issues[0].outcome)

    async def test_untyped_resolver_exception_is_unknown(self):
        _, issues = await resolve_artifact_refs(
            [REF],
            resolver=FakeResolver(RuntimeError("database unavailable")),
            root_trace_id="root-1",
            uid="user-1",
        )
        self.assertEqual("unknown", issues[0].outcome)

    async def test_tool_result_preserves_only_standard_artifact_refs(self):
        registry = ToolRegistry()

        async def produce_script():
            return ToolResult(
                title="script",
                output="created",
                metadata={
                    "artifact_refs": [REF.model_dump(mode="json")],
                    "private": "must-not-leak",
                },
            )

        registry.register(produce_script)
        result = await registry.execute("produce_script", {})
        self.assertEqual([REF.model_dump(mode="json")], result["artifact_refs"])
        self.assertNotIn("private", json.dumps(result))


if __name__ == "__main__":
    unittest.main()