From 3b2007eb2c7f17c3519af898fc23fb0ec58cd477 Mon Sep 17 00:00:00 2001 From: malladi nagarjuna Date: Sun, 9 Aug 2026 20:10:07 +0530 Subject: [PATCH 1/4] 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"); + } } From a630e996f22940b5b2c790a1afaff64c7dbf66a3 Mon Sep 17 00:00:00 2001 From: malladi nagarjuna Date: Mon, 10 Aug 2026 03:11:56 +0530 Subject: [PATCH 2/4] feat: make AgentExecutor a functional interface with default cancel method This fixes #1009 --- .../a2aproject/sdk/server/agentexecution/AgentExecutor.java | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/server-common/src/main/java/org/a2aproject/sdk/server/agentexecution/AgentExecutor.java b/server-common/src/main/java/org/a2aproject/sdk/server/agentexecution/AgentExecutor.java index aba33913a..a80c13879 100644 --- a/server-common/src/main/java/org/a2aproject/sdk/server/agentexecution/AgentExecutor.java +++ b/server-common/src/main/java/org/a2aproject/sdk/server/agentexecution/AgentExecutor.java @@ -97,6 +97,7 @@ * @see org.a2aproject.sdk.server.requesthandlers.DefaultRequestHandler * @see org.a2aproject.sdk.spec.AgentCard */ +@FunctionalInterface public interface AgentExecutor { /** * Executes the agent's business logic for a message. @@ -147,5 +148,7 @@ public interface AgentExecutor { * @throws org.a2aproject.sdk.spec.TaskNotCancelableError if this agent does not support cancellation * @throws A2AError if cancellation is supported but failed to execute */ - void cancel(RequestContext context, AgentEmitter emitter) throws A2AError; + default void cancel(RequestContext context, AgentEmitter emitter) throws A2AError { + throw new org.a2aproject.sdk.spec.TaskNotCancelableError("Agent execution is not cancelable."); + } } From 4930f649f87c7d11e89ee937f47241a16c4b7e1a Mon Sep 17 00:00:00 2001 From: malladi nagarjuna Date: Fri, 14 Aug 2026 12:31:59 +0530 Subject: [PATCH 3/4] fix(server): detect async agents via terminal state instead of a new keepAlive() API Signed-off-by: malladi nagarjuna --- .../sdk/server/events/EventConsumer.java | 47 ++++++- .../DefaultRequestHandler.java | 70 ++++++---- .../sdk/server/tasks/AgentEmitter.java | 23 ++-- .../META-INF/a2a-defaults.properties | 5 + .../sdk/server/events/EventConsumerTest.java | 53 ++++++++ .../DefaultRequestHandlerTest.java | 122 +++++++++++++++++- 6 files changed, 277 insertions(+), 43 deletions(-) diff --git a/server-common/src/main/java/org/a2aproject/sdk/server/events/EventConsumer.java b/server-common/src/main/java/org/a2aproject/sdk/server/events/EventConsumer.java index 18da23b22..6f40fb46b 100644 --- a/server-common/src/main/java/org/a2aproject/sdk/server/events/EventConsumer.java +++ b/server-common/src/main/java/org/a2aproject/sdk/server/events/EventConsumer.java @@ -6,10 +6,12 @@ import org.a2aproject.sdk.spec.A2AError; import org.a2aproject.sdk.spec.A2AServerException; import org.a2aproject.sdk.spec.Event; +import org.a2aproject.sdk.spec.InternalError; import org.a2aproject.sdk.spec.Message; import org.a2aproject.sdk.spec.Task; import org.a2aproject.sdk.spec.TaskState; import org.a2aproject.sdk.spec.TaskStatusUpdateEvent; +import org.a2aproject.sdk.server.tasks.AgentEmitter; import mutiny.zero.BackpressureStrategy; import mutiny.zero.TubeConfiguration; import mutiny.zero.ZeroPublisher; @@ -27,6 +29,9 @@ public class EventConsumer { private volatile int pollTimeoutsAfterAgentCompleted = 0; private volatile @Nullable TaskState lastSeenTaskState = null; private volatile int pollTimeoutsWhileAwaitingFinal = 0; + // True once an agent returns without reaching a terminal state (handed work off elsewhere). + private volatile boolean asyncAgentPending = false; + private volatile int pollTimeoutsWhileAsyncPending = 0; private static final String ERROR_MSG = "Agent did not return any response"; private static final int NO_WAIT = -1; @@ -41,6 +46,12 @@ public class EventConsumer { private static final int MAX_AWAITING_FINAL_TIMEOUT_MS = 3000; private static final int MAX_POLL_TIMEOUTS_AWAITING_FINAL = (MAX_AWAITING_FINAL_TIMEOUT_MS + QUEUE_WAIT_MILLISECONDS - 1) / QUEUE_WAIT_MILLISECONDS; + // Bounds otherwise-unbounded polling if an async agent's background work never completes. + private static final int MAX_ASYNC_PENDING_TIMEOUT_MS = 60_000; + private static final int MAX_POLL_TIMEOUTS_ASYNC_PENDING = + (MAX_ASYNC_PENDING_TIMEOUT_MS + QUEUE_WAIT_MILLISECONDS - 1) / QUEUE_WAIT_MILLISECONDS; + // Not final so tests can shrink this below the 60s default instead of waiting for it. + int maxPollTimeoutsAsyncPending = MAX_POLL_TIMEOUTS_ASYNC_PENDING; // Delay between tube.send(finalEvent) and tube.complete() to allow the SSE transport // layer to flush the write before the stream-end signal arrives. Mutiny's internal // demand management can call request(1) on the underlying publisher independently of @@ -55,6 +66,16 @@ public EventConsumer(EventQueue queue, Executor executor) { LOGGER.debug("EventConsumer created with queue {}", System.identityHashCode(queue)); } + /** + * Configures how long to wait for an async agent to reach a terminal state before giving up. + * + * @param seconds the timeout in seconds + */ + public void setAsyncAgentTimeoutSeconds(int seconds) { + this.maxPollTimeoutsAsyncPending = + (seconds * 1000 + QUEUE_WAIT_MILLISECONDS - 1) / QUEUE_WAIT_MILLISECONDS; + } + public Event consumeOne() throws A2AServerException, EventQueueClosedException { EventQueueItem item = queue.dequeueEventItem(NO_WAIT); if (item == null) { @@ -171,12 +192,28 @@ public Flow.Publisher consumeAll() { LOGGER.debug("Agent completed, awaiting final event (timeout {}/{}), continuing to poll (queue={})", pollTimeoutsWhileAwaitingFinal, MAX_POLL_TIMEOUTS_AWAITING_FINAL, System.identityHashCode(queue)); pollTimeoutsAfterAgentCompleted = 0; // Reset counter while awaiting final + } else if (asyncAgentPending && isInterruptedState) { + // Interrupted (non-terminal) state, not async hand-off - timeout doesn't apply. + pollTimeoutsWhileAsyncPending = 0; + } else if (asyncAgentPending && queueSize == 0 && !awaitingFinal) { + pollTimeoutsWhileAsyncPending++; + if (pollTimeoutsWhileAsyncPending >= maxPollTimeoutsAsyncPending) { + LOGGER.warn("Async agent did not reach a terminal state within {}ms, closing queue (queue={})", + MAX_ASYNC_PENDING_TIMEOUT_MS, System.identityHashCode(queue)); + queue.close(); + completed = true; + tube.fail(new InternalError("Agent execution did not complete within the allotted time")); + return; + } + } else if (asyncAgentPending && (queueSize > 0 || awaitingFinal)) { + pollTimeoutsWhileAsyncPending = 0; } continue; } // Event received - reset timeout counters pollTimeoutsAfterAgentCompleted = 0; pollTimeoutsWhileAwaitingFinal = 0; + pollTimeoutsWhileAsyncPending = 0; event = item.getEvent(); LOGGER.debug("EventConsumer received event: {} (queue={})", event.getClass().getSimpleName(), System.identityHashCode(queue)); @@ -294,8 +331,14 @@ public EnhancedRunnable.DoneCallback createAgentRunnableDoneCallback() { error = agentRunnable.getError(); LOGGER.debug("EventConsumer: Set error field from agent callback"); } else { - agentCompleted = true; - LOGGER.debug("EventConsumer: Agent completed successfully, set agentCompleted=true, will close queue after draining"); + AgentEmitter emitter = agentRunnable.getEmitter(); + if (emitter != null && !emitter.isTerminalStateReached()) { + asyncAgentPending = true; + LOGGER.debug("EventConsumer: Agent returned without reaching a terminal state, awaiting final event asynchronously"); + } else { + agentCompleted = true; + LOGGER.debug("EventConsumer: Agent completed successfully, set agentCompleted=true, will close queue after draining"); + } } }; } 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 b0fe17602..1de4d902b 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 @@ -196,6 +196,7 @@ public class DefaultRequestHandler implements RequestHandler { private static final String A2A_BLOCKING_AGENT_TIMEOUT_SECONDS = "a2a.blocking.agent.timeout.seconds"; private static final String A2A_BLOCKING_CONSUMPTION_TIMEOUT_SECONDS = "a2a.blocking.consumption.timeout.seconds"; private static final String A2A_BLOCKING_RECONCILIATION_TIMEOUT_SECONDS = "a2a.blocking.reconciliation.timeout.seconds"; + private static final String A2A_ASYNC_AGENT_TIMEOUT_SECONDS = "a2a.async-agent.timeout.seconds"; private static final String A2A_REQUEST_CONTEXT_POPULATE_REFERRED_TASKS = "a2a.request-context.populate-referred-tasks"; @Inject @@ -245,6 +246,18 @@ public class DefaultRequestHandler implements RequestHandler { */ int reconciliationTimeoutSeconds; + /** + * Timeout in seconds to wait for an async agent (one that returns from {@code execute()} + * without reaching a terminal {@link AgentEmitter} state) to eventually complete. Bounds + * otherwise-unbounded polling if the agent's background work crashes or hangs. + *

+ * Property: {@code a2a.async-agent.timeout.seconds}
+ * Default: 60 seconds
+ * Note: Property override requires a configurable {@link A2AConfigProvider} on the classpath + * (e.g., MicroProfileConfigProvider in reference implementations). + */ + int asyncAgentTimeoutSeconds; + // Fields set by constructor injection cannot be final. We need a noargs constructor for // Jakarta compatibility, and it seems that making fields set by constructor injection // final, is not proxyable in all runtimes @@ -308,6 +321,8 @@ void initConfig() { configProvider.getValue(A2A_BLOCKING_CONSUMPTION_TIMEOUT_SECONDS)); reconciliationTimeoutSeconds = Integer.parseInt( configProvider.getValue(A2A_BLOCKING_RECONCILIATION_TIMEOUT_SECONDS)); + asyncAgentTimeoutSeconds = Integer.parseInt( + configProvider.getValue(A2A_ASYNC_AGENT_TIMEOUT_SECONDS)); if (authorizationProviderInstance != null && authorizationProviderInstance.isResolvable()) { authorizationProvider = authorizationProviderInstance.get(); } @@ -392,6 +407,7 @@ public DefaultRequestHandler build() { handler.agentCompletionTimeoutSeconds = 5; handler.consumptionCompletionTimeoutSeconds = 2; handler.reconciliationTimeoutSeconds = 1; + handler.asyncAgentTimeoutSeconds = 60; handler.authorizationProvider = authorizationProvider; handler.requestContextBuilder = () -> new SimpleRequestContextBuilder(taskStore, populateReferredTasks, authorizationProvider); @@ -486,6 +502,7 @@ public Task onCancelTask(CancelTaskParams params, ServerCallContext context) thr EventQueue queue = queueManager.createOrTap(task.id()); EventConsumer consumer = new EventConsumer(queue, eventConsumerExecutor); + consumer.setAsyncAgentTimeoutSeconds(asyncAgentTimeoutSeconds); // Call agentExecutor.cancel() to enqueue the CANCELED event RequestContext cancelRequestContext = requestContextBuilder.get() @@ -577,6 +594,7 @@ public EventKind onMessageSend(MessageSendParams params, ServerCallContext conte // Create consumer BEFORE starting agent - callback is registered inside registerAndExecuteAgentAsync EventConsumer consumer = new EventConsumer(queue, eventConsumerExecutor); + consumer.setAsyncAgentTimeoutSeconds(asyncAgentTimeoutSeconds); EnhancedRunnable producerRunnable = registerAndExecuteAgentAsync(queueTaskId, mss.requestContext, queue, consumer.createAgentRunnableDoneCallback()); @@ -635,10 +653,10 @@ public EventKind onMessageSend(MessageSendParams params, ServerCallContext conte // 5. Fetch current task state from TaskStore (includes all consumed & persisted events) LOGGER.debug("DefaultRequestHandler: Entering blocking fire-and-forget handling for task {}", taskId.get()); + boolean isAsync = isAgentAsync(producerRunnable); try { - // Step 1: Wait for agent to finish (with configurable timeout) - // Note: We evaluate isAgentAsync dynamically because the agent sets it inside its run() - if (agentFuture != null && !(producerRunnable.getEmitter() != null && producerRunnable.getEmitter().isAsync())) { + // Step 1: Wait for the agent to finish, unless already handed off asynchronously. + if (agentFuture != null && !isAsync) { try { agentFuture.get(agentCompletionTimeoutSeconds, SECONDS); LOGGER.debug("DefaultRequestHandler: Step 1 - Agent completed for task {}", taskId.get()); @@ -647,17 +665,13 @@ 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 - // 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) { + // Step 2: Close the queue to signal consumption can complete (fire-and-forget + // tasks have no final event otherwise). Re-check isAsync since Step 1 may have + // changed it. + isAsync = isAgentAsync(producerRunnable); + if (!isAsync) { 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 { @@ -691,10 +705,16 @@ public EventKind onMessageSend(MessageSendParams params, ServerCallContext conte LOGGER.warn(msg, e.getCause()); throw new InternalError(msg); } catch (TimeoutException e) { - // Timeout from consumption future.get() - different from finalization timeout - String msg = String.format("Timeout waiting for task %s consumption", taskId.get()); - LOGGER.warn(msg, e); - throw new InternalError(msg); + // For an async agent this is expected (its work legitimately hasn't finished); + // EventConsumer's own fallback timeout still bounds it, so return the task's + // current state instead of failing the call. + if (!isAsync) { + String msg = String.format("Timeout waiting for task %s consumption", taskId.get()); + LOGGER.warn(msg, e); + throw new InternalError(msg); + } + LOGGER.debug("DefaultRequestHandler: Step 3 - Async agent for task {} still running after {}s, returning current state", + taskId.get(), consumptionCompletionTimeoutSeconds); } // Step 5: Fetch the current task state from TaskStore @@ -789,6 +809,7 @@ public Flow.Publisher onMessageSendStream( // Create consumer BEFORE starting agent - callback is registered inside registerAndExecuteAgentAsync EventConsumer consumer = new EventConsumer(queue, eventConsumerExecutor); + consumer.setAsyncAgentTimeoutSeconds(asyncAgentTimeoutSeconds); EnhancedRunnable producerRunnable = registerAndExecuteAgentAsync(queueTaskId, mss.requestContext, queue, consumer.createAgentRunnableDoneCallback()); @@ -992,6 +1013,7 @@ public Flow.Publisher onSubscribeToTask(TaskIdParams params, // Instead of enqueuing and hoping EventConsumer polls it in time, we prepend it // directly to the Publisher stream, ensuring synchronous delivery to subscriber EventConsumer consumer = new EventConsumer(queue, eventConsumerExecutor); + consumer.setAsyncAgentTimeoutSeconds(asyncAgentTimeoutSeconds); Flow.Publisher results = resultAggregator.consumeAndEmit(consumer); LOGGER.debug("onSubscribeToTask - prepending initial task snapshot to stream, taskId: {}", params.id()); return insertingProcessor( @@ -1032,6 +1054,15 @@ private boolean shouldAddPushInfo(MessageSendParams params) { return pushConfigStore != null && params.configuration() != null && params.configuration().taskPushNotificationConfig() != null; } + /** + * Returns whether the agent has handed execution off asynchronously: {@code execute()} + * returned without its {@link AgentEmitter} reaching a terminal state. + */ + private static boolean isAgentAsync(EnhancedRunnable producerRunnable) { + AgentEmitter emitter = producerRunnable.getEmitter(); + return emitter != null && !emitter.isTerminalStateReached(); + } + /** * Register and execute the agent asynchronously in the agent-executor thread pool. * @@ -1103,12 +1134,7 @@ public void run() { // Queue lifecycle is managed by EventConsumer.consumeAll() // which closes the queue on final events. logThreadStats("AGENT COMPLETE END"); - 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); - } + runnable.invokeDoneCallbacks(); }); 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 b925e69d5..017881749 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,7 +102,6 @@ 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. @@ -117,22 +116,16 @@ public AgentEmitter(RequestContext context, EventQueue eventQueue) { } /** - * 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. + * Returns whether a terminal status ({@link #complete}, {@link #fail}, {@link #cancel}, or + * {@link #reject}) has been reached. + *

+ * If {@link org.a2aproject.sdk.server.agentexecution.AgentExecutor#execute} returns without + * this being true, the agent handed work off to another thread. * - * @return true if keepAlive() has been called + * @return true if a terminal status has been reached */ - public boolean isAsync() { - return isAsync.get(); + public boolean isTerminalStateReached() { + return terminalStateReached.get(); } /** diff --git a/server-common/src/main/resources/META-INF/a2a-defaults.properties b/server-common/src/main/resources/META-INF/a2a-defaults.properties index e1a71fe24..781be3b1a 100644 --- a/server-common/src/main/resources/META-INF/a2a-defaults.properties +++ b/server-common/src/main/resources/META-INF/a2a-defaults.properties @@ -14,6 +14,11 @@ a2a.blocking.consumption.timeout.seconds=5 # When in-memory event capture is empty, polls TaskStore with this bounded timeout a2a.blocking.reconciliation.timeout.seconds=1 +# Timeout for an async agent (execute() returns without reaching a terminal AgentEmitter +# state, e.g. real work continues on another thread) to eventually complete (seconds) +# Bounds otherwise-unbounded polling if the agent's background work crashes or hangs +a2a.async-agent.timeout.seconds=60 + # AsyncExecutorProducer - Thread pool configuration # Core pool size for async agent execution a2a.executor.core-pool-size=5 diff --git a/server-common/src/test/java/org/a2aproject/sdk/server/events/EventConsumerTest.java b/server-common/src/test/java/org/a2aproject/sdk/server/events/EventConsumerTest.java index 41389c3ec..049292ac5 100644 --- a/server-common/src/test/java/org/a2aproject/sdk/server/events/EventConsumerTest.java +++ b/server-common/src/test/java/org/a2aproject/sdk/server/events/EventConsumerTest.java @@ -16,6 +16,8 @@ import java.util.concurrent.atomic.AtomicReference; import org.a2aproject.sdk.jsonrpc.common.json.JsonProcessingException; +import org.a2aproject.sdk.server.agentexecution.RequestContext; +import org.a2aproject.sdk.server.tasks.AgentEmitter; import org.a2aproject.sdk.server.tasks.InMemoryTaskStore; import org.a2aproject.sdk.server.tasks.PushNotificationSender; import org.a2aproject.sdk.spec.A2AError; @@ -546,6 +548,57 @@ public void onComplete() { assertEquals(0, receivedEvents.size(), "QueueClosedEvent should be intercepted, not delivered"); } + @Test + public void testConsumeAllFailsAfterAsyncPendingTimeout() throws Exception { + // An agent that returned without reaching a terminal state (e.g. its background + // work hung or crashed) must not be polled forever - it should eventually fail. + EventQueue queue = EventQueueUtil.getEventQueueBuilder(mainEventBus) + .taskId(TASK_ID) + .mainEventBus(mainEventBus) + .build().tap(); + EventConsumer consumer = new EventConsumer(queue, Runnable::run); + consumer.setAsyncAgentTimeoutSeconds(1); // ~1s instead of the 60s default + + RequestContext context = new RequestContext.Builder().setTaskId(TASK_ID).setContextId("ctx-1").build(); + AgentEmitter emitter = new AgentEmitter(context, queue); + EnhancedRunnable agentRunnable = new EnhancedRunnable() { + @Override + public void run() { + } + }; + agentRunnable.setEmitter(emitter); + consumer.createAgentRunnableDoneCallback().done(agentRunnable); + + Flow.Publisher publisher = consumer.consumeAll(); + final AtomicReference error = new AtomicReference<>(); + final CountDownLatch completionLatch = new CountDownLatch(1); + + publisher.subscribe(new Flow.Subscriber<>() { + @Override + public void onSubscribe(Flow.Subscription subscription) { + subscription.request(Long.MAX_VALUE); + } + + @Override + public void onNext(EventQueueItem item) { + } + + @Override + public void onError(Throwable throwable) { + error.set(throwable); + completionLatch.countDown(); + } + + @Override + public void onComplete() { + completionLatch.countDown(); + } + }); + + assertTrue(completionLatch.await(5, TimeUnit.SECONDS), "Test timed out waiting for the fallback timeout to fire."); + assertNotNull(error.get(), "Expected the stream to fail once the async agent's fallback timeout elapsed"); + } + private void enqueueAndConsumeOneEvent(Event event) throws Exception { // Use callback to wait for event processing waitForEventProcessing(() -> eventQueue.enqueueEvent(event)); 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 df23a560e..0c0cc9d89 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 @@ -9,6 +9,7 @@ import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; +import java.util.ArrayList; import java.util.List; import java.util.Map; import java.util.Set; @@ -162,6 +163,7 @@ void testInitConfigReadsBlockingTimeouts() { when(configProvider.getValue("a2a.blocking.agent.timeout.seconds")).thenReturn("30"); when(configProvider.getValue("a2a.blocking.consumption.timeout.seconds")).thenReturn("5"); when(configProvider.getValue("a2a.blocking.reconciliation.timeout.seconds")).thenReturn("7"); + when(configProvider.getValue("a2a.async-agent.timeout.seconds")).thenReturn("90"); DefaultRequestHandler handler = new DefaultRequestHandler(); handler.configProvider = configProvider; @@ -170,6 +172,7 @@ void testInitConfigReadsBlockingTimeouts() { assertEquals(30, handler.agentCompletionTimeoutSeconds); assertEquals(5, handler.consumptionCompletionTimeoutSeconds); assertEquals(7, handler.reconciliationTimeoutSeconds); + assertEquals(90, handler.asyncAgentTimeoutSeconds); } /** @@ -1148,14 +1151,13 @@ public void onComplete() { } @Test - void testAsyncAgentWithKeepAlive_Blocking_WaitsForCompletion() throws Exception { - // Arrange: Agent uses keepAlive and completes asynchronously + void testAsyncAgent_Blocking_WaitsForCompletion() throws Exception { + // Arrange: agent hands off to a background thread and returns without reaching a + // terminal state; no opt-in call needed for this to be treated as async. 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 @@ -1221,4 +1223,116 @@ void testAsyncAgentWithKeepAlive_Blocking_WaitsForCompletion() throws Exception // 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"); } + + @Test + void testAsyncAgent_Blocking_DirectCall_ReturnsCompletedState() throws Exception { + // Mimics real-world code like RxJava's Single.subscribeOn(Schedulers.io()): execute() + // returns almost immediately, completion happens shortly after on another thread - + // comfortably inside the test fixture's 2s consumptionCompletionTimeoutSeconds. + CountDownLatch agentBackgroundThreadStarted = new CountDownLatch(1); + + agentExecutorExecute = (context, emitter) -> { + emitter.startWork(); + internalExecutor.execute(() -> { + agentBackgroundThreadStarted.countDown(); + try { + Thread.sleep(300); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + emitter.complete(); + }); + // execute() returns immediately, well before the background thread finishes + }; + + Message initialMessage = Message.builder() + .messageId("msg-async-repro-1") + .role(Message.Role.ROLE_USER) + .parts(new TextPart("start async task")) + .build(); + + MessageSendParams initialParams = MessageSendParams.builder() + .message(initialMessage) + .configuration(MessageSendConfiguration.builder() + .returnImmediately(false) + .acceptedOutputModes(List.of()) + .build()) + .build(); + + EventKind result = requestHandler.onMessageSend(initialParams, NULL_CONTEXT); + + assertTrue(agentBackgroundThreadStarted.await(5, TimeUnit.SECONDS)); + assertInstanceOf(Task.class, result); + Task task = (Task) result; + assertEquals(TaskState.TASK_STATE_COMPLETED, task.status().state(), + "Task should reflect the background thread's completion, not close early"); + } + + @Test + void testAsyncAgent_Streaming_DeliversFinalEventAfterBackgroundCompletion() throws Exception { + // Arrange: same async hand-off pattern, but over the streaming (SSE) path - + // the stream must stay open until the background thread completes the task. + CountDownLatch agentBackgroundThreadStarted = new CountDownLatch(1); + + agentExecutorExecute = (context, emitter) -> { + emitter.startWork(); + internalExecutor.execute(() -> { + agentBackgroundThreadStarted.countDown(); + try { + Thread.sleep(300); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + emitter.complete(); + }); + }; + + Message initialMessage = Message.builder() + .messageId("msg-async-stream-1") + .role(Message.Role.ROLE_USER) + .parts(new TextPart("start async streaming task")) + .build(); + + MessageSendParams params = MessageSendParams.builder() + .message(initialMessage) + .configuration(MessageSendConfiguration.builder() + .returnImmediately(true) + .acceptedOutputModes(List.of()) + .build()) + .build(); + + Flow.Publisher publisher = requestHandler.onMessageSendStream(params, contextWithVersion("1.0")); + + List received = new ArrayList<>(); + CountDownLatch streamDone = new CountDownLatch(1); + publisher.subscribe(new Flow.Subscriber<>() { + @Override + public void onSubscribe(Flow.Subscription s) { + s.request(Long.MAX_VALUE); + } + + @Override + public void onNext(StreamingEventKind item) { + received.add(item); + } + + @Override + public void onError(Throwable t) { + streamDone.countDown(); + } + + @Override + public void onComplete() { + streamDone.countDown(); + } + }); + + assertTrue(agentBackgroundThreadStarted.await(5, TimeUnit.SECONDS)); + assertTrue(streamDone.await(10, TimeUnit.SECONDS), "Stream should complete once the background thread finishes"); + + boolean sawCompleted = received.stream().anyMatch(e -> + (e instanceof Task t && t.status().state() == TaskState.TASK_STATE_COMPLETED) + || (e instanceof TaskStatusUpdateEvent u && u.status().state() == TaskState.TASK_STATE_COMPLETED)); + assertTrue(sawCompleted, "Stream should deliver the COMPLETED event from the background thread, not close early"); + } } From ead9373ad99ade57dca6e93bfa90b617787ea29f Mon Sep 17 00:00:00 2001 From: malladi nagarjuna Date: Fri, 14 Aug 2026 20:04:47 +0530 Subject: [PATCH 4/4] exclude interrupted states from async-agent detection --- .../DefaultRequestHandler.java | 6 +++-- .../DefaultRequestHandlerTest.java | 22 +++++++++++++++++++ 2 files changed, 26 insertions(+), 2 deletions(-) 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 d74fdd8be..52db80996 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 @@ -685,7 +685,9 @@ public EventKind onMessageSend(MessageSendParams params, ServerCallContext conte // 5. Fetch current task state from TaskStore (includes all consumed & persisted events) LOGGER.debug("DefaultRequestHandler: Entering blocking fire-and-forget handling for task {}", taskId.get()); - boolean isAsync = isAgentAsync(producerRunnable); + boolean isInterruptedState = kind instanceof Task interruptedTask + && interruptedTask.status().state().isInterrupted(); + boolean isAsync = isAgentAsync(producerRunnable) && !isInterruptedState; try { // Step 1: Wait for the agent to finish, unless already handed off asynchronously. if (agentFuture != null && !isAsync) { @@ -702,7 +704,7 @@ public EventKind onMessageSend(MessageSendParams params, ServerCallContext conte // Step 2: Close the queue to signal consumption can complete (fire-and-forget // tasks have no final event otherwise). Re-check isAsync since Step 1 may have // changed it. - isAsync = isAgentAsync(producerRunnable); + isAsync = isAgentAsync(producerRunnable) && !isInterruptedState; if (!isAsync) { 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()); 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 62c127ef1..41e1e8283 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 @@ -1343,6 +1343,28 @@ public void onComplete() { assertTrue(sawCompleted, "Stream should deliver the COMPLETED event from the background thread, not close early"); } + @Test + void testInputRequired_Blocking_ReturnsQuickly() throws Exception { + agentExecutorExecute = (context, emitter) -> emitter.requiresInput(); + + MessageSendParams params = MessageSendParams.builder() + .message(MESSAGE) + .configuration(MessageSendConfiguration.builder() + .returnImmediately(false) + .acceptedOutputModes(List.of()) + .build()) + .build(); + + long start = System.nanoTime(); + EventKind result = requestHandler.onMessageSend(params, NULL_CONTEXT); + long elapsedMs = (System.nanoTime() - start) / 1_000_000; + + assertInstanceOf(Task.class, result); + assertEquals(TaskState.TASK_STATE_INPUT_REQUIRED, ((Task) result).status().state()); + assertTrue(elapsedMs < 1500, + "Blocking send should return quickly for INPUT_REQUIRED, not wait out the consumption timeout (took " + elapsedMs + "ms)"); + } + @Test void testOnGetTaskHistoryLengthLimitsHistory() throws Exception { Task task = taskWithHistory("task-hl-limit");