Fix: Making unmarshalResponse type-safe

Signed-off-by: Emmanuel Hugonnet <[email protected]>

Author: ehsavoie

client/transport/jsonrpc/src/main/java/io/a2a/client/transport/jsonrpc/JSONRPCTransport.java

@@ -25,6 +25,8 @@
 import io.a2a.spec.CancelTaskResponse;
 
 import io.a2a.spec.DeleteTaskPushNotificationConfigParams;
+import io.a2a.spec.DeleteTaskPushNotificationConfigRequest;
+import io.a2a.spec.DeleteTaskPushNotificationConfigResponse;
 import io.a2a.spec.EventKind;
 import io.a2a.spec.GetAuthenticatedExtendedCardRequest;
 import io.a2a.spec.GetAuthenticatedExtendedCardResponse;
@@ -43,7 +45,6 @@
 import io.a2a.spec.ListTasksRequest;
 import io.a2a.spec.ListTasksResponse;
 import io.a2a.spec.ListTasksResult;
-import io.a2a.spec.DeleteTaskPushNotificationConfigRequest;
 import io.a2a.spec.MessageSendParams;
 import io.a2a.spec.SendMessageResponse;
 import io.a2a.spec.SendStreamingMessageRequest;
@@ -96,7 +97,7 @@ public EventKind sendMessage(MessageSendParams request, @Nullable ClientCallCont
 
         try {
             String httpResponseBody = sendPostRequest(payloadAndHeaders,io.a2a.spec.SendMessageRequest.METHOD);
-            SendMessageResponse response = (SendMessageResponse) unmarshalResponse(httpResponseBody, io.a2a.spec.SendMessageRequest.METHOD);
+            SendMessageResponse response = unmarshalResponse(httpResponseBody, io.a2a.spec.SendMessageRequest.METHOD);
             return response.getResult();
         } catch (A2AClientException e) {
             throw e;
@@ -140,7 +141,7 @@ public Task getTask(TaskQueryParams request, @Nullable ClientCallContext context
 
         try {
             String httpResponseBody = sendPostRequest(payloadAndHeaders, io.a2a.spec.GetTaskRequest.METHOD);
-            GetTaskResponse response = (GetTaskResponse) unmarshalResponse(httpResponseBody, io.a2a.spec.GetTaskRequest.METHOD);
+            GetTaskResponse response = unmarshalResponse(httpResponseBody, io.a2a.spec.GetTaskRequest.METHOD);
             return response.getResult();
         } catch (A2AClientException e) {
             throw e;
@@ -157,7 +158,7 @@ public Task cancelTask(TaskIdParams request, @Nullable ClientCallContext context
 
         try {
             String httpResponseBody = sendPostRequest(payloadAndHeaders, CancelTaskRequest.METHOD);
-            CancelTaskResponse response = (CancelTaskResponse) unmarshalResponse(httpResponseBody, CancelTaskRequest.METHOD);
+            CancelTaskResponse response = unmarshalResponse(httpResponseBody, CancelTaskRequest.METHOD);
             return response.getResult();
         } catch (A2AClientException e) {
             throw e;
@@ -173,7 +174,7 @@ public ListTasksResult listTasks(ListTasksParams request, @Nullable ClientCallCo
                 agentCard, context);
         try {
             String httpResponseBody = sendPostRequest(payloadAndHeaders, ListTasksRequest.METHOD);
-            ListTasksResponse response = (ListTasksResponse) unmarshalResponse(httpResponseBody, ListTasksRequest.METHOD);
+            ListTasksResponse response = unmarshalResponse(httpResponseBody, ListTasksRequest.METHOD);
             return response.getResult();
         } catch (IOException | InterruptedException e) {
             throw new A2AClientException("Failed to list tasks: " + e, e);
@@ -189,7 +190,7 @@ public TaskPushNotificationConfig setTaskPushNotificationConfiguration(TaskPushN
 
         try {
             String httpResponseBody = sendPostRequest(payloadAndHeaders, SetTaskPushNotificationConfigRequest.METHOD);
-            SetTaskPushNotificationConfigResponse response = (SetTaskPushNotificationConfigResponse) unmarshalResponse(httpResponseBody,
+            SetTaskPushNotificationConfigResponse response = unmarshalResponse(httpResponseBody,
                     SetTaskPushNotificationConfigRequest.METHOD);
             return response.getResult();
         } catch (A2AClientException e) {
@@ -208,7 +209,7 @@ public TaskPushNotificationConfig getTaskPushNotificationConfiguration(GetTaskPu
 
         try {
             String httpResponseBody = sendPostRequest(payloadAndHeaders,GetTaskPushNotificationConfigRequest.METHOD);
-            GetTaskPushNotificationConfigResponse response = (GetTaskPushNotificationConfigResponse) unmarshalResponse(httpResponseBody,
+            GetTaskPushNotificationConfigResponse response = unmarshalResponse(httpResponseBody,
                     GetTaskPushNotificationConfigRequest.METHOD);
             return response.getResult();
         } catch (A2AClientException e) {
@@ -228,7 +229,7 @@ public List<TaskPushNotificationConfig> listTaskPushNotificationConfigurations(
 
         try {
             String httpResponseBody = sendPostRequest(payloadAndHeaders, ListTaskPushNotificationConfigRequest.METHOD);
-            ListTaskPushNotificationConfigResponse response = (ListTaskPushNotificationConfigResponse) unmarshalResponse(httpResponseBody,
+            ListTaskPushNotificationConfigResponse response = unmarshalResponse(httpResponseBody,
                     ListTaskPushNotificationConfigRequest.METHOD);
             return response.getResult();
         } catch (A2AClientException e) {
@@ -247,7 +248,8 @@ public void deleteTaskPushNotificationConfigurations(DeleteTaskPushNotificationC
 
         try {
             String httpResponseBody = sendPostRequest(payloadAndHeaders,DeleteTaskPushNotificationConfigRequest.METHOD);
-            unmarshalResponse(httpResponseBody, DeleteTaskPushNotificationConfigRequest.METHOD);
+            DeleteTaskPushNotificationConfigResponse response = unmarshalResponse(httpResponseBody, DeleteTaskPushNotificationConfigRequest.METHOD);
+            // Response validated (no error), but no result to return
         } catch (A2AClientException e) {
             throw e;
         } catch (IOException | InterruptedException e) {
@@ -306,7 +308,7 @@ public AgentCard getAgentCard(@Nullable ClientCallContext context) throws A2ACli
 
             try {
                 String httpResponseBody = sendPostRequest(payloadAndHeaders,GetAuthenticatedExtendedCardRequest.METHOD);
-                GetAuthenticatedExtendedCardResponse response = (GetAuthenticatedExtendedCardResponse) unmarshalResponse(httpResponseBody,
+                GetAuthenticatedExtendedCardResponse response = unmarshalResponse(httpResponseBody,
                         GetAuthenticatedExtendedCardRequest.METHOD);
                 agentCard = response.getResult();
                 needsExtendedCard = false;
@@ -360,14 +362,30 @@ private A2AHttpClient.PostBuilder createPostBuilder(PayloadAndHeaders payloadAnd
         return postBuilder;
     }
 
-    private JSONRPCResponse<?> unmarshalResponse(String response, String method)
+    /**
+     * Unmarshals a JSON-RPC response string into a type-safe response object.
+     * <p>
+     * This method parses the JSON-RPC response body and returns the appropriate
+     * response type based on the method parameter. If the response contains an error,
+     * an A2AClientException is thrown.
+     *
+     * @param <T> the expected response type, must extend JSONRPCResponse
+     * @param response the JSON-RPC response body as a string
+     * @param method the method name used to determine the response type
+     * @return the parsed response object of type T
+     * @throws A2AClientException if the response contains an error or parsing fails
+     * @throws JsonProcessingException if the JSON cannot be processed
+     */
+    @SuppressWarnings("unchecked")
+    private <T extends JSONRPCResponse<?>> T unmarshalResponse(String response, String method)
             throws A2AClientException, JsonProcessingException {
         JSONRPCResponse<?> value = JSONRPCUtils.parseResponseBody(response, method);
         JSONRPCError error = value.getError();
         if (error != null) {
             throw new A2AClientException(error.getMessage() + (error.getData() != null ? ": " + error.getData() : ""), error);
         }
-        return value;
+        // Safe cast: JSONRPCUtils.parseResponseBody returns the correct concrete type based on method
+        return (T) value;
     }
 
     private @Nullable Map<String, String> getHttpHeaders(@Nullable ClientCallContext context) {