Skip to content

Commit c12df84

Browse files
tzolovtzolov
committed
Replace McpServerFeatures.CompletionRefKey by McpSchemaCompleteReference
Signed-off-by: Christian Tzolov <christian.tzolov@broadcom.com>
1 parent 784ef34 commit c12df84

5 files changed

Lines changed: 14 additions & 54 deletions

File tree

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -779,7 +779,7 @@ void testCompletionShouldReturnExpectedSuggestions(String clientType) {
779779
new Prompt("code_review", "this is code review prompt", List.of()),
780780
(mcpSyncServerExchange, getPromptRequest) -> null))
781781
.completions(new McpServerFeatures.SyncCompletionSpecification(
782-
new McpServerFeatures.CompletionRefKey("ref/prompt", "code_review"), completionHandler))
782+
new McpSchema.PromptReference("ref/prompt", "code_review"), completionHandler))
783783
.build();
784784

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

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

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -270,7 +270,7 @@ private static class AsyncServerImpl extends McpAsyncServer {
270270
// broadcasting loggingNotification.
271271
private LoggingLevel minLoggingLevel = LoggingLevel.DEBUG;
272272

273-
private final ConcurrentHashMap<McpServerFeatures.CompletionRefKey, McpServerFeatures.AsyncCompletionSpecification> completions = new ConcurrentHashMap<>();
273+
private final ConcurrentHashMap<McpSchema.CompleteReference, McpServerFeatures.AsyncCompletionSpecification> completions = new ConcurrentHashMap<>();
274274

275275
private List<String> protocolVersions = List.of(McpSchema.LATEST_PROTOCOL_VERSION);
276276

@@ -745,11 +745,10 @@ private McpServerSession.RequestHandler<McpSchema.CompleteResult> completionComp
745745
}
746746
}
747747

748-
McpServerFeatures.CompletionRefKey key = McpServerFeatures.CompletionRefKey.from(request);
749-
McpServerFeatures.AsyncCompletionSpecification specification = this.completions.get(key);
748+
McpServerFeatures.AsyncCompletionSpecification specification = this.completions.get(request.ref());
750749

751750
if (specification == null) {
752-
return Mono.error(new McpError("AsyncCompletionSpecification not found: " + key));
751+
return Mono.error(new McpError("AsyncCompletionSpecification not found: " + request.ref()));
753752
}
754753

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

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -193,7 +193,7 @@ class AsyncSpecification {
193193
*/
194194
private final Map<String, McpServerFeatures.AsyncPromptSpecification> prompts = new HashMap<>();
195195

196-
private final Map<McpServerFeatures.CompletionRefKey, McpServerFeatures.AsyncCompletionSpecification> completions = new HashMap<>();
196+
private final Map<McpSchema.CompleteReference, McpServerFeatures.AsyncCompletionSpecification> completions = new HashMap<>();
197197

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

@@ -639,7 +639,7 @@ class SyncSpecification {
639639
*/
640640
private final Map<String, McpServerFeatures.SyncPromptSpecification> prompts = new HashMap<>();
641641

642-
private final Map<McpServerFeatures.CompletionRefKey, McpServerFeatures.SyncCompletionSpecification> completions = new HashMap<>();
642+
private final Map<McpSchema.CompleteReference, McpServerFeatures.SyncCompletionSpecification> completions = new HashMap<>();
643643

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

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

Lines changed: 7 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@ record Async(McpSchema.Implementation serverInfo, McpSchema.ServerCapabilities s
4242
List<McpServerFeatures.AsyncToolSpecification> tools, Map<String, AsyncResourceSpecification> resources,
4343
List<McpSchema.ResourceTemplate> resourceTemplates,
4444
Map<String, McpServerFeatures.AsyncPromptSpecification> prompts,
45-
Map<CompletionRefKey, McpServerFeatures.AsyncCompletionSpecification> completions,
45+
Map<McpSchema.CompleteReference, McpServerFeatures.AsyncCompletionSpecification> completions,
4646
List<BiFunction<McpAsyncServerExchange, List<McpSchema.Root>, Mono<Void>>> rootsChangeConsumers,
4747
String instructions) {
4848

@@ -62,7 +62,7 @@ record Async(McpSchema.Implementation serverInfo, McpSchema.ServerCapabilities s
6262
List<McpServerFeatures.AsyncToolSpecification> tools, Map<String, AsyncResourceSpecification> resources,
6363
List<McpSchema.ResourceTemplate> resourceTemplates,
6464
Map<String, McpServerFeatures.AsyncPromptSpecification> prompts,
65-
Map<CompletionRefKey, McpServerFeatures.AsyncCompletionSpecification> completions,
65+
Map<McpSchema.CompleteReference, McpServerFeatures.AsyncCompletionSpecification> completions,
6666
List<BiFunction<McpAsyncServerExchange, List<McpSchema.Root>, Mono<Void>>> rootsChangeConsumers,
6767
String instructions) {
6868

@@ -114,7 +114,7 @@ static Async fromSync(Sync syncSpec) {
114114
prompts.put(key, AsyncPromptSpecification.fromSync(prompt));
115115
});
116116

117-
Map<CompletionRefKey, McpServerFeatures.AsyncCompletionSpecification> completions = new HashMap<>();
117+
Map<McpSchema.CompleteReference, McpServerFeatures.AsyncCompletionSpecification> completions = new HashMap<>();
118118
syncSpec.completions().forEach((key, completion) -> {
119119
completions.put(key, AsyncCompletionSpecification.fromSync(completion));
120120
});
@@ -150,7 +150,7 @@ record Sync(McpSchema.Implementation serverInfo, McpSchema.ServerCapabilities se
150150
Map<String, McpServerFeatures.SyncResourceSpecification> resources,
151151
List<McpSchema.ResourceTemplate> resourceTemplates,
152152
Map<String, McpServerFeatures.SyncPromptSpecification> prompts,
153-
Map<CompletionRefKey, McpServerFeatures.SyncCompletionSpecification> completions,
153+
Map<McpSchema.CompleteReference, McpServerFeatures.SyncCompletionSpecification> completions,
154154
List<BiConsumer<McpSyncServerExchange, List<McpSchema.Root>>> rootsChangeConsumers, String instructions) {
155155

156156
/**
@@ -170,7 +170,7 @@ record Sync(McpSchema.Implementation serverInfo, McpSchema.ServerCapabilities se
170170
Map<String, McpServerFeatures.SyncResourceSpecification> resources,
171171
List<McpSchema.ResourceTemplate> resourceTemplates,
172172
Map<String, McpServerFeatures.SyncPromptSpecification> prompts,
173-
Map<CompletionRefKey, McpServerFeatures.SyncCompletionSpecification> completions,
173+
Map<McpSchema.CompleteReference, McpServerFeatures.SyncCompletionSpecification> completions,
174174
List<BiConsumer<McpSyncServerExchange, List<McpSchema.Root>>> rootsChangeConsumers,
175175
String instructions) {
176176

@@ -355,7 +355,7 @@ static AsyncPromptSpecification fromSync(SyncPromptSpecification prompt) {
355355
* {@link McpAsyncServerExchange} used to interact with the client. The second
356356
* argument is a {@link io.modelcontextprotocol.spec.McpSchema.CompleteRequest}.
357357
*/
358-
public record AsyncCompletionSpecification(CompletionRefKey referenceKey,
358+
public record AsyncCompletionSpecification(McpSchema.CompleteReference referenceKey,
359359
BiFunction<McpAsyncServerExchange, McpSchema.CompleteRequest, Mono<McpSchema.CompleteResult>> completionHandler) {
360360

361361
/**
@@ -492,46 +492,8 @@ public record SyncPromptSpecification(McpSchema.Prompt prompt,
492492
* {@link McpSyncServerExchange} used to interact with the client. The second argument
493493
* is a {@link io.modelcontextprotocol.spec.McpSchema.CompleteRequest}.
494494
*/
495-
public record SyncCompletionSpecification(CompletionRefKey referenceKey,
495+
public record SyncCompletionSpecification(McpSchema.CompleteReference referenceKey,
496496
BiFunction<McpSyncServerExchange, McpSchema.CompleteRequest, McpSchema.CompleteResult> completionHandler) {
497497
}
498498

499-
/**
500-
* A unique key representing a completion reference, composed of its type and
501-
* identifier. This key is used to look up asynchronous completion specifications in a
502-
* map-like structure.
503-
*
504-
* <p>
505-
* The {@code type} typically corresponds to the kind of reference, such as
506-
* {@code "ref/prompt"} or {@code "ref/resource"}, while the {@code identifier} is the
507-
* name or URI associated with the specific reference.
508-
*
509-
* @param type the reference type (e.g., "ref/prompt", "ref/resource")
510-
* @param identifier the reference identifier (e.g., prompt name or resource URI)
511-
*/
512-
public record CompletionRefKey(String type, String identifier) {
513-
514-
/**
515-
* Creates a {@code CompletionRefKey} from a {@link McpSchema.CompleteRequest}.
516-
* The key is derived from the request's reference type and its associated name or
517-
* URI.
518-
* @param request the {@code CompleteRequest} containing a prompt or resource
519-
* reference
520-
* @return a unique key based on the request's reference
521-
* @throws IllegalArgumentException if the reference type is unsupported
522-
*/
523-
public static CompletionRefKey from(McpSchema.CompleteRequest request) {
524-
var ref = request.ref();
525-
if (ref instanceof McpSchema.PromptReference pr) {
526-
return new CompletionRefKey(ref.type(), pr.name());
527-
}
528-
else if (ref instanceof McpSchema.ResourceReference rr) {
529-
return new CompletionRefKey(ref.type(), rr.uri());
530-
}
531-
else {
532-
throw new IllegalArgumentException("Unsupported reference type: " + ref);
533-
}
534-
}
535-
}
536-
537499
}

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

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1239,8 +1239,7 @@ public record CompleteArgument(
12391239

12401240
@JsonInclude(JsonInclude.Include.NON_ABSENT)
12411241
@JsonIgnoreProperties(ignoreUnknown = true)
1242-
public record CompleteResult(
1243-
@JsonProperty("values") CompleteCompletion completion) { // @formatter:off
1242+
public record CompleteResult(@JsonProperty("values") CompleteCompletion completion) { // @formatter:off
12441243

12451244
public record CompleteCompletion(
12461245
@JsonProperty("values") List<String> values,

0 commit comments

Comments
 (0)