From 01bb8c57fb4bde6dd5685f46f5604ebc923a5a3a Mon Sep 17 00:00:00 2001 From: Kabir Khan Date: Wed, 18 Feb 2026 13:33:33 +0000 Subject: [PATCH] fix: flaky testInputRequiredWorkflow test --- .../io/a2a/server/apps/common/AbstractA2AServerTest.java | 6 +++++- .../io/a2a/server/apps/common/AgentExecutorProducer.java | 7 +++++-- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/tests/server-common/src/test/java/io/a2a/server/apps/common/AbstractA2AServerTest.java b/tests/server-common/src/test/java/io/a2a/server/apps/common/AbstractA2AServerTest.java index ae6b8443a..7ecd64cac 100644 --- a/tests/server-common/src/test/java/io/a2a/server/apps/common/AbstractA2AServerTest.java +++ b/tests/server-common/src/test/java/io/a2a/server/apps/common/AbstractA2AServerTest.java @@ -1538,6 +1538,7 @@ private boolean awaitChildQueueCountStable(String taskId, int expectedCount, lon @Timeout(value = 1, unit = TimeUnit.MINUTES) public void testInputRequiredWorkflow() throws Exception { String inputRequiredTaskId = "input-required-test-" + java.util.UUID.randomUUID(); + boolean taskCreated = false; try { // 1. Send initial message - AgentExecutor will transition task to INPUT_REQUIRED Message initialMessage = Message.builder(MESSAGE) @@ -1572,6 +1573,7 @@ public void testInputRequiredWorkflow() throws Exception { assertTrue(initialLatch.await(10, TimeUnit.SECONDS)); assertFalse(initialUnexpectedEvent.get()); assertEquals(TaskState.TASK_STATE_INPUT_REQUIRED, initialState.get()); + taskCreated = true; // 2. Send input message - AgentExecutor will complete the task Message inputMessage = Message.builder(MESSAGE) @@ -1608,7 +1610,9 @@ public void testInputRequiredWorkflow() throws Exception { assertEquals(TaskState.TASK_STATE_COMPLETED, completedState.get()); } finally { - deleteTaskInTaskStore(inputRequiredTaskId); + if (taskCreated) { + deleteTaskInTaskStore(inputRequiredTaskId); + } } } diff --git a/tests/server-common/src/test/java/io/a2a/server/apps/common/AgentExecutorProducer.java b/tests/server-common/src/test/java/io/a2a/server/apps/common/AgentExecutorProducer.java index d62391bf8..b82041143 100644 --- a/tests/server-common/src/test/java/io/a2a/server/apps/common/AgentExecutorProducer.java +++ b/tests/server-common/src/test/java/io/a2a/server/apps/common/AgentExecutorProducer.java @@ -79,7 +79,9 @@ public void execute(RequestContext context, AgentEmitter agentEmitter) throws A2 if (taskId != null && taskId.startsWith("input-required-test")) { // First call: context.getTask() == null (new task) if (context.getTask() == null) { - agentEmitter.startWork(); + // Go directly to INPUT_REQUIRED without intermediate WORKING state + // This avoids race condition where blocking call interrupts on WORKING + // before INPUT_REQUIRED is persisted to TaskStore agentEmitter.requiresInput(agentEmitter.newAgentMessage( List.of(new TextPart("Please provide additional information")), context.getMessage().metadata())); @@ -91,7 +93,8 @@ public void execute(RequestContext context, AgentEmitter agentEmitter) throws A2 throw new InvalidParamsError("We didn't get the expected input"); } // Second call: context.getTask() != null (input provided) - agentEmitter.startWork(); + // Go directly to COMPLETED without intermediate WORKING state + // This avoids the same race condition as the first call agentEmitter.complete(); return; }