Merge pull request '0.2.0: fix lost responses and close handling, publish to Gitea' (#2) from fix-pending-requests into master

Reviewed-on: #2

.github/workflows/gradle-build.yml

@@ -0,0 +1,34 @@
+name: Gradle Build
+
+on:
+  push:
+    tags:
+      - 'v*'       # Run when a version tag like v0.1.0 is pushed
+  pull_request:   # Optional: test PRs too
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+
+    steps:
+      # 1. Checkout code
+      - uses: actions/checkout@v4
+
+      # 2. Set up JDK 17
+      - name: Set up JDK 17
+        uses: actions/setup-java@v4
+        with:
+          java-version: '17'
+          distribution: 'temurin'
+
+      # 3. Set up Gradle
+      - name: Setup Gradle
+        uses: gradle/actions/setup-gradle@v4
+
+      # 4. Build project
+      - name: Build with Gradle
+        run: ./gradlew build
+
+      # 5. Run tests
+      - name: Run tests
+        run: ./gradlew test

.github/workflows/gradle-publish.yml

@@ -1,44 +0,0 @@
-# This workflow uses actions that are not certified by GitHub.
-# They are provided by a third-party and are governed by
-# separate terms of service, privacy policy, and support
-# documentation.
-# This workflow will build a package using Gradle and then publish it to GitHub packages when a release is created
-# For more information see: https://github.com/actions/setup-java/blob/main/docs/advanced-usage.md#Publishing-using-gradle
-
-name: Gradle Package
-
-on:
-  release:
-    types: [created]
-
-jobs:
-  build:
-
-    runs-on: ubuntu-latest
-    permissions:
-      contents: read
-      packages: write
-
-    steps:
-    - uses: actions/checkout@v4
-    - name: Set up JDK 17
-      uses: actions/setup-java@v4
-      with:
-        java-version: '17'
-        distribution: 'temurin'
-        server-id: github # Value of the distributionManagement/repository/id field of the pom.xml
-        settings-path: ${{ github.workspace }} # location for the settings.xml file
-
-    - name: Setup Gradle
-      uses: gradle/actions/setup-gradle@af1da67850ed9a4cedd57bfd976089dd991e2582 # v4.0.0
-
-    - name: Build with Gradle
-      run: ./gradlew build
-
-    # The USERNAME and TOKEN need to correspond to the credentials environment variables used in
-    # the publishing section of your build.gradle
-    - name: Publish to GitHub Packages
-      run: ./gradlew publish
-      env:
-        USERNAME: ${{ github.actor }}
-        TOKEN: ${{ secrets.GITHUB_TOKEN }}

README.md

@@ -17,67 +17,69 @@ It ensures compile-time guarantees for message structure, reduces runtime errors
 
 - Fully typed WebSocket sessions
 - Type-safe request/response mechanism
-- Optional typed notifications (Todo)
+- Typed notifications
 - Lightweight, minimal runtime overhead
 
 ## Installation
 
-#### Add the repository
+`session-java` is published to the Gitea package registry at `git.selimaj.dev`.
 
 ```groovy
 repositories {
-    mavenCentral()
-
-    maven {
-        url = uri("https://maven.pkg.github.com/selimaj-dev/session-java")
-        credentials {
-            username = project.findProperty("gpr.user") ?: System.getenv("GITHUB_USERNAME")
-            password = project.findProperty("gpr.key") ?: System.getenv("GITHUB_TOKEN")
-        }
-    }
+    maven { url 'https://git.selimaj.dev/api/packages/selimaj-dev/maven' }
 }
-```
 
-#### Add the dependency
-```groovy
 dependencies {
-    implementation "dev.selimaj.session:session-java:0.1.0"
+    implementation 'dev.selimaj:session-java:0.2.0'
 }
 ```
 
----
+Versions up to 0.1.6 were published through JitPack as `com.github.selimaj-dev:session-java`.
 
-### **Basic Example (client)**
+## Usage
 
 ```java
-Session session = Session.connect("ws://localhost:8080/");
+static final Method<String, String, String> DATA =
+        new Method<>("data", String.class, String.class, String.class);
 
-TextNode response = session.request(Methods.Data, TextNode.valueOf("Hello from client")).get();
+Session session = Session.connect("ws://localhost:8080/", 10, TimeUnit.SECONDS);
 
-System.out.println(response);
+// Handle requests and notifications from the server.
+session.onRequest(DATA, (id, data) -> SessionResult.ok("echo: " + data));
+session.onNotification(DATA, data -> System.out.println("notified: " + data));
+
+// Find out when the connection ends, from either side.
+session.onClose(() -> System.out.println("disconnected"));
+
+// Send a request, with or without a deadline.
+String reply = session.request(DATA, "Hello", Duration.ofSeconds(10)).get();
 
 session.close();
 ```
 
-### **Basic Example (server)**
+### Semantics
 
-#### ⚠️ Server functionality not done yet this is only for future examples
+- A request completes exceptionally with:
+  - `SessionErrorException` when the peer answers with an error response (`getError()` has the payload)
+  - `SessionClosedException` when the session closes before the response arrives
+  - `TimeoutException` when the timeout passes, if one was given
+- `onClose` handlers run exactly once, whether the session was closed locally, by the peer, or by a network failure. A handler registered after close runs immediately.
+- Sends are serialized, so the session is safe to use from multiple threads.
+- A request for an unknown method is answered with an error response instead of no reply.
 
-```java
-SessionServer server = SessionServer.bind("ws://localhost:8080/");
+## Publishing
 
-server.onClient((session, addr) -> {
-    session.onRequest(Methods.Data, (id, data) -> {
-        return SessionResult.ok(TextNode.valueOf("Response from server"));
-    });
-});
+```sh
+GITEA_TOKEN=<token with package:write> ./gradlew publish
 ```
 
 ## Protocol
 
+Every message is a JSON text frame with a `type` tag.
+
 #### Request
 
-The request `id` is separated from the peer, and will increment only on it's requests.
+The `id` is chosen by the sender and increments per peer.
 
 ```json
 { "type": "request", "id": 1, "method": "data", "data": "Hello from client" }
@@ -85,16 +87,22 @@ The request `id` is separated from the peer, and will increment only on it's req
 
 #### Response
 
-The response `id` **must** remain the same as the request.
+A response **must** carry the id of the request it answers.
 
 ```json
 { "type": "response", "id": 1, "result": "Hello from server" }
 ```
 
-#### Notifications
-
-A notification is a method that doesn't need validation or output, it simply notifies a peer for a specific information
+#### Error response
 
 ```json
-{ "type": "notification", "result": "Hello from server" }
+{ "type": "errorresponse", "id": 1, "error": "Invalid data" }
+```
+
+#### Notification
+
+A notification is fire-and-forget and gets no response.
+
+```json
+{ "type": "notification", "method": "data", "data": "Hello from server" }
 ```

build.gradle

@@ -3,13 +3,14 @@ plugins {
     id 'maven-publish'
 }
 
-group = "dev.selimaj.session"
-version = "0.1.0"
+group = "dev.selimaj"
+version = "0.2.0"
 
 java {
     toolchain {
         languageVersion = JavaLanguageVersion.of(17)
     }
+    withSourcesJar()
 }
 
 repositories {
@@ -17,18 +18,15 @@ repositories {
 }
 
 dependencies {
-    // JSON parsing and serialization
-    implementation 'com.fasterxml.jackson.core:jackson-databind:2.17.0'
-    implementation 'com.fasterxml.jackson.core:jackson-annotations:2.17.0'
-    implementation 'com.fasterxml.jackson.core:jackson-core:2.17.0'
+    // JSON parsing and serialization; JsonNode is part of the public API.
+    api 'com.fasterxml.jackson.core:jackson-databind:2.17.0'
+    api 'com.fasterxml.jackson.core:jackson-annotations:2.17.0'
+    api 'com.fasterxml.jackson.core:jackson-core:2.17.0'
 
-    // WebSocket API
-    implementation 'jakarta.websocket:jakarta.websocket-api:2.1.1'
-    // WebSocket client runtime
-    implementation 'org.glassfish.tyrus:tyrus-client:2.1.2'
-
-    // For testing
+    // For testing: a real WebSocket server for the client to talk to.
     testImplementation 'org.junit.jupiter:junit-jupiter:5.10.0'
+    testRuntimeOnly 'org.junit.platform:junit-platform-launcher'
+    testImplementation 'org.java-websocket:Java-WebSocket:1.6.0'
 }
 
 test {
@@ -36,21 +34,25 @@ test {
 }
 
 publishing {
-   publications {
-        gpr(MavenPublication) {
-            from(components.java)
-            groupId = project.group
-            artifactId = "session-java"
-            version = project.version
+    publications {
+        mavenJava(MavenPublication) {
+            artifactId = 'session-java'
+            from components.java
         }
     }
+
     repositories {
+        // Gitea package registry. Publish with a token that has package:write:
+        //   GITEA_TOKEN=... ./gradlew publish
         maven {
-            name = "GitHubPackages"
-            url = uri("https://maven.pkg.github.com/selimaj-dev/session-java")
-            credentials {
-                username = project.findProperty("gpr.user") ?: System.getenv("USERNAME")
-                password = project.findProperty("gpr.key") ?: System.getenv("TOKEN")
+            name = 'Gitea'
+            url = uri('https://git.selimaj.dev/api/packages/selimaj-dev/maven')
+            credentials(HttpHeaderCredentials) {
+                name = 'Authorization'
+                value = "token ${findProperty('giteaToken') ?: System.getenv('GITEA_TOKEN') ?: ''}"
+            }
+            authentication {
+                header(HttpHeaderAuthentication)
             }
         }
     }

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

@@ -3,15 +3,19 @@ package dev.selimaj.session;
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.node.TextNode;
 
+import dev.selimaj.session.types.SessionClosedException;
 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;
 import java.net.http.WebSocket;
+import java.time.Duration;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.TimeUnit;
 
 public final class Session {
 
@@ -34,41 +38,104 @@ public final class Session {
         return new Session(ws, listener);
     }
 
-    public void send(Message msg) throws Exception {
-        listener.send(ws, msg);
+    public static Session connect(String uri, long timeout, TimeUnit unit) throws Exception {
+        SessionListener listener = new SessionListener();
+
+        WebSocket ws = HttpClient.newHttpClient()
+                .newWebSocketBuilder()
+                .buildAsync(URI.create(uri), listener)
+                .orTimeout(timeout, unit)
+                .join();
+
+        return new Session(ws, listener);
     }
 
-    public void respond(int to, JsonNode val) throws Exception {
-        listener.respond(ws, to, val);
+    public CompletableFuture<Void> send(Message msg) {
+        return listener.send(ws, msg);
     }
 
-    public void respondError(int to, JsonNode val) throws Exception {
-        listener.respondError(ws, to, val);
+    public CompletableFuture<Void> respond(int to, JsonNode val) {
+        return listener.respond(ws, to, val);
     }
 
-    public void notify(String method, Object data) throws Exception {
-        listener.notify(ws, method, data);
+    public CompletableFuture<Void> respondError(int to, JsonNode val) {
+        return listener.respondError(ws, to, val);
     }
 
-    <Req extends JsonNode, Res extends JsonNode, Err extends JsonNode> CompletableFuture<Res> request(
-            Method<Req, Res, Err> method, Req req) throws Exception {
+    public CompletableFuture<Void> notify(String method, Object data) {
+        return listener.notify(ws, method, data);
+    }
+
+    /**
+     * Sends a request and completes with the peer's response. Completes
+     * exceptionally with {@link dev.selimaj.session.types.SessionErrorException}
+     * on an error response, or {@link SessionClosedException} if the session
+     * closes first.
+     */
+    public <Req, Res, Err> CompletableFuture<Res> request(
+            Method<Req, Res, Err> method, Req req) {
         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) -> {
+    /**
+     * {@link #request(Method, Object)} that completes exceptionally with a
+     * {@link java.util.concurrent.TimeoutException} after {@code timeout}.
+     */
+    public <Req, Res, Err> CompletableFuture<Res> request(
+            Method<Req, Res, Err> method, Req req, Duration timeout) {
+        return request(method, req).orTimeout(timeout.toMillis(), TimeUnit.MILLISECONDS);
+    }
+
+    public <Req, Res, Err> void onRequest(Method<Req, Res, Err> method,
+            MethodHandler<Req, Res, Err> handler) {
+        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()));
+                return SessionResult.error(TextNode.valueOf(String.valueOf(e.getMessage())));
             }
         };
 
         this.listener.methods.put(method.getName(), wrapper);
     }
 
-    public void close() throws Exception {
-        ws.sendClose(0, "Session closed");
+    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);
+    }
+
+    /**
+     * Runs {@code handler} once when the session closes, whether closed locally,
+     * by the peer, or by a connection failure. Runs immediately if the session
+     * is already closed.
+     */
+    public void onClose(Runnable handler) {
+        listener.closeHandlers.add(handler);
+
+        if (listener.isClosed()) {
+            listener.runCloseHandler(handler);
+        }
+    }
+
+    public boolean isOpen() {
+        return !listener.isClosed() && !ws.isInputClosed() && !ws.isOutputClosed();
+    }
+
+    public void close() {
+        listener.markClosed(new SessionClosedException("Session closed"));
+
+        if (!ws.isOutputClosed()) {
+            ws.sendClose(WebSocket.NORMAL_CLOSURE, "Session closed");
+        }
     }
 }

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

@@ -1,50 +1,92 @@
 package dev.selimaj.session;
 
 import java.net.http.WebSocket;
+import java.nio.ByteBuffer;
+import java.util.List;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.node.TextNode;
 
 import dev.selimaj.session.types.Message;
 import dev.selimaj.session.types.MethodHandler;
+import dev.selimaj.session.types.NotificationHandler;
+import dev.selimaj.session.types.SessionClosedException;
+import dev.selimaj.session.types.SessionErrorException;
 
 public class SessionListener implements WebSocket.Listener {
-    public 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 List<Runnable> closeHandlers = new CopyOnWriteArrayList<>();
 
     final ObjectMapper mapper = new ObjectMapper();
     private final ConcurrentHashMap<Integer, CompletableFuture<JsonNode>> pending = new ConcurrentHashMap<>();
     private final AtomicInteger id = new AtomicInteger();
+    private final AtomicBoolean closed = new AtomicBoolean();
+
+    /** Text frames of a message still being received. */
+    private final StringBuilder partial = new StringBuilder();
+
+    /**
+     * java.net.http.WebSocket allows only one outstanding text send, so every
+     * send is chained after the previous one.
+     */
+    private final Object sendLock = new Object();
+    private CompletableFuture<?> lastSend = CompletableFuture.completedFuture(null);
+
+    @Override
+    public void onOpen(WebSocket ws) {
+        ws.request(1);
+    }
 
     @Override
     public CompletionStage<?> onText(WebSocket ws, CharSequence data, boolean last) {
-        Message msg = null;
+        partial.append(data);
+
+        if (last) {
+            String text = partial.toString();
+            partial.setLength(0);
+            handleMessage(ws, text);
+        }
+
+        ws.request(1);
+        return null;
+    }
+
+    private void handleMessage(WebSocket ws, String text) {
+        Message msg;
 
         try {
-            msg = mapper.readValue(data.toString(), Message.class);
+            msg = mapper.readValue(text, Message.class);
         } catch (Exception e) {
             e.printStackTrace();
-            return null;
+            return;
         }
 
         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())
-                        .thenAccept(res -> {
-                            try {
-                                if (res.isError()) {
-                                    respondError(ws, r.id(), res.value());
-                                } else {
-                                    respond(ws, r.id(), res.value());
-                                }
-                            } catch (Exception ignored) {
-                            }
-                        });
+            if (handler == null) {
+                respondError(ws, r.id(), TextNode.valueOf("Unknown method: " + r.method()));
+                return;
+            }
+
+            try {
+                var res = handler.handle(r.id(), r.data());
+
+                if (res.isError()) {
+                    respondError(ws, r.id(), res.getJsonNode(mapper));
+                } else {
+                    respond(ws, r.id(), res.getJsonNode(mapper));
+                }
+            } catch (Exception e) {
+                respondError(ws, r.id(), TextNode.valueOf(String.valueOf(e.getMessage())));
             }
         } else if (msg instanceof Message.Response r) {
             CompletableFuture<JsonNode> fut = pending.remove(r.id());
@@ -53,51 +95,144 @@ public class SessionListener implements WebSocket.Listener {
         } else if (msg instanceof Message.ErrorResponse r) {
             CompletableFuture<JsonNode> fut = pending.remove(r.id());
             if (fut != null)
-                fut.completeExceptionally(
-                        new RuntimeException(r.error().toString()));
-        }
+                fut.completeExceptionally(new SessionErrorException(r.error()));
+        } 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();
+                }
+            }
+        }
+    }
+
+    @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;
     }
 
-    void send(WebSocket ws, Message msg) throws Exception {
-        ws.sendText(mapper.writeValueAsString(msg), true);
+    @Override
+    public CompletionStage<?> onClose(WebSocket ws, int statusCode, String reason) {
+        markClosed(new SessionClosedException("Closed by peer: " + statusCode + " " + reason));
+        return null;
     }
 
-    void respond(WebSocket ws, int to, JsonNode val) throws Exception {
-        send(ws, new Message.Response(to, val));
+    @Override
+    public void onError(WebSocket ws, Throwable error) {
+        markClosed(new SessionClosedException("Connection failed", error));
     }
 
-    void respondError(WebSocket ws, int to, JsonNode val) throws Exception {
-        send(ws, new Message.ErrorResponse(to, val));
+    boolean isClosed() {
+        return closed.get();
     }
 
-    void notify(WebSocket ws, String method, Object data) throws Exception {
-        send(ws, new Message.Notification(method, mapper.valueToTree(data)));
+    /**
+     * Fails every pending request and runs the close handlers, exactly once.
+     */
+    void markClosed(SessionClosedException cause) {
+        if (!closed.compareAndSet(false, true)) {
+            return;
+        }
+
+        pending.values().forEach(fut -> fut.completeExceptionally(cause));
+        pending.clear();
+
+        for (Runnable handler : closeHandlers) {
+            runCloseHandler(handler);
+        }
     }
 
-    <Req, Res> CompletableFuture<Res> request(
+    /** Runs a close handler unless another thread already claimed it. */
+    void runCloseHandler(Runnable handler) {
+        if (!closeHandlers.remove(handler)) {
+            return;
+        }
+
+        try {
+            handler.run();
+        } catch (Exception e) {
+            e.printStackTrace();
+        }
+    }
+
+    CompletableFuture<Void> send(WebSocket ws, Message msg) {
+        if (closed.get()) {
+            return CompletableFuture.failedFuture(new SessionClosedException("Session is closed"));
+        }
+
+        String text;
+        try {
+            text = mapper.writeValueAsString(msg);
+        } catch (Exception e) {
+            return CompletableFuture.failedFuture(e);
+        }
+
+        synchronized (sendLock) {
+            CompletableFuture<Void> sent = lastSend
+                    .handle((r, e) -> null)
+                    .thenCompose(v -> ws.sendText(text, true))
+                    .thenApply(w -> null);
+            lastSend = sent;
+            return sent;
+        }
+    }
+
+    CompletableFuture<Void> respond(WebSocket ws, int to, JsonNode val) {
+        return send(ws, new Message.Response(to, val));
+    }
+
+    CompletableFuture<Void> respondError(WebSocket ws, int to, JsonNode val) {
+        return send(ws, new Message.ErrorResponse(to, val));
+    }
+
+    CompletableFuture<Void> notify(WebSocket ws, String method, Object data) {
+        return send(ws, new Message.Notification(method, mapper.valueToTree(data)));
+    }
+
+    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> response = new CompletableFuture<>();
 
-        CompletableFuture<JsonNode> fut = new CompletableFuture<>();
-        pending.put(id, fut);
+        // Registered before sending so a fast response can't be missed.
+        pending.put(id, response);
 
-        send(ws, new Message.Request(
-                id,
-                method,
-                mapper.valueToTree(data)));
+        if (closed.get()) {
+            pending.remove(id);
+            return CompletableFuture.failedFuture(new SessionClosedException("Session is closed"));
+        }
 
-        return fut.thenApply(json -> {
+        send(ws, new Message.Request(id, method, mapper.valueToTree(data))).whenComplete((v, e) -> {
+            if (e != null) {
+                pending.remove(id);
+                response.completeExceptionally(e);
+            }
+        });
+
+        // Drop the entry however the caller's future ends (e.g. orTimeout).
+        CompletableFuture<Res> result = response.thenApply(json -> {
             try {
                 return mapper.treeToValue(json, resType);
             } catch (Exception e) {
-                throw new RuntimeException(e);
+                throw new java.util.concurrent.CompletionException(e);
             }
         });
+        result.whenComplete((r, e) -> pending.remove(id));
+
+        return result;
     }
 }

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

@@ -1,8 +1,6 @@
 package dev.selimaj.session.types;
 
-import com.fasterxml.jackson.databind.JsonNode;
-
-public class Method<Req extends JsonNode, Res extends JsonNode, Err extends JsonNode> {
+public class Method<Req, Res, Err> {
     private final String name;
     private final Class<Req> reqClass;
     private final Class<Res> resClass;

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

@@ -1,10 +1,6 @@
 package dev.selimaj.session.types;
 
-import java.util.concurrent.CompletableFuture;
-
-import com.fasterxml.jackson.databind.JsonNode;
-
 @FunctionalInterface
-public interface MethodHandler<Req extends JsonNode> {
-    CompletableFuture<SessionResult> handle(int id, Req data);
+public interface MethodHandler<Req, Res, Err> {
+    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/SessionClosedException.java

@@ -0,0 +1,12 @@
+package dev.selimaj.session.types;
+
+/** The session closed (or failed) before the operation could complete. */
+public class SessionClosedException extends RuntimeException {
+    public SessionClosedException(String message) {
+        super(message);
+    }
+
+    public SessionClosedException(String message, Throwable cause) {
+        super(message, cause);
+    }
+}

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

@@ -0,0 +1,18 @@
+package dev.selimaj.session.types;
+
+import com.fasterxml.jackson.databind.JsonNode;
+
+/** The peer answered a request with an error response. */
+public class SessionErrorException extends RuntimeException {
+    private final JsonNode error;
+
+    public SessionErrorException(JsonNode error) {
+        super(error == null ? "null" : error.isTextual() ? error.asText() : error.toString());
+        this.error = error;
+    }
+
+    /** The error payload exactly as the peer sent it. */
+    public JsonNode getError() {
+        return error;
+    }
+}

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 static CompletableFuture<SessionResult> ok(JsonNode value) {
-        return CompletableFuture.completedFuture(new SessionResult(false, value));
+public record SessionResult<T, E>(boolean isError, Object value) {
+    public static <T, E> SessionResult<T, E> ok(T value) {
+        return new SessionResult<>(false, value);
     }
 
-    public static CompletableFuture<SessionResult> error(JsonNode value) {
-        return CompletableFuture.completedFuture(new SessionResult(false, value));
+    public static <T, E> SessionResult<T, E> error(E err) {
+        return new SessionResult<>(true, 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);
     }
 }

src/test/java/dev/selimaj/session/Methods.java

@@ -1,10 +0,0 @@
-package dev.selimaj.session;
-
-import com.fasterxml.jackson.databind.node.TextNode;
-
-import dev.selimaj.session.types.Method;
-
-public class Methods {
-    public static final Method<TextNode, TextNode, TextNode> Data = new Method<>("data",
-            TextNode.class, TextNode.class, TextNode.class);
-}

src/test/java/dev/selimaj/session/SessionTest.java

@@ -1,34 +1,195 @@
 package dev.selimaj.session;
 
+import static org.junit.jupiter.api.Assertions.*;
+
+import java.nio.ByteBuffer;
+import java.nio.charset.StandardCharsets;
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import org.java_websocket.enums.Opcode;
+import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.Test;
 
+import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.node.TextNode;
+
+import dev.selimaj.session.types.Method;
+import dev.selimaj.session.types.SessionClosedException;
+import dev.selimaj.session.types.SessionErrorException;
 import dev.selimaj.session.types.SessionResult;
 
-public class SessionTest {
-    static void testRequest() throws Exception {
-        Session session = Session.connect("ws://localhost:8080/");
+class SessionTest {
+    static final Method<String, String, String> ECHO = new Method<>("echo", String.class, String.class,
+            String.class);
+    static final Method<String, String, String> FAIL = new Method<>("fail", String.class, String.class,
+            String.class);
 
-        session.onRequest(Methods.Data, (id, data) -> {
-            return SessionResult.ok(TextNode.valueOf("Response from client"));
-        });
+    private final TestServer server = new TestServer();
+    private Session session;
 
-        TextNode response = session.request(Methods.Data, TextNode.valueOf("Hello from client")).get();
+    @AfterEach
+    void tearDown() throws Exception {
+        if (session != null) {
+            session.close();
+        }
+        server.stop(1000);
+    }
 
-        System.out.println(response);
+    private Session connect(String uri) throws Exception {
+        session = Session.connect(uri, 5, TimeUnit.SECONDS);
+        return session;
+    }
 
-        session.close();
+    private static Throwable cause(CompletableFuture<?> future) {
+        ExecutionException e = assertThrows(ExecutionException.class, () -> future.get(5, TimeUnit.SECONDS));
+        return e.getCause();
     }
 
     @Test
-    void requestTest() throws Exception {
+    void requestAndResponse() throws Exception {
+        Session s = connect(server.start(TestServer::echo));
+
+        assertEquals("hi", s.request(ECHO, "hi").get(5, TimeUnit.SECONDS));
     }
 
-    public static void main(String[] args) {
-        try {
-            testRequest();
-        } catch (Exception e) {
-            e.printStackTrace();
+    @Test
+    void errorResponseCompletesWithSessionErrorException() throws Exception {
+        Session s = connect(server.start((conn, msg) -> TestServer.send(conn,
+                "{\"type\":\"errorresponse\",\"id\":" + msg.get("id").asInt() + ",\"error\":\"nope\"}")));
+
+        Throwable cause = cause(s.request(FAIL, "x"));
+        SessionErrorException error = assertInstanceOf(SessionErrorException.class, cause);
+        assertEquals("nope", error.getMessage());
+        assertEquals(TextNode.valueOf("nope"), error.getError());
+    }
+
+    @Test
+    void manyConcurrentRequestsAllComplete() throws Exception {
+        // The server answers synchronously, so responses race the client's
+        // bookkeeping, and requests are sent from many threads at once.
+        Session s = connect(server.start(TestServer::echo));
+
+        List<CompletableFuture<String>> futures = new ArrayList<>();
+        List<Thread> threads = new ArrayList<>();
+        for (int t = 0; t < 8; t++) {
+            int thread = t;
+            Thread th = new Thread(() -> {
+                for (int i = 0; i < 100; i++) {
+                    CompletableFuture<String> f = s.request(ECHO, thread + ":" + i);
+                    synchronized (futures) {
+                        futures.add(f);
+                    }
+                }
+            });
+            threads.add(th);
+            th.start();
         }
+        for (Thread th : threads) {
+            th.join();
+        }
+
+        CompletableFuture.allOf(futures.toArray(CompletableFuture[]::new)).get(10, TimeUnit.SECONDS);
+        assertEquals(800, futures.size());
+    }
+
+    @Test
+    void pendingRequestsFailAndOnCloseRunsOnceWhenServerCloses() throws Exception {
+        Session s = connect(server.start((conn, msg) -> {
+        }));
+        AtomicInteger closes = new AtomicInteger();
+        s.onClose(closes::incrementAndGet);
+
+        CompletableFuture<String> pending = s.request(ECHO, "never answered");
+        server.received.poll(5, TimeUnit.SECONDS);
+        server.client.close(1001, "going away");
+
+        assertInstanceOf(SessionClosedException.class, cause(pending));
+        assertEquals(1, closes.get());
+        assertFalse(s.isOpen());
+        assertInstanceOf(SessionClosedException.class, cause(s.request(ECHO, "late")));
+    }
+
+    @Test
+    void requestTimeout() throws Exception {
+        Session s = connect(server.start((conn, msg) -> {
+        }));
+
+        assertInstanceOf(TimeoutException.class, cause(s.request(ECHO, "x", Duration.ofMillis(100))));
+    }
+
+    @Test
+    void fragmentedMessageIsReassembled() throws Exception {
+        Session s = connect(server.start((conn, msg) -> {
+            byte[] json = ("{\"type\":\"response\",\"id\":" + msg.get("id").asInt()
+                    + ",\"result\":\"reassembled\"}").getBytes(StandardCharsets.UTF_8);
+            int third = json.length / 3;
+            conn.sendFragmentedFrame(Opcode.TEXT, ByteBuffer.wrap(json, 0, third), false);
+            conn.sendFragmentedFrame(Opcode.TEXT, ByteBuffer.wrap(json, third, third), false);
+            conn.sendFragmentedFrame(Opcode.TEXT, ByteBuffer.wrap(json, 2 * third, json.length - 2 * third),
+                    true);
+        }));
+
+        assertEquals("reassembled", s.request(ECHO, "x").get(5, TimeUnit.SECONDS));
+    }
+
+    @Test
+    void handlerResultsAndUnknownMethodsAreAnswered() throws Exception {
+        Session s = connect(server.start((conn, msg) -> {
+        }));
+        s.onRequest(ECHO, (id, data) -> SessionResult.ok("echo:" + data));
+        s.onRequest(FAIL, (id, data) -> SessionResult.error("failed:" + data));
+
+        TestServer.send(server.client, "{\"type\":\"request\",\"id\":1,\"method\":\"echo\",\"data\":\"a\"}");
+        TestServer.send(server.client, "{\"type\":\"request\",\"id\":2,\"method\":\"fail\",\"data\":\"b\"}");
+        TestServer.send(server.client, "{\"type\":\"request\",\"id\":3,\"method\":\"nope\",\"data\":null}");
+
+        JsonNode ok = server.received.poll(5, TimeUnit.SECONDS);
+        JsonNode err = server.received.poll(5, TimeUnit.SECONDS);
+        JsonNode unknown = server.received.poll(5, TimeUnit.SECONDS);
+
+        assertEquals("response", ok.get("type").asText());
+        assertEquals("echo:a", ok.get("result").asText());
+        assertEquals("errorresponse", err.get("type").asText());
+        assertEquals("failed:b", err.get("error").asText());
+        assertEquals("errorresponse", unknown.get("type").asText());
+        assertEquals("Unknown method: nope", unknown.get("error").asText());
+    }
+
+    @Test
+    void notificationsAreDelivered() throws Exception {
+        Session s = connect(server.start((conn, msg) -> {
+        }));
+        CompletableFuture<String> received = new CompletableFuture<>();
+        s.onNotification(ECHO, received::complete);
+
+        TestServer.send(server.client, "{\"type\":\"notification\",\"method\":\"echo\",\"data\":\"hello\"}");
+
+        assertEquals("hello", received.get(5, TimeUnit.SECONDS));
+    }
+
+    @Test
+    void closeSendsNormalClosureAndFailsLaterRequests() throws Exception {
+        Session s = connect(server.start(TestServer::echo));
+        CountDownLatch closed = new CountDownLatch(1);
+        s.onClose(closed::countDown);
+
+        s.close();
+
+        assertTrue(closed.await(5, TimeUnit.SECONDS));
+        assertEquals(1000, server.closeCodes.poll(5, TimeUnit.SECONDS));
+        assertInstanceOf(SessionClosedException.class, cause(s.request(ECHO, "late")));
+
+        // Registered after close: runs immediately, once.
+        AtomicInteger late = new AtomicInteger();
+        s.onClose(late::incrementAndGet);
+        assertEquals(1, late.get());
     }
 }

src/test/java/dev/selimaj/session/TestServer.java

@@ -0,0 +1,84 @@
+package dev.selimaj.session;
+
+import java.net.InetSocketAddress;
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.TimeUnit;
+import java.util.function.BiConsumer;
+
+import org.java_websocket.WebSocket;
+import org.java_websocket.handshake.ClientHandshake;
+import org.java_websocket.server.WebSocketServer;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
+
+/** An in-process protocol peer. Each received JSON message goes to {@link #handler}. */
+class TestServer extends WebSocketServer {
+    static final ObjectMapper MAPPER = new ObjectMapper();
+
+    private final CountDownLatch started = new CountDownLatch(1);
+    final BlockingQueue<JsonNode> received = new LinkedBlockingQueue<>();
+    final BlockingQueue<Integer> closeCodes = new LinkedBlockingQueue<>();
+    volatile WebSocket client;
+    volatile BiConsumer<WebSocket, JsonNode> handler = (conn, msg) -> {
+    };
+
+    TestServer() {
+        super(new InetSocketAddress("127.0.0.1", 0));
+        setReuseAddr(true);
+    }
+
+    String start(BiConsumer<WebSocket, JsonNode> handler) throws InterruptedException {
+        this.handler = handler;
+        start();
+        if (!started.await(5, TimeUnit.SECONDS)) {
+            throw new IllegalStateException("test server did not start");
+        }
+        return "ws://127.0.0.1:" + getPort() + "/";
+    }
+
+    /** Echoes request data back as the result. */
+    static void echo(WebSocket conn, JsonNode msg) {
+        if ("request".equals(msg.path("type").asText())) {
+            send(conn, "{\"type\":\"response\",\"id\":" + msg.get("id").asInt()
+                    + ",\"result\":" + msg.get("data") + "}");
+        }
+    }
+
+    static void send(WebSocket conn, String json) {
+        conn.send(json);
+    }
+
+    @Override
+    public void onOpen(WebSocket conn, ClientHandshake handshake) {
+        client = conn;
+    }
+
+    @Override
+    public void onClose(WebSocket conn, int code, String reason, boolean remote) {
+        closeCodes.add(code);
+    }
+
+    @Override
+    public void onMessage(WebSocket conn, String message) {
+        try {
+            JsonNode msg = MAPPER.readTree(message);
+            received.add(msg);
+            handler.accept(conn, msg);
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    @Override
+    public void onError(WebSocket conn, Exception ex) {
+        ex.printStackTrace();
+    }
+
+    @Override
+    public void onStart() {
+        started.countDown();
+    }
+}