Hide TaskUpdater.updateStatus() methods, and add isFinal field to TaskState"

Author: kabir

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

@@ -25,20 +25,15 @@ public TaskUpdater(EventQueue eventQueue, RequestContext context) {
         this.contextId = context.getContextId();
     }
 
-    public void updateStatus(TaskState taskState) {
+    private void updateStatus(TaskState taskState) {
         updateStatus(taskState, null);
     }
 
-    public void updateStatus(TaskState state, Message message) {
-        updateStatus(state, message, false);
-    }
-
-    public void updateStatus(TaskState state, Message message, boolean isFinal) {
-
+    private void updateStatus(TaskState state, Message message) {
         TaskStatusUpdateEvent event = new TaskStatusUpdateEvent.Builder()
                 .taskId(taskId)
                 .contextId(contextId)
-                .isFinal(isFinal)
+                .isFinal(state.isFinal())
                 .status(new TaskStatus(state, message, null))
                 .build();
         eventQueue.enqueueEvent(event);
@@ -68,15 +63,15 @@ public void complete() {
     }
 
     public void complete(Message message) {
-        updateStatus(TaskState.COMPLETED, message, true);
+        updateStatus(TaskState.COMPLETED, message);
     }
 
     public void fail() {
         fail(null);
     }
 
     public void fail(Message message) {
-        updateStatus(TaskState.FAILED, message, true);
+        updateStatus(TaskState.FAILED, message);
     }
 
     public void submit() {
@@ -100,7 +95,7 @@ public void cancel() {
     }
 
     public void cancel(Message message) {
-        updateStatus(TaskState.CANCELED, message, true);
+        updateStatus(TaskState.CANCELED, message);
     }
 
     public Message newAgentMessage(List<Part<?>> parts, Map<String, Object> metadata) {

src/main/java/io/a2a/spec/TaskState.java

@@ -11,23 +11,33 @@ public enum TaskState {
     WORKING("working"),
     INPUT_REQUIRED("input-required"),
     AUTH_REQUIRED("auth-required"),
-    COMPLETED("completed"),
-    CANCELED("canceled"),
-    FAILED("failed"),
-    REJECTED("rejected"),
-    UNKNOWN("unknown");
+    COMPLETED("completed", true),
+    CANCELED("canceled", true),
+    FAILED("failed", true),
+    REJECTED("rejected", true),
+    UNKNOWN("unknown", true);
 
     private final String state;
+    private final boolean isFinal;
 
     TaskState(String state) {
+        this(state, false);
+    }
+
+    TaskState(String state, boolean isFinal) {
         this.state = state;
+        this.isFinal = isFinal;
     }
 
     @JsonValue
     public String asString() {
         return state;
     }
 
+    public boolean isFinal(){
+        return isFinal;
+    }
+
     @JsonCreator
     public static TaskState fromString(String state) {
         switch (state) {

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

@@ -6,7 +6,6 @@
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertSame;
-import static org.junit.jupiter.api.Assertions.assertTrue;
 
 import java.util.List;
 import java.util.Map;
@@ -51,30 +50,6 @@ public void init() {
         taskUpdater = new TaskUpdater(eventQueue, context);
     }
 
-    //@Test
-    //public void testInit() {
-    //    // Python has a unit test testing that the constructor works. Not really relevant
-    //}
-
-    @Test
-    public void testUpdateStatusWithoutMessage() throws Exception {
-        taskUpdater.updateStatus(TaskState.WORKING);
-        checkTaskStatusUpdateEventOnQueue(false, TaskState.WORKING, null);
-
-    }
-
-    @Test
-    public void testUpdateStatusWithMessage() throws Exception {
-        taskUpdater.updateStatus(TaskState.WORKING, SAMPLE_MESSAGE);
-        checkTaskStatusUpdateEventOnQueue(false, TaskState.WORKING, SAMPLE_MESSAGE);
-    }
-
-    @Test
-    public void testUpdateStatusFinal() throws Exception {
-        taskUpdater.updateStatus(TaskState.COMPLETED, null, true);
-        checkTaskStatusUpdateEventOnQueue(true, TaskState.COMPLETED, null);
-    }
-
     @Test
     public void testAddArtifactWithCustomIdAndName() throws Exception {
         taskUpdater.addArtifact(SAMPLE_PARTS, "custom-artifact-id", "Custom Artifact", null);