WealthySmart/backend/app/main.py

import asyncio
import json
import uuid
from http.cookies import SimpleCookie
from contextlib import asynccontextmanager

from agent_framework_ag_ui import add_agent_framework_fastapi_endpoint
from fastapi import FastAPI, Request, Response
from fastapi.middleware.cors import CORSMiddleware
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 (
    ALGORITHM,
    check_login_rate_limit,
    create_access_token,
    verify_admin_credentials,
)
from app.config import settings
from app.db import get_session, run_alembic_upgrade
from app.seed import seed_db
from app.timeutil import reset_client_timezone, set_client_timezone
from app.services.exchange_rate import refresh_rates_periodically


AGENT_PATH = "/api/v1/agent/agui"


def _pair_orphan_tool_calls(messages: list) -> list:
    """Inject synthetic tool responses for any assistant tool_calls that have
    no matching tool message.  OpenAI rejects histories where a tool_calls
    entry is not immediately followed by the corresponding tool response."""
    out: list = []
    pending: list[str] = []

    def flush():
        for call_id in pending:
            out.append({"role": "tool", "tool_call_id": call_id, "content": ""})
        pending.clear()

    for msg in messages:
        role = msg.get("role", "")
        if role == "tool":
            call_id = msg.get("tool_call_id") or msg.get("toolCallId")
            if call_id and call_id in pending:
                pending.remove(call_id)
            out.append(msg)
            continue
        if role == "assistant":
            flush()
            out.append(msg)
            for tc in msg.get("tool_calls") or msg.get("toolCalls") or []:
                tc_id = tc.get("id") if isinstance(tc, dict) else None
                if tc_id:
                    pending.append(tc_id)
            continue
        flush()
        out.append(msg)

    flush()
    return out


def _drop_stale_client_history(session, thread_id, messages: list) -> list:
    """Conversation history is server-owned (the AG-UI snapshot store).

    A request that carries assistant/tool history for a thread with NO
    stored snapshot comes from a stale client — typically a second tab that
    was open when 'Nueva conversación' cleared the thread. Trusting it
    would re-persist the cleared conversation ('chunks' resurrecting after
    every clear). Keep only the last user turn; MAF starts the thread
    fresh from there.
    """
    if not thread_id or not messages:
        return messages
    has_history = any(m.get("role") in ("assistant", "tool") for m in messages)
    if not has_history:
        return messages
    from sqlmodel import select

    from app.models.models import ChatThreadSnapshot

    row = session.exec(
        select(ChatThreadSnapshot.id).where(
            ChatThreadSnapshot.thread_id == str(thread_id)
        )
    ).first()
    if row is not None:
        return messages  # thread exists; history is legitimate
    last_user = None
    for i in range(len(messages) - 1, -1, -1):
        if messages[i].get("role") == "user":
            last_user = i
            break
    return messages[last_user:] if last_user is not None else []


@asynccontextmanager
async def lifespan(app: FastAPI):
    run_alembic_upgrade()
    seed_db()
    rate_refresh_task = asyncio.create_task(refresh_rates_periodically())
    try:
        yield
    finally:
        rate_refresh_task.cancel()
        try:
            await rate_refresh_task
        except asyncio.CancelledError:
            pass


app = FastAPI(title="WealthySmart API", version="0.1.0", lifespan=lifespan)

app.add_middleware(
    CORSMiddleware,
    allow_origins=settings.cors_origins_list,
    allow_credentials=True,
    allow_methods=["GET", "POST", "PATCH", "PUT", "DELETE", "OPTIONS"],
    allow_headers=["content-type", "authorization"],
)


@app.middleware("http")
async def agent_auth_and_session(request: Request, call_next):
    """For the AG-UI route, validate the JWT, repair message history, and
    bind a DB session to a ContextVar so agent tools can query without going
    through Depends."""
    if not request.url.path.startswith(AGENT_PATH):
        return await call_next(request)

    if request.method == "OPTIONS":
        return await call_next(request)

    auth_header = request.headers.get("authorization", "")
    token: str | None = None
    if auth_header.lower().startswith("bearer "):
        token = auth_header.split(" ", 1)[1].strip()
    else:
        cookies = SimpleCookie()
        cookies.load(request.headers.get("cookie", ""))
        if "ws_token" in cookies:
            token = cookies["ws_token"].value

    if not token:
        return Response(status_code=401, content="Missing auth")
    try:
        payload = jwt.decode(token, settings.SECRET_KEY, algorithms=[ALGORITHM])
        if not payload.get("sub"):
            return Response(status_code=401, content="Invalid token")
    except JWTError:
        return Response(status_code=401, content="Invalid token")

    session_gen = get_session()
    session = next(session_gen)
    token_var = set_session(session)

    # Repair orphan tool_calls before the MAF agent sees the message history.
    if request.method == "POST" and "application/json" in request.headers.get("content-type", ""):
        raw = await request.body()
        try:
            body = json.loads(raw)
            if isinstance(body.get("messages"), list):
                body["messages"] = _drop_stale_client_history(
                    session,
                    body.get("threadId") or body.get("thread_id"),
                    body["messages"],
                )
                body["messages"] = _pair_orphan_tool_calls(body["messages"])
                raw = json.dumps(body).encode()
        except Exception:
            pass
        # Starlette caches the body; replace it so call_next sees the fixed bytes.
        request._body = raw  # type: ignore[attr-defined]
    # Browser's IANA timezone, forwarded by the BFF — lets tools resolve
    # "today" wherever the user is (falls back to Costa Rica).
    tz_token = set_client_timezone(request.headers.get("x-client-timezone"))
    try:
        return await call_next(request)
    finally:
        reset_client_timezone(tz_token)
        reset_session(token_var)
        try:
            next(session_gen)
        except StopIteration:
            pass


# Register app routes
app.include_router(api_router)

# 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("/")
def root():
    return {"app": "WealthySmart", "version": "0.1.0"}


# ── Cookie-based auth endpoints (used by the Vite SPA) ──────────────────────

class LoginRequest(BaseModel):
    username: str
    password: str


@app.post("/api/auth/login")
def cookie_login(body: LoginRequest, request: Request, response: Response):
    check_login_rate_limit(request)
    verify_admin_credentials(body.username, body.password)
    token = create_access_token(body.username)
    response.set_cookie(
        key="ws_token",
        value=token,
        httponly=True,
        samesite="lax",
        max_age=settings.ACCESS_TOKEN_EXPIRE_MINUTES * 60,
        secure=settings.COOKIE_SECURE,
    )
    return {"ok": True}


@app.post("/api/auth/logout", status_code=204)
def cookie_logout(response: Response):
    response.delete_cookie("ws_token")


@app.get("/api/health")
def health():
    return {"ok": True}