WealthySmart/backend/tests/test_chat_persistence.py

"""Chat thread persistence: snapshot store round-trip + clear endpoint."""

import asyncio

from agent_framework_ag_ui import AGUIThreadSnapshot
from sqlmodel import select

import app.agent.snapshot_store as snapshot_store_module
from app.agent.snapshot_store import PostgresSnapshotStore
from app.api.v1.endpoints.chat import clear_thread
from app.models.models import ChatThreadSnapshot


def _store(monkeypatch, engine):
    # The store opens its own sessions from app.db.engine; point it at the
    # test engine instead.
    monkeypatch.setattr(snapshot_store_module, "engine", engine)
    return PostgresSnapshotStore()


def test_save_get_roundtrip_normalizes_and_filters(monkeypatch, engine):
    store = _store(monkeypatch, engine)
    snapshot = AGUIThreadSnapshot(
        messages=[
            {"id": "u1", "role": "user", "content": "hola"},
            {
                "id": "a1",
                "role": "assistant",
                "content": "",
                # MAF emits snake_case; the store must return camelCase so
                # @ag-ui/client's Zod schema doesn't strip it.
                "tool_calls": [{"id": "t1", "type": "function"}],
            },
            {"id": "t1r", "role": "tool", "content": "ok", "tool_call_id": "t1"},
            # MAF writes tool_calls: null on plain-text assistant messages —
            # must not survive as toolCalls: null in the replay.
            {"id": "a2", "role": "assistant", "content": "done", "tool_calls": None},
            {"id": "ctx-123", "role": "system", "content": "run context"},
        ],
        state={"k": "v"},
    )
    asyncio.run(store.save(scope="default", thread_id="main", snapshot=snapshot))

    loaded = asyncio.run(store.get(scope="default", thread_id="main"))
    assert loaded is not None
    assert [m["id"] for m in loaded.messages] == ["u1", "a1", "t1r", "a2"]
    assistant = loaded.messages[1]
    assert "tool_calls" not in assistant
    assert assistant["toolCalls"] == [{"id": "t1", "type": "function"}]
    tool_msg = loaded.messages[2]
    assert "tool_call_id" not in tool_msg
    assert tool_msg["toolCallId"] == "t1"
    text_msg = loaded.messages[3]
    assert "toolCalls" not in text_msg and "tool_calls" not in text_msg
    assert loaded.state == {"k": "v"}


def test_save_upserts_single_row(monkeypatch, engine, session):
    store = _store(monkeypatch, engine)
    for n in range(3):
        snap = AGUIThreadSnapshot(
            messages=[{"id": f"u{i}", "role": "user", "content": str(i)} for i in range(n + 1)]
        )
        asyncio.run(store.save(scope="default", thread_id="main", snapshot=snap))

    rows = session.exec(select(ChatThreadSnapshot)).all()
    assert len(rows) == 1
    assert len(rows[0].messages) == 3


def test_get_unknown_thread_returns_none(monkeypatch, engine):
    store = _store(monkeypatch, engine)
    assert asyncio.run(store.get(scope="default", thread_id="nope")) is None


def test_delete_and_scope_isolation(monkeypatch, engine):
    store = _store(monkeypatch, engine)
    snap = AGUIThreadSnapshot(messages=[{"id": "u1", "role": "user", "content": "x"}])
    asyncio.run(store.save(scope="default", thread_id="main", snapshot=snap))

    assert asyncio.run(store.get(scope="other", thread_id="main")) is None
    assert asyncio.run(store.delete(scope="default", thread_id="main")) is True
    assert asyncio.run(store.delete(scope="default", thread_id="main")) is False
    assert asyncio.run(store.get(scope="default", thread_id="main")) is None


def test_save_dedupes_multi_run_turn(monkeypatch, engine):
    """Mirror of MAF's duplicate merge on frontend-tool roundtrips: repeated
    message ids, a re-issued identical render call, and extra tool results."""
    store = _store(monkeypatch, engine)

    def call(cid, name, args='{"a":1}'):
        return {"id": cid, "type": "function", "function": {"name": name, "arguments": args}}

    snapshot = AGUIThreadSnapshot(
        messages=[
            {"id": "u1", "role": "user", "content": "resumen"},
            {"id": "a1", "role": "assistant", "content": "",
             "tool_calls": [call("c1", "get_cycle_summary"), call("c2", "render_spending_summary")]},
            {"id": "t1", "role": "tool", "content": "data", "tool_call_id": "c1"},
            # MAF re-records the resent suffix verbatim…
            {"id": "u1", "role": "user", "content": "resumen"},
            {"id": "t1", "role": "tool", "content": "data", "tool_call_id": "c1"},
            {"id": "t2", "role": "tool", "content": "ok", "tool_call_id": "c2"},
            # …plus a synthetic second result for an answered call…
            {"id": "t3", "role": "tool", "content": "", "tool_call_id": "c1"},
            # …and the model re-issues the same render call with a new id.
            {"id": "a2", "role": "assistant", "content": "",
             "tool_calls": [call("c3", "render_spending_summary")]},
            {"id": "t4", "role": "tool", "content": "ok", "tool_call_id": "c3"},
        ]
    )
    asyncio.run(store.save(scope="default", thread_id="main", snapshot=snapshot))

    loaded = asyncio.run(store.get(scope="default", thread_id="main"))
    assert [m["id"] for m in loaded.messages] == ["u1", "a1", "t1", "t2"]
    # The duplicate render call (c3) and its result are gone; each surviving
    # call has exactly one result.
    assert [tc["id"] for tc in loaded.messages[1]["toolCalls"]] == ["c1", "c2"]
    assert [(m["toolCallId"]) for m in loaded.messages if m["role"] == "tool"] == ["c1", "c2"]


def test_clear_thread_endpoint(monkeypatch, engine, session):
    store = _store(monkeypatch, engine)
    assert clear_thread(session=session, _user="testadmin") == {"cleared": False}

    snap = AGUIThreadSnapshot(messages=[{"id": "u1", "role": "user", "content": "x"}])
    asyncio.run(store.save(scope="default", thread_id="main", snapshot=snap))

    assert clear_thread(session=session, _user="testadmin") == {"cleared": True}
    assert asyncio.run(store.get(scope="default", thread_id="main")) is None