DocsGPT/tests/storage/db/test_migration_0029.py

"""Migration round-trip + persistence tests for 0029_agent_guardrails."""

from __future__ import annotations

import os
import subprocess
import sys
from pathlib import Path

import pytest
from sqlalchemy import text


pytestmark = pytest.mark.integration

_0028 = "0028_user_logs_agent_lookup_idx"
_0029 = "0029_agent_guardrails"


def _alembic_ini() -> Path:
    return Path(__file__).resolve().parents[3] / "docsgpt" / "alembic.ini"


def _run_alembic(url: str, *args: str) -> None:
    subprocess.check_call(
        [sys.executable, "-m", "alembic", "-c", str(_alembic_ini()), *args],
        timeout=60,
        env={**os.environ, "POSTGRES_URI": url},
    )


def _alembic_history(url: str) -> str:
    return subprocess.check_output(
        [sys.executable, "-m", "alembic", "-c", str(_alembic_ini()), "history"],
        timeout=60,
        env={**os.environ, "POSTGRES_URI": url},
        text=True,
    )


def _alembic_heads(url: str) -> list[str]:
    out = subprocess.check_output(
        [sys.executable, "-m", "alembic", "-c", str(_alembic_ini()), "heads"],
        timeout=60,
        env={**os.environ, "POSTGRES_URI": url},
        text=True,
    )
    return [line for line in out.splitlines() if line.strip()]


def _column_exists(conn, table: str, column: str) -> bool:
    return bool(
        conn.execute(
            text(
                "SELECT 1 FROM information_schema.columns "
                "WHERE table_name = :t AND column_name = :c"
            ),
            {"t": table, "c": column},
        ).scalar()
    )


def _table_exists(conn, table: str) -> bool:
    return bool(
        conn.execute(
            text("SELECT to_regclass(:t)"), {"t": f"public.{table}"}
        ).scalar()
    )


class TestMigration0029RoundTrip:
    def test_single_head(self, pg_engine):
        """One head, and 0029 still on the path to it.

        Deliberately does NOT assert 0029 *is* the head — every later
        migration would break that — only that the chain stays linear and
        0029 has not been orphaned off it.
        """
        url = pg_engine.url.render_as_string(hide_password=False)
        heads = _alembic_heads(url)
        assert len(heads) == 1, f"expected one alembic head, got {heads}"
        assert _0029 in _alembic_history(url)

    def test_upgrade_creates_column_and_table(self, pg_engine):
        with pg_engine.connect() as conn:
            assert _column_exists(conn, "agents", "config")
            assert _table_exists(conn, "guardrail_events")

    def test_existing_agents_default_to_empty_config(self, pg_engine):
        """The server default must backfill, so old agents parse as disabled."""
        from docsgpt.guardrails.config import AgentConfig

        with pg_engine.begin() as conn:
            conn.execute(
                text(
                    "INSERT INTO agents (user_id, name, status) "
                    "VALUES ('u-29', 'legacy', 'published')"
                )
            )
            stored = conn.execute(
                text("SELECT config FROM agents WHERE user_id = 'u-29'")
            ).scalar()
        assert stored == {}
        assert AgentConfig.parse(stored).guardrails.enabled is False

    def test_downgrade_then_upgrade_is_clean(self, pg_engine):
        url = pg_engine.url.render_as_string(hide_password=False)
        _run_alembic(url, "downgrade", _0028)
        with pg_engine.connect() as conn:
            assert not _column_exists(conn, "agents", "config")
            assert not _table_exists(conn, "guardrail_events")
        _run_alembic(url, "upgrade", "head")
        with pg_engine.connect() as conn:
            assert _column_exists(conn, "agents", "config")
            assert _table_exists(conn, "guardrail_events")

    def test_events_survive_message_deletion(self, pg_engine):
        """ON DELETE SET NULL: the compliance trail outlives the conversation."""
        with pg_engine.begin() as conn:
            conv_id = conn.execute(
                text(
                    "INSERT INTO conversations (user_id, name) "
                    "VALUES ('u-29b', 'c') RETURNING id"
                )
            ).scalar()
            msg_id = conn.execute(
                text(
                    "INSERT INTO conversation_messages "
                    "(conversation_id, position, prompt, response, user_id) "
                    "VALUES (:c, 0, 'q', 'a', 'u-29b') RETURNING id"
                ),
                {"c": conv_id},
            ).scalar()
            conn.execute(
                text(
                    "INSERT INTO guardrail_events "
                    "(user_id, message_id, stage, check_name, detector_type, "
                    " action, outcome) "
                    "VALUES ('u-29b', :m, 'output', 'pii', 'PII', 'redact', "
                    "'triggered')"
                ),
                {"m": msg_id},
            )
            conn.execute(
                text("DELETE FROM conversation_messages WHERE id = :m"),
                {"m": msg_id},
            )
            row = conn.execute(
                text(
                    "SELECT message_id, check_name FROM guardrail_events "
                    "WHERE user_id = 'u-29b'"
                )
            ).fetchone()
        assert row is not None, "the audit row must survive"
        assert row[0] is None
        assert row[1] == "pii"


class TestAgentConfigPersistence:
    def test_config_round_trips_through_the_repository(self, pg_engine):
        from docsgpt.storage.db.repositories.agents import AgentsRepository

        config = {
            "guardrails": {
                "enabled": True,
                "mode": "scan_all",
                "fail_open": False,
                "timeout_ms": 1500,
                "block_message": "Nope.",
                "controls": [
                    {
                        "check": "pii",
                        "stage": "input",
                        "action": "redact",
                        "enabled": True,
                        "settings": {"entities": ["EMAIL"]},
                    }
                ],
            }
        }
        with pg_engine.begin() as conn:
            repo = AgentsRepository(conn)
            created = repo.create("u-cfg", "a", "published", config=config)
            assert created["config"] == config, "JSONB must not double-encode"

            repo.update(str(created["id"]), "u-cfg", {"config": {"guardrails": {}}})
            reread = repo.get(str(created["id"]), "u-cfg")
        assert reread["config"] == {"guardrails": {}}

    def test_config_survives_a_parse_round_trip(self, pg_engine):
        from docsgpt.guardrails.config import AgentConfig
        from docsgpt.storage.db.repositories.agents import AgentsRepository

        raw = AgentConfig.model_validate(
            {
                "guardrails": {
                    "enabled": True,
                    "controls": [{"check": "secrets", "stage": "output",
                                  "action": "redact"}],
                }
            }
        ).model_dump(mode="json")
        with pg_engine.begin() as conn:
            repo = AgentsRepository(conn)
            created = repo.create("u-cfg2", "a", "published", config=raw)
            reread = repo.get(str(created["id"]), "u-cfg2")
        parsed = AgentConfig.parse(reread["config"]).guardrails
        assert parsed.enabled is True
        assert parsed.controls[0].check == "secrets"


class TestGuardrailEventsRepository:
    def test_record_and_read_back(self, pg_engine):
        from docsgpt.storage.db.repositories.guardrail_events import (
            GuardrailEventsRepository,
        )
        from docsgpt.storage.db.repositories.agents import AgentsRepository

        with pg_engine.begin() as conn:
            agent = AgentsRepository(conn).create("u-ev", "a", "published")
            agent_id = str(agent["id"])
            repo = GuardrailEventsRepository(conn)
            written = repo.record_many(
                [
                    {
                        "user_id": "u-ev",
                        "agent_id": agent_id,
                        "stage": "input",
                        "check_name": "denylist",
                        "detector_type": "DENYLIST",
                        "action": "block",
                        "outcome": "triggered",
                        "category": "BANNED_TERM",
                        "match_count": 2,
                        "detail": "2 banned term match(es)",
                    },
                    {
                        "user_id": "u-ev",
                        "agent_id": agent_id,
                        "stage": "output",
                        "check_name": "policy",
                        "detector_type": "POLICY",
                        "action": "flag",
                        "outcome": "not_evaluated",
                        "detail": "timeout",
                    },
                ]
            )
            assert written == 2
            events = repo.list_for_agent(agent_id, "u-ev")
            summary = repo.summary_for_user("u-ev")

        assert len(events) == 2
        assert summary["totals"]["blocked"] == 1
        assert summary["totals"]["not_evaluated"] == 1
        assert summary["totals"]["flagged"] == 0

    def test_empty_batch_is_a_noop(self, pg_engine):
        from docsgpt.storage.db.repositories.guardrail_events import (
            GuardrailEventsRepository,
        )

        with pg_engine.begin() as conn:
            assert GuardrailEventsRepository(conn).record_many([]) == 0

    def test_events_are_scoped_to_the_requesting_user(self, pg_engine):
        from docsgpt.storage.db.repositories.agents import AgentsRepository
        from docsgpt.storage.db.repositories.guardrail_events import (
            GuardrailEventsRepository,
        )

        with pg_engine.begin() as conn:
            agent = AgentsRepository(conn).create("owner", "a", "published")
            agent_id = str(agent["id"])
            repo = GuardrailEventsRepository(conn)
            repo.record_many(
                [
                    {"user_id": "owner", "agent_id": agent_id, "stage": "input",
                     "check_name": "pii", "detector_type": "PII",
                     "action": "flag", "outcome": "triggered"},
                    {"user_id": "someone-else", "agent_id": agent_id,
                     "stage": "input", "check_name": "secrets",
                     "detector_type": "SECRETS", "action": "flag",
                     "outcome": "triggered"},
                ]
            )
            mine = repo.list_for_agent(agent_id, "owner")
        assert len(mine) == 1
        assert mine[0]["check_name"] == "pii"

    def test_listing_never_exposes_the_key_or_scanned_text(self, pg_engine):
        """``api_key`` is the agent's raw key and ``matched_value`` is raw PII."""
        from docsgpt.storage.db.repositories.agents import AgentsRepository
        from docsgpt.storage.db.repositories.guardrail_events import (
            GuardrailEventsRepository,
        )

        with pg_engine.begin() as conn:
            agent = AgentsRepository(conn).create("u-leak", "a", "published")
            agent_id = str(agent["id"])
            repo = GuardrailEventsRepository(conn)
            repo.record_many(
                [{"user_id": "u-leak", "agent_id": agent_id,
                  "api_key": "super-secret-agent-key", "stage": "output",
                  "check_name": "pii", "detector_type": "PII",
                  "action": "redact", "outcome": "triggered",
                  "matched_value": "ada@example.com"}]
            )
            rows = repo.list_for_agent(agent_id, "u-leak")
        assert rows
        assert "api_key" not in rows[0]
        assert "matched_value" not in rows[0]
        assert "super-secret-agent-key" not in str(rows[0])
        assert "ada@example.com" not in str(rows[0])

    def test_purge_respects_the_window(self, pg_engine):
        from docsgpt.storage.db.repositories.guardrail_events import (
            GuardrailEventsRepository,
        )

        with pg_engine.begin() as conn:
            conn.execute(
                text(
                    "INSERT INTO guardrail_events "
                    "(user_id, stage, check_name, detector_type, action, "
                    " outcome, created_at) "
                    "VALUES ('u-purge', 'input', 'pii', 'PII', 'flag', "
                    "'triggered', NOW() - INTERVAL '90 days')"
                )
            )
            conn.execute(
                text(
                    "INSERT INTO guardrail_events "
                    "(user_id, stage, check_name, detector_type, action, "
                    " outcome) "
                    "VALUES ('u-purge', 'input', 'pii', 'PII', 'flag', "
                    "'triggered')"
                )
            )
            deleted = GuardrailEventsRepository(conn).purge_older_than(30)
            remaining = conn.execute(
                text(
                    "SELECT COUNT(*) FROM guardrail_events "
                    "WHERE user_id = 'u-purge'"
                )
            ).scalar()
        assert deleted == 1
        assert remaining == 1