From 3b2007eb2c7f17c3519af898fc23fb0ec58cd477 Mon Sep 17 00:00:00 2001 From: malladi nagarjuna Date: Sun, 9 Aug 2026 20:10:07 +0530 Subject: [PATCH] fix(server): support asynchronous agent execution via keepAlive This fixes #872 --- .../sdk/server/events/EnhancedRunnable.java | 10 +++ .../DefaultRequestHandler.java | 24 +++++- .../sdk/server/tasks/AgentEmitter.java | 20 +++++ .../DefaultRequestHandlerTest.java | 75 +++++++++++++++++++ 4 files changed, 125 insertions(+), 4 deletions(-) diff --git a/server-common/src/main/java/org/a2aproject/sdk/server/events/EnhancedRunnable.java b/server-common/src/main/java/org/a2aproject/sdk/server/events/EnhancedRunnable.java index 0650e1284..be7686b74 100644 --- a/server-common/src/main/java/org/a2aproject/sdk/server/events/EnhancedRunnable.java +++ b/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 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( diff --git a/server-common/src/main/java/org/a2aproject/sdk/server/requesthandlers/DefaultRequestHandler.java b/server-common/src/main/java/org/a2aproject/sdk/server/requesthandlers/DefaultRequestHandler.java index 522145bb7..b0fe17602 100644 --- a/server-common/src/main/java/org/a2aproject/sdk/server/requesthandlers/DefaultRequestHandler.java +++ b/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()); diff --git a/server-common/src/main/java/org/a2aproject/sdk/server/tasks/AgentEmitter.java b/server-common/src/main/java/org/a2aproject/sdk/server/tasks/AgentEmitter.java index 19d458fd6..b925e69d5 100644 --- a/server-common/src/main/java/org/a2aproject/sdk/server/tasks/AgentEmitter.java +++ b/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. * diff --git a/server-common/src/test/java/org/a2aproject/sdk/server/requesthandlers/DefaultRequestHandlerTest.java b/server-common/src/test/java/org/a2aproject/sdk/server/requesthandlers/DefaultRequestHandlerTest.java index 12b7ebeed..df23a560e 100644 --- a/server-common/src/test/java/org/a2aproject/sdk/server/requesthandlers/DefaultRequestHandlerTest.java +++ b/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 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"); + } }