perf(api): skip saved prompts when loading chat history (#44521)

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

api/core/memory/token_buffer_memory.py

@@ -3,7 +3,7 @@ from collections.abc import Sequence
 from dataclasses import dataclass
 
 from sqlalchemy import select
-from sqlalchemy.orm import Session, sessionmaker
+from sqlalchemy.orm import Session, load_only, sessionmaker
 
 from core.app.app_config.features.file_upload.manager import FileUploadConfigManager
 from core.app.file_access import DatabaseFileAccessController
@@ -215,7 +215,24 @@ class TokenBufferMemory:
         require the app. An omitted workflow repository uses the configured factory.
         """
         # fetch limited messages, and return reversed
-        stmt = select(Message).where(Message.conversation_id == conversation.id).order_by(Message.created_at.desc())
+        # Load only the columns used to rebuild history. Message.message stores the full prompt of each turn,
+        # earlier history included, so loading it makes memory grow quadratically with conversation length.
+        stmt = (
+            select(Message)
+            .options(
+                load_only(
+                    Message.id,
+                    Message.parent_message_id,
+                    Message.query,
+                    Message.answer,
+                    Message.answer_tokens,
+                    Message.workflow_run_id,
+                    Message.created_at,
+                )
+            )
+            .where(Message.conversation_id == conversation.id)
+            .order_by(Message.created_at.desc())
+        )
 
         message_limit = min(message_limit, 500) if message_limit and message_limit > 0 else 500
 

api/core/prompt/utils/get_thread_messages_length.py

@@ -1,5 +1,5 @@
 from sqlalchemy import select
-from sqlalchemy.orm import Session
+from sqlalchemy.orm import Session, load_only
 
 from core.prompt.utils.extract_thread_messages import extract_thread_messages
 from models.model import Message
@@ -9,8 +9,14 @@ def get_thread_messages_length(conversation_id: str, *, session: Session) -> int
     """
     Get the number of thread messages based on the parent message id.
     """
-    # Fetch all messages related to the conversation
-    stmt = select(Message).where(Message.conversation_id == conversation_id).order_by(Message.created_at.desc())
+    # Fetch all messages related to the conversation, loading only the columns needed to walk the thread.
+    # Message.message stores the full prompt of each turn, so loading it grows quadratically with the conversation.
+    stmt = (
+        select(Message)
+        .options(load_only(Message.id, Message.parent_message_id, Message.answer))
+        .where(Message.conversation_id == conversation_id)
+        .order_by(Message.created_at.desc())
+    )
 
     messages = session.scalars(stmt).all()
 

api/tests/unit_tests/core/memory/test_token_buffer_memory.py

@@ -1,6 +1,7 @@
 """Comprehensive SQLite-backed tests for token-buffer memory."""
 
 import json
+import re
 from collections.abc import Iterator, Mapping, Sequence
 from dataclasses import dataclass
 from datetime import UTC, datetime, timedelta
@@ -146,6 +147,7 @@ def _persist_message(
     answer_tokens: int = 5,
     created_at: datetime | None = None,
     workflow_run_id: str | None = None,
+    parent_message_id: str | None = None,
 ) -> Message:
     message = Message(
         id=str(uuid4()),
@@ -154,6 +156,7 @@ def _persist_message(
         _inputs={},
         query=query,
         message={},
+        parent_message_id=parent_message_id,
         message_unit_price=Decimal(0),
         answer=answer,
         answer_tokens=answer_tokens,
@@ -375,6 +378,38 @@ class TestGetHistoryPromptMessages:
         assert sum("FROM apps" in sql for sql in selects) == 1
         assert len(result) == 10
 
+    def test_history_query_skips_saved_prompt_column(self, database: Database) -> None:
+        """Message.message stores each turn's full prompt, so history must neither select nor lazy-load it."""
+        mem = self._make_memory(database)
+        base_time = datetime.now(UTC).replace(tzinfo=None)
+        parent_message_id: str | None = None
+        for index in range(3):
+            parent_message_id = _persist_message(
+                database,
+                mem.conversation.id,
+                query=f"query-{index}",
+                answer=f"answer-{index}",
+                created_at=base_time + timedelta(seconds=index),
+                parent_message_id=parent_message_id,
+            ).id
+        # Reload rows from the database instead of reusing fully loaded instances from the identity map.
+        database.session.expunge_all()
+        before = len(database.statements)
+
+        result = mem.get_history_prompt_messages()
+
+        message_selects = [sql for sql, _ in database.statements[before:] if "FROM messages" in sql]
+        assert len(message_selects) == 1
+        assert re.search(r"\bmessages\.message\b", message_selects[0]) is None
+        assert [prompt.content for prompt in result] == [
+            "query-0",
+            "answer-0",
+            "query-1",
+            "answer-1",
+            "query-2",
+            "answer-2",
+        ]
+
     @pytest.mark.parametrize(
         ("token_values", "max_token_limit", "expected_length"),
         [

api/tests/unit_tests/core/prompt/test_extract_thread_messages.py

@@ -1,8 +1,10 @@
+import re
 from datetime import datetime, timedelta
 from decimal import Decimal
 from uuid import uuid4
 
 import pytest
+from sqlalchemy import Engine, event
 from sqlalchemy.orm import Session
 
 from constants import UUID_NIL
@@ -198,3 +200,45 @@ def test_get_thread_messages_length_keeps_non_empty_latest_answer(sqlite_session
     length = get_thread_messages_length("conversation-2", session=sqlite_session)
 
     assert length == 2
+
+
+def test_get_thread_messages_length_skips_saved_prompt_column(sqlite_engine: Engine, sqlite_session: Session):
+    """Message.message stores each turn's full prompt, so counting the thread must neither select nor lazy-load it."""
+    id1, id2 = str(uuid4()), str(uuid4())
+    now = datetime.now()
+    sqlite_session.add_all(
+        [
+            _persisted_message(
+                message_id=id2,
+                conversation_id="conversation-3",
+                parent_message_id=id1,
+                answer="latest-answer",
+                created_at=now,
+            ),
+            _persisted_message(
+                message_id=id1,
+                conversation_id="conversation-3",
+                parent_message_id=UUID_NIL,
+                answer="older-answer",
+                created_at=now - timedelta(seconds=1),
+            ),
+        ]
+    )
+    sqlite_session.commit()
+    # Reload rows from the database instead of reusing fully loaded instances from the identity map.
+    sqlite_session.expunge_all()
+    statements: list[str] = []
+
+    def record_statement(_connection, _cursor, statement, _parameters, _context, _executemany) -> None:
+        statements.append(statement)
+
+    event.listen(sqlite_engine, "before_cursor_execute", record_statement)
+    try:
+        length = get_thread_messages_length("conversation-3", session=sqlite_session)
+    finally:
+        event.remove(sqlite_engine, "before_cursor_execute", record_statement)
+
+    assert length == 2
+    message_selects = [sql for sql in statements if "FROM messages" in sql]
+    assert len(message_selects) == 1
+    assert re.search(r"\bmessages\.message\b", message_selects[0]) is None