fix: flaky testInputRequiredWorkflow test (#671)

Author: kabir

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);
+            }
         }
     }
 

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;
                     }