process_inserted_message() now does one branch fetch

riturajFi · riturajFi · commit b57c833a68e0 · 2026-04-02T01:03:51.000+05:30
diff --git a/src/server/core/acontext_core/service/controller/message.py b/src/server/core/acontext_core/service/controller/message.py
@@ -1,6 +1,7 @@
 from ..data import message as MD
 from ..data import learning_space as LS
 from ...infra.db import DB_CLIENT
+from ...schema.orm import Message
 from ...schema.session.task import TaskStatus
 from ...schema.session.message import MessageBlob
 from ...schema.utils import asUUID
@@ -19,11 +20,11 @@
 }
 
 
-def _claimable_branch_message_ids(message_rows: list[dict]) -> list[asUUID]:
+def _claimable_branch_message_ids(messages: list[Message]) -> list[asUUID]:
     return [
-        row["id"]
-        for row in message_rows
-        if row["session_task_process_status"] in _CLAIMABLE_BRANCH_STATUSES
+        message.id
+        for message in messages
+        if message.session_task_process_status in _CLAIMABLE_BRANCH_STATUSES
     ]
 
 
@@ -50,15 +51,18 @@ async def process_inserted_message(
     wide = get_wide_event()
     disabled = await get_metrics(project_id, ExcessMetricTags.new_task_created)
     target_message_ids: list[asUUID] = []
+    branch_messages: list[Message] = []
 
     try:
         async with DB_CLIENT.get_session_context() as session:
-            r = await MD.fetch_message_branch_path_rows(session, message_id, session_id)
-            message_rows, eil = r.unpack()
+            r = await MD.fetch_message_branch_path_messages(
+                session, message_id, session_id
+            )
+            branch_messages, eil = r.unpack()
             if eil:
                 return r
 
-            target_message_ids = _claimable_branch_message_ids(message_rows)
+            target_message_ids = _claimable_branch_message_ids(branch_messages)
             wide["claimed_branch_message_count"] = len(target_message_ids)
 
             if not target_message_ids:
@@ -78,25 +82,22 @@ async def process_inserted_message(
                 session, target_message_ids, TaskStatus.RUNNING
             )
 
-        async with DB_CLIENT.get_session_context() as session:
-            r = await MD.fetch_message_branch_path_data(
-                session, message_id, session_id, user_kek=user_kek
-            )
-            messages, eil = r.unpack()
-            if eil:
-                await _try_rollback_to_failed(target_message_ids)
-                return r
+        r = await MD.hydrate_message_parts(branch_messages, user_kek=user_kek)
+        messages, eil = r.unpack()
+        if eil:
+            await _try_rollback_to_failed(target_message_ids)
+            return r
 
-            messages_data = [
-                MessageBlob(
-                    message_id=m.id,
-                    parent_id=m.parent_id,
-                    role=m.role,
-                    parts=m.parts,
-                    task_id=m.task_id,
-                )
-                for m in messages
-            ]
+        messages_data = [
+            MessageBlob(
+                message_id=m.id,
+                parent_id=m.parent_id,
+                role=m.role,
+                parts=m.parts,
+                task_id=m.task_id,
+            )
+            for m in messages
+        ]
 
         async with DB_CLIENT.get_session_context() as session:
             r = await LS.get_learning_space_for_session(session, session_id)
@@ -139,6 +140,8 @@ async def process_inserted_message(
         wide["task_agent_outcome"] = "exception"
         await _try_rollback_to_failed(target_message_ids)
         raise
+
+
 async def process_session_pending_message(
     project_config: ProjectConfig,
     project_id: asUUID,
diff --git a/src/server/core/acontext_core/service/data/message.py b/src/server/core/acontext_core/service/data/message.py
@@ -109,28 +109,41 @@ async def fetch_messages_data_by_ids(
                 f"Some messages({message_ids}) not found in database: {e}"
             )
 
-        if not ordered_messages:
+        return await hydrate_message_parts(ordered_messages, user_kek=user_kek)
+
+    except Exception as e:
+        return Result.reject(f"Error fetching messages by IDs {message_ids}: {e}")
+
+
+async def hydrate_message_parts(
+    messages: List[Message],
+    user_kek: bytes | None = None,
+) -> Result[List[Message]]:
+    """
+    Load parts from S3 for already-fetched message rows.
+
+    Mutates each message in place by setting `message.parts`.
+    """
+    try:
+        if not messages:
             return Result.resolve([])
 
-        # Fetch parts concurrently for all messages
         parts_tasks = [
             _fetch_message_parts(message.parts_asset_meta, user_kek=user_kek)
-            for message in ordered_messages
+            for message in messages
         ]
         parts_results = await asyncio.gather(*parts_tasks)
 
-        # Assign parts to messages
-        for message, parts_result in zip(ordered_messages, parts_results):
+        for message, parts_result in zip(messages, parts_results):
             d, eil = parts_result.unpack()
             if eil:
                 message.parts = None
                 continue
             message.parts = d
 
-        return Result.resolve(ordered_messages)
-
+        return Result.resolve(messages)
     except Exception as e:
-        return Result.reject(f"Error fetching messages by IDs {message_ids}: {e}")
+        return Result.reject(f"Error hydrating message parts: {e}")
 
 
 async def fetch_message_branch_path_ids(
@@ -226,6 +239,69 @@ async def fetch_message_branch_path_rows(
         )
 
 
+async def fetch_message_branch_path_messages(
+    db_session: AsyncSession,
+    message_id: asUUID,
+    session_id: asUUID | None = None,
+) -> Result[List[Message]]:
+    """
+    Fetch one message's branch path as ordered Message rows.
+
+    Uses a recursive CTE to walk parent_id upward in one query.
+    """
+    try:
+        query = text(
+            """
+            WITH RECURSIVE message_path AS (
+                SELECT id, parent_id, session_id, 0 AS depth
+                FROM messages
+                WHERE id = :message_id
+
+                UNION ALL
+
+                SELECT
+                    parent.id,
+                    parent.parent_id,
+                    parent.session_id,
+                    child.depth + 1 AS depth
+                FROM messages AS parent
+                JOIN message_path AS child
+                  ON parent.id = child.parent_id
+            )
+            SELECT m.*
+            FROM message_path AS mp
+            JOIN messages AS m ON m.id = mp.id
+            ORDER BY mp.depth DESC, m.id ASC
+            """
+        )
+        result = await db_session.execute(
+            select(Message).from_statement(query),
+            {"message_id": message_id},
+        )
+        messages = list(result.scalars().all())
+
+        if not messages:
+            return Result.reject(f"Message {message_id} doesn't exist")
+
+        path_session_ids = {message.session_id for message in messages}
+
+        if session_id is not None and path_session_ids != {session_id}:
+            return Result.reject(
+                f"Message {message_id} does not belong to session {session_id}"
+            )
+
+        if len(path_session_ids) != 1:
+            return Result.reject(
+                f"Message {message_id} has an invalid cross-session parent chain"
+            )
+
+        return Result.resolve(messages)
+    except Exception as e:
+        return Result.reject(
+            f"Error fetching branch path messages for message {message_id}: {e}"
+        )
+
+
 async def branch_pending_message_length(
     db_session: AsyncSession,
     message_id: asUUID,
@@ -267,13 +343,11 @@ async def fetch_message_branch_path_data(
     """
     Fetch one message's branch path with parts loaded from S3.
     """
-    r = await fetch_message_branch_path_ids(db_session, message_id, session_id)
-    message_ids, eil = r.unpack()
+    r = await fetch_message_branch_path_messages(db_session, message_id, session_id)
+    messages, eil = r.unpack()
     if eil:
         return Result.reject(str(eil))
-    return await fetch_messages_data_by_ids(
-        db_session, message_ids, user_kek=user_kek
-    )
+    return await hydrate_message_parts(messages, user_kek=user_kek)
 
 
 async def fetch_session_messages(
diff --git a/src/server/core/tests/service/test_message_data_branch_path.py b/src/server/core/tests/service/test_message_data_branch_path.py
@@ -1,5 +1,7 @@
+import json
 import pytest
 from uuid import uuid4
+from unittest.mock import AsyncMock, patch
 
 from acontext_core.schema.orm import Message, Project, Session
 from acontext_core.service.data import message as MD
@@ -167,6 +169,136 @@ async def test_fetch_message_branch_path_rows_returns_statuses(db_client):
         ]
 
 
+@pytest.mark.asyncio
+async def test_fetch_message_branch_path_messages_returns_ordered_messages(db_client):
+    _, acontext_session = await _create_project_and_session(db_client)
+
+    async with db_client.get_session_context() as session:
+        root = Message(
+            session_id=acontext_session.id,
+            role="user",
+            parts_asset_meta={
+                "bucket": "test-bucket",
+                "s3_key": "parts/root.json",
+                "etag": "etag-root",
+                "sha256": "sha-root",
+                "mime": "application/json",
+                "size_b": 10,
+            },
+            session_task_process_status="success",
+        )
+        session.add(root)
+        await session.flush()
+
+        child = Message(
+            session_id=acontext_session.id,
+            role="assistant",
+            parts_asset_meta={
+                "bucket": "test-bucket",
+                "s3_key": "parts/child.json",
+                "etag": "etag-child",
+                "sha256": "sha-child",
+                "mime": "application/json",
+                "size_b": 11,
+            },
+            parent_id=root.id,
+            session_task_process_status="pending",
+        )
+        session.add(child)
+        await session.flush()
+
+        leaf = Message(
+            session_id=acontext_session.id,
+            role="user",
+            parts_asset_meta={
+                "bucket": "test-bucket",
+                "s3_key": "parts/leaf.json",
+                "etag": "etag-leaf",
+                "sha256": "sha-leaf",
+                "mime": "application/json",
+                "size_b": 12,
+            },
+            parent_id=child.id,
+            session_task_process_status="running",
+        )
+        session.add(leaf)
+        await session.flush()
+
+        r = await MD.fetch_message_branch_path_messages(
+            session, leaf.id, acontext_session.id
+        )
+        messages, eil = r.unpack()
+
+        assert eil is None
+        assert [message.id for message in messages] == [root.id, child.id, leaf.id]
+        assert [message.session_task_process_status for message in messages] == [
+            "success",
+            "pending",
+            "running",
+        ]
+
+
+@pytest.mark.asyncio
+async def test_hydrate_message_parts_uses_existing_branch_message_order(db_client):
+    _, acontext_session = await _create_project_and_session(db_client)
+
+    async with db_client.get_session_context() as session:
+        root = Message(
+            session_id=acontext_session.id,
+            role="user",
+            parts_asset_meta={
+                "bucket": "test-bucket",
+                "s3_key": "parts/root.json",
+                "etag": "etag-root",
+                "sha256": "sha-root",
+                "mime": "application/json",
+                "size_b": 10,
+            },
+        )
+        session.add(root)
+        await session.flush()
+
+        leaf = Message(
+            session_id=acontext_session.id,
+            role="assistant",
+            parts_asset_meta={
+                "bucket": "test-bucket",
+                "s3_key": "parts/leaf.json",
+                "etag": "etag-leaf",
+                "sha256": "sha-leaf",
+                "mime": "application/json",
+                "size_b": 11,
+            },
+            parent_id=root.id,
+        )
+        session.add(leaf)
+        await session.flush()
+
+        branch_result = await MD.fetch_message_branch_path_messages(
+            session, leaf.id, acontext_session.id
+        )
+        messages, eil = branch_result.unpack()
+
+        assert eil is None
+
+        with patch(
+            "acontext_core.service.data.message.S3_CLIENT.download_object",
+            new_callable=AsyncMock,
+            side_effect=[
+                json.dumps([{"type": "text", "text": "root"}]).encode("utf-8"),
+                json.dumps([{"type": "text", "text": "leaf"}]).encode("utf-8"),
+            ],
+        ) as mock_download:
+            hydrated_result = await MD.hydrate_message_parts(messages)
+            hydrated_messages, hydrate_error = hydrated_result.unpack()
+
+        assert hydrate_error is None
+        assert [message.id for message in hydrated_messages] == [root.id, leaf.id]
+        assert hydrated_messages[0].parts[0].text == "root"
+        assert hydrated_messages[1].parts[0].text == "leaf"
+        assert mock_download.await_count == 2
+
+
 @pytest.mark.asyncio
 async def test_branch_pending_message_length_ignores_sibling_branches(db_client):
     _, acontext_session = await _create_project_and_session(db_client)
diff --git a/src/server/core/tests/service/test_process_inserted_message.py b/src/server/core/tests/service/test_process_inserted_message.py
