state-hub/tests/test_message_feed.py

"""STATE-WP-0098-T05: inbox change feed (docs/inbox-change-feed-contract.md)."""
from __future__ import annotations

import importlib.util
import os
import uuid
from pathlib import Path

import pytest
import pytest_asyncio
import sqlalchemy
from alembic.migration import MigrationContext
from alembic.operations import Operations
from sqlalchemy import select, text
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker

from api.models.agent_message import AgentMessage, MessageReceipt
from api.services.message_receipts import upsert_receipt

CONTRACT = "statehub.inbox-feed.v1"


@pytest_asyncio.fixture(autouse=True)
async def _clean_tombstones(test_engine):
    # conftest deletes agent_messages after each test, which (correctly) leaves
    # tombstones behind; start every feed test from an empty tombstone table.
    async with test_engine.begin() as conn:
        await conn.execute(text("DELETE FROM message_tombstones"))
    yield


@pytest.fixture
def factory(test_engine):
    return async_sessionmaker(test_engine, class_=AsyncSession, expire_on_commit=False)


async def _send(client, to_agent="state-hub", **payload):
    body = {"from_agent": "the-custodian", "subject": "s", "body": "b", "to_agent": to_agent, **payload}
    resp = await client.post("/messages/", json=body)
    assert resp.status_code == 201, resp.text
    return resp.json()


async def _revisions(factory, model=AgentMessage):
    async with factory() as session:
        if model is AgentMessage:
            rows = await session.execute(select(AgentMessage.id, AgentMessage.revision))
            return {str(i): r for i, r in rows.all()}
        rows = await session.execute(
            select(MessageReceipt.message_id, MessageReceipt.agent, MessageReceipt.revision)
        )
        return {f"{m}:{a}": r for m, a, r in rows.all()}


async def _snapshot(client):
    resp = await client.get("/messages/feed/snapshot")
    assert resp.status_code == 200, resp.text
    assert resp.headers["cache-control"] == "no-store"
    return resp.json()


async def _page(client, after, limit=500, expect=200, **params):
    resp = await client.get("/messages/feed", params={"after": after, "limit": limit, **params})
    assert resp.status_code == expect, resp.text
    assert resp.headers["cache-control"] == "no-store"
    return resp.json()


async def _drain(client, after, limit=2):
    pages = []
    while True:
        page = await _page(client, after, limit=limit)
        pages.append(page)
        after = page["next_cursor"]
        if not page["has_more"]:
            return pages


def _rev(cursor):
    return int(cursor.split(":", 1)[1])


@pytest.mark.asyncio
async def test_one_revision_per_transaction(factory):
    async with factory() as session:
        a = AgentMessage(from_agent="x", to_agent="y", subject="a", body="a")
        b = AgentMessage(from_agent="x", to_agent="broadcast", subject="b", body="b", kind="news")
        session.add_all([a, b])
        await session.flush()
        await upsert_receipt(session, [b.id], "reader-1", delivered=True)
        await session.commit()
        assert a.revision == b.revision  # eager_defaults: RETURNING, no lazy load
    revs = await _revisions(factory)
    receipt_revs = await _revisions(factory, MessageReceipt)
    assert set(revs.values()) == {a.revision}
    assert set(receipt_revs.values()) == {a.revision}

    async with factory() as session:
        c = AgentMessage(from_agent="x", to_agent="y", subject="c", body="c")
        session.add(c)
        await session.commit()
    assert c.revision > a.revision


@pytest.mark.asyncio
async def test_no_revision_bump_on_noop_update(client, factory):
    direct = await _send(client)
    news = await _send(client, to_agent="broadcast")
    assert (await client.patch(f"/messages/{direct['id']}/read", json={})).status_code == 200
    resp = await client.patch(f"/messages/{news['id']}/read", json={"reader": "reader-1"})
    assert resp.status_code == 200
    before = await _revisions(factory)
    receipts_before = await _revisions(factory, MessageReceipt)

    # COALESCE upsert that changes nothing, and a raw UPDATE setting equal values.
    async with factory() as session:
        await upsert_receipt(session, [uuid.UUID(news["id"])], "reader-1", read=True, delivered=False)
        await session.execute(
            text("UPDATE agent_messages SET read_at = read_at, subject = subject")
        )
        await session.commit()
    assert await _revisions(factory) == before
    assert await _revisions(factory, MessageReceipt) == receipts_before

    # A real change bumps exactly the touched receipt.
    async with factory() as session:
        await upsert_receipt(session, [uuid.UUID(news["id"])], "reader-1", acknowledged=True)
        await session.commit()
    after = await _revisions(factory, MessageReceipt)
    key = f"{news['id']}:reader-1"
    assert after[key] > receipts_before[key]


@pytest.mark.asyncio
async def test_archive_is_first_write_wins(client, factory):
    msg = await _send(client)
    first = await client.patch(f"/messages/{msg['id']}/archive")
    assert first.status_code == 200
    rev = (await _revisions(factory))[msg["id"]]
    second = await client.patch(f"/messages/{msg['id']}/archive")
    assert second.status_code == 200
    assert second.json()["archived_at"] == first.json()["archived_at"]
    assert (await _revisions(factory))[msg["id"]] == rev


@pytest.mark.asyncio
async def test_delete_writes_tombstone_in_same_revision(client, factory):
    root = await _send(client)
    reply = (
        await client.post(f"/messages/{root['id']}/reply", json={"from_agent": "state-hub", "body": "r"})
    ).json()
    news = await _send(client, to_agent="broadcast")
    await upsert_receipt_via_api(client, news["id"])
    cursor = (await _snapshot(client))["cursor"]

    async with factory() as session:
        await session.execute(text("SET LOCAL statehub.delete_reason = 'test cleanup'"))
        await session.execute(
            text("DELETE FROM agent_messages WHERE id IN (:a, :b)"),
            {"a": root["id"], "b": news["id"]},
        )
        await session.commit()

    page = await _page(client, cursor)
    revisions = {c["revision"] for c in page["changes"]}
    assert len(revisions) == 1
    by_entity = {(c["entity"], c["key"]): c for c in page["changes"]}
    root_tomb = by_entity[("tombstone", root["id"])]
    assert root_tomb["op"] == "delete"
    assert root_tomb["state"]["reason"] == "test cleanup"
    assert "body" not in root_tomb["state"]
    assert ("tombstone", news["id"]) in by_entity
    # The cascade SET NULL on the reply arrives as an upsert in the same revision.
    reply_change = by_entity[("message", reply["id"])]
    assert reply_change["state"]["thread_id"] is None
    # Cascaded receipts are gone from the snapshot.
    snap = await _snapshot(client)
    assert snap["receipts"] == []
    assert {t["message_id"] for t in snap["tombstones"]} == {root["id"], news["id"]}


async def upsert_receipt_via_api(client, message_id):
    resp = await client.post(f"/messages/{message_id}/ack", json={"agent": "reader-1"})
    assert resp.status_code in (200, 201), resp.text


@pytest.mark.asyncio
async def test_snapshot_and_feed_parity(client, factory):
    first = await _snapshot(client)
    assert first["contract"] == CONTRACT
    assert first["messages"] == [] and first["count"] == 0

    m1 = await _send(client)
    m2 = await _send(client, to_agent="broadcast", kind="standing")
    await _send(client, to_agent="broadcast", kind="standing", supersedes_id=m2["id"])
    await client.patch(f"/messages/{m1['id']}/read", json={})
    await client.get("/messages/", params={"to_agent": "reader-2", "unread_only": "true"})
    await upsert_receipt_via_api(client, m2["id"])
    await client.post(f"/messages/{m1['id']}/reply", json={"from_agent": "state-hub", "body": "r"})
    doomed = await _send(client)
    async with factory() as session:
        await session.execute(text("DELETE FROM agent_messages WHERE id = :i"), {"i": doomed["id"]})
        await session.commit()

    snap = await _snapshot(client)
    pages = await _drain(client, "v1:0", limit=2)
    assert len(pages) > 1
    messages, receipts, tombstones = {}, {}, {}
    seen = []
    for page in pages:
        assert page["contract"] == CONTRACT
        for change in page["changes"]:
            seen.append(change["revision"])
            if change["entity"] == "message":
                messages[change["key"]] = change["state"]
            elif change["entity"] == "receipt":
                receipts[change["key"]] = change["state"]
            else:
                tombstones[change["key"]] = change["state"]
                messages.pop(change["key"], None)
                for key in [k for k in receipts if k.startswith(change["key"] + ":")]:
                    receipts.pop(key)
    assert seen == sorted(seen)
    assert pages[-1]["next_cursor"] == snap["cursor"] == pages[-1]["high_water"]
    assert sorted(messages.values(), key=lambda m: m["id"]) == snap["messages"]
    assert sorted(receipts.values(), key=lambda r: (r["message_id"], r["agent"])) == snap["receipts"]
    assert list(tombstones.values()) == snap["tombstones"]
    assert snap["count"] == len(snap["messages"]) == 4
    assert len(snap["content_hash"]) == 64
    # Same state, same hash; a change moves it.
    assert (await _snapshot(client))["content_hash"] == snap["content_hash"]
    await client.patch(f"/messages/{m2['id']}/read", json={"reader": "reader-3"})
    assert (await _snapshot(client))["content_hash"] != snap["content_hash"]


@pytest.mark.asyncio
async def test_page_never_splits_a_revision(factory, client):
    async with factory() as session:
        session.add_all(
            [AgentMessage(from_agent="x", to_agent="y", subject=str(i), body="b") for i in range(3)]
        )
        await session.commit()
    await _send(client)

    page = await _page(client, "v1:0", limit=1)
    assert len(page["changes"]) == 3
    assert len({c["revision"] for c in page["changes"]}) == 1
    assert page["has_more"] is True
    assert page["next_cursor"] == f"v1:{page['changes'][0]['revision']}"
    last = await _page(client, page["next_cursor"], limit=1)
    assert len(last["changes"]) == 1 and last["has_more"] is False
    assert last["next_cursor"] == last["high_water"]


@pytest.mark.asyncio
async def test_replay_and_empty_page_heartbeat(client):
    await _send(client)
    await _send(client)
    snap = await _snapshot(client)
    a = await _page(client, "v1:0")
    b = await _page(client, "v1:0")
    assert a["changes"] == b["changes"] and a["next_cursor"] == b["next_cursor"]

    empty = await _page(client, snap["cursor"])
    assert empty["changes"] == [] and empty["has_more"] is False
    assert empty["next_cursor"] == snap["cursor"] == empty["high_water"]
    assert empty["source_time"] and "source_instance" in empty


@pytest.mark.asyncio
async def test_cursor_older_than_tombstone_horizon_is_410(client, factory):
    await _send(client)
    async with factory() as session:
        await session.execute(
            text(
                "INSERT INTO message_tombstones (message_id, revision, deleted_at) "
                "VALUES (:i, nextval('inbox_revision_seq'), now() - interval '91 days')"
            ),
            {"i": str(uuid.uuid4())},
        )
        await session.commit()
    fresh = await _send(client)
    snap = await _snapshot(client)
    compacted = _rev(snap["compacted_before"])
    assert compacted > 0
    assert snap["tombstones"] == []  # out of horizon, not served

    gone = await _page(client, f"v1:{compacted - 1}", expect=410)
    assert gone["resnapshot_required"] is True
    ok = await _page(client, f"v1:{compacted}")
    assert [c["key"] for c in ok["changes"]] == [fresh["id"]]
    assert ok["compacted_before"] == snap["compacted_before"]


@pytest.mark.asyncio
async def test_foreign_cursor_is_409_and_malformed_is_400(client):
    await _send(client)
    snap = await _snapshot(client)
    body = await _page(client, "v2:1", expect=409)
    assert body["resnapshot_required"] is True
    await _page(client, snap["cursor"], source_instance="some-other-hub", expect=409)
    await _page(client, f"v1:{_rev(snap['cursor']) + 1000}", expect=409)
    await _page(client, "garbage", expect=400)
    if snap["source_instance"] is not None:
        await _page(client, snap["cursor"], source_instance=snap["source_instance"])


@pytest.mark.asyncio
async def test_feed_routes_do_not_collide_with_message_routes(client):
    resp = await client.post("/messages/feed/snapshot")
    assert resp.status_code == 405
    assert (await client.get("/messages/feed")).status_code == 422  # after is required


def _load_migration():
    path = Path(__file__).parent.parent / "migrations/versions/c4d5e6f7a8b9_inbox_feed_revision.py"
    spec = importlib.util.spec_from_file_location("inbox_feed_migration", path)
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


_FUNCS = "('inbox_feed_revision','inbox_feed_stamp_revision','inbox_feed_tombstone')"


def test_migration_downgrade_upgrade_and_ddl_parity():
    """Run the migration on the test DB inside a rolled-back transaction."""
    from api.models import agent_message

    migration = _load_migration()
    assert migration.FUNCTIONS_SQL == agent_message.INBOX_FEED_FUNCTIONS_SQL
    url = os.getenv(
        "TEST_DATABASE_URL", "postgresql+asyncpg://custodian:changeme@127.0.0.1:5432/custodian_test"
    ).replace("+asyncpg", "+psycopg2")
    engine = sqlalchemy.create_engine(url)
    try:
        with engine.connect() as conn:
            trans = conn.begin()
            try:
                conn.execute(
                    text(
                        "INSERT INTO agent_messages (id, from_agent, to_agent, subject, body) "
                        "VALUES (gen_random_uuid(), 'x', 'y', 's', 'b'), "
                        "(gen_random_uuid(), 'x', 'y', 's2', 'b')"
                    )
                )
                ops = Operations(MigrationContext.configure(conn))
                migration.op = ops
                migration.downgrade()
                cols = {c["name"] for c in sqlalchemy.inspect(conn).get_columns("agent_messages")}
                assert "revision" not in cols
                assert not sqlalchemy.inspect(conn).has_table("message_tombstones")
                migration.upgrade()
                revs = conn.execute(text("SELECT DISTINCT revision FROM agent_messages")).scalars().all()
                assert len(revs) == 1 and revs[0] is not None  # one backfill revision
                triggers = set(
                    conn.execute(
                        text(
                            "SELECT tgname FROM pg_trigger WHERE NOT tgisinternal "
                            "AND tgrelid IN ('agent_messages'::regclass, 'message_receipts'::regclass)"
                        )
                    ).scalars()
                )
                assert triggers == {
                    "agent_messages_inbox_revision",
                    "agent_messages_inbox_tombstone",
                    "message_receipts_inbox_revision",
                }
                conn.execute(text("DELETE FROM agent_messages"))
                assert conn.execute(text("SELECT count(*) FROM message_tombstones")).scalar() == 2
            finally:
                trans.rollback()
    finally:
        engine.dispose()


@pytest.mark.asyncio
async def test_snapshot_transaction_is_repeatable_read_read_only(factory):
    from api.routers.message_feed import _begin_consistent_read

    async with factory() as session:
        await _begin_consistent_read(session)
        assert (await session.execute(text("SHOW transaction_isolation"))).scalar() == "repeatable read"
        assert (await session.execute(text("SHOW transaction_read_only"))).scalar() == "on"


@pytest.mark.asyncio
async def test_advisory_lock_serializes_writers_in_revision_order(factory):
    import asyncio

    first = factory()
    second = factory()
    try:
        a = AgentMessage(from_agent="x", to_agent="y", subject="a", body="a")
        first.add(a)
        await first.flush()  # holds the inbox lock until commit

        async def write_b():
            b = AgentMessage(from_agent="x", to_agent="y", subject="b", body="b")
            second.add(b)
            await second.commit()
            return b

        task = asyncio.create_task(write_b())
        await asyncio.sleep(0.5)
        assert not task.done()  # blocked on pg_advisory_xact_lock
        await first.commit()
        b = await asyncio.wait_for(task, timeout=10)
        assert b.revision > a.revision
    finally:
        await first.close()
        await second.close()