state-hub/api/routers/message_feed.py

"""Inbox change feed (STATE-WP-0098-T05).

Read-only pull feed for derived inbox copies (hub-core). Spec:
docs/inbox-change-feed-contract.md §2, §3, §5. Revisions are assigned by the
database trigger from migration c4d5e6f7a8b9; one write transaction = one
revision, and an advisory lock makes commit order equal revision order, so the
maximum visible revision is a safe watermark.
"""
from __future__ import annotations

import hashlib
import json
import uuid
from datetime import datetime
from typing import Any

from fastapi import APIRouter, Depends, Query, Response
from fastapi.responses import JSONResponse
from sqlalchemy import text
from sqlalchemy.ext.asyncio import AsyncSession

from api.config import settings
from api.database import get_session
from api.schemas.inbox_feed import (
    CONTRACT,
    CursorError,
    FeedChange,
    FeedPage,
    FeedSnapshot,
    ForeignCursorError,
    decode_cursor,
    encode_cursor,
)

router = APIRouter(prefix="/messages/feed", tags=["messages", "inbox-feed"])

# Contract D1: tombstones are served for 90 days.
TOMBSTONE_HORIZON_DAYS = 90
DEFAULT_LIMIT = 500
MAX_LIMIT = 5000
NO_STORE = {"Cache-Control": "no-store"}

_MESSAGE_COLUMNS = (
    "id, from_agent, to_agent, subject, body, thread_id, read_at, archived_at, "
    "kind, expires_at, supersedes_id, created_at, revision"
)
_RECEIPT_COLUMNS = "message_id, agent, delivered_at, read_at, acknowledged_at, revision"
_TOMBSTONE_COLUMNS = "message_id, revision, deleted_at, reason"

# Revisions <= this value carry tombstones that fell out of the horizon; a
# cursor below it can no longer be served losslessly (410).
_COMPACTED_SQL = f"""
SELECT coalesce(max(revision), 0) FROM message_tombstones
WHERE deleted_at < now() - interval '{TOMBSTONE_HORIZON_DAYS} days'
"""
_HIGH_WATER_SQL = """
SELECT greatest(
    (SELECT coalesce(max(revision), 0) FROM agent_messages),
    (SELECT coalesce(max(revision), 0) FROM message_receipts),
    (SELECT coalesce(max(revision), 0) FROM message_tombstones)
)
"""
_REVISIONS_AFTER_SQL = """
SELECT revision FROM agent_messages WHERE revision > :after
UNION ALL
SELECT revision FROM message_receipts WHERE revision > :after
UNION ALL
SELECT revision FROM message_tombstones WHERE revision > :after AND revision > :compacted
"""


def _jsonable(value: Any) -> Any:
    if isinstance(value, uuid.UUID):
        return str(value)
    if isinstance(value, datetime):
        return value.isoformat()
    return value


def _row(mapping) -> dict[str, Any]:
    return {key: _jsonable(value) for key, value in mapping.items()}


def receipt_key(state: dict[str, Any]) -> str:
    return f"{state['message_id']}:{state['agent']}"


def content_hash(messages, receipts, tombstones) -> str:
    payload = json.dumps(
        {"messages": messages, "receipts": receipts, "tombstones": tombstones},
        sort_keys=True,
        separators=(",", ":"),
    )
    return hashlib.sha256(payload.encode()).hexdigest()


async def _begin_consistent_read(session: AsyncSession) -> None:
    """Start one REPEATABLE READ READ ONLY transaction for the whole response."""
    await session.connection(
        execution_options={"isolation_level": "REPEATABLE READ", "postgresql_readonly": True}
    )


async def _scalar(session: AsyncSession, sql: str, **params) -> Any:
    return (await session.execute(text(sql), params)).scalar_one()


async def _rows(session: AsyncSession, sql: str, **params) -> list[dict[str, Any]]:
    return [_row(m) for m in (await session.execute(text(sql), params)).mappings().all()]


@router.get("/snapshot", response_model=FeedSnapshot)
async def feed_snapshot(
    response: Response, session: AsyncSession = Depends(get_session)
) -> FeedSnapshot:
    response.headers.update(NO_STORE)
    await _begin_consistent_read(session)
    source_time = await _scalar(session, "SELECT now()")
    compacted = await _scalar(session, _COMPACTED_SQL)
    high_water = await _scalar(session, _HIGH_WATER_SQL)
    messages = await _rows(
        session, f"SELECT {_MESSAGE_COLUMNS} FROM agent_messages ORDER BY id"
    )
    receipts = await _rows(
        session,
        f"SELECT {_RECEIPT_COLUMNS} FROM message_receipts ORDER BY message_id, agent",
    )
    tombstones = await _rows(
        session,
        f"SELECT {_TOMBSTONE_COLUMNS} FROM message_tombstones "
        "WHERE revision > :compacted ORDER BY message_id",
        compacted=compacted,
    )
    return FeedSnapshot(
        source_instance=settings.state_hub_instance_label,
        source_time=source_time,
        cursor=encode_cursor(high_water),
        compacted_before=encode_cursor(compacted),
        messages=messages,
        receipts=receipts,
        tombstones=tombstones,
        count=len(messages),
        content_hash=content_hash(messages, receipts, tombstones),
    )


def _error(status_code: int, detail: str, **extra) -> JSONResponse:
    return JSONResponse(
        status_code=status_code,
        content={"contract": CONTRACT, "detail": detail, **extra},
        headers=NO_STORE,
    )


@router.get(
    "",
    response_model=FeedPage,
    responses={400: {}, 409: {}, 410: {}},
)
async def feed_page(
    response: Response,
    after: str = Query(..., description="Cursor from a snapshot or previous page"),
    limit: int = Query(DEFAULT_LIMIT, ge=1, le=MAX_LIMIT),
    source_instance: str | None = Query(
        None, description="source_instance the cursor came from; 409 if foreign"
    ),
    session: AsyncSession = Depends(get_session),
):
    response.headers.update(NO_STORE)
    instance = settings.state_hub_instance_label
    try:
        after_rev = decode_cursor(after)
    except ForeignCursorError as exc:
        return _error(409, str(exc), resnapshot_required=True)
    except CursorError as exc:
        return _error(400, str(exc))
    if source_instance is not None and source_instance != instance:
        return _error(
            409,
            f"cursor from source_instance {source_instance!r}, this is {instance!r}",
            resnapshot_required=True,
            source_instance=instance,
        )

    await _begin_consistent_read(session)
    source_time = await _scalar(session, "SELECT now()")
    compacted = await _scalar(session, _COMPACTED_SQL)
    high_water = await _scalar(session, _HIGH_WATER_SQL)
    if after_rev > high_water:
        # A cursor from the future: another instance, or a restored source.
        return _error(
            409,
            "cursor is ahead of this source's high_water; re-snapshot",
            resnapshot_required=True,
            source_instance=instance,
            high_water=encode_cursor(high_water),
        )
    if after_rev < compacted:
        return _error(
            410,
            "cursor is older than the tombstone horizon; re-snapshot",
            resnapshot_required=True,
            source_instance=instance,
            compacted_before=encode_cursor(compacted),
        )

    params = {"after": after_rev, "compacted": compacted}
    boundary = await _scalar(
        session,
        f"SELECT max(revision) FROM (SELECT revision FROM ({_REVISIONS_AFTER_SQL}) r "
        "ORDER BY revision LIMIT :limit) page",
        limit=limit,
        **params,
    )
    changes: list[FeedChange] = []
    has_more = False
    if boundary is not None:
        # A page never splits a revision: take every change up to the boundary.
        window = {**params, "upto": boundary}
        for state in await _rows(
            session,
            f"SELECT {_MESSAGE_COLUMNS} FROM agent_messages "
            "WHERE revision > :after AND revision <= :upto "
            "ORDER BY revision, created_at, id",
            **window,
        ):
            changes.append(
                FeedChange(
                    revision=state["revision"], entity="message",
                    key=state["id"], op="upsert", state=state,
                )
            )
        for state in await _rows(
            session,
            f"SELECT {_RECEIPT_COLUMNS} FROM message_receipts "
            "WHERE revision > :after AND revision <= :upto "
            "ORDER BY revision, message_id, agent",
            **window,
        ):
            changes.append(
                FeedChange(
                    revision=state["revision"], entity="receipt",
                    key=receipt_key(state), op="upsert", state=state,
                )
            )
        for state in await _rows(
            session,
            f"SELECT {_TOMBSTONE_COLUMNS} FROM message_tombstones "
            "WHERE revision > :after AND revision > :compacted AND revision <= :upto "
            "ORDER BY revision, message_id",
            **window,
        ):
            changes.append(
                FeedChange(
                    revision=state["revision"], entity="tombstone",
                    key=state["message_id"], op="delete", state=state,
                )
            )
        # Stable sort: revision order; within a revision messages, receipts, tombstones.
        changes.sort(key=lambda change: change.revision)
        has_more = (
            await _scalar(
                session,
                f"SELECT exists(SELECT 1 FROM ({_REVISIONS_AFTER_SQL}) r "
                "WHERE revision > :upto)",
                **{**params, "upto": boundary},
            )
        )
    return FeedPage(
        source_instance=instance,
        source_time=source_time,
        high_water=encode_cursor(high_water),
        next_cursor=encode_cursor(boundary if boundary is not None else after_rev),
        has_more=has_more,
        compacted_before=encode_cursor(compacted),
        changes=changes,
    )