Assistant: persist the chat thread in Postgres

Back MAF's opt-in AG-UI snapshot persistence with a DB store so the
conversation survives page reloads and restarts:

- ChatThreadSnapshot: one upserted row per (scope, thread_id), JSONB
  messages/state/interrupt (migration 794baf50635b)
- PostgresSnapshotStore implements the AGUIThreadSnapshotStore protocol.
  It opens its own sessions (MAF saves during SSE streaming, after the
  request middleware has closed the ContextVar session) and sanitizes on
  save: camelCase toolCalls/toolCallId (the @ag-ui/client Zod schema
  strips snake_case, which would lose cards on replay), drop BFF ctx-*
  context messages, dedupe MAF's re-recorded multi-run turns, and drop
  orphaned tool results
- Endpoint mount gains snapshot_store + a constant scope resolver
  (single-user app; auth already enforced by agent_auth_and_session)
- DELETE /api/v1/chat/thread resets the conversation

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

backend/alembic/versions/794baf50635b_chat_thread_snapshot.py

@@ -0,0 +1,45 @@
+"""chat thread snapshot
+
+Revision ID: 794baf50635b
+Revises: 961802a2f50d
+Create Date: 2026-07-04 17:49:14.914103
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+import sqlmodel
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision: str = '794baf50635b'
+down_revision: Union[str, Sequence[str], None] = '961802a2f50d'
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    """Upgrade schema."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table('chatthreadsnapshot',
+    sa.Column('id', sa.Integer(), nullable=False),
+    sa.Column('scope', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
+    sa.Column('thread_id', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
+    sa.Column('messages', sa.JSON().with_variant(postgresql.JSONB(astext_type=sa.Text()), 'postgresql'), server_default='[]', nullable=False),
+    sa.Column('state', sa.JSON().with_variant(postgresql.JSONB(astext_type=sa.Text()), 'postgresql'), nullable=True),
+    sa.Column('interrupt', sa.JSON().with_variant(postgresql.JSONB(astext_type=sa.Text()), 'postgresql'), nullable=True),
+    sa.Column('updated_at', sa.DateTime(), nullable=False),
+    sa.PrimaryKeyConstraint('id'),
+    sa.UniqueConstraint('scope', 'thread_id')
+    )
+    op.create_index(op.f('ix_chatthreadsnapshot_scope'), 'chatthreadsnapshot', ['scope'], unique=False)
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    """Downgrade schema."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_index(op.f('ix_chatthreadsnapshot_scope'), table_name='chatthreadsnapshot')
+    op.drop_table('chatthreadsnapshot')
+    # ### end Alembic commands ###

backend/app/agent/snapshot_store.py

@@ -0,0 +1,185 @@
+"""Postgres-backed AG-UI thread snapshot store.
+
+Implements the AGUIThreadSnapshotStore protocol from agent-framework-ag-ui so
+the assistant conversation survives page reloads: MAF calls `save` with the
+full cumulative message list at each run end, and `get` on the next request
+(server-side context rebuild, or a client hydration run with empty messages).
+
+Two deliberate choices:
+- Own sessions per call: MAF saves during SSE streaming, after the
+  agent_auth_and_session middleware has already closed the request-bound
+  ContextVar session. Sync SQLModel work runs in a thread so the event loop
+  keeps streaming.
+- `save` normalizes assistant `tool_calls` -> `toolCalls`: @ag-ui/client's
+  Zod schema silently strips the snake_case field, so a replayed snapshot
+  would lose tool calls (and the spending-summary card) on the way to the
+  browser. MAF reads both casings back, so storing camelCase is safe.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from typing import Any
+
+from agent_framework_ag_ui import AGUIThreadSnapshot
+from sqlmodel import Session, select
+
+from app.db import engine
+from app.models.models import ChatThreadSnapshot
+from app.timeutil import utcnow
+
+
+def _normalize_message(msg: dict[str, Any]) -> dict[str, Any]:
+    out = dict(msg)
+    tool_calls = out.pop("tool_calls", None)
+    if tool_calls is None:
+        tool_calls = out.pop("toolCalls", None)
+    else:
+        out.pop("toolCalls", None)
+    # MAF writes tool_calls: null on plain-text assistant messages; only a
+    # non-empty list is meaningful downstream.
+    if isinstance(tool_calls, list) and tool_calls:
+        out["toolCalls"] = tool_calls
+    if "tool_call_id" in out:
+        out.setdefault("toolCallId", out["tool_call_id"])
+        del out["tool_call_id"]
+    return out
+
+
+def _clean_messages(messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
+    """Sanitize the cumulative snapshot before storing it.
+
+    MAF's history merge re-records the resent suffix of multi-run turns
+    (frontend-tool roundtrips) without deduplicating, and the model
+    occasionally re-issues an identical render tool call on the follow-up
+    run. Left as-is, the snapshot grows duplicates that replay as repeated
+    bubbles/cards after a reload.
+    """
+    # Pass 1: drop BFF-injected ctx-* system messages, duplicate message
+    # ids, and repeated identical tool calls (same function + arguments)
+    # within one user turn.
+    seen_ids: set[str] = set()
+    turn_calls: set[tuple[Any, Any]] = set()
+    kept_call_ids: set[str] = set()
+    first_pass: list[dict[str, Any]] = []
+    for raw in messages:
+        m = _normalize_message(raw)
+        mid = str(m.get("id", ""))
+        if mid.startswith("ctx-"):
+            continue
+        if mid and mid in seen_ids:
+            continue
+        role = m.get("role")
+        if role == "user":
+            turn_calls = set()
+        if role == "assistant" and isinstance(m.get("toolCalls"), list):
+            kept: list[dict[str, Any]] = []
+            for tc in m["toolCalls"]:
+                fn = tc.get("function") or {}
+                key = (fn.get("name"), fn.get("arguments"))
+                if key in turn_calls:
+                    continue
+                turn_calls.add(key)
+                kept.append(tc)
+                if tc.get("id"):
+                    kept_call_ids.add(str(tc["id"]))
+            if kept:
+                m["toolCalls"] = kept
+            else:
+                m.pop("toolCalls", None)
+                if not m.get("content"):
+                    continue  # nothing left worth replaying
+        if mid:
+            seen_ids.add(mid)
+        first_pass.append(m)
+
+    # Pass 2: one tool result per surviving call id. Results whose call was
+    # deduplicated away must go too — an orphan tool message in the stored
+    # history would 400 the next OpenAI request when MAF rebuilds context.
+    answered: set[str] = set()
+    out: list[dict[str, Any]] = []
+    for m in first_pass:
+        if m.get("role") == "tool":
+            tcid = str(m.get("toolCallId") or "")
+            if tcid and (tcid not in kept_call_ids or tcid in answered):
+                continue
+            if tcid:
+                answered.add(tcid)
+        out.append(m)
+    return out
+
+
+class PostgresSnapshotStore:
+    """One upserted row per (scope, thread_id); latest snapshot only."""
+
+    async def save(
+        self, *, scope: str, thread_id: str, snapshot: AGUIThreadSnapshot
+    ) -> None:
+        def _save() -> None:
+            with Session(engine) as session:
+                row = session.exec(
+                    select(ChatThreadSnapshot).where(
+                        ChatThreadSnapshot.scope == scope,
+                        ChatThreadSnapshot.thread_id == thread_id,
+                    )
+                ).first()
+                if row is None:
+                    row = ChatThreadSnapshot(scope=scope, thread_id=thread_id)
+                row.messages = _clean_messages(snapshot.messages)
+                row.state = snapshot.state
+                row.interrupt = snapshot.interrupt
+                row.updated_at = utcnow()
+                session.add(row)
+                session.commit()
+
+        await asyncio.to_thread(_save)
+
+    async def get(
+        self, *, scope: str, thread_id: str
+    ) -> AGUIThreadSnapshot | None:
+        def _get() -> AGUIThreadSnapshot | None:
+            with Session(engine) as session:
+                row = session.exec(
+                    select(ChatThreadSnapshot).where(
+                        ChatThreadSnapshot.scope == scope,
+                        ChatThreadSnapshot.thread_id == thread_id,
+                    )
+                ).first()
+                if row is None:
+                    return None
+                return AGUIThreadSnapshot(
+                    messages=row.messages or [],
+                    state=row.state,
+                    interrupt=row.interrupt,
+                )
+
+        return await asyncio.to_thread(_get)
+
+    async def delete(self, *, scope: str, thread_id: str) -> bool:
+        def _delete() -> bool:
+            with Session(engine) as session:
+                row = session.exec(
+                    select(ChatThreadSnapshot).where(
+                        ChatThreadSnapshot.scope == scope,
+                        ChatThreadSnapshot.thread_id == thread_id,
+                    )
+                ).first()
+                if row is None:
+                    return False
+                session.delete(row)
+                session.commit()
+                return True
+
+        return await asyncio.to_thread(_delete)
+
+    async def clear(self, *, scope: str | None = None) -> None:
+        def _clear() -> None:
+            with Session(engine) as session:
+                stmt = select(ChatThreadSnapshot)
+                if scope is not None:
+                    stmt = stmt.where(ChatThreadSnapshot.scope == scope)
+                for row in session.exec(stmt).all():
+                    session.delete(row)
+                session.commit()
+
+        await asyncio.to_thread(_clear)

backend/app/api/v1/endpoints/chat.py

@@ -0,0 +1,41 @@
+"""Assistant chat thread management.
+
+The conversation itself is persisted by the AG-UI snapshot store
+(app/agent/snapshot_store.py) under a single ("default", "main") thread;
+this router only exposes the reset used by "Nueva conversación".
+"""
+
+from fastapi import APIRouter, Depends
+from pydantic import BaseModel
+from sqlmodel import Session, select
+
+from app.auth import get_current_user
+from app.db import get_session
+from app.models.models import ChatThreadSnapshot
+
+router = APIRouter(prefix="/chat", tags=["chat"])
+
+CHAT_SCOPE = "default"
+CHAT_THREAD_ID = "main"
+
+
+class ClearThreadResponse(BaseModel):
+    cleared: bool
+
+
+@router.delete("/thread", response_model=ClearThreadResponse)
+def clear_thread(
+    session: Session = Depends(get_session),
+    _user: str = Depends(get_current_user),
+):
+    row = session.exec(
+        select(ChatThreadSnapshot).where(
+            ChatThreadSnapshot.scope == CHAT_SCOPE,
+            ChatThreadSnapshot.thread_id == CHAT_THREAD_ID,
+        )
+    ).first()
+    if row is None:
+        return {"cleared": False}
+    session.delete(row)
+    session.commit()
+    return {"cleared": True}

backend/app/api/v1/router.py

@@ -6,6 +6,7 @@ from app.api.v1.endpoints import (
     auth,
     budget,
     categories,
+    chat,
     exchange_rate,
     import_transactions,
     installments,
@@ -38,3 +39,4 @@ api_router.include_router(pensions.router)
 api_router.include_router(municipal_receipts.router)
 api_router.include_router(savings_accrual.router)
 api_router.include_router(sync_status.router)
+api_router.include_router(chat.router)

backend/app/main.py

@@ -11,6 +11,7 @@ from jose import JWTError, jwt
 from pydantic import BaseModel
 
 from app.agent.agent import build_agent
+from app.agent.snapshot_store import PostgresSnapshotStore
 from app.agent.tools import reset_session, set_session
 from app.api.v1.router import api_router
 from app.auth import (
@@ -148,8 +149,17 @@ async def agent_auth_and_session(request: Request, call_next):
 # Register app routes
 app.include_router(api_router)
 
-# Mount the AG-UI agent endpoint.
-add_agent_framework_fastapi_endpoint(app, build_agent(), AGENT_PATH)
+# Mount the AG-UI agent endpoint. The snapshot store persists the chat
+# thread across page reloads; scope is a constant because this is a
+# single-user app and auth is already enforced by agent_auth_and_session
+# (the scope resolver only sees the AGUIRequest, which carries no identity).
+add_agent_framework_fastapi_endpoint(
+    app,
+    build_agent(),
+    AGENT_PATH,
+    snapshot_store=PostgresSnapshotStore(),
+    snapshot_scope_resolver=lambda _request: "default",
+)
 
 
 @app.get("/")

backend/app/models/models.py

@@ -537,3 +537,32 @@ class WaterMeterReading(WaterMeterReadingBase, table=True):
 class WaterMeterReadingRead(WaterMeterReadingBase):
     id: int
     created_at: datetime
+
+
+# --- Assistant Chat Thread ---
+
+
+class ChatThreadSnapshot(SQLModel, table=True):
+    """Latest AG-UI thread snapshot per (scope, thread_id).
+
+    Written by app/agent/snapshot_store.py: MAF overwrites the full
+    cumulative message list at each run end — one row per thread, not an
+    append log.
+    """
+
+    __table_args__ = (UniqueConstraint("scope", "thread_id"),)
+
+    id: Optional[int] = Field(default=None, primary_key=True)
+    scope: str = Field(index=True)
+    thread_id: str
+    messages: list = Field(
+        default_factory=list,
+        sa_column=Column(JSON_OR_JSONB, nullable=False, server_default="[]"),
+    )
+    state: Optional[dict] = Field(
+        default=None, sa_column=Column(JSON_OR_JSONB, nullable=True)
+    )
+    interrupt: Optional[list] = Field(
+        default=None, sa_column=Column(JSON_OR_JSONB, nullable=True)
+    )
+    updated_at: datetime = Field(default_factory=utcnow)

backend/tests/test_chat_persistence.py

@@ -0,0 +1,131 @@
+"""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