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 == 2 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_cursor_reads_mixed_v1_v2_journal(tmp_path): store = FileSystemTaskStore(str(tmp_path)) ledger = TaskLedger(root_trace_id="mixed", mission="mixed") await store.commit( ledger, expected_revision=-1, event=EventDraft("legacy_event", {"legacy": True}), ) path = store.ledger_path("mixed") raw = json.loads(path.read_text(encoding="utf-8")) raw["_events"][0]["schema_version"] = 1 path.write_text(json.dumps(raw), encoding="utf-8") current = await store.load("mixed") await store.commit( current, expected_revision=current.revision, event=EventDraft("current_event", {"task_id": "task"}), ) first = await store.list_events("mixed", limit=1) second = await store.list_events("mixed", cursor=first.next_cursor, limit=1) assert [first.events[0].schema_version, second.events[0].schema_version] == [1, 2] assert [first.events[0].sequence, second.events[0].sequence] == [1, 2] @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)