fix(server): support asynchronous agent execution via keepAlive

server-common/src/main/java/org/a2aproject/sdk/server/events/EnhancedRunnable.java

@@ -5,9 +5,11 @@
 import java.util.concurrent.atomic.AtomicBoolean;
 
 import org.jspecify.annotations.Nullable;
+import org.a2aproject.sdk.server.tasks.AgentEmitter;
 
 public abstract class EnhancedRunnable implements Runnable {
     private volatile @Nullable Throwable error;
+    private volatile @Nullable AgentEmitter emitter;
     private final List<DoneCallback> doneCallbacks = new CopyOnWriteArrayList<>();
     private final AtomicBoolean started = new AtomicBoolean(false);
 
@@ -19,6 +21,14 @@ public void setError(Throwable error) {
         this.error = error;
     }
 
+    public @Nullable AgentEmitter getEmitter() {
+        return emitter;
+    }
+
+    public void setEmitter(AgentEmitter emitter) {
+        this.emitter = emitter;
+    }
+
     public void addDoneCallback(DoneCallback doneCallback) {
         if (started.get()) {
             throw new IllegalStateException(

server-common/src/main/java/org/a2aproject/sdk/server/requesthandlers/DefaultRequestHandler.java

@@ -637,7 +637,8 @@ public EventKind onMessageSend(MessageSendParams params, ServerCallContext conte
 
                 try {
                     // Step 1: Wait for agent to finish (with configurable timeout)
-                    if (agentFuture != null) {
+                    // Note: We evaluate isAgentAsync dynamically because the agent sets it inside its run()
+                    if (agentFuture != null && !(producerRunnable.getEmitter() != null && producerRunnable.getEmitter().isAsync())) {
                         try {
                             agentFuture.get(agentCompletionTimeoutSeconds, SECONDS);
                             LOGGER.debug("DefaultRequestHandler: Step 1 - Agent completed for task {}", taskId.get());
@@ -646,13 +647,22 @@ public EventKind onMessageSend(MessageSendParams params, ServerCallContext conte
                             LOGGER.debug("DefaultRequestHandler: Step 1 - Agent still running for task {} after {}s timeout",
                                 taskId.get(), agentCompletionTimeoutSeconds);
                         }
+                    } else if (producerRunnable.getEmitter() != null && producerRunnable.getEmitter().isAsync()) {
+                        LOGGER.debug("DefaultRequestHandler: Step 1 - Agent is async, skipping agentFuture wait for task {}", taskId.get());
                     }
 
                     // Step 2: Close the queue to signal consumption can complete
                     // For fire-and-forget tasks, there's no final event, so we need to close the queue
                     // This allows EventConsumer.consumeAll() to exit
-                    queue.close(false, false);  // graceful close, don't notify parent yet
-                    LOGGER.debug("DefaultRequestHandler: Step 2 - Closed queue for task {} to allow consumption completion", taskId.get());
+                    // If the agent is async, it promises to emit a final event, so we don't close the queue here
+                    // Re-evaluate isAsync because the agent might have set it during Step 1
+                    boolean isFinallyAsync = producerRunnable.getEmitter() != null && producerRunnable.getEmitter().isAsync();
+                    if (!isFinallyAsync) {
+                        queue.close(false, false);  // graceful close, don't notify parent yet
+                        LOGGER.debug("DefaultRequestHandler: Step 2 - Closed queue for task {} to allow consumption completion", taskId.get());
+                    } else {
+                        LOGGER.debug("DefaultRequestHandler: Step 2 - Agent is async, keeping queue open to await final event for task {}", taskId.get());
+                    }
 
                     // Step 3: Wait for consumption to complete (now that queue is closed)
                     if (etai.consumptionFuture() != null) {
@@ -1043,6 +1053,7 @@ private EnhancedRunnable registerAndExecuteAgentAsync(String taskId, RequestCont
             public void run() {
                 LOGGER.debug("Agent execution starting for task {}", taskId);
                 AgentEmitter emitter = new AgentEmitter(requestContext, queue);
+                setEmitter(emitter);
                 try {
                     agentExecutor.execute(requestContext, emitter);
                 } catch (A2AError e) {
@@ -1092,7 +1103,12 @@ public void run() {
                     // Queue lifecycle is managed by EventConsumer.consumeAll()
                     // which closes the queue on final events.
                     logThreadStats("AGENT COMPLETE END");
-                    runnable.invokeDoneCallbacks();
+                    AgentEmitter emitter = runnable.getEmitter();
+                    if (emitter == null || !emitter.isAsync()) {
+                        runnable.invokeDoneCallbacks();
+                    } else {
+                        LOGGER.debug("Agent is marked as async, keeping queue open for task {}", taskId);
+                    }
                 });
         runningAgents.put(taskId, cf);
         LOGGER.debug("Registered agent for task {}, runningAgents.size() after: {}", taskId, runningAgents.size());

server-common/src/main/java/org/a2aproject/sdk/server/tasks/AgentEmitter.java

@@ -102,6 +102,7 @@ public class AgentEmitter {
     private final String taskId;
     private final String contextId;
     private final AtomicBoolean terminalStateReached = new AtomicBoolean(false);
+    private final AtomicBoolean isAsync = new AtomicBoolean(false);
 
     /**
      * Creates a new AgentEmitter for the given request context and event queue.
@@ -115,6 +116,25 @@ public AgentEmitter(RequestContext context, EventQueue eventQueue) {
         this.contextId = context.getContextId();
     }
 
+    /**
+     * Marks this agent execution as asynchronous, preventing premature queue closure
+     * before a terminal event is explicitly emitted.
+     *
+     * @since 1.0.0
+     */
+    public void keepAlive() {
+        this.isAsync.set(true);
+    }
+
+    /**
+     * Returns whether this emitter has been marked for asynchronous execution.
+     *
+     * @return true if keepAlive() has been called
+     */
+    public boolean isAsync() {
+        return isAsync.get();
+    }
+
     /**
      * Updates the task status to the given state with an optional message.
      *

server-common/src/test/java/org/a2aproject/sdk/server/requesthandlers/DefaultRequestHandlerTest.java

@@ -1146,4 +1146,79 @@ public void onComplete() {
         assertEquals("1.0", pushConfigStore.getProtocolVersion(taskId, taskId),
             "Protocol version should be stored when push config is provided via onMessageSendStream");
     }
+
+    @Test
+    void testAsyncAgentWithKeepAlive_Blocking_WaitsForCompletion() throws Exception {
+        // Arrange: Agent uses keepAlive and completes asynchronously
+        CountDownLatch agentBackgroundThreadStarted = new CountDownLatch(1);
+        CountDownLatch agentRelease = new CountDownLatch(1);
+
+        agentExecutorExecute = (context, emitter) -> {
+            // Signal that we are going to run asynchronously
+            emitter.keepAlive();
+            emitter.startWork();
+
+            // Simulate RxJava / async background thread
+            internalExecutor.execute(() -> {
+                agentBackgroundThreadStarted.countDown();
+                try {
+                    agentRelease.await(10, TimeUnit.SECONDS);
+                } catch (InterruptedException e) {
+                    Thread.currentThread().interrupt();
+                }
+                emitter.complete();
+            });
+            // execute() returns immediately!
+        };
+
+        Message initialMessage = Message.builder()
+            .messageId("msg-async-1")
+            .role(Message.Role.ROLE_USER)
+            .parts(new TextPart("start async task"))
+            .build();
+
+        // Blocking call (returnImmediately = false)
+        MessageSendParams initialParams = MessageSendParams.builder()
+            .message(initialMessage)
+            .configuration(MessageSendConfiguration.builder()
+                .returnImmediately(false)
+                .acceptedOutputModes(List.of())
+                .build())
+            .build();
+
+        // Use a background thread to call onMessageSend since it should block
+        AtomicReference<EventKind> resultRef = new AtomicReference<>();
+        CountDownLatch callComplete = new CountDownLatch(1);
+
+        internalExecutor.execute(() -> {
+            try {
+                EventKind result = requestHandler.onMessageSend(initialParams, NULL_CONTEXT);
+                resultRef.set(result);
+            } catch (Exception e) {
+                e.printStackTrace();
+            } finally {
+                callComplete.countDown();
+            }
+        });
+
+        // Wait for the background thread to start
+        assertTrue(agentBackgroundThreadStarted.await(5, TimeUnit.SECONDS));
+
+        // The requestHandler should be blocked, so callComplete should NOT have counted down
+        assertFalse(callComplete.await(1, TimeUnit.SECONDS), "Client call should block while async agent runs");
+
+        // Release the agent so it can emit completion
+        agentRelease.countDown();
+
+        // Now the client call should complete
+        assertTrue(callComplete.await(5, TimeUnit.SECONDS), "Client call should complete after agent finishes");
+
+        EventKind result = resultRef.get();
+        assertNotNull(result);
+        assertInstanceOf(Task.class, result);
+        Task task = (Task) result;
+
+        // Since it's a blocking non-streaming call, the final state should be returned
+        assertEquals(TaskState.TASK_STATE_COMPLETED, task.status().state(), "Task should be in COMPLETED state");
+    }
 }