"""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 def test_drop_stale_client_history(session, monkeypatch, engine): from app.main import _drop_stale_client_history history = [ {"role": "user", "content": "vieja pregunta"}, {"role": "assistant", "content": "vieja respuesta"}, {"role": "tool", "content": "x", "toolCallId": "t1"}, {"role": "user", "content": "nueva pregunta"}, ] # No stored snapshot for the thread → stale client: keep last user turn. assert _drop_stale_client_history(session, "main", list(history)) == [ {"role": "user", "content": "nueva pregunta"} ] # Fresh conversation (no assistant/tool yet) passes through. fresh = [{"role": "user", "content": "hola"}] assert _drop_stale_client_history(session, "main", list(fresh)) == fresh # Empty messages (hydration trigger) pass through. assert _drop_stale_client_history(session, "main", []) == [] # With a stored snapshot the full history is legitimate. store = _store(monkeypatch, engine) asyncio.run( store.save( scope="default", thread_id="main", snapshot=AGUIThreadSnapshot( messages=[{"id": "u1", "role": "user", "content": "hola"}] ), ) ) assert _drop_stale_client_history(session, "main", list(history)) == history