Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -1608,7 +1610,9 @@ public void testInputRequiredWorkflow() throws Exception {
assertEquals(TaskState.TASK_STATE_COMPLETED, completedState.get());

} finally {
deleteTaskInTaskStore(inputRequiredTaskId);
if (taskCreated) {
deleteTaskInTaskStore(inputRequiredTaskId);
}
}
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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()));
Expand All @@ -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;
}
Expand Down