Merge pull request #105 from kabir/tck

Fix Quality and Mandatory TCK tests

Author: fjuma

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

@@ -4,8 +4,10 @@
 import static io.a2a.util.AsyncUtils.createTubeConfig;
 import static io.a2a.util.AsyncUtils.processor;
 
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.concurrent.Executor;
@@ -30,9 +32,11 @@
 import io.a2a.server.tasks.ResultAggregator;
 import io.a2a.server.tasks.TaskManager;
 import io.a2a.server.tasks.TaskStore;
+import io.a2a.spec.Artifact;
 import io.a2a.spec.EventKind;
 import io.a2a.spec.InternalError;
 import io.a2a.spec.JSONRPCError;
+import io.a2a.spec.Message;
 import io.a2a.spec.MessageSendParams;
 import io.a2a.spec.PushNotificationConfig;
 import io.a2a.spec.StreamingEventKind;
@@ -85,6 +89,21 @@ public Task onGetTask(TaskQueryParams params) throws JSONRPCError {
             log.debug("No task found for {}. Throwing TaskNotFoundError", params.id());
             throw new TaskNotFoundError();
         }
+        if (params.historyLength() != null && task.getHistory() != null && params.historyLength() < task.getHistory().size()) {
+            List<Message> history;
+            if (params.historyLength() <= 0) {
+                history = new ArrayList<>();
+            } else {
+                history = task.getHistory().subList(
+                        task.getHistory().size() - params.historyLength(),
+                        task.getHistory().size() - 1);
+            }
+
+            task = new Task.Builder(task)
+                    .history(history)
+                    .build();
+        }
+
         log.debug("Task found {}", task);
         return task;
     }

src/main/java/io/a2a/server/tasks/TaskManager.java

@@ -44,15 +44,19 @@ public Task getTask() {
         if (taskId == null) {
             return null;
         }
-        return taskStore.get(taskId);
+        if (currentTask != null) {
+            return currentTask;
+        }
+        currentTask = taskStore.get(taskId);
+        return currentTask;
     }
 
-    public void saveTaskEvent(Task task) throws A2AServerException {
+    Task saveTaskEvent(Task task) throws A2AServerException {
         checkIdsAndUpdateIfNecessary(task.getId(), task.getContextId());
-        saveTask(task);
+        return saveTask(task);
     }
 
-    public void saveTaskEvent(TaskStatusUpdateEvent event) throws A2AServerException {
+    Task saveTaskEvent(TaskStatusUpdateEvent event) throws A2AServerException {
         checkIdsAndUpdateIfNecessary(event.getTaskId(), event.getContextId());
         Task task = ensureTask(event.getTaskId(), event.getContextId());
 
@@ -67,10 +71,10 @@ public void saveTaskEvent(TaskStatusUpdateEvent event) throws A2AServerException
         }
 
         task = builder.build();
-        saveTask(task);
+        return saveTask(task);
     }
 
-    public void saveTaskEvent(TaskArtifactUpdateEvent event) throws A2AServerException {
+    Task saveTaskEvent(TaskArtifactUpdateEvent event) throws A2AServerException {
         checkIdsAndUpdateIfNecessary(event.getTaskId(), event.getContextId());
         Task task = ensureTask(event.getTaskId(), event.getContextId());
 
@@ -122,7 +126,7 @@ public void saveTaskEvent(TaskArtifactUpdateEvent event) throws A2AServerExcepti
                 .artifacts(artifacts)
                 .build();
 
-        saveTask(task);
+        return saveTask(task);
     }
 
     public Event process(Event event) throws A2AServerException {
@@ -145,7 +149,7 @@ public Task updateWithMessage(Message message, Task task) {
         task = new Task.Builder(task)
                 .history(history)
                 .build();
-        currentTask = task;
+        saveTask(task);
         return task;
     }
 
@@ -164,7 +168,11 @@ private void checkIdsAndUpdateIfNecessary(String eventTaskId, String eventContex
     }
 
     private Task ensureTask(String eventTaskId, String eventContextId) {
-        Task task = taskStore.get(taskId);
+        Task task = currentTask;
+        if (task != null) {
+            return task;
+        }
+        task = taskStore.get(taskId);
         if (task == null) {
             task = createTask(eventTaskId, eventContextId);
             saveTask(task);
@@ -182,11 +190,13 @@ private Task createTask(String taskId, String contextId) {
                 .build();
     }
 
-    private void saveTask(Task task) {
+    private Task saveTask(Task task) {
         taskStore.save(task);
         if (taskId == null) {
             taskId = task.getId();
             contextId = task.getContextId();
         }
+        currentTask = task;
+        return currentTask;
     }
 }

src/test/java/io/a2a/server/tasks/TaskManagerTest.java

@@ -57,9 +57,10 @@ public void testGetTaskNonExistent() {
 
     @Test
     public void testSaveTaskEventNewTask() throws A2AServerException {
-        taskManager.saveTaskEvent(minimalTask);
+        Task saved = taskManager.saveTaskEvent(minimalTask);
         Task retrieved = taskManager.getTask();
         assertSame(minimalTask, retrieved);
+        assertSame(retrieved, saved);
     }
 
     @Test
@@ -83,10 +84,11 @@ public void testSaveTaskEventStatusUpdate() throws A2AServerException {
                 new HashMap<>());
 
 
-        taskManager.saveTaskEvent(event);
+        Task saved = taskManager.saveTaskEvent(event);
         Task updated = taskManager.getTask();
 
         assertNotSame(initialTask, updated);
+        assertSame(updated, saved);
 
         assertEquals(initialTask.getId(), updated.getId());
         assertEquals(initialTask.getContextId(), updated.getContextId());
@@ -108,9 +110,11 @@ public void testSaveTaskEventArtifactUpdate() throws A2AServerException {
                 .contextId(minimalTask.getContextId())
                 .artifact(newArtifact)
                 .build();
-        taskManager.saveTaskEvent(event);
+        Task saved = taskManager.saveTaskEvent(event);
 
         Task updatedTask = taskManager.getTask();
+        assertSame(updatedTask, saved);
+
         assertNotSame(initialTask, updatedTask);
         assertEquals(initialTask.getId(), updatedTask.getId());
         assertEquals(initialTask.getContextId(), updatedTask.getContextId());
@@ -136,14 +140,15 @@ public void testEnsureTaskNonExistentForStatusUpdate() throws A2AServerException
                 .isFinal(false)
                 .build();
 
-        taskManagerWithoutId.saveTaskEvent(event);
+        Task task = taskManagerWithoutId.saveTaskEvent(event);
         assertEquals(event.getTaskId(), taskManagerWithoutId.getTaskId());
         assertEquals(event.getContextId(), taskManagerWithoutId.getContextId());
 
         Task newTask = taskManagerWithoutId.getTask();
         assertEquals(event.getTaskId(), newTask.getId());
         assertEquals(event.getContextId(), newTask.getContextId());
         assertEquals(TaskState.SUBMITTED, newTask.getStatus().state());
+        assertSame(newTask, task);
     }
 
     @Test
@@ -155,12 +160,13 @@ public void testSaveTaskEventNewTaskNoTaskId() throws A2AServerException {
                 .status(new TaskStatus(TaskState.WORKING))
                 .build();
 
-        taskManagerWithoutId.saveTaskEvent(task);
+        Task saved = taskManagerWithoutId.saveTaskEvent(task);
         assertEquals(task.getId(), taskManagerWithoutId.getTaskId());
         assertEquals(task.getContextId(), taskManagerWithoutId.getContextId());
 
         Task retrieved = taskManagerWithoutId.getTask();
         assertSame(task, retrieved);
+        assertSame(retrieved, saved);
     }
 
     @Test