diff --git a/org.restlet.ext.mcp/src/main/java/org/restlet/ext/mcp/StdioServerHelper.java b/org.restlet.ext.mcp/src/main/java/org/restlet/ext/mcp/StdioServerHelper.java index 552cb3f47a..9b74f584f1 100644 --- a/org.restlet.ext.mcp/src/main/java/org/restlet/ext/mcp/StdioServerHelper.java +++ b/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()); + }; } /** diff --git a/org.restlet.ext.mcp/src/main/java/org/restlet/ext/mcp/internal/McpWorkers.java b/org.restlet.ext.mcp/src/main/java/org/restlet/ext/mcp/internal/McpWorkers.java index cf74b4d8f7..2a3b39b153 100644 --- a/org.restlet.ext.mcp/src/main/java/org/restlet/ext/mcp/internal/McpWorkers.java +++ b/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 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 sink, Operation operation, AtomicReference 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 { diff --git a/org.restlet.ext.mcp/src/test/java/org/restlet/ext/mcp/McpModernStdioTestCase.java b/org.restlet.ext.mcp/src/test/java/org/restlet/ext/mcp/McpModernStdioTestCase.java index 0d804b924e..a3992e3e09 100644 --- a/org.restlet.ext.mcp/src/test/java/org/restlet/ext/mcp/McpModernStdioTestCase.java +++ b/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 meta = meta(); + meta.put(MetaKeys.PROGRESS_TOKEN, "p4"); + Map 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 meta = meta(); + meta.put(MetaKeys.PROGRESS_TOKEN, "p5"); + Map 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 request( int id, String method, Map params, Map meta) { Map withMeta = new LinkedHashMap<>(params); @@ -323,6 +370,11 @@ private void send(Map message) throws IOException { toServer.flush(); } + private void sendRaw(String line) throws IOException { + toServer.write((line + "\n").getBytes(StandardCharsets.UTF_8)); + toServer.flush(); + } + private Map read() throws Exception { String line = lines.poll(10, TimeUnit.SECONDS); assertNotNull(line, "no server output"); diff --git a/org.restlet.ext.mcp/src/test/java/org/restlet/ext/mcp/McpWorkersTestCase.java b/org.restlet.ext.mcp/src/test/java/org/restlet/ext/mcp/McpWorkersTestCase.java new file mode 100644 index 0000000000..d78e967cbe --- /dev/null +++ b/org.restlet.ext.mcp/src/test/java/org/restlet/ext/mcp/McpWorkersTestCase.java @@ -0,0 +1,93 @@ +/** + * Copyright 2005-2026 Qlik + *

+ * 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 + *

+ * 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 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(); + } + } +}