Skip to content

Commit 2b87d65

Browse files
google-genai-botcopybara-github
google-genai-bot
authored and
copybara-github
committed
fix: update stream completion check in A2A SDK to handle all terminal and interrupted task states
PiperOrigin-RevId: 964059169
1 parent 3e6b915 commit 2b87d65

2 files changed

Lines changed: 143 additions & 4 deletions

File tree

‎a2a/src/main/java/com/google/adk/a2a/agent/RemoteA2AAgent.java‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@
3636
import com.google.genai.types.Part;
3737
import io.a2a.client.Client;
3838
import io.a2a.client.ClientEvent;
39+
import io.a2a.client.MessageEvent;
3940
import io.a2a.client.TaskEvent;
4041
import io.a2a.client.TaskUpdateEvent;
4142
import io.a2a.spec.A2AClientException;
@@ -580,13 +581,15 @@ private Event createAggregatedEvent(Content content, @Nullable ClientEvent trigg
580581
}
581582

582583
private static boolean isCompleted(ClientEvent event) {
583-
TaskState executionState = TaskState.UNKNOWN;
584+
TaskState state;
584585
if (event instanceof TaskEvent taskEvent) {
585-
executionState = taskEvent.getTask().getStatus().state();
586+
state = taskEvent.getTask().getStatus().state();
586587
} else if (event instanceof TaskUpdateEvent updateEvent) {
587-
executionState = updateEvent.getTask().getStatus().state();
588+
state = updateEvent.getTask().getStatus().state();
589+
} else {
590+
return false;
588591
}
589-
return executionState.equals(TaskState.COMPLETED);
592+
return state.isFinal() || state == TaskState.INPUT_REQUIRED || state == TaskState.AUTH_REQUIRED;
590593
}
591594

592595
private static ImmutableList<Part> eventParts(Event event) {

‎a2a/src/test/java/com/google/adk/a2a/agent/RemoteA2AAgentTest.java‎

Lines changed: 136 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@
4646
import com.google.genai.types.Part;
4747
import io.a2a.client.Client;
4848
import io.a2a.client.ClientEvent;
49+
import io.a2a.client.MessageEvent;
4950
import io.a2a.client.TaskEvent;
5051
import io.a2a.client.TaskUpdateEvent;
5152
import io.a2a.spec.AgentCapabilities;
@@ -889,6 +890,47 @@ private ClientEvent createFinalEvent(String text) {
889890
return createTestEvent(new TextPart(text), TaskState.COMPLETED, false, false);
890891
}
891892

893+
private ClientEvent createTerminalStatusUpdateEvent(String message, TaskState state) {
894+
Task task =
895+
new Task.Builder()
896+
.id("task-id-1")
897+
.contextId("context-1")
898+
.status(new TaskStatus(state))
899+
.build();
900+
TaskStatusUpdateEvent statusUpdate =
901+
new TaskStatusUpdateEvent.Builder()
902+
.taskId("task-id-1")
903+
.contextId("context-1")
904+
.status(
905+
new TaskStatus(
906+
state,
907+
new Message.Builder()
908+
.role(Message.Role.AGENT)
909+
.parts(ImmutableList.of(new TextPart(message)))
910+
.build(),
911+
null))
912+
.build();
913+
return new TaskUpdateEvent(task, statusUpdate);
914+
}
915+
916+
private ClientEvent createFailedEvent(String errorMessage) {
917+
return createTerminalStatusUpdateEvent(errorMessage, TaskState.FAILED);
918+
}
919+
920+
private ClientEvent createCanceledEvent(String message) {
921+
return createTerminalStatusUpdateEvent(message, TaskState.CANCELED);
922+
}
923+
924+
private ClientEvent createMessageEvent(String text) {
925+
Message message =
926+
new Message.Builder()
927+
.messageId("msg-id-1")
928+
.role(Message.Role.AGENT)
929+
.parts(new TextPart(text))
930+
.build();
931+
return new MessageEvent(message);
932+
}
933+
892934
private ClientEvent createTestEvent(
893935
io.a2a.spec.Part<?> part, TaskState state, boolean append, boolean lastChunk) {
894936
Artifact artifact =
@@ -916,6 +958,100 @@ private ClientEvent createTestEvent(
916958
return new TaskUpdateEvent(task, updateEvent);
917959
}
918960

961+
@Test
962+
public void runAsync_terminatesOnFailureTaskState() {
963+
RemoteA2AAgent agent = createAgent();
964+
mockStreamResponse(
965+
consumer -> {
966+
consumer.accept(createPartialEvent("Processing data...", true, false), agentCard);
967+
consumer.accept(createFailedEvent("Internal Server Error"), agentCard);
968+
});
969+
970+
List<Event> events = agent.runAsync(invocationContext).toList().blockingGet();
971+
972+
// The stream must terminate cleanly and include the final event with error info.
973+
assertThat(events).hasSize(2);
974+
assertText(events.get(0), "Processing data...");
975+
assertText(events.get(1), "Processing data..."); // aggregated/merged
976+
assertThat(events.get(1).errorMessage().orElse(null)).isEqualTo("Internal Server Error");
977+
}
978+
979+
@Test
980+
public void runAsync_terminatesOnCanceledTaskState() {
981+
RemoteA2AAgent agent = createAgent();
982+
mockStreamResponse(
983+
consumer -> {
984+
consumer.accept(createPartialEvent("Processing data...", true, false), agentCard);
985+
consumer.accept(createCanceledEvent("Execution Canceled"), agentCard);
986+
});
987+
988+
List<Event> events = agent.runAsync(invocationContext).toList().blockingGet();
989+
990+
// The stream must terminate cleanly and include the final event with cancellation info.
991+
assertThat(events).hasSize(3);
992+
assertText(events.get(0), "Processing data...");
993+
assertText(events.get(1), "Processing data..."); // aggregated
994+
assertText(events.get(2), "Execution Canceled"); // terminal canceled event
995+
}
996+
997+
@Test
998+
public void runAsync_terminatesOnInputRequiredTaskState() {
999+
RemoteA2AAgent agent = createAgent();
1000+
mockStreamResponse(
1001+
consumer -> {
1002+
consumer.accept(createPartialEvent("Processing data...", true, false), agentCard);
1003+
consumer.accept(
1004+
createTerminalStatusUpdateEvent("User Action Needed", TaskState.INPUT_REQUIRED),
1005+
agentCard);
1006+
});
1007+
1008+
List<Event> events = agent.runAsync(invocationContext).toList().blockingGet();
1009+
1010+
assertThat(events).hasSize(3);
1011+
assertText(events.get(0), "Processing data...");
1012+
assertText(events.get(1), "Processing data..."); // aggregated
1013+
assertText(events.get(2), "User Action Needed");
1014+
}
1015+
1016+
@Test
1017+
public void runAsync_terminatesOnRejectedTaskState() {
1018+
RemoteA2AAgent agent = createAgent();
1019+
mockStreamResponse(
1020+
consumer -> {
1021+
consumer.accept(createPartialEvent("Analyzing request...", true, false), agentCard);
1022+
consumer.accept(
1023+
createTerminalStatusUpdateEvent("Execution Rejected by Policy", TaskState.REJECTED),
1024+
agentCard);
1025+
});
1026+
1027+
List<Event> events = agent.runAsync(invocationContext).toList().blockingGet();
1028+
1029+
assertThat(events).hasSize(3);
1030+
assertText(events.get(0), "Analyzing request...");
1031+
assertText(events.get(1), "Analyzing request..."); // aggregated
1032+
assertText(events.get(2), "Execution Rejected by Policy");
1033+
}
1034+
1035+
@Test
1036+
public void runAsync_doesNotTerminateOnMessageEvent() {
1037+
RemoteA2AAgent agent = createAgent();
1038+
mockStreamResponse(
1039+
consumer -> {
1040+
consumer.accept(createPartialEvent("Processing data...", true, false), agentCard);
1041+
consumer.accept(createMessageEvent("Standard chat update message"), agentCard);
1042+
consumer.accept(createFinalEvent("Done"), agentCard);
1043+
});
1044+
1045+
List<Event> events = agent.runAsync(invocationContext).toList().blockingGet();
1046+
1047+
// The stream must not terminate early on MessageEvent, and run until createFinalEvent.
1048+
assertThat(events).hasSize(4);
1049+
assertText(events.get(0), "Processing data...");
1050+
assertText(events.get(1), "Processing data..."); // aggregated (flushed)
1051+
assertText(events.get(2), "Standard chat update message"); // message event
1052+
assertText(events.get(3), "Done"); // terminal completed event
1053+
}
1054+
9191055
private RemoteA2AAgent.Builder getAgentBuilder() {
9201056
return RemoteA2AAgent.builder().name("remote-agent").a2aClient(mockClient).agentCard(agentCard);
9211057
}

0 commit comments

Comments
 (0)