Skip to content

fix: Make onSendMessage() handling less aggressive #448

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 1 commit into from
Closed

fix: Make onSendMessage() handling less aggressive #448

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
38 changes: 28 additions & 10 deletions server-common/src/main/java/io/a2a/server/requesthandlers/DefaultRequestHandler.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
import static java.util.concurrent.TimeUnit.*;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Objects;
import java.util.Optional;
Expand Down Expand Up @@ -56,6 +57,8 @@
import io.a2a.spec.TaskPushNotificationConfig;
import io.a2a.spec.TaskQueryParams;
import io.a2a.spec.TaskState;
import io.a2a.spec.TaskStatus;
import io.a2a.spec.TextPart;
import io.a2a.spec.UnsupportedOperationError;
import org.eclipse.microprofile.config.inject.ConfigProperty;
import org.slf4j.Logger;
Expand Down Expand Up @@ -262,7 +265,8 @@ public EventKind onMessageSend(MessageSendParams params, ServerCallContext conte
// 2. Close the queue to signal consumption can complete
// 3. Wait for consumption to finish processing events
// 4. Fetch final task state from TaskStore

boolean failed = false;
String error = null;
try {
// Step 1: Wait for agent to finish (with configurable timeout)
if (agentFuture != null) {
Expand All @@ -288,22 +292,36 @@ public EventKind onMessageSend(MessageSendParams params, ServerCallContext conte
}
} catch (InterruptedException e) {
Thread.currentThread().interrupt();
String msg = String.format("Error waiting for task %s completion", taskId);
LOGGER.warn(msg, e);
throw new InternalError(msg);
error = String.format("Error waiting for task %s completion", taskId);
LOGGER.warn(error, e);
failed = true;
} catch (java.util.concurrent.ExecutionException e) {
String msg = String.format("Error during task %s execution", taskId);
LOGGER.warn(msg, e.getCause());
throw new InternalError(msg);
error = String.format("Error during task %s execution", taskId);
LOGGER.warn(error, e.getCause());
failed = true;
} catch (java.util.concurrent.TimeoutException e) {
String msg = String.format("Timeout waiting for consumption to complete for task %s", taskId);
LOGGER.warn(msg, taskId);
throw new InternalError(msg);
error = String.format("Timeout waiting for consumption to complete for task %s", taskId);
LOGGER.warn(error, e);
failed = true;
}

// Step 4: Fetch the final task state from TaskStore (all events have been processed)
Task updatedTask = taskStore.get(taskId);
if (updatedTask != null) {
if (failed && !updatedTask.getStatus().state().isFinal()) {
// If we had exceptions above, update the Task to reflect that
Message msg = new Message.Builder()
.role(Message.Role.AGENT)
.parts(Collections.singletonList(new TextPart(error)))
.contextId(mss.requestContext.getContextId())
.taskId(taskId)
.build();

updatedTask = new Task.Builder(updatedTask)
.status(new TaskStatus(TaskState.FAILED, msg, null))
.build();
taskStore.save(updatedTask);
}
kind = updatedTask;
if (LOGGER.isDebugEnabled()) {
LOGGER.debug("Fetched final task for {} with state {} and {} artifacts",
Expand Down