Skip to content

Commit dc477fc

Browse files
bzsurbhitzolov
bzsurbhi
authored and
tzolov
committed
feat: add context support to CompleteRequest for dependent completions
- Add CompleteContext record to support contextual completion scenarios - Update CompleteRequest to include context and meta fields with backward-compatible constructors - Modify McpAsyncServer to extract and pass context/meta from request parameters - Add tests for context-aware completions including dependent scenarios - Maintain backward compatibility for existing completion handlers Signed-off-by: Christian Tzolov <[email protected]>
1 parent c8c08dd commit dc477fc

File tree

4 files changed

+336
-12
lines changed

4 files changed

+336
-12
lines changed

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

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -721,6 +721,8 @@ private McpSchema.CompleteRequest parseCompletionParams(Object object) {
721721
Map<String, Object> params = (Map<String, Object>) object;
722722
Map<String, Object> refMap = (Map<String, Object>) params.get("ref");
723723
Map<String, Object> argMap = (Map<String, Object>) params.get("argument");
724+
Map<String, Object> contextMap = (Map<String, Object>) params.get("context");
725+
Map<String, Object> meta = (Map<String, Object>) params.get("_meta");
724726

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

@@ -736,7 +738,13 @@ private McpSchema.CompleteRequest parseCompletionParams(Object object) {
736738
McpSchema.CompleteRequest.CompleteArgument argument = new McpSchema.CompleteRequest.CompleteArgument(argName,
737739
argValue);
738740

739-
return new McpSchema.CompleteRequest(ref, argument);
741+
McpSchema.CompleteRequest.CompleteContext context = null;
742+
if (contextMap != null) {
743+
Map<String, String> arguments = (Map<String, String>) contextMap.get("arguments");
744+
context = new McpSchema.CompleteRequest.CompleteContext(arguments);
745+
}
746+
747+
return new McpSchema.CompleteRequest(ref, argument, meta, context);
740748
}
741749

742750
/**

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

Lines changed: 23 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1755,18 +1755,31 @@ public String identifier() {
17551755
@JsonInclude(JsonInclude.Include.NON_ABSENT)
17561756
@JsonIgnoreProperties(ignoreUnknown = true)
17571757
public record CompleteRequest(// @formatter:off
1758-
@JsonProperty("ref") McpSchema.CompleteReference ref,
1759-
@JsonProperty("argument") CompleteArgument argument,
1760-
@JsonProperty("_meta") Map<String, Object> meta) implements Request {
1758+
@JsonProperty("ref") McpSchema.CompleteReference ref,
1759+
@JsonProperty("argument") CompleteArgument argument,
1760+
@JsonProperty("_meta") Map<String, Object> meta,
1761+
@JsonProperty("context") CompleteContext context) implements Request {
17611762

1762-
public CompleteRequest(McpSchema.CompleteReference ref, CompleteArgument argument) {
1763-
this(ref, argument, null);
1764-
}
1763+
public CompleteRequest(McpSchema.CompleteReference ref, CompleteArgument argument, Map<String, Object> meta) {
1764+
this(ref, argument, meta, null);
1765+
}
17651766

1766-
public record CompleteArgument(
1767-
@JsonProperty("name") String name,
1768-
@JsonProperty("value") String value) {
1769-
}// @formatter:on
1767+
public CompleteRequest(McpSchema.CompleteReference ref, CompleteArgument argument, CompleteContext context) {
1768+
this(ref, argument, null, context);
1769+
}
1770+
1771+
public CompleteRequest(McpSchema.CompleteReference ref, CompleteArgument argument) {
1772+
this(ref, argument, null, null);
1773+
}
1774+
1775+
public record CompleteArgument(
1776+
@JsonProperty("name") String name,
1777+
@JsonProperty("value") String value) {
1778+
}
1779+
1780+
public record CompleteContext(
1781+
@JsonProperty("arguments") Map<String, String> arguments) {
1782+
}// @formatter:on
17701783
}
17711784

17721785
@JsonInclude(JsonInclude.Include.NON_ABSENT)

‎mcp/src/test/java/io/modelcontextprotocol/server/McpCompletionTests.java

Lines changed: 303 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,303 @@
1+
package io.modelcontextprotocol.server;
2+
3+
import java.util.List;
4+
import java.util.Map;
5+
import java.util.concurrent.atomic.AtomicReference;
6+
import java.util.function.BiFunction;
7+
8+
import org.apache.catalina.LifecycleException;
9+
import org.apache.catalina.LifecycleState;
10+
import org.apache.catalina.startup.Tomcat;
11+
import static org.assertj.core.api.Assertions.assertThat;
12+
import org.junit.jupiter.api.AfterEach;
13+
import org.junit.jupiter.api.BeforeEach;
14+
import org.junit.jupiter.api.Test;
15+
import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
16+
17+
import com.fasterxml.jackson.databind.ObjectMapper;
18+
19+
import io.modelcontextprotocol.client.McpClient;
20+
import io.modelcontextprotocol.client.transport.HttpClientSseClientTransport;
21+
import io.modelcontextprotocol.server.transport.HttpServletSseServerTransportProvider;
22+
import io.modelcontextprotocol.server.transport.TomcatTestUtil;
23+
import io.modelcontextprotocol.spec.McpSchema;
24+
import io.modelcontextprotocol.spec.McpSchema.CompleteRequest;
25+
import io.modelcontextprotocol.spec.McpSchema.CompleteResult;
26+
import io.modelcontextprotocol.spec.McpSchema.InitializeResult;
27+
import io.modelcontextprotocol.spec.McpSchema.Prompt;
28+
import io.modelcontextprotocol.spec.McpSchema.PromptArgument;
29+
import io.modelcontextprotocol.spec.McpSchema.ReadResourceResult;
30+
import io.modelcontextprotocol.spec.McpSchema.ResourceReference;
31+
import io.modelcontextprotocol.spec.McpSchema.PromptReference;
32+
import io.modelcontextprotocol.spec.McpSchema.ServerCapabilities;
33+
import io.modelcontextprotocol.spec.McpError;
34+
35+
/**
36+
* Tests for completion functionality with context support.
37+
*
38+
* @author Surbhi Bansal
39+
*/
40+
class McpCompletionTests {
41+
42+
private HttpServletSseServerTransportProvider mcpServerTransportProvider;
43+
44+
private static final String CUSTOM_MESSAGE_ENDPOINT = "/otherPath/mcp/message";
45+
46+
McpClient.SyncSpec clientBuilder;
47+
48+
private Tomcat tomcat;
49+
50+
@BeforeEach
51+
public void before() {
52+
// Create and con figure the transport provider
53+
mcpServerTransportProvider = HttpServletSseServerTransportProvider.builder()
54+
.objectMapper(new ObjectMapper())
55+
.messageEndpoint(CUSTOM_MESSAGE_ENDPOINT)
56+
.build();
57+
58+
tomcat = TomcatTestUtil.createTomcatServer("", 3400, mcpServerTransportProvider);
59+
try {
60+
tomcat.start();
61+
assertThat(tomcat.getServer().getState()).isEqualTo(LifecycleState.STARTED);
62+
}
63+
catch (Exception e) {
64+
throw new RuntimeException("Failed to start Tomcat", e);
65+
}
66+
67+
this.clientBuilder = McpClient.sync(HttpClientSseClientTransport.builder("http://localhost:" + 3400).build());
68+
}
69+
70+
@AfterEach
71+
public void after() {
72+
if (mcpServerTransportProvider != null) {
73+
mcpServerTransportProvider.closeGracefully().block();
74+
}
75+
if (tomcat != null) {
76+
try {
77+
tomcat.stop();
78+
tomcat.destroy();
79+
}
80+
catch (LifecycleException e) {
81+
throw new RuntimeException("Failed to stop Tomcat", e);
82+
}
83+
}
84+
}
85+
86+
@Test
87+
void testCompletionHandlerReceivesContext() {
88+
AtomicReference<CompleteRequest> receivedRequest = new AtomicReference<>();
89+
BiFunction<McpSyncServerExchange, CompleteRequest, CompleteResult> completionHandler = (exchange, request) -> {
90+
receivedRequest.set(request);
91+
return new CompleteResult(new CompleteResult.CompleteCompletion(List.of("test-completion"), 1, false));
92+
};
93+
94+
ResourceReference resourceRef = new ResourceReference("ref/resource", "test://resource/{param}");
95+
96+
McpSchema.Resource resource = new McpSchema.Resource("test://resource/{param}", "Test Resource",
97+
"A resource for testing", "text/plain", 123L, null);
98+
99+
var mcpServer = McpServer.sync(mcpServerTransportProvider)
100+
.capabilities(ServerCapabilities.builder().completions().build())
101+
.resources(new McpServerFeatures.SyncResourceSpecification(resource,
102+
(exchange, req) -> new ReadResourceResult(List.of())))
103+
.completions(new McpServerFeatures.SyncCompletionSpecification(resourceRef, completionHandler))
104+
.build();
105+
106+
try (var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample " + "client", "0.0.0"))
107+
.build();) {
108+
InitializeResult initResult = mcpClient.initialize();
109+
assertThat(initResult).isNotNull();
110+
111+
// Test with context
112+
CompleteRequest request = new CompleteRequest(resourceRef,
113+
new CompleteRequest.CompleteArgument("param", "test"), null,
114+
new CompleteRequest.CompleteContext(Map.of("previous", "value")));
115+
116+
CompleteResult result = mcpClient.completeCompletion(request);
117+
118+
// Verify handler received the context
119+
assertThat(receivedRequest.get().context()).isNotNull();
120+
assertThat(receivedRequest.get().context().arguments()).containsEntry("previous", "value");
121+
assertThat(result.completion().values()).containsExactly("test-completion");
122+
}
123+
124+
mcpServer.close();
125+
}
126+
127+
@Test
128+
void testCompletionBackwardCompatibility() {
129+
AtomicReference<Boolean> contextWasNull = new AtomicReference<>(false);
130+
BiFunction<McpSyncServerExchange, CompleteRequest, CompleteResult> completionHandler = (exchange, request) -> {
131+
contextWasNull.set(request.context() == null);
132+
return new CompleteResult(
133+
new CompleteResult.CompleteCompletion(List.of("no-context-completion"), 1, false));
134+
};
135+
136+
McpSchema.Prompt prompt = new Prompt("test-prompt", "this is a test prompt",
137+
List.of(new PromptArgument("arg", "string", false)));
138+
139+
var mcpServer = McpServer.sync(mcpServerTransportProvider)
140+
.capabilities(ServerCapabilities.builder().completions().build())
141+
.prompts(new McpServerFeatures.SyncPromptSpecification(prompt,
142+
(mcpSyncServerExchange, getPromptRequest) -> null))
143+
.completions(new McpServerFeatures.SyncCompletionSpecification(
144+
new PromptReference("ref/prompt", "test-prompt"), completionHandler))
145+
.build();
146+
147+
try (var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample " + "client", "0.0.0"))
148+
.build();) {
149+
InitializeResult initResult = mcpClient.initialize();
150+
assertThat(initResult).isNotNull();
151+
152+
// Test without context
153+
CompleteRequest request = new CompleteRequest(new PromptReference("ref/prompt", "test-prompt"),
154+
new CompleteRequest.CompleteArgument("arg", "val"));
155+
156+
CompleteResult result = mcpClient.completeCompletion(request);
157+
158+
// Verify context was null
159+
assertThat(contextWasNull.get()).isTrue();
160+
assertThat(result.completion().values()).containsExactly("no-context-completion");
161+
}
162+
163+
mcpServer.close();
164+
}
165+
166+
@Test
167+
void testDependentCompletionScenario() {
168+
BiFunction<McpSyncServerExchange, CompleteRequest, CompleteResult> completionHandler = (exchange, request) -> {
169+
// Simulate database/table completion scenario
170+
if (request.ref() instanceof ResourceReference resourceRef) {
171+
if ("db://{database}/{table}".equals(resourceRef.uri())) {
172+
if ("database".equals(request.argument().name())) {
173+
// Complete database names
174+
return new CompleteResult(new CompleteResult.CompleteCompletion(
175+
List.of("users_db", "products_db", "analytics_db"), 3, false));
176+
}
177+
else if ("table".equals(request.argument().name())) {
178+
// Complete table names based on selected database
179+
if (request.context() != null && request.context().arguments() != null) {
180+
String db = request.context().arguments().get("database");
181+
if ("users_db".equals(db)) {
182+
return new CompleteResult(new CompleteResult.CompleteCompletion(
183+
List.of("users", "sessions", "permissions"), 3, false));
184+
}
185+
else if ("products_db".equals(db)) {
186+
return new CompleteResult(new CompleteResult.CompleteCompletion(
187+
List.of("products", "categories", "inventory"), 3, false));
188+
}
189+
}
190+
}
191+
}
192+
}
193+
return new CompleteResult(new CompleteResult.CompleteCompletion(List.of(), 0, false));
194+
};
195+
196+
McpSchema.Resource resource = new McpSchema.Resource("db://{database}/{table}", "Database Table",
197+
"Resource representing a table in a database", "application/json", 456L, null);
198+
199+
var mcpServer = McpServer.sync(mcpServerTransportProvider)
200+
.capabilities(ServerCapabilities.builder().completions().build())
201+
.resources(new McpServerFeatures.SyncResourceSpecification(resource,
202+
(exchange, req) -> new ReadResourceResult(List.of())))
203+
.completions(new McpServerFeatures.SyncCompletionSpecification(
204+
new ResourceReference("ref/resource", "db://{database}/{table}"), completionHandler))
205+
.build();
206+
207+
try (var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample " + "client", "0.0.0"))
208+
.build();) {
209+
InitializeResult initResult = mcpClient.initialize();
210+
assertThat(initResult).isNotNull();
211+
212+
// First, complete database
213+
CompleteRequest dbRequest = new CompleteRequest(
214+
new ResourceReference("ref/resource", "db://{database}/{table}"),
215+
new CompleteRequest.CompleteArgument("database", ""));
216+
217+
CompleteResult dbResult = mcpClient.completeCompletion(dbRequest);
218+
assertThat(dbResult.completion().values()).contains("users_db", "products_db");
219+
220+
// Then complete table with database context
221+
CompleteRequest tableRequest = new CompleteRequest(
222+
new ResourceReference("ref/resource", "db://{database}/{table}"),
223+
new CompleteRequest.CompleteArgument("table", ""),
224+
new CompleteRequest.CompleteContext(Map.of("database", "users_db")));
225+
226+
CompleteResult tableResult = mcpClient.completeCompletion(tableRequest);
227+
assertThat(tableResult.completion().values()).containsExactly("users", "sessions", "permissions");
228+
229+
// Different database gives different tables
230+
CompleteRequest tableRequest2 = new CompleteRequest(
231+
new ResourceReference("ref/resource", "db://{database}/{table}"),
232+
new CompleteRequest.CompleteArgument("table", ""),
233+
new CompleteRequest.CompleteContext(Map.of("database", "products_db")));
234+
235+
CompleteResult tableResult2 = mcpClient.completeCompletion(tableRequest2);
236+
assertThat(tableResult2.completion().values()).containsExactly("products", "categories", "inventory");
237+
}
238+
239+
mcpServer.close();
240+
}
241+
242+
@Test
243+
void testCompletionErrorOnMissingContext() {
244+
BiFunction<McpSyncServerExchange, CompleteRequest, CompleteResult> completionHandler = (exchange, request) -> {
245+
if (request.ref() instanceof ResourceReference resourceRef) {
246+
if ("db://{database}/{table}".equals(resourceRef.uri())) {
247+
if ("table".equals(request.argument().name())) {
248+
// Check if database context is provided
249+
if (request.context() == null || request.context().arguments() == null
250+
|| !request.context().arguments().containsKey("database")) {
251+
throw new McpError("Please select a database first to see available tables");
252+
}
253+
// Normal completion if context is provided
254+
String db = request.context().arguments().get("database");
255+
if ("test_db".equals(db)) {
256+
return new CompleteResult(new CompleteResult.CompleteCompletion(
257+
List.of("users", "orders", "products"), 3, false));
258+
}
259+
}
260+
}
261+
}
262+
return new CompleteResult(new CompleteResult.CompleteCompletion(List.of(), 0, false));
263+
};
264+
265+
McpSchema.Resource resource = new McpSchema.Resource("db://{database}/{table}", "Database Table",
266+
"Resource representing a table in a database", "application/json", 456L, null);
267+
268+
var mcpServer = McpServer.sync(mcpServerTransportProvider)
269+
.capabilities(ServerCapabilities.builder().completions().build())
270+
.resources(new McpServerFeatures.SyncResourceSpecification(resource,
271+
(exchange, req) -> new ReadResourceResult(List.of())))
272+
.completions(new McpServerFeatures.SyncCompletionSpecification(
273+
new ResourceReference("ref/resource", "db://{database}/{table}"), completionHandler))
274+
.build();
275+
276+
try (var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample" + "client", "0.0.0"))
277+
.build();) {
278+
InitializeResult initResult = mcpClient.initialize();
279+
assertThat(initResult).isNotNull();
280+
281+
// Try to complete table without database context - should raise error
282+
CompleteRequest requestWithoutContext = new CompleteRequest(
283+
new ResourceReference("ref/resource", "db://{database}/{table}"),
284+
new CompleteRequest.CompleteArgument("table", ""));
285+
286+
assertThatExceptionOfType(McpError.class)
287+
.isThrownBy(() -> mcpClient.completeCompletion(requestWithoutContext))
288+
.withMessageContaining("Please select a database first");
289+
290+
// Now complete with proper context - should work normally
291+
CompleteRequest requestWithContext = new CompleteRequest(
292+
new ResourceReference("ref/resource", "db://{database}/{table}"),
293+
new CompleteRequest.CompleteArgument("table", ""),
294+
new CompleteRequest.CompleteContext(Map.of("database", "test_db")));
295+
296+
CompleteResult resultWithContext = mcpClient.completeCompletion(requestWithContext);
297+
assertThat(resultWithContext.completion().values()).containsExactly("users", "orders", "products");
298+
}
299+
300+
mcpServer.close();
301+
}
302+
303+
}

‎mcp/src/test/java/io/modelcontextprotocol/spec/McpSchemaTests.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1266,7 +1266,7 @@ void testCompleteRequestWithMeta() throws Exception {
12661266
Map<String, Object> meta = new HashMap<>();
12671267
meta.put("progressToken", "complete-progress-789");
12681268

1269-
McpSchema.CompleteRequest request = new McpSchema.CompleteRequest(resourceRef, argument, meta);
1269+
McpSchema.CompleteRequest request = new McpSchema.CompleteRequest(resourceRef, argument, meta, null);
12701270

12711271
String value = mapper.writeValueAsString(request);
12721272
assertThatJson(value).when(Option.IGNORING_ARRAY_ORDER)

0 commit comments

Comments
 (0)