fix: return TaskNotFoundError when SendMessage references non-existen… · a2aproject/a2a-java@5156deb · GitHub
Skip to content

Commit 5156deb

Browse files
yyy9942yyy9942
committed
fix: return TaskNotFoundError when SendMessage references non-existent taskId
Per A2A spec section 3.4.2, when a client includes a taskId in a Message, it MUST reference an existing task. Previously, the SDK silently created a new task using the client-provided taskId, which the spec explicitly forbids. Applies to all three transports (JSON-RPC, gRPC, HTTP+JSON) via the common DefaultRequestHandler.initMessageSend() code path. Existing tests that relied on the buggy behaviour have been updated to either pre-populate the task store (when testing follow-up messages) or omit the client-provided taskId on the initial message (letting the server generate it, matching the spec-correct flow). This fixes #766
1 parent 98d2bdb commit 5156deb

10 files changed

Lines changed: 438 additions & 258 deletions

File tree

‎extras/push-notification-config-store-database-jpa/src/test/java/org/a2aproject/sdk/extras/pushnotificationconfigstore/database/jpa/JpaDatabasePushNotificationConfigStoreIntegrationTest.java‎

Lines changed: 21 additions & 8 deletions

‎extras/queue-manager-replicated/tests-multi-instance/tests/src/test/java/org/a2aproject/sdk/extras/queuemanager/replicated/tests/multiinstance/MultiInstanceReplicationTest.java‎

Lines changed: 23 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,9 @@
99
import java.time.Duration;
1010
import java.util.Collections;
1111
import java.util.List;
12+
import java.util.concurrent.CountDownLatch;
1213
import java.util.concurrent.CopyOnWriteArrayList;
14+
import java.util.concurrent.TimeUnit;
1315
import java.util.concurrent.atomic.AtomicBoolean;
1416
import java.util.concurrent.atomic.AtomicInteger;
1517
import java.util.concurrent.atomic.AtomicReference;
@@ -19,6 +21,7 @@
1921
import org.a2aproject.sdk.A2A;
2022
import org.a2aproject.sdk.client.Client;
2123
import org.a2aproject.sdk.client.ClientEvent;
24+
import org.a2aproject.sdk.client.TaskEvent;
2225
import org.a2aproject.sdk.client.config.ClientConfig;
2326
import org.a2aproject.sdk.client.transport.jsonrpc.JSONRPCTransport;
2427
import org.a2aproject.sdk.client.transport.jsonrpc.JSONRPCTransportConfig;
@@ -254,15 +257,13 @@ public void testInfrastructureStartup() {
254257
*/
255258
@Test
256259
public void testMultiInstanceEventReplication() throws Exception {
257-
final String taskId = "replication-test-task-" + System.currentTimeMillis();
258-
final String contextId = "replication-test-context";
259-
260260
Throwable testFailure = null;
261+
final String[] taskIdHolder = {null};
262+
final String[] contextIdHolder = {null};
263+
261264
try {
262-
// Step 1: Send initial message NON-streaming to create task
265+
// Step 1: Send initial message NON-streaming to create task (no client-provided taskId)
263266
Message initialMessage = Message.builder(A2A.toUserMessage("Initial test message"))
264-
.taskId(taskId)
265-
.contextId(contextId)
266267
.build();
267268

268269
// Use NON-streaming client to create the task
@@ -272,13 +273,22 @@ public void testMultiInstanceEventReplication() throws Exception {
272273
.withTransport(JSONRPCTransport.class, new JSONRPCTransportConfig())
273274
.build();
274275

275-
Task createdTask = null;
276276
try {
277-
nonStreamingClient.sendMessage(initialMessage, null);
277+
CountDownLatch createLatch = new CountDownLatch(1);
278+
AtomicReference<Task> taskRef = new AtomicReference<>();
279+
nonStreamingClient.sendMessage(initialMessage, List.of((ClientEvent event, AgentCard card) -> {
280+
if (event instanceof TaskEvent te) {
281+
taskRef.set(te.getTask());
282+
}
283+
createLatch.countDown();
284+
}), (Throwable err) -> createLatch.countDown());
285+
assertTrue(createLatch.await(15, TimeUnit.SECONDS), "Task creation timed out");
278286

279-
// Retrieve the task to verify it was created
280-
createdTask = nonStreamingClient.getTask(new TaskQueryParams(taskId), null);
287+
Task createdTask = taskRef.get();
281288
assertNotNull(createdTask, "Task should be created");
289+
taskIdHolder[0] = createdTask.id();
290+
contextIdHolder[0] = createdTask.contextId();
291+
assertNotNull(taskIdHolder[0], "Server-generated task ID should not be null");
282292

283293
// Task should be in a non-final state (SUBMITTED or WORKING are both valid)
284294
TaskState state = createdTask.status().state();
@@ -297,6 +307,9 @@ public void testMultiInstanceEventReplication() throws Exception {
297307
throw e;
298308
}
299309

310+
final String taskId = taskIdHolder[0];
311+
final String contextId = contextIdHolder[0];
312+
300313
// Step 2: Subscribe from both app1 and app2 with proper latches
301314

302315
// We need to wait for at least 3 new events after resubscription:

‎extras/queue-manager-replicated/tests-single-instance/src/test/java/org/a2aproject/sdk/extras/queuemanager/replicated/tests/KafkaReplicationIntegrationTest.java‎

Lines changed: 13 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -130,19 +130,14 @@ public void tearDown() throws Exception {
130130

131131
@Test
132132
public void testA2AMessageReplicatedToKafka() throws Exception {
133-
String taskId = "kafka-replication-test-" + System.currentTimeMillis();
134-
String contextId = "test-context-" + System.currentTimeMillis();
135-
136133
// Clear any previous events
137134
testConsumer.clear();
138135

139-
// Send A2A message that should trigger events and replication
136+
// Send A2A message that should trigger events and replication (no client-provided taskId)
140137
Message message = Message.builder()
141138
.role(Message.Role.ROLE_USER)
142139
.parts(List.of(new TextPart("create")))
143-
.taskId(taskId)
144140
.messageId("test-msg-" + System.currentTimeMillis())
145-
.contextId(contextId)
146141
.build();
147142

148143
CountDownLatch a2aLatch = new CountDownLatch(1);
@@ -166,7 +161,9 @@ public void testA2AMessageReplicatedToKafka() throws Exception {
166161

167162
Task task = createdTask.get();
168163
assertNotNull(task, "Task should be created");
169-
assertEquals(taskId, task.id());
164+
String taskId = task.id();
165+
String contextId = task.contextId();
166+
assertNotNull(taskId, "Server-generated task ID should not be null");
170167
assertEquals(TaskState.TASK_STATE_SUBMITTED, task.status().state());
171168

172169
// Wait for the event to be replicated to Kafka
@@ -193,19 +190,14 @@ public void testA2AMessageReplicatedToKafka() throws Exception {
193190

194191
@Test
195192
public void testKafkaEventReceivedByA2AServer() throws Exception {
196-
String taskId = "kafka-to-a2a-test-" + System.currentTimeMillis();
197-
String contextId = "test-context-" + System.currentTimeMillis();
198-
199193
// Clear any previous events
200194
testConsumer.clear();
201195

202-
// First create a task in the A2A system using non-streaming client
196+
// First create a task in the A2A system using non-streaming client (no client-provided taskId)
203197
Message createMessage = Message.builder()
204198
.role(Message.Role.ROLE_USER)
205199
.parts(List.of(new TextPart("create")))
206-
.taskId(taskId)
207200
.messageId("create-msg-" + System.currentTimeMillis())
208-
.contextId(contextId)
209201
.build();
210202

211203
CountDownLatch createLatch = new CountDownLatch(1);
@@ -223,6 +215,9 @@ public void testKafkaEventReceivedByA2AServer() throws Exception {
223215
assertTrue(createLatch.await(15, TimeUnit.SECONDS), "Task creation timed out");
224216
Task initialTask = createdTask.get();
225217
assertNotNull(initialTask, "Task should be created");
218+
String taskId = initialTask.id();
219+
String contextId = initialTask.contextId();
220+
assertNotNull(taskId, "Server-generated task ID should not be null");
226221
assertEquals(TaskState.TASK_STATE_SUBMITTED, initialTask.status().state(), "Initial task should be in SUBMITTED state");
227222

228223
// Add a small delay to ensure the task is fully processed before resubscription
@@ -312,20 +307,15 @@ public void testKafkaEventReceivedByA2AServer() throws Exception {
312307

313308
@Test
314309
public void testQueueClosedEventTerminatesRemoteSubscribers() throws Exception {
315-
String taskId = "queue-closed-test-" + System.currentTimeMillis();
316-
String contextId = "test-context-" + System.currentTimeMillis();
317-
318310
// Clear any previous events
319311
testConsumer.clear();
320312

321-
// Use polling (non-blocking) client with "working" command
313+
// Use polling (non-blocking) client with "working" command (no client-provided taskId)
322314
// This creates task in WORKING state (non-final) and keeps queue alive
323315
Message workingMessage = Message.builder()
324316
.role(Message.Role.ROLE_USER)
325317
.parts(List.of(new TextPart("working")))
326-
.taskId(taskId)
327318
.messageId("working-msg-" + System.currentTimeMillis())
328-
.contextId(contextId)
329319
.build();
330320

331321
CountDownLatch workingLatch = new CountDownLatch(1);
@@ -348,7 +338,7 @@ public void testQueueClosedEventTerminatesRemoteSubscribers() throws Exception {
348338
assertTrue(workingLatch.await(15, TimeUnit.SECONDS), "Task creation timed out");
349339
String createdTaskId = taskIdRef.get();
350340
assertNotNull(createdTaskId, "Task should be created");
351-
assertEquals(taskId, createdTaskId);
341+
String taskId = createdTaskId;
352342

353343
// Set up streaming resubscription to listen for the QueueClosedEvent
354344
CountDownLatch streamCompletedLatch = new CountDownLatch(1);
@@ -408,19 +398,14 @@ public void testQueueClosedEventTerminatesRemoteSubscribers() throws Exception {
408398

409399
@Test
410400
public void testPoisonPillGenerationOnTaskFinalization() throws Exception {
411-
String taskId = "poison-pill-gen-test-" + System.currentTimeMillis();
412-
String contextId = "test-context-" + System.currentTimeMillis();
413-
414401
// Clear any previous events
415402
testConsumer.clear();
416403

417-
// Create a task that will be completed (finalized)
404+
// Create a task that will be completed (finalized) - no client-provided taskId
418405
Message completeMessage = Message.builder()
419406
.role(Message.Role.ROLE_USER)
420407
.parts(List.of(new TextPart("complete")))
421-
.taskId(taskId)
422408
.messageId("complete-msg-" + System.currentTimeMillis())
423-
.contextId(contextId)
424409
.build();
425410

426411
CountDownLatch completeLatch = new CountDownLatch(1);
@@ -439,6 +424,8 @@ public void testPoisonPillGenerationOnTaskFinalization() throws Exception {
439424
assertTrue(completeLatch.await(15, TimeUnit.SECONDS), "Task creation timed out");
440425
Task createdTask = finalTask.get();
441426
assertNotNull(createdTask, "Task should be created");
427+
String taskId = createdTask.id();
428+
assertNotNull(taskId, "Server-generated task ID should not be null");
442429

443430
// The task should complete very quickly since it's a simple operation
444431
// Wait a moment to ensure all events have been enqueued

‎extras/task-store-database-jpa/src/test/java/org/a2aproject/sdk/extras/taskstore/database/jpa/JpaDatabaseTaskStoreIntegrationTest.java‎

Lines changed: 4 additions & 8 deletions

0 commit comments

Comments
 (0)