@@ -875,6 +875,125 @@ public void testSubscribeExistingTaskSuccessWithClientConsumers() throws Excepti
875875 }
876876 }
877877
878+ /**
879+ * Tests that SubscribeToTask stream stays open for interrupted states (INPUT_REQUIRED, AUTH_REQUIRED)
880+ * and only terminates on terminal states.
881+ * <p>
882+ * Per A2A Protocol Specification 3.1.6 (SubscribeToTask):
883+ * "The stream MUST terminate when the task reaches a terminal state (completed, failed, canceled, or rejected)."
884+ * <p>
885+ * Interrupted states are NOT terminal - the stream should remain open to deliver future state updates.
886+ * <p>
887+ * This test addresses issue #754: Stream was incorrectly closing immediately for INPUT_REQUIRED state.
888+ */
889+ @ Test
890+ @ Timeout (value = 3 , unit = TimeUnit .MINUTES )
891+ public void testSubscribeToTaskWithInterruptedStateKeepsStreamOpen () throws Exception {
892+ // Create task in INPUT_REQUIRED state (interrupted, not terminal)
893+ Task inputRequiredTask = Task .builder ()
894+ .id ("task-input-required-" + UUID .randomUUID ())
895+ .contextId ("session-xyz" )
896+ .status (new TaskStatus (TaskState .TASK_STATE_INPUT_REQUIRED ))
897+ .build ();
898+
899+ saveTaskInTaskStore (inputRequiredTask );
900+ try {
901+ ensureQueueForTask (inputRequiredTask .id ());
902+
903+ // Track events received through the stream
904+ CopyOnWriteArrayList <io .a2a .spec .UpdateEvent > receivedEvents = new CopyOnWriteArrayList <>();
905+ AtomicBoolean receivedInitialTask = new AtomicBoolean (false );
906+ AtomicBoolean streamClosedPrematurely = new AtomicBoolean (false );
907+ AtomicReference <Throwable > errorRef = new AtomicReference <>();
908+ CountDownLatch completionLatch = new CountDownLatch (1 );
909+
910+ // Consumer to track all events
911+ BiConsumer <ClientEvent , AgentCard > consumer = (event , agentCard ) -> {
912+ if (event instanceof TaskEvent taskEvent ) {
913+ if (!receivedInitialTask .get ()) {
914+ receivedInitialTask .set (true );
915+ // First event should be the initial task snapshot
916+ return ;
917+ }
918+ } else if (event instanceof TaskUpdateEvent taskUpdateEvent ) {
919+ io .a2a .spec .UpdateEvent updateEvent = taskUpdateEvent .getUpdateEvent ();
920+ receivedEvents .add (updateEvent );
921+
922+ // Check if this is the final terminal state
923+ if (updateEvent instanceof TaskStatusUpdateEvent tue && tue .isFinal ()) {
924+ completionLatch .countDown ();
925+ }
926+ }
927+ };
928+
929+ // Error handler to detect premature stream closure
930+ Consumer <Throwable > errorHandler = error -> {
931+ if (!isStreamClosedError (error )) {
932+ errorRef .set (error );
933+ }
934+ // If completion latch hasn't been counted down yet, stream closed prematurely
935+ if (completionLatch .getCount () > 0 ) {
936+ streamClosedPrematurely .set (true );
937+ }
938+ completionLatch .countDown ();
939+ };
940+
941+ // Subscribe to the task
942+ CountDownLatch subscriptionLatch = new CountDownLatch (1 );
943+ awaitStreamingSubscription ()
944+ .whenComplete ((unused , throwable ) -> subscriptionLatch .countDown ());
945+
946+ getClient ().subscribeToTask (new TaskIdParams (inputRequiredTask .id ()), List .of (consumer ), errorHandler );
947+
948+ // Wait for subscription to be established
949+ assertTrue (subscriptionLatch .await (15 , TimeUnit .SECONDS ), "Subscription should be established" );
950+
951+ // Wait a bit to ensure stream doesn't close prematurely
952+ Thread .sleep (500 );
953+
954+ // Verify stream is still open (no premature closure)
955+ assertFalse (streamClosedPrematurely .get (),
956+ "Stream should NOT close for INPUT_REQUIRED state (interrupted, not terminal)" );
957+
958+ // Send status update to WORKING (still non-terminal)
959+ enqueueEventOnServer (TaskStatusUpdateEvent .builder ()
960+ .taskId (inputRequiredTask .id ())
961+ .contextId (inputRequiredTask .contextId ())
962+ .status (new TaskStatus (TaskState .TASK_STATE_WORKING ))
963+ .build ());
964+
965+ // Wait a bit and verify stream is still open
966+ Thread .sleep (500 );
967+ assertFalse (streamClosedPrematurely .get (),
968+ "Stream should remain open after transitioning to WORKING" );
969+
970+ // Send terminal status update to COMPLETED
971+ enqueueEventOnServer (TaskStatusUpdateEvent .builder ()
972+ .taskId (inputRequiredTask .id ())
973+ .contextId (inputRequiredTask .contextId ())
974+ .status (new TaskStatus (TaskState .TASK_STATE_COMPLETED ))
975+ .build ());
976+
977+ // Now stream should close
978+ assertTrue (completionLatch .await (30 , TimeUnit .SECONDS ),
979+ "Stream should close after terminal state" );
980+
981+ // Verify we received both updates before stream closed
982+ assertEquals (2 , receivedEvents .size (),
983+ "Should receive both status updates before stream closes" );
984+
985+ TaskStatusUpdateEvent firstUpdate = (TaskStatusUpdateEvent ) receivedEvents .get (0 );
986+ assertEquals (TaskState .TASK_STATE_WORKING , firstUpdate .status ().state ());
987+
988+ TaskStatusUpdateEvent secondUpdate = (TaskStatusUpdateEvent ) receivedEvents .get (1 );
989+ assertEquals (TaskState .TASK_STATE_COMPLETED , secondUpdate .status ().state ());
990+
991+ assertNull (errorRef .get (), "Should not have any errors" );
992+ } finally {
993+ deleteTaskInTaskStore (inputRequiredTask .id ());
994+ }
995+ }
996+
878997 @ Test
879998 public void testSubscribeNoExistingTaskError () throws Exception {
880999 CountDownLatch errorLatch = new CountDownLatch (1 );
0 commit comments