Skip to content

fix: ensure correct task/contextId in emitted Messages #976

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

Merged
jmesnil merged 1 commit into from
Jul 9, 2026
+71 −7
Merged

fix: ensure correct task/contextId in emitted Messages #976

Show file tree
Hide file tree
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
12 changes: 12 additions & 0 deletions server-common/src/main/java/org/a2aproject/sdk/server/tasks/AgentEmitter.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,8 @@
import java.util.concurrent.atomic.AtomicBoolean;

import org.a2aproject.sdk.server.agentexecution.RequestContext;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.a2aproject.sdk.server.events.EventQueue;
import org.a2aproject.sdk.spec.A2AError;
import org.a2aproject.sdk.spec.Artifact;
Expand Down Expand Up @@ -94,6 +96,8 @@
* @since 1.0.0
*/
public class AgentEmitter {
private static final Logger LOGGER = LoggerFactory.getLogger(AgentEmitter.class);

private final EventQueue eventQueue;
private final String taskId;
private final String contextId;
Expand Down Expand Up @@ -508,6 +512,14 @@ public void sendMessage(List<Part<?>> parts, @Nullable Map<String, Object> metad
* @since 1.0.0
*/
public void sendMessage(Message message) {
if (message.taskId() != null && !message.taskId().equals(taskId)) {
LOGGER.error("Message taskId mismatch: expected={}, actual={}", taskId, message.taskId());
throw new IllegalArgumentException("Message taskId does not match the emitter's taskId");
}
if (message.contextId() != null && !message.contextId().equals(contextId)) {
LOGGER.error("Message contextId mismatch: expected={}, actual={}", contextId, message.contextId());
throw new IllegalArgumentException("Message contextId does not match the emitter's contextId");
}
eventQueue.enqueueEvent(message);
}

Expand Down
66 changes: 59 additions & 7 deletions server-common/src/test/java/org/a2aproject/sdk/server/tasks/AgentEmitterTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,7 @@ public class AgentEmitterTest {
private static final List<Part<?>> SAMPLE_PARTS = List.of(new TextPart("Test message"));

private static final PushNotificationSender NOOP_PUSHNOTIFICATION_SENDER = (event, snapshot) -> {};
public static final int WAIT_MILLI_SECONDS = 5000;

EventQueue eventQueue;
private MainEventBus mainEventBus;
Expand Down Expand Up @@ -82,7 +83,7 @@ public void cleanup() {
@Test
public void testAddArtifactWithCustomIdAndName() throws Exception {
agentEmitter.addArtifact(SAMPLE_PARTS, "custom-artifact-id", "Custom Artifact", null);
EventQueueItem item = eventQueue.dequeueEventItem(5000);
EventQueueItem item = eventQueue.dequeueEventItem(WAIT_MILLI_SECONDS);
assertNotNull(item);
Event event = item.getEvent();
assertNotNull(event);
Expand Down Expand Up @@ -267,7 +268,7 @@ public void testNewAgentMessageWithMetadata() throws Exception {
@Test
public void testAddArtifactWithAppendTrue() throws Exception {
agentEmitter.addArtifact(SAMPLE_PARTS, "artifact-id", "Test Artifact", null, true, null);
EventQueueItem item = eventQueue.dequeueEventItem(5000);
EventQueueItem item = eventQueue.dequeueEventItem(WAIT_MILLI_SECONDS);
assertNotNull(item);
Event event = item.getEvent();
assertNotNull(event);
Expand All @@ -288,7 +289,7 @@ public void testAddArtifactWithAppendTrue() throws Exception {
@Test
public void testAddArtifactWithLastChunkTrue() throws Exception {
agentEmitter.addArtifact(SAMPLE_PARTS, "artifact-id", "Test Artifact", null, null, true);
EventQueueItem item = eventQueue.dequeueEventItem(5000);
EventQueueItem item = eventQueue.dequeueEventItem(WAIT_MILLI_SECONDS);
assertNotNull(item);
Event event = item.getEvent();
assertNotNull(event);
Expand All @@ -305,7 +306,7 @@ public void testAddArtifactWithLastChunkTrue() throws Exception {
@Test
public void testAddArtifactWithAppendAndLastChunk() throws Exception {
agentEmitter.addArtifact(SAMPLE_PARTS, "artifact-id", "Test Artifact", null, true, false);
EventQueueItem item = eventQueue.dequeueEventItem(5000);
EventQueueItem item = eventQueue.dequeueEventItem(WAIT_MILLI_SECONDS);
assertNotNull(item);
Event event = item.getEvent();
assertNotNull(event);
Expand All @@ -321,7 +322,7 @@ public void testAddArtifactWithAppendAndLastChunk() throws Exception {
@Test
public void testAddArtifactGeneratesIdWhenNull() throws Exception {
agentEmitter.addArtifact(SAMPLE_PARTS, null, "Test Artifact", null);
EventQueueItem item = eventQueue.dequeueEventItem(5000);
EventQueueItem item = eventQueue.dequeueEventItem(WAIT_MILLI_SECONDS);
assertNotNull(item);
Event event = item.getEvent();
assertNotNull(event);
Expand Down Expand Up @@ -419,7 +420,7 @@ public void testConcurrentCompletionAttempts() throws Exception {
thread2.join();

// Exactly one event should have been queued
EventQueueItem item = eventQueue.dequeueEventItem(5000);
EventQueueItem item = eventQueue.dequeueEventItem(WAIT_MILLI_SECONDS);
assertNotNull(item);
Event event = item.getEvent();
assertNotNull(event);
Expand All @@ -433,9 +434,60 @@ public void testConcurrentCompletionAttempts() throws Exception {
assertNull(eventQueue.dequeueEventItem(0));
}

@Test
public void sendMessageWithMatchingIdsSucceeds() throws Exception {
agentEmitter.sendMessage(SAMPLE_MESSAGE);

EventQueueItem item = eventQueue.dequeueEventItem(WAIT_MILLI_SECONDS);
assertNotNull(item);
assertInstanceOf(Message.class, item.getEvent());
Message message = (Message) item.getEvent();
assertEquals(TEST_TASK_ID, message.taskId());
assertEquals(TEST_TASK_CONTEXT_ID, message.contextId());
}

@Test
public void sendMessageWithNullIdsSucceeds() throws Exception {
Message message = Message.builder()
.role(ROLE_AGENT)
.parts(new TextPart("no ids"))
.build();
agentEmitter.sendMessage(message);

EventQueueItem item = eventQueue.dequeueEventItem(WAIT_MILLI_SECONDS);
assertNotNull(item);
assertInstanceOf(Message.class, item.getEvent());
}

@Test
public void sendMessageWithMismatchedTaskIdThrows() {
Message mismatchedMessage = Message.builder()
.taskId("wrong-task-id")
.contextId(TEST_TASK_CONTEXT_ID)
.role(ROLE_AGENT)
.parts(new TextPart("mismatched"))
.build();
IllegalArgumentException ex = assertThrows(IllegalArgumentException.class,
() -> agentEmitter.sendMessage(mismatchedMessage));
assertTrue(ex.getMessage().contains("Message taskId does not match"));
}

@Test
public void sendMessageWithMismatchedContextIdThrows() {
Message mismatchedMessage = Message.builder()
.taskId(TEST_TASK_ID)
.contextId("wrong-context-id")
.role(ROLE_AGENT)
.parts(new TextPart("mismatched"))
.build();
IllegalArgumentException ex = assertThrows(IllegalArgumentException.class,
() -> agentEmitter.sendMessage(mismatchedMessage));
assertTrue(ex.getMessage().contains("Message contextId does not match"));
}

private TaskStatusUpdateEvent checkTaskStatusUpdateEventOnQueue(boolean isFinal, TaskState state, Message statusMessage) throws Exception {
// Wait up to 5 seconds for event (async MainEventBusProcessor needs time to distribute)
EventQueueItem item = eventQueue.dequeueEventItem(5000);
EventQueueItem item = eventQueue.dequeueEventItem(WAIT_MILLI_SECONDS);
assertNotNull(item);
Event event = item.getEvent();

Expand Down
Loading