Skip to content

implement completions #141

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.BiFunction;
import java.util.function.Function;
import java.util.stream.Collectors;

Expand All @@ -18,19 +19,12 @@
import io.modelcontextprotocol.client.transport.WebFluxSseClientTransport;
import io.modelcontextprotocol.server.McpServer;
import io.modelcontextprotocol.server.McpServerFeatures;
import io.modelcontextprotocol.server.McpSyncServerExchange;
import io.modelcontextprotocol.server.transport.WebFluxSseServerTransportProvider;
import io.modelcontextprotocol.spec.McpError;
import io.modelcontextprotocol.spec.McpSchema;
import io.modelcontextprotocol.spec.McpSchema.CallToolResult;
import io.modelcontextprotocol.spec.McpSchema.ClientCapabilities;
import io.modelcontextprotocol.spec.McpSchema.CreateMessageRequest;
import io.modelcontextprotocol.spec.McpSchema.CreateMessageResult;
import io.modelcontextprotocol.spec.McpSchema.InitializeResult;
import io.modelcontextprotocol.spec.McpSchema.ModelPreferences;
import io.modelcontextprotocol.spec.McpSchema.Role;
import io.modelcontextprotocol.spec.McpSchema.Root;
import io.modelcontextprotocol.spec.McpSchema.ServerCapabilities;
import io.modelcontextprotocol.spec.McpSchema.Tool;
import io.modelcontextprotocol.spec.McpSchema.*;
import io.modelcontextprotocol.spec.McpSchema.ServerCapabilities.CompletionCapabilities;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.params.ParameterizedTest;
Expand Down Expand Up @@ -620,4 +614,48 @@ void testLoggingNotification(String clientType) {
mcpServer.close();
}

}
@ParameterizedTest(name = "{0} : Completion call")
@ValueSource(strings = { "httpclient", "webflux" })
void testCompletionShouldReturnExpectedSuggestions(String clientType) {
var clientBuilder = clientBuilders.get(clientType);

var expectedValues = List.of("python", "pytorch", "pyside");
var completionResponse = new McpSchema.CompleteResult(new CompleteResult.CompleteCompletion(expectedValues, 10, // total
true // hasMore
));

BiFunction<McpSyncServerExchange, CompleteRequest, CompleteResult> completionHandler = (mcpSyncServerExchange,
request) -> {
assertThat(request.argument().name()).isEqualTo("language");
assertThat(request.argument().value()).isEqualTo("py");
assertThat(request.ref().type()).isEqualTo("ref/prompt");
return completionResponse;
};

var mcpServer = McpServer.sync(mcpServerTransportProvider)
.capabilities(ServerCapabilities.builder().completions(new CompletionCapabilities()).build())
.prompts(new McpServerFeatures.SyncPromptSpecification(
new Prompt("code_review", "this is code review prompt", List.of()),
(mcpSyncServerExchange, getPromptRequest) -> null))
.completions(new McpServerFeatures.SyncCompletionSpecification(
new McpServerFeatures.CompletionRefKey("ref/prompt", "code_review"), completionHandler))
.build();

try (var mcpClient = clientBuilder.build()) {

InitializeResult initResult = mcpClient.initialize();
assertThat(initResult).isNotNull();

CompleteRequest request = new CompleteRequest(
new CompleteRequest.PromptReference("ref/prompt", "code_review"),
new CompleteRequest.CompleteArgument("language", "py"));

CompleteResult result = mcpClient.completeCompletion(request);

assertThat(result).isNotNull();
}

mcpServer.close();
}

}
Original file line number Diff line number Diff line change
Expand Up @@ -300,7 +300,7 @@ private ServerResponse handleMessage(ServerRequest request) {
return ServerResponse.status(HttpStatus.SERVICE_UNAVAILABLE).body("Server is shutting down");
}

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

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -71,6 +71,7 @@
*
* @author Dariusz Jędrzejczyk
* @author Christian Tzolov
* @author Jihoon Kim
* @see McpClient
* @see McpSchema
* @see McpClientSession
Expand Down Expand Up @@ -801,4 +802,25 @@ void setProtocolVersions(List<String> protocolVersions) {
this.protocolVersions = protocolVersions;
}

// --------------------------
// Completions
// --------------------------
private static final TypeReference<McpSchema.CompleteResult> COMPLETION_COMPLETE_RESULT_TYPE_REF = new TypeReference<>() {
};

/**
* Sends a completion/complete request to generate value suggestions based on a given
* reference and argument. This is typically used to provide auto-completion options
* for user input fields.
* @param completeRequest The request containing the prompt or resource reference and
* argument for which to generate completions.
* @return A Mono that completes with the result containing completion suggestions.
* @see McpSchema.CompleteRequest
* @see McpSchema.CompleteResult
*/
public Mono<McpSchema.CompleteResult> completeCompletion(McpSchema.CompleteRequest completeRequest) {
return this.withInitializationCheck("complete completions", initializedResult -> this.mcpSession
.sendRequest(McpSchema.METHOD_COMPLETION_COMPLETE, completeRequest, COMPLETION_COMPLETE_RESULT_TYPE_REF));
}

}
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@
*
* @author Dariusz Jędrzejczyk
* @author Christian Tzolov
* @author Jihoon Kim
* @see McpClient
* @see McpAsyncClient
* @see McpSchema
Expand Down Expand Up @@ -325,4 +326,14 @@ public void setLoggingLevel(McpSchema.LoggingLevel loggingLevel) {
this.delegate.setLoggingLevel(loggingLevel).block();
}

/**
* Send a completion/complete request.
* @param completeRequest the completion request contains the prompt or resource
* reference and arguments for generating suggestions.
* @return the completion result containing suggested values.
*/
public McpSchema.CompleteResult completeCompletion(McpSchema.CompleteRequest completeRequest) {
return this.delegate.completeCompletion(completeRequest).block();
}

}
Original file line number Diff line number Diff line change
Expand Up @@ -68,6 +68,7 @@
*
* @author Christian Tzolov
* @author Dariusz Jędrzejczyk
* @author Jihoon Kim
* @see McpServer
* @see McpSchema
* @see McpClientSession
Expand Down Expand Up @@ -264,6 +265,8 @@ private static class AsyncServerImpl extends McpAsyncServer {

private final ConcurrentHashMap<String, McpServerFeatures.AsyncPromptSpecification> prompts = new ConcurrentHashMap<>();

private final ConcurrentHashMap<McpServerFeatures.CompletionRefKey, McpServerFeatures.AsyncCompletionSpecification> completions = new ConcurrentHashMap<>();

// FIXME: this field is deprecated and should be remvoed together with the
// broadcasting loggingNotification.
private LoggingLevel minLoggingLevel = LoggingLevel.DEBUG;
Expand All @@ -281,6 +284,7 @@ private static class AsyncServerImpl extends McpAsyncServer {
this.resources.putAll(features.resources());
this.resourceTemplates.addAll(features.resourceTemplates());
this.prompts.putAll(features.prompts());
this.completions.putAll(features.completions());

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

Expand Down Expand Up @@ -313,6 +317,11 @@ private static class AsyncServerImpl extends McpAsyncServer {
requestHandlers.put(McpSchema.METHOD_LOGGING_SET_LEVEL, setLoggerRequestHandler());
}

// Add completion API handlers if the completion capability is enabled
if (this.serverCapabilities.completions() != null) {
requestHandlers.put(McpSchema.METHOD_COMPLETION_COMPLETE, completionCompleteRequestHandler());
}

Map<String, McpServerSession.NotificationHandler> notificationHandlers = new HashMap<>();

notificationHandlers.put(McpSchema.METHOD_NOTIFICATION_INITIALIZED, (exchange, params) -> Mono.empty());
Expand Down Expand Up @@ -705,6 +714,85 @@ private McpServerSession.RequestHandler<Object> setLoggerRequestHandler() {
};
}

private McpServerSession.RequestHandler<McpSchema.CompleteResult> completionCompleteRequestHandler() {
return (exchange, params) -> {
McpSchema.CompleteRequest request = parseCompletionParams(params);

if (request.ref() == null) {
return Mono.error(new McpError("ref must not be null"));
}

if (request.ref().type() == null) {
return Mono.error(new McpError("type must not be null"));
}

String type = request.ref().type();

// check if the referenced resource exists
if (type.equals("ref/prompt")
&& request.ref() instanceof McpSchema.CompleteRequest.PromptReference promptReference) {
McpServerFeatures.AsyncPromptSpecification prompt = this.prompts.get(promptReference.name());
if (prompt == null) {
return Mono.error(new McpError("Prompt not found: " + promptReference.name()));
}
}

if (type.equals("ref/resource")
&& request.ref() instanceof McpSchema.CompleteRequest.ResourceReference resourceReference) {
McpServerFeatures.AsyncResourceSpecification resource = this.resources.get(resourceReference.uri());
if (resource == null) {
return Mono.error(new McpError("Resource not found: " + resourceReference.uri()));
}
}

McpServerFeatures.CompletionRefKey key = McpServerFeatures.CompletionRefKey.from(request);
McpServerFeatures.AsyncCompletionSpecification specification = this.completions.get(key);

if (specification == null) {
return Mono.error(new McpError("AsyncCompletionSpecification not found: " + key));
}

return specification.completionHandler().apply(exchange, request);
};
}

/**
* Parses the raw JSON-RPC request parameters into a
* {@link McpSchema.CompleteRequest} object.
* <p>
* This method manually extracts the `ref` and `argument` fields from the input
* map, determines the correct reference type (either prompt or resource), and
* constructs a fully-typed {@code CompleteRequest} instance.
* @param object the raw request parameters, expected to be a Map containing "ref"
* and "argument" entries.
* @return a {@link McpSchema.CompleteRequest} representing the structured
* completion request.
* @throws IllegalArgumentException if the "ref" type is not recognized.
*/
@SuppressWarnings("unchecked")
private McpSchema.CompleteRequest parseCompletionParams(Object object) {
Map<String, Object> params = (Map<String, Object>) object;
Map<String, Object> refMap = (Map<String, Object>) params.get("ref");
Map<String, Object> argMap = (Map<String, Object>) params.get("argument");

String refType = (String) refMap.get("type");

McpSchema.CompleteRequest.PromptOrResourceReference ref = switch (refType) {
case "ref/prompt" ->
new McpSchema.CompleteRequest.PromptReference(refType, (String) refMap.get("name"));
case "ref/resource" ->
new McpSchema.CompleteRequest.ResourceReference(refType, (String) refMap.get("uri"));
default -> throw new IllegalArgumentException("Invalid ref type: " + refType);
};

String argName = (String) argMap.get("name");
String argValue = (String) argMap.get("value");
McpSchema.CompleteRequest.CompleteArgument argument = new McpSchema.CompleteRequest.CompleteArgument(
argName, argValue);

return new McpSchema.CompleteRequest(ref, argument);
}

// ---------------------------------------
// Sampling
// ---------------------------------------
Expand Down
43 changes: 40 additions & 3 deletions mcp/src/main/java/io/modelcontextprotocol/server/McpServer.java
Original file line number Diff line number Diff line change
Expand Up @@ -114,6 +114,7 @@
*
* @author Christian Tzolov
* @author Dariusz Jędrzejczyk
* @author Jihoon Kim
* @see McpAsyncServer
* @see McpSyncServer
* @see McpServerTransportProvider
Expand Down Expand Up @@ -191,6 +192,8 @@ class AsyncSpecification {
*/
private final Map<String, McpServerFeatures.AsyncPromptSpecification> prompts = new HashMap<>();

private final Map<McpServerFeatures.CompletionRefKey, McpServerFeatures.AsyncCompletionSpecification> completions = new HashMap<>();

private final List<BiFunction<McpAsyncServerExchange, List<McpSchema.Root>, Mono<Void>>> rootsChangeHandlers = new ArrayList<>();

private AsyncSpecification(McpServerTransportProvider transportProvider) {
Expand Down Expand Up @@ -563,7 +566,8 @@ public AsyncSpecification objectMapper(ObjectMapper objectMapper) {
*/
public McpAsyncServer build() {
var features = new McpServerFeatures.Async(this.serverInfo, this.serverCapabilities, this.tools,
this.resources, this.resourceTemplates, this.prompts, this.rootsChangeHandlers, this.instructions);
this.resources, this.resourceTemplates, this.prompts, this.completions, this.rootsChangeHandlers,
this.instructions);
var mapper = this.objectMapper != null ? this.objectMapper : new ObjectMapper();
return new McpAsyncServer(this.transportProvider, mapper, features);
}
Expand Down Expand Up @@ -617,6 +621,8 @@ class SyncSpecification {
*/
private final Map<String, McpServerFeatures.SyncPromptSpecification> prompts = new HashMap<>();

private final Map<McpServerFeatures.CompletionRefKey, McpServerFeatures.SyncCompletionSpecification> completions = new HashMap<>();

private final List<BiConsumer<McpSyncServerExchange, List<McpSchema.Root>>> rootsChangeHandlers = new ArrayList<>();

private SyncSpecification(McpServerTransportProvider transportProvider) {
Expand Down Expand Up @@ -922,6 +928,37 @@ public SyncSpecification prompts(McpServerFeatures.SyncPromptSpecification... pr
return this;
}

/**
* Registers multiple completions with their handlers using a List. This method is
* useful when completions need to be added in bulk from a collection.
* @param completions List of completion specifications. Must not be null.
* @return This builder instance for method chaining
* @throws IllegalArgumentException if completions is null
* @see #completions(McpServerFeatures.SyncCompletionSpecification...)
*/
public SyncSpecification completions(List<McpServerFeatures.SyncCompletionSpecification> completions) {
Assert.notNull(completions, "Completions list must not be null");
for (McpServerFeatures.SyncCompletionSpecification completion : completions) {
this.completions.put(completion.referenceKey(), completion);
}
return this;
}

/**
* Registers multiple completions with their handlers using varargs. This method
* is useful when completions are defined inline and added directly.
* @param completions Array of completion specifications. Must not be null.
* @return This builder instance for method chaining
* @throws IllegalArgumentException if completions is null
*/
public SyncSpecification completions(McpServerFeatures.SyncCompletionSpecification... completions) {
Assert.notNull(completions, "Completions list must not be null");
for (McpServerFeatures.SyncCompletionSpecification completion : completions) {
this.completions.put(completion.referenceKey(), completion);
}
return this;
}

/**
* Registers a consumer that will be notified when the list of roots changes. This
* is useful for updating resource availability dynamically, such as when new
Expand Down Expand Up @@ -988,8 +1025,8 @@ public SyncSpecification objectMapper(ObjectMapper objectMapper) {
*/
public McpSyncServer build() {
McpServerFeatures.Sync syncFeatures = new McpServerFeatures.Sync(this.serverInfo, this.serverCapabilities,
this.tools, this.resources, this.resourceTemplates, this.prompts, this.rootsChangeHandlers,
this.instructions);
this.tools, this.resources, this.resourceTemplates, this.prompts, this.completions,
this.rootsChangeHandlers, this.instructions);
McpServerFeatures.Async asyncFeatures = McpServerFeatures.Async.fromSync(syncFeatures);
var mapper = this.objectMapper != null ? this.objectMapper : new ObjectMapper();
var asyncServer = new McpAsyncServer(this.transportProvider, mapper, asyncFeatures);
Expand Down
Loading