From 7e2369f2590d60ef25e8d2385ac0b47ebc0d4ff6 Mon Sep 17 00:00:00 2001
From: Greg Oledzki <greg.oledzki@moderne.io>
Date: Fri, 8 May 2026 08:16:29 +0200
Subject: [PATCH] Fix shutdown race condition

---
 src/main/java/io/moderne/jsonrpc/JsonRpc.java | 17 ++++++++++++++---
 1 file changed, 14 insertions(+), 3 deletions(-)

diff --git a/src/main/java/io/moderne/jsonrpc/JsonRpc.java b/src/main/java/io/moderne/jsonrpc/JsonRpc.java
index 7850422..76028c0 100644
--- a/src/main/java/io/moderne/jsonrpc/JsonRpc.java
+++ b/src/main/java/io/moderne/jsonrpc/JsonRpc.java
@@ -56,6 +56,16 @@ public <P> JsonRpc rpc(String name, JsonRpcMethod<P> method) {
     public CompletableFuture<JsonRpcSuccess> send(JsonRpcRequest request) {
         CompletableFuture<JsonRpcSuccess> response = new CompletableFuture<>();
         openRequests.put(request.getId(), response);
+        if (shutdown) {
+            // Reader loop already exited (peer EOF or explicit shutdown) and
+            // may have drained openRequests before our put. Fail the future
+            // ourselves; completeExceptionally is a no-op if the reader did
+            // observe our entry and failed it first.
+            openRequests.remove(request.getId());
+            response.completeExceptionally(new JsonRpcException(
+                    JsonRpcError.internalError(null, "JSON-RPC peer closed the stream")));
+            return response;
+        }
         messageHandler.send(request, formatter);
         return response;
     }
@@ -113,15 +123,16 @@ protected void compute() {
                         }
                     } catch (EOFException e) {
                         // Peer closed the stream — there's nothing more to read.
-                        // Fail any in-flight requests once and exit the loop
-                        // instead of spinning on the closed pipe.
+                        // Set shutdown FIRST so a concurrent send() observes it
+                        // and fails its own future after put; otherwise a request
+                        // registered after this drain would be stranded.
+                        shutdown = true;
                         JsonRpcException eof = new JsonRpcException(
                                 JsonRpcError.internalError(null, "JSON-RPC peer closed the stream"));
                         for (CompletableFuture<JsonRpcSuccess> future : openRequests.values()) {
                             future.completeExceptionally(eof);
                         }
                         openRequests.clear();
-                        shutdown = true;
                     } catch (JsonRpcReceiveException e) {
                         // Frame- or parse-level failure on an inbound message.
                         // Send the error back to the peer; do NOT touch