script-build-and-framework-tmp/agent/tests/test_trace_attachments.py

from __future__ import annotations

import base64
import inspect
from hashlib import sha256

import pytest

from agent import AttachmentRef, TraceAttachmentStore
from agent.core.runner import AgentRunner, RunConfig
from agent.trace.models import Message, Trace
from agent.trace.store import FileSystemTraceStore


@pytest.mark.asyncio
async def test_filesystem_trace_attachment_roundtrip_and_digest_verification(tmp_path):
    store = FileSystemTraceStore(str(tmp_path))
    await store.create_trace(Trace(trace_id="trace", mode="agent"))
    await store.add_message(
        Message.create(
            trace_id="trace",
            role="tool",
            sequence=1,
            content="attachment source",
        )
    )
    content = b"\x89PNG\r\ncontrolled-image"
    digest = f"sha256:{sha256(content).hexdigest()}"

    ref = await store.store_message_attachment(
        trace_id="trace",
        message_id="trace-0001",
        media_type="image/png",
        content=content,
        sha256=digest,
    )

    assert isinstance(store, TraceAttachmentStore)
    assert isinstance(ref, AttachmentRef)
    assert ref.size_bytes == len(content)
    assert await store.read_message_attachment(ref) == content
    reloaded = FileSystemTraceStore(str(tmp_path))
    assert await reloaded.read_message_attachment(ref) == content
    assert (
        await store.store_message_attachment(
            trace_id="trace",
            message_id="trace-0001",
            media_type="image/png",
            content=content,
            sha256=digest,
        )
        == ref
    )

    with pytest.raises(ValueError, match="digest mismatch"):
        await store.store_message_attachment(
            trace_id="trace",
            message_id="message-2",
            media_type="image/png",
            content=content,
            sha256=f"sha256:{'0' * 64}",
        )

    forged_media_type = AttachmentRef(
        trace_id=ref.trace_id,
        message_id=ref.message_id,
        media_type="image/jpeg",
        sha256=ref.sha256,
        size_bytes=ref.size_bytes,
    )
    with pytest.raises(ValueError, match="does not match metadata"):
        await reloaded.read_message_attachment(forged_media_type)

    payload_path = (
        tmp_path
        / ref.trace_id
        / "attachments"
        / ref.message_id
        / f"{ref.sha256.removeprefix('sha256:')}.bin"
    )
    payload_path.write_bytes(b"X" * len(content))
    with pytest.raises(ValueError, match="integrity check failed"):
        await reloaded.read_message_attachment(ref)


@pytest.mark.asyncio
async def test_runner_uses_optional_attachment_port_and_old_store_stays_compatible(
    tmp_path,
):
    content = b"image bytes"
    image = {
        "type": "base64",
        "media_type": "image/png",
        "data": base64.b64encode(content).decode("ascii"),
    }
    store = FileSystemTraceStore(str(tmp_path))
    await store.create_trace(Trace(trace_id="trace", mode="agent"))
    await store.add_message(
        Message.create(
            trace_id="trace",
            role="tool",
            sequence=1,
            content="attachment source",
        )
    )
    runner = AgentRunner(trace_store=store)

    await runner._persist_tool_attachments("trace", "trace-0001", [image])

    ref = AttachmentRef(
        trace_id="trace",
        message_id="trace-0001",
        media_type="image/png",
        sha256=f"sha256:{sha256(content).hexdigest()}",
        size_bytes=len(content),
    )
    assert await store.read_message_attachment(ref) == content

    class LegacyStore:
        pass

    legacy_runner = AgentRunner(trace_store=LegacyStore())
    await legacy_runner._persist_tool_attachments("trace", "message", [image])

    with pytest.raises(ValueError, match="invalid base64"):
        await runner._persist_tool_attachments(
            "trace",
            "bad-message",
            [{"type": "base64", "media_type": "image/png", "data": "%%%"}],
        )


@pytest.mark.asyncio
async def test_non_multimodal_run_works_with_trace_store_without_attachment_port(
    tmp_path,
):
    inner = FileSystemTraceStore(str(tmp_path))

    class LegacyTraceStore:
        def __getattr__(self, name):
            if name in {"store_message_attachment", "read_message_attachment"}:
                raise AttributeError(name)
            return getattr(inner, name)

    async def llm_call(**_kwargs):
        return {"content": "done", "tool_calls": None, "finish_reason": "stop"}

    legacy = LegacyTraceStore()
    assert not isinstance(legacy, TraceAttachmentStore)
    runner = AgentRunner(trace_store=legacy, llm_call=llm_call)

    result = await runner.run_result(
        [{"role": "user", "content": "plain text task"}],
        RunConfig(name="legacy text run", max_iterations=1),
    )

    assert result["status"] == "completed"
    assert "_get_messages_dir" not in inspect.getsource(AgentRunner)


def test_attachment_ref_rejects_unsafe_or_invalid_fields():
    digest = f"sha256:{'0' * 64}"
    with pytest.raises(ValueError, match="media_type"):
        AttachmentRef("trace", "message", "", digest, 1)
    with pytest.raises(ValueError, match="size_bytes"):
        AttachmentRef("trace", "message", "image/png", digest, -1)