Merge pull request #106 from kabir/tck

Remove concurrent executor from TCK AgentExecutor

Author: fjuma

src/main/java/io/a2a/server/events/EnhancedRunnable.java

@@ -21,7 +21,7 @@ public void addDoneCallback(DoneCallback doneCallback) {
         }
     }
 
-    protected void invokeDoneCallbacks() {
+    public void invokeDoneCallbacks() {
         synchronized (doneCallbacks) {
             for (DoneCallback doneCallback : doneCallbacks) {
                 doneCallback.done(this);

src/main/java/io/a2a/server/requesthandlers/DefaultRequestHandler.java

@@ -10,6 +10,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Executor;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Flow;
@@ -32,7 +33,6 @@
 import io.a2a.server.tasks.ResultAggregator;
 import io.a2a.server.tasks.TaskManager;
 import io.a2a.server.tasks.TaskStore;
-import io.a2a.spec.Artifact;
 import io.a2a.spec.EventKind;
 import io.a2a.spec.InternalError;
 import io.a2a.spec.JSONRPCError;
@@ -47,7 +47,6 @@
 import io.a2a.spec.TaskQueryParams;
 import io.a2a.spec.UnsupportedOperationError;
 import io.a2a.util.TempLoggerWrapper;
-
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -62,8 +61,7 @@ public class DefaultRequestHandler implements RequestHandler {
     private final PushNotifier pushNotifier;
     private final Supplier<RequestContext.Builder> requestContextBuilder;
 
-    // TODO the value upstream is asyncio.Task. Trying a Runnable
-    private final Map<String, EnhancedRunnable> runningAgents = Collections.synchronizedMap(new HashMap<>());
+    private final Map<String, CompletableFuture<Void>> runningAgents = Collections.synchronizedMap(new HashMap<>());
 
     private final Executor executor = Executors.newCachedThreadPool();
 
@@ -134,7 +132,10 @@ public Task onCancelTask(TaskIdParams params) throws JSONRPCError {
                         .build(),
                 queue);
 
-        // TODO need to cancel the asyncio.Task looked up from runningAgents
+        CompletableFuture<Void> cf = runningAgents.get(task.getId());
+        if (cf != null) {
+            cf.cancel(true);
+        }
 
         EventConsumer consumer = new EventConsumer(queue);
         EventKind type = resultAggregator.consumeAll(consumer);
@@ -203,9 +204,9 @@ public EventKind onMessageSend(MessageSendParams params) throws JSONRPCError {
         } finally {
             if (interrupted) {
                 // TODO Make this async
-                cleanupProducer(producerRunnable, taskId);
+                cleanupProducer(taskId);
             } else {
-                cleanupProducer(producerRunnable, taskId);
+                cleanupProducer(taskId);
             }
         }
 
@@ -244,7 +245,8 @@ public Flow.Publisher<StreamingEventKind> onMessageSendStream(MessageSendParams
         EnhancedRunnable producerRunnable = registerAndExecuteAgentAsync(taskId.get(), requestContext, queue);
 
         EventConsumer consumer = new EventConsumer(queue);
-        // TODO https://github.com/fjuma/a2a-java-sdk/issues/62 Add this callback
+
+        producerRunnable.addDoneCallback(consumer.createAgentRunnableDoneCallback());
 
         try {
             Flow.Publisher<Event> results = resultAggregator.consumeAndEmit(consumer);
@@ -286,7 +288,7 @@ public Flow.Publisher<StreamingEventKind> onMessageSendStream(MessageSendParams
 
             return convertingProcessor(eventPublisher, event -> (StreamingEventKind) event);
         } finally {
-            cleanupProducer(producerRunnable, taskId.get());
+            cleanupProducer(taskId.get());
         }
     }
 
@@ -347,42 +349,40 @@ private boolean shouldAddPushInfo(MessageSendParams params) {
         return pushNotifier != null && params.configuration() != null && params.configuration().pushNotification() != null;
     }
 
-    private void runEventStream(RequestContext requestContext, EventQueue queue)  throws JSONRPCError {
-        agentExecutor.execute(requestContext, queue);
-        // TODO this is in the Python implementation, but enabling it causes test hangs
-        try {
-            queueManager.signalPollingStarted(queue);
-        } catch (InterruptedException e) {
-            Thread.currentThread().interrupt();
-        } finally {
-            queue.close();
-        }
-    }
-
-    private EnhancedRunnable registerAndExecuteAgentAsync(String taskId, RequestContext requestContext, EventQueue eventQueue) {
+    private EnhancedRunnable registerAndExecuteAgentAsync(String taskId, RequestContext requestContext, EventQueue queue) {
         EnhancedRunnable runnable = new EnhancedRunnable() {
             @Override
             public void run() {
+                agentExecutor.execute(requestContext, queue);
                 try {
-                    runEventStream(requestContext, eventQueue);
-                } catch (Throwable throwable) {
-                    setError(throwable);
-                } finally {
-                    invokeDoneCallbacks();
+                    queueManager.signalPollingStarted(queue);
+                } catch (InterruptedException e) {
+                    Thread.currentThread().interrupt();
                 }
-
             }
         };
-        runningAgents.put(taskId, runnable);
-        executor.execute(runnable);
+
+        CompletableFuture<Void> cf = CompletableFuture.runAsync(runnable, executor)
+                .whenComplete((v, err) -> {
+                    if (err != null) {
+                        runnable.setError(err);
+                    }
+                    queue.close();
+                    runnable.invokeDoneCallbacks();
+                });
+        runningAgents.put(taskId, cf);
         return runnable;
     }
 
-    private void cleanupProducer(Runnable producerRunnable, String taskId) {
+    private void cleanupProducer(String taskId) {
         // TODO the Python implementation waits for the producerRunnable
-
-        queueManager.close(taskId);
-        runningAgents.remove(taskId);
+        CompletableFuture<Void> cf = runningAgents.get(taskId);
+        if (cf != null) {
+            cf.whenComplete((v, t) -> {
+                queueManager.close(taskId);
+                runningAgents.remove(taskId);
+            });
+        }
     }
 
 }

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

@@ -1,15 +1,8 @@
 package io.a2a.examples.helloworld.server;
 
-import java.util.Map;
-import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.ExecutorService;
-import java.util.concurrent.Executors;
-import java.util.concurrent.TimeUnit;
-
+import jakarta.annotation.PreDestroy;
 import jakarta.enterprise.context.ApplicationScoped;
 import jakarta.enterprise.inject.Produces;
-import jakarta.annotation.PreDestroy;
 
 import io.a2a.server.agentexecution.AgentExecutor;
 import io.a2a.server.agentexecution.RequestContext;
@@ -31,16 +24,6 @@ public AgentExecutor agentExecutor() {
     }
 
     private static class FireAndForgetAgentExecutor implements AgentExecutor {
-        // Dedicated thread pool for background task execution
-        private final ExecutorService taskExecutor = Executors.newCachedThreadPool(r -> {
-            Thread t = new Thread(r, "AgentTask-" + System.currentTimeMillis());
-            t.setDaemon(true); // Don't prevent JVM shutdown
-            return t;
-        });
-        
-        // Track running tasks for cancellation - store the future reference
-        private final Map<String, CompletableFuture<Void>> runningTasks = new ConcurrentHashMap<>();
-
         @Override
         public void execute(RequestContext context, EventQueue eventQueue) throws JSONRPCError {
             Task task = context.getTask();
@@ -68,12 +51,12 @@ public void execute(RequestContext context, EventQueue eventQueue) throws JSONRP
 
             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);
+//            // 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);
+//            // 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");
@@ -99,16 +82,6 @@ public void cancel(RequestContext context, EventQueue eventQueue) throws JSONRPC
                 throw new TaskNotCancelableError();
             }
 
-            // Cancel the CompletableFuture
-            CompletableFuture<Void> taskFuture = runningTasks.get(task.getId());
-            if (taskFuture != null) {
-                boolean cancelled = taskFuture.cancel(true); // mayInterruptIfRunning = true
-                System.out.println("====> cancellation attempted, success: " + cancelled);
-            }
-            
-            // Remove from running tasks and update status
-            runningTasks.remove(task.getId());
-            
             eventQueue.enqueueEvent(new TaskStatusUpdateEvent.Builder()
                     .taskId(task.getId())
                     .contextId(task.getContextId())
@@ -129,21 +102,9 @@ private void executeTaskInBackground(RequestContext context, EventQueue eventQue
             try {
                 System.out.println("====> background execution started for task: " + taskId);
 
-                // Check if task was cancelled before we even started
-                if (!runningTasks.containsKey(taskId)) {
-                    System.out.println("====> task was cancelled before background execution started");
-                    return;
-                }
-                
                 // Perform the actual work
                 Object result = performActualWork(context);
 
-                // Check again if task was cancelled during execution
-                if (!runningTasks.containsKey(taskId)) {
-                    System.out.println("====> task was cancelled during execution");
-                    return;
-                }
-                
                 // Task completed successfully
                 eventQueue.enqueueEvent(new TaskStatusUpdateEvent.Builder()
                         .taskId(taskId)
@@ -159,33 +120,13 @@ private void executeTaskInBackground(RequestContext context, EventQueue eventQue
                 System.out.println("====> background task was interrupted: " + taskId);
                 Thread.currentThread().interrupt();
 
-                // Only send CANCELED event if task is still tracked (not already cancelled)
-                if (runningTasks.containsKey(taskId)) {
-                    eventQueue.enqueueEvent(new TaskStatusUpdateEvent.Builder()
-                            .taskId(taskId)
-                            .contextId(context.getContextId())
-                            .status(new TaskStatus(TaskState.CANCELED))
-                            .isFinal(true)
-                            .build());
-                }
-                
             } catch (Exception e) {
                 // Task failed
                 System.err.println("====> background task failed: " + taskId);
                 e.printStackTrace();
 
-                if (runningTasks.containsKey(taskId)) {
-                    eventQueue.enqueueEvent(new TaskStatusUpdateEvent.Builder()
-                            .taskId(taskId)
-                            .contextId(context.getContextId())
-                            .status(new TaskStatus(TaskState.FAILED))
-                            .isFinal(true)
-                            .build());
-                }
-                
             } finally {
-                // Always clean up - remove from running tasks
-                runningTasks.remove(taskId);
+                // Always clean up
                 System.out.println("====> background task cleanup completed: " + taskId);
             }
         }
@@ -220,15 +161,6 @@ private Object performActualWork(RequestContext context) throws InterruptedExcep
         @PreDestroy
         public void cleanup() {
             System.out.println("====> shutting down task executor");
-            taskExecutor.shutdown();
-            try {
-                if (!taskExecutor.awaitTermination(5, TimeUnit.SECONDS)) {
-                    taskExecutor.shutdownNow();
-                }
-            } catch (InterruptedException e) {
-                taskExecutor.shutdownNow();
-                Thread.currentThread().interrupt();
-            }
-        }
+         }
     }
 }