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

import asyncio
import json
import multiprocessing
from copy import deepcopy

import pytest

from agent.orchestration.models import (
    ArtifactRef,
    AttemptSubmission,
    BackgroundOperation,
    CommandRecord,
    EventDraft,
    OperationKind,
    OperationStatus,
    TaskLedger,
    TaskStatus,
)
from agent.orchestration.state_machine import InvalidTaskTransition, transition
from agent.orchestration.store import (
    ArtifactConflict,
    FileSystemArtifactStore,
    FileSystemTaskStore,
    RevisionConflict,
    TaskStoreError,
)


def _commit_from_process(base_path, mission, barrier, results):
    async def run():
        store = FileSystemTaskStore(base_path)
        ledger = await store.load("root")
        ledger.mission = mission
        barrier.wait()
        try:
            await store.commit(ledger, expected_revision=0)
        except RevisionConflict:
            results.put("conflict")
        else:
            results.put("committed")

    asyncio.run(run())


def _freeze_from_process(base_path, barrier, results):
    async def run():
        store = FileSystemArtifactStore(base_path)
        barrier.wait()
        snapshot = await store.freeze("root", "attempt", AttemptSubmission(summary="same"))
        results.put(snapshot.snapshot_id)

    asyncio.run(run())


@pytest.mark.asyncio
async def test_ledger_atomic_revision_and_roundtrip(tmp_path):
    store = FileSystemTaskStore(str(tmp_path))
    ledger = TaskLedger(root_trace_id="root", mission="mission")
    committed = await store.commit(ledger, expected_revision=-1)
    assert committed.revision == 0
    loaded = await store.load("root")
    assert loaded.mission == "mission"
    assert loaded.revision == 0

    stale = TaskLedger(root_trace_id="root", mission="stale", revision=-1)
    with pytest.raises(RevisionConflict):
        await store.commit(stale, expected_revision=-1)

    on_disk = json.loads((tmp_path / "root" / "orchestration" / "ledger.json").read_text())
    assert on_disk["revision"] == 0


@pytest.mark.asyncio
async def test_snapshot_is_immutable_and_hash_is_stable(tmp_path):
    store = FileSystemArtifactStore(str(tmp_path)).for_root("root")
    submission = AttemptSubmission(
        summary=" delivered ",
        artifact_refs=[ArtifactRef(uri="file:///result", version="v1")],
        evidence_refs=[ArtifactRef(uri="db://row/1", digest="abc")],
    )
    first = await store.freeze("attempt-1", submission)
    second = await store.freeze("attempt-2", submission)
    assert first.sha256 == second.sha256
    assert (await store.get(first.snapshot_id)).normalized_content["summary"] == "delivered"

    with pytest.raises(ValueError, match="version or digest"):
        ArtifactRef(uri="file:///mutable")


@pytest.mark.asyncio
async def test_same_attempt_snapshot_is_first_write_wins(tmp_path):
    store = FileSystemArtifactStore(str(tmp_path))
    submission = AttemptSubmission(
        summary="delivered",
        artifact_refs=[ArtifactRef(uri="file:///result", digest="abc")],
    )

    first = await store.freeze("root", "attempt-1", submission)
    replay = await store.freeze("root", "attempt-1", submission)

    assert replay == first
    assert await store.get_for_attempt("root", "attempt-1") == first
    with pytest.raises(ArtifactConflict, match="different artifact snapshot"):
        await store.freeze(
            "root",
            "attempt-1",
            AttemptSubmission(summary="changed", artifact_refs=submission.artifact_refs),
        )


@pytest.mark.asyncio
async def test_concurrent_same_attempt_freeze_creates_one_snapshot(tmp_path):
    first_store = FileSystemArtifactStore(str(tmp_path))
    second_store = FileSystemArtifactStore(str(tmp_path))
    submission = AttemptSubmission(summary="same")

    first, second = await asyncio.gather(
        first_store.freeze("root", "attempt-1", submission),
        second_store.freeze("root", "attempt-1", submission),
    )

    assert first == second
    artifact_files = list((tmp_path / "root" / "orchestration" / "artifacts").glob("*.json"))
    assert len(artifact_files) == 1


@pytest.mark.asyncio
async def test_artifact_orphans_can_be_listed_and_explicitly_cleaned(tmp_path):
    store = FileSystemArtifactStore(str(tmp_path))
    kept = await store.freeze("root", "attempt-kept", AttemptSubmission(summary="kept"))
    orphan = await store.freeze("root", "attempt-orphan", AttemptSubmission(summary="orphan"))

    assert await store.list_orphans("root", ["attempt-kept"]) == [orphan]
    assert await store.cleanup_orphans("root", ["attempt-kept"]) == [orphan.snapshot_id]
    assert await store.get("root", kept.snapshot_id) == kept
    with pytest.raises(FileNotFoundError):
        await store.get("root", orphan.snapshot_id)


@pytest.mark.asyncio
async def test_operations_commands_and_old_ledgers_roundtrip(tmp_path):
    store = FileSystemTaskStore(str(tmp_path))
    ledger = TaskLedger(root_trace_id="root", mission="mission")
    operation = BackgroundOperation(
        operation_id="op-1",
        root_trace_id="root",
        kind=OperationKind.DISPATCH,
        request={"task_ids": ["task-1"]},
        request_fingerprint="fingerprint",
        status=OperationStatus.RUNNING,
    )
    ledger.operations[operation.operation_id] = operation
    ledger.command_records["start:key"] = CommandRecord(
        command_id="start:key",
        operation="start_dispatch",
        fingerprint="fingerprint",
        result_ref={"operation_id": "op-1"},
        committed_revision=0,
        operation_id="op-1",
    )
    await store.commit(ledger, -1)

    loaded = await store.load("root")
    assert loaded.operations["op-1"].kind is OperationKind.DISPATCH
    assert loaded.command_records["start:key"].result_ref == {"operation_id": "op-1"}
    assert [item.operation_id for item in await store.list_recoverable("root")] == ["op-1"]

    old_path = tmp_path / "old-root" / "orchestration" / "ledger.json"
    old_path.parent.mkdir(parents=True)
    old_path.write_text(json.dumps({"root_trace_id": "old-root", "mission": "old"}), encoding="utf-8")
    old = await store.load("old-root")
    assert old.operations == {}
    assert old.command_records == {}


@pytest.mark.asyncio
async def test_events_commit_atomically_and_page_with_opaque_cursor(tmp_path):
    store = FileSystemTaskStore(str(tmp_path))
    ledger = TaskLedger(root_trace_id="root", mission="mission")
    await store.commit(
        ledger,
        -1,
        event=EventDraft("ledger_created", {"mission": "mission"}, command_id="command-1"),
    )
    ledger.mission = "updated"
    await store.commit(
        ledger,
        0,
        event=EventDraft("mission_updated", {"mission": "updated"}, command_id="command-2"),
    )

    first_page = await store.list_events("root", limit=1)
    assert [event.sequence for event in first_page.events] == [1]
    assert first_page.events[0].schema_version == 1
    assert first_page.has_more is True
    assert first_page.next_cursor and "command" not in first_page.next_cursor
    second_page = await store.list_events("root", cursor=first_page.next_cursor)
    assert [event.event_type for event in second_page.events] == ["mission_updated"]
    assert second_page.events[0].ledger_revision == 1

    raw = json.loads((tmp_path / "root" / "orchestration" / "ledger.json").read_text())
    assert len(raw["_events"]) == 2
    assert "_events" not in (await store.load("root")).to_dict()
    with pytest.raises(ValueError, match="Invalid event cursor"):
        await store.list_events("other-root", cursor=first_page.next_cursor)


@pytest.mark.asyncio
async def test_event_command_replay_does_not_duplicate_event(tmp_path):
    store = FileSystemTaskStore(str(tmp_path))
    ledger = TaskLedger(root_trace_id="root", mission="mission")
    await store.commit(ledger, -1, event=EventDraft("created", {}, command_id="command-1"))
    ledger.mission = "same command replay"
    await store.commit(ledger, 0, event=EventDraft("created", {}, command_id="command-1"))

    page = await store.list_events("root")
    assert len(page.events) == 1


@pytest.mark.asyncio
async def test_two_store_instances_enforce_compare_and_swap(tmp_path):
    first_store = FileSystemTaskStore(str(tmp_path))
    second_store = FileSystemTaskStore(str(tmp_path))
    ledger = TaskLedger(root_trace_id="root", mission="mission")
    await first_store.commit(ledger, -1)
    first_copy = await first_store.load("root")
    stale_copy = deepcopy(first_copy)
    first_copy.mission = "winner"
    stale_copy.mission = "stale"
    await first_store.commit(first_copy, 0)
    with pytest.raises(RevisionConflict):
        await second_store.commit(stale_copy, 0)
    assert (await second_store.load("root")).mission == "winner"


def test_cross_process_lock_allows_only_one_cas_winner(tmp_path):
    asyncio.run(
        FileSystemTaskStore(str(tmp_path)).commit(
            TaskLedger(root_trace_id="root", mission="initial"),
            -1,
        )
    )
    context = multiprocessing.get_context("spawn")
    barrier = context.Barrier(2)
    results = context.Queue()
    processes = [
        context.Process(
            target=_commit_from_process,
            args=(str(tmp_path), mission, barrier, results),
        )
        for mission in ("first", "second")
    ]
    for process in processes:
        process.start()
    for process in processes:
        process.join(timeout=10)
        assert process.exitcode == 0

    assert sorted([results.get(timeout=1), results.get(timeout=1)]) == ["committed", "conflict"]
    assert asyncio.run(FileSystemTaskStore(str(tmp_path)).load("root")).revision == 1


def test_cross_process_same_attempt_freeze_writes_one_snapshot(tmp_path):
    context = multiprocessing.get_context("spawn")
    barrier = context.Barrier(2)
    results = context.Queue()
    processes = [
        context.Process(
            target=_freeze_from_process,
            args=(str(tmp_path), barrier, results),
        )
        for _index in range(2)
    ]
    for process in processes:
        process.start()
    for process in processes:
        process.join(timeout=10)
        assert process.exitcode == 0

    assert results.get(timeout=1) == results.get(timeout=1)
    artifact_files = list((tmp_path / "root" / "orchestration" / "artifacts").glob("*.json"))
    assert len(artifact_files) == 1


@pytest.mark.asyncio
async def test_store_rejects_paths_outside_base(tmp_path):
    task_store = FileSystemTaskStore(str(tmp_path))
    artifact_store = FileSystemArtifactStore(str(tmp_path))
    with pytest.raises(ValueError, match="escapes"):
        await task_store.load("../escape")
    with pytest.raises(ValueError, match="escapes"):
        await artifact_store.get("root", "../../escape")


@pytest.mark.asyncio
async def test_store_reports_corrupt_ledger_and_event_journal(tmp_path):
    store = FileSystemTaskStore(str(tmp_path))
    path = tmp_path / "root" / "orchestration" / "ledger.json"
    path.parent.mkdir(parents=True)
    path.write_text("{truncated", encoding="utf-8")
    with pytest.raises(TaskStoreError, match="Cannot load task ledger"):
        await store.load("root")

    path.write_text(
        json.dumps({"root_trace_id": "root", "mission": "mission", "_events": {}}),
        encoding="utf-8",
    )
    with pytest.raises(TaskStoreError, match="event journal is invalid"):
        await store.list_events("root")

    with pytest.raises(ValueError, match="between 1 and 1000"):
        await store.list_events("root", limit=0)


@pytest.mark.asyncio
async def test_store_wraps_nested_model_and_event_parse_errors(tmp_path):
    store = FileSystemTaskStore(str(tmp_path))
    path = tmp_path / "root" / "orchestration" / "ledger.json"
    path.parent.mkdir(parents=True)
    path.write_text(
        json.dumps(
            {
                "root_trace_id": "root",
                "mission": "mission",
                "tasks": {"broken": {"task_id": "broken", "status": "not-a-status"}},
            }
        ),
        encoding="utf-8",
    )
    with pytest.raises(TaskStoreError, match="Cannot parse task ledger"):
        await store.load("root")

    path.write_text(
        json.dumps({"root_trace_id": "root", "mission": "mission", "tasks": []}),
        encoding="utf-8",
    )
    with pytest.raises(TaskStoreError, match="Cannot parse task ledger"):
        await store.load("root")

    path.write_text(
        json.dumps(
            {
                "root_trace_id": "root",
                "mission": "mission",
                "_events": [{"sequence": 1}],
            }
        ),
        encoding="utf-8",
    )
    with pytest.raises(TaskStoreError, match="contains an invalid event"):
        await store.list_events("root")


@pytest.mark.asyncio
async def test_failed_atomic_write_does_not_advance_caller_ledger(tmp_path, monkeypatch):
    store = FileSystemTaskStore(str(tmp_path))
    ledger = TaskLedger(root_trace_id="root", mission="mission")
    await store.commit(ledger, expected_revision=-1, event=EventDraft("created"))
    original_updated_at = ledger.updated_at
    ledger.mission = "uncommitted"

    def fail_write(_path, _data):
        raise OSError("injected replace failure")

    monkeypatch.setattr(FileSystemTaskStore, "_atomic_json_write", staticmethod(fail_write))
    with pytest.raises(OSError, match="injected replace failure"):
        await store.commit(ledger, expected_revision=0, event=EventDraft("updated"))

    assert ledger.revision == 0
    assert ledger.updated_at == original_updated_at
    loaded = await store.load("root")
    assert loaded.mission == "mission"
    assert [event.event_type for event in (await store.list_events("root")).events] == ["created"]


@pytest.mark.asyncio
async def test_orphan_cleanup_never_trusts_snapshot_id_as_path(tmp_path):
    store = FileSystemArtifactStore(str(tmp_path))
    orchestration_dir = tmp_path / "root" / "orchestration"
    artifacts_dir = orchestration_dir / "artifacts"
    artifacts_dir.mkdir(parents=True)
    ledger_path = orchestration_dir / "ledger.json"
    ledger_path.write_text("sentinel", encoding="utf-8")
    (artifacts_dir / "malicious.json").write_text(
        json.dumps(
            {
                "snapshot_id": "../ledger",
                "attempt_id": "orphan",
                "normalized_content": {},
                "sha256": "digest",
                "artifact_refs": [],
                "evidence_refs": [],
            }
        ),
        encoding="utf-8",
    )

    assert await store.cleanup_orphans("root", []) == ["../ledger"]
    assert ledger_path.read_text(encoding="utf-8") == "sentinel"
    assert not (artifacts_dir / "malicious.json").exists()


def test_state_machine_rejects_shortcut_to_completed():
    assert transition(TaskStatus.PENDING, TaskStatus.RUNNING) == TaskStatus.RUNNING
    with pytest.raises(InvalidTaskTransition):
        transition(TaskStatus.RUNNING, TaskStatus.COMPLETED)