DocsGPT/docsgpt/storage/db/repositories/conversations.py

"""Repository for the ``conversations`` and ``conversation_messages`` tables.

Covers every operation the legacy Mongo code performs on
``conversations_collection``:

- create / get / list / delete conversations
- append message (transactional position allocation)
- update message at index (overwrite + optional truncation)
- set / unset feedback on a message
- rename conversation
- update compression metadata
- shared_with access checks
"""

from __future__ import annotations

import json
import logging
from enum import Enum
from typing import Optional

from sqlalchemy import Connection, text
from sqlalchemy.dialects.postgresql import insert as pg_insert

from docsgpt.storage.db.base_repository import looks_like_uuid, row_to_dict
from docsgpt.storage.db.models import conversations_table, conversation_messages_table
from docsgpt.storage.db.serialization import PGNativeJSONEncoder

logger = logging.getLogger(__name__)


class MessageUpdateOutcome(str, Enum):
    """Discriminated result of ``update_message_by_id``.

    Distinguishes the row-actually-updated case from the row-already-at-
    the-requested-terminal-state case so an abort handler can journal
    ``end`` instead of ``error`` when the normal-path finalize already
    flipped the row to ``complete``.
    """

    UPDATED = "updated"
    ALREADY_COMPLETE = "already_complete"
    ALREADY_FAILED = "already_failed"
    NOT_FOUND = "not_found"
    INVALID = "invalid"


class HeartbeatState(str, Enum):
    """What a liveness heartbeat found when it tried to stamp a row.

    ``TERMINAL`` and ``MISSING`` both mean "not stamped", but they want
    opposite handling: a terminal row may still be reclaimed by a finishing
    stream, whereas a missing one can never be written again.
    """

    STAMPED = "stamped"
    TERMINAL = "terminal"
    MISSING = "missing"


def _message_row_to_dict(row) -> dict:
    """Like ``row_to_dict`` but renames the DB column ``message_metadata``
    back to the public API key ``metadata`` so callers keep the Mongo-era
    shape. See migration 0001_initial for the column rename rationale."""
    out = row_to_dict(row)
    if "message_metadata" in out:
        out["metadata"] = out.pop("message_metadata")
    return out


class ConversationsRepository:
    def __init__(self, conn: Connection) -> None:
        self._conn = conn

    # ------------------------------------------------------------------
    # Reference translation helpers
    # ------------------------------------------------------------------
    #
    # During the Mongo→Postgres dual-write window, callers routinely
    # hand us Mongo ObjectId strings (24-char hex) for fields that are
    # UUID FKs in Postgres (``agent_id``, ``attachments`` entries, ...).
    # Casting those straight to ``uuid`` raises and the outer dual-write
    # shim swallows the exception, so the write silently drops. These
    # helpers translate via the ``legacy_mongo_id`` columns we added
    # precisely for this purpose.

    def _resolve_agent_ref(self, agent_id_raw: str | None) -> str | None:
        """Translate ``agent_id_raw`` to a Postgres UUID string.

        - ``None``/empty → ``None`` (no agent).
        - Already-UUID-shaped → returned as-is.
        - Otherwise treated as a Mongo ObjectId and looked up via
          ``agents.legacy_mongo_id``. Returns ``None`` if no PG row
          exists yet (e.g. the agent was created before the backfill
          ran).
        """
        if not agent_id_raw:
            return None
        value = str(agent_id_raw)
        if looks_like_uuid(value):
            return value
        result = self._conn.execute(
            text("SELECT id FROM agents WHERE legacy_mongo_id = :lid LIMIT 1"),
            {"lid": value},
        )
        row = result.fetchone()
        return str(row[0]) if row is not None else None

    def _resolve_attachment_refs(
        self, ids: list[str] | None,
    ) -> list[str]:
        """Translate a list of attachment ids to canonical PG
        ``attachments.id`` UUIDs.

        Inputs may be:

        - A Mongo ObjectId string (24-hex), legacy dual-write era —
          must be looked up via ``attachments.legacy_mongo_id``.
        - A UUID string that is a real ``attachments.id`` PK.
        - A UUID string that is *only* present as
          ``attachments.legacy_mongo_id`` — this is the post-cutover
          shape: ``/store_attachment`` mints a UUID, hands it to the
          worker, and the worker stashes it in ``legacy_mongo_id``
          while the row gets a freshly-generated PK. Trusting the
          input UUID as a PK here orphans the array entry: the column
          is ``uuid[]`` (no FK), so PG accepts the bad value and all
          downstream reads via ``AttachmentsRepository.get_any`` miss.

        Resolution therefore tries ``legacy_mongo_id`` first for every
        id (UUID-shaped or not), then falls back to the direct PK
        match. Unknown ids are dropped — they'd have failed the
        ``uuid[]`` cast otherwise and the whole row would have vanished
        via dual-write's exception swallow.
        """
        if not ids:
            return []
        # Defer to AttachmentsRepository for the batched lookup so the
        # legacy-first semantics live in one place.
        from docsgpt.storage.db.repositories.attachments import (
            AttachmentsRepository,
        )

        clean: list[str] = [str(raw) for raw in ids if raw is not None]
        if not clean:
            return []
        repo = AttachmentsRepository(self._conn)
        mapping = repo.resolve_ids(clean)
        out: list[str] = []
        for value in clean:
            mapped = mapping.get(value)
            if mapped is not None:
                out.append(mapped)
        return out

    # ------------------------------------------------------------------
    # Conversation CRUD
    # ------------------------------------------------------------------

    def create(
        self,
        user_id: str,
        name: str | None = None,
        *,
        agent_id: str | None = None,
        api_key: str | None = None,
        is_shared_usage: bool = False,
        shared_token: str | None = None,
        visibility: str = "listed",
        legacy_mongo_id: str | None = None,
    ) -> dict:
        """Create a new conversation.

        ``legacy_mongo_id`` is used by the dual-write shim so that a
        Postgres row inserted *after* a successful Mongo insert carries
        the Mongo ``_id`` as a lookup key. Subsequent appends/updates
        can then resolve the PG row by that id via
        :meth:`get_by_legacy_id`.
        """
        values: dict = {
            "user_id": user_id,
            "name": name,
            "visibility": visibility,
        }
        # ``agent_id`` may arrive as a Mongo ObjectId during the dual-write
        # window; resolve to a UUID (or drop silently if not yet backfilled).
        resolved_agent_id = self._resolve_agent_ref(agent_id)
        if resolved_agent_id:
            values["agent_id"] = resolved_agent_id
        if api_key:
            values["api_key"] = api_key
        if is_shared_usage:
            values["is_shared_usage"] = True
        if shared_token:
            values["shared_token"] = shared_token
        if legacy_mongo_id:
            values["legacy_mongo_id"] = legacy_mongo_id

        stmt = pg_insert(conversations_table).values(**values).returning(conversations_table)
        result = self._conn.execute(stmt)
        return row_to_dict(result.fetchone())

    def get_by_legacy_id(
        self, legacy_mongo_id: str, user_id: str | None = None,
    ) -> Optional[dict]:
        """Look up a conversation by the original Mongo ObjectId string.

        Used by the dual-write helpers to translate a Mongo ``_id`` into
        the Postgres UUID for follow-up writes. When ``user_id`` is
        provided, the lookup is scoped to rows owned by that user so
        callers can't accidentally resolve another user's conversation.
        """
        legacy_mongo_id = str(legacy_mongo_id) if legacy_mongo_id is not None else None
        if user_id is not None:
            result = self._conn.execute(
                text(
                    "SELECT * FROM conversations "
                    "WHERE legacy_mongo_id = :legacy_id AND user_id = :user_id"
                ),
                {"legacy_id": legacy_mongo_id, "user_id": user_id},
            )
        else:
            result = self._conn.execute(
                text(
                    "SELECT * FROM conversations WHERE legacy_mongo_id = :legacy_id"
                ),
                {"legacy_id": legacy_mongo_id},
            )
        row = result.fetchone()
        return row_to_dict(row) if row is not None else None

    def get(self, conversation_id: str, user_id: str) -> Optional[dict]:
        """Fetch a conversation the user owns or has shared access to."""
        result = self._conn.execute(
            text(
                "SELECT * FROM conversations "
                "WHERE id = CAST(:id AS uuid) "
                "AND (user_id = :user_id OR :user_id = ANY(shared_with))"
            ),
            {"id": conversation_id, "user_id": user_id},
        )
        row = result.fetchone()
        return row_to_dict(row) if row is not None else None

    def get_any(self, conversation_id: str, user_id: str) -> Optional[dict]:
        """Resolve a conversation by either PG UUID or legacy Mongo ObjectId string.

        Returns a conversation the user owns or has shared access to.
        """
        if looks_like_uuid(conversation_id):
            row = self.get(conversation_id, user_id)
            if row is not None:
                return row
        return self.get_by_legacy_id(conversation_id, user_id)

    def get_owned(self, conversation_id: str, user_id: str) -> Optional[dict]:
        """Fetch a conversation owned by the user (no shared access)."""
        result = self._conn.execute(
            text(
                "SELECT * FROM conversations "
                "WHERE id = CAST(:id AS uuid) AND user_id = :user_id"
            ),
            {"id": conversation_id, "user_id": user_id},
        )
        row = result.fetchone()
        return row_to_dict(row) if row is not None else None

    def list_for_user(self, user_id: str, limit: int = 30) -> list[dict]:
        """List a user's sidebar conversations, most recent first.

        Only ``visibility = 'listed'`` rows surface; agent/API/OpenAI-compat
        traffic persists as ``'hidden'`` and is excluded.
        """
        result = self._conn.execute(
            text(
                "SELECT * FROM conversations "
                "WHERE user_id = :user_id "
                "AND visibility = 'listed' "
                "ORDER BY date DESC LIMIT :limit"
            ),
            {"user_id": user_id, "limit": limit},
        )
        return [row_to_dict(r) for r in result.fetchall()]

    def search_for_user(
        self, user_id: str, query: str, limit: int = 30,
    ) -> list[dict]:
        """Search a user's conversations by name or message content.

        Same visibility filter as :meth:`list_for_user`. Matches against
        ``conversations.name`` or any of the conversation's messages'
        ``prompt`` / ``response`` columns (case-insensitive substring).

        Each returned row includes ``match_field`` (one of ``name``,
        ``prompt``, ``response``) and ``match_text`` (the full text of the
        first matching field, ``name`` taking precedence over messages,
        ``prompt`` over ``response``) so callers can render a snippet.
        """
        if not query:
            return []
        result = self._conn.execute(
            text(
                "SELECT c.*, mt.match_field, mt.match_text "
                "FROM conversations c "
                "JOIN LATERAL ( "
                "  SELECT field AS match_field, txt AS match_text "
                "  FROM ( "
                "    SELECT 'name'::text AS field, c.name AS txt, 0 AS prio "
                "    WHERE c.name ILIKE :pattern "
                "    UNION ALL "
                "    SELECT 'prompt'::text, m.prompt, 1 "
                "    FROM conversation_messages m "
                "    WHERE m.conversation_id = c.id "
                "      AND m.prompt ILIKE :pattern "
                "    UNION ALL "
                "    SELECT 'response'::text, m.response, 2 "
                "    FROM conversation_messages m "
                "    WHERE m.conversation_id = c.id "
                "      AND m.response ILIKE :pattern "
                "  ) s "
                "  ORDER BY prio "
                "  LIMIT 1 "
                ") mt ON TRUE "
                "WHERE c.user_id = :user_id "
                "AND c.visibility = 'listed' "
                "ORDER BY c.date DESC LIMIT :limit"
            ),
            {
                "user_id": user_id,
                "pattern": f"%{query}%",
                "limit": limit,
            },
        )
        return [row_to_dict(r) for r in result.fetchall()]

    def rename(self, conversation_id: str, user_id: str, name: str) -> bool:
        # Shape-gate so a non-UUID id (legacy Mongo ObjectId still floating
        # around in client-side state during the cutover) never reaches the
        # ``CAST(:id AS uuid)`` — that cast raises on the server and poisons
        # the enclosing transaction, making every subsequent query on the
        # same connection fail.
        if not looks_like_uuid(conversation_id):
            return False
        result = self._conn.execute(
            text(
                "UPDATE conversations SET name = :name, updated_at = now() "
                "WHERE id = CAST(:id AS uuid) AND user_id = :user_id"
            ),
            {"id": conversation_id, "user_id": user_id, "name": name},
        )
        return result.rowcount > 0

    def add_shared_user(self, conversation_id: str, user_to_add: str) -> bool:
        """Idempotently append ``user_to_add`` to ``shared_with``.

        Accepts either a PG UUID or a legacy Mongo ObjectId as the
        conversation id. Mirrors Mongo ``$addToSet`` semantics via the
        ``NOT (:user = ANY(shared_with))`` guard.
        """
        if not user_to_add:
            return False
        if looks_like_uuid(conversation_id):
            sql = (
                "UPDATE conversations "
                "SET shared_with = array_append(shared_with, :user), "
                "    updated_at = now() "
                "WHERE id = CAST(:id AS uuid) "
                "AND NOT (:user = ANY(shared_with))"
            )
        else:
            sql = (
                "UPDATE conversations "
                "SET shared_with = array_append(shared_with, :user), "
                "    updated_at = now() "
                "WHERE legacy_mongo_id = :id "
                "AND NOT (:user = ANY(shared_with))"
            )
        result = self._conn.execute(
            text(sql), {"id": conversation_id, "user": user_to_add},
        )
        return result.rowcount > 0

    def remove_shared_user(self, conversation_id: str, user_to_remove: str) -> bool:
        """Remove ``user_to_remove`` from ``shared_with``. Mirror of Mongo ``$pull``."""
        if not user_to_remove:
            return False
        if looks_like_uuid(conversation_id):
            sql = (
                "UPDATE conversations "
                "SET shared_with = array_remove(shared_with, :user), "
                "    updated_at = now() "
                "WHERE id = CAST(:id AS uuid) "
                "AND :user = ANY(shared_with)"
            )
        else:
            sql = (
                "UPDATE conversations "
                "SET shared_with = array_remove(shared_with, :user), "
                "    updated_at = now() "
                "WHERE legacy_mongo_id = :id "
                "AND :user = ANY(shared_with)"
            )
        result = self._conn.execute(
            text(sql), {"id": conversation_id, "user": user_to_remove},
        )
        return result.rowcount > 0

    def set_shared_token(self, conversation_id: str, user_id: str, token: str) -> bool:
        # Shape-gate: see ``rename`` — prevents transaction poisoning when
        # a non-UUID id reaches this code path.
        if not looks_like_uuid(conversation_id):
            return False
        result = self._conn.execute(
            text(
                "UPDATE conversations SET shared_token = :token, updated_at = now() "
                "WHERE id = CAST(:id AS uuid) AND user_id = :user_id"
            ),
            {"id": conversation_id, "user_id": user_id, "token": token},
        )
        return result.rowcount > 0

    def update_compression_metadata(
        self, conversation_id: str, user_id: str, metadata: dict,
    ) -> bool:
        """Replace the entire ``compression_metadata`` JSONB blob.

        Prefer :meth:`append_compression_point` + :meth:`set_compression_flags`
        to match the Mongo service semantics exactly (those two mirror
        ``$set`` + ``$push $slice``). This method is retained for callers
        that already compute the full merged blob client-side.
        """
        # Shape-gate: see ``rename`` — prevents transaction poisoning when
        # a non-UUID id reaches this code path.
        if not looks_like_uuid(conversation_id):
            return False
        result = self._conn.execute(
            text(
                "UPDATE conversations "
                "SET compression_metadata = CAST(:meta AS jsonb), updated_at = now() "
                "WHERE id = CAST(:id AS uuid) AND user_id = :user_id"
            ),
            {"id": conversation_id, "user_id": user_id, "meta": json.dumps(metadata)},
        )
        return result.rowcount > 0

    def set_compression_flags(
        self,
        conversation_id: str,
        *,
        is_compressed: bool,
        last_compression_at,
    ) -> bool:
        """Update ``compression_metadata.is_compressed`` and
        ``compression_metadata.last_compression_at`` without touching
        ``compression_points``.

        Mirrors the Mongo ``$set`` on those two subfields in
        ``ConversationService.update_compression_metadata``. Initialises
        the surrounding object when the row has no ``compression_metadata``
        yet.
        """
        # Shape-gate: the streaming pipeline may pass through a legacy id
        # that ``get_by_legacy_id`` couldn't resolve; in that case the id
        # remains a non-UUID string and the CAST would poison the txn.
        if not looks_like_uuid(conversation_id):
            return False
        result = self._conn.execute(
            text(
                """
                UPDATE conversations SET
                    compression_metadata = jsonb_set(
                        jsonb_set(
                            COALESCE(compression_metadata, '{}'::jsonb),
                            '{is_compressed}',
                            to_jsonb(CAST(:is_compressed AS boolean)), true
                        ),
                        '{last_compression_at}',
                        to_jsonb(CAST(:last_compression_at AS text)), true
                    ),
                    updated_at = now()
                WHERE id = CAST(:id AS uuid)
                """
            ),
            {
                "id": conversation_id,
                "is_compressed": bool(is_compressed),
                "last_compression_at": (
                    str(last_compression_at) if last_compression_at is not None else None
                ),
            },
        )
        return result.rowcount > 0

    def append_compression_point(
        self,
        conversation_id: str,
        point: dict,
        *,
        max_points: int,
    ) -> bool:
        """Append one compression point, keeping at most ``max_points``.

        Mirrors Mongo's ``$push {"$each": [point], "$slice": -max_points}``
        on ``compression_metadata.compression_points``. Preserves the
        other top-level keys in ``compression_metadata``.
        """
        # Shape-gate: see ``set_compression_flags``.
        if not looks_like_uuid(conversation_id):
            return False
        # Idempotent on the same point: two persist paths (the mid-execution
        # handler and the route's end-of-stream save) used to both append it.
        # The duplicate check lives inside the UPDATE so it runs under the
        # row lock — a separate SELECT would let two sessions both pass it.
        # A point without ``query_index``/``compressed_summary`` never matches
        # (NULL comparisons are false), so it is always appended.
        dup_index = point.get("query_index")
        dup_summary = point.get("compressed_summary")
        result = self._conn.execute(
            text(
                """
                UPDATE conversations SET
                    compression_metadata = CASE
                        WHEN (compression_metadata -> 'compression_points' -> -1
                              ->> 'query_index') = CAST(:dup_index AS text)
                         AND (compression_metadata -> 'compression_points' -> -1
                              ->> 'compressed_summary') = CAST(:dup_summary AS text)
                        THEN compression_metadata
                        ELSE jsonb_set(
                            COALESCE(compression_metadata, '{}'::jsonb),
                            '{compression_points}',
                            COALESCE(
                                (
                                    SELECT jsonb_agg(elem ORDER BY rn)
                                    FROM (
                                        SELECT
                                            elem,
                                            row_number() OVER () AS rn,
                                            count(*) OVER () AS cnt
                                        FROM jsonb_array_elements(
                                            COALESCE(
                                                compression_metadata -> 'compression_points',
                                                '[]'::jsonb
                                            ) || jsonb_build_array(CAST(:point AS jsonb))
                                        ) AS elem
                                    ) ranked
                                    WHERE rn > cnt - :max_points
                                ),
                                '[]'::jsonb
                            ),
                            true
                        )
                    END,
                    updated_at = now()
                WHERE id = CAST(:id AS uuid)
                """
            ),
            {
                "id": conversation_id,
                "point": json.dumps(point, cls=PGNativeJSONEncoder),
                "max_points": int(max_points),
                "dup_index": None if dup_index is None else str(dup_index),
                "dup_summary": None if dup_summary is None else str(dup_summary),
            },
        )
        return result.rowcount > 0

    def delete(self, conversation_id: str, user_id: str) -> bool:
        # Shape-gate: see ``rename`` — prevents transaction poisoning when
        # a non-UUID id reaches this code path.
        if not looks_like_uuid(conversation_id):
            return False
        # Artifacts are parent-derived: deleting a conversation must take its
        # artifacts with it (rows + bytes + the quota they consume). conversation_id
        # is a bare uuid (no FK cascade), so delete the rows explicitly and reap the
        # stored bytes best-effort.
        from docsgpt.storage.db.repositories.artifacts import ArtifactsRepository

        artifacts = ArtifactsRepository(self._conn)
        paths = artifacts.storage_paths_for_conversation(conversation_id)
        # Traces linked to a message go with it (ON DELETE CASCADE). This
        # takes the rest -- scheduled runs, stateless /v1 rounds, turns whose
        # message was never reserved -- whose tool results and chunk previews
        # the user expects gone with the conversation.
        self._conn.execute(
            text(
                "DELETE FROM request_traces WHERE conversation_id = CAST(:id AS uuid) "
                "AND EXISTS (SELECT 1 FROM conversations "
                "WHERE id = CAST(:id AS uuid) AND user_id = :user_id)"
            ),
            {"id": conversation_id, "user_id": user_id},
        )
        result = self._conn.execute(
            text(
                "DELETE FROM conversations "
                "WHERE id = CAST(:id AS uuid) AND user_id = :user_id"
            ),
            {"id": conversation_id, "user_id": user_id},
        )
        deleted = result.rowcount > 0
        if deleted:
            artifacts.delete_for_conversation(conversation_id)
            self._reap_artifact_storage(paths)
        return deleted

    def delete_all_for_user(self, user_id: str) -> int:
        from docsgpt.storage.db.repositories.artifacts import ArtifactsRepository

        # Reap artifacts BEFORE the conversations: delete_for_user_conversations
        # resolves them via a subquery over the still-present conversation rows.
        artifacts = ArtifactsRepository(self._conn)
        paths = artifacts.storage_paths_for_user_conversations(user_id)
        artifacts.delete_for_user_conversations(user_id)
        # See ``delete``: traces not linked to a message are removed explicitly.
        self._conn.execute(
            text(
                "DELETE FROM request_traces WHERE conversation_id IN "
                "(SELECT id FROM conversations WHERE user_id = :user_id)"
            ),
            {"user_id": user_id},
        )
        result = self._conn.execute(
            text("DELETE FROM conversations WHERE user_id = :user_id"),
            {"user_id": user_id},
        )
        self._reap_artifact_storage(paths)
        return result.rowcount

    def reassign_api_key(self, *, old_key: str, new_key: str) -> int:
        """Re-point conversation rows from ``old_key`` to ``new_key``.

        Called when an agent's key is rotated. Per-agent analytics match
        ``c.api_key = :api_key OR c.agent_id = :agent_pg_id``, but a
        conversation created from api-key traffic may have ``api_key`` set and
        ``agent_id`` NULL (``conversation_service`` only sets ``agent_id`` when
        the caller passes one). Rewriting ``api_key`` keeps those rows attached
        to the agent after rotation instead of orphaning them.

        Matched by ``api_key`` only — ``agents.key`` is globally unique so a key
        maps to exactly one agent. Returns the number of rows updated.
        """
        if not old_key or not new_key:
            return 0
        result = self._conn.execute(
            text("UPDATE conversations SET api_key = :new_key WHERE api_key = :old_key"),
            {"old_key": old_key, "new_key": new_key},
        )
        return result.rowcount

    @staticmethod
    def _reap_artifact_storage(paths: list) -> None:
        """Best-effort delete artifact bytes whose rows were cascade-removed; never raises."""
        if not paths:
            return
        try:
            from docsgpt.storage.storage_creator import StorageCreator

            storage = StorageCreator.get_storage()
        except Exception:
            logger.warning("conversation delete: storage unavailable for artifact cleanup", exc_info=True)
            return
        for path in paths:
            try:
                storage.delete_file(path)
            except Exception:
                logger.warning("conversation delete: failed to delete artifact bytes %s", path, exc_info=True)

    # ------------------------------------------------------------------
    # Messages
    # ------------------------------------------------------------------

    def get_messages(self, conversation_id: str) -> list[dict]:
        result = self._conn.execute(
            text(
                "SELECT * FROM conversation_messages "
                "WHERE conversation_id = CAST(:conv_id AS uuid) "
                "ORDER BY position ASC"
            ),
            {"conv_id": conversation_id},
        )
        return [_message_row_to_dict(r) for r in result.fetchall()]

    def first_n_message_ids(self, conversation_id: str, first_n: int) -> set[str]:
        """Return the ids of the first ``first_n`` messages (by position) of the conversation.

        Same ordering a share snapshot uses (``get_messages``' ``position ASC``
        then ``[:first_n]``); empty set for ``first_n <= 0`` or a non-UUID id.
        """
        if first_n <= 0 or not looks_like_uuid(conversation_id):
            return set()
        result = self._conn.execute(
            text(
                "SELECT id FROM conversation_messages "
                "WHERE conversation_id = CAST(:cid AS uuid) "
                "ORDER BY position ASC LIMIT :n"
            ),
            {"cid": conversation_id, "n": int(first_n)},
        )
        return {str(r[0]) for r in result.fetchall()}

    def message_in_first_n(self, conversation_id: str, message_id: str, first_n: int) -> bool:
        """True if message_id is among the first ``first_n`` messages (by position) of the conversation."""
        if not message_id or first_n <= 0:
            return False
        # Shape-gate both ids: a non-UUID reaching CAST(... AS uuid) would raise
        # and poison the enclosing transaction (see ``rename``).
        if not looks_like_uuid(conversation_id) or not looks_like_uuid(message_id):
            return False
        row = self._conn.execute(
            text(
                "SELECT 1 FROM ( "
                "  SELECT id FROM conversation_messages "
                "  WHERE conversation_id = CAST(:cid AS uuid) "
                "  ORDER BY position ASC LIMIT :n "
                ") t WHERE t.id = CAST(:mid AS uuid)"
            ),
            {"cid": conversation_id, "mid": message_id, "n": int(first_n)},
        ).fetchone()
        return row is not None

    def get_message_at(self, conversation_id: str, position: int) -> Optional[dict]:
        # Shape-gate: see ``rename``. Callers today always pass a resolved
        # UUID (via ``get_any`` first), but the guard costs nothing and
        # keeps future callers safe from txn-poisoning.
        if not looks_like_uuid(conversation_id):
            return None
        result = self._conn.execute(
            text(
                "SELECT * FROM conversation_messages "
                "WHERE conversation_id = CAST(:conv_id AS uuid) "
                "AND position = :pos"
            ),
            {"conv_id": conversation_id, "pos": position},
        )
        row = result.fetchone()
        return _message_row_to_dict(row) if row is not None else None

    def append_message(self, conversation_id: str, message: dict) -> dict:
        """Append a message to a conversation.

        Uses ``SELECT ... FOR UPDATE`` to allocate the next position
        atomically. The caller must be inside a transaction.

        Mirrors Mongo's ``$push`` on the ``queries`` array.
        """
        # Lock the parent conversation row to serialize concurrent appends.
        self._conn.execute(
            text(
                "SELECT id FROM conversations "
                "WHERE id = CAST(:conv_id AS uuid) FOR UPDATE"
            ),
            {"conv_id": conversation_id},
        )
        next_pos_result = self._conn.execute(
            text(
                "SELECT COALESCE(MAX(position), -1) + 1 AS next_pos "
                "FROM conversation_messages "
                "WHERE conversation_id = CAST(:conv_id AS uuid)"
            ),
            {"conv_id": conversation_id},
        )
        next_pos = next_pos_result.scalar()

        values = {
            "conversation_id": conversation_id,
            "position": next_pos,
            "prompt": message.get("prompt"),
            "response": message.get("response"),
            "thought": message.get("thought"),
            "sources": message.get("sources") or [],
            "tool_calls": message.get("tool_calls") or [],
            "model_id": message.get("model_id"),
            "message_metadata": message.get("metadata") or {},
        }
        # Callers that know the turn failed (e.g. an agent that yielded a
        # terminal error) must be able to say so; without this the column
        # default silently made every appended row "complete".
        if message.get("status") is not None:
            values["status"] = message["status"]
        if message.get("timestamp") is not None:
            values["timestamp"] = message["timestamp"]

        attachments = message.get("attachments")
        if attachments:
            # Attachment ids may arrive as Mongo ObjectIds during the
            # dual-write window — resolve each to a PG UUID or drop it.
            resolved = self._resolve_attachment_refs(
                [str(a) for a in attachments],
            )
            if resolved:
                values["attachments"] = resolved

        stmt = (
            pg_insert(conversation_messages_table)
            .values(**values)
            .returning(conversation_messages_table)
        )
        result = self._conn.execute(stmt)
        # Touch the parent conversation's updated_at.
        self._conn.execute(
            text(
                "UPDATE conversations SET updated_at = now() "
                "WHERE id = CAST(:id AS uuid)"
            ),
            {"id": conversation_id},
        )
        return _message_row_to_dict(result.fetchone())

    def update_message_at(
        self, conversation_id: str, position: int, fields: dict,
    ) -> bool:
        """Update specific fields on a message at a given position.

        Mirrors Mongo's ``$set`` on ``queries.{index}.*``.
        """
        allowed = {
            "prompt", "response", "thought", "sources", "tool_calls",
            "attachments", "model_id", "metadata", "timestamp", "status",
            # Feedback can be re-set in rare continuation flows; without
            # it in the whitelist an upstream re-append that happens to
            # carry feedback would silently lose it. Mirrors
            # ``set_feedback`` — column is JSONB.
            "feedback", "feedback_timestamp",
        }
        filtered = {k: v for k, v in fields.items() if k in allowed}
        if not filtered:
            return False

        # Map public API key ``metadata`` → DB column ``message_metadata``.
        api_to_col = {"metadata": "message_metadata"}

        set_parts = []
        params: dict = {"conv_id": conversation_id, "pos": position}
        for key, val in filtered.items():
            col = api_to_col.get(key, key)
            if key in ("sources", "tool_calls", "metadata", "feedback"):
                set_parts.append(f"{col} = CAST(:{col} AS jsonb)")
                if val is None:
                    params[col] = None
                else:
                    params[col] = (
                        json.dumps(val) if not isinstance(val, str) else val
                    )
            elif key == "attachments":
                # Attachment ids may be Mongo ObjectIds during the
                # dual-write window; translate via attachments.legacy_mongo_id.
                set_parts.append(f"{col} = CAST(:{col} AS uuid[])")
                params[col] = self._resolve_attachment_refs(
                    [str(a) for a in val] if val else [],
                )
            else:
                set_parts.append(f"{col} = :{col}")
                params[col] = val

        if "timestamp" not in filtered:
            set_parts.append("timestamp = now()")
        sql = (
            f"UPDATE conversation_messages SET {', '.join(set_parts)} "
            "WHERE conversation_id = CAST(:conv_id AS uuid) AND position = :pos"
        )
        result = self._conn.execute(text(sql), params)
        return result.rowcount > 0

    def reserve_message(
        self,
        conversation_id: str,
        *,
        prompt: str,
        placeholder_response: str,
        request_id: str | None = None,
        status: str = "pending",
        attachments: list[str] | None = None,
        model_id: str | None = None,
        metadata: dict | None = None,
    ) -> dict:
        """Pre-persist a placeholder assistant message before the LLM call."""
        self._conn.execute(
            text(
                "SELECT id FROM conversations "
                "WHERE id = CAST(:conv_id AS uuid) FOR UPDATE"
            ),
            {"conv_id": conversation_id},
        )
        next_pos = self._conn.execute(
            text(
                "SELECT COALESCE(MAX(position), -1) + 1 AS next_pos "
                "FROM conversation_messages "
                "WHERE conversation_id = CAST(:conv_id AS uuid)"
            ),
            {"conv_id": conversation_id},
        ).scalar()

        values = {
            "conversation_id": conversation_id,
            "position": next_pos,
            "prompt": prompt,
            "response": placeholder_response,
            "status": status,
            "request_id": request_id,
            "model_id": model_id,
            "message_metadata": metadata or {},
        }
        if attachments:
            resolved = self._resolve_attachment_refs(
                [str(a) for a in attachments],
            )
            if resolved:
                values["attachments"] = resolved

        stmt = (
            pg_insert(conversation_messages_table)
            .values(**values)
            .returning(conversation_messages_table)
        )
        result = self._conn.execute(stmt)
        self._conn.execute(
            text(
                "UPDATE conversations SET updated_at = now() "
                "WHERE id = CAST(:id AS uuid)"
            ),
            {"id": conversation_id},
        )
        return _message_row_to_dict(result.fetchone())

    def update_message_by_id(
        self, message_id: str, fields: dict,
        *, only_if_non_terminal: bool = False,
        reclaim_reconciler_failed: bool = False,
    ) -> MessageUpdateOutcome:
        """Update specific fields on a message identified by its UUID.

        ``metadata`` is merged into the existing JSONB rather than
        overwritten, so a reconciler-set ``reconcile_attempts`` survives
        a successful late finalize. When ``only_if_non_terminal`` is
        True, the update is gated so a late finalize cannot retract a
        reconciler-set ``failed`` (or a prior ``complete``).

        ``reclaim_reconciler_failed`` punches two narrow holes in that gate,
        both for rows whose ``failed`` stamp a later successful answer
        disproves:

        1. A row the **reconciler** failed for staleness. Such a stream
           demonstrably was not dead — production showed a 20-minute tool loop
           swept at minute 6 and then completing normally at minute 20, its
           answer refused and replaced by "Response was terminated prior to
           completion". Matched on the reconciler's own error string, and
           skipping rows whose awaiting-approval prompt it already revoked.
        2. A row marked ``resume_retryable`` by the stream handler, which sets
           it only when it also released the continuation claim — i.e. it
           explicitly handed the turn back for a retry that reuses this same
           row id.

        Everything else stays terminal: genuine failures (client disconnect,
        provider errors, the pending-tool cleanup task) are unaffected.

        Reclaiming also strips the reconciler's bookkeeping keys, which
        otherwise survive the metadata merge and leave a ``complete`` row
        carrying a stale ``error``.

        The return value discriminates "I updated the row" from "the
        row was already at a terminal state" so the abort handler can
        journal ``end`` when the normal-path finalize already ran. A reclaim
        returns plain ``UPDATED``: callers that branch on the outcome should
        treat it as an ordinary successful finalize.

        Args:
            message_id: UUID of the message row.
            fields: Field/value pairs to write.
            only_if_non_terminal: Gate the write on a non-terminal status.
            reclaim_reconciler_failed: Also allow rows the reconciler failed
                for staleness, and rows a failed resume marked retryable.

        Returns:
            MessageUpdateOutcome: What happened to the row.
        """
        if not looks_like_uuid(message_id):
            return MessageUpdateOutcome.INVALID
        allowed = {
            "prompt", "response", "thought", "sources", "tool_calls",
            "attachments", "model_id", "metadata", "timestamp", "status",
            "request_id", "feedback", "feedback_timestamp",
        }
        filtered = {k: v for k, v in fields.items() if k in allowed}
        if not filtered:
            return MessageUpdateOutcome.INVALID

        api_to_col = {"metadata": "message_metadata"}

        # On reclaim the reconciler's bookkeeping must not survive into the
        # finished row: ``error`` would show a stale failure in the API
        # response, and leaving the counters behind corrupts any "how often
        # does this fire" query.
        reclaim_strip = (
            " - 'error' - 'reconcile_attempts' - 'last_reconcile_seen_heartbeat'"
            " - 'resume_retryable'"
            if reclaim_reconciler_failed
            else ""
        )

        set_parts = []
        params: dict = {"id": message_id}
        for key, val in filtered.items():
            col = api_to_col.get(key, key)
            if key == "metadata":
                if val is None:
                    set_parts.append(f"{col} = NULL")
                else:
                    set_parts.append(
                        f"{col} = (COALESCE({col}, '{{}}'::jsonb){reclaim_strip}) "
                        f"|| CAST(:{col} AS jsonb)"
                    )
                    params[col] = (
                        json.dumps(val) if not isinstance(val, str) else val
                    )
            elif key in ("sources", "tool_calls", "feedback"):
                set_parts.append(f"{col} = CAST(:{col} AS jsonb)")
                if val is None:
                    params[col] = None
                else:
                    params[col] = (
                        json.dumps(val) if not isinstance(val, str) else val
                    )
            elif key == "attachments":
                set_parts.append(f"{col} = CAST(:{col} AS uuid[])")
                params[col] = self._resolve_attachment_refs(
                    [str(a) for a in val] if val else [],
                )
            else:
                set_parts.append(f"{col} = :{col}")
                params[col] = val

        set_parts.append("updated_at = now()")
        # Metadata may not be among the written fields (a finalize with no
        # query metadata passes none), but the strip must still happen.
        if reclaim_reconciler_failed and "metadata" not in filtered:
            set_parts.append(
                f"message_metadata = COALESCE(message_metadata, '{{}}'::jsonb)"
                f"{reclaim_strip}"
            )
        update_where = ["id = CAST(:id AS uuid)"]
        if only_if_non_terminal:
            terminal_gate = "status NOT IN ('complete', 'failed')"
            if reclaim_reconciler_failed:
                terminal_gate = (
                    "(" + terminal_gate + " OR ("
                    "status = 'failed' "
                    "AND message_metadata->>'error' LIKE "
                    "'reconciler: stuck in pending/streaming%' "
                    "AND COALESCE("
                    "(message_metadata->>'reconciler_cleared_approval')::boolean,"
                    " false) = false"
                    ") OR ("
                    # Second hole, same principle: the stream handler that
                    # failed a resume also handed the turn back by releasing
                    # the continuation claim, so it stamped the row
                    # ``resume_retryable``. The retry reuses this very row
                    # (the reserved id lives in the persisted agent_config);
                    # without this it would land ALREADY_FAILED and the
                    # user's second, successful answer would be discarded.
                    "status = 'failed' "
                    "AND COALESCE("
                    "(message_metadata->>'resume_retryable')::boolean,"
                    " false) = true"
                    "))"
                )
            update_where.append(terminal_gate)
        # Single-statement attempt + prior-status probe. Both CTEs see
        # the same MVCC snapshot, so ``prior.status`` reflects the row
        # state before the UPDATE — exactly what we need to tell
        # ``ALREADY_COMPLETE`` apart from ``ALREADY_FAILED`` apart from
        # ``NOT_FOUND`` without a follow-up SELECT.
        sql = (
            "WITH attempted AS ("
            f"  UPDATE conversation_messages SET {', '.join(set_parts)} "
            f"  WHERE {' AND '.join(update_where)} "
            "  RETURNING 1 AS updated"
            "), "
            "prior AS ("
            "  SELECT status FROM conversation_messages "
            "  WHERE id = CAST(:id AS uuid)"
            ") "
            "SELECT (SELECT updated FROM attempted) AS updated, "
            "       (SELECT status FROM prior) AS prior_status"
        )
        row = self._conn.execute(text(sql), params).fetchone()
        if row is None:
            return MessageUpdateOutcome.NOT_FOUND
        updated, prior_status = row[0], row[1]
        if updated:
            if reclaim_reconciler_failed and prior_status == "failed":
                # Distinct, greppable signal: each of these is one user-visible
                # answer that would previously have been discarded. Should
                # trend to zero once false sweeps stop.
                logger.info(
                    "finalize_message: reclaimed reconciler-failed row "
                    "message_id=%s",
                    message_id,
                )
            return MessageUpdateOutcome.UPDATED
        if prior_status is None:
            return MessageUpdateOutcome.NOT_FOUND
        if prior_status == "complete":
            return MessageUpdateOutcome.ALREADY_COMPLETE
        if prior_status == "failed":
            return MessageUpdateOutcome.ALREADY_FAILED
        # ``only_if_non_terminal=False`` always updates an existing row,
        # so reaching here means the gate excluded it for some status
        # the terminal set doesn't cover — treat as "not found" rather
        # than inventing a new variant.
        return MessageUpdateOutcome.NOT_FOUND

    def update_message_status(
        self, message_id: str, status: str,
    ) -> bool:
        """Cheap status-only transition (e.g. pending → streaming).

        Only flips non-terminal rows: a reconciler-set ``failed`` row
        stays put so the late streaming chunk doesn't silently retract
        the alert.
        """
        if not looks_like_uuid(message_id):
            return False
        result = self._conn.execute(
            text(
                "UPDATE conversation_messages SET status = :status, "
                "updated_at = now() "
                "WHERE id = CAST(:id AS uuid) "
                "AND status NOT IN ('complete', 'failed')"
            ),
            {"id": message_id, "status": status},
        )
        return result.rowcount > 0

    def heartbeat_message_state(self, message_id: str) -> HeartbeatState:
        """Stamp the liveness heartbeat and report what the row looks like.

        One round trip, three outcomes. The distinction matters because the
        two "not stamped" cases want opposite handling:

        - ``MISSING`` — the row was deleted mid-stream (``truncate_after`` on
          a retry/edit). Nothing this stream produces can ever be read, so the
          caller should stop working.
        - ``TERMINAL`` — the row exists but is ``complete``/``failed``,
          typically because the reconciler swept it. The stream must NOT be
          cancelled: if it finishes, ``finalize_message`` is allowed to reclaim
          a reconciler-failed row and land the real answer.
        - ``STAMPED`` — normal in-flight row, freshness extended.

        The stamp deliberately avoids ``timestamp`` (creation time, used for
        history sort) and ``status`` (the WAL marker), and skips terminal rows
        so a late heartbeat can't retract a reconciler-set ``failed``.

        Args:
            message_id: UUID of the message row.

        Returns:
            HeartbeatState: What the row was at stamp time.
        """
        if not looks_like_uuid(message_id):
            return HeartbeatState.MISSING
        # Both CTEs read the same MVCC snapshot, so ``present`` describes the
        # row as it was when the UPDATE ran — no follow-up SELECT needed.
        row = self._conn.execute(
            text(
                """
                WITH stamped AS (
                    UPDATE conversation_messages
                    SET message_metadata = jsonb_set(
                        COALESCE(message_metadata, '{}'::jsonb),
                        '{last_heartbeat_at}',
                        to_jsonb(clock_timestamp())
                    )
                    WHERE id = CAST(:id AS uuid)
                      AND status NOT IN ('complete', 'failed')
                    RETURNING 1 AS ok
                )
                SELECT (SELECT ok FROM stamped) AS stamped,
                       EXISTS(
                           SELECT 1 FROM conversation_messages
                           WHERE id = CAST(:id AS uuid)
                       ) AS present
                """
            ),
            {"id": message_id},
        ).fetchone()
        if row is not None and row[0]:
            return HeartbeatState.STAMPED
        if row is not None and row[1]:
            return HeartbeatState.TERMINAL
        return HeartbeatState.MISSING

    def heartbeat_message(self, message_id: str) -> bool:
        """Stamp the heartbeat; True when an in-flight row was updated."""
        return self.heartbeat_message_state(message_id) is HeartbeatState.STAMPED

    def confirm_executed_tool_calls(self, message_id: str) -> int:
        """Flip ``tool_call_attempts.status='executed' → 'confirmed'`` for the message."""
        if not looks_like_uuid(message_id):
            return 0
        result = self._conn.execute(
            text(
                "UPDATE tool_call_attempts SET status = 'confirmed', "
                "updated_at = now() "
                "WHERE message_id = CAST(:mid AS uuid) AND status = 'executed'"
            ),
            {"mid": message_id},
        )
        return result.rowcount or 0

    def truncate_after(self, conversation_id: str, keep_up_to: int) -> int:
        """Delete messages with position > keep_up_to.

        Mirrors Mongo's ``$push`` + ``$slice`` that trims queries after an
        index-based update.

        Rows being deleted may still be **owned by a running stream**: this
        runs on the retry/edit path, and a stream survives client disconnect
        (the SSE keepalive pump drains the generator in a daemon thread), while
        a stream paused on a client-side tool can resume many minutes later.
        Deleting such a row cascades its ``message_events`` away and makes
        every subsequent journal and ``tool_call_attempts`` write FK-fail — in
        production one paused stream, deleted 9 minutes before it resumed,
        produced 342 WARNs plus 60 ERRORs and finalized into nothing.

        Deleting is still correct — the user explicitly asked to replace these
        turns — so the rows go, but any non-terminal one is first stamped
        terminal, which leaves an audit trail of what was superseded.

        That stamp does NOT survive: the UPDATE's predicate is a subset of the
        DELETE's (it also filters ``status NOT IN ('complete','failed')``), and
        both run in one transaction, so every stamped row is removed
        microseconds later and the owning stream's late ``finalize_message``
        correctly reports ``NOT_FOUND``, not ``ALREADY_FAILED``. Do not rely on
        the stamp for anything but the ``RETURNING id`` log line.

        What DOES survive is the ``superseded_messages`` tombstone written
        below. A stream superseded less than ``STREAM_HEARTBEAT_INTERVAL``
        before it finishes is never cancelled — the ticker in
        ``api/answer/routes/base.py`` is the only thing that sets the cancel
        flag, and it cannot be set from here because it is a ``threading.Event``
        in another worker process — so the stream runs to completion and
        finalizes into a row that is gone. Without the tombstone that logged an
        ERROR ``answer_persist_failed``, identical to a genuinely orphaned
        answer, every time a user hit retry on an answer shorter than the
        heartbeat interval.

        Args:
            conversation_id: Conversation whose tail is being trimmed.
            keep_up_to: Highest position to keep.

        Returns:
            Number of rows deleted.
        """
        # Shape-gate: see ``rename`` — prevents transaction poisoning when
        # a non-UUID id reaches this code path.
        if not looks_like_uuid(conversation_id):
            return 0
        superseded = self._conn.execute(
            text(
                """
                UPDATE conversation_messages
                SET status = 'failed',
                    message_metadata = jsonb_set(
                        COALESCE(message_metadata, '{}'::jsonb),
                        '{error}',
                        to_jsonb(
                            CAST('superseded by a newer turn at this position'
                                 AS text)
                        )
                    )
                WHERE conversation_id = CAST(:conv_id AS uuid)
                  AND position > :pos
                  AND status NOT IN ('complete', 'failed')
                RETURNING id
                """
            ),
            {"conv_id": conversation_id, "pos": keep_up_to},
        ).fetchall()
        if superseded:
            logger.info(
                "truncate_after: superseding %d in-flight message(s) in "
                "conversation %s: %s",
                len(superseded),
                conversation_id,
                ", ".join(str(r[0]) for r in superseded),
            )
        # Tombstone before the delete, in the same transaction: the late
        # finalize usually runs in a different worker process, so the signal
        # has to be durable. No FK, or the CASCADE below would remove it.
        self._conn.execute(
            text(
                """
                INSERT INTO superseded_messages (message_id, conversation_id)
                SELECT id, conversation_id
                  FROM conversation_messages
                 WHERE conversation_id = CAST(:conv_id AS uuid)
                   AND position > :pos
                ON CONFLICT (message_id) DO NOTHING
                """
            ),
            {"conv_id": conversation_id, "pos": keep_up_to},
        )
        result = self._conn.execute(
            text(
                "DELETE FROM conversation_messages "
                "WHERE conversation_id = CAST(:conv_id AS uuid) "
                "AND position > :pos"
            ),
            {"conv_id": conversation_id, "pos": keep_up_to},
        )
        return result.rowcount

    def was_superseded(self, message_id: str) -> bool:
        """True when ``message_id`` was deleted by a retry/edit truncation.

        Lets a stream that outlived its own row tell an expected supersede from
        a genuinely orphaned answer. See :meth:`truncate_after`.

        Args:
            message_id: The reserved message id the stream was writing to.

        Returns:
            True if a tombstone exists for that message.
        """
        if not looks_like_uuid(message_id):
            return False
        row = self._conn.execute(
            text(
                "SELECT 1 FROM superseded_messages "
                "WHERE message_id = CAST(:mid AS uuid)"
            ),
            {"mid": message_id},
        ).fetchone()
        return row is not None

    def cleanup_superseded_older_than(self, days: int) -> int:
        """Delete ``superseded_messages`` tombstones older than ``days``.

        The tombstone only has to outlive the stream that was superseded, which
        is bounded by the request timeout — days of retention is already
        generous, and the table has no other reader.

        Args:
            days: Retention window in days.

        Returns:
            Number of rows deleted.
        """
        result = self._conn.execute(
            text(
                "DELETE FROM superseded_messages "
                "WHERE superseded_at < clock_timestamp() "
                "- make_interval(days => :days)"
            ),
            {"days": days},
        )
        return result.rowcount or 0

    def set_feedback(
        self, conversation_id: str, position: int, feedback: dict | None,
    ) -> bool:
        """Set or unset feedback on a message.

        ``feedback`` is a JSONB value, e.g. ``{"text": "thumbs_up",
        "timestamp": "..."}`` or ``None`` to unset.
        """
        # Shape-gate: see ``rename`` — prevents transaction poisoning when
        # a non-UUID id reaches this code path.
        if not looks_like_uuid(conversation_id):
            return False
        fb_json = json.dumps(feedback) if feedback is not None else None
        result = self._conn.execute(
            text(
                "UPDATE conversation_messages "
                "SET feedback = CAST(:fb AS jsonb) "
                "WHERE conversation_id = CAST(:conv_id AS uuid) AND position = :pos"
            ),
            {"conv_id": conversation_id, "pos": position, "fb": fb_json},
        )
        return result.rowcount > 0

    def message_count(self, conversation_id: str) -> int:
        result = self._conn.execute(
            text(
                "SELECT COUNT(*) FROM conversation_messages "
                "WHERE conversation_id = CAST(:conv_id AS uuid)"
            ),
            {"conv_id": conversation_id},
        )
        return result.scalar() or 0