MCP: admit calls chained on another's end, match only in-range stdio ids

org.restlet.ext.mcp/src/main/java/org/restlet/ext/mcp/StdioServerHelper.java

@@ -14,6 +14,8 @@
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.OutputStream;
+import java.math.BigDecimal;
+import java.math.BigInteger;
 import java.nio.charset.StandardCharsets;
 import java.time.Duration;
 import java.util.Map;
@@ -265,11 +267,41 @@ private static boolean modern(Map<?, ?> envelope) {
                                         .PROTOCOL_VERSION));
     }
 
-    /** A request id comparable across its JSON integer forms, such as 1 and 1.0. */
+    /**
+     * A request id comparable across its JSON number forms, such as 1 and 1.0: a long when the
+     * number is an integer in its range, exactly, else the number itself, normalised.
+     */
     private static Object key(Object id) {
-        return id instanceof Number number && number.doubleValue() == number.longValue()
-                ? (Object) number.longValue()
-                : id;
+        if (!(id instanceof Number number)) {
+            return id;
+        }
+        BigDecimal exact = exact(number);
+        if (exact == null) {
+            return id; // NaN or infinite: compared as given.
+        }
+        try {
+            return exact.longValueExact();
+        } catch (ArithmeticException _) {
+            return exact.stripTrailingZeros();
+        }
+    }
+
+    /**
+     * The exact value of a number: a floating-point one by its binary value, not its shortest
+     * decimal form, which may round an integer such as 1000000000000000128.0.
+     *
+     * @return The value, or null for NaN or an infinity.
+     */
+    private static BigDecimal exact(Number number) {
+        return switch (number) {
+            case BigDecimal decimal -> decimal;
+            case BigInteger integer -> new BigDecimal(integer);
+            case Double _, Float _ ->
+                    Double.isFinite(number.doubleValue())
+                            ? new BigDecimal(number.doubleValue())
+                            : null;
+            default -> BigDecimal.valueOf(number.longValue());
+        };
     }
 
     /**

org.restlet.ext.mcp/src/main/java/org/restlet/ext/mcp/internal/McpWorkers.java

@@ -24,7 +24,16 @@
 public final class McpWorkers {
     private final ExecutorService executor = Executors.newVirtualThreadPerTaskExecutor();
     private final Semaphore admission;
-    private final int maxCalls;
+
+    /** Guards {@link #running}, and signals when it drops to zero. */
+    private final Object idle = new Object();
+
+    /**
+     * The calls admitted and not yet done signalling their end. A call frees its admission permit
+     * before it signals, so that a call chained on its end is admitted, but counts here until its
+     * subscribers are done.
+     */
+    private int running;
 
     /** The time limit of each call, or null for none. */
     private final Duration timeout;
@@ -53,7 +62,6 @@ public McpWorkers(int maxCalls, Duration timeout, Duration closeTimeout) {
             throw new IllegalArgumentException("Positive call limit and timeouts required");
         }
         this.admission = new Semaphore(maxCalls);
-        this.maxCalls = maxCalls;
         this.timeout = timeout == null ? null : Duration.ofNanos(timeout.toNanos());
         this.closeTimeout = Duration.ofNanos(closeTimeout.toNanos());
     }
@@ -86,40 +94,56 @@ private void admit(MonoSink<Void> sink, Operation operation, Runnable cancel) {
                 sink.error(new IOException("MCP transport closed or call limit exceeded"));
                 return;
             }
+            synchronized (idle) {
+                running++;
+            }
             executor.execute(() -> work(sink, operation, worker, cancelled));
         }
     }
 
     /**
-     * Runs an admitted call and signals its end. The permit is released after the signal, so a
-     * caller of {@link #awaitIdle(Duration)} sees what the call's subscribers did when it ended.
+     * Runs an admitted call and signals its end. The permit is released before the signal, so a
+     * call its subscribers chain on its end is admitted; the call still counts as running until
+     * they are done, so a caller of {@link #awaitIdle(Duration)} sees what they did.
      */
+    // An Error is caught to end the call and leave the running count, then rethrown.
+    @SuppressWarnings("java:S1181")
     private void work(
             MonoSink<Void> sink,
             Operation operation,
             AtomicReference<Thread> worker,
             AtomicBoolean cancelled) {
         worker.set(Thread.currentThread());
+        Throwable error = null;
         try {
-            if (cancelled.get()) {
-                return;
-            }
-            Exception error = null;
+            boolean ran = false;
             try {
-                operation.run();
-            } catch (Exception e) {
+                if (!cancelled.get()) {
+                    ran = true;
+                    operation.run();
+                }
+            } catch (Exception | Error e) {
+                // An Error, such as a failed assertion, ends the call too, and is rethrown.
                 error = e;
+            } finally {
+                admission.release();
             }
-            if (cancelled.get()) {
-                return;
-            }
-            if (error == null) {
-                sink.success();
-            } else {
-                sink.error(error);
+            if (ran && !cancelled.get()) {
+                if (error == null) {
+                    sink.success();
+                } else {
+                    sink.error(error);
+                }
             }
         } finally {
-            admission.release();
+            synchronized (idle) {
+                if (--running == 0) {
+                    idle.notifyAll();
+                }
+            }
+        }
+        if (error instanceof Error fatal) {
+            throw fatal;
         }
     }
 
@@ -131,11 +155,17 @@ private void work(
      * @throws InterruptedException If interrupted while waiting.
      */
     public boolean awaitIdle(Duration wait) throws InterruptedException {
-        if (!admission.tryAcquire(maxCalls, wait.toNanos(), TimeUnit.NANOSECONDS)) {
-            return false;
+        long deadline = System.nanoTime() + wait.toNanos();
+        synchronized (idle) {
+            while (running > 0) {
+                long remaining = deadline - System.nanoTime();
+                if (remaining <= 0) {
+                    return false;
+                }
+                TimeUnit.NANOSECONDS.timedWait(idle, remaining);
+            }
+            return true;
         }
-        admission.release(maxCalls);
-        return true;
     }
 
     public void close() throws IOException {

org.restlet.ext.mcp/src/test/java/org/restlet/ext/mcp/McpModernStdioTestCase.java

@@ -298,6 +298,53 @@ void anInvalidIdIsNotEchoed() throws Exception {
         assertTrue(reply.containsKey("error"), reply.toString());
     }
 
+    /** A cancellation naming an id beyond the long range does not match the largest long. */
+    @Test
+    void anOutOfRangeIdCancelsNothing() throws Exception {
+        Map<String, Object> meta = meta();
+        meta.put(MetaKeys.PROGRESS_TOKEN, "p4");
+        Map<String, Object> call =
+                request(0, "tools/call", Map.of("name", "endless", "arguments", Map.of()), meta);
+        call.put("id", Long.MAX_VALUE);
+        send(call);
+        assertEquals("notifications/progress", read().get("method"));
+
+        sendRaw(
+                "{\"jsonrpc\":\"2.0\",\"method\":\"notifications/cancelled\","
+                        + "\"params\":{\"requestId\":9223372036854775808.0}}");
+        send(
+                request(
+                        10,
+                        "tools/call",
+                        Map.of("name", "echo", "arguments", Map.of("text", "still")),
+                        meta()));
+        assertEquals("echo:still", firstText(result(read())));
+        assertEquals(1, cancelled.getCount(), "another request was cancelled");
+
+        sendRaw(
+                "{\"jsonrpc\":\"2.0\",\"method\":\"notifications/cancelled\","
+                        + "\"params\":{\"requestId\":9223372036854775807}}");
+        assertTrue(cancelled.await(10, TimeUnit.SECONDS), "the request was not cancelled");
+    }
+
+    /** A cancellation naming an integral floating-point id matches its exact value. */
+    @Test
+    void anExactFloatingPointIdCancelsItsRequest() throws Exception {
+        Map<String, Object> meta = meta();
+        meta.put(MetaKeys.PROGRESS_TOKEN, "p5");
+        Map<String, Object> call =
+                request(0, "tools/call", Map.of("name", "endless", "arguments", Map.of()), meta);
+        call.put("id", 1000000000000000128L);
+        send(call);
+        assertEquals("notifications/progress", read().get("method"));
+
+        sendRaw(
+                "{\"jsonrpc\":\"2.0\",\"method\":\"notifications/cancelled\","
+                        + "\"params\":{\"requestId\":1000000000000000128.0}}");
+
+        assertTrue(cancelled.await(10, TimeUnit.SECONDS), "the request was not cancelled");
+    }
+
     private static Map<String, Object> request(
             int id, String method, Map<String, Object> params, Map<String, Object> meta) {
         Map<String, Object> withMeta = new LinkedHashMap<>(params);
@@ -323,6 +370,11 @@ private void send(Map<String, Object> message) throws IOException {
         toServer.flush();
     }
 
+    private void sendRaw(String line) throws IOException {
+        toServer.write((line + "\n").getBytes(StandardCharsets.UTF_8));
+        toServer.flush();
+    }
+
     private Map<String, Object> read() throws Exception {
         String line = lines.poll(10, TimeUnit.SECONDS);
         assertNotNull(line, "no server output");

org.restlet.ext.mcp/src/test/java/org/restlet/ext/mcp/McpWorkersTestCase.java

@@ -0,0 +1,93 @@
+/**
+ * Copyright 2005-2026 Qlik
+ *<p>
+ * The content of this file is subject to the terms of the Apache 2.0 open
+ * source license available at https://www.opensource.org/licenses/apache-2.0
+ *<p>
+ * Restlet is a registered trademark of QlikTech International AB.
+ */
+package org.restlet.ext.mcp;
+
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.time.Duration;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import org.junit.jupiter.api.Test;
+import org.restlet.ext.mcp.internal.McpWorkers;
+
+/** The admission workers shared by the MCP transports. */
+class McpWorkersTestCase {
+
+    /** A call chained on the end of another is admitted, with a single slot. */
+    @Test
+    void aChainedCallIsAdmittedWithOneSlot() throws Exception {
+        McpWorkers workers = new McpWorkers(1, Duration.ofSeconds(5));
+        AtomicInteger calls = new AtomicInteger();
+        try {
+            assertDoesNotThrow(
+                    () ->
+                            workers.submit(calls::incrementAndGet, () -> {})
+                                    .then(workers.submit(calls::incrementAndGet, () -> {}))
+                                    .block(Duration.ofSeconds(5)));
+            assertEquals(2, calls.get());
+        } finally {
+            workers.close();
+        }
+    }
+
+    /** A call failing with an Error ends with it, and no longer counts as running. */
+    @Test
+    void aCallFailingWithAnErrorEnds() throws Exception {
+        McpWorkers workers = new McpWorkers(1, Duration.ofSeconds(5));
+        try {
+            reactor.core.publisher.Mono<Void> call =
+                    workers.submit(
+                            () -> {
+                                throw new AssertionError("broken");
+                            },
+                            () -> {});
+            RuntimeException blocked =
+                    org.junit.jupiter.api.Assertions.assertThrows(
+                            RuntimeException.class, () -> call.block(Duration.ofSeconds(5)));
+            Throwable failure = reactor.core.Exceptions.unwrap(blocked);
+            assertTrue(failure instanceof AssertionError, failure.toString());
+            assertEquals("broken", failure.getMessage());
+            assertTrue(workers.awaitIdle(Duration.ofSeconds(5)));
+        } finally {
+            workers.close();
+        }
+    }
+
+    /** A call is running until the subscribers of its end are done. */
+    @Test
+    void awaitIdleWaitsForTheSubscribersOfTheEnd() throws Exception {
+        McpWorkers workers = new McpWorkers(1, Duration.ofSeconds(5));
+        CountDownLatch entered = new CountDownLatch(1);
+        CountDownLatch release = new CountDownLatch(1);
+        try {
+            workers.submit(() -> {}, () -> {})
+                    .doOnSuccess(
+                            ignored -> {
+                                entered.countDown();
+                                try {
+                                    release.await(5, TimeUnit.SECONDS);
+                                } catch (InterruptedException _) {
+                                    Thread.currentThread().interrupt();
+                                }
+                            })
+                    .subscribe();
+            assertTrue(entered.await(5, TimeUnit.SECONDS));
+
+            assertFalse(workers.awaitIdle(Duration.ofMillis(50)));
+            release.countDown();
+            assertTrue(workers.awaitIdle(Duration.ofSeconds(5)));
+        } finally {
+            workers.close();
+        }
+    }
+}