feat: Support multiple payload types in push notifications

Expand PushNotificationSender to support all StreamingEventKind payload types
as defined in the A2A specification, not just Task objects.

Fixes: #490
Signed-off-by: Jeff Mesnil <[email protected]>

Author: jmesnil

extras/push-notification-config-store-database-jpa/src/test/java/io/a2a/extras/pushnotificationconfigstore/database/jpa/JpaDatabasePushNotificationConfigStoreIntegrationTest.java

@@ -25,6 +25,7 @@
 import io.a2a.spec.GetTaskPushNotificationConfigParams;
 import io.a2a.spec.Message;
 import io.a2a.spec.PushNotificationConfig;
+import io.a2a.spec.StreamingEventKind;
 import io.a2a.spec.Task;
 import io.a2a.spec.TaskPushNotificationConfig;
 import io.a2a.spec.TextPart;
@@ -88,9 +89,9 @@ public void testDirectNotificationTrigger() {
         mockPushNotificationSender.sendNotification(testTask);
 
         // Verify it was captured
-        Queue<Task> captured = mockPushNotificationSender.getCapturedTasks();
+        Queue<StreamingEventKind> captured = mockPushNotificationSender.getCapturedEvents();
         assertEquals(1, captured.size());
-        assertEquals("direct-test-task", captured.peek().getId());
+        assertEquals("direct-test-task", ((Task)captured.peek()).getId());
     }
 
     @Test
@@ -151,7 +152,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,10 +162,12 @@ public void testJpaDatabasePushNotificationConfigStoreIntegration() throws Excep
         assertTrue(notificationReceived, "Timeout waiting for push notification.");
 
         // Step 6: Verify the captured notification
-        Queue<Task> capturedTasks = mockPushNotificationSender.getCapturedTasks();
+        Queue<StreamingEventKind> capturedTasks = mockPushNotificationSender.getCapturedEvents();
 
         // Verify the notification contains the correct task with artifacts
         Task notifiedTaskWithArtifact = capturedTasks.stream()
+            .filter(e -> Task.TASK.equals(e.getKind()))
+            .map(e -> (Task)e)
             .filter(t -> taskId.equals(t.getId()) && t.getArtifacts() != null && t.getArtifacts().size() > 0)
             .findFirst()
             .orElse(null);

extras/push-notification-config-store-database-jpa/src/test/java/io/a2a/extras/pushnotificationconfigstore/database/jpa/MockPushNotificationSender.java

@@ -8,7 +8,7 @@
 import jakarta.enterprise.inject.Alternative;
 
 import io.a2a.server.tasks.PushNotificationSender;
-import io.a2a.spec.Task;
+import io.a2a.spec.StreamingEventKind;
 
 /**
  * Mock implementation of PushNotificationSender for integration testing.
@@ -19,18 +19,18 @@
 @Priority(100)
 public class MockPushNotificationSender implements PushNotificationSender {
 
-    private final Queue<Task> capturedTasks = new ConcurrentLinkedQueue<>();
+    private final Queue<StreamingEventKind> capturedEvents = new ConcurrentLinkedQueue<>();
 
     @Override
-    public void sendNotification(Task task) {
-        capturedTasks.add(task);
+    public void sendNotification(StreamingEventKind kind) {
+        capturedEvents.add(kind);
     }
 
-    public Queue<Task> getCapturedTasks() {
-        return capturedTasks;
+    public Queue<StreamingEventKind> getCapturedEvents() {
+        return capturedEvents;
     }
 
     public void clear() {
-        capturedTasks.clear();
+        capturedEvents.clear();
     }
 }

server-common/src/main/java/io/a2a/server/tasks/BasePushNotificationSender.java

@@ -3,6 +3,10 @@
 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 java.io.IOException;
 import java.util.List;
@@ -18,8 +22,12 @@
 
 import io.a2a.client.http.A2AHttpClient;
 import io.a2a.client.http.JdkA2AHttpClient;
+import io.a2a.spec.Message;
 import io.a2a.spec.PushNotificationConfig;
+import io.a2a.spec.StreamingEventKind;
 import io.a2a.spec.Task;
+import io.a2a.spec.TaskArtifactUpdateEvent;
+import io.a2a.spec.TaskStatusUpdateEvent;
 import io.a2a.util.Utils;
 
 @ApplicationScoped
@@ -42,34 +50,44 @@ public BasePushNotificationSender(PushNotificationConfigStore configStore, A2AHt
     }
 
     @Override
-    public void sendNotification(Task task) {
-        List<PushNotificationConfig> 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<PushNotificationConfig> pushConfigs = configStore.getInfo(taskId);
         if (pushConfigs == null || pushConfigs.isEmpty()) {
             return;
         }
 
         List<CompletableFuture<Boolean>> dispatchResults = pushConfigs
                 .stream()
-                .map(pushConfig -> dispatch(task, pushConfig))
+                .map(pushConfig -> dispatch(kind, pushConfig))
                 .toList();
         CompletableFuture<Void> allFutures = CompletableFuture.allOf(dispatchResults.toArray(new CompletableFuture[0]));
         CompletableFuture<Boolean> 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());
+                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<Boolean> dispatch(Task task, PushNotificationConfig pushInfo) {
-        return CompletableFuture.supplyAsync(() -> dispatchNotification(task, pushInfo));
+    private CompletableFuture<Boolean> 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 +98,7 @@ private boolean dispatchNotification(Task task, PushNotificationConfig pushInfo)
 
         String body;
         try {
-            body = Utils.marshalFrom(task);
+            body = Utils.marshalFrom(kind);
         } catch (JsonProcessingException e) {
             LOGGER.debug("Error writing value as string: {}", e.getMessage(), e);
             return false;

server-common/src/main/java/io/a2a/server/tasks/PushNotificationSender.java

@@ -1,15 +1,15 @@
 package io.a2a.server.tasks;
 
-import io.a2a.spec.Task;
+import io.a2a.spec.StreamingEventKind;
 
 /**
  * Interface for sending push notifications for tasks.
  */
 public interface PushNotificationSender {
 
     /**
-     * Sends a push notification containing the latest task state.
-     * @param task the task
+     * Sends a push notification containing payload about a task.
+     * @param kind the payload to push
      */
-    void sendNotification(Task task);
+    void sendNotification(StreamingEventKind kind);
 }

server-common/src/test/java/io/a2a/server/requesthandlers/AbstractA2ARequestHandlerTest.java

@@ -16,7 +16,6 @@
 
 import jakarta.enterprise.context.Dependent;
 
-import com.fasterxml.jackson.databind.JsonNode;
 import io.quarkus.arc.profile.IfBuildProfile;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.Assertions;
@@ -200,10 +199,7 @@ public PostBuilder body(String body) {
 
             @Override
             public A2AHttpResponse post() throws IOException, InterruptedException {
-                JsonNode root = Utils.OBJECT_MAPPER.readTree(body);
-                // This will need to be updated for #490 to unmarshall based on the kind of payload
-                JsonNode taskNode = root.elements().next();
-                Task task = Utils.OBJECT_MAPPER.treeToValue(taskNode, Task.TYPE_REFERENCE);
+                Task task = Utils.unmarshalStreamingEventKindFrom(body);
                 tasks.add(task);
                 try {
                     return new A2AHttpResponse() {