fix: Ensure onSetTaskPushNotificationConfig returns the updated pushNotificationConfig

Author: fjuma

server-common/src/main/java/io/a2a/server/requesthandlers/DefaultRequestHandler.java

@@ -273,9 +273,8 @@ public TaskPushNotificationConfig onSetTaskPushNotificationConfig(
             throw new TaskNotFoundError();
         }
 
-        pushConfigStore.setInfo(params.taskId(), params.pushNotificationConfig());
-
-        return params;
+        PushNotificationConfig pushNotificationConfig = pushConfigStore.setInfo(params.taskId(), params.pushNotificationConfig());
+        return new TaskPushNotificationConfig(params.taskId(), pushNotificationConfig);
     }
 
     @Override

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

@@ -27,7 +27,7 @@ public InMemoryPushNotificationConfigStore() {
     }
 
     @Override
-    public void setInfo(String taskId, PushNotificationConfig notificationConfig) {
+    public PushNotificationConfig setInfo(String taskId, PushNotificationConfig notificationConfig) {
         List<PushNotificationConfig> notificationConfigList = pushNotificationInfos.getOrDefault(taskId, new ArrayList<>());
         PushNotificationConfig.Builder builder = new PushNotificationConfig.Builder(notificationConfig);
         if (notificationConfig.id() == null) {
@@ -45,6 +45,7 @@ public void setInfo(String taskId, PushNotificationConfig notificationConfig) {
         }
         notificationConfigList.add(notificationConfig);
         pushNotificationInfos.put(taskId, notificationConfigList);
+        return notificationConfig;
     }
 
     @Override

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

@@ -13,8 +13,9 @@ public interface PushNotificationConfigStore {
      * Sets or updates the push notification configuration for a task.
      * @param taskId the task ID
      * @param notificationConfig the push notification configuration
+     * @return the potentially updated push notification configuration
      */
-    void setInfo(String taskId, PushNotificationConfig notificationConfig);
+    PushNotificationConfig setInfo(String taskId, PushNotificationConfig notificationConfig);
 
     /**
      * Retrieves the push notification configuration for a task.

spec-grpc/src/main/java/io/a2a/grpc/utils/ProtoUtils.java

@@ -687,24 +687,38 @@ public static MessageSendParams messageSendParams(io.a2a.grpc.SendMessageRequest
             return builder.build();
         }
 
-        public static TaskPushNotificationConfig taskPushNotificationConfig(io.a2a.grpc.CreateTaskPushNotificationConfigRequest request) {
-            return taskPushNotificationConfig(request.getConfig());
+        public static TaskPushNotificationConfig taskPushNotificationConfig(io.a2a.grpc.CreateTaskPushNotificationConfigRequestOrBuilder request) {
+            return taskPushNotificationConfig(request.getConfig(), true);
         }
 
-        public static TaskPushNotificationConfig taskPushNotificationConfig(io.a2a.grpc.TaskPushNotificationConfig config) {
+        public static TaskPushNotificationConfig taskPushNotificationConfig(io.a2a.grpc.TaskPushNotificationConfigOrBuilder config) {
+            return taskPushNotificationConfig(config, false);
+        }
+
+        private static TaskPushNotificationConfig taskPushNotificationConfig(io.a2a.grpc.TaskPushNotificationConfigOrBuilder config, boolean create) {
             String name = config.getName(); // "tasks/{id}/pushNotificationConfigs/{push_id}"
             String[] parts = name.split("/");
-            if (parts.length < 4) {
-                throw new IllegalArgumentException("Invalid name format for TaskPushNotificationConfig: " + name);
+            String configId = "";
+            if (create) {
+                if (parts.length < 3) {
+                    throw new IllegalArgumentException("Invalid name format for TaskPushNotificationConfig: " + name);
+                }
+                if (parts.length == 4) {
+                    configId = parts[3];
+                }
+            } else {
+                if (parts.length < 4) {
+                    throw new IllegalArgumentException("Invalid name format for TaskPushNotificationConfig: " + name);
+                }
+                configId = parts[3];
             }
             String taskId = parts[1];
-            String configId = parts[3];
             PushNotificationConfig pnc = pushNotification(config.getPushNotificationConfig(), configId);
             return new TaskPushNotificationConfig(taskId, pnc);
         }
 
         public static GetTaskPushNotificationConfigParams getTaskPushNotificationConfigParams(io.a2a.grpc.GetTaskPushNotificationConfigRequest request) {
-            String name = request.getName(); // "tasks/{id}/pushNotificationConfigs/{push_id}" or /tasks/{id}
+            String name = request.getName(); // "tasks/{id}/pushNotificationConfigs/{push_id}"
             String[] parts = name.split("/");
             String taskId = parts[1];
             String configId;
@@ -775,9 +789,6 @@ private static PushNotificationConfig pushNotification(io.a2a.grpc.PushNotificat
         }
 
         private static PushNotificationConfig pushNotification(io.a2a.grpc.PushNotificationConfig pushNotification) {
-            /*if (pushNotification == null) {
-                return null;
-            }*/
             return pushNotification(pushNotification, pushNotification.getId());
         }
 

transport/jsonrpc/src/test/java/io/a2a/transport/jsonrpc/handler/JSONRPCHandlerTest.java

@@ -557,7 +557,10 @@ public void testSetPushNotificationConfigSuccess() {
                         MINIMAL_TASK.getId(), new PushNotificationConfig.Builder().url("http://example.com").build());
         SetTaskPushNotificationConfigRequest request = new SetTaskPushNotificationConfigRequest("1", taskPushConfig);
         SetTaskPushNotificationConfigResponse response = handler.setPushNotificationConfig(request, callContext);
-        Assertions.assertSame(taskPushConfig, response.getResult());
+        TaskPushNotificationConfig taskPushConfigResult =
+                new TaskPushNotificationConfig(
+                        MINIMAL_TASK.getId(), new PushNotificationConfig.Builder().url("http://example.com").id(MINIMAL_TASK.getId()).build());
+        Assertions.assertEquals(taskPushConfigResult, response.getResult());
     }
 
     @Test