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

from __future__ import annotations

import json

import pytest

from agent.trace.models import Trace
from agent.trace.store import FileSystemTraceStore


@pytest.mark.asyncio
async def test_cognition_facade_preserves_new_event_flow(tmp_path) -> None:
    store = FileSystemTraceStore(str(tmp_path))
    await store.create_trace(Trace(trace_id="trace-new", mode="agent"))

    await store.append_cognition_event(
        "trace-new",
        {
            "type": "query",
            "sequence": 4,
            "goal_id": "goal-1",
            "source_ids": ["knowledge-1"],
            "sources": [
                {
                    "id": "knowledge-1",
                    "task": "task",
                    "content": "content",
                }
            ],
        },
    )

    pending = await store.get_pending_knowledge_entries("trace-new")
    assert [item["knowledge_id"] for item in pending] == ["knowledge-1"]

    await store.update_knowledge_evaluation(
        "trace-new",
        "knowledge-1",
        {"eval_status": "helpful", "reason": "used"},
        "goal_complete",
    )

    assert await store.get_pending_knowledge_entries("trace-new") == []
    compatibility_log = await store.get_knowledge_log("trace-new")
    assert compatibility_log["entries"] is compatibility_log["events"]


@pytest.mark.asyncio
async def test_cognition_facade_preserves_legacy_log_update(tmp_path) -> None:
    store = FileSystemTraceStore(str(tmp_path))
    await store.create_trace(Trace(trace_id="trace-legacy", mode="agent"))
    legacy_file = tmp_path / "trace-legacy" / "knowledge_log.json"
    legacy_file.write_text(
        json.dumps(
            {
                "trace_id": "trace-legacy",
                "entries": [
                    {
                        "knowledge_id": "knowledge-legacy",
                        "injected_at_sequence": 2,
                        "eval_result": None,
                    }
                ],
            }
        ),
        encoding="utf-8",
    )

    await store.update_knowledge_evaluation(
        "trace-legacy",
        "knowledge-legacy",
        {"eval_status": "unused", "reason": "not needed"},
        "goal_complete",
    )

    persisted = json.loads(legacy_file.read_text(encoding="utf-8"))
    assert persisted["entries"][0]["eval_result"]["eval_status"] == "unused"