feat: Add support for notification tokens in BasePushNotificationSender (a2aproject#308)

fjuma · web-flow · commit e20a3120fbe3 · 2025-09-29T17:38:29.000+01:00
# Description Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly: - [x] Follow the [`CONTRIBUTING` Guide](../CONTRIBUTING.md). - [x] Make your Pull Request title in the <https://www.conventionalcommits.org/> specification. - Important Prefixes for [release-please](https://github.com/googleapis/release-please): - `fix:` which represents bug fixes, and correlates to a [SemVer](https://semver.org/) patch. - `feat:` represents a new feature, and correlates to a SemVer minor. - `feat!:`, or `fix!:`, `refactor!:`, etc., which represent a breaking change (indicated by the `!`) and will result in a SemVer major. - [x] Ensure the tests pass - [x] Appropriate READMEs were updated (if necessary) Fixes a2aproject#307 🦕
diff --git a/common/src/main/java/io/a2a/common/A2AHeaders.java b/common/src/main/java/io/a2a/common/A2AHeaders.java
@@ -10,6 +10,11 @@ public final class A2AHeaders {
      * Used to communicate which extensions are requested by the client.
      */
     public static final String X_A2A_EXTENSIONS = "X-A2A-Extensions";
+
+    /**
+     * HTTP header name for a push notification token.
+     */
+    public static final String X_A2A_NOTIFICATION_TOKEN = "X-A2A-Notification-Token";
     
     private A2AHeaders() {
         // Utility class
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
@@ -1,5 +1,6 @@
 package io.a2a.server.tasks;
 
+import static io.a2a.common.A2AHeaders.X_A2A_NOTIFICATION_TOKEN;
 import jakarta.enterprise.context.ApplicationScoped;
 import jakarta.inject.Inject;
 
@@ -68,10 +69,12 @@ private CompletableFuture<Boolean> dispatch(Task task, PushNotificationConfig pu
 
     private boolean dispatchNotification(Task task, PushNotificationConfig pushInfo) {
         String url = pushInfo.url();
+        String token = pushInfo.token();
 
-        // TODO: Implement authentication and token header support
-        // The Python implementation adds X-A2A-Notification-Token header when pushInfo.token is present
-        // See: https://github.com/a2aproject/a2a-python/blob/main/src/a2a/server/tasks/base_push_notification_sender.py#L55-57
+        A2AHttpClient.PostBuilder postBuilder = httpClient.createPost();
+        if (token != null && !token.isBlank()) {
+            postBuilder.addHeader(X_A2A_NOTIFICATION_TOKEN, token);
+        }
 
         String body;
         try {
@@ -85,7 +88,7 @@ private boolean dispatchNotification(Task task, PushNotificationConfig pushInfo)
         }
 
         try {
-            httpClient.createPost()
+            postBuilder
                     .url(url)
                     .body(body)
                     .post();
diff --git a/server-common/src/test/java/io/a2a/server/tasks/InMemoryPushNotificationConfigStoreTest.java b/server-common/src/test/java/io/a2a/server/tasks/InMemoryPushNotificationConfigStoreTest.java
@@ -14,13 +14,13 @@
 import java.util.List;
 
 import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Disabled;
 import org.junit.jupiter.api.Test;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
 import io.a2a.client.http.A2AHttpClient;
 import io.a2a.client.http.A2AHttpResponse;
+import io.a2a.common.A2AHeaders;
 import io.a2a.spec.PushNotificationConfig;
 import io.a2a.spec.Task;
 import io.a2a.spec.TaskState;
@@ -47,6 +47,29 @@ public void setUp() {
         notificationSender = new BasePushNotificationSender(configStore, mockHttpClient);
     }
 
+    private void setupBasicMockHttpResponse() throws Exception {
+        when(mockHttpClient.createPost()).thenReturn(mockPostBuilder);
+        when(mockPostBuilder.url(any(String.class))).thenReturn(mockPostBuilder);
+        when(mockPostBuilder.body(any(String.class))).thenReturn(mockPostBuilder);
+        when(mockPostBuilder.post()).thenReturn(mockHttpResponse);
+        when(mockHttpResponse.success()).thenReturn(true);
+    }
+
+    private void verifyHttpCallWithoutToken(PushNotificationConfig config, Task task, String expectedToken) throws Exception {
+        ArgumentCaptor<String> bodyCaptor = ArgumentCaptor.forClass(String.class);
+        verify(mockHttpClient).createPost();
+        verify(mockPostBuilder).url(config.url());
+        verify(mockPostBuilder).body(bodyCaptor.capture());
+        verify(mockPostBuilder).post();
+        // Verify that addHeader was never called for authentication token
+        verify(mockPostBuilder, never()).addHeader(A2AHeaders.X_A2A_NOTIFICATION_TOKEN, expectedToken);
+        
+        // Verify the request body contains the task data
+        String sentBody = bodyCaptor.getValue();
+        assertTrue(sentBody.contains(task.getId()));
+        assertTrue(sentBody.contains(task.getStatus().state().asString()));
+    }
+
     private Task createSampleTask(String taskId, TaskState state) {
         return new Task.Builder()
                 .id(taskId)
@@ -229,7 +252,6 @@ public void testSendNotificationSuccess() throws Exception {
     }
 
     @Test
-    @Disabled("Token authentication is not yet implemented in BasePushNotificationSender (TODO auth)")
     public void testSendNotificationWithToken() throws Exception {
         String taskId = "task_send_with_token";
         Task task = createSampleTask(taskId, TaskState.COMPLETED);
@@ -240,21 +262,19 @@ public void testSendNotificationWithToken() throws Exception {
         when(mockHttpClient.createPost()).thenReturn(mockPostBuilder);
         when(mockPostBuilder.url(any(String.class))).thenReturn(mockPostBuilder);
         when(mockPostBuilder.body(any(String.class))).thenReturn(mockPostBuilder);
+        when(mockPostBuilder.addHeader(any(String.class), any(String.class))).thenReturn(mockPostBuilder);
         when(mockPostBuilder.post()).thenReturn(mockHttpResponse);
         when(mockHttpResponse.success()).thenReturn(true);
 
         notificationSender.sendNotification(task);
 
-        // TODO: Once token authentication is implemented, verify that:
-        // 1. The token is included in request headers (e.g., X-A2A-Notification-Token)
-        // 2. The HTTP client is called with proper authentication
-        // 3. The token from the config is actually used
-        
-        // For now, just verify basic HTTP client interaction
+        // Verify HTTP client was called with proper authentication
         ArgumentCaptor<String> bodyCaptor = ArgumentCaptor.forClass(String.class);
         verify(mockHttpClient).createPost();
         verify(mockPostBuilder).url(config.url());
         verify(mockPostBuilder).body(bodyCaptor.capture());
+        // Verify that the token is included in request headers as X-A2A-Notification-Token
+        verify(mockPostBuilder).addHeader(A2AHeaders.X_A2A_NOTIFICATION_TOKEN, config.token());
         verify(mockPostBuilder).post();
         
         // Verify the request body contains the task data
@@ -274,6 +294,30 @@ public void testSendNotificationNoConfig() throws Exception {
         verify(mockHttpClient, never()).createPost();
     }
 
+    @Test
+    public void testSendNotificationWithEmptyToken() throws Exception {
+        String taskId = "task_send_empty_token";
+        Task task = createSampleTask(taskId, TaskState.COMPLETED);
+        PushNotificationConfig config = createSamplePushConfig("http://notify.me/here", "cfg1", "");
+        configStore.setInfo(taskId, config);
+
+        setupBasicMockHttpResponse();
+        notificationSender.sendNotification(task);
+        verifyHttpCallWithoutToken(config, task, "");
+    }
+
+    @Test
+    public void testSendNotificationWithBlankToken() throws Exception {
+        String taskId = "task_send_blank_token";
+        Task task = createSampleTask(taskId, TaskState.COMPLETED);
+        PushNotificationConfig config = createSamplePushConfig("http://notify.me/here", "cfg1", "   ");
+        configStore.setInfo(taskId, config);
+
+        setupBasicMockHttpResponse();
+        notificationSender.sendNotification(task);
+        verifyHttpCallWithoutToken(config, task, "   ");
+    }
+
     @Test
     public void testMultipleConfigsForSameTask() {
         String taskId = "task_multiple";
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
@@ -18,6 +18,7 @@
 
 import io.a2a.client.http.A2AHttpClient;
 import io.a2a.client.http.A2AHttpResponse;
+import io.a2a.common.A2AHeaders;
 import io.a2a.util.Utils;
 import io.a2a.spec.PushNotificationConfig;
 import io.a2a.spec.Task;
@@ -133,6 +134,33 @@ public void setUp() {
         sender = new BasePushNotificationSender(configStore, testHttpClient);
     }
 
+    private void testSendNotificationWithInvalidToken(String token, String testName) throws InterruptedException {
+        String taskId = testName;
+        Task taskData = createSampleTask(taskId, TaskState.COMPLETED);
+        PushNotificationConfig config = createSamplePushConfig("http://notify.me/here", "cfg1", 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);
+
+        sender.sendNotification(taskData);
+
+        // 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 no authentication header was sent (invalid token should not add header)
+        assertEquals(1, testHttpClient.headers.size());
+        Map<String, String> sentHeaders = testHttpClient.headers.get(0);
+        assertTrue(sentHeaders.isEmpty(), "No headers should be sent when token is invalid");
+    }
+
     private Task createSampleTask(String taskId, TaskState state) {
         return new Task.Builder()
                 .id(taskId)
@@ -196,8 +224,11 @@ public void testSendNotificationWithTokenSuccess() throws InterruptedException {
         Task sentTask = testHttpClient.tasks.get(0);
         assertEquals(taskData.getId(), sentTask.getId());
         
-        // TODO: When authentication is implemented in BasePushNotificationSender, verify that the
-        // X-A2A-Notification-Token header is sent.
+        // Verify that the X-A2A-Notification-Token header is sent with the correct token
+        assertEquals(1, testHttpClient.headers.size());
+        Map<String, String> sentHeaders = testHttpClient.headers.get(0);
+        assertTrue(sentHeaders.containsKey(A2AHeaders.X_A2A_NOTIFICATION_TOKEN));
+        assertEquals(config.token(), sentHeaders.get(A2AHeaders.X_A2A_NOTIFICATION_TOKEN));
     }
 
     @Test
@@ -212,6 +243,16 @@ public void testSendNotificationNoConfig() {
         assertEquals(0, testHttpClient.tasks.size());
     }
 
+    @Test
+    public void testSendNotificationWithEmptyToken() throws InterruptedException {
+        testSendNotificationWithInvalidToken("", "task_send_empty_token");
+    }
+
+    @Test
+    public void testSendNotificationWithBlankToken() throws InterruptedException {
+        testSendNotificationWithInvalidToken("   ", "task_send_blank_token");
+    }
+
     @Test
     public void testSendNotificationMultipleConfigs() throws InterruptedException {
         String taskId = "task_multiple_configs";
