Tighten up TaskManager

Author: kabir

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