Get rid of the background processing

Author: kabir

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

@@ -108,7 +108,6 @@ public class DefaultRequestHandler implements RequestHandler {
     private final Supplier<RequestContext.Builder> requestContextBuilder;
 
     private final ConcurrentMap<String, CompletableFuture<Void>> runningAgents = new ConcurrentHashMap<>();
-    private final Set<CompletableFuture<Void>> backgroundTasks = ConcurrentHashMap.newKeySet();
 
     private final Executor executor;
 
@@ -374,9 +373,9 @@ public EventKind onMessageSend(MessageSendParams params, ServerCallContext conte
             CompletableFuture<Void> agentFuture = runningAgents.remove(taskId);
             LOGGER.debug("Removed agent for task {} from runningAgents in finally block, size after: {}", taskId, runningAgents.size());
 
-            // Track cleanup as background task to avoid blocking Vert.x threads
+            // Cleanup as background task to avoid blocking Vert.x threads
             // Pass the consumption future to ensure cleanup waits for background consumption to complete
-            trackBackgroundTask(cleanupProducer(agentFuture, etai != null ? etai.consumptionFuture() : null, taskId, queue, false));
+            cleanupProducer(agentFuture, etai != null ? etai.consumptionFuture() : null, taskId, queue, false);
         }
 
         LOGGER.debug("Returning: {}", kind);
@@ -386,8 +385,8 @@ public EventKind onMessageSend(MessageSendParams params, ServerCallContext conte
     @Override
     public Flow.Publisher<StreamingEventKind> onMessageSendStream(
             MessageSendParams params, ServerCallContext context) throws JSONRPCError {
-        LOGGER.debug("onMessageSendStream START - task: {}; context: {}; runningAgents: {}; backgroundTasks: {}",
-                params.message().getTaskId(), params.message().getContextId(), runningAgents.size(), backgroundTasks.size());
+        LOGGER.debug("onMessageSendStream START - task: {}; context: {}; runningAgents: {}",
+                params.message().getTaskId(), params.message().getContextId(), runningAgents.size());
         MessageSendSetup mss = initMessageSend(params, context);
 
         AtomicReference<String> taskId = new AtomicReference<>(mss.requestContext.getTaskId());
@@ -398,12 +397,9 @@ public Flow.Publisher<StreamingEventKind> onMessageSendStream(
         EnhancedRunnable producerRunnable = registerAndExecuteAgentAsync(taskId.get(), mss.requestContext, queue);
 
         // Move consumer creation and callback registration outside try block
-        // so consumer is available for background consumption on client disconnect
         EventConsumer consumer = new EventConsumer(queue);
         producerRunnable.addDoneCallback(consumer.createAgentRunnableDoneCallback());
 
-        AtomicBoolean backgroundConsumeStarted = new AtomicBoolean(false);
-
         try {
             Flow.Publisher<EventQueueItem> results = resultAggregator.consumeAndEmit(consumer);
 
@@ -444,7 +440,8 @@ public Flow.Publisher<StreamingEventKind> onMessageSendStream(
 
             Flow.Publisher<StreamingEventKind> finalPublisher = convertingProcessor(eventPublisher, event -> (StreamingEventKind) event);
 
-            // Wrap publisher to detect client disconnect and continue background consumption
+            // Wrap publisher to detect client disconnect and immediately close ChildQueue
+            // This prevents ChildQueue backpressure from blocking MainEventBusProcessor
             return subscriber -> {
                 LOGGER.debug("Creating subscription wrapper for task {}", taskId.get());
                 finalPublisher.subscribe(new Flow.Subscriber<StreamingEventKind>() {
@@ -464,8 +461,10 @@ public void request(long n) {
 
                             @Override
                             public void cancel() {
-                                LOGGER.debug("Client cancelled subscription for task {}, starting background consumption", taskId.get());
-                                startBackgroundConsumption();
+                                LOGGER.debug("Client cancelled subscription for task {}, closing ChildQueue immediately", taskId.get());
+                                // Close ChildQueue immediately to prevent backpressure
+                                // (clears queue and releases semaphore permits)
+                                queue.close(true);  // immediate=true
                                 subscription.cancel();
                             }
                         });
@@ -490,45 +489,26 @@ public void onComplete() {
                             subscriber.onComplete();
                         } catch (IllegalStateException e) {
                             // Client already disconnected and response closed - this is expected
-                            // for streaming responses where client disconnect triggers background
-                            // consumption. Log and ignore.
+                            // for streaming responses where client disconnect closes ChildQueue.
+                            // Log and ignore.
                             if (e.getMessage() != null && e.getMessage().contains("Response has already been written")) {
                                 LOGGER.debug("Client disconnected before onComplete, response already closed for task {}", taskId.get());
                             } else {
                                 throw e;
                             }
                         }
                     }
-
-                    private void startBackgroundConsumption() {
-                        if (backgroundConsumeStarted.compareAndSet(false, true)) {
-                            LOGGER.debug("Starting background consumption for task {}", taskId.get());
-                            // Client disconnected: continue consuming and persisting events in background
-                            CompletableFuture<Void> bgTask = CompletableFuture.runAsync(() -> {
-                                try {
-                                    LOGGER.debug("Background consumption thread started for task {}", taskId.get());
-                                    resultAggregator.consumeAll(consumer);
-                                    LOGGER.debug("Background consumption completed for task {}", taskId.get());
-                                } catch (Exception e) {
-                                    LOGGER.error("Error during background consumption for task {}", taskId.get(), e);
-                                }
-                            }, executor);
-                            trackBackgroundTask(bgTask);
-                        } else {
-                            LOGGER.debug("Background consumption already started for task {}", taskId.get());
-                        }
-                    }
                 });
             };
         } finally {
-            LOGGER.debug("onMessageSendStream FINALLY - task: {}; runningAgents: {}; backgroundTasks: {}",
-                    taskId.get(), runningAgents.size(), backgroundTasks.size());
+            LOGGER.debug("onMessageSendStream FINALLY - task: {}; runningAgents: {}",
+                    taskId.get(), runningAgents.size());
 
             // Remove agent from map immediately to prevent accumulation
             CompletableFuture<Void> agentFuture = runningAgents.remove(taskId.get());
             LOGGER.debug("Removed agent for task {} from runningAgents in finally block, size after: {}", taskId.get(), runningAgents.size());
 
-            trackBackgroundTask(cleanupProducer(agentFuture, null, taskId.get(), queue, true));
+            cleanupProducer(agentFuture, null, taskId.get(), queue, true);
         }
     }
 
@@ -695,47 +675,6 @@ public void run() {
         return runnable;
     }
 
-    private void trackBackgroundTask(CompletableFuture<Void> task) {
-        backgroundTasks.add(task);
-        LOGGER.debug("Tracking background task (total: {}): {}", backgroundTasks.size(), task);
-
-        task.whenComplete((result, throwable) -> {
-            try {
-                if (throwable != null) {
-                    // Unwrap CompletionException to check for CancellationException
-                    Throwable cause = throwable;
-                    if (throwable instanceof java.util.concurrent.CompletionException && throwable.getCause() != null) {
-                        cause = throwable.getCause();
-                    }
-
-                    if (cause instanceof java.util.concurrent.CancellationException) {
-                        LOGGER.debug("Background task cancelled: {}", task);
-                    } else {
-                        LOGGER.error("Background task failed", throwable);
-                    }
-                }
-            } finally {
-                backgroundTasks.remove(task);
-                LOGGER.debug("Removed background task (remaining: {}): {}", backgroundTasks.size(), task);
-            }
-        });
-    }
-
-    /**
-     * Wait for all background tasks to complete.
-     * Useful for testing to ensure cleanup completes before assertions.
-     *
-     * @return CompletableFuture that completes when all background tasks finish
-     */
-    public CompletableFuture<Void> waitForBackgroundTasks() {
-        CompletableFuture<?>[] tasks = backgroundTasks.toArray(new CompletableFuture[0]);
-        if (tasks.length == 0) {
-            return CompletableFuture.completedFuture(null);
-        }
-        LOGGER.debug("Waiting for {} background tasks to complete", tasks.length);
-        return CompletableFuture.allOf(tasks);
-    }
-
     private CompletableFuture<Void> cleanupProducer(CompletableFuture<Void> agentFuture, CompletableFuture<Void> consumptionFuture, String taskId, EventQueue queue, boolean isStreaming) {
         LOGGER.debug("Starting cleanup for task {} (streaming={})", taskId, isStreaming);
         logThreadStats("CLEANUP START");
@@ -837,7 +776,6 @@ private void logThreadStats(String label) {
         LOGGER.debug("=== THREAD STATS: {} ===", label);
         LOGGER.debug("Active threads: {}", activeThreads);
         LOGGER.debug("Running agents: {}", runningAgents.size());
-        LOGGER.debug("Background tasks: {}", backgroundTasks.size());
         LOGGER.debug("Queue manager active queues: {}", queueManager.getClass().getSimpleName());
 
         // List running agents
@@ -848,13 +786,6 @@ private void logThreadStats(String label) {
             );
         }
 
-        // List background tasks
-        if (!backgroundTasks.isEmpty()) {
-            LOGGER.debug("Background tasks:");
-            backgroundTasks.forEach(task ->
-                LOGGER.debug("  - {}: {}", task, task.isDone() ? "DONE" : "RUNNING")
-            );
-        }
         LOGGER.debug("=== END THREAD STATS ===");
     }