state-hub/api/models/agent_message.py

import uuid
from datetime import datetime

from sqlalchemy import (
    DDL,
    BigInteger,
    DateTime,
    FetchedValue,
    ForeignKey,
    Index,
    Sequence,
    String,
    Text,
    event,
    text,
)
from sqlalchemy.dialects.postgresql import UUID
from sqlalchemy.orm import Mapped, mapped_column, relationship

from api.models.base import Base, new_uuid


BROADCAST = "broadcast"
MESSAGE_KINDS = ("message", "news", "standing")

# STATE-WP-0098-T05: inbox change feed (docs/inbox-change-feed-contract.md).
# ``revision`` is owned by the database trigger below; application code never
# sets it. Every inbox write transaction gets exactly one revision.
INBOX_REVISION_SEQ = Sequence("inbox_revision_seq", metadata=Base.metadata)


def _revision_column() -> Mapped[int]:
    return mapped_column(
        BigInteger,
        nullable=False,
        index=True,
        server_default=FetchedValue(),
        server_onupdate=FetchedValue(),
    )


class AgentMessage(Base):
    __tablename__ = "agent_messages"
    __table_args__ = (
        Index("ix_agent_messages_to_kind_archived", "to_agent", "kind", "archived_at"),
    )
    __mapper_args__ = {"eager_defaults": True}

    id: Mapped[uuid.UUID] = mapped_column(
        UUID(as_uuid=True), primary_key=True, default=new_uuid
    )
    from_agent: Mapped[str] = mapped_column(String(100), nullable=False)
    to_agent: Mapped[str] = mapped_column(String(100), nullable=False, index=True)
    subject: Mapped[str] = mapped_column(String(500), nullable=False)
    body: Mapped[str] = mapped_column(Text, nullable=False)
    thread_id: Mapped[uuid.UUID | None] = mapped_column(
        UUID(as_uuid=True),
        ForeignKey("agent_messages.id", ondelete="SET NULL"),
        nullable=True,
        index=True,
    )
    read_at: Mapped[datetime | None] = mapped_column(
        DateTime(timezone=True), nullable=True
    )
    archived_at: Mapped[datetime | None] = mapped_column(
        DateTime(timezone=True), nullable=True
    )
    # STATE-WP-0093: kind is only meaningful for broadcasts
    # (message | news | standing); direct messages stay "message".
    kind: Mapped[str] = mapped_column(
        String(20), nullable=False, default="message", server_default="message"
    )
    expires_at: Mapped[datetime | None] = mapped_column(
        DateTime(timezone=True), nullable=True
    )
    supersedes_id: Mapped[uuid.UUID | None] = mapped_column(
        UUID(as_uuid=True),
        ForeignKey("agent_messages.id", ondelete="SET NULL"),
        nullable=True,
    )
    created_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True),
        server_default=text("now()"),
        nullable=False,
    )
    revision: Mapped[int] = _revision_column()

    thread_root: Mapped["AgentMessage | None"] = relationship(
        "AgentMessage",
        remote_side="AgentMessage.id",
        foreign_keys=[thread_id],
        lazy="select",
    )


class MessageReceipt(Base):
    """Per-reader state for broadcast messages (STATE-WP-0093).

    Written only for broadcasts; direct messages keep the global
    ``AgentMessage.read_at``.
    """

    __tablename__ = "message_receipts"
    __mapper_args__ = {"eager_defaults": True}

    message_id: Mapped[uuid.UUID] = mapped_column(
        UUID(as_uuid=True),
        ForeignKey("agent_messages.id", ondelete="CASCADE"),
        primary_key=True,
    )
    agent: Mapped[str] = mapped_column(String(100), primary_key=True, index=True)
    delivered_at: Mapped[datetime | None] = mapped_column(
        DateTime(timezone=True), nullable=True
    )
    read_at: Mapped[datetime | None] = mapped_column(
        DateTime(timezone=True), nullable=True
    )
    acknowledged_at: Mapped[datetime | None] = mapped_column(
        DateTime(timezone=True), nullable=True
    )
    revision: Mapped[int] = _revision_column()


class MessageTombstone(Base):
    """A deleted ``agent_messages`` row, written by the delete trigger.

    Bodies are never kept. Served by the inbox feed within the 90-day horizon
    (contract D1).
    """

    __tablename__ = "message_tombstones"

    message_id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), primary_key=True)
    revision: Mapped[int] = mapped_column(BigInteger, nullable=False, index=True)
    deleted_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False)
    reason: Mapped[str | None] = mapped_column(Text, nullable=True)


# Trigger DDL for ``Base.metadata.create_all`` (the test suite). Production
# gets the same objects from migration c4d5e6f7a8b9_inbox_feed_revision, which
# keeps its own frozen copy; tests/test_message_feed.py checks they agree.
INBOX_ADVISORY_LOCK_KEY = 5_348_490_098

INBOX_FEED_FUNCTIONS_SQL = f"""
CREATE OR REPLACE FUNCTION inbox_feed_revision() RETURNS bigint
LANGUAGE plpgsql AS $fn$
DECLARE
    tx text := txid_current()::text;
    cached text := coalesce(current_setting('statehub.inbox_rev', true), '');
    rev bigint;
BEGIN
    PERFORM pg_advisory_xact_lock({INBOX_ADVISORY_LOCK_KEY});
    IF split_part(cached, ':', 1) = tx THEN
        RETURN split_part(cached, ':', 2)::bigint;
    END IF;
    rev := nextval('inbox_revision_seq');
    PERFORM set_config('statehub.inbox_rev', tx || ':' || rev::text, true);
    RETURN rev;
END
$fn$;

CREATE OR REPLACE FUNCTION inbox_feed_stamp_revision() RETURNS trigger
LANGUAGE plpgsql AS $fn$
BEGIN
    IF TG_OP = 'UPDATE' THEN
        IF NEW IS NOT DISTINCT FROM OLD THEN
            RETURN NEW;
        END IF;
    END IF;
    NEW.revision := inbox_feed_revision();
    RETURN NEW;
END
$fn$;

CREATE OR REPLACE FUNCTION inbox_feed_tombstone() RETURNS trigger
LANGUAGE plpgsql AS $fn$
BEGIN
    INSERT INTO message_tombstones (message_id, revision, deleted_at, reason)
    VALUES (
        OLD.id,
        inbox_feed_revision(),
        now(),
        nullif(current_setting('statehub.delete_reason', true), '')
    )
    ON CONFLICT (message_id) DO UPDATE
        SET revision = EXCLUDED.revision,
            deleted_at = EXCLUDED.deleted_at,
            reason = EXCLUDED.reason;
    RETURN OLD;
END
$fn$;
"""

event.listen(
    AgentMessage.__table__,
    "after_create",
    DDL(
        INBOX_FEED_FUNCTIONS_SQL
        + """
CREATE TRIGGER agent_messages_inbox_revision
    BEFORE INSERT OR UPDATE ON agent_messages
    FOR EACH ROW EXECUTE FUNCTION inbox_feed_stamp_revision();
CREATE TRIGGER agent_messages_inbox_tombstone
    BEFORE DELETE ON agent_messages
    FOR EACH ROW EXECUTE FUNCTION inbox_feed_tombstone();
"""
    ),
)
event.listen(
    MessageReceipt.__table__,
    "after_create",
    DDL(
        INBOX_FEED_FUNCTIONS_SQL
        + """
CREATE TRIGGER message_receipts_inbox_revision
    BEFORE INSERT OR UPDATE ON message_receipts
    FOR EACH ROW EXECUTE FUNCTION inbox_feed_stamp_revision();
"""
    ),
)