WealthySmart/backend/tests/test_chat_audit.py

"""Chat-run audit persistence and redacted HTTP access logs."""

import asyncio
from datetime import timedelta

from agent_framework_ag_ui import AGUIThreadSnapshot
from starlette.requests import Request
from starlette.responses import Response
from sqlmodel import select

import app.agent.chat_audit as chat_audit
import app.agent.snapshot_store as snapshot_store_module
from app.agent.snapshot_store import PostgresSnapshotStore
from app.models.models import ChatRunLog
from app.timeutil import utcnow


def _audit_store(monkeypatch, engine):
    monkeypatch.setattr(chat_audit, "engine", engine)
    monkeypatch.setattr(snapshot_store_module, "engine", engine)
    return PostgresSnapshotStore()


def test_snapshot_completes_chat_run_with_tools_and_response(monkeypatch, engine, session):
    store = _audit_store(monkeypatch, engine)
    run_id = chat_audit.create_chat_run(
        request_id="req-1",
        agui_run_id="run-1",
        thread_id="main",
        messages=[{"role": "user", "content": "¿Cuál es mi saldo?"}],
        client_ip="203.0.113.7",
        user_agent="Mozilla/5.0 TestBrowser/1.0",
        browser="TestBrowser",
    )
    token = chat_audit.set_current_chat_run(run_id)
    try:
        asyncio.run(
            store.save(
                scope="default",
                thread_id="resp_123",
                snapshot=AGUIThreadSnapshot(
                    messages=[
                        {"id": "u1", "role": "user", "content": "¿Cuál es mi saldo?"},
                        {
                            "id": "a1",
                            "role": "assistant",
                            "content": None,
                            "tool_calls": [
                                {
                                    "id": "call-1",
                                    "type": "function",
                                    "function": {"name": "get_net_worth", "arguments": "{}"},
                                }
                            ],
                        },
                        {"id": "t1", "role": "tool", "tool_call_id": "call-1", "content": "{\"net_crc\": 1}"},
                        {"id": "a2", "role": "assistant", "content": "Tu saldo es ₡1."},
                    ]
                ),
            )
        )
    finally:
        chat_audit.reset_current_chat_run(token)

    row = session.get(ChatRunLog, run_id)
    assert row is not None
    assert row.status == "completed"
    assert row.agui_run_id == "run-1"
    assert row.response_thread_id == "resp_123"
    assert row.user_prompt == "¿Cuál es mi saldo?"
    assert row.client_ip == "203.0.113.7"
    assert row.user_agent == "Mozilla/5.0 TestBrowser/1.0"
    assert row.browser == "TestBrowser"
    assert row.assistant_response == "Tu saldo es ₡1."
    assert row.tool_calls == [
        {
            "id": "call-1",
            "name": "get_net_worth",
            "arguments": "{}",
            "result": '{"net_crc": 1}',
        }
    ]


def test_failure_and_retention_cleanup(monkeypatch, engine, session):
    _audit_store(monkeypatch, engine)
    failed_id = chat_audit.create_chat_run(
        request_id="req-failed", agui_run_id="run-failed", thread_id="main", messages=[]
    )
    chat_audit.mark_chat_run_failed(failed_id, "RuntimeError: provider unavailable")
    assert session.get(ChatRunLog, failed_id).status == "failed"

    expired = ChatRunLog(
        request_id="req-expired",
        model="gpt-5.6-luna",
        expires_at=utcnow() - timedelta(seconds=1),
    )
    session.add(expired)
    session.commit()

    assert chat_audit.purge_expired_chat_runs() == 1
    assert session.exec(select(ChatRunLog).where(ChatRunLog.request_id == "req-expired")).first() is None


def test_api_access_log_never_contains_request_body(monkeypatch):
    from app.main import log_api_request

    logged = []
    monkeypatch.setattr("app.main.log_event", lambda event, **fields: logged.append((event, fields)))
    scope = {
        "type": "http",
        "method": "POST",
        "path": "/api/auth/login",
        "raw_path": b"/api/auth/login",
        "query_string": b"",
        "headers": [
            (b"x-real-ip", b"203.0.113.9"),
            (b"user-agent", b"Mozilla/5.0 Chrome/138.0.0.0 Safari/537.36"),
        ],
        "scheme": "http",
        "server": ("testserver", 80),
        "client": ("testclient", 1234),
    }
    request = Request(scope)

    async def call_next(_request):
        return Response(status_code=201)

    response = asyncio.run(log_api_request(request, call_next))
    assert response.headers["X-Request-ID"]
    assert logged[0][0] == "api_request"
    fields = logged[0][1]
    assert fields["path"] == "/api/auth/login"
    assert fields["status_code"] == 201
    assert fields["client_ip"] == "203.0.113.9"
    assert fields["browser"] == "Chrome"
    assert fields["user_agent"] == "Mozilla/5.0 Chrome/138.0.0.0 Safari/537.36"
    assert "password" not in fields
    assert "body" not in fields