Cleanup, and remove unneded checks from TCK AgentExecutor.

Move AgentExecutor to use TaskUpdater, and add TaskUpdater.cancel() method.

Author: kabir

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

@@ -70,11 +70,11 @@ public void complete(Message message) {
         updateStatus(TaskState.COMPLETED, message, true);
     }
 
-    public void failed() {
-        failed(null);
+    public void fail() {
+        fail(null);
     }
 
-    public void failed(Message message) {
+    public void fail(Message message) {
         updateStatus(TaskState.FAILED, message, true);
     }
 
@@ -94,6 +94,13 @@ public void startWork(Message message) {
         updateStatus(TaskState.WORKING, message);
     }
 
+    public void cancel() {
+        cancel(null);
+    }
+
+    public void cancel(Message message) {
+        updateStatus(TaskState.CANCELED, message, true);
+    }
 
     public Message newAgentMessage(List<Part<?>> parts, Map<String, Object> metadata) {
         return new Message.Builder()

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

@@ -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";
@@ -128,16 +126,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(eventQueue, context.getTaskId(), context.getTaskId());
+
             // 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(eventQueue, context.getTaskId(), context.getTaskId());
+            updater.cancel();
             eventQueue.enqueueEvent(new TaskStatusUpdateEvent.Builder()
                     .taskId(task.getId())
                     .contextId(task.getContextId())