diff --git a/extras/push-notification-config-store-database-jpa/src/test/java/io/a2a/extras/pushnotificationconfigstore/database/jpa/JpaDatabasePushNotificationConfigStoreIntegrationTest.java b/extras/push-notification-config-store-database-jpa/src/test/java/io/a2a/extras/pushnotificationconfigstore/database/jpa/JpaDatabasePushNotificationConfigStoreIntegrationTest.java index 58681cf44..57e817d5d 100644 --- a/extras/push-notification-config-store-database-jpa/src/test/java/io/a2a/extras/pushnotificationconfigstore/database/jpa/JpaDatabasePushNotificationConfigStoreIntegrationTest.java +++ b/extras/push-notification-config-store-database-jpa/src/test/java/io/a2a/extras/pushnotificationconfigstore/database/jpa/JpaDatabasePushNotificationConfigStoreIntegrationTest.java @@ -88,7 +88,7 @@ public void testDirectNotificationTrigger() { mockPushNotificationSender.sendNotification(testTask); // Verify it was captured - Queue captured = mockPushNotificationSender.getCapturedTasks(); + Queue captured = mockPushNotificationSender.getCapturedEvents(); assertEquals(1, captured.size()); assertEquals("direct-test-task", captured.peek().getId()); } @@ -151,7 +151,7 @@ public void testJpaDatabasePushNotificationConfigStoreIntegration() throws Excep boolean notificationReceived = false; while (System.currentTimeMillis() < end) { - if (!mockPushNotificationSender.getCapturedTasks().isEmpty()) { + if (!mockPushNotificationSender.getCapturedEvents().isEmpty()) { notificationReceived = true; break; } @@ -161,7 +161,7 @@ public void testJpaDatabasePushNotificationConfigStoreIntegration() throws Excep assertTrue(notificationReceived, "Timeout waiting for push notification."); // Step 6: Verify the captured notification - Queue capturedTasks = mockPushNotificationSender.getCapturedTasks(); + Queue capturedTasks = mockPushNotificationSender.getCapturedEvents(); // Verify the notification contains the correct task with artifacts Task notifiedTaskWithArtifact = capturedTasks.stream() diff --git a/extras/push-notification-config-store-database-jpa/src/test/java/io/a2a/extras/pushnotificationconfigstore/database/jpa/MockPushNotificationSender.java b/extras/push-notification-config-store-database-jpa/src/test/java/io/a2a/extras/pushnotificationconfigstore/database/jpa/MockPushNotificationSender.java index 0a6bba415..568d02e04 100644 --- a/extras/push-notification-config-store-database-jpa/src/test/java/io/a2a/extras/pushnotificationconfigstore/database/jpa/MockPushNotificationSender.java +++ b/extras/push-notification-config-store-database-jpa/src/test/java/io/a2a/extras/pushnotificationconfigstore/database/jpa/MockPushNotificationSender.java @@ -3,12 +3,12 @@ import java.util.Queue; import java.util.concurrent.ConcurrentLinkedQueue; +import io.a2a.spec.StreamingEventKind; import jakarta.annotation.Priority; import jakarta.enterprise.context.ApplicationScoped; import jakarta.enterprise.inject.Alternative; import io.a2a.server.tasks.PushNotificationSender; -import io.a2a.spec.Task; /** * Mock implementation of PushNotificationSender for integration testing. @@ -19,18 +19,18 @@ @Priority(100) public class MockPushNotificationSender implements PushNotificationSender { - private final Queue capturedTasks = new ConcurrentLinkedQueue<>(); + private final Queue capturedEvents = new ConcurrentLinkedQueue<>(); @Override - public void sendNotification(Task task) { - capturedTasks.add(task); + public void sendNotification(StreamingEventKind kind) { + capturedEvents.add(kind); } - public Queue getCapturedTasks() { - return capturedTasks; + public Queue getCapturedEvents() { + return capturedEvents; } public void clear() { - capturedTasks.clear(); + capturedEvents.clear(); } } diff --git a/server-common/src/main/java/io/a2a/server/tasks/BasePushNotificationSender.java b/server-common/src/main/java/io/a2a/server/tasks/BasePushNotificationSender.java index e400b7ba0..1ac6aba47 100644 --- a/server-common/src/main/java/io/a2a/server/tasks/BasePushNotificationSender.java +++ b/server-common/src/main/java/io/a2a/server/tasks/BasePushNotificationSender.java @@ -3,6 +3,15 @@ import static io.a2a.client.http.A2AHttpClient.APPLICATION_JSON; import static io.a2a.client.http.A2AHttpClient.CONTENT_TYPE; import static io.a2a.common.A2AHeaders.X_A2A_NOTIFICATION_TOKEN; +import static io.a2a.spec.Message.MESSAGE; +import static io.a2a.spec.Task.TASK; +import static io.a2a.spec.TaskArtifactUpdateEvent.ARTIFACT_UPDATE; +import static io.a2a.spec.TaskStatusUpdateEvent.STATUS_UPDATE; + +import io.a2a.spec.Message; +import io.a2a.spec.StreamingEventKind; +import io.a2a.spec.TaskArtifactUpdateEvent; +import io.a2a.spec.TaskStatusUpdateEvent; import jakarta.enterprise.context.ApplicationScoped; import jakarta.inject.Inject; @@ -42,34 +51,45 @@ public BasePushNotificationSender(PushNotificationConfigStore configStore, A2AHt } @Override - public void sendNotification(Task task) { - List pushConfigs = configStore.getInfo(task.getId()); + public void sendNotification(StreamingEventKind kind) { + String taskId = switch (kind.getKind()) { + case TASK -> ((Task) kind).getId(); + case MESSAGE -> ((Message)kind).getTaskId(); + case STATUS_UPDATE -> ((TaskStatusUpdateEvent)kind).getTaskId(); + case ARTIFACT_UPDATE -> ((TaskArtifactUpdateEvent)kind).getTaskId(); + default -> null; + }; + if (taskId == null) { + return; + } + + List pushConfigs = configStore.getInfo(taskId); if (pushConfigs == null || pushConfigs.isEmpty()) { return; } List> dispatchResults = pushConfigs .stream() - .map(pushConfig -> dispatch(task, pushConfig)) + .map(pushConfig -> dispatch(kind, pushConfig)) .toList(); CompletableFuture allFutures = CompletableFuture.allOf(dispatchResults.toArray(new CompletableFuture[0])); CompletableFuture dispatchResult = allFutures.thenApply(v -> dispatchResults.stream() .allMatch(CompletableFuture::join)); try { boolean allSent = dispatchResult.get(); - if (! allSent) { - LOGGER.warn("Some push notifications failed to send for taskId: " + task.getId()); + if (!allSent) { + LOGGER.warn("Some push notifications failed to send for taskId: " + taskId); } } catch (InterruptedException | ExecutionException e) { - LOGGER.warn("Some push notifications failed to send for taskId " + task.getId() + ": {}", e.getMessage(), e); + LOGGER.warn("Some push notifications failed to send for taskId " + taskId + ": {}", e.getMessage(), e); } } - private CompletableFuture dispatch(Task task, PushNotificationConfig pushInfo) { - return CompletableFuture.supplyAsync(() -> dispatchNotification(task, pushInfo)); + private CompletableFuture dispatch(StreamingEventKind kind, PushNotificationConfig pushInfo) { + return CompletableFuture.supplyAsync(() -> dispatchNotification(kind, pushInfo)); } - private boolean dispatchNotification(Task task, PushNotificationConfig pushInfo) { + private boolean dispatchNotification(StreamingEventKind kind, PushNotificationConfig pushInfo) { String url = pushInfo.url(); String token = pushInfo.token(); @@ -80,7 +100,7 @@ private boolean dispatchNotification(Task task, PushNotificationConfig pushInfo) String body; try { - body = Utils.OBJECT_MAPPER.writeValueAsString(task); + body = Utils.OBJECT_MAPPER.writeValueAsString(kind); } catch (JsonProcessingException e) { LOGGER.debug("Error writing value as string: {}", e.getMessage(), e); return false; diff --git a/server-common/src/main/java/io/a2a/server/tasks/PushNotificationSender.java b/server-common/src/main/java/io/a2a/server/tasks/PushNotificationSender.java index 81d577f46..c780d4b75 100644 --- a/server-common/src/main/java/io/a2a/server/tasks/PushNotificationSender.java +++ b/server-common/src/main/java/io/a2a/server/tasks/PushNotificationSender.java @@ -1,6 +1,6 @@ package io.a2a.server.tasks; -import io.a2a.spec.Task; +import io.a2a.spec.StreamingEventKind; /** * Interface for sending push notifications for tasks. @@ -8,8 +8,8 @@ public interface PushNotificationSender { /** - * Sends a push notification containing the latest task state. - * @param task the task + * Sends a push notification with a payload related to the task. + * @param kind the payload to push */ - void sendNotification(Task task); + void sendNotification(StreamingEventKind kind); } diff --git a/server-common/src/test/java/io/a2a/server/tasks/PushNotificationSenderTest.java b/server-common/src/test/java/io/a2a/server/tasks/PushNotificationSenderTest.java index c52a4cf20..c84605451 100644 --- a/server-common/src/test/java/io/a2a/server/tasks/PushNotificationSenderTest.java +++ b/server-common/src/test/java/io/a2a/server/tasks/PushNotificationSenderTest.java @@ -16,6 +16,10 @@ import java.util.concurrent.TimeUnit; import java.util.function.Consumer; +import io.a2a.spec.Message; +import io.a2a.spec.Part; +import io.a2a.spec.StreamingEventKind; +import io.a2a.spec.TextPart; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -67,6 +71,7 @@ class TestPostBuilder implements A2AHttpClient.PostBuilder { @Override public PostBuilder body(String body) { this.body = body; + System.out.println("body = " + body); return this; } @@ -80,6 +85,7 @@ public A2AHttpResponse post() throws IOException, InterruptedException { Task task = Utils.OBJECT_MAPPER.readValue(body, Task.TYPE_REFERENCE); tasks.add(task); urls.add(url); + System.out.println(requestHeaders); headers.add(new java.util.HashMap<>(requestHeaders)); return new A2AHttpResponse() { @@ -95,7 +101,7 @@ public boolean success() { @Override public String body() { - return ""; + return body; } }; } finally { @@ -316,4 +322,45 @@ public void testSendNotificationHttpError() { // Verify no tasks were successfully processed due to the error assertEquals(0, testHttpClient.tasks.size()); } + + @Test + public void testSendNotificationWithMessage() throws InterruptedException { + String taskId = "task_send_notification_with_message"; + Task taskData = createSampleTask(taskId, TaskState.COMPLETED); + PushNotificationConfig config = createSamplePushConfig("http://notify.me/here", "cfg1", "unique_token"); + + // Set up the configuration in the store + configStore.setInfo(taskId, config); + + // Set up latch to wait for async completion + testHttpClient.latch = new CountDownLatch(1); + + Message message = new Message.Builder() + .taskId(taskId) + .messageId("task_push_notification_message") + .parts(Collections.singletonList(new TextPart("Message for task " + taskId))) + .role(Message.Role.USER) + .build(); + sender.sendNotification(message); + + // Wait for the async operation to complete + assertTrue(testHttpClient.latch.await(5, TimeUnit.SECONDS), "HTTP call should complete within 5 seconds"); + + // Verify the task was sent via HTTP + assertEquals(1, testHttpClient.tasks.size()); + Task sentTask = testHttpClient.tasks.get(0); + assertEquals(taskData.getId(), sentTask.getId()); + + // Verify that the X-A2A-Notification-Token header is sent with the correct token + assertEquals(1, testHttpClient.headers.size()); + Map sentHeaders = testHttpClient.headers.get(0); + assertEquals(2, sentHeaders.size()); + assertTrue(sentHeaders.containsKey(A2AHeaders.X_A2A_NOTIFICATION_TOKEN)); + assertEquals(config.token(), sentHeaders.get(A2AHeaders.X_A2A_NOTIFICATION_TOKEN)); + // Content-Type header should always be present + assertTrue(sentHeaders.containsKey(CONTENT_TYPE)); + assertEquals(APPLICATION_JSON, sentHeaders.get(CONTENT_TYPE)); + + } + }