From fb66ed08e1b6a3ceb5dcfef09ffda339aadd9c4d Mon Sep 17 00:00:00 2001
From: Klesti Selimaj <klesti.selimay@icloud.com>
Date: Thu, 19 Feb 2026 18:08:33 +0100
Subject: [PATCH] Request handling

---
 src/main/java/dev/selimaj/session/Session.java   | 16 ++++++++++++++++
 .../dev/selimaj/session/SessionListener.java     |  6 +++---
 .../dev/selimaj/session/types/HandlerResult.java |  6 ------
 .../dev/selimaj/session/types/MethodHandler.java |  8 ++++----
 .../java/dev/selimaj/session/SessionTest.java    |  7 +++++--
 5 files changed, 28 insertions(+), 15 deletions(-)
 delete mode 100644 src/main/java/dev/selimaj/session/types/HandlerResult.java

diff --git a/src/main/java/dev/selimaj/session/Session.java b/src/main/java/dev/selimaj/session/Session.java
index 32909d1..47e3236 100644
--- a/src/main/java/dev/selimaj/session/Session.java
+++ b/src/main/java/dev/selimaj/session/Session.java
@@ -1,9 +1,12 @@
 package dev.selimaj.session;
 
 import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.node.TextNode;
 
+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 java.net.URI;
 import java.net.http.HttpClient;
@@ -52,6 +55,19 @@ public final class Session {
         return listener.request(ws, method.getName(), req, method.getResClass());
     }
 
+    <Req extends JsonNode, Res extends JsonNode, Err extends JsonNode> void onRequest(Method<Req, Res, Err> method,
+            MethodHandler<Req> handler) {
+        MethodHandler<JsonNode> wrapper = (id, value) -> {
+            try {
+                return handler.handle(id, listener.mapper.treeToValue(value, method.getReqClass()));
+            } catch (Exception e) {
+                return SessionResult.error(TextNode.valueOf(e.getMessage()));
+            }
+        };
+
+        this.listener.methods.put(method.getName(), wrapper);
+    }
+
     public void close() throws Exception {
         ws.sendClose(0, "Session closed");
     }
diff --git a/src/main/java/dev/selimaj/session/SessionListener.java b/src/main/java/dev/selimaj/session/SessionListener.java
index 4889ad4..c274cf5 100644
--- a/src/main/java/dev/selimaj/session/SessionListener.java
+++ b/src/main/java/dev/selimaj/session/SessionListener.java
@@ -13,9 +13,9 @@ import dev.selimaj.session.types.Message;
 import dev.selimaj.session.types.MethodHandler;
 
 public class SessionListener implements WebSocket.Listener {
-    public final ConcurrentHashMap<String, MethodHandler> methods = new ConcurrentHashMap<>();
+    public final ConcurrentHashMap<String, MethodHandler<JsonNode>> methods = new ConcurrentHashMap<>();
 
-    private final ObjectMapper mapper = new ObjectMapper();
+    final ObjectMapper mapper = new ObjectMapper();
     private final ConcurrentHashMap<Integer, CompletableFuture<JsonNode>> pending = new ConcurrentHashMap<>();
     private final AtomicInteger id = new AtomicInteger();
 
@@ -31,7 +31,7 @@ public class SessionListener implements WebSocket.Listener {
         }
 
         if (msg instanceof Message.Request r) {
-            MethodHandler handler = methods.get(r.method());
+            MethodHandler<JsonNode> handler = methods.get(r.method());
 
             if (handler != null) {
                 handler.handle(r.id(), r.data())
diff --git a/src/main/java/dev/selimaj/session/types/HandlerResult.java b/src/main/java/dev/selimaj/session/types/HandlerResult.java
deleted file mode 100644
index d2eaf00..0000000
--- a/src/main/java/dev/selimaj/session/types/HandlerResult.java
+++ /dev/null
@@ -1,6 +0,0 @@
-package dev.selimaj.session.types;
-
-import com.fasterxml.jackson.databind.JsonNode;
-
-public record HandlerResult(boolean isError, JsonNode value) {
-}
diff --git a/src/main/java/dev/selimaj/session/types/MethodHandler.java b/src/main/java/dev/selimaj/session/types/MethodHandler.java
index 46ea0b3..b1c9879 100644
--- a/src/main/java/dev/selimaj/session/types/MethodHandler.java
+++ b/src/main/java/dev/selimaj/session/types/MethodHandler.java
@@ -1,10 +1,10 @@
 package dev.selimaj.session.types;
 
-import com.fasterxml.jackson.databind.JsonNode;
-
 import java.util.concurrent.CompletableFuture;
 
+import com.fasterxml.jackson.databind.JsonNode;
+
 @FunctionalInterface
-public interface MethodHandler {
-    CompletableFuture<HandlerResult> handle(int id, JsonNode data);
+public interface MethodHandler<Req extends JsonNode> {
+    CompletableFuture<SessionResult> handle(int id, Req data);
 }
diff --git a/src/test/java/dev/selimaj/session/SessionTest.java b/src/test/java/dev/selimaj/session/SessionTest.java
index 14a4093..422c0ea 100644
--- a/src/test/java/dev/selimaj/session/SessionTest.java
+++ b/src/test/java/dev/selimaj/session/SessionTest.java
@@ -3,17 +3,20 @@ package dev.selimaj.session;
 import org.junit.jupiter.api.Test;
 
 import com.fasterxml.jackson.databind.node.TextNode;
+import dev.selimaj.session.types.SessionResult;
 
 public class SessionTest {
     static void testRequest() throws Exception {
         Session session = Session.connect("ws://localhost:8080/");
 
+        session.onRequest(Methods.Data, (id, data) -> {
+            return SessionResult.ok(TextNode.valueOf("Response from client"));
+        });
+
         TextNode response = session.request(Methods.Data, TextNode.valueOf("Hello from client")).get();
 
         System.out.println(response);
 
-        Thread.sleep(1000);
-
         session.close();
     }