chore: sync changes to existing tasks tests from Python (#299)

Author: kabir

server-common/src/main/java/io/a2a/server/events/EventQueue.java

@@ -17,28 +17,45 @@ public abstract class EventQueue implements AutoCloseable {
 
     private static final Logger LOGGER = LoggerFactory.getLogger(EventQueue.class);
 
-    // TODO decide on a capacity
-    private static final int queueSize = 1000;
+    public static final int DEFAULT_QUEUE_SIZE = 1000;
 
+    private final int queueSize;
     private final BlockingQueue<Event> queue = new LinkedBlockingDeque<>();
-    private final Semaphore semaphore = new Semaphore(queueSize, true);
+    private final Semaphore semaphore;
     private volatile boolean closed = false;
 
 
 
     protected EventQueue() {
-        this(null);
+        this(DEFAULT_QUEUE_SIZE);
+    }
+
+    protected EventQueue(int queueSize) {
+        if (queueSize <= 0) {
+            throw new IllegalArgumentException("Queue size must be greater than 0");
+        }
+        this.queueSize = queueSize;
+        this.semaphore = new Semaphore(queueSize, true);
+        LOGGER.trace("Creating {} with queue size: {}", this, queueSize);
     }
 
     protected EventQueue(EventQueue parent) {
+        this(DEFAULT_QUEUE_SIZE);
         LOGGER.trace("Creating {}, parent: {}", this, parent);
     }
 
     public static EventQueue create() {
-        
         return new MainQueue();
     }
 
+    public static EventQueue create(int queueSize) {
+        return new MainQueue(queueSize);
+    }
+
+    public int getQueueSize() {
+        return queueSize;
+    }
+
     public abstract void awaitQueuePollerStart() throws InterruptedException ;
 
     abstract void signalQueuePollerStarted();
@@ -132,6 +149,14 @@ static class MainQueue extends EventQueue {
         private final CountDownLatch pollingStartedLatch = new CountDownLatch(1);
         private final AtomicBoolean pollingStarted = new AtomicBoolean(false);
 
+        MainQueue() {
+            super();
+        }
+
+        MainQueue(int queueSize) {
+            super(queueSize);
+        }
+
         EventQueue tap() {
             ChildQueue child = new ChildQueue(this);
             children.add(child);

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

@@ -76,6 +76,11 @@ Task saveTaskEvent(TaskStatusUpdateEvent event) throws A2AServerException {
             builder.history(newHistory);
         }
 
+        // Handle metadata from the event
+        if (event.getMetadata() != null) {
+            builder.metadata(event.getMetadata());
+        }
+
         task = builder.build();
         return saveTask(task);
     }

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

@@ -3,6 +3,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.UUID;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 import io.a2a.server.agentexecution.RequestContext;
 import io.a2a.server.events.EventQueue;
@@ -18,6 +19,8 @@ public class TaskUpdater {
     private final EventQueue eventQueue;
     private final String taskId;
     private final String contextId;
+    private final AtomicBoolean terminalStateReached = new AtomicBoolean(false);
+    private final Object stateLock = new Object();
 
     public TaskUpdater(RequestContext context, EventQueue eventQueue) {
         this.eventQueue = eventQueue;
@@ -26,20 +29,41 @@ public TaskUpdater(RequestContext context, EventQueue eventQueue) {
     }
 
     private void updateStatus(TaskState taskState) {
-        updateStatus(taskState, null);
+        updateStatus(taskState, null, taskState.isFinal());
     }
 
-    private void updateStatus(TaskState state, Message message) {
-        TaskStatusUpdateEvent event = new TaskStatusUpdateEvent.Builder()
-                .taskId(taskId)
-                .contextId(contextId)
-                .isFinal(state.isFinal())
-                .status(new TaskStatus(state, message, null))
-                .build();
-        eventQueue.enqueueEvent(event);
+    private void updateStatus(TaskState taskState, Message message) {
+        updateStatus(taskState, message, taskState.isFinal());
+    }
+
+    private void updateStatus(TaskState state, Message message, boolean isFinal) {
+        synchronized (stateLock) {
+            // Check if we're already in a terminal state
+            if (terminalStateReached.get()) {
+                throw new IllegalStateException("Cannot update task status - terminal state already reached");
+            }
+            
+            // If this is a final state, set the flag
+            if (isFinal) {
+                terminalStateReached.set(true);
+            }
+            
+            TaskStatusUpdateEvent event = new TaskStatusUpdateEvent.Builder()
+                    .taskId(taskId)
+                    .contextId(contextId)
+                    .isFinal(isFinal)
+                    .status(new TaskStatus(state, message, null))
+                    .build();
+            eventQueue.enqueueEvent(event);
+        }
     }
 
     public void addArtifact(List<Part<?>> parts, String artifactId, String name, Map<String, Object> metadata) {
+        addArtifact(parts, artifactId, name, metadata, null, null);
+    }
+
+    public void addArtifact(List<Part<?>> parts, String artifactId, String name, Map<String, Object> metadata, 
+                            Boolean append, Boolean lastChunk) {
         if (artifactId == null) {
             artifactId = UUID.randomUUID().toString();
         }
@@ -54,6 +78,8 @@ public void addArtifact(List<Part<?>> parts, String artifactId, String name, Map
                                 .metadata(metadata)
                                 .build()
                 )
+                .append(append)
+                .lastChunk(lastChunk)
                 .build();
         eventQueue.enqueueEvent(event);
     }
@@ -98,6 +124,46 @@ public void cancel(Message message) {
         updateStatus(TaskState.CANCELED, message);
     }
 
+    public void reject() {
+        reject(null);
+    }
+
+    public void reject(Message message) {
+        updateStatus(TaskState.REJECTED, message);
+    }
+
+    public void requiresInput() {
+        requiresInput(null, false);
+    }
+
+    public void requiresInput(Message message) {
+        requiresInput(message, false);
+    }
+
+    public void requiresInput(boolean isFinal) {
+        requiresInput(null, isFinal);
+    }
+
+    public void requiresInput(Message message, boolean isFinal) {
+        updateStatus(TaskState.INPUT_REQUIRED, message, isFinal);
+    }
+
+    public void requiresAuth() {
+        requiresAuth(null, false);
+    }
+
+    public void requiresAuth(Message message) {
+        requiresAuth(message, false);
+    }
+
+    public void requiresAuth(boolean isFinal) {
+        requiresAuth(null, isFinal);
+    }
+
+    public void requiresAuth(Message message, boolean isFinal) {
+        updateStatus(TaskState.AUTH_REQUIRED, message, isFinal);
+    }
+
     public Message newAgentMessage(List<Part<?>> parts, Map<String, Object> metadata) {
         return new Message.Builder()
                 .role(Message.Role.AGENT)