fix(backend): a re-captured or bot-edited message keeps its earlier text as a version, an older copy never overwrites a newer edit

backend/src/userbot/modules/capture/message.py

@@ -1,11 +1,14 @@
+from datetime import datetime
+
 from pyrogram import Client
 from pyrogram.types import Message
 
 from userbot.modules.capture import repository
 from userbot.modules.capture.context import CaptureContext
+from userbot.modules.capture.repository import MessageRow
 from userbot.modules.links import extract_links
 from userbot.modules.links import repository as links_repository
-from userbot.modules.media import capture_media, self_destruct_ttl
+from userbot.modules.media import capture_media, media_unique_id, self_destruct_ttl
 from userbot.modules.raw_json import raw_json
 from utils.policy.models import CaptureToggles
 
@@ -34,27 +37,35 @@ def callbacks(message: Message) -> list[tuple[int, str | None, bytes | None]]:
     return buttons
 
 
+def message_row(message: Message, chat_id: int, date: datetime) -> MessageRow:
+    return MessageRow(
+        chat_id=chat_id,
+        message_id=message.id,
+        date=date,
+        sender_id=sender_id(message),
+        text=message.text or message.caption,
+        raw=raw_json(message),
+        has_media=message.media is not None,
+        is_self_destruct=self_destruct_ttl(message) is not None,
+        edit_date=message.edit_date,
+        media_unique_id=media_unique_id(message),
+    )
+
+
 async def capture_message(
     client: Client, message: Message, ctx: CaptureContext, toggles: CaptureToggles
 ) -> None:
     if message.empty or message.chat is None or message.date is None:
         return
     chat_id = message.chat.id or 0
-    await repository.upsert_message(
-        ctx.pool,
-        ctx.account_id,
-        chat_id,
-        message.id,
-        message.date,
-        sender_id(message),
-        message.text or message.caption,
-        raw_json(message),
-        has_media=message.media is not None,
-        is_self_destruct=self_destruct_ttl(message) is not None,
+    current = await repository.upsert_message(
+        ctx.pool, ctx.account_id, message_row(message, chat_id, message.date)
     )
     from userbot.modules.capture.identity import capture_identity  # noqa: PLC0415
 
     await capture_identity(message, ctx)
+    if not current:
+        return
     await capture_media(client, message, ctx, chat_id, message.id, toggles)
     buttons = callbacks(message)
     if buttons:

backend/src/userbot/modules/capture/repository.py

@@ -1,3 +1,4 @@
+from dataclasses import dataclass
 from datetime import datetime
 
 import asyncpg
@@ -21,13 +22,19 @@ _TOUCH_EDITED = """
 INSERT INTO messages
     (account_id, chat_id, message_id, date, sender_id, text, raw,
      has_media, is_self_destruct, edited_at)
-VALUES ($1, $2, $3, $4, $5, $6, $7::jsonb, $8, $9, now())
+VALUES ($1, $2, $3, $4, $5, $6, $7::jsonb, $8, $9, CASE WHEN $10 THEN now() END)
 ON CONFLICT (account_id, chat_id, message_id, date) DO UPDATE SET
     text = EXCLUDED.text,
     raw = EXCLUDED.raw,
     has_media = EXCLUDED.has_media,
     is_self_destruct = EXCLUDED.is_self_destruct,
-    edited_at = now()
+    edited_at = coalesce(EXCLUDED.edited_at, messages.edited_at)
+"""
+
+_LOCK_MESSAGE = """
+SELECT pg_advisory_xact_lock(
+    hashtextextended(concat_ws(':', 'message', $1::int, $2::bigint, $3::bigint), 0)
+)
 """
 
 _FIND_DELETED = """
@@ -61,9 +68,32 @@ WHERE m.account_id = $1 AND m.chat_id = $2 AND m.message_id = $3
 ON CONFLICT DO NOTHING
 """
 
+_INSERT_PAST_VERSION = """
+INSERT INTO message_versions
+    (account_id, chat_id, message_id, observed_at, edit_date, text, raw)
+SELECT $1::int, $2::bigint, $3::bigint,
+    coalesce(
+        (SELECT min(v.observed_at) - interval '1 microsecond'
+         FROM message_versions v
+         WHERE v.account_id = $1 AND v.chat_id = $2 AND v.message_id = $3
+           AND v.edit_date > $4::timestamptz),
+        $4::timestamptz,
+        $5::timestamptz
+    ),
+    $4::timestamptz, $6::text, $7::jsonb
+WHERE NOT EXISTS (
+    SELECT 1 FROM message_versions v
+    WHERE v.account_id = $1 AND v.chat_id = $2 AND v.message_id = $3
+      AND v.edit_date IS NOT DISTINCT FROM $4::timestamptz
+      AND v.text IS NOT DISTINCT FROM $6::text
+)
+ON CONFLICT DO NOTHING
+"""
+
 _CURRENT_CONTENT = """
 SELECT
     m.text AS text,
+    m.raw->>'edit_date' AS edit_date,
     (SELECT d.unique_id FROM media d
      WHERE d.account_id = m.account_id AND d.chat_id = m.chat_id
        AND d.message_id = m.message_id) AS media_unique_id
@@ -121,33 +151,108 @@ WHERE evicted_key = $2 OR (unique_id = $1 AND evicted_at IS NOT NULL)
 """
 
 
-async def upsert_message(  # noqa: PLR0913
-    pool: asyncpg.Pool,
-    account_id: int,
-    chat_id: int,
-    message_id: int,
-    date: datetime,
-    sender_id: int | None,
-    text: str | None,
-    raw: str,
-    *,
-    has_media: bool,
-    is_self_destruct: bool,
-) -> None:
-    await pool.execute(
-        _UPSERT_MESSAGE,
-        account_id,
-        chat_id,
-        message_id,
-        date,
-        sender_id,
-        text,
-        raw,
-        has_media,
-        is_self_destruct,
+@dataclass(frozen=True, slots=True)
+class MessageRow:
+    chat_id: int
+    message_id: int
+    date: datetime
+    sender_id: int | None
+    text: str | None
+    raw: str
+    has_media: bool
+    is_self_destruct: bool
+    edit_date: datetime | None
+    media_unique_id: str | None
+
+    @property
+    def columns(self) -> tuple:
+        return (
+            self.chat_id,
+            self.message_id,
+            self.date,
+            self.sender_id,
+            self.text,
+            self.raw,
+            self.has_media,
+            self.is_self_destruct,
+        )
+
+
+async def _lock_current(
+    conn: asyncpg.Connection, account_id: int, row: MessageRow
+) -> asyncpg.Record | None:
+    await conn.execute(_LOCK_MESSAGE, account_id, row.chat_id, row.message_id)
+    return await conn.fetchrow(
+        _CURRENT_CONTENT, account_id, row.chat_id, row.message_id
     )
 
 
+def _is_stale(current: asyncpg.Record, row: MessageRow) -> bool:
+    stored = current["edit_date"]
+    if stored is None:
+        return False
+    if row.edit_date is None:
+        return True
+    try:
+        return row.edit_date < datetime.fromisoformat(stored)
+    except (TypeError, ValueError):
+        return False
+
+
+async def _record_past(
+    conn: asyncpg.Connection, account_id: int, row: MessageRow
+) -> None:
+    await conn.execute(
+        _INSERT_PAST_VERSION,
+        account_id,
+        row.chat_id,
+        row.message_id,
+        row.edit_date,
+        row.date,
+        row.text,
+        row.raw,
+    )
+
+
+async def _record_edit(
+    conn: asyncpg.Connection, account_id: int, row: MessageRow, *, mark_edited: bool
+) -> None:
+    await conn.execute(_SNAPSHOT_ORIGINAL, account_id, row.chat_id, row.message_id)
+    await conn.execute(_TOUCH_EDITED, account_id, *row.columns, mark_edited)
+    await conn.execute(
+        _INSERT_VERSION,
+        account_id,
+        row.chat_id,
+        row.message_id,
+        row.edit_date,
+        row.text,
+        row.raw,
+    )
+
+
+async def upsert_message(pool: asyncpg.Pool, account_id: int, row: MessageRow) -> bool:
+    async with pool.acquire() as conn, conn.transaction():
+        current = await _lock_current(conn, account_id, row)
+        if current is None:
+            await conn.execute(_UPSERT_MESSAGE, account_id, *row.columns)
+            return True
+        stored_media = current["media_unique_id"]
+        changed = current["text"] != row.text or (
+            stored_media is not None and stored_media != row.media_unique_id
+        )
+        if _is_stale(current, row):
+            if changed:
+                await _record_past(conn, account_id, row)
+            return False
+        if changed:
+            await _record_edit(
+                conn, account_id, row, mark_edited=row.edit_date is not None
+            )
+        else:
+            await conn.execute(_UPSERT_MESSAGE, account_id, *row.columns)
+    return True
+
+
 async def max_message_id(
     pool: asyncpg.Pool, account_id: int, chat_id: int
 ) -> int | None:
@@ -188,45 +293,20 @@ async def mark_deleted_channel(
     await _mark_deleted(pool, account_id, rows)
 
 
-async def add_version(  # noqa: PLR0913
-    pool: asyncpg.Pool,
-    account_id: int,
-    chat_id: int,
-    message_id: int,
-    date: datetime,
-    sender_id: int | None,
-    text: str | None,
-    raw: str,
-    edit_date: datetime | None,
-    media_unique_id: str | None,
-    *,
-    has_media: bool,
-    is_self_destruct: bool,
-) -> bool:
+async def add_version(pool: asyncpg.Pool, account_id: int, row: MessageRow) -> bool:
     async with pool.acquire() as conn, conn.transaction():
-        current = await conn.fetchrow(_CURRENT_CONTENT, account_id, chat_id, message_id)
-        text_changed = current is None or current["text"] != text
-        media_changed = (
-            current is not None and current["media_unique_id"] != media_unique_id
-        )
-        if not (text_changed or media_changed):
-            return False
-        await conn.execute(_SNAPSHOT_ORIGINAL, account_id, chat_id, message_id)
-        await conn.execute(
-            _TOUCH_EDITED,
-            account_id,
-            chat_id,
-            message_id,
-            date,
-            sender_id,
-            text,
-            raw,
-            has_media,
-            is_self_destruct,
-        )
-        await conn.execute(
-            _INSERT_VERSION, account_id, chat_id, message_id, edit_date, text, raw
-        )
+        current = await _lock_current(conn, account_id, row)
+        if current is not None:
+            changed = (
+                current["text"] != row.text
+                or current["media_unique_id"] != row.media_unique_id
+            )
+            if not changed:
+                return False
+            if _is_stale(current, row):
+                await _record_past(conn, account_id, row)
+                return False
+        await _record_edit(conn, account_id, row, mark_edited=True)
     return True