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."); + } } 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/events/EventConsumer.java b/server-common/src/main/java/org/a2aproject/sdk/server/events/EventConsumer.java index 531421aca..51713cb12 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 @@ -83,6 +94,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) { @@ -199,12 +220,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)); @@ -327,8 +364,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 4b63c6716..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 @@ -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 @@ -320,6 +333,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(); } @@ -404,6 +419,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); @@ -518,6 +534,7 @@ private Task doCancelTask(CancelTaskParams params, ServerCallContext context) th 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() @@ -609,6 +626,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()); @@ -667,9 +685,12 @@ 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 isInterruptedState = kind instanceof Task interruptedTask + && interruptedTask.status().state().isInterrupted(); + boolean isAsync = isAgentAsync(producerRunnable) && !isInterruptedState; try { - // Step 1: Wait for agent to finish (with configurable timeout) - if (agentFuture != null) { + // 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()); @@ -680,11 +701,16 @@ public EventKind onMessageSend(MessageSendParams params, ServerCallContext conte } } - // 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()); + // 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) && !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()); + } 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) { @@ -713,10 +739,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 @@ -811,6 +843,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()); @@ -1014,6 +1047,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( @@ -1054,6 +1088,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. * @@ -1075,6 +1118,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) { 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..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 @@ -115,6 +115,19 @@ public AgentEmitter(RequestContext context, EventQueue eventQueue) { this.contextId = context.getContextId(); } + /** + * 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 a terminal status has been reached + */ + public boolean isTerminalStateReached() { + return terminalStateReached.get(); + } + /** * Updates the task status to the given state with an optional message. * 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 0ab5b660d..bfcbc07cb 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 54882a078..6f37270dc 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; @@ -590,6 +592,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 48205fe7b..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 @@ -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; @@ -169,6 +170,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; @@ -177,6 +179,7 @@ void testInitConfigReadsBlockingTimeouts() { assertEquals(30, handler.agentCompletionTimeoutSeconds); assertEquals(5, handler.consumptionCompletionTimeoutSeconds); assertEquals(7, handler.reconciliationTimeoutSeconds); + assertEquals(90, handler.asyncAgentTimeoutSeconds); } /** @@ -1154,6 +1157,214 @@ public void onComplete() { "Protocol version should be stored when push config is provided via onMessageSendStream"); } + @Test + 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) -> { + 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"); + } + + @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"); + } + + @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");