Paginate chat history so long sessions open quickly.

Load the latest 80 messages first and fetch older pages on scroll-up via before_id, instead of pulling thousands of rows up front.

Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
oliver 2026-08-11 10:13:18 +08:00
parent 2a52349668
commit c3f3992a51
7 changed files with 256 additions and 17 deletions

View file

@ -81,6 +81,8 @@ class AssistantStoreProtocol(Protocol):
def get_llm_profile_secret(self, profile_id: 'str') -> 'Optional[str]': ...
def get_messages(self, session_id: 'str', limit: 'int' = 200) -> 'list[ChatMessage]': ...
def get_messages_after_id(self, *, session_id: 'str', after_id: 'int', limit: 'int' = 200) -> 'list[ChatMessage]': ...
def get_messages_before_id(self, *, session_id: 'str', before_id: 'int', limit: 'int' = 200) -> 'list[ChatMessage]': ...
def exists_message_before_id(self, *, session_id: 'str', before_id: 'int') -> 'bool': ...
def get_or_create_channel_session(self, *, tenant_id: 'str', channel: 'str', external_chat_id: 'str', external_user_id: 'str', session_title: 'str') -> 'str': ...
def get_or_create_channel_session_v2(self, *, tenant_id: 'str', channel: 'str', account_id: 'str', external_chat_id: 'str', external_user_id: 'str', session_title: 'str') -> 'str': ...
def get_secret(self, key: 'str') -> 'Optional[str]': ...

View file

@ -237,6 +237,102 @@ class ChatMessagesSaRepository:
rows.insert(0, dict(arow))
return [_row_to_chat_message(r) for r in rows]
def get_messages_before_id(
self, *, session_id: str, before_id: int, limit: int
) -> list[ChatMessage]:
"""Return up to ``limit`` messages with id < ``before_id``, ASC within the window.
Same orphan-tool prepend behavior as :meth:`get_messages_recent_asc` so a page
that starts mid-tool-turn still includes the parent assistant row.
"""
if limit <= 0:
return []
sid = str(session_id or "").strip()
if not sid:
return []
bid = int(before_id or 0)
if bid <= 0:
return []
lim = max(1, min(int(limit), 2000))
ids_sq = (
select(chat_message.c.id)
.where(chat_message.c.session_id == sid, chat_message.c.id < bid)
.order_by(chat_message.c.id.desc())
.limit(lim)
.scalar_subquery()
)
stmt = (
select(
chat_message.c.id,
chat_message.c.session_id,
chat_message.c.role,
chat_message.c.content,
chat_message.c.tool_calls,
chat_message.c.attachments,
chat_message.c.turn_uuid,
chat_message.c.event_type,
chat_message.c.event_payload,
chat_message.c.timestamp,
)
.where(chat_message.c.session_id == sid, chat_message.c.id.in_(ids_sq))
.order_by(chat_message.c.id.asc())
)
prepended: set[int] = set()
with self._engine.connect() as conn:
rows: list[dict[str, Any]] = [dict(r) for r in conn.execute(stmt).mappings().all()]
while rows:
first = rows[0]
if str(first.get("role") or "") != "tool":
break
aid = _tool_row_assistant_message_id(first.get("tool_calls"))
if aid is None:
break
first_id = int(first["id"])
if aid >= first_id:
break
if any(int(r["id"]) == int(aid) for r in rows):
break
if int(aid) in prepended:
break
arow = conn.execute(
select(
chat_message.c.id,
chat_message.c.session_id,
chat_message.c.role,
chat_message.c.content,
chat_message.c.tool_calls,
chat_message.c.attachments,
chat_message.c.turn_uuid,
chat_message.c.event_type,
chat_message.c.event_payload,
chat_message.c.timestamp,
)
.where(chat_message.c.session_id == sid, chat_message.c.id == int(aid))
.limit(1)
).mappings().first()
if not arow:
break
prepended.add(int(aid))
rows.insert(0, dict(arow))
return [_row_to_chat_message(r) for r in rows]
def exists_message_before_id(self, *, session_id: str, before_id: int) -> bool:
sid = str(session_id or "").strip()
bid = int(before_id or 0)
if not sid or bid <= 0:
return False
with self._engine.connect() as conn:
return bool(
conn.execute(
select(
exists().where(
chat_message.c.session_id == sid,
chat_message.c.id < bid,
)
)
).scalar()
)
def get_messages_after_id(
self, *, session_id: str, after_id: int, limit: int
) -> list[ChatMessage]:

View file

@ -2282,6 +2282,26 @@ class SqliteStore(ScheduledJobStoreMixin):
session_id=sid, after_id=aid, limit=lim
)
def get_messages_before_id(self, *, session_id: str, before_id: int, limit: int = 200) -> list[ChatMessage]:
"""Return up to ``limit`` messages with id < ``before_id``, ASC within the window."""
sid = str(session_id or "").strip()
if not sid:
return []
bid = int(before_id or 0)
if bid <= 0:
return []
lim = max(1, min(int(limit), 2000))
return self._chat_messages_repo().get_messages_before_id(
session_id=sid, before_id=bid, limit=lim
)
def exists_message_before_id(self, *, session_id: str, before_id: int) -> bool:
sid = str(session_id or "").strip()
bid = int(before_id or 0)
if not sid or bid <= 0:
return False
return self._chat_messages_repo().exists_message_before_id(session_id=sid, before_id=bid)
def add_tool_log(
self,
session_id: str,