deerflow-code/offline-backend-20260512/backend/packages/harness/deerflow/persistence/maf_roundtable/messages.py

"""Message repository for the MAF roundtable V2 (stable conversation history).

Every user-visible bubble (user / manager / seat / system / clarification /
approval / result / tool_progress) has exactly one row here. ``seq`` is a
conversation-global monotonic counter (unique per conversation) allocated via
``SELECT MAX(seq) + 1`` with ``IntegrityError`` retry — the same portable
approach as ``maf_roundtable_events`` (works on sqlite / postgres / mysql).

The message history is the durable record the frontend rebuilds the timeline
from on refresh; the SSE conversation-event stream (§4.3) is the real-time
channel. Both share the conversation-global seq space so a refresh can resume
the SSE cursor from the last persisted message/event seq.
"""

from __future__ import annotations

from datetime import UTC, datetime
from typing import Any

from sqlalchemy import func, select
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker

from deerflow.persistence.maf_roundtable.model import MafRoundtableMessageRow


def _now() -> datetime:
    return datetime.now(UTC)


def message_to_dict(row: MafRoundtableMessageRow) -> dict[str, Any]:
    """Internal serialisation — includes all columns. Never expose to the browser."""
    return {
        "id": row.id,
        "conversation_id": row.conversation_id,
        "seq": row.seq,
        "message_kind": row.message_kind,
        "role": row.role,
        "agent_id": row.agent_id,
        "run_id": row.run_id,
        "message_ref": row.message_ref,
        "content": row.content,
        "metadata": row.metadata_json,
        "created_at": row.created_at.isoformat() if isinstance(row.created_at, datetime) else row.created_at,
    }


_METADATA_WHITELIST_BY_KIND: dict[str, frozenset[str]] = {
    "user": frozenset(),
    "manager": frozenset({"intent_type", "dispatch_plan_summary"}),
    "seat": frozenset({"seat_name", "delivery_status", "attempt_number"}),
    "system": frozenset({"reason"}),
    "clarification": frozenset({"clarification_id", "options", "allow_multiple", "status"}),
    "approval": frozenset({"approval_id", "tool_name", "risk_level", "summary", "status"}),
    "result": frozenset({"result_type", "has_controlled_gap", "gap_summary"}),
    "tool_progress": frozenset({"tool_name", "progress_pct", "status"}),
}


def message_public_to_dict(row: MafRoundtableMessageRow) -> dict[str, Any]:
    """Public DTO — only display-safe fields. No run_id, no unfiltered metadata."""
    raw_meta = row.metadata_json
    safe_meta: dict[str, Any] | None = None
    if isinstance(raw_meta, dict):
        allowed = _METADATA_WHITELIST_BY_KIND.get(row.message_kind, frozenset())
        safe_meta = {k: v for k, v in raw_meta.items() if k in allowed} or None
    return {
        "id": row.id,
        "conversation_id": row.conversation_id,
        "seq": row.seq,
        "message_kind": row.message_kind,
        "role": row.role,
        "agent_id": row.agent_id,
        "message_ref": row.message_ref,
        "content": row.content,
        "metadata": safe_meta,
        "created_at": row.created_at.isoformat() if isinstance(row.created_at, datetime) else row.created_at,
    }


class MafMessageRepository:
    """Conversation-scoped message history with a monotonic per-conversation seq."""

    def __init__(self, session_factory: async_sessionmaker[AsyncSession]) -> None:
        self._sf = session_factory

    async def append_message(
        self,
        *,
        conversation_id: str,
        message_kind: str,
        role: str,
        content: str = "",
        agent_id: str | None = None,
        run_id: str | None = None,
        message_ref: str | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> dict[str, Any]:
        """Reject an unsafe standalone write.

        Message and event writes must use
        :meth:`MafRoundtableStore.append_conversation_item` inside one
        transaction.  Separate ``MAX(seq) + 1`` allocators would violate the
        conversation's single SSE cursor.
        """
        del conversation_id, message_kind, role, content, agent_id, run_id, message_ref, metadata
        raise RuntimeError(
            "Conversation message writes must use "
            "MafRoundtableStore.append_conversation_item"
        )

    async def list_messages(
        self,
        conversation_id: str,
        *,
        after_seq: int = 0,
        before_seq: int | None = None,
        limit: int = 50,
    ) -> list[dict[str, Any]]:
        """Paginated message history in ascending seq order.

        ``after_seq`` / ``before_seq`` implement cursor paging: the frontend
        loads the latest page on entry (``after_seq=0`` + reverse via
        :meth:`list_messages_desc`) and older pages on scroll-up.
        """
        stmt = (
            select(MafRoundtableMessageRow)
            .where(
                MafRoundtableMessageRow.conversation_id == conversation_id,
                MafRoundtableMessageRow.seq > after_seq,
            )
        )
        if before_seq is not None:
            stmt = stmt.where(MafRoundtableMessageRow.seq < before_seq)
        stmt = stmt.order_by(MafRoundtableMessageRow.seq.asc()).limit(limit)
        async with self._sf() as session:
            rows = (await session.execute(stmt)).scalars().all()
            return [message_to_dict(r) for r in rows]

    async def list_messages_desc(
        self, conversation_id: str, *, limit: int = 50, before_seq: int | None = None
    ) -> list[dict[str, Any]]:
        """Latest ``limit`` messages, newest-first (for initial page load).

        Returned newest-first; callers that want chronological order reverse
        the list. ``before_seq`` loads the next older page.
        """
        stmt = select(MafRoundtableMessageRow).where(
            MafRoundtableMessageRow.conversation_id == conversation_id
        )
        if before_seq is not None:
            stmt = stmt.where(MafRoundtableMessageRow.seq < before_seq)
        stmt = stmt.order_by(MafRoundtableMessageRow.seq.desc()).limit(limit)
        async with self._sf() as session:
            rows = (await session.execute(stmt)).scalars().all()
            return [message_to_dict(r) for r in rows]

    async def latest_seq(self, conversation_id: str) -> int:
        async with self._sf() as session:
            max_seq = (
                await session.execute(
                    select(func.max(MafRoundtableMessageRow.seq)).where(
                        MafRoundtableMessageRow.conversation_id == conversation_id
                    )
                )
            ).scalar_one()
            return int(max_seq or 0)

    async def count_messages(self, conversation_id: str) -> int:
        async with self._sf() as session:
            return int((
                await session.execute(
                    select(func.count()).select_from(MafRoundtableMessageRow).where(
                        MafRoundtableMessageRow.conversation_id == conversation_id
                    )
                )
            ).scalar_one())