From a421f3ccac683c22dd96deaa2ee07c269fc80a4c Mon Sep 17 00:00:00 2001
From: Kabir Khan <kkhan@redhat.com>
Date: Thu, 5 Feb 2026 16:44:32 +0000
Subject: [PATCH] fix: Eliminate race condition in testInputRequiredWorkflow

The test consumers were counting down their latches on the FIRST TaskEvent
received, which could be the intermediate WORKING state instead of the
expected terminal state (INPUT_REQUIRED or COMPLETED).

This caused intermittent test failures with:
  expected: <INPUT_REQUIRED> but was: <WORKING>

Fix: Both consumers now only count down when receiving the expected
terminal state:
- initialConsumer waits for INPUT_REQUIRED (not WORKING)
- completionConsumer waits for COMPLETED (not WORKING)

This matches the agent's event emission pattern:
  agentEmitter.startWork()    // WORKING state
  agentEmitter.requiresInput(...) // INPUT_REQUIRED state
  ...later...
  agentEmitter.startWork()    // WORKING state
  agentEmitter.complete()     // COMPLETED state
---
 .../apps/common/AbstractA2AServerTest.java    | 24 +++++++++++++++----
 1 file changed, 20 insertions(+), 4 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 6b99776c7..a15f2b087 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
@@ -1512,9 +1512,17 @@ public void testInputRequiredWorkflow() throws Exception {
             AtomicBoolean initialUnexpectedEvent = new AtomicBoolean(false);
 
             BiConsumer<ClientEvent, AgentCard> initialConsumer = (event, agentCard) -> {
+                // Idempotency guard: prevent late events from modifying state after latch countdown
+                if (initialLatch.getCount() == 0) {
+                    return;
+                }
                 if (event instanceof TaskEvent te) {
-                    initialState.set(te.getTask().status().state());
-                    initialLatch.countDown();
+                    TaskState state = te.getTask().status().state();
+                    initialState.set(state);
+                    // Only count down when we receive INPUT_REQUIRED, not intermediate states like WORKING
+                    if (state == TaskState.INPUT_REQUIRED) {
+                        initialLatch.countDown();
+                    }
                 } else {
                     initialUnexpectedEvent.set(true);
                 }
@@ -1538,9 +1546,17 @@ public void testInputRequiredWorkflow() throws Exception {
             AtomicBoolean completionUnexpectedEvent = new AtomicBoolean(false);
 
             BiConsumer<ClientEvent, AgentCard> completionConsumer = (event, agentCard) -> {
+                // Idempotency guard: prevent late events from modifying state after latch countdown
+                if (completionLatch.getCount() == 0) {
+                    return;
+                }
                 if (event instanceof TaskEvent te) {
-                    completedState.set(te.getTask().status().state());
-                    completionLatch.countDown();
+                    TaskState state = te.getTask().status().state();
+                    completedState.set(state);
+                    // Only count down when we receive COMPLETED, not intermediate states like WORKING
+                    if (state == TaskState.COMPLETED) {
+                        completionLatch.countDown();
+                    }
                 } else {
                     completionUnexpectedEvent.set(true);
                 }