Skip to content

Commit 757ed8a

Browse files
kabirkabir
committed
Gemini review
1 parent 856115a commit 757ed8a

File tree

5 files changed

+198
-17
lines changed

5 files changed

+198
-17
lines changed

‎extras/task-store-database-jpa/src/main/java/io/a2a/extras/taskstore/database/jpa/JpaDatabaseTaskStore.java‎

Lines changed: 37 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -233,9 +233,17 @@ public ListTasksResult list(ListTasksParams params) {
233233
countQueryBuilder.append(" AND t.statusTimestamp > :lastUpdatedAfter");
234234
}
235235

236-
// Apply pagination cursor (tasks after pageToken)
236+
// Apply pagination cursor using keyset pagination for composite sort (timestamp DESC, id ASC)
237+
// PageToken format: "timestamp_millis:taskId" (e.g., "1699999999000:task-123")
237238
if (params.pageToken() != null && !params.pageToken().isEmpty()) {
238-
queryBuilder.append(" AND t.id > :pageToken");
239+
String[] tokenParts = params.pageToken().split(":", 2);
240+
if (tokenParts.length == 2) {
241+
// Keyset pagination: get tasks where timestamp < tokenTimestamp OR (timestamp = tokenTimestamp AND id > tokenId)
242+
queryBuilder.append(" AND (t.statusTimestamp < :tokenTimestamp OR (t.statusTimestamp = :tokenTimestamp AND t.id > :tokenId))");
243+
} else {
244+
// Fallback for legacy pageToken format (ID only) - for backward compatibility during transition
245+
queryBuilder.append(" AND t.id > :pageToken");
246+
}
239247
}
240248

241249
// Sort by status timestamp descending (most recent first), then by ID for stable ordering
@@ -255,7 +263,23 @@ public ListTasksResult list(ListTasksParams params) {
255263
query.setParameter("lastUpdatedAfter", params.lastUpdatedAfter());
256264
}
257265
if (params.pageToken() != null && !params.pageToken().isEmpty()) {
258-
query.setParameter("pageToken", params.pageToken());
266+
String[] tokenParts = params.pageToken().split(":", 2);
267+
if (tokenParts.length == 2) {
268+
// Parse keyset pagination parameters
269+
try {
270+
long timestampMillis = Long.parseLong(tokenParts[0]);
271+
Instant tokenTimestamp = Instant.ofEpochMilli(timestampMillis);
272+
String tokenId = tokenParts[1];
273+
query.setParameter("tokenTimestamp", tokenTimestamp);
274+
query.setParameter("tokenId", tokenId);
275+
} catch (NumberFormatException e) {
276+
// Invalid timestamp format, fall back to ID-only
277+
query.setParameter("pageToken", params.pageToken());
278+
}
279+
} else {
280+
// Legacy format (ID only)
281+
query.setParameter("pageToken", params.pageToken());
282+
}
259283
}
260284

261285
// Apply page size limit (+1 to check for next page)
@@ -295,10 +319,18 @@ public ListTasksResult list(ListTasksParams params) {
295319
}
296320
}
297321

298-
// Determine next page token (ID of last task if there are more results)
322+
// Determine next page token (timestamp:ID of last task if there are more results)
323+
// Format: "timestamp_millis:taskId" for keyset pagination
299324
String nextPageToken = null;
300325
if (hasMore && !tasks.isEmpty()) {
301-
nextPageToken = tasks.get(tasks.size() - 1).getId();
326+
Task lastTask = tasks.get(tasks.size() - 1);
327+
if (lastTask.getStatus() != null && lastTask.getStatus().timestamp() != null) {
328+
long timestampMillis = lastTask.getStatus().timestamp().toInstant().toEpochMilli();
329+
nextPageToken = timestampMillis + ":" + lastTask.getId();
330+
} else {
331+
// Fallback to ID-only if timestamp is missing (shouldn't happen with proper data)
332+
nextPageToken = lastTask.getId();
333+
}
302334
}
303335

304336
// Apply post-processing transformations (history limiting, artifact removal)

‎extras/task-store-database-jpa/src/main/java/io/a2a/extras/taskstore/database/jpa/JpaTask.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -145,8 +145,9 @@ private void updateDenormalizedFields(Task task) {
145145
io.a2a.spec.TaskState taskState = task.getStatus().state();
146146
this.state = (taskState != null) ? taskState.asString() : null;
147147
// Extract status timestamp for efficient querying and sorting
148+
// Truncate to milliseconds for keyset pagination consistency (pageToken uses millis)
148149
this.statusTimestamp = (task.getStatus().timestamp() != null)
149-
? task.getStatus().timestamp().toInstant()
150+
? task.getStatus().timestamp().toInstant().truncatedTo(java.time.temporal.ChronoUnit.MILLIS)
150151
: null;
151152
} else {
152153
this.state = null;

‎extras/task-store-database-jpa/src/test/java/io/a2a/extras/taskstore/database/jpa/JpaDatabaseTaskStoreTest.java‎

Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -468,6 +468,122 @@ public void testListTasksPagination() {
468468
assertNull(result3.nextPageToken(), "Last page should have no next page token");
469469
}
470470

471+
@Test
472+
@Transactional
473+
public void testListTasksPaginationWithDifferentTimestamps() {
474+
// Create tasks with different timestamps to verify keyset pagination
475+
// with composite sort (timestamp DESC, id ASC)
476+
OffsetDateTime now = OffsetDateTime.now(java.time.ZoneOffset.UTC);
477+
478+
// Task 1: 10 minutes ago, ID="task-diff-a"
479+
Task task1 = new Task.Builder()
480+
.id("task-diff-a")
481+
.contextId("context-diff-timestamps")
482+
.status(new TaskStatus(TaskState.WORKING, null, now.minusMinutes(10)))
483+
.build();
484+
taskStore.save(task1);
485+
486+
// Task 2: 5 minutes ago, ID="task-diff-b"
487+
Task task2 = new Task.Builder()
488+
.id("task-diff-b")
489+
.contextId("context-diff-timestamps")
490+
.status(new TaskStatus(TaskState.WORKING, null, now.minusMinutes(5)))
491+
.build();
492+
taskStore.save(task2);
493+
494+
// Task 3: 5 minutes ago, ID="task-diff-c" (same timestamp as task2, tests ID tie-breaker)
495+
Task task3 = new Task.Builder()
496+
.id("task-diff-c")
497+
.contextId("context-diff-timestamps")
498+
.status(new TaskStatus(TaskState.WORKING, null, now.minusMinutes(5)))
499+
.build();
500+
taskStore.save(task3);
501+
502+
// Task 4: Now, ID="task-diff-d"
503+
Task task4 = new Task.Builder()
504+
.id("task-diff-d")
505+
.contextId("context-diff-timestamps")
506+
.status(new TaskStatus(TaskState.WORKING, null, now))
507+
.build();
508+
taskStore.save(task4);
509+
510+
// Task 5: 1 minute ago, ID="task-diff-e"
511+
Task task5 = new Task.Builder()
512+
.id("task-diff-e")
513+
.contextId("context-diff-timestamps")
514+
.status(new TaskStatus(TaskState.WORKING, null, now.minusMinutes(1)))
515+
.build();
516+
taskStore.save(task5);
517+
518+
// Expected order (timestamp DESC, id ASC):
519+
// 1. task-diff-d (now)
520+
// 2. task-diff-e (1 min ago)
521+
// 3. task-diff-b (5 min ago, ID 'b')
522+
// 4. task-diff-c (5 min ago, ID 'c')
523+
// 5. task-diff-a (10 min ago)
524+
525+
// Page 1: Get first 2 tasks
526+
ListTasksParams params1 = new ListTasksParams.Builder()
527+
.contextId("context-diff-timestamps")
528+
.pageSize(2)
529+
.build();
530+
ListTasksResult result1 = taskStore.list(params1);
531+
532+
assertEquals(5, result1.totalSize());
533+
assertEquals(2, result1.pageSize());
534+
assertNotNull(result1.nextPageToken(), "Should have next page token");
535+
536+
// Verify first page order
537+
assertEquals("task-diff-d", result1.tasks().get(0).getId(), "First task should be most recent");
538+
assertEquals("task-diff-e", result1.tasks().get(1).getId(), "Second task should be 1 min ago");
539+
540+
// Verify pageToken format: "timestamp_millis:taskId"
541+
assertTrue(result1.nextPageToken().contains(":"), "PageToken should have format timestamp:id");
542+
String[] tokenParts = result1.nextPageToken().split(":", 2);
543+
assertEquals(2, tokenParts.length, "PageToken should have exactly 2 parts");
544+
assertEquals("task-diff-e", tokenParts[1], "PageToken should contain last task ID");
545+
546+
// Page 2: Get next 2 tasks
547+
ListTasksParams params2 = new ListTasksParams.Builder()
548+
.contextId("context-diff-timestamps")
549+
.pageSize(2)
550+
.pageToken(result1.nextPageToken())
551+
.build();
552+
ListTasksResult result2 = taskStore.list(params2);
553+
554+
assertEquals(5, result2.totalSize());
555+
assertEquals(2, result2.pageSize());
556+
assertNotNull(result2.nextPageToken(), "Should have next page token");
557+
558+
// Verify second page order (tasks with same timestamp, sorted by ID)
559+
assertEquals("task-diff-b", result2.tasks().get(0).getId(), "Third task should be 5 min ago, ID 'b'");
560+
assertEquals("task-diff-c", result2.tasks().get(1).getId(), "Fourth task should be 5 min ago, ID 'c'");
561+
562+
// Page 3: Get last task
563+
ListTasksParams params3 = new ListTasksParams.Builder()
564+
.contextId("context-diff-timestamps")
565+
.pageSize(2)
566+
.pageToken(result2.nextPageToken())
567+
.build();
568+
ListTasksResult result3 = taskStore.list(params3);
569+
570+
assertEquals(5, result3.totalSize());
571+
assertEquals(1, result3.pageSize());
572+
assertNull(result3.nextPageToken(), "Last page should have no next page token");
573+
574+
// Verify last task
575+
assertEquals("task-diff-a", result3.tasks().get(0).getId(), "Last task should be oldest");
576+
577+
// Verify no duplicates across all pages
578+
List<String> allTaskIds = new ArrayList<>();
579+
allTaskIds.addAll(result1.tasks().stream().map(Task::getId).toList());
580+
allTaskIds.addAll(result2.tasks().stream().map(Task::getId).toList());
581+
allTaskIds.addAll(result3.tasks().stream().map(Task::getId).toList());
582+
583+
assertEquals(5, allTaskIds.size(), "Should have exactly 5 tasks across all pages");
584+
assertEquals(5, allTaskIds.stream().distinct().count(), "Should have no duplicate tasks");
585+
}
586+
471587
@Test
472588
@Transactional
473589
public void testListTasksHistoryLimiting() {

‎reference/rest/src/main/java/io/a2a/server/rest/quarkus/A2AServerRoutes.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -144,7 +144,7 @@ public void listTasks(RoutingContext rc) {
144144

145145
Boolean includeArtifacts = null;
146146
if (includeArtifactsStr != null && !includeArtifactsStr.isEmpty()) {
147-
includeArtifacts = Boolean.parseBoolean(includeArtifactsStr);
147+
includeArtifacts = Boolean.valueOf(includeArtifactsStr);
148148
}
149149

150150
response = jsonRestHandler.listTasks(contextId, statusStr, pageSize, pageToken,

‎server-common/src/main/java/io/a2a/server/tasks/InMemoryTaskStore.java‎

Lines changed: 42 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -58,27 +58,59 @@ public ListTasksResult list(ListTasksParams params) {
5858
int pageSize = params.getEffectivePageSize();
5959
int startIndex = 0;
6060

61-
// Handle page token (cursor: task ID from previous page)
62-
// Since we're sorted by timestamp DESC then ID ASC, we can't use binary search
63-
// Instead, find the task with the matching ID using linear search
61+
// Handle page token using keyset pagination (format: "timestamp_millis:taskId")
62+
// Find the first task that comes after the pageToken position in sorted order
6463
if (params.pageToken() != null && !params.pageToken().isEmpty()) {
65-
for (int i = 0; i < allFilteredTasks.size(); i++) {
66-
if (allFilteredTasks.get(i).getId().equals(params.pageToken())) {
67-
startIndex = i + 1;
68-
break;
64+
String[] tokenParts = params.pageToken().split(":", 2);
65+
if (tokenParts.length == 2) {
66+
try {
67+
long tokenTimestampMillis = Long.parseLong(tokenParts[0]);
68+
java.time.Instant tokenTimestamp = java.time.Instant.ofEpochMilli(tokenTimestampMillis);
69+
String tokenId = tokenParts[1];
70+
71+
// Find first task where: timestamp < tokenTimestamp OR (timestamp == tokenTimestamp AND id > tokenId)
72+
for (int i = 0; i < allFilteredTasks.size(); i++) {
73+
Task task = allFilteredTasks.get(i);
74+
if (task.getStatus() != null && task.getStatus().timestamp() != null) {
75+
java.time.Instant taskTimestamp = task.getStatus().timestamp().toInstant();
76+
int timestampCompare = taskTimestamp.compareTo(tokenTimestamp);
77+
78+
if (timestampCompare < 0 || (timestampCompare == 0 && task.getId().compareTo(tokenId) > 0)) {
79+
startIndex = i;
80+
break;
81+
}
82+
}
83+
}
84+
} catch (NumberFormatException e) {
85+
// Invalid pageToken format, start from beginning
86+
startIndex = 0;
87+
}
88+
} else {
89+
// Legacy format (ID only) - fallback for backward compatibility
90+
for (int i = 0; i < allFilteredTasks.size(); i++) {
91+
if (allFilteredTasks.get(i).getId().equals(params.pageToken())) {
92+
startIndex = i + 1;
93+
break;
94+
}
6995
}
7096
}
71-
// If not found, startIndex remains 0 (start from beginning)
7297
}
7398

7499
// Get the page of tasks
75100
int endIndex = Math.min(startIndex + pageSize, allFilteredTasks.size());
76101
List<Task> pageTasks = allFilteredTasks.subList(startIndex, endIndex);
77102

78-
// Determine next page token
103+
// Determine next page token (format: "timestamp_millis:taskId")
79104
String nextPageToken = null;
80105
if (endIndex < allFilteredTasks.size()) {
81-
nextPageToken = allFilteredTasks.get(endIndex - 1).getId();
106+
Task lastTask = allFilteredTasks.get(endIndex - 1);
107+
if (lastTask.getStatus() != null && lastTask.getStatus().timestamp() != null) {
108+
long timestampMillis = lastTask.getStatus().timestamp().toInstant().toEpochMilli();
109+
nextPageToken = timestampMillis + ":" + lastTask.getId();
110+
} else {
111+
// Fallback to ID-only if timestamp is missing
112+
nextPageToken = lastTask.getId();
113+
}
82114
}
83115

84116
// Transform tasks: limit history and optionally remove artifacts

0 commit comments

Comments
 (0)