Skip to content

Commit 4366fe4

Browse files
tzolovtzolov
committed
feat(mcp): Add granular control over MCP server capabilities
- Add capability configuration to enable/disable tools, resources, prompts, and completions individually - Refactor server configuration to conditionally register capabilities based on configuration - Add support for server instructions configuration - Update documentation with new capability options and completion management - Add tests for new capabilities and configurations Resolves #3207 Signed-off-by: Christian Tzolov <[email protected]>
1 parent c0b9240 commit 4366fe4

File tree

4 files changed

+365
-49
lines changed

4 files changed

+365
-49
lines changed

‎auto-configurations/mcp/spring-ai-autoconfigure-mcp-server/src/main/java/org/springframework/ai/mcp/server/autoconfigure/McpServerAutoConfiguration.java‎

Lines changed: 84 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -149,11 +149,11 @@ private List<McpServerFeatures.SyncToolSpecification> toSyncToolSpecifications(L
149149

150150
// De-duplicate tools by their name, keeping the first occurrence of each tool
151151
// name
152-
return tools.stream()
153-
.collect(Collectors.toMap(tool -> tool.getToolDefinition().name(), // Key:
154-
// tool
155-
// name
156-
tool -> tool, // Value: the tool itself
152+
return tools.stream() // Key: tool name
153+
.collect(Collectors.toMap(tool -> tool.getToolDefinition().name(), tool -> tool, // Value:
154+
// the
155+
// tool
156+
// itself
157157
(existing, replacement) -> existing)) // On duplicate key, keep the
158158
// existing tool
159159
.values()
@@ -185,47 +185,66 @@ public McpSyncServer mcpSyncServer(McpServerTransportProvider transportProvider,
185185
// Create the server with both tool and resource capabilities
186186
SyncSpecification serverBuilder = McpServer.sync(transportProvider).serverInfo(serverInfo);
187187

188-
List<SyncToolSpecification> toolSpecifications = new ArrayList<>(tools.stream().flatMap(List::stream).toList());
188+
// Tools
189+
if (serverProperties.getCapabilities().isTool()) {
190+
logger.info("Enable tools capabilities, notification: " + serverProperties.isToolChangeNotification());
191+
capabilitiesBuilder.tools(serverProperties.isToolChangeNotification());
189192

190-
List<ToolCallback> providerToolCallbacks = toolCallbackProvider.stream()
191-
.map(pr -> List.of(pr.getToolCallbacks()))
192-
.flatMap(List::stream)
193-
.filter(fc -> fc instanceof ToolCallback)
194-
.map(fc -> (ToolCallback) fc)
195-
.toList();
193+
List<SyncToolSpecification> toolSpecifications = new ArrayList<>(
194+
tools.stream().flatMap(List::stream).toList());
196195

197-
toolSpecifications.addAll(this.toSyncToolSpecifications(providerToolCallbacks, serverProperties));
196+
List<ToolCallback> providerToolCallbacks = toolCallbackProvider.stream()
197+
.map(pr -> List.of(pr.getToolCallbacks()))
198+
.flatMap(List::stream)
199+
.filter(fc -> fc instanceof ToolCallback)
200+
.map(fc -> (ToolCallback) fc)
201+
.toList();
198202

199-
if (!CollectionUtils.isEmpty(toolSpecifications)) {
200-
serverBuilder.tools(toolSpecifications);
201-
capabilitiesBuilder.tools(serverProperties.isToolChangeNotification());
202-
logger.info("Registered tools: " + toolSpecifications.size() + ", notification: "
203-
+ serverProperties.isToolChangeNotification());
203+
toolSpecifications.addAll(this.toSyncToolSpecifications(providerToolCallbacks, serverProperties));
204+
205+
if (!CollectionUtils.isEmpty(toolSpecifications)) {
206+
serverBuilder.tools(toolSpecifications);
207+
logger.info("Registered tools: " + toolSpecifications.size());
208+
}
204209
}
205210

206-
List<SyncResourceSpecification> resourceSpecifications = resources.stream().flatMap(List::stream).toList();
207-
if (!CollectionUtils.isEmpty(resourceSpecifications)) {
208-
serverBuilder.resources(resourceSpecifications);
211+
// Resources
212+
if (serverProperties.getCapabilities().isResource()) {
213+
logger.info(
214+
"Enable resources capabilities, notification: " + serverProperties.isResourceChangeNotification());
209215
capabilitiesBuilder.resources(false, serverProperties.isResourceChangeNotification());
210-
logger.info("Registered resources: " + resourceSpecifications.size() + ", notification: "
211-
+ serverProperties.isResourceChangeNotification());
216+
217+
List<SyncResourceSpecification> resourceSpecifications = resources.stream().flatMap(List::stream).toList();
218+
if (!CollectionUtils.isEmpty(resourceSpecifications)) {
219+
serverBuilder.resources(resourceSpecifications);
220+
logger.info("Registered resources: " + resourceSpecifications.size());
221+
}
212222
}
213223

214-
List<SyncPromptSpecification> promptSpecifications = prompts.stream().flatMap(List::stream).toList();
215-
if (!CollectionUtils.isEmpty(promptSpecifications)) {
216-
serverBuilder.prompts(promptSpecifications);
224+
// Prompts
225+
if (serverProperties.getCapabilities().isPrompt()) {
226+
logger.info("Enable prompts capabilities, notification: " + serverProperties.isPromptChangeNotification());
217227
capabilitiesBuilder.prompts(serverProperties.isPromptChangeNotification());
218-
logger.info("Registered prompts: " + promptSpecifications.size() + ", notification: "
219-
+ serverProperties.isPromptChangeNotification());
228+
229+
List<SyncPromptSpecification> promptSpecifications = prompts.stream().flatMap(List::stream).toList();
230+
if (!CollectionUtils.isEmpty(promptSpecifications)) {
231+
serverBuilder.prompts(promptSpecifications);
232+
logger.info("Registered prompts: " + promptSpecifications.size());
233+
}
220234
}
221235

222-
List<SyncCompletionSpecification> completionSpecifications = completions.stream()
223-
.flatMap(List::stream)
224-
.toList();
225-
if (!CollectionUtils.isEmpty(completionSpecifications)) {
226-
serverBuilder.completions(completionSpecifications);
236+
// Completions
237+
if (serverProperties.getCapabilities().isCompletion()) {
238+
logger.info("Enable completions capabilities");
227239
capabilitiesBuilder.completions();
228-
logger.info("Registered completions: " + completionSpecifications.size());
240+
241+
List<SyncCompletionSpecification> completionSpecifications = completions.stream()
242+
.flatMap(List::stream)
243+
.toList();
244+
if (!CollectionUtils.isEmpty(completionSpecifications)) {
245+
serverBuilder.completions(completionSpecifications);
246+
logger.info("Registered completions: " + completionSpecifications.size());
247+
}
229248
}
230249

231250
rootsChangeConsumers.ifAvailable(consumer -> {
@@ -257,11 +276,11 @@ private List<McpServerFeatures.AsyncToolSpecification> toAsyncToolSpecification(
257276
McpServerProperties serverProperties) {
258277
// De-duplicate tools by their name, keeping the first occurrence of each tool
259278
// name
260-
return tools.stream()
261-
.collect(Collectors.toMap(tool -> tool.getToolDefinition().name(), // Key:
262-
// tool
263-
// name
264-
tool -> tool, // Value: the tool itself
279+
return tools.stream() // Key: tool name
280+
.collect(Collectors.toMap(tool -> tool.getToolDefinition().name(), tool -> tool, // Value:
281+
// the
282+
// tool
283+
// itself
265284
(existing, replacement) -> existing)) // On duplicate key, keep the
266285
// existing tool
267286
.values()
@@ -303,35 +322,51 @@ public McpAsyncServer mcpAsyncServer(McpServerTransportProvider transportProvide
303322

304323
toolSpecifications.addAll(this.toAsyncToolSpecification(providerToolCallbacks, serverProperties));
305324

325+
// Tools
326+
if (serverProperties.getCapabilities().isTool()) {
327+
logger.info("Enable tools capabilities, notification: " + serverProperties.isToolChangeNotification());
328+
capabilitiesBuilder.tools(serverProperties.isToolChangeNotification());
329+
}
330+
306331
if (!CollectionUtils.isEmpty(toolSpecifications)) {
307332
serverBuilder.tools(toolSpecifications);
308-
capabilitiesBuilder.tools(serverProperties.isToolChangeNotification());
309-
logger.info("Registered tools: " + toolSpecifications.size() + ", notification: "
310-
+ serverProperties.isToolChangeNotification());
333+
logger.info("Registered tools: " + toolSpecifications.size());
334+
}
335+
336+
// Resources
337+
if (serverProperties.getCapabilities().isResource()) {
338+
logger.info(
339+
"Enable resources capabilities, notification: " + serverProperties.isResourceChangeNotification());
340+
capabilitiesBuilder.resources(false, serverProperties.isResourceChangeNotification());
311341
}
312342

313343
List<AsyncResourceSpecification> resourceSpecifications = resources.stream().flatMap(List::stream).toList();
314344
if (!CollectionUtils.isEmpty(resourceSpecifications)) {
315345
serverBuilder.resources(resourceSpecifications);
316-
capabilitiesBuilder.resources(false, serverProperties.isResourceChangeNotification());
317-
logger.info("Registered resources: " + resourceSpecifications.size() + ", notification: "
318-
+ serverProperties.isResourceChangeNotification());
346+
logger.info("Registered resources: " + resourceSpecifications.size());
319347
}
320348

349+
// Prompts
350+
if (serverProperties.getCapabilities().isPrompt()) {
351+
logger.info("Enable prompts capabilities, notification: " + serverProperties.isPromptChangeNotification());
352+
capabilitiesBuilder.prompts(serverProperties.isPromptChangeNotification());
353+
}
321354
List<AsyncPromptSpecification> promptSpecifications = prompts.stream().flatMap(List::stream).toList();
322355
if (!CollectionUtils.isEmpty(promptSpecifications)) {
323356
serverBuilder.prompts(promptSpecifications);
324-
capabilitiesBuilder.prompts(serverProperties.isPromptChangeNotification());
325-
logger.info("Registered prompts: " + promptSpecifications.size() + ", notification: "
326-
+ serverProperties.isPromptChangeNotification());
357+
logger.info("Registered prompts: " + promptSpecifications.size());
327358
}
328359

360+
// Completions
361+
if (serverProperties.getCapabilities().isCompletion()) {
362+
logger.info("Enable completions capabilities");
363+
capabilitiesBuilder.completions();
364+
}
329365
List<AsyncCompletionSpecification> completionSpecifications = completions.stream()
330366
.flatMap(List::stream)
331367
.toList();
332368
if (!CollectionUtils.isEmpty(completionSpecifications)) {
333369
serverBuilder.completions(completionSpecifications);
334-
capabilitiesBuilder.completions();
335370
logger.info("Registered completions: " + completionSpecifications.size());
336371
}
337372

‎auto-configurations/mcp/spring-ai-autoconfigure-mcp-server/src/main/java/org/springframework/ai/mcp/server/autoconfigure/McpServerProperties.java‎

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -132,6 +132,56 @@ public class McpServerProperties {
132132
*/
133133
private ServerType type = ServerType.SYNC;
134134

135+
private Capabilities capabilities = new Capabilities();
136+
137+
public static class Capabilities {
138+
139+
private boolean resource = true;
140+
141+
private boolean tool = true;
142+
143+
private boolean prompt = true;
144+
145+
private boolean completion = true;
146+
147+
public boolean isResource() {
148+
return resource;
149+
}
150+
151+
public void setResource(boolean resource) {
152+
this.resource = resource;
153+
}
154+
155+
public boolean isTool() {
156+
return tool;
157+
}
158+
159+
public void setTool(boolean tool) {
160+
this.tool = tool;
161+
}
162+
163+
public boolean isPrompt() {
164+
return prompt;
165+
}
166+
167+
public void setPrompt(boolean prompt) {
168+
this.prompt = prompt;
169+
}
170+
171+
public boolean isCompletion() {
172+
return completion;
173+
}
174+
175+
public void setCompletion(boolean completion) {
176+
this.completion = completion;
177+
}
178+
179+
}
180+
181+
public Capabilities getCapabilities() {
182+
return capabilities;
183+
}
184+
135185
/**
136186
* Server types supported by the MCP server.
137187
*/

0 commit comments

Comments
 (0)