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

from datetime import datetime, timedelta
import json
from types import SimpleNamespace

import pytest

from agent.orchestration.store import FileSystemArtifactStore, FileSystemTaskStore
from agent.orchestration.wiring import wire_orchestration
from agent.trace.models import Trace
from agent.trace.store import FileSystemTraceStore

from test_coordinator_integration import FakeExecutor, ROOT_TASK_SPEC, make_coordinator


@pytest.mark.asyncio
async def test_event_v2_payloads_are_refs_and_noop_stage_start_is_silent(tmp_path):
    coordinator, store, _ = await make_coordinator(tmp_path, FakeExecutor([]))
    root_id = (await store.load("root")).root_task_id
    reserved = await coordinator._create_attempt("root", root_id, "worker", None)
    attempt_id = reserved["attempt_id"]
    await coordinator._mark_stage_started("root", None, 0, attempt_id=attempt_id)
    before = await store.load("root")
    page_before = await store.list_events("root", limit=1000)
    await coordinator._mark_stage_started("root", None, 0, attempt_id=attempt_id)
    after = await store.load("root")
    page_after = await store.list_events("root", limit=1000)

    assert after.revision == before.revision
    assert len(page_after.events) == len(page_before.events)
    attempt_event = next(
        event for event in page_after.events if event.event_type == "attempt_created"
    )
    assert attempt_event.schema_version == 2
    assert attempt_event.payload == {"task_id": root_id, "attempt_id": attempt_id}
    stage_event = next(
        event for event in page_after.events if event.event_type == "stage_started"
    )
    assert stage_event.payload == {"task_id": root_id, "attempt_id": attempt_id}
    mirrored = [
        json.loads(line)
        for line in (
            tmp_path / "root" / "orchestration" / "events.jsonl"
        ).read_text(encoding="utf-8").splitlines()
    ]
    assert mirrored[-1]["event_id"] == stage_event.event_id
    assert mirrored[-1]["schema_version"] == 2


@pytest.mark.asyncio
async def test_default_wiring_keeps_durable_journal_without_legacy_mirror(tmp_path):
    trace_store = FileSystemTraceStore(str(tmp_path))
    runner = SimpleNamespace(trace_store=trace_store, task_coordinator=None)
    task_store = FileSystemTaskStore(str(tmp_path))
    coordinator = wire_orchestration(
        runner,
        task_store,
        FileSystemArtifactStore(str(tmp_path)),
    )
    await coordinator.ensure_ledger("root", ROOT_TASK_SPEC)
    assert task_store.ledger_path("root").exists()
    assert not (tmp_path / "root" / "orchestration" / "events.jsonl").exists()


@pytest.mark.asyncio
async def test_trace_correlation_filters_apply_before_limit(tmp_path):
    store = FileSystemTraceStore(str(tmp_path))
    target_time = datetime(2020, 1, 1)
    await store.create_trace(
        Trace(
            trace_id="target",
            mode="agent",
            agent_role="validator",
            parent_trace_id="parent",
            created_at=target_time,
            context={
                "root_trace_id": "root",
                "task_id": "task",
                "attempt_id": "attempt",
                "validation_id": "validation",
                "operation_id": "operation",
            },
        )
    )
    for index in range(1001):
        await store.create_trace(
            Trace(
                trace_id=f"unrelated-{index:04d}",
                mode="agent",
                created_at=target_time + timedelta(days=1, seconds=index),
            )
        )

    filters = {
        "agent_role": "validator",
        "parent_trace_id": "parent",
        "root_trace_id": "root",
        "task_id": "task",
        "attempt_id": "attempt",
        "validation_id": "validation",
        "operation_id": "operation",
    }
    for key, value in filters.items():
        traces = await store.list_traces(limit=1, **{key: value})
        assert [trace.trace_id for trace in traces] == ["target"]

    await store.create_trace(
        Trace(trace_id="root", mode="agent", agent_role="planner")
    )
    assert {
        trace.trace_id
        for trace in await store.list_traces(root_trace_id="root", limit=10)
    } == {"root", "target"}