Agent/examples/application_reference/providers/artifacts.py

from __future__ import annotations

from dataclasses import dataclass
from typing import Any

from cyber_agent.core.artifacts import (
    ArtifactRef,
    ValidationMaterial,
    material_content_hash,
)


@dataclass(frozen=True)
class _StoredArtifact:
    ref: ArtifactRef
    root_trace_id: str
    uid: str | None
    content: Any


class ReferenceArtifactStore:
    """Example-owned content storage exposed through the framework resolver port."""

    def __init__(self) -> None:
        self._items: dict[tuple[str, str], _StoredArtifact] = {}

    def put(
        self,
        *,
        artifact_id: str,
        version: str,
        kind: str,
        root_trace_id: str,
        uid: str | None,
        content: Any,
    ) -> ArtifactRef:
        ref = ArtifactRef(
            artifact_id=artifact_id,
            version=version,
            content_hash=material_content_hash(content),
            kind=kind,
            mime_type="application/json",
        )
        self._items[(artifact_id, version)] = _StoredArtifact(
            ref=ref,
            root_trace_id=root_trace_id,
            uid=uid,
            content=content,
        )
        return ref

    async def resolve(self, ref, root_trace_id, uid):
        item = self._items[(ref.artifact_id, ref.version)]
        if item.ref != ref:
            raise ValueError("ArtifactRef content hash does not match storage")
        if item.root_trace_id != root_trace_id or item.uid != uid:
            raise ValueError("Artifact belongs to another root or owner")
        return ValidationMaterial(
            **ref.model_dump(mode="json"),
            root_trace_id=root_trace_id,
            uid=uid,
            content=item.content,
        )