refactor: Refactor SSEEventListener to support accepting JSONRPCResponse as a parameter.

Signed-off-by: Sun Yuhan <sunyuhan1998@users.noreply.github.com>

Author: sunyuhan1998

client/src/main/java/io/a2a/client/A2AClient.java

@@ -413,15 +413,7 @@ public void sendStreamingMessage(String requestId, MessageSendParams messageSend
         SendStreamingMessageRequest sendStreamingMessageRequest = sendStreamingMessageRequestBuilder.build();
         try {
             transport.sendMessageStreaming(sendStreamingMessageRequest, agentUrl, SEND_MESSAGE_RESPONSE_REFERENCE,
-                    response -> {
-                        String msg;
-                        try {
-                            msg = OBJECT_MAPPER.writeValueAsString(response);
-                        } catch (JsonProcessingException e) {
-                            throw new RuntimeException(e);
-                        }
-                        sseEventListener.onMessage(msg, ref.get());
-                    },
+                    response -> sseEventListener.onMessage(response, ref.get()),
                     throwable -> sseEventListener.onError(throwable, ref.get()),
                     () -> {
                         // We don't need to do anything special on completion
@@ -478,15 +470,7 @@ public void resubscribeToTask(String requestId, TaskIdParams taskIdParams, Consu
         TaskResubscriptionRequest taskResubscriptionRequest = taskResubscriptionRequestBuilder.build();
         try {
             transport.sendMessageStreaming(taskResubscriptionRequest, agentUrl, GET_TASK_RESPONSE_REFERENCE,
-                    response -> {
-                        String msg;
-                        try {
-                            msg = OBJECT_MAPPER.writeValueAsString(response);
-                        } catch (JsonProcessingException e) {
-                            throw new RuntimeException(e);
-                        }
-                        sseEventListener.onMessage(msg, ref.get());
-                    },
+                    response -> sseEventListener.onMessage(response, ref.get()),
                     throwable -> sseEventListener.onError(throwable, ref.get()),
                     () -> {
                         // We don't need to do anything special on completion

client/src/main/java/io/a2a/client/sse/SSEEventListener.java

@@ -7,8 +7,9 @@
 import java.util.logging.Logger;
 
 import com.fasterxml.jackson.core.JsonProcessingException;
-import com.fasterxml.jackson.databind.JsonNode;
 import io.a2a.spec.JSONRPCError;
+import io.a2a.spec.JSONRPCResponse;
+import io.a2a.spec.SendStreamingMessageResponse;
 import io.a2a.spec.StreamingEventKind;
 import io.a2a.spec.TaskStatusUpdateEvent;
 
@@ -26,36 +27,34 @@ public SSEEventListener(Consumer<StreamingEventKind> eventHandler, Consumer<JSON
 
     public void onMessage(String message, Future<Void> completableFuture) {
         try {
-            handleMessage(OBJECT_MAPPER.readTree(message),completableFuture);
+            SendStreamingMessageResponse sendStreamingMessageResponse = OBJECT_MAPPER.readValue(message, SendStreamingMessageResponse.class);
+            handleMessage(sendStreamingMessageResponse,completableFuture);
         } catch (JsonProcessingException e) {
             log.warning("Failed to parse JSON message: " + message);
         }
     }
 
+    public void onMessage(JSONRPCResponse<?> response, Future<Void> completableFuture) {
+        handleMessage(response,completableFuture);
+    }
+
     public void onError(Throwable throwable, Future<Void> future) {
         failureHandler.run();
         future.cancel(true); // close SSE channel
     }
 
-    private void handleMessage(JsonNode jsonNode, Future<Void> future) {
-        try {
-            if (jsonNode.has("error")) {
-                JSONRPCError error = OBJECT_MAPPER.treeToValue(jsonNode.get("error"), JSONRPCError.class);
-                errorHandler.accept(error);
-            } else if (jsonNode.has("result")) {
-                // result can be a Task, Message, TaskStatusUpdateEvent, or TaskArtifactUpdateEvent
-                JsonNode result = jsonNode.path("result");
-                StreamingEventKind event = OBJECT_MAPPER.treeToValue(result, StreamingEventKind.class);
-                eventHandler.accept(event);
-                if (event instanceof TaskStatusUpdateEvent && ((TaskStatusUpdateEvent) event).isFinal()) {
-                    future.cancel(true); // close SSE channel
-                }
-            } else {
-                throw new IllegalArgumentException("Unknown message type");
+    private void handleMessage(JSONRPCResponse<?> response, Future<Void> future) {
+        if (null != response.getError()) {
+            errorHandler.accept(response.getError());
+        } else if (null != response.getResult()) {
+            // result can be a Task, Message, TaskStatusUpdateEvent, or TaskArtifactUpdateEvent
+            StreamingEventKind event = (StreamingEventKind) response.getResult();
+            eventHandler.accept(event);
+            if (event instanceof TaskStatusUpdateEvent && ((TaskStatusUpdateEvent) event).isFinal()) {
+                future.cancel(true); // close SSE channel
             }
-        } catch (JsonProcessingException e) {
-            throw new RuntimeException(e);
+        } else {
+            throw new IllegalArgumentException("Unknown message type");
         }
     }
-
 }