From 6bce5539f74eb11c961a1af62c07782d236c4fe2 Mon Sep 17 00:00:00 2001 From: Carlos Escalante Date: Sat, 4 Jul 2026 18:09:55 -0600 Subject: [PATCH] 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 --- .../794baf50635b_chat_thread_snapshot.py | 45 +++++ backend/app/agent/snapshot_store.py | 185 ++++++++++++++++++ backend/app/api/v1/endpoints/chat.py | 41 ++++ backend/app/api/v1/router.py | 2 + backend/app/main.py | 14 +- backend/app/models/models.py | 29 +++ backend/tests/test_chat_persistence.py | 131 +++++++++++++ frontend/src/lib/api-types.gen.ts | 55 ++++++ 8 files changed, 500 insertions(+), 2 deletions(-) create mode 100644 backend/alembic/versions/794baf50635b_chat_thread_snapshot.py create mode 100644 backend/app/agent/snapshot_store.py create mode 100644 backend/app/api/v1/endpoints/chat.py create mode 100644 backend/tests/test_chat_persistence.py diff --git a/backend/alembic/versions/794baf50635b_chat_thread_snapshot.py b/backend/alembic/versions/794baf50635b_chat_thread_snapshot.py new file mode 100644 index 0000000..08b3da8 --- /dev/null +++ b/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 ### diff --git a/backend/app/agent/snapshot_store.py b/backend/app/agent/snapshot_store.py new file mode 100644 index 0000000..4cfe5eb --- /dev/null +++ b/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) diff --git a/backend/app/api/v1/endpoints/chat.py b/backend/app/api/v1/endpoints/chat.py new file mode 100644 index 0000000..9739574 --- /dev/null +++ b/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} diff --git a/backend/app/api/v1/router.py b/backend/app/api/v1/router.py index e78194d..924e0f3 100644 --- a/backend/app/api/v1/router.py +++ b/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) diff --git a/backend/app/main.py b/backend/app/main.py index 1738c08..d1b09a3 100644 --- a/backend/app/main.py +++ b/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("/") diff --git a/backend/app/models/models.py b/backend/app/models/models.py index f51db01..405ed4c 100644 --- a/backend/app/models/models.py +++ b/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) diff --git a/backend/tests/test_chat_persistence.py b/backend/tests/test_chat_persistence.py new file mode 100644 index 0000000..0f00334 --- /dev/null +++ b/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 diff --git a/frontend/src/lib/api-types.gen.ts b/frontend/src/lib/api-types.gen.ts index 2a9d6fb..8b16a25 100644 --- a/frontend/src/lib/api-types.gen.ts +++ b/frontend/src/lib/api-types.gen.ts @@ -345,6 +345,23 @@ export interface paths { patch: operations["update_category_api_v1_categories__category_id__patch"]; trace?: never; }; + "/api/v1/chat/thread": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + post?: never; + /** Clear Thread */ + delete: operations["clear_thread_api_v1_chat_thread_delete"]; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/api/v1/exchange-rate/": { parameters: { query?: never; @@ -1200,6 +1217,11 @@ export interface components { /** Name */ name?: string | null; }; + /** ClearThreadResponse */ + ClearThreadResponse: { + /** Cleared */ + cleared: boolean; + }; /** * Currency * @enum {string} @@ -3023,6 +3045,39 @@ export interface operations { }; }; }; + clear_thread_api_v1_chat_thread_delete: { + parameters: { + query?: never; + header?: { + authorization?: string | null; + }; + path?: never; + cookie?: { + ws_token?: string | null; + }; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["ClearThreadResponse"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; current_rate_api_v1_exchange_rate__get: { parameters: { query?: never;