WealthySmart/backend/app/agent/snapshot_store.py

"""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)