From d7cd22bf494f56a7091283a8d3aa330c745cc827 Mon Sep 17 00:00:00 2001 From: Kabir Khan Date: Thu, 9 Jul 2026 12:28:38 +0100 Subject: [PATCH] fix: ensure correct task/contextId in emitted Messages --- .../sdk/server/tasks/AgentEmitter.java | 12 ++++ .../sdk/server/tasks/AgentEmitterTest.java | 66 +++++++++++++++++-- 2 files changed, 71 insertions(+), 7 deletions(-) diff --git a/server-common/src/main/java/org/a2aproject/sdk/server/tasks/AgentEmitter.java b/server-common/src/main/java/org/a2aproject/sdk/server/tasks/AgentEmitter.java index c062b6989..b15e99116 100644 --- a/server-common/src/main/java/org/a2aproject/sdk/server/tasks/AgentEmitter.java +++ b/server-common/src/main/java/org/a2aproject/sdk/server/tasks/AgentEmitter.java @@ -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; @@ -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; @@ -508,6 +512,14 @@ public void sendMessage(List> parts, @Nullable Map 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); } diff --git a/server-common/src/test/java/org/a2aproject/sdk/server/tasks/AgentEmitterTest.java b/server-common/src/test/java/org/a2aproject/sdk/server/tasks/AgentEmitterTest.java index b63f9b54c..3799b3e1b 100644 --- a/server-common/src/test/java/org/a2aproject/sdk/server/tasks/AgentEmitterTest.java +++ b/server-common/src/test/java/org/a2aproject/sdk/server/tasks/AgentEmitterTest.java @@ -44,6 +44,7 @@ public class AgentEmitterTest { private static final List> 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; @@ -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); @@ -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); @@ -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); @@ -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); @@ -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); @@ -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); @@ -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();