Cursor-based pagination and update time handling

Author: lkawka

src/a2a/server/tasks/database_task_store.py

@@ -1,8 +1,17 @@
 import logging
 
+from datetime import datetime, timezone
+
 
 try:
-    from sqlalchemy import Table, delete, func, select
+    from sqlalchemy import (
+        Table,
+        and_,
+        delete,
+        func,
+        or_,
+        select,
+    )
     from sqlalchemy.ext.asyncio import (
         AsyncEngine,
         AsyncSession,
@@ -24,6 +33,7 @@
 from a2a.server.tasks.task_store import TaskStore, TasksPage
 from a2a.types import ListTasksParams, Task
 from a2a.utils.constants import DEFAULT_LIST_TASKS_PAGE_SIZE
+from a2a.utils.task import decode_page_token, encode_page_token
 
 
 logger = logging.getLogger(__name__)
@@ -154,44 +164,85 @@ async def list(
         """Retrieves all tasks from the database."""
         await self._ensure_initialized()
         async with self.async_session_maker() as session:
-            page_number = int(params.page_token) if params.page_token else 0
-            page_size = params.page_size or DEFAULT_LIST_TASKS_PAGE_SIZE
-            offset = page_number * page_size
-
-            # Base query for filtering
             base_stmt = select(self.task_model)
+
+            # Add filters
             if params.context_id:
                 base_stmt = base_stmt.where(
                     self.task_model.context_id == params.context_id
                 )
-            if params.status is not None:
+            if params.status and params.status != 'unknown':
                 base_stmt = base_stmt.where(
                     self.task_model.status['state'].as_string()
                     == params.status.value
                 )
+            if params.last_updated_after:
+                last_updated_after_iso = datetime.fromtimestamp(
+                    params.last_updated_after / 1000, tz=timezone.utc
+                ).isoformat()
+                base_stmt = base_stmt.where(
+                    self.task_model.status['timestamp'].as_string()
+                    >= last_updated_after_iso
+                )
 
             # Get total count
             count_stmt = select(func.count()).select_from(base_stmt.alias())
             total_count = (await session.execute(count_stmt)).scalar_one()
 
-            # Get paginated results
-            stmt = (
-                base_stmt.order_by(self.task_model.id.desc())
-                .limit(page_size)
-                .offset(offset)
+            stmt = base_stmt.order_by(
+                self.task_model.status['timestamp']
+                .as_string()
+                .desc()
+                .nulls_last(),
+                self.task_model.id.desc(),
             )
+
+            # Get paginated results
+            if params.page_token:
+                start_task_id = decode_page_token(params.page_token)
+                start_task = (
+                    await session.execute(
+                        select(self.task_model).where(
+                            self.task_model.id == start_task_id
+                        )
+                    )
+                ).scalar_one_or_none()
+                if not start_task:
+                    raise ValueError(f'Invalid page token: {params.page_token}')
+                if start_task.status.timestamp:
+                    stmt = stmt.where(
+                        or_(
+                            self.task_model.status['timestamp']
+                            .as_string()
+                            .is_(None),
+                            self.task_model.status['timestamp'].as_string()
+                            >= start_task.status.timestamp,
+                        )
+                    )
+                else:
+                    stmt = stmt.where(
+                        and_(
+                            self.task_model.status['timestamp']
+                            .as_string()
+                            .is_(None),
+                            self.task_model.id <= start_task.id,
+                        )
+                    )
+            page_size = params.page_size or DEFAULT_LIST_TASKS_PAGE_SIZE
+            stmt = stmt.limit(page_size + 1)  # Add 1 for next page token
+
             result = await session.execute(stmt)
             tasks_models = result.scalars().all()
             tasks = [self._from_orm(task_model) for task_model in tasks_models]
 
             next_page_token = (
-                str(page_number + 1)
-                if total_count > (page_number + 1) * page_size
+                encode_page_token(tasks[-1].id)
+                if len(tasks) == page_size + 1
                 else None
             )
 
             return TasksPage(
-                tasks=tasks,
+                tasks=tasks[:page_size],
                 total_size=total_count,
                 next_page_token=next_page_token,
             )

src/a2a/server/tasks/inmemory_task_store.py

@@ -1,10 +1,13 @@
 import asyncio
 import logging
 
+from datetime import datetime, timezone
+
 from a2a.server.context import ServerCallContext
 from a2a.server.tasks.task_store import TaskStore, TasksPage
 from a2a.types import ListTasksParams, Task
 from a2a.utils.constants import DEFAULT_LIST_TASKS_PAGE_SIZE
+from a2a.utils.task import decode_page_token, encode_page_token
 
 
 logger = logging.getLogger(__name__)
@@ -53,27 +56,58 @@ async def list(
         async with self.lock:
             tasks = list(self.tasks.values())
 
-        # Apply filtering
+        # Filter tasks
         if params.context_id:
             tasks = [
                 task for task in tasks if task.context_id == params.context_id
             ]
-        if params.status is not None:
+        if params.status and params.status != 'unknown':
             tasks = [
                 task for task in tasks if task.status.state == params.status
             ]
+        if params.last_updated_after:
+            last_updated_after_iso = datetime.fromtimestamp(
+                params.last_updated_after / 1000, tz=timezone.utc
+            ).isoformat()
+            tasks = [
+                task
+                for task in tasks
+                if (
+                    task.status.timestamp
+                    and task.status.timestamp >= last_updated_after_iso
+                )
+            ]
 
-        # Apply pagination
-        total_size = len(tasks)
-        page_token = int(params.page_token) if params.page_token else 0
-        page_size = params.page_size or DEFAULT_LIST_TASKS_PAGE_SIZE
-        tasks = tasks[page_token * page_size : (page_token + 1) * page_size]
+        # Order tasks by last update time. To ensure stable sorting, in cases where timestamps are null or not unique, do a second order comparison of IDs.
+        tasks.sort(
+            key=lambda task: (
+                task.status.timestamp is not None,
+                task.status.timestamp,
+                task.id,
+            ),
+            reverse=True,
+        )
 
+        # Paginate tasks
+        total_size = len(tasks)
+        start_idx = 0
+        if params.page_token:
+            start_task_id = decode_page_token(params.page_token)
+            valid_token = False
+            for i, task in enumerate(tasks):
+                if task.id == start_task_id:
+                    start_idx = i
+                    valid_token = True
+                    break
+            if not valid_token:
+                raise ValueError(f'Invalid page token: {params.page_token}')
+        end_idx = start_idx + (params.page_size or DEFAULT_LIST_TASKS_PAGE_SIZE)
         next_page_token = (
-            str(page_token + 1)
-            if (page_token + 1) * page_size < total_size
+            encode_page_token(tasks[end_idx].id)
+            if end_idx < total_size
             else None
         )
+        tasks = tasks[start_idx:end_idx]
 
         return TasksPage(
             next_page_token=next_page_token,

src/a2a/utils/task.py

@@ -1,7 +1,10 @@
 """Utility functions for creating A2A Task objects."""
 
+import binascii
 import uuid
 
+from base64 import b64decode, b64encode
+
 from a2a.types import Artifact, Message, Task, TaskState, TaskStatus, TextPart
 
 
@@ -92,3 +95,39 @@ def apply_history_length(task: Task, history_length: int | None) -> Task:
         return task.model_copy(update={'history': limited_history})
 
     return task
+
+
+_ENCODING = 'utf-8'
+
+
+def encode_page_token(task_id: str) -> str:
+    """Encodes page token for tasks pagination.
+
+    Args:
+        task_id: The ID of the task.
+
+    Returns:
+        The encoded page token.
+    """
+    return b64encode(task_id.encode(_ENCODING)).decode(_ENCODING)
+
+
+def decode_page_token(page_token: str) -> str:
+    """Decodes page token for tasks pagination.
+
+    Args:
+        page_token: The encoded page token.
+
+    Returns:
+        The decoded task ID.
+    """
+    encoded_str = page_token
+    missing_padding = len(encoded_str) % 4
+    if missing_padding:
+        encoded_str += '=' * (4 - missing_padding)
+    print(f'input: {encoded_str}')
+    try:
+        decoded = b64decode(encoded_str.encode(_ENCODING)).decode(_ENCODING)
+    except (binascii.Error, UnicodeDecodeError) as e:
+        raise ValueError('Token is not a valid base64-encoded cursor.') from e
+    return decoded