diff --git a/roboco/api/routes/messages.py b/roboco/api/routes/messages.py index 073b8ff7..9414ba5e 100644 --- a/roboco/api/routes/messages.py +++ b/roboco/api/routes/messages.py @@ -25,6 +25,7 @@ from roboco.services.messaging import ( MessageCreateRequest as ServiceMessageRequest, ) from roboco.services.messaging import ( + MessageCursor, get_messaging_service, ) @@ -45,10 +46,20 @@ async def list_messages( """List messages in a session (session existence is checked in the service).""" messaging = get_messaging_service(db) try: + before_cursor = ( + MessageCursor(params.before, params.before_id) + if params.before is not None + else None + ) + after_cursor = ( + MessageCursor(params.after, params.after_id) + if params.after is not None + else None + ) messages, has_more = await messaging.list_messages_for_session( session_id=params.session_id, - before=params.before, - after=params.after, + before=before_cursor, + after=after_cursor, message_type=params.type_filter, limit=params.limit, ) diff --git a/roboco/api/schemas/messages.py b/roboco/api/schemas/messages.py index 18b8aeeb..f45f9f09 100644 --- a/roboco/api/schemas/messages.py +++ b/roboco/api/schemas/messages.py @@ -23,6 +23,12 @@ class ListMessagesParams(BaseModel): session_id: UUID before: datetime | None = None after: datetime | None = None + # Keyset-pagination tie-breakers: pass the last (before) / first (after) + # message's id alongside its timestamp so equal-timestamp messages are not + # skipped across pages. Optional — a timestamp-only cursor keeps the + # legacy strict-inequality behavior. + before_id: UUID | None = None + after_id: UUID | None = None type_filter: MessageType | None = None limit: int = Field(50, ge=1, le=100) diff --git a/roboco/services/messaging.py b/roboco/services/messaging.py index 07158193..647ec813 100644 --- a/roboco/services/messaging.py +++ b/roboco/services/messaging.py @@ -16,7 +16,7 @@ from datetime import UTC, datetime, timedelta from typing import Any, ClassVar, cast from uuid import UUID -from sqlalchemy import select +from sqlalchemy import and_, or_, select from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import joinedload, selectinload @@ -76,6 +76,23 @@ def _minutes_to_timedelta(value: int | None) -> timedelta | None: # ============================================================================= +@dataclass(frozen=True, slots=True) +class MessageCursor: + """A keyset-pagination cursor over messages: a ``(timestamp, id)`` pair. + + The ``id`` tie-breaks equal timestamps so paging across same-timestamp + messages skips nothing: the next page resumes past the cursor's id at + the shared timestamp rather than being excluded by a strict + ``timestamp < cursor`` / ``timestamp > cursor`` inequality (which drops + every row sharing the cursor's timestamp — they are cut by ``limit`` on + the prior page and never returned). ``id`` is ``None`` for a legacy + timestamp-only cursor (strict inequality, the prior behavior). + """ + + timestamp: datetime + id: UUID | None = None + + class MessagingService(BaseService): """ Service for managing all messaging operations. @@ -1571,8 +1588,8 @@ class MessagingService(BaseService): self, *, session_id: UUID, - before: datetime | None, - after: datetime | None, + before: MessageCursor | None, + after: MessageCursor | None, message_type: MessageType | None, limit: int, ) -> tuple[list[MessageTable], bool]: @@ -1580,7 +1597,10 @@ class MessagingService(BaseService): Existing `get_messages` skips the session-existence check — this variant raises NotFoundError when the session is missing so routes - can return a clean 404 without issuing their own query. + can return a clean 404 without issuing their own query. The ``before`` + / ``after`` cursors carry the keyset ``(timestamp, id)`` pair so + callers can paginate without skipping equal-timestamp messages across + pages. """ await self.get_session_or_raise(session_id) return await self.get_messages( @@ -1668,8 +1688,8 @@ class MessagingService(BaseService): async def get_messages( self, session_id: UUID, - before: datetime | None = None, - after: datetime | None = None, + before: MessageCursor | None = None, + after: MessageCursor | None = None, message_type: MessageType | None = None, limit: int = 50, ) -> tuple[list[MessageTable], bool]: @@ -1678,8 +1698,12 @@ class MessagingService(BaseService): Args: session_id: Session to get messages from - before: Get messages before this timestamp - after: Get messages after this timestamp + before: Keyset cursor — return messages older than (or, when the + cursor carries an id, equal-timestamp-but-smaller-id than) this + position. ``None`` for no upper bound. + after: Keyset cursor — return messages newer than (or + equal-timestamp-but-larger-id than) this position. ``None`` for + no lower bound. message_type: Filter by message type limit: Maximum messages to return @@ -1688,15 +1712,48 @@ class MessagingService(BaseService): """ query = select(MessageTable).where(MessageTable.session_id == session_id) - if before: - query = query.where(MessageTable.timestamp < before) - if after: - query = query.where(MessageTable.timestamp > after) + if before is not None: + if before.id is not None: + # Compound keyset cursor: resume past (before.timestamp, + # before.id) in desc order. Rows where (timestamp, id) < + # (before.timestamp, before.id): strictly older, OR same + # timestamp with a smaller id. Without this tie-break, + # equal-timestamp rows cut by ``limit`` on the prior page are + # excluded by the strict ``< before.timestamp`` and vanish. + query = query.where( + or_( + MessageTable.timestamp < before.timestamp, + and_( + MessageTable.timestamp == before.timestamp, + MessageTable.id < before.id, + ), + ) + ) + else: + query = query.where(MessageTable.timestamp < before.timestamp) + if after is not None: + if after.id is not None: + query = query.where( + or_( + MessageTable.timestamp > after.timestamp, + and_( + MessageTable.timestamp == after.timestamp, + MessageTable.id > after.id, + ), + ) + ) + else: + query = query.where(MessageTable.timestamp > after.timestamp) if message_type: query = query.where(MessageTable.type == message_type) - # Get one extra to check if there are more - query = query.order_by(MessageTable.timestamp.desc()).limit(limit + 1) + # Deterministic total order: timestamp desc, then id desc so the + # "last item" cursor is unambiguous for keyset pagination (without the + # id tie-break, same-timestamp rows have a non-deterministic order and + # the cursor is ambiguous). + query = query.order_by( + MessageTable.timestamp.desc(), MessageTable.id.desc() + ).limit(limit + 1) result = await self.session.execute(query) messages = list(result.scalars().all()) diff --git a/tests/integration/test_messaging_service.py b/tests/integration/test_messaging_service.py index 6f20dc68..00c09fe9 100644 --- a/tests/integration/test_messaging_service.py +++ b/tests/integration/test_messaging_service.py @@ -7,7 +7,7 @@ from datetime import UTC, datetime, timedelta from types import SimpleNamespace from typing import TYPE_CHECKING, Any, cast from unittest.mock import AsyncMock, patch -from uuid import uuid4 +from uuid import UUID, uuid4 from uuid import uuid4 as _u import pytest @@ -38,6 +38,7 @@ from roboco.models.session import ( from roboco.services.base import ConflictError, NotFoundError from roboco.services.messaging import ( ApiSessionCreate, + MessageCursor, MessagingService, get_messaging_service, ) @@ -1106,8 +1107,8 @@ async def test_get_messages_with_filters(msg_setup: dict) -> None: cutoff = datetime.now(UTC) - timedelta(hours=1) msgs, _ = await svc.get_messages( sess.id, - before=datetime.now(UTC) + timedelta(hours=1), - after=cutoff, + before=MessageCursor(datetime.now(UTC) + timedelta(hours=1)), + after=MessageCursor(cutoff), message_type=MessageType.DIALOGUE, ) assert isinstance(msgs, list) @@ -1130,6 +1131,113 @@ async def test_get_messages_with_limit(msg_setup: dict) -> None: assert has_more is True +async def _seed_messages_same_timestamp( + svc: MessagingService, session: Any, aid: UUID, count: int +) -> tuple[UUID, list[UUID]]: + """Create ``count`` messages in one session and force them ALL to the same + timestamp so the equal-timestamp pagination skip is reproducible. Returns + ``(session_id, message_ids)``.""" + ch = await svc.create_channel(_channel_req(uuid4().hex[:6])) + grp = await svc.create_group(GroupCreateRequest(name="g1", channel_id=ch.id)) + sess = await svc.create_session(SessionCreateRequest(group_id=grp.id)) + ids: list[UUID] = [] + for i in range(count): + m = await svc.send_message( + MessageCreateRequest(agent_id=aid, session_id=sess.id, content=f"m-{i}") + ) + ids.append(m.id) + fixed = datetime.now(UTC) + rows = ( + ( + await session.execute( + select(MessageTable).where(MessageTable.session_id == sess.id) + ) + ) + .scalars() + .all() + ) + for row in rows: + row.timestamp = fixed + await session.flush() + return sess.id, ids + + +@pytest.mark.asyncio +async def test_get_messages_compound_before_cursor_no_skip_on_equal_timestamps( + msg_setup: dict, +) -> None: + """Equal-timestamp messages must not be skipped across pages (F106). + + With a strict ``timestamp < before`` cursor and ``order_by(timestamp.desc())``, + messages sharing the page's last timestamp are cut by ``limit`` on page 1 + and excluded (``< T``) from page 2 — they vanish. The compound + ``(timestamp, id)`` cursor tie-breaks on the unique id so the next page + resumes exactly past the last id (no skip, no duplicate). Requires the + deterministic ``order_by(timestamp.desc(), id.desc())`` so the "last item" + cursor is unambiguous. + """ + svc = msg_setup["svc"] + session = svc.session + aid = msg_setup["agent_id"] + total = 5 + sess_id, ids = await _seed_messages_same_timestamp(svc, session, aid, count=total) + + page_size = 3 + page1, has_more = await svc.get_messages(sess_id, limit=page_size) + assert has_more is True + assert len(page1) == page_size + # Deterministic order: same timestamp → by id.desc() + page1_ids = [m.id for m in page1] + assert page1_ids == sorted(page1_ids, reverse=True) + + last = page1[-1] + page2, has_more2 = await svc.get_messages( + sess_id, before=MessageCursor(last.timestamp, last.id), limit=page_size + ) + # All 5 covered, no skip, no duplicate. + seen = {m.id for m in page1} | {m.id for m in page2} + assert seen == set(ids) + assert len(page1) + len(page2) == total + assert has_more2 is False + + +@pytest.mark.asyncio +async def test_get_messages_compound_after_cursor_no_skip_on_equal_timestamps( + msg_setup: dict, +) -> None: + """Forward pagination (``after``) with the compound ``(timestamp, id)`` + cursor tie-breaks on id so newer-direction pagination across equal + timestamps skips nothing either (F106). + + With a strict ``timestamp > after`` cursor, every row sharing the cursor's + timestamp is EXCLUDED — so forward-paginating from a middle message would + return nothing (all rows are at the same timestamp, none strictly greater). + The compound cursor includes same-timestamp rows with a larger id. + """ + svc = msg_setup["svc"] + session = svc.session + aid = msg_setup["agent_id"] + total = 5 + fetch_limit = 10 + mid_index = 2 + newer_count = 2 + sess_id, _ids = await _seed_messages_same_timestamp(svc, session, aid, count=total) + + # Full desc ordering: [id_largest, ..., id_smallest]. + all_msgs, _ = await svc.get_messages(sess_id, limit=fetch_limit) + assert len(all_msgs) == total + # Pick the middle row as the forward cursor. + mid = all_msgs[mid_index] + newer = {m.id for m in all_msgs[:newer_count]} # the rows above mid (larger id) + + page, _ = await svc.get_messages( + sess_id, after=MessageCursor(mid.timestamp, mid.id), limit=fetch_limit + ) + # The same-timestamp rows newer than mid are returned — NOT skipped. + assert {m.id for m in page} == newer + assert len(page) == newer_count + + @pytest.mark.asyncio async def test_get_message_context_redirects_when_session_closed( msg_setup: dict,