Enforce required client capabilities in the modern server

conformance-tests/server-servlet-modern/src/main/java/io/modelcontextprotocol/conformance/server/modern/ModernConformanceServlet.java

@@ -11,7 +11,6 @@
 import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
 import io.modelcontextprotocol.json.McpJsonDefaults;
 import io.modelcontextprotocol.json.McpJsonMapper;
-import io.modelcontextprotocol.modern.McpException;
 import io.modelcontextprotocol.modern.McpSchema.AudioContent;
 import io.modelcontextprotocol.modern.McpSchema.BlobResourceContents;
 import io.modelcontextprotocol.modern.McpSchema.CallToolOutcome;
@@ -255,6 +254,15 @@ private static Handler<CallToolRequest, CallToolOutcome> stream(StreamingTool to
 		return (ctx, request) -> McpSyncResponse.streaming(notifier -> tool.run(ctx, request, notifier));
 	}
 
+	private static final ClientCapabilities.Elicitation FORM = new ClientCapabilities.Elicitation(
+			new ClientCapabilities.Elicitation.Form(), null);
+
+	private static final ClientCapabilities ELICITATION = ClientCapabilities.builder().elicitation(FORM).build();
+
+	private static final ClientCapabilities SAMPLING = ClientCapabilities.builder().sampling().build();
+
+	private static final ClientCapabilities ROOTS = ClientCapabilities.builder().roots().build();
+
 	private static final class ConformanceTools implements McpSyncToolRepository {
 
 		private final Map<String, McpSyncTool> tools = new LinkedHashMap<>();
@@ -316,12 +324,8 @@ private static final class ConformanceTools implements McpSyncToolRepository {
 				.build(), respond((ctx, req) -> text("ok")));
 
 			// SEP-2575 diagnostic tools
-			add("test_missing_capability", "Test tool requiring sampling", (ctx, req) -> {
-				if (!ctx.clientCapabilities().supportsSampling()) {
-					throw McpException.missingClientCapability(ClientCapabilities.builder().sampling().build());
-				}
-				return McpSyncResponse.result(text("Success"));
-			});
+			add("test_missing_capability", "Test tool requiring sampling", SAMPLING,
+					respond((ctx, req) -> text("Success")));
 			add("test_streaming_elicitation", "Diagnostic tool validating response progress streams",
 					stream((ctx, req, notifier) -> {
 						notifier.progress(50, 100.0, null);
@@ -342,7 +346,7 @@ private static final class ConformanceTools implements McpSyncToolRepository {
 
 			// SEP-2322 MRTR tools
 			add("test_input_required_result_elicitation", "MRTR: returns InputRequiredResult with elicitation request",
-					respond((ctx, req) -> {
+					ELICITATION, respond((ctx, req) -> {
 						Map<String, Object> content = acceptedContent(req.inputResponses(), "user_name");
 						if (content == null || !(content.get("name") instanceof String name)) {
 							return InputRequiredResult.builder()
@@ -352,7 +356,7 @@ private static final class ConformanceTools implements McpSyncToolRepository {
 						return text("Hello, " + name + "!");
 					}));
 			add("test_input_required_result_sampling", "MRTR: returns InputRequiredResult with sampling request",
-					respond((ctx, req) -> {
+					SAMPLING, respond((ctx, req) -> {
 						Map<String, Object> response = responseObject(req.inputResponses(), "capital_question");
 						if (response == null) {
 							return InputRequiredResult.builder()
@@ -364,15 +368,15 @@ private static final class ConformanceTools implements McpSyncToolRepository {
 						return text("Sampling result: " + answer);
 					}));
 			add("test_input_required_result_list_roots", "MRTR: returns InputRequiredResult with roots/list request",
-					respond((ctx, req) -> {
+					ROOTS, respond((ctx, req) -> {
 						Map<String, Object> response = responseObject(req.inputResponses(), "client_roots");
 						if (response == null || !(response.get("roots") instanceof List<?> roots)) {
 							return InputRequiredResult.builder().listRoots("client_roots").build();
 						}
 						return text("Found " + roots.size() + " root(s): " + roots);
 					}));
 			add("test_input_required_result_request_state", "MRTR: returns InputRequiredResult with requestState",
-					respond((ctx, req) -> {
+					ELICITATION, respond((ctx, req) -> {
 						Map<String, Object> content = acceptedContent(req.inputResponses(), "confirm");
 						if ("request-state".equals(req.requestState()) && content != null
 								&& Boolean.TRUE.equals(content.get("ok"))) {
@@ -384,7 +388,8 @@ private static final class ConformanceTools implements McpSyncToolRepository {
 							.build();
 					}));
 			add("test_input_required_result_multiple_inputs",
-					"MRTR: returns InputRequiredResult with multiple input requests", respond((ctx, req) -> {
+					"MRTR: returns InputRequiredResult with multiple input requests",
+					ClientCapabilities.builder().elicitation(FORM).sampling().roots().build(), respond((ctx, req) -> {
 						Map<String, Object> user = acceptedContent(req.inputResponses(), "user_name");
 						Map<String, Object> greeting = responseObject(req.inputResponses(), "greeting");
 						Map<String, Object> roots = responseObject(req.inputResponses(), "client_roots");
@@ -400,7 +405,7 @@ private static final class ConformanceTools implements McpSyncToolRepository {
 							.requestState("multiple-inputs")
 							.build();
 					}));
-			add("test_input_required_result_multi_round", "MRTR: multi-round InputRequiredResult workflow",
+			add("test_input_required_result_multi_round", "MRTR: multi-round InputRequiredResult workflow", ELICITATION,
 					respond((ctx, req) -> {
 						String state = req.requestState();
 						if (state != null && state.startsWith("round-1")) {
@@ -427,7 +432,7 @@ private static final class ConformanceTools implements McpSyncToolRepository {
 			// Integrity protection comes from the server's RequestStateCodec: a tampered
 			// requestState is rejected before this handler runs.
 			add("test_input_required_result_tampered_state", "MRTR: HMAC-signed requestState integrity test",
-					respond((ctx, req) -> {
+					ELICITATION, respond((ctx, req) -> {
 						if ("tamper-test".equals(req.requestState())
 								&& responseObject(req.inputResponses(), "confirm") != null) {
 							return text("integrity-ok: state verified");
@@ -459,11 +464,21 @@ && responseObject(req.inputResponses(), "confirm") != null) {
 		}
 
 		private void add(String name, String description, Handler<CallToolRequest, CallToolOutcome> handler) {
-			add(Tool.builder(name, EMPTY_JSON_SCHEMA).description(description).build(), handler);
+			add(name, description, ClientCapabilities.NONE, handler);
+		}
+
+		private void add(String name, String description, ClientCapabilities requiredCapabilities,
+				Handler<CallToolRequest, CallToolOutcome> handler) {
+			add(Tool.builder(name, EMPTY_JSON_SCHEMA).description(description).build(), requiredCapabilities, handler);
 		}
 
 		private void add(Tool tool, Handler<CallToolRequest, CallToolOutcome> handler) {
-			this.tools.put(tool.name(), McpSyncTool.of(tool, handler::handle));
+			add(tool, ClientCapabilities.NONE, handler);
+		}
+
+		private void add(Tool tool, ClientCapabilities requiredCapabilities,
+				Handler<CallToolRequest, CallToolOutcome> handler) {
+			this.tools.put(tool.name(), McpSyncTool.of(tool, requiredCapabilities, handler::handle));
 		}
 
 		@Override
@@ -608,7 +623,7 @@ private static final class ConformancePrompts implements McpSyncPromptRepository
 								.build())));
 			add(Prompt.builder("test_input_required_result_prompt")
 				.description("MRTR: prompt that requires elicitation input")
-				.build(), respond((ctx, req) -> {
+				.build(), ELICITATION, respond((ctx, req) -> {
 					Map<String, Object> content = acceptedContent(req.inputResponses(), "user_context");
 					if (content == null || !(content.get("context") instanceof String context)) {
 						return InputRequiredResult.builder()
@@ -622,7 +637,12 @@ private static final class ConformancePrompts implements McpSyncPromptRepository
 		}
 
 		private void add(Prompt prompt, Handler<GetPromptRequest, GetPromptOutcome> handler) {
-			this.prompts.put(prompt.name(), McpSyncPrompt.of(prompt, handler::handle));
+			add(prompt, ClientCapabilities.NONE, handler);
+		}
+
+		private void add(Prompt prompt, ClientCapabilities requiredCapabilities,
+				Handler<GetPromptRequest, GetPromptOutcome> handler) {
+			this.prompts.put(prompt.name(), McpSyncPrompt.of(prompt, requiredCapabilities, handler::handle));
 		}
 
 		private static GetPromptResult messages(PromptMessage... messages) {

mcp-core/src/main/java/io/modelcontextprotocol/modern/McpException.java

@@ -42,6 +42,13 @@ public static McpException invalidParams(String message, Object data) {
 		return new McpException(ErrorCodes.INVALID_PARAMS, message, data);
 	}
 
+	/**
+	 * The request cannot be served without {@code required}, which the client did not
+	 * declare. Answered with {@code -32021}. Tools, prompts and resources declare such
+	 * capabilities up front instead (see
+	 * {@link io.modelcontextprotocol.modern.server.feature.McpAsyncTool}); this is for
+	 * extension methods.
+	 */
 	public static McpException missingClientCapability(ClientCapabilities required) {
 		return new McpException(ErrorCodes.MISSING_REQUIRED_CLIENT_CAPABILITY, "Missing required client capability",
 				new MissingRequiredClientCapabilityData(required));

mcp-core/src/main/java/io/modelcontextprotocol/modern/McpSchema.java

@@ -192,6 +192,11 @@ public interface CacheableResult extends Result {
 	 * A result asking the client for more input before the request can complete. An
 	 * extension method supporting multi round-trip requests returns its own record
 	 * implementing this interface and the method's outcome type.
+	 * <p>
+	 * Each input request needs a client capability: {@code elicitation} in the request's
+	 * mode, {@code sampling} or {@code roots}. Only input whose capability the client
+	 * declared for this request may be requested; a result asking for anything else is
+	 * answered with an internal error.
 	 */
 	public interface InputRequired extends Result {
 
@@ -1410,7 +1415,8 @@ public String mode() {
 	/**
 	 * A result signalling that the server needs more input before it can complete the
 	 * original request. Only {@code tools/call}, {@code resources/read} and
-	 * {@code prompts/get} may return it.
+	 * {@code prompts/get} may return it, and only with the input requests
+	 * {@link InputRequired} allows.
 	 */
 	@JsonInclude(JsonInclude.Include.NON_ABSENT)
 	@JsonIgnoreProperties(ignoreUnknown = true)

mcp-core/src/main/java/io/modelcontextprotocol/modern/server/McpFeature.java

@@ -7,6 +7,7 @@
 import java.util.Set;
 
 import io.modelcontextprotocol.modern.McpException;
+import io.modelcontextprotocol.modern.McpSchema.InputRequired;
 import io.modelcontextprotocol.modern.McpSchema.Result;
 import io.modelcontextprotocol.modern.McpSchema.ServerCapabilities;
 import reactor.core.publisher.Mono;
@@ -38,8 +39,8 @@ default void capabilities(ServerCapabilities.Builder builder) {
 
 	/**
 	 * The methods that support multi round-trip requests: they may answer with an
-	 * {@code InputRequired} result and accept its {@code requestState} on retry. Default:
-	 * none.
+	 * {@link InputRequired} result, within the input requests it allows, and accept its
+	 * {@code requestState} on retry. Default: none.
 	 */
 	default Set<String> inputRequiredMethods() {
 		return Set.of();

mcp-core/src/main/java/io/modelcontextprotocol/modern/server/McpRequestContext.java

@@ -9,6 +9,7 @@
 import io.modelcontextprotocol.common.McpTransportContext;
 import io.modelcontextprotocol.modern.McpSchema.ClientCapabilities;
 import io.modelcontextprotocol.modern.McpSchema.Implementation;
+import io.modelcontextprotocol.modern.McpSchema.InputRequired;
 import io.modelcontextprotocol.util.Assert;
 
 /**
@@ -76,7 +77,10 @@ public String protocolVersion() {
 		return this.protocolVersion;
 	}
 
-	/** The client capabilities declared in {@code _meta} for this request. */
+	/**
+	 * The client capabilities declared in {@code _meta} for this request. Input requests
+	 * may only use these; see {@link InputRequired}.
+	 */
 	public ClientCapabilities clientCapabilities() {
 		return this.clientCapabilities;
 	}

mcp-core/src/main/java/io/modelcontextprotocol/modern/server/McpServer.java

@@ -27,6 +27,7 @@
 import io.modelcontextprotocol.modern.McpSchema.CacheScope;
 import io.modelcontextprotocol.modern.McpSchema.CacheableResult;
 import io.modelcontextprotocol.modern.McpSchema.ClientCapabilities;
+import io.modelcontextprotocol.modern.McpSchema.ElicitUrlRequest;
 import io.modelcontextprotocol.modern.McpSchema.ErrorCodes;
 import io.modelcontextprotocol.modern.McpSchema.Implementation;
 import io.modelcontextprotocol.modern.McpSchema.InputRequired;
@@ -238,6 +239,17 @@ private JSONRPCResponse resultResponse(McpRequestContext ctx, Result result) {
 			return bug(ctx, "Method does not support input-required results");
 		}
 		Map<String, Object> resultMap = new LinkedHashMap<>(this.jsonMapper.convertValue(result, MAP_TYPE_REF));
+		if (result instanceof InputRequired) {
+			// Checked on the serialized form, so a request's elicitation mode is read the
+			// same way the client will read it.
+			String undeclared = undeclaredInputRequest(ctx.clientCapabilities(), resultMap.get("inputRequests"));
+			if (undeclared != null) {
+				return bug(ctx,
+						"Input request '" + undeclared
+								+ "' needs a client capability the request did not declare; require the capability"
+								+ " up front or only request input the client declared");
+			}
+		}
 		Map<String, Object> resultMeta = new LinkedHashMap<>();
 		if (resultMap.get("_meta") instanceof Map<?, ?> existingMeta) {
 			existingMeta.forEach((k, v) -> resultMeta.put(String.valueOf(k), v));
@@ -267,6 +279,31 @@ private static JSONRPCResponse errorResponse(Object id, Throwable throwable) {
 		return JSONRPCResponse.error(id, new JSONRPCError(ErrorCodes.INTERNAL_ERROR, "Internal error"));
 	}
 
+	// The key of the first input request the client cannot answer, or null if it can
+	// answer all of them. Methods outside the spec's three are not checked.
+	private static String undeclaredInputRequest(ClientCapabilities declared, Object inputRequests) {
+		if (!(inputRequests instanceof Map<?, ?> requests)) {
+			return null;
+		}
+		for (Map.Entry<?, ?> entry : requests.entrySet()) {
+			if (!(entry.getValue() instanceof Map<?, ?> request)) {
+				continue;
+			}
+			boolean supported = switch (String.valueOf(request.get("method"))) {
+				case McpSchema.METHOD_ELICITATION_CREATE -> request.get("params") instanceof Map<?, ?> params
+						&& ElicitUrlRequest.MODE.equals(params.get("mode")) ? declared.supportsElicitationUrl()
+								: declared.supportsElicitationForm();
+				case McpSchema.METHOD_SAMPLING_CREATE_MESSAGE -> declared.supportsSampling();
+				case McpSchema.METHOD_ROOTS_LIST -> declared.supportsRoots();
+				default -> true;
+			};
+			if (!supported) {
+				return String.valueOf(entry.getKey());
+			}
+		}
+		return null;
+	}
+
 	private static JSONRPCResponse bug(McpRequestContext ctx, String problem) {
 		logger.warn("Request {} ({}): {}", ctx.requestId(), ctx.method(), problem);
 		return JSONRPCResponse.error(ctx.requestId(), new JSONRPCError(ErrorCodes.INTERNAL_ERROR, "Internal error"));

mcp-core/src/main/java/io/modelcontextprotocol/modern/server/feature/McpAsyncPrompt.java

@@ -4,6 +4,7 @@
 
 package io.modelcontextprotocol.modern.server.feature;
 
+import io.modelcontextprotocol.modern.McpSchema.ClientCapabilities;
 import io.modelcontextprotocol.modern.McpSchema.GetPromptOutcome;
 import io.modelcontextprotocol.modern.McpSchema.GetPromptRequest;
 import io.modelcontextprotocol.modern.McpSchema.Prompt;
@@ -14,30 +15,55 @@
 
 /**
  * A prompt's definition together with the handler that renders it.
+ * <p>
+ * Client capabilities the prompt cannot be rendered without belong in
+ * {@link #requiredCapabilities()}: a client that did not declare them gets {@code -32021}
+ * and the handler does not run. Input the prompt can do without, or can get another way,
+ * is requested only when {@link McpRequestContext#clientCapabilities()} declares its
+ * capability. Requesting input the client did not declare is answered with an internal
+ * error.
  *
  * @author Dariusz Jędrzejczyk
  */
 public final class McpAsyncPrompt {
 
 	private final Prompt definition;
 
+	private final ClientCapabilities requiredCapabilities;
+
 	private final Handler handler;
 
-	private McpAsyncPrompt(Prompt definition, Handler handler) {
+	private McpAsyncPrompt(Prompt definition, ClientCapabilities requiredCapabilities, Handler handler) {
 		Assert.notNull(definition, "definition must not be null");
+		Assert.notNull(requiredCapabilities, "requiredCapabilities must not be null");
 		Assert.notNull(handler, "handler must not be null");
 		this.definition = definition;
+		this.requiredCapabilities = requiredCapabilities;
 		this.handler = handler;
 	}
 
+	/** A prompt that requires no client capabilities. */
 	public static McpAsyncPrompt of(Prompt definition, Handler handler) {
-		return new McpAsyncPrompt(definition, handler);
+		return new McpAsyncPrompt(definition, ClientCapabilities.NONE, handler);
+	}
+
+	/**
+	 * A prompt that only works with clients declaring {@code requiredCapabilities}. A
+	 * request from any other client is rejected with {@code -32021} without running the
+	 * handler.
+	 */
+	public static McpAsyncPrompt of(Prompt definition, ClientCapabilities requiredCapabilities, Handler handler) {
+		return new McpAsyncPrompt(definition, requiredCapabilities, handler);
 	}
 
 	public Prompt definition() {
 		return this.definition;
 	}
 
+	public ClientCapabilities requiredCapabilities() {
+		return this.requiredCapabilities;
+	}
+
 	public Handler handler() {
 		return this.handler;
 	}