Skip to content

Commit cc846e7

Browse files
kabirkabir
committed
More review fixes
1 parent 757ed8a commit cc846e7

File tree

3 files changed

+94
-44
lines changed

3 files changed

+94
-44
lines changed

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

Lines changed: 14 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -239,10 +239,12 @@ public ListTasksResult list(ListTasksParams params) {
239239
String[] tokenParts = params.pageToken().split(":", 2);
240240
if (tokenParts.length == 2) {
241241
// Keyset pagination: get tasks where timestamp < tokenTimestamp OR (timestamp = tokenTimestamp AND id > tokenId)
242+
// All tasks have timestamps (TaskStatus canonical constructor ensures this)
242243
queryBuilder.append(" AND (t.statusTimestamp < :tokenTimestamp OR (t.statusTimestamp = :tokenTimestamp AND t.id > :tokenId))");
243244
} else {
244-
// Fallback for legacy pageToken format (ID only) - for backward compatibility during transition
245-
queryBuilder.append(" AND t.id > :pageToken");
245+
// Legacy ID-only pageToken format is not supported with timestamp-based sorting
246+
// Throw error to prevent incorrect pagination results
247+
throw new io.a2a.spec.InvalidParamsError(null, "Invalid pageToken format: expected 'timestamp:id'", null);
246248
}
247249
}
248250

@@ -268,18 +270,19 @@ public ListTasksResult list(ListTasksParams params) {
268270
// Parse keyset pagination parameters
269271
try {
270272
long timestampMillis = Long.parseLong(tokenParts[0]);
271-
Instant tokenTimestamp = Instant.ofEpochMilli(timestampMillis);
272273
String tokenId = tokenParts[1];
274+
275+
// All tasks have timestamps (TaskStatus canonical constructor ensures this)
276+
Instant tokenTimestamp = Instant.ofEpochMilli(timestampMillis);
273277
query.setParameter("tokenTimestamp", tokenTimestamp);
274278
query.setParameter("tokenId", tokenId);
275279
} catch (NumberFormatException e) {
276-
// Invalid timestamp format, fall back to ID-only
277-
query.setParameter("pageToken", params.pageToken());
280+
// Malformed timestamp in pageToken
281+
throw new io.a2a.spec.InvalidParamsError(null,
282+
"Invalid pageToken format: timestamp must be numeric milliseconds", null);
278283
}
279-
} else {
280-
// Legacy format (ID only)
281-
query.setParameter("pageToken", params.pageToken());
282284
}
285+
// Note: Legacy ID-only format already rejected in query building phase
283286
}
284287

285288
// Apply page size limit (+1 to check for next page)
@@ -324,13 +327,9 @@ public ListTasksResult list(ListTasksParams params) {
324327
String nextPageToken = null;
325328
if (hasMore && !tasks.isEmpty()) {
326329
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-
}
330+
// All tasks have timestamps (TaskStatus canonical constructor ensures this)
331+
long timestampMillis = lastTask.getStatus().timestamp().toInstant().toEpochMilli();
332+
nextPageToken = timestampMillis + ":" + lastTask.getId();
334333
}
335334

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

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

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -692,6 +692,49 @@ public void testListTasksDefaultPageSize() {
692692
assertNotNull(result.nextPageToken(), "Should have next page");
693693
}
694694

695+
@Test
696+
@Transactional
697+
public void testListTasksInvalidPageTokenFormat() {
698+
// Create a task
699+
Task task = new Task.Builder()
700+
.id("task-invalid-token")
701+
.contextId("context-invalid-token")
702+
.status(new TaskStatus(TaskState.WORKING))
703+
.build();
704+
taskStore.save(task);
705+
706+
// Test 1: Legacy ID-only pageToken should throw InvalidParamsError
707+
ListTasksParams params1 = new ListTasksParams.Builder()
708+
.contextId("context-invalid-token")
709+
.pageToken("task-invalid-token") // ID-only format (legacy)
710+
.build();
711+
712+
try {
713+
taskStore.list(params1);
714+
throw new AssertionError("Expected InvalidParamsError for legacy ID-only pageToken");
715+
} catch (io.a2a.spec.InvalidParamsError e) {
716+
// Expected - legacy format not supported
717+
assertTrue(e.getMessage().contains("Invalid pageToken format"),
718+
"Error message should mention invalid format");
719+
}
720+
721+
// Test 2: Malformed timestamp in pageToken should throw InvalidParamsError
722+
ListTasksParams params2 = new ListTasksParams.Builder()
723+
.contextId("context-invalid-token")
724+
.pageToken("not-a-number:task-id") // Invalid timestamp
725+
.build();
726+
727+
try {
728+
taskStore.list(params2);
729+
throw new AssertionError("Expected InvalidParamsError for malformed timestamp");
730+
} catch (io.a2a.spec.InvalidParamsError e) {
731+
// Expected - malformed timestamp
732+
assertTrue(e.getMessage().contains("timestamp must be numeric"),
733+
"Error message should mention numeric timestamp requirement");
734+
}
735+
}
736+
737+
695738
@Test
696739
@Transactional
697740
public void testListTasksOrderingById() {

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

Lines changed: 37 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,11 @@ public ListTasksResult list(ListTasksParams params) {
4747
(task.getStatus() != null &&
4848
task.getStatus().timestamp() != null &&
4949
task.getStatus().timestamp().toInstant().isAfter(params.lastUpdatedAfter())))
50-
.sorted(Comparator.comparing((Task t) -> t.getStatus().timestamp(),
50+
.sorted(Comparator.comparing(
51+
(Task t) -> (t.getStatus() != null && t.getStatus().timestamp() != null)
52+
// Truncate to milliseconds for consistency with pageToken precision
53+
? t.getStatus().timestamp().toInstant().truncatedTo(java.time.temporal.ChronoUnit.MILLIS)
54+
: null,
5155
Comparator.nullsLast(Comparator.reverseOrder()))
5256
.thenComparing(Task::getId))
5357
.toList();
@@ -59,7 +63,7 @@ public ListTasksResult list(ListTasksParams params) {
5963
int startIndex = 0;
6064

6165
// Handle page token using keyset pagination (format: "timestamp_millis:taskId")
62-
// Find the first task that comes after the pageToken position in sorted order
66+
// Use binary search to efficiently find the first task after the pageToken position (O(log N))
6367
if (params.pageToken() != null && !params.pageToken().isEmpty()) {
6468
String[] tokenParts = params.pageToken().split(":", 2);
6569
if (tokenParts.length == 2) {
@@ -68,31 +72,39 @@ public ListTasksResult list(ListTasksParams params) {
6872
java.time.Instant tokenTimestamp = java.time.Instant.ofEpochMilli(tokenTimestampMillis);
6973
String tokenId = tokenParts[1];
7074

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-
}
75+
// Binary search for first task where: timestamp < tokenTimestamp OR (timestamp == tokenTimestamp AND id > tokenId)
76+
// Since list is sorted (timestamp DESC, id ASC), we search for the insertion point
77+
int left = 0;
78+
int right = allFilteredTasks.size();
79+
80+
while (left < right) {
81+
int mid = left + (right - left) / 2;
82+
Task task = allFilteredTasks.get(mid);
83+
84+
// All tasks have timestamps (TaskStatus canonical constructor ensures this)
85+
// Truncate to milliseconds for consistency with pageToken precision
86+
java.time.Instant taskTimestamp = task.getStatus().timestamp().toInstant()
87+
.truncatedTo(java.time.temporal.ChronoUnit.MILLIS);
88+
int timestampCompare = taskTimestamp.compareTo(tokenTimestamp);
89+
90+
if (timestampCompare < 0 || (timestampCompare == 0 && task.getId().compareTo(tokenId) > 0)) {
91+
// This task is after the token, search left half
92+
right = mid;
93+
} else {
94+
// This task is before or equal to token, search right half
95+
left = mid + 1;
8296
}
8397
}
98+
startIndex = left;
8499
} catch (NumberFormatException e) {
85-
// Invalid pageToken format, start from beginning
86-
startIndex = 0;
100+
// Malformed timestamp in pageToken
101+
throw new io.a2a.spec.InvalidParamsError(null,
102+
"Invalid pageToken format: timestamp must be numeric milliseconds", null);
87103
}
88104
} 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-
}
95-
}
105+
// Legacy ID-only pageToken format is not supported with timestamp-based sorting
106+
// Throw error to prevent incorrect pagination results
107+
throw new io.a2a.spec.InvalidParamsError(null, "Invalid pageToken format: expected 'timestamp:id'", null);
96108
}
97109
}
98110

@@ -104,13 +116,9 @@ public ListTasksResult list(ListTasksParams params) {
104116
String nextPageToken = null;
105117
if (endIndex < allFilteredTasks.size()) {
106118
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-
}
119+
// All tasks have timestamps (TaskStatus canonical constructor ensures this)
120+
long timestampMillis = lastTask.getStatus().timestamp().toInstant().toEpochMilli();
121+
nextPageToken = timestampMillis + ":" + lastTask.getId();
114122
}
115123

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

0 commit comments

Comments
 (0)