fix: clear status.message after adding to history in TaskManager.updateWithMessage #424 (#425)

Author: flex-myeonghyeon

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

@@ -104,12 +104,16 @@ public Event process(Event event) throws A2AServerException {
     }
 
     public Task updateWithMessage(Message message, Task task) {
-        List<Message> history = task.getHistory() == null ? new ArrayList<>() : new ArrayList<>(task.getHistory());
-        if (task.getStatus().message() != null) {
-            history.add(task.getStatus().message());
+        List<Message> history = new ArrayList<>(task.getHistory());
+
+        TaskStatus status = task.getStatus();
+        if (status.message() != null) {
+            history.add(status.message());
+            status = new TaskStatus(status.state(), null, status.timestamp());
         }
         history.add(message);
         task = new Task.Builder(task)
+                .status(status)
                 .history(history)
                 .build();
         saveTask(task);

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

@@ -692,4 +692,48 @@ public void testSaveTaskInternal() throws A2AServerException {
         assertEquals("test-context", taskManagerWithoutId.getContextId());
         assertSame(savedTask, taskManagerWithoutId.getTask());
     }
+
+    @Test
+    public void testUpdateWithMessage() throws A2AServerException {
+        Message initialMessage = new Message.Builder()
+                .role(Message.Role.USER)
+                .parts(Collections.singletonList(new TextPart("initial message")))
+                .messageId("initial-msg-id")
+                .build();
+
+        TaskManager taskManagerWithInitialMessage = new TaskManager(null, null, taskStore, initialMessage);
+
+        Message taskMessage = new Message.Builder()
+                .role(Message.Role.AGENT)
+                .parts(Collections.singletonList(new TextPart("task message")))
+                .messageId("task-msg-id")
+                .build();
+
+        TaskStatusUpdateEvent event = new TaskStatusUpdateEvent.Builder()
+                .taskId("new-task-id")
+                .contextId("some-context")
+                .status(new TaskStatus(TaskState.SUBMITTED, taskMessage, null))
+                .isFinal(false)
+                .build();
+
+        Task saved = taskManagerWithInitialMessage.saveTaskEvent(event);
+
+        Message updateMessage = new Message.Builder()
+                .role(Message.Role.USER)
+                .parts(Collections.singletonList(new TextPart("update message")))
+                .messageId("update-msg-id")
+                .build();
+
+        Task updated = taskManagerWithInitialMessage.updateWithMessage(updateMessage, saved);
+
+        // There should now be a history containing the initialMessage, task message and update message
+        assertNotNull(updated.getHistory());
+        assertEquals(3, updated.getHistory().size());
+        assertEquals("initial message", ((TextPart) updated.getHistory().get(0).getParts().get(0)).getText());
+
+        // The message in the current state should be null
+        assertNull(updated.getStatus().message());
+        assertEquals("task message", ((TextPart) updated.getHistory().get(1).getParts().get(0)).getText());
+        assertEquals("update message", ((TextPart) updated.getHistory().get(2).getParts().get(0)).getText());
+    }
 }