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()