Merge pull request #109 from kabir/cleanup

Cleanup the AgentExecutors and use TaskUpdater where possible

Author: fjuma

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

@@ -4,6 +4,7 @@
 import java.util.Map;
 import java.util.UUID;
 
+import io.a2a.server.agentexecution.RequestContext;
 import io.a2a.server.events.EventQueue;
 import io.a2a.spec.Artifact;
 import io.a2a.spec.Message;
@@ -18,26 +19,21 @@ public class TaskUpdater {
     private final String taskId;
     private final String contextId;
 
-    public TaskUpdater(EventQueue eventQueue, String taskId, String contextId) {
+    public TaskUpdater(RequestContext context, EventQueue eventQueue) {
         this.eventQueue = eventQueue;
-        this.taskId = taskId;
-        this.contextId = contextId;
+        this.taskId = context.getTaskId();
+        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);
@@ -67,15 +63,15 @@ public void complete() {
     }
 
     public void complete(Message message) {
-        updateStatus(TaskState.COMPLETED, message, true);
+        updateStatus(TaskState.COMPLETED, message);
     }
 
-    public void failed() {
-        failed(null);
+    public void fail() {
+        fail(null);
     }
 
-    public void failed(Message message) {
-        updateStatus(TaskState.FAILED, message, true);
+    public void fail(Message message) {
+        updateStatus(TaskState.FAILED, message);
     }
 
     public void submit() {
@@ -94,6 +90,13 @@ public void startWork(Message message) {
         updateStatus(TaskState.WORKING, message);
     }
 
+    public void cancel() {
+        cancel(null);
+    }
+
+    public void cancel(Message message) {
+        updateStatus(TaskState.CANCELED, message);
+    }
 
     public Message newAgentMessage(List<Part<?>> parts, Map<String, Object> metadata) {
         return new Message.Builder()

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/apps/AgentExecutorProducer.java

@@ -6,6 +6,7 @@
 import io.a2a.server.agentexecution.AgentExecutor;
 import io.a2a.server.agentexecution.RequestContext;
 import io.a2a.server.events.EventQueue;
+import io.a2a.server.tasks.TaskUpdater;
 import io.a2a.spec.JSONRPCError;
 import io.a2a.spec.Task;
 import io.a2a.spec.TaskState;
@@ -29,12 +30,8 @@ public void execute(RequestContext context, EventQueue eventQueue) throws JSONRP
             @Override
             public void cancel(RequestContext context, EventQueue eventQueue) throws JSONRPCError {
                 if (context.getTask().getId().equals("cancel-task-123")) {
-                    Task task = context.getTask();
-                    Task updated = new Task.Builder(task)
-                            .status(new TaskStatus(TaskState.CANCELED))
-                            .build();
-
-                    eventQueue.enqueueEvent(updated);
+                    TaskUpdater taskUpdater = new TaskUpdater(context, eventQueue);
+                    taskUpdater.cancel();
                 } else if (context.getTask().getId().equals("cancel-task-not-supported-123")) {
                     throw new UnsupportedOperationError();
                 }

src/test/java/io/a2a/server/requesthandlers/JSONRPCHandlerTest.java

@@ -35,6 +35,7 @@
 import io.a2a.server.tasks.PushNotifier;
 import io.a2a.server.tasks.ResultAggregator;
 import io.a2a.server.tasks.TaskStore;
+import io.a2a.server.tasks.TaskUpdater;
 import io.a2a.spec.AgentCapabilities;
 import io.a2a.spec.AgentCard;
 import io.a2a.spec.Artifact;
@@ -166,11 +167,8 @@ public void testOnCancelTaskSuccess() throws Exception {
             // Looking at the Python implementation, they typically use AgentExecutors that
             // don't support cancellation. So my theory is the Agent updates the task to the CANCEL status
             Task task = context.getTask();
-            Task updated = new Task.Builder(task)
-                    .status(new TaskStatus(TaskState.CANCELED))
-                    .build();
-
-            eventQueue.enqueueEvent(updated);
+            TaskUpdater taskUpdater = new TaskUpdater(context, eventQueue);
+            taskUpdater.cancel();
         };
 
         CancelTaskRequest request = new CancelTaskRequest("111", new TaskIdParams(MINIMAL_TASK.getId()));

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

@@ -6,11 +6,11 @@
 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;
 
+import io.a2a.server.agentexecution.RequestContext;
 import io.a2a.server.events.Event;
 import io.a2a.server.events.EventQueue;
 import io.a2a.spec.Message;
@@ -20,9 +20,7 @@
 import io.a2a.spec.TaskStatusUpdateEvent;
 import io.a2a.spec.TextPart;
 import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Disabled;
 import org.junit.jupiter.api.Test;
-import org.mockito.Mockito;
 
 public class TaskUpdaterTest {
     public static final String TEST_TASK_ID = "test-task-id";
@@ -45,31 +43,11 @@ public class TaskUpdaterTest {
     @BeforeEach
     public void init() {
         eventQueue = EventQueue.create();
-        taskUpdater = new TaskUpdater(eventQueue, TEST_TASK_ID, TEST_TASK_CONTEXT_ID);
-    }
-
-    //@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);
+        RequestContext context = new RequestContext.Builder()
+                .setTaskId(TEST_TASK_ID)
+                .setContextId(TEST_TASK_CONTEXT_ID)
+                .build();
+        taskUpdater = new TaskUpdater(context, eventQueue);
     }
 
     @Test
@@ -128,16 +106,28 @@ public void testStartWorkWithMessage() throws Exception {
 
     @Test
     public void testFailedWithoutMessage() throws Exception {
-        taskUpdater.failed();
+        taskUpdater.fail();
         checkTaskStatusUpdateEventOnQueue(true, TaskState.FAILED, null);
     }
 
     @Test
     public void testFailedWithMessage() throws Exception {
-        taskUpdater.failed(SAMPLE_MESSAGE);
+        taskUpdater.fail(SAMPLE_MESSAGE);
         checkTaskStatusUpdateEventOnQueue(true, TaskState.FAILED, SAMPLE_MESSAGE);
     }
 
+    @Test
+    public void testCanceledWithoutMessage() throws Exception {
+        taskUpdater.cancel();
+        checkTaskStatusUpdateEventOnQueue(true, TaskState.CANCELED, null);
+    }
+
+    @Test
+    public void testCanceledWithMessage() throws Exception {
+        taskUpdater.cancel(SAMPLE_MESSAGE);
+        checkTaskStatusUpdateEventOnQueue(true, TaskState.CANCELED, SAMPLE_MESSAGE);
+    }
+
     @Test
     public void testNewAgentMessage() throws Exception {
         Message message = taskUpdater.newAgentMessage(SAMPLE_PARTS, null);

tck/src/main/java/io/a2a/tck/server/AgentExecutorProducer.java

@@ -7,6 +7,7 @@
 import io.a2a.server.agentexecution.AgentExecutor;
 import io.a2a.server.agentexecution.RequestContext;
 import io.a2a.server.events.EventQueue;
+import io.a2a.server.tasks.TaskUpdater;
 import io.a2a.spec.JSONRPCError;
 import io.a2a.spec.Task;
 import io.a2a.spec.TaskNotCancelableError;
@@ -42,22 +43,12 @@ public void execute(RequestContext context, EventQueue eventQueue) throws JSONRP
                 eventQueue.enqueueEvent(task);
             }
 
+            TaskUpdater updater = new TaskUpdater(context, eventQueue);
+
             // Immediately set to WORKING state
-            eventQueue.enqueueEvent(new TaskStatusUpdateEvent.Builder()
-                    .taskId(context.getTaskId())
-                    .contextId(context.getContextId())
-                    .status(new TaskStatus(TaskState.WORKING))
-                    .build());
-            
+            updater.startWork();
             System.out.println("====> task set to WORKING, starting background execution");
 
-//            // Fire and forget - start the task but don't wait for it
-//            CompletableFuture<Void> taskFuture = CompletableFuture
-//                .runAsync(() -> executeTaskInBackground(context, eventQueue), taskExecutor);
-            
-//            // Store the future for potential cancellation
-//            runningTasks.put(context.getTaskId(), taskFuture);
-            
             // Method returns immediately - task continues in background
             System.out.println("====> execute() method returning immediately, task running in background");
         }
@@ -66,12 +57,7 @@ public void execute(RequestContext context, EventQueue eventQueue) throws JSONRP
         public void cancel(RequestContext context, EventQueue eventQueue) throws JSONRPCError {
             System.out.println("====> task cancel request received");
             Task task = context.getTask();
-            
-            if (task == null) {
-                System.out.println("====> task not found");
-                throw new TaskNotFoundError();
-            }
-            
+
             if (task.getStatus().state() == TaskState.CANCELED) {
                 System.out.println("====> task already canceled");
                 throw new TaskNotCancelableError();
@@ -82,6 +68,8 @@ public void cancel(RequestContext context, EventQueue eventQueue) throws JSONRPC
                 throw new TaskNotCancelableError();
             }
 
+            TaskUpdater updater = new TaskUpdater(context, eventQueue);
+            updater.cancel();
             eventQueue.enqueueEvent(new TaskStatusUpdateEvent.Builder()
                     .taskId(task.getId())
                     .contextId(task.getContextId())
@@ -92,69 +80,6 @@ public void cancel(RequestContext context, EventQueue eventQueue) throws JSONRPC
             System.out.println("====> task canceled");
         }
 
-        /**
-         * This method runs completely in the background.
-         * The main execute() method has already returned.
-         */
-        private void executeTaskInBackground(RequestContext context, EventQueue eventQueue) {
-            String taskId = context.getTaskId();
-            
-            try {
-                System.out.println("====> background execution started for task: " + taskId);
-                
-                // Perform the actual work
-                Object result = performActualWork(context);
-                
-                // Task completed successfully
-                eventQueue.enqueueEvent(new TaskStatusUpdateEvent.Builder()
-                        .taskId(taskId)
-                        .contextId(context.getContextId())
-                        .status(new TaskStatus(TaskState.COMPLETED))
-                        .isFinal(true)
-                        .build());
-                
-                System.out.println("====> background task completed successfully: " + taskId);
-                
-            } catch (InterruptedException e) {
-                // Task was interrupted (cancelled)
-                System.out.println("====> background task was interrupted: " + taskId);
-                Thread.currentThread().interrupt();
-                
-            } catch (Exception e) {
-                // Task failed
-                System.err.println("====> background task failed: " + taskId);
-                e.printStackTrace();
-                
-            } finally {
-                // Always clean up
-                System.out.println("====> background task cleanup completed: " + taskId);
-            }
-        }
-
-        /**
-         * This method represents the actual work that needs to be done.
-         * Replace this with your real business logic.
-         */
-        private Object performActualWork(RequestContext context) throws InterruptedException {
-            
-            
-            System.out.println("====> starting actual work for task: " + context.getTaskId());
-            
-            // Simulate work that can be interrupted
-            for (int i = 0; i < 10; i++) {
-                // Check for interruption regularly during long-running work
-                if (Thread.currentThread().isInterrupted()) {
-                    throw new InterruptedException("Task was cancelled during execution");
-                }
-                
-                Thread.sleep(200); // Simulate work chunks
-                System.out.println("====> work progress for task " + context.getTaskId() + ": " + ((i + 1) * 10) + "%");
-            }
-            
-            System.out.println("====> finished actual work for task: " + context.getTaskId());
-            return "Task completed successfully";
-        }
-
         /**
          * Cleanup method for proper resource management
          */