Fixed notification handling

src/main/java/dev/selimaj/session/Session.java

@@ -7,6 +7,7 @@ import dev.selimaj.session.types.SessionResult;
 import dev.selimaj.session.types.Message;
 import dev.selimaj.session.types.Method;
 import dev.selimaj.session.types.MethodHandler;
+import dev.selimaj.session.types.NotificationHandler;
 
 import java.net.URI;
 import java.net.http.HttpClient;
@@ -63,16 +64,17 @@ public final class Session {
         listener.notify(ws, method, data);
     }
 
-    <Req, Res, Err> CompletableFuture<Res> request(
+    public <Req, Res, Err> CompletableFuture<Res> request(
             Method<Req, Res, Err> method, Req req) throws Exception {
         return listener.request(ws, method.getName(), req, method.getResClass());
     }
 
-    <Req, Res, Err> void onRequest(Method<Req, Res, Err> method,
+    public <Req, Res, Err> void onRequest(Method<Req, Res, Err> method,
-            MethodHandler<Req> handler) {
+            MethodHandler<Req, Res, Err> handler) {
-        MethodHandler<JsonNode> wrapper = (id, value) -> {
+        MethodHandler<JsonNode, JsonNode, JsonNode> wrapper = (id, value) -> {
             try {
-                return handler.handle(id, listener.mapper.treeToValue(value, method.getReqClass()));
+                return handler.handle(id, listener.mapper.treeToValue(value, method.getReqClass()))
+                        .intoJSON(listener.mapper);
             } catch (Exception e) {
                 return SessionResult.error(TextNode.valueOf(e.getMessage()));
             }
@@ -81,6 +83,20 @@ public final class Session {
         this.listener.methods.put(method.getName(), wrapper);
     }
 
+    public <Req, Res, Err> void onNotification(Method<Req, Res, Err> method, NotificationHandler<Req> handler) {
+        NotificationHandler<JsonNode> wrapper = (value) -> {
+            CompletableFuture.runAsync(() -> {
+                try {
+                    handler.handle(listener.mapper.treeToValue(value, method.getReqClass()));
+                } catch (Exception e) {
+                    e.printStackTrace();
+                }
+            });
+        };
+
+        this.listener.notificationHandlers.put(method.getName(), wrapper);
+    }
+
     public void close() throws Exception {
         ws.sendClose(0, "Session closed");
     }

src/main/java/dev/selimaj/session/SessionListener.java

@@ -3,6 +3,7 @@ package dev.selimaj.session;
 import java.net.http.WebSocket;
 import java.nio.ByteBuffer;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
 import java.util.concurrent.CompletionStage;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -12,9 +13,11 @@ import com.fasterxml.jackson.databind.ObjectMapper;
 
 import dev.selimaj.session.types.Message;
 import dev.selimaj.session.types.MethodHandler;
+import dev.selimaj.session.types.NotificationHandler;
 
 public class SessionListener implements WebSocket.Listener {
-    final ConcurrentHashMap<String, MethodHandler<JsonNode>> methods = new ConcurrentHashMap<>();
+    final ConcurrentHashMap<String, MethodHandler<JsonNode, JsonNode, JsonNode>> methods = new ConcurrentHashMap<>();
+    final ConcurrentHashMap<String, NotificationHandler<JsonNode>> notificationHandlers = new ConcurrentHashMap<>();
 
     final ObjectMapper mapper = new ObjectMapper();
     private final ConcurrentHashMap<Integer, CompletableFuture<JsonNode>> pending = new ConcurrentHashMap<>();
@@ -32,20 +35,19 @@ public class SessionListener implements WebSocket.Listener {
         }
 
         if (msg instanceof Message.Request r) {
-            MethodHandler<JsonNode> handler = methods.get(r.method());
+            MethodHandler<JsonNode, JsonNode, JsonNode> handler = methods.get(r.method());
 
             if (handler != null) {
-                handler.handle(r.id(), r.data())
+                try {
-                        .thenAccept(res -> {
+                    var res = handler.handle(r.id(), r.data());
-                            try {
+
-                                if (res.isError()) {
+                    if (res.isError()) {
-                                    respondError(ws, r.id(), res.value());
+                        respondError(ws, r.id(), res.getJsonNode(mapper));
-                                } else {
+                    } else {
-                                    respond(ws, r.id(), res.value());
+                        respond(ws, r.id(), res.getJsonNode(mapper));
-                                }
+                    }
-                            } catch (Exception ignored) {
+                } catch (Exception ignored) {
-                            }
+                }
-                        });
             }
         } else if (msg instanceof Message.Response r) {
             CompletableFuture<JsonNode> fut = pending.remove(r.id());
@@ -56,16 +58,45 @@ public class SessionListener implements WebSocket.Listener {
             if (fut != null)
                 fut.completeExceptionally(
                         new RuntimeException(r.error().toString()));
+        } else if (msg instanceof Message.Notification r) {
+            NotificationHandler<JsonNode> handler = notificationHandlers.get(r.method());
+
+            if (handler != null) {
+                try {
+                    handler.handle(r.data());
+                } catch (Exception e) {
+                    e.printStackTrace();
+                }
+            }
         }
 
+        ws.request(1);
         return null;
     }
 
+    @Override
+    public void onOpen(WebSocket ws) {
+        ws.request(1);
+    }
+
     @Override
     public CompletionStage<?> onPing(WebSocket ws, ByteBuffer message) {
+        ws.request(1);
         return ws.sendPong(message);
     }
 
+    @Override
+    public CompletionStage<?> onPong(WebSocket ws, ByteBuffer message) {
+        ws.request(1);
+        return null;
+    }
+
+    @Override
+    public CompletionStage<?> onClose(WebSocket ws, int statusCode, String reason) {
+        ws.request(1);
+        return null;
+    }
+
     void send(WebSocket ws, Message msg) throws Exception {
         ws.sendText(mapper.writeValueAsString(msg), true);
     }
@@ -82,28 +113,36 @@ public class SessionListener implements WebSocket.Listener {
         send(ws, new Message.Notification(method, mapper.valueToTree(data)));
     }
 
-    <Req, Res> CompletableFuture<Res> request(
+    public <Req, Res> CompletableFuture<Res> request(
             WebSocket ws,
             String method,
             Req data,
-            Class<Res> resType) throws Exception {
+            Class<Res> resType) {
 
         int id = this.id.incrementAndGet();
 
-        CompletableFuture<JsonNode> fut = new CompletableFuture<>();
+        return CompletableFuture
-        pending.put(id, fut);
+                .runAsync(() -> {
-
+                    try {
-        send(ws, new Message.Request(
+                        send(ws, new Message.Request(
-                id,
+                                id,
-                method,
+                                method,
-                mapper.valueToTree(data)));
+                                mapper.valueToTree(data)));
-
+                    } catch (Exception e) {
-        return fut.thenApply(json -> {
+                        throw new CompletionException(e);
-            try {
+                    }
-                return mapper.treeToValue(json, resType);
+                })
-            } catch (Exception e) {
+                .thenCompose(v -> {
-                throw new RuntimeException(e);
+                    CompletableFuture<JsonNode> fut = new CompletableFuture<>();
-            }
+                    pending.put(id, fut);
-        });
+                    return fut;
+                })
+                .thenApply(json -> {
+                    try {
+                        return mapper.treeToValue(json, resType);
+                    } catch (Exception e) {
+                        throw new CompletionException(e);
+                    }
+                });
     }
 }

src/main/java/dev/selimaj/session/types/MethodHandler.java

@@ -1,8 +1,6 @@
 package dev.selimaj.session.types;
 
-import java.util.concurrent.CompletableFuture;
-
 @FunctionalInterface
-public interface MethodHandler<Req> {
+public interface MethodHandler<Req, Res, Err> {
-    CompletableFuture<SessionResult> handle(int id, Req data);
+    SessionResult<Res, Err> handle(int id, Req data);
 }

src/main/java/dev/selimaj/session/types/NotificationHandler.java

@@ -0,0 +1,6 @@
+package dev.selimaj.session.types;
+
+@FunctionalInterface
+public interface NotificationHandler<Req> {
+    void handle(Req data);
+}

src/main/java/dev/selimaj/session/types/SessionResult.java

@@ -1,15 +1,22 @@
 package dev.selimaj.session.types;
 
-import java.util.concurrent.CompletableFuture;
-
 import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
 
-public record SessionResult(boolean isError, JsonNode value) {
+public record SessionResult<T, E>(boolean isError, Object value) {
-    public static CompletableFuture<SessionResult> ok(JsonNode value) {
+    public static <T, E> SessionResult<T, E> ok(T value) {
-        return CompletableFuture.completedFuture(new SessionResult(false, value));
+        return new SessionResult<>(false, value);
     }
 
-    public static CompletableFuture<SessionResult> error(JsonNode value) {
+    public static <T, E> SessionResult<T, E> error(E err) {
-        return CompletableFuture.completedFuture(new SessionResult(false, value));
+        return new SessionResult<>(false, err);
+    }
+
+    public SessionResult<JsonNode, JsonNode> intoJSON(ObjectMapper mapper) {
+        return new SessionResult<>(this.isError, mapper.valueToTree(this.value));
+    }
+
+    public JsonNode getJsonNode(ObjectMapper mapper) {
+        return mapper.valueToTree(this.value);
     }
 }