Review fixes

Author: kabir

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

@@ -6,6 +6,7 @@
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
+import java.time.OffsetDateTime;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
@@ -420,7 +421,7 @@ public void testListTasksCombinedFilters() {
     public void testListTasksPagination() {
         // Create 5 tasks with same timestamp to ensure ID-based pagination works
         // (With timestamp DESC sorting, same timestamps allow ID ASC tie-breaking)
-        java.time.OffsetDateTime sameTimestamp = java.time.OffsetDateTime.now(java.time.ZoneOffset.UTC);
+        OffsetDateTime sameTimestamp = OffsetDateTime.now(java.time.ZoneOffset.UTC);
         for (int i = 1; i <= 5; i++) {
             Task task = new Task.Builder()
                     .id("task-page-" + i)
@@ -580,7 +581,7 @@ public void testListTasksDefaultPageSize() {
     public void testListTasksOrderingById() {
         // Create tasks with same timestamp to test ID-based tie-breaking
         // (spec requires sorting by timestamp DESC, then ID ASC)
-        java.time.OffsetDateTime sameTimestamp = java.time.OffsetDateTime.now(java.time.ZoneOffset.UTC);
+        OffsetDateTime sameTimestamp = OffsetDateTime.now(java.time.ZoneOffset.UTC);
 
         Task task1 = new Task.Builder()
                 .id("task-order-a")

server-common/src/main/java/io/a2a/server/requesthandlers/DefaultRequestHandler.java

@@ -5,8 +5,11 @@
 import static io.a2a.server.util.async.AsyncUtils.processor;
 import static java.util.concurrent.TimeUnit.*;
 
+import java.time.Instant;
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
@@ -172,9 +175,9 @@ public ListTasksResult onListTasks(ListTasksParams params, ServerCallContext con
         // Validate lastUpdatedAfter timestamp if provided
         if (params.lastUpdatedAfter() != null) {
             // Check if timestamp is in the future (optional validation per spec)
-            java.time.Instant now = java.time.Instant.now();
+            Instant now = Instant.now();
             if (params.lastUpdatedAfter().isAfter(now)) {
-                java.util.Map<String, Object> errorData = new java.util.HashMap<>();
+                Map<String, Object> errorData = new HashMap<>();
                 errorData.put("parameter", "lastUpdatedAfter");
                 errorData.put("reason", "Timestamp cannot be in the future");
                 throw new InvalidParamsError(null, "Invalid params", errorData);

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

@@ -38,28 +38,17 @@ public void delete(String taskId) {
 
     @Override
     public ListTasksResult list(ListTasksParams params) {
-        Stream<Task> taskStream = tasks.values().stream();
-
-        // Apply filters
-        if (params.contextId() != null) {
-            taskStream = taskStream.filter(task -> params.contextId().equals(task.getContextId()));
-        }
-        if (params.status() != null) {
-            taskStream = taskStream.filter(task ->
-                task.getStatus() != null && params.status().equals(task.getStatus().state())
-            );
-        }
-        if (params.lastUpdatedAfter() != null) {
-            taskStream = taskStream.filter(task ->
-                task.getStatus() != null &&
-                task.getStatus().timestamp() != null &&
-                task.getStatus().timestamp().toInstant().isAfter(params.lastUpdatedAfter())
-            );
-        }
-
-        // Sort by status timestamp descending (most recent first), then by ID ascending for stable ordering
-        List<Task> allFilteredTasks = taskStream
-                .sorted(Comparator.comparing((Task t) -> t.getStatus().timestamp(), Comparator.nullsLast(Comparator.reverseOrder()))
+        // Filter and sort tasks in a single stream pipeline
+        List<Task> allFilteredTasks = tasks.values().stream()
+                .filter(task -> params.contextId() == null || params.contextId().equals(task.getContextId()))
+                .filter(task -> params.status() == null ||
+                        (task.getStatus() != null && params.status().equals(task.getStatus().state())))
+                .filter(task -> params.lastUpdatedAfter() == null ||
+                        (task.getStatus() != null &&
+                         task.getStatus().timestamp() != null &&
+                         task.getStatus().timestamp().toInstant().isAfter(params.lastUpdatedAfter())))
+                .sorted(Comparator.comparing((Task t) -> t.getStatus().timestamp(),
+                        Comparator.nullsLast(Comparator.reverseOrder()))
                         .thenComparing(Task::getId))
                 .toList();
 

transport/rest/src/main/java/io/a2a/transport/rest/handler/RestHandler.java

@@ -11,7 +11,11 @@
 import jakarta.enterprise.context.ApplicationScoped;
 import jakarta.inject.Inject;
 
+import java.time.Instant;
+import java.time.format.DateTimeParseException;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.concurrent.Flow;
 
 import io.a2a.server.PublicAgentCard;
@@ -205,9 +209,9 @@ public HTTPRestResponse listTasks(@Nullable String contextId, @Nullable String s
             }
             if (lastUpdatedAfter != null) {
                 try {
-                    paramsBuilder.lastUpdatedAfter(java.time.Instant.parse(lastUpdatedAfter));
-                } catch (java.time.format.DateTimeParseException e) {
-                    java.util.Map<String, Object> errorData = new java.util.HashMap<>();
+                    paramsBuilder.lastUpdatedAfter(Instant.parse(lastUpdatedAfter));
+                } catch (DateTimeParseException e) {
+                    Map<String, Object> errorData = new HashMap<>();
                     errorData.put("parameter", "lastUpdatedAfter");
                     errorData.put("reason", "Must be valid ISO-8601 timestamp");
                     throw new InvalidParamsError(null, "Invalid params", errorData);