Skip to content

Commit e0cb707

Browse files
jitokimjitokim
committed
implement handler to support serverCapabilities completions
Signed-off-by: jitokim <pigberger70@gmail.com>
1 parent 2895d15 commit e0cb707

File tree

7 files changed

+138
-11
lines changed

7 files changed

+138
-11
lines changed

‎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/server/McpAsyncServer.java

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,7 @@
6767
*
6868
* @author Christian Tzolov
6969
* @author Dariusz Jędrzejczyk
70+
* @author Jihoon Kim
7071
* @see McpServer
7172
* @see McpSchema
7273
* @see McpClientSession
@@ -257,6 +258,8 @@ private static class AsyncServerImpl extends McpAsyncServer {
257258

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

261+
private final ConcurrentHashMap<McpServerFeatures.CompletionRefKey, McpServerFeatures.AsyncCompletionSpecification> completions = new ConcurrentHashMap<>();
262+
260263
private LoggingLevel minLoggingLevel = LoggingLevel.DEBUG;
261264

262265
private List<String> protocolVersions = List.of(McpSchema.LATEST_PROTOCOL_VERSION);
@@ -272,6 +275,7 @@ private static class AsyncServerImpl extends McpAsyncServer {
272275
this.resources.putAll(features.resources());
273276
this.resourceTemplates.addAll(features.resourceTemplates());
274277
this.prompts.putAll(features.prompts());
278+
this.completions.putAll(features.completions());
275279

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

@@ -304,6 +308,11 @@ private static class AsyncServerImpl extends McpAsyncServer {
304308
requestHandlers.put(McpSchema.METHOD_LOGGING_SET_LEVEL, setLoggerRequestHandler());
305309
}
306310

311+
// Add completion API handlers if the completion capability is enabled
312+
if (this.serverCapabilities.completions() != null) {
313+
requestHandlers.put(McpSchema.METHOD_COMPLETION_COMPLETE, completionCompleteRequestHandler());
314+
}
315+
307316
Map<String, McpServerSession.NotificationHandler> notificationHandlers = new HashMap<>();
308317

309318
notificationHandlers.put(McpSchema.METHOD_NOTIFICATION_INITIALIZED, (exchange, params) -> Mono.empty());
@@ -686,6 +695,47 @@ private McpServerSession.RequestHandler<Void> setLoggerRequestHandler() {
686695
};
687696
}
688697

698+
private McpServerSession.RequestHandler<McpSchema.CompleteResult> completionCompleteRequestHandler() {
699+
return (exchange, params) -> {
700+
McpSchema.CompleteRequest request = objectMapper.convertValue(params, McpSchema.CompleteRequest.class);
701+
702+
if (request.ref() == null) {
703+
return Mono.error(new McpError("ref must not be null"));
704+
}
705+
706+
if (request.ref().type() == null) {
707+
return Mono.error(new McpError("type must not be null"));
708+
}
709+
710+
String type = request.ref().type();
711+
712+
if (type.equals("ref/prompt")
713+
&& request.ref() instanceof McpSchema.CompleteRequest.PromptReference promptReference) {
714+
McpServerFeatures.AsyncPromptSpecification prompt = this.prompts.get(promptReference.name());
715+
if (prompt == null) {
716+
return Mono.error(new McpError("Prompt not found: " + promptReference.name()));
717+
}
718+
}
719+
720+
if (type.equals("ref/resource")
721+
&& request.ref() instanceof McpSchema.CompleteRequest.ResourceReference resourceReference) {
722+
McpServerFeatures.AsyncResourceSpecification resource = this.resources.get(resourceReference.uri());
723+
if (resource == null) {
724+
return Mono.error(new McpError("Resource not found: " + resourceReference.uri()));
725+
}
726+
}
727+
728+
McpServerFeatures.CompletionRefKey key = McpServerFeatures.CompletionRefKey.from(request);
729+
McpServerFeatures.AsyncCompletionSpecification specification = this.completions.get(key);
730+
731+
if (specification == null) {
732+
return Mono.error(new McpError("AsyncCompletionSpecification not found: " + key));
733+
}
734+
735+
return specification.completionHandler().apply(exchange, request);
736+
};
737+
}
738+
689739
// ---------------------------------------
690740
// Sampling
691741
// ---------------------------------------

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

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

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

196199
private AsyncSpecification(McpServerTransportProvider transportProvider) {
@@ -563,7 +566,8 @@ public AsyncSpecification objectMapper(ObjectMapper objectMapper) {
563566
*/
564567
public McpAsyncServer build() {
565568
var features = new McpServerFeatures.Async(this.serverInfo, this.serverCapabilities, this.tools,
566-
this.resources, this.resourceTemplates, this.prompts, this.rootsChangeHandlers, this.instructions);
569+
this.resources, this.resourceTemplates, this.prompts, this.completions, this.rootsChangeHandlers,
570+
this.instructions);
567571
var mapper = this.objectMapper != null ? this.objectMapper : new ObjectMapper();
568572
return new McpAsyncServer(this.transportProvider, mapper, features);
569573
}
@@ -617,6 +621,8 @@ class SyncSpecification {
617621
*/
618622
private final Map<String, McpServerFeatures.SyncPromptSpecification> prompts = new HashMap<>();
619623

624+
private final Map<McpServerFeatures.CompletionRefKey, McpServerFeatures.SyncCompletionSpecification> completions = new HashMap<>();
625+
620626
private final List<BiConsumer<McpSyncServerExchange, List<McpSchema.Root>>> rootsChangeHandlers = new ArrayList<>();
621627

622628
private SyncSpecification(McpServerTransportProvider transportProvider) {
@@ -922,6 +928,22 @@ public SyncSpecification prompts(McpServerFeatures.SyncPromptSpecification... pr
922928
return this;
923929
}
924930

931+
public SyncSpecification completions(List<McpServerFeatures.SyncCompletionSpecification> completions) {
932+
Assert.notNull(completions, "Completions list must not be null");
933+
for (McpServerFeatures.SyncCompletionSpecification completion : completions) {
934+
this.completions.put(completion.referenceKey(), completion);
935+
}
936+
return this;
937+
}
938+
939+
public SyncSpecification completions(McpServerFeatures.SyncCompletionSpecification... completions) {
940+
Assert.notNull(completions, "Completions list must not be null");
941+
for (McpServerFeatures.SyncCompletionSpecification completion : completions) {
942+
this.completions.put(completion.referenceKey(), completion);
943+
}
944+
return this;
945+
}
946+
925947
/**
926948
* Registers a consumer that will be notified when the list of roots changes. This
927949
* is useful for updating resource availability dynamically, such as when new
@@ -988,8 +1010,8 @@ public SyncSpecification objectMapper(ObjectMapper objectMapper) {
9881010
*/
9891011
public McpSyncServer build() {
9901012
McpServerFeatures.Sync syncFeatures = new McpServerFeatures.Sync(this.serverInfo, this.serverCapabilities,
991-
this.tools, this.resources, this.resourceTemplates, this.prompts, this.rootsChangeHandlers,
992-
this.instructions);
1013+
this.tools, this.resources, this.resourceTemplates, this.prompts, this.completions,
1014+
this.rootsChangeHandlers, this.instructions);
9931015
McpServerFeatures.Async asyncFeatures = McpServerFeatures.Async.fromSync(syncFeatures);
9941016
var mapper = this.objectMapper != null ? this.objectMapper : new ObjectMapper();
9951017
var asyncServer = new McpAsyncServer(this.transportProvider, mapper, asyncFeatures);

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

Lines changed: 49 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
* MCP server features specification that a particular server can choose to support.
2222
*
2323
* @author Dariusz Jędrzejczyk
24+
* @author Jihoon Kim
2425
*/
2526
public class McpServerFeatures {
2627

@@ -41,6 +42,7 @@ record Async(McpSchema.Implementation serverInfo, McpSchema.ServerCapabilities s
4142
List<McpServerFeatures.AsyncToolSpecification> tools, Map<String, AsyncResourceSpecification> resources,
4243
List<McpSchema.ResourceTemplate> resourceTemplates,
4344
Map<String, McpServerFeatures.AsyncPromptSpecification> prompts,
45+
Map<CompletionRefKey, McpServerFeatures.AsyncCompletionSpecification> completions,
4446
List<BiFunction<McpAsyncServerExchange, List<McpSchema.Root>, Mono<Void>>> rootsChangeConsumers,
4547
String instructions) {
4648

@@ -60,14 +62,16 @@ record Async(McpSchema.Implementation serverInfo, McpSchema.ServerCapabilities s
6062
List<McpServerFeatures.AsyncToolSpecification> tools, Map<String, AsyncResourceSpecification> resources,
6163
List<McpSchema.ResourceTemplate> resourceTemplates,
6264
Map<String, McpServerFeatures.AsyncPromptSpecification> prompts,
65+
Map<CompletionRefKey, McpServerFeatures.AsyncCompletionSpecification> completions,
6366
List<BiFunction<McpAsyncServerExchange, List<McpSchema.Root>, Mono<Void>>> rootsChangeConsumers,
6467
String instructions) {
6568

6669
Assert.notNull(serverInfo, "Server info must not be null");
6770

6871
this.serverInfo = serverInfo;
6972
this.serverCapabilities = (serverCapabilities != null) ? serverCapabilities
70-
: new McpSchema.ServerCapabilities(null, // experimental
73+
: new McpSchema.ServerCapabilities(null, // completions
74+
null, // experimental
7175
new McpSchema.ServerCapabilities.LoggingCapabilities(), // Enable
7276
// logging
7377
// by
@@ -81,6 +85,7 @@ record Async(McpSchema.Implementation serverInfo, McpSchema.ServerCapabilities s
8185
this.resources = (resources != null) ? resources : Map.of();
8286
this.resourceTemplates = (resourceTemplates != null) ? resourceTemplates : List.of();
8387
this.prompts = (prompts != null) ? prompts : Map.of();
88+
this.completions = (completions != null) ? completions : Map.of();
8489
this.rootsChangeConsumers = (rootsChangeConsumers != null) ? rootsChangeConsumers : List.of();
8590
this.instructions = instructions;
8691
}
@@ -109,6 +114,11 @@ static Async fromSync(Sync syncSpec) {
109114
prompts.put(key, AsyncPromptSpecification.fromSync(prompt));
110115
});
111116

117+
Map<CompletionRefKey, McpServerFeatures.AsyncCompletionSpecification> completions = new HashMap<>();
118+
syncSpec.completions().forEach((key, completion) -> {
119+
completions.put(key, AsyncCompletionSpecification.fromSync(completion));
120+
});
121+
112122
List<BiFunction<McpAsyncServerExchange, List<McpSchema.Root>, Mono<Void>>> rootChangeConsumers = new ArrayList<>();
113123

114124
for (var rootChangeConsumer : syncSpec.rootsChangeConsumers()) {
@@ -118,7 +128,7 @@ static Async fromSync(Sync syncSpec) {
118128
}
119129

120130
return new Async(syncSpec.serverInfo(), syncSpec.serverCapabilities(), tools, resources,
121-
syncSpec.resourceTemplates(), prompts, rootChangeConsumers, syncSpec.instructions());
131+
syncSpec.resourceTemplates(), prompts, completions, rootChangeConsumers, syncSpec.instructions());
122132
}
123133
}
124134

@@ -140,6 +150,7 @@ record Sync(McpSchema.Implementation serverInfo, McpSchema.ServerCapabilities se
140150
Map<String, McpServerFeatures.SyncResourceSpecification> resources,
141151
List<McpSchema.ResourceTemplate> resourceTemplates,
142152
Map<String, McpServerFeatures.SyncPromptSpecification> prompts,
153+
Map<CompletionRefKey, McpServerFeatures.SyncCompletionSpecification> completions,
143154
List<BiConsumer<McpSyncServerExchange, List<McpSchema.Root>>> rootsChangeConsumers, String instructions) {
144155

145156
/**
@@ -159,14 +170,16 @@ record Sync(McpSchema.Implementation serverInfo, McpSchema.ServerCapabilities se
159170
Map<String, McpServerFeatures.SyncResourceSpecification> resources,
160171
List<McpSchema.ResourceTemplate> resourceTemplates,
161172
Map<String, McpServerFeatures.SyncPromptSpecification> prompts,
173+
Map<CompletionRefKey, McpServerFeatures.SyncCompletionSpecification> completions,
162174
List<BiConsumer<McpSyncServerExchange, List<McpSchema.Root>>> rootsChangeConsumers,
163175
String instructions) {
164176

165177
Assert.notNull(serverInfo, "Server info must not be null");
166178

167179
this.serverInfo = serverInfo;
168180
this.serverCapabilities = (serverCapabilities != null) ? serverCapabilities
169-
: new McpSchema.ServerCapabilities(null, // experimental
181+
: new McpSchema.ServerCapabilities(null, // completions
182+
null, // experimental
170183
new McpSchema.ServerCapabilities.LoggingCapabilities(), // Enable
171184
// logging
172185
// by
@@ -180,6 +193,7 @@ record Sync(McpSchema.Implementation serverInfo, McpSchema.ServerCapabilities se
180193
this.resources = (resources != null) ? resources : new HashMap<>();
181194
this.resourceTemplates = (resourceTemplates != null) ? resourceTemplates : new ArrayList<>();
182195
this.prompts = (prompts != null) ? prompts : new HashMap<>();
196+
this.completions = (completions != null) ? completions : new HashMap<>();
183197
this.rootsChangeConsumers = (rootsChangeConsumers != null) ? rootsChangeConsumers : new ArrayList<>();
184198
this.instructions = instructions;
185199
}
@@ -325,6 +339,19 @@ static AsyncPromptSpecification fromSync(SyncPromptSpecification prompt) {
325339
}
326340
}
327341

342+
public record AsyncCompletionSpecification(
343+
BiFunction<McpAsyncServerExchange, McpSchema.CompleteRequest, Mono<McpSchema.CompleteResult>> completionHandler) {
344+
345+
static AsyncCompletionSpecification fromSync(SyncCompletionSpecification completion) {
346+
if (completion == null) {
347+
return null;
348+
}
349+
return new AsyncCompletionSpecification((exchange, request) -> Mono
350+
.fromCallable(() -> completion.completionHandler().apply(new McpSyncServerExchange(exchange), request))
351+
.subscribeOn(Schedulers.boundedElastic()));
352+
}
353+
}
354+
328355
/**
329356
* Specification of a tool with its synchronous handler function. Tools are the
330357
* primary way for MCP servers to expose functionality to AI models. Each tool
@@ -431,4 +458,23 @@ public record SyncPromptSpecification(McpSchema.Prompt prompt,
431458
BiFunction<McpSyncServerExchange, McpSchema.GetPromptRequest, McpSchema.GetPromptResult> promptHandler) {
432459
}
433460

461+
public record SyncCompletionSpecification(CompletionRefKey referenceKey,
462+
BiFunction<McpSyncServerExchange, McpSchema.CompleteRequest, McpSchema.CompleteResult> completionHandler) {
463+
}
464+
465+
public record CompletionRefKey(String type, String identifier) {
466+
public static CompletionRefKey from(McpSchema.CompleteRequest request) {
467+
var ref = request.ref();
468+
if (ref instanceof McpSchema.CompleteRequest.PromptReference pr) {
469+
return new CompletionRefKey(ref.type(), pr.name());
470+
}
471+
else if (ref instanceof McpSchema.CompleteRequest.ResourceReference rr) {
472+
return new CompletionRefKey(ref.type(), rr.uri());
473+
}
474+
else {
475+
throw new IllegalArgumentException("Unsupported reference type: " + ref);
476+
}
477+
}
478+
}
479+
434480
}

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

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,7 @@
44

55
package io.modelcontextprotocol.server;
66

7-
import io.modelcontextprotocol.spec.McpError;
87
import io.modelcontextprotocol.spec.McpSchema;
9-
import io.modelcontextprotocol.spec.McpSchema.ClientCapabilities;
108
import io.modelcontextprotocol.spec.McpSchema.LoggingMessageNotification;
119
import io.modelcontextprotocol.util.Assert;
1210

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

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
11
package io.modelcontextprotocol.server;
22

3-
import com.fasterxml.jackson.core.type.TypeReference;
43
import io.modelcontextprotocol.spec.McpSchema;
54

65
/**

‎mcp/src/main/java/io/modelcontextprotocol/spec/McpSchema.java

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,8 @@ private McpSchema() {
7979

8080
public static final String METHOD_NOTIFICATION_PROMPTS_LIST_CHANGED = "notifications/prompts/list_changed";
8181

82+
public static final String METHOD_COMPLETION_COMPLETE = "completions/complete";
83+
8284
// Logging Methods
8385
public static final String METHOD_LOGGING_SET_LEVEL = "logging/setLevel";
8486

@@ -314,12 +316,16 @@ public ClientCapabilities build() {
314316
@JsonInclude(JsonInclude.Include.NON_ABSENT)
315317
@JsonIgnoreProperties(ignoreUnknown = true)
316318
public record ServerCapabilities( // @formatter:off
319+
@JsonProperty("completions") CompletionCapabilities completions,
317320
@JsonProperty("experimental") Map<String, Object> experimental,
318321
@JsonProperty("logging") LoggingCapabilities logging,
319322
@JsonProperty("prompts") PromptCapabilities prompts,
320323
@JsonProperty("resources") ResourceCapabilities resources,
321324
@JsonProperty("tools") ToolCapabilities tools) {
322325

326+
@JsonInclude(JsonInclude.Include.NON_ABSENT)
327+
public record CompletionCapabilities() {
328+
}
323329

324330
@JsonInclude(JsonInclude.Include.NON_ABSENT)
325331
public record LoggingCapabilities() {
@@ -347,12 +353,18 @@ public static Builder builder() {
347353

348354
public static class Builder {
349355

356+
private CompletionCapabilities completions;
350357
private Map<String, Object> experimental;
351358
private LoggingCapabilities logging = new LoggingCapabilities();
352359
private PromptCapabilities prompts;
353360
private ResourceCapabilities resources;
354361
private ToolCapabilities tools;
355362

363+
public Builder completions(CompletionCapabilities completions) {
364+
this.completions = completions;
365+
return this;
366+
}
367+
356368
public Builder experimental(Map<String, Object> experimental) {
357369
this.experimental = experimental;
358370
return this;
@@ -379,7 +391,7 @@ public Builder tools(Boolean listChanged) {
379391
}
380392

381393
public ServerCapabilities build() {
382-
return new ServerCapabilities(experimental, logging, prompts, resources, tools);
394+
return new ServerCapabilities(completions, experimental, logging, prompts, resources, tools);
383395
}
384396
}
385397
} // @formatter:on

0 commit comments

Comments
 (0)