Skip to content

Commit 7853efe

Browse files
jitokimtzolov
jitokim
authored and
tzolov
committed
feat(completion): implement completion API support (#141)
Add completion API support to the MCP protocol implementation: - Add CompleteRequest and CompleteResult schema classes - Implement completion capabilities in ServerCapabilities - Add completion handlers in McpAsyncServer and McpServer - Add completion client methods in McpAsyncClient and McpSyncClient - Add CompletionRefKey and completion specifications in McpServerFeatures - Add integration test for completion functionality - Fix isPresent() check to use isEmpty() in WebMvcSseServerTransportProvider - Replace McpServerFeatures.CompletionRefKey by McpSchemaCompleteReference Co-authored-by: Christian Tzolov <[email protected]> Signed-off-by: jitokim <[email protected]>
1 parent 84adde1 commit 7853efe

File tree

9 files changed

+340
-33
lines changed

9 files changed

+340
-33
lines changed

‎mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/WebFluxSseIntegrationTests.java‎

Lines changed: 54 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
import java.util.concurrent.ConcurrentHashMap;
1111
import java.util.concurrent.TimeUnit;
1212
import java.util.concurrent.atomic.AtomicReference;
13+
import java.util.function.BiFunction;
1314
import java.util.function.Function;
1415
import java.util.stream.Collectors;
1516

@@ -20,19 +21,12 @@
2021
import io.modelcontextprotocol.server.McpServer;
2122
import io.modelcontextprotocol.server.McpServerFeatures;
2223
import io.modelcontextprotocol.server.TestUtil;
24+
import io.modelcontextprotocol.server.McpSyncServerExchange;
2325
import io.modelcontextprotocol.server.transport.WebFluxSseServerTransportProvider;
2426
import io.modelcontextprotocol.spec.McpError;
2527
import io.modelcontextprotocol.spec.McpSchema;
26-
import io.modelcontextprotocol.spec.McpSchema.CallToolResult;
27-
import io.modelcontextprotocol.spec.McpSchema.ClientCapabilities;
28-
import io.modelcontextprotocol.spec.McpSchema.CreateMessageRequest;
29-
import io.modelcontextprotocol.spec.McpSchema.CreateMessageResult;
30-
import io.modelcontextprotocol.spec.McpSchema.InitializeResult;
31-
import io.modelcontextprotocol.spec.McpSchema.ModelPreferences;
32-
import io.modelcontextprotocol.spec.McpSchema.Role;
33-
import io.modelcontextprotocol.spec.McpSchema.Root;
34-
import io.modelcontextprotocol.spec.McpSchema.ServerCapabilities;
35-
import io.modelcontextprotocol.spec.McpSchema.Tool;
28+
import io.modelcontextprotocol.spec.McpSchema.*;
29+
import io.modelcontextprotocol.spec.McpSchema.ServerCapabilities.CompletionCapabilities;
3630
import org.junit.jupiter.api.AfterEach;
3731
import org.junit.jupiter.api.BeforeEach;
3832
import org.junit.jupiter.params.ParameterizedTest;
@@ -759,4 +753,53 @@ void testLoggingNotification(String clientType) {
759753
mcpServer.close();
760754
}
761755

762-
}
756+
// ---------------------------------------
757+
// Completion Tests
758+
// ---------------------------------------
759+
@ParameterizedTest(name = "{0} : Completion call")
760+
@ValueSource(strings = { "httpclient", "webflux" })
761+
void testCompletionShouldReturnExpectedSuggestions(String clientType) {
762+
var clientBuilder = clientBuilders.get(clientType);
763+
764+
var expectedValues = List.of("python", "pytorch", "pyside");
765+
var completionResponse = new McpSchema.CompleteResult(new CompleteResult.CompleteCompletion(expectedValues, 10, // total
766+
true // hasMore
767+
));
768+
769+
AtomicReference<CompleteRequest> samplingRequest = new AtomicReference<>();
770+
BiFunction<McpSyncServerExchange, CompleteRequest, CompleteResult> completionHandler = (mcpSyncServerExchange,
771+
request) -> {
772+
samplingRequest.set(request);
773+
return completionResponse;
774+
};
775+
776+
var mcpServer = McpServer.sync(mcpServerTransportProvider)
777+
.capabilities(ServerCapabilities.builder().completions(new CompletionCapabilities()).build())
778+
.prompts(new McpServerFeatures.SyncPromptSpecification(
779+
new Prompt("code_review", "this is code review prompt", List.of()),
780+
(mcpSyncServerExchange, getPromptRequest) -> null))
781+
.completions(new McpServerFeatures.SyncCompletionSpecification(
782+
new McpSchema.PromptReference("ref/prompt", "code_review"), completionHandler))
783+
.build();
784+
785+
try (var mcpClient = clientBuilder.build()) {
786+
787+
InitializeResult initResult = mcpClient.initialize();
788+
assertThat(initResult).isNotNull();
789+
790+
CompleteRequest request = new CompleteRequest(new PromptReference("ref/prompt", "code_review"),
791+
new CompleteRequest.CompleteArgument("language", "py"));
792+
793+
CompleteResult result = mcpClient.completeCompletion(request);
794+
795+
assertThat(result).isNotNull();
796+
797+
assertThat(samplingRequest.get().argument().name()).isEqualTo("language");
798+
assertThat(samplingRequest.get().argument().value()).isEqualTo("py");
799+
assertThat(samplingRequest.get().ref().type()).isEqualTo("ref/prompt");
800+
}
801+
802+
mcpServer.close();
803+
}
804+
805+
}

‎mcp-spring/mcp-spring-webmvc/src/main/java/io/modelcontextprotocol/server/transport/WebMvcSseServerTransportProvider.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -300,7 +300,7 @@ private ServerResponse handleMessage(ServerRequest request) {
300300
return ServerResponse.status(HttpStatus.SERVICE_UNAVAILABLE).body("Server is shutting down");
301301
}
302302

303-
if (!request.param("sessionId").isPresent()) {
303+
if (request.param("sessionId").isEmpty()) {
304304
return ServerResponse.badRequest().body(new McpError("Session ID missing in message endpoint"));
305305
}
306306

‎mcp/src/main/java/io/modelcontextprotocol/client/McpAsyncClient.java‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,7 @@
7171
*
7272
* @author Dariusz Jędrzejczyk
7373
* @author Christian Tzolov
74+
* @author Jihoon Kim
7475
* @see McpClient
7576
* @see McpSchema
7677
* @see McpClientSession
@@ -816,4 +817,25 @@ void setProtocolVersions(List<String> protocolVersions) {
816817
this.protocolVersions = protocolVersions;
817818
}
818819

820+
// --------------------------
821+
// Completions
822+
// --------------------------
823+
private static final TypeReference<McpSchema.CompleteResult> COMPLETION_COMPLETE_RESULT_TYPE_REF = new TypeReference<>() {
824+
};
825+
826+
/**
827+
* Sends a completion/complete request to generate value suggestions based on a given
828+
* reference and argument. This is typically used to provide auto-completion options
829+
* for user input fields.
830+
* @param completeRequest The request containing the prompt or resource reference and
831+
* argument for which to generate completions.
832+
* @return A Mono that completes with the result containing completion suggestions.
833+
* @see McpSchema.CompleteRequest
834+
* @see McpSchema.CompleteResult
835+
*/
836+
public Mono<McpSchema.CompleteResult> completeCompletion(McpSchema.CompleteRequest completeRequest) {
837+
return this.withInitializationCheck("complete completions", initializedResult -> this.mcpSession
838+
.sendRequest(McpSchema.METHOD_COMPLETION_COMPLETE, completeRequest, COMPLETION_COMPLETE_RESULT_TYPE_REF));
839+
}
840+
819841
}

‎mcp/src/main/java/io/modelcontextprotocol/client/McpSyncClient.java‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@
4646
*
4747
* @author Dariusz Jędrzejczyk
4848
* @author Christian Tzolov
49+
* @author Jihoon Kim
4950
* @see McpClient
5051
* @see McpAsyncClient
5152
* @see McpSchema
@@ -334,4 +335,14 @@ public void setLoggingLevel(McpSchema.LoggingLevel loggingLevel) {
334335
this.delegate.setLoggingLevel(loggingLevel).block();
335336
}
336337

338+
/**
339+
* Send a completion/complete request.
340+
* @param completeRequest the completion request contains the prompt or resource
341+
* reference and arguments for generating suggestions.
342+
* @return the completion result containing suggested values.
343+
*/
344+
public McpSchema.CompleteResult completeCompletion(McpSchema.CompleteRequest completeRequest) {
345+
return this.delegate.completeCompletion(completeRequest).block();
346+
}
347+
337348
}

‎mcp/src/main/java/io/modelcontextprotocol/server/McpAsyncServer.java‎

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,7 @@
6969
*
7070
* @author Christian Tzolov
7171
* @author Dariusz Jędrzejczyk
72+
* @author Jihoon Kim
7273
* @see McpServer
7374
* @see McpSchema
7475
* @see McpClientSession
@@ -269,6 +270,8 @@ private static class AsyncServerImpl extends McpAsyncServer {
269270
// broadcasting loggingNotification.
270271
private LoggingLevel minLoggingLevel = LoggingLevel.DEBUG;
271272

273+
private final ConcurrentHashMap<McpSchema.CompleteReference, McpServerFeatures.AsyncCompletionSpecification> completions = new ConcurrentHashMap<>();
274+
272275
private List<String> protocolVersions = List.of(McpSchema.LATEST_PROTOCOL_VERSION);
273276

274277
AsyncServerImpl(McpServerTransportProvider mcpTransportProvider, ObjectMapper objectMapper,
@@ -282,6 +285,7 @@ private static class AsyncServerImpl extends McpAsyncServer {
282285
this.resources.putAll(features.resources());
283286
this.resourceTemplates.addAll(features.resourceTemplates());
284287
this.prompts.putAll(features.prompts());
288+
this.completions.putAll(features.completions());
285289

286290
Map<String, McpServerSession.RequestHandler<?>> requestHandlers = new HashMap<>();
287291

@@ -314,6 +318,11 @@ private static class AsyncServerImpl extends McpAsyncServer {
314318
requestHandlers.put(McpSchema.METHOD_LOGGING_SET_LEVEL, setLoggerRequestHandler());
315319
}
316320

321+
// Add completion API handlers if the completion capability is enabled
322+
if (this.serverCapabilities.completions() != null) {
323+
requestHandlers.put(McpSchema.METHOD_COMPLETION_COMPLETE, completionCompleteRequestHandler());
324+
}
325+
317326
Map<String, McpServerSession.NotificationHandler> notificationHandlers = new HashMap<>();
318327

319328
notificationHandlers.put(McpSchema.METHOD_NOTIFICATION_INITIALIZED, (exchange, params) -> Mono.empty());
@@ -706,6 +715,81 @@ private McpServerSession.RequestHandler<Object> setLoggerRequestHandler() {
706715
};
707716
}
708717

718+
private McpServerSession.RequestHandler<McpSchema.CompleteResult> completionCompleteRequestHandler() {
719+
return (exchange, params) -> {
720+
McpSchema.CompleteRequest request = parseCompletionParams(params);
721+
722+
if (request.ref() == null) {
723+
return Mono.error(new McpError("ref must not be null"));
724+
}
725+
726+
if (request.ref().type() == null) {
727+
return Mono.error(new McpError("type must not be null"));
728+
}
729+
730+
String type = request.ref().type();
731+
732+
// check if the referenced resource exists
733+
if (type.equals("ref/prompt") && request.ref() instanceof McpSchema.PromptReference promptReference) {
734+
McpServerFeatures.AsyncPromptSpecification prompt = this.prompts.get(promptReference.name());
735+
if (prompt == null) {
736+
return Mono.error(new McpError("Prompt not found: " + promptReference.name()));
737+
}
738+
}
739+
740+
if (type.equals("ref/resource")
741+
&& request.ref() instanceof McpSchema.ResourceReference resourceReference) {
742+
McpServerFeatures.AsyncResourceSpecification resource = this.resources.get(resourceReference.uri());
743+
if (resource == null) {
744+
return Mono.error(new McpError("Resource not found: " + resourceReference.uri()));
745+
}
746+
}
747+
748+
McpServerFeatures.AsyncCompletionSpecification specification = this.completions.get(request.ref());
749+
750+
if (specification == null) {
751+
return Mono.error(new McpError("AsyncCompletionSpecification not found: " + request.ref()));
752+
}
753+
754+
return specification.completionHandler().apply(exchange, request);
755+
};
756+
}
757+
758+
/**
759+
* Parses the raw JSON-RPC request parameters into a
760+
* {@link McpSchema.CompleteRequest} object.
761+
* <p>
762+
* This method manually extracts the `ref` and `argument` fields from the input
763+
* map, determines the correct reference type (either prompt or resource), and
764+
* constructs a fully-typed {@code CompleteRequest} instance.
765+
* @param object the raw request parameters, expected to be a Map containing "ref"
766+
* and "argument" entries.
767+
* @return a {@link McpSchema.CompleteRequest} representing the structured
768+
* completion request.
769+
* @throws IllegalArgumentException if the "ref" type is not recognized.
770+
*/
771+
@SuppressWarnings("unchecked")
772+
private McpSchema.CompleteRequest parseCompletionParams(Object object) {
773+
Map<String, Object> params = (Map<String, Object>) object;
774+
Map<String, Object> refMap = (Map<String, Object>) params.get("ref");
775+
Map<String, Object> argMap = (Map<String, Object>) params.get("argument");
776+
777+
String refType = (String) refMap.get("type");
778+
779+
McpSchema.CompleteReference ref = switch (refType) {
780+
case "ref/prompt" -> new McpSchema.PromptReference(refType, (String) refMap.get("name"));
781+
case "ref/resource" -> new McpSchema.ResourceReference(refType, (String) refMap.get("uri"));
782+
default -> throw new IllegalArgumentException("Invalid ref type: " + refType);
783+
};
784+
785+
String argName = (String) argMap.get("name");
786+
String argValue = (String) argMap.get("value");
787+
McpSchema.CompleteRequest.CompleteArgument argument = new McpSchema.CompleteRequest.CompleteArgument(
788+
argName, argValue);
789+
790+
return new McpSchema.CompleteRequest(ref, argument);
791+
}
792+
709793
// ---------------------------------------
710794
// Sampling
711795
// ---------------------------------------

‎mcp/src/main/java/io/modelcontextprotocol/server/McpServer.java‎

Lines changed: 40 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -115,6 +115,7 @@
115115
*
116116
* @author Christian Tzolov
117117
* @author Dariusz Jędrzejczyk
118+
* @author Jihoon Kim
118119
* @see McpAsyncServer
119120
* @see McpSyncServer
120121
* @see McpServerTransportProvider
@@ -192,6 +193,8 @@ class AsyncSpecification {
192193
*/
193194
private final Map<String, McpServerFeatures.AsyncPromptSpecification> prompts = new HashMap<>();
194195

196+
private final Map<McpSchema.CompleteReference, McpServerFeatures.AsyncCompletionSpecification> completions = new HashMap<>();
197+
195198
private final List<BiFunction<McpAsyncServerExchange, List<McpSchema.Root>, Mono<Void>>> rootsChangeHandlers = new ArrayList<>();
196199

197200
private Duration requestTimeout = Duration.ofSeconds(10); // Default timeout
@@ -581,7 +584,8 @@ public AsyncSpecification objectMapper(ObjectMapper objectMapper) {
581584
*/
582585
public McpAsyncServer build() {
583586
var features = new McpServerFeatures.Async(this.serverInfo, this.serverCapabilities, this.tools,
584-
this.resources, this.resourceTemplates, this.prompts, this.rootsChangeHandlers, this.instructions);
587+
this.resources, this.resourceTemplates, this.prompts, this.completions, this.rootsChangeHandlers,
588+
this.instructions);
585589
var mapper = this.objectMapper != null ? this.objectMapper : new ObjectMapper();
586590
return new McpAsyncServer(this.transportProvider, mapper, features, this.requestTimeout);
587591
}
@@ -635,6 +639,8 @@ class SyncSpecification {
635639
*/
636640
private final Map<String, McpServerFeatures.SyncPromptSpecification> prompts = new HashMap<>();
637641

642+
private final Map<McpSchema.CompleteReference, McpServerFeatures.SyncCompletionSpecification> completions = new HashMap<>();
643+
638644
private final List<BiConsumer<McpSyncServerExchange, List<McpSchema.Root>>> rootsChangeHandlers = new ArrayList<>();
639645

640646
private Duration requestTimeout = Duration.ofSeconds(10); // Default timeout
@@ -957,6 +963,37 @@ public SyncSpecification prompts(McpServerFeatures.SyncPromptSpecification... pr
957963
return this;
958964
}
959965

966+
/**
967+
* Registers multiple completions with their handlers using a List. This method is
968+
* useful when completions need to be added in bulk from a collection.
969+
* @param completions List of completion specifications. Must not be null.
970+
* @return This builder instance for method chaining
971+
* @throws IllegalArgumentException if completions is null
972+
* @see #completions(McpServerFeatures.SyncCompletionSpecification...)
973+
*/
974+
public SyncSpecification completions(List<McpServerFeatures.SyncCompletionSpecification> completions) {
975+
Assert.notNull(completions, "Completions list must not be null");
976+
for (McpServerFeatures.SyncCompletionSpecification completion : completions) {
977+
this.completions.put(completion.referenceKey(), completion);
978+
}
979+
return this;
980+
}
981+
982+
/**
983+
* Registers multiple completions with their handlers using varargs. This method
984+
* is useful when completions are defined inline and added directly.
985+
* @param completions Array of completion specifications. Must not be null.
986+
* @return This builder instance for method chaining
987+
* @throws IllegalArgumentException if completions is null
988+
*/
989+
public SyncSpecification completions(McpServerFeatures.SyncCompletionSpecification... completions) {
990+
Assert.notNull(completions, "Completions list must not be null");
991+
for (McpServerFeatures.SyncCompletionSpecification completion : completions) {
992+
this.completions.put(completion.referenceKey(), completion);
993+
}
994+
return this;
995+
}
996+
960997
/**
961998
* Registers a consumer that will be notified when the list of roots changes. This
962999
* is useful for updating resource availability dynamically, such as when new
@@ -1023,8 +1060,8 @@ public SyncSpecification objectMapper(ObjectMapper objectMapper) {
10231060
*/
10241061
public McpSyncServer build() {
10251062
McpServerFeatures.Sync syncFeatures = new McpServerFeatures.Sync(this.serverInfo, this.serverCapabilities,
1026-
this.tools, this.resources, this.resourceTemplates, this.prompts, this.rootsChangeHandlers,
1027-
this.instructions);
1063+
this.tools, this.resources, this.resourceTemplates, this.prompts, this.completions,
1064+
this.rootsChangeHandlers, this.instructions);
10281065
McpServerFeatures.Async asyncFeatures = McpServerFeatures.Async.fromSync(syncFeatures);
10291066
var mapper = this.objectMapper != null ? this.objectMapper : new ObjectMapper();
10301067
var asyncServer = new McpAsyncServer(this.transportProvider, mapper, asyncFeatures, this.requestTimeout);

0 commit comments

Comments
 (0)