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 <christian.tzolov@broadcom.com>

Signed-off-by: jitokim <pigberger70@gmail.com>
This commit is contained in:
jitokim 2025-04-11 13:36:37 +09:00 • committed by Christian Tzolov
commit 7853efefda
2 changed files with 340 additions and 33 deletions

View file

@ -10,6 +10,7 @@ import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.BiFunction;
import java.util.function.Function;
import java.util.stream.Collectors;
@ -20,19 +21,12 @@ import io.modelcontextprotocol.client.transport.WebFluxSseClientTransport;
import io.modelcontextprotocol.server.McpServer;
import io.modelcontextprotocol.server.McpServerFeatures;
import io.modelcontextprotocol.server.TestUtil;
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;
@ -759,4 +753,53 @@ class WebFluxSseIntegrationTests {
mcpServer.close();
}
}
// ---------------------------------------
// Completion Tests
// ---------------------------------------
@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
));
AtomicReference<CompleteRequest> samplingRequest = new AtomicReference<>();
BiFunction<McpSyncServerExchange, CompleteRequest, CompleteResult> completionHandler = (mcpSyncServerExchange,
request) -> {
samplingRequest.set(request);
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 McpSchema.PromptReference("ref/prompt", "code_review"), completionHandler))
.build();
try (var mcpClient = clientBuilder.build()) {
InitializeResult initResult = mcpClient.initialize();
assertThat(initResult).isNotNull();
CompleteRequest request = new CompleteRequest(new PromptReference("ref/prompt", "code_review"),
new CompleteRequest.CompleteArgument("language", "py"));
CompleteResult result = mcpClient.completeCompletion(request);
assertThat(result).isNotNull();
assertThat(samplingRequest.get().argument().name()).isEqualTo("language");
assertThat(samplingRequest.get().argument().value()).isEqualTo("py");
assertThat(samplingRequest.get().ref().type()).isEqualTo("ref/prompt");
}
mcpServer.close();
}
}

View file

@ -300,7 +300,7 @@ public class WebMvcSseServerTransportProvider implements McpServerTransportProvi
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"));
}