Skip to content

Commit a5394ff

Browse files
tzolovtzolov
committed
efactor(mcp): reorganize McpAsyncServer and enhance error handling
- Reorganize McpAsyncServer code into logical sections (Lifecycle, Tool, Resource, Prompt Management) - (bug) Fix resource registration to use URI instead of name as map key - Improve error logging in StdioServerTransport with helper methods - Clean up documentation and add transport mode descriptions - Update sample server configuration with clearer method names and organization
1 parent a2a8be2 commit a5394ff

File tree

5 files changed

+195
-180
lines changed

5 files changed

+195
-180
lines changed

‎mcp/src/main/java/org/springframework/ai/mcp/server/McpAsyncServer.java

Lines changed: 147 additions & 137 deletions
Original file line numberDiff line numberDiff line change
@@ -129,6 +129,52 @@ public McpAsyncServer(McpTransport mcpTransport, McpSchema.Implementation server
129129
notificationHandlers);
130130
}
131131

132+
// ---------------------------------------
133+
// Lifecycle Management
134+
// ---------------------------------------
135+
private DefaultMcpSession.RequestHandler initializeRequestHandler() {
136+
return params -> {
137+
McpSchema.InitializeRequest initializeRequest = transport.unmarshalFrom(params,
138+
new TypeReference<McpSchema.InitializeRequest>() {
139+
});
140+
141+
logger.info("Client initialize request - Protocol: {}, Capabilities: {}, Info: {}",
142+
initializeRequest.protocolVersion(), initializeRequest.capabilities(),
143+
initializeRequest.clientInfo());
144+
145+
if (!McpSchema.LATEST_PROTOCOL_VERSION.equals(initializeRequest.protocolVersion())) {
146+
return Mono
147+
.<Object>error(new McpError(
148+
"Unsupported protocol version from client: " + initializeRequest.protocolVersion()))
149+
.publishOn(Schedulers.boundedElastic());
150+
}
151+
152+
return Mono
153+
.<Object>just(new McpSchema.InitializeResult(McpSchema.LATEST_PROTOCOL_VERSION, this.serverCapabilities,
154+
this.serverInfo, null))
155+
.publishOn(Schedulers.boundedElastic());
156+
};
157+
}
158+
159+
/**
160+
* Gracefully closes the server, allowing any in-progress operations to complete.
161+
* @return A Mono that completes when the server has been closed
162+
*/
163+
public Mono<Void> closeGracefully() {
164+
return this.mcpSession.closeGracefully();
165+
}
166+
167+
/**
168+
* Close the server immediately.
169+
*/
170+
public void close() {
171+
this.mcpSession.close();
172+
}
173+
174+
// ---------------------------------------
175+
// Tool Management
176+
// ---------------------------------------
177+
132178
/**
133179
* Add a new tool registration at runtime.
134180
* @param toolRegistration The tool registration to add
@@ -185,6 +231,53 @@ public Mono<Void> removeTool(String toolName) {
185231
return Mono.error(new McpError("Tool with name '" + toolName + "' not found"));
186232
}
187233

234+
/**
235+
* Notifies clients that the list of available tools has changed.
236+
* @return A Mono that completes when all clients have been notified
237+
*/
238+
public Mono<Void> notifyToolsListChanged() {
239+
return this.mcpSession.sendNotification("notifications/tools/list_changed", null);
240+
}
241+
242+
private DefaultMcpSession.RequestHandler toolsListRequestHandler() {
243+
return params -> {
244+
McpSchema.PaginatedRequest request = transport.unmarshalFrom(params,
245+
new TypeReference<McpSchema.PaginatedRequest>() {
246+
});
247+
248+
List<Tool> tools = this.tools.stream().map(toolRegistration -> {
249+
return toolRegistration.tool();
250+
}).toList();
251+
252+
logger.info("Client tools list request - Cursor: {}", request.cursor());
253+
return Mono.just(new McpSchema.ListToolsResult(tools, null));
254+
};
255+
}
256+
257+
private DefaultMcpSession.RequestHandler toolsCallRequestHandler() {
258+
return params -> {
259+
McpSchema.CallToolRequest callToolRequest = transport.unmarshalFrom(params,
260+
new TypeReference<McpSchema.CallToolRequest>() {
261+
});
262+
263+
Optional<ToolRegistration> toolRegistration = this.tools.stream()
264+
.filter(tr -> callToolRequest.name().equals(tr.tool().name()))
265+
.findAny();
266+
267+
if (toolRegistration.isEmpty()) {
268+
return Mono.<Object>error(new McpError("Tool not found: " + callToolRequest.name()));
269+
}
270+
271+
CallToolResult callResponse = toolRegistration.get().call().apply(callToolRequest.arguments());
272+
273+
return Mono.just(callResponse);
274+
};
275+
}
276+
277+
// ---------------------------------------
278+
// Resource Management
279+
// ---------------------------------------
280+
188281
/**
189282
* Add a new resource handler at runtime.
190283
* @param resourceHandler The resource handler to add
@@ -236,6 +329,43 @@ public Mono<Void> removeResource(String resourceUri) {
236329
return Mono.error(new McpError("Resource with URI '" + resourceUri + "' not found"));
237330
}
238331

332+
/**
333+
* Notifies clients that the list of available resources has changed.
334+
* @return A Mono that completes when all clients have been notified
335+
*/
336+
public Mono<Void> notifyResourcesListChanged() {
337+
return this.mcpSession.sendNotification("notifications/resources/list_changed", null);
338+
}
339+
340+
private DefaultMcpSession.RequestHandler resourcesListRequestHandler() {
341+
return params -> {
342+
var resourceList = this.resources.values().stream().map(ResourceRegistration::resource).toList();
343+
return Mono.just(new McpSchema.ListResourcesResult(resourceList, null));
344+
};
345+
}
346+
347+
private DefaultMcpSession.RequestHandler resourceTemplateListRequestHandler() {
348+
return params -> Mono.just(new McpSchema.ListResourceTemplatesResult(this.resourceTemplates, null));
349+
350+
}
351+
352+
private DefaultMcpSession.RequestHandler resourcesReadRequestHandler() {
353+
return params -> {
354+
McpSchema.ReadResourceRequest resourceRequest = transport.unmarshalFrom(params,
355+
new TypeReference<McpSchema.ReadResourceRequest>() {
356+
});
357+
var resourceUri = resourceRequest.uri();
358+
if (this.resources.containsKey(resourceUri)) {
359+
return Mono.just(this.resources.get(resourceUri).readHandler().apply(resourceRequest));
360+
}
361+
return Mono.error(new McpError("Resource not found: " + resourceUri));
362+
};
363+
}
364+
365+
// ---------------------------------------
366+
// Prompt Management
367+
// ---------------------------------------
368+
239369
/**
240370
* Add a new prompt handler at runtime.
241371
* @param promptRegistration The prompt handler to add
@@ -255,7 +385,11 @@ public Mono<Void> addPrompt(PromptRegistration promptRegistration) {
255385
}
256386

257387
this.prompts.put(promptRegistration.propmpt().name(), promptRegistration);
388+
258389
logger.info("Added prompt handler: {}", promptRegistration.propmpt().name());
390+
391+
// Servers that declared the listChanged capability SHOULD send a notification,
392+
// when the list of available prompts changes
259393
if (this.serverCapabilities.prompts().listChanged()) {
260394
return notifyPromptsListChanged();
261395
}
@@ -276,8 +410,11 @@ public Mono<Void> removePrompt(String promptName) {
276410
}
277411

278412
PromptRegistration removed = this.prompts.remove(promptName);
413+
279414
if (removed != null) {
280415
logger.info("Removed prompt handler: {}", promptName);
416+
// Servers that declared the listChanged capability SHOULD send a
417+
// notification, when the list of available prompts changes
281418
if (this.serverCapabilities.prompts().listChanged()) {
282419
return this.notifyPromptsListChanged();
283420
}
@@ -286,100 +423,12 @@ public Mono<Void> removePrompt(String promptName) {
286423
return Mono.error(new McpError("Prompt with name '" + promptName + "' not found"));
287424
}
288425

289-
// ---------------------------------------
290-
// Request Handlers
291-
// ---------------------------------------
292-
private DefaultMcpSession.RequestHandler initializeRequestHandler() {
293-
return params -> {
294-
McpSchema.InitializeRequest request = transport.unmarshalFrom(params,
295-
new TypeReference<McpSchema.InitializeRequest>() {
296-
});
297-
298-
logger.info("Client initialize request - Protocol: {}, Capabilities: {}, Info: {}",
299-
request.protocolVersion(), request.capabilities(), request.clientInfo());
300-
301-
if (!McpSchema.LATEST_PROTOCOL_VERSION.equals(request.protocolVersion())) {
302-
return Mono
303-
.<Object>error(
304-
new McpError("Unsupported protocol version from client: " + request.protocolVersion()))
305-
.publishOn(Schedulers.boundedElastic());
306-
}
307-
308-
return Mono
309-
.<Object>just(new McpSchema.InitializeResult(McpSchema.LATEST_PROTOCOL_VERSION, this.serverCapabilities,
310-
this.serverInfo, null))
311-
.publishOn(Schedulers.boundedElastic());
312-
};
313-
}
314-
315-
private DefaultMcpSession.RequestHandler toolsListRequestHandler() {
316-
return params -> {
317-
McpSchema.PaginatedRequest request = transport.unmarshalFrom(params,
318-
new TypeReference<McpSchema.PaginatedRequest>() {
319-
});
320-
321-
List<Tool> tools = this.tools.stream().map(toolRegistration -> {
322-
return toolRegistration.tool();
323-
}).toList();
324-
325-
logger.info("Client tools list request - Cursor: {}", request.cursor());
326-
return Mono.just(new McpSchema.ListToolsResult(tools, null));
327-
};
328-
}
329-
330-
private DefaultMcpSession.RequestHandler toolsCallRequestHandler() {
331-
return params -> {
332-
McpSchema.CallToolRequest callToolRequest = transport.unmarshalFrom(params,
333-
new TypeReference<McpSchema.CallToolRequest>() {
334-
});
335-
336-
Optional<ToolRegistration> toolRegistration = this.tools.stream()
337-
.filter(tr -> callToolRequest.name().equals(tr.tool().name()))
338-
.findAny();
339-
340-
if (toolRegistration.isEmpty()) {
341-
return Mono.<Object>error(new McpError("Tool not found: " + callToolRequest.name()));
342-
}
343-
344-
CallToolResult callResponse = toolRegistration.get().call().apply(callToolRequest.arguments());
345-
346-
return Mono.just(callResponse);
347-
};
348-
}
349-
350-
private DefaultMcpSession.RequestHandler resourcesListRequestHandler() {
351-
return params -> {
352-
// McpSchema.PaginatedRequest request = transport.unmarshalFrom(params,
353-
// new TypeReference<McpSchema.PaginatedRequest>() {
354-
// });
355-
356-
var resourceList = this.resources.values().stream().map(ResourceRegistration::resource).toList();
357-
358-
return Mono.just(new McpSchema.ListResourcesResult(resourceList, null));
359-
};
360-
}
361-
362-
private DefaultMcpSession.RequestHandler resourceTemplateListRequestHandler() {
363-
return params -> {
364-
// McpSchema.PaginatedRequest request = transport.unmarshalFrom(params,
365-
// new TypeReference<McpSchema.PaginatedRequest>() {
366-
// });
367-
368-
return Mono.just(new McpSchema.ListResourceTemplatesResult(this.resourceTemplates, null));
369-
};
370-
}
371-
372-
private DefaultMcpSession.RequestHandler resourcesReadRequestHandler() {
373-
return params -> {
374-
McpSchema.ReadResourceRequest request = transport.unmarshalFrom(params,
375-
new TypeReference<McpSchema.ReadResourceRequest>() {
376-
});
377-
var resourceUri = request.uri();
378-
if (this.resources.containsKey(resourceUri)) {
379-
return Mono.just(this.resources.get(resourceUri).readHandler().apply(request));
380-
}
381-
return Mono.error(new McpError("Resource not found: " + resourceUri));
382-
};
426+
/**
427+
* Notifies clients that the list of available prompts has changed.
428+
* @return A Mono that completes when all clients have been notified
429+
*/
430+
public Mono<Void> notifyPromptsListChanged() {
431+
return this.mcpSession.sendNotification("notifications/prompts/list_changed", null);
383432
}
384433

385434
private DefaultMcpSession.RequestHandler promptsListRequestHandler() {
@@ -396,56 +445,17 @@ private DefaultMcpSession.RequestHandler promptsListRequestHandler() {
396445

397446
private DefaultMcpSession.RequestHandler promptsGetRequestHandler() {
398447
return params -> {
399-
McpSchema.GetPromptRequest request = transport.unmarshalFrom(params,
448+
McpSchema.GetPromptRequest promptRequest = transport.unmarshalFrom(params,
400449
new TypeReference<McpSchema.GetPromptRequest>() {
401450
});
402451

403452
// Implement prompt retrieval logic here
404-
if (this.prompts.containsKey(request.name())) {
405-
return Mono.just(this.prompts.get(request.name()).promptHandler().apply(request));
453+
if (this.prompts.containsKey(promptRequest.name())) {
454+
return Mono.just(this.prompts.get(promptRequest.name()).promptHandler().apply(promptRequest));
406455
}
407456

408-
return Mono.error(new McpError("Prompt not found: " + request.name()));
457+
return Mono.error(new McpError("Prompt not found: " + promptRequest.name()));
409458
};
410459
}
411460

412-
/**
413-
* Notifies clients that the list of available tools has changed.
414-
* @return A Mono that completes when all clients have been notified
415-
*/
416-
public Mono<Void> notifyToolsListChanged() {
417-
return this.mcpSession.sendNotification("notifications/tools/list_changed", null);
418-
}
419-
420-
/**
421-
* Notifies clients that the list of available resources has changed.
422-
* @return A Mono that completes when all clients have been notified
423-
*/
424-
public Mono<Void> notifyResourcesListChanged() {
425-
return this.mcpSession.sendNotification("notifications/resources/list_changed", null);
426-
}
427-
428-
/**
429-
* Notifies clients that the list of available prompts has changed.
430-
* @return A Mono that completes when all clients have been notified
431-
*/
432-
public Mono<Void> notifyPromptsListChanged() {
433-
return this.mcpSession.sendNotification("notifications/prompts/list_changed", null);
434-
}
435-
436-
/**
437-
* Gracefully closes the server, allowing any in-progress operations to complete.
438-
* @return A Mono that completes when the server has been closed
439-
*/
440-
public Mono<Void> closeGracefully() {
441-
return this.mcpSession.closeGracefully();
442-
}
443-
444-
/**
445-
* Close the server immediately.
446-
*/
447-
public void close() {
448-
this.mcpSession.close();
449-
}
450-
451461
}

‎mcp/src/main/java/org/springframework/ai/mcp/server/McpServer.java

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -392,7 +392,7 @@ public Builder resources(Map<String, ResourceRegistration> resourceRegsitrations
392392
public Builder resources(List<ResourceRegistration> resourceRegsitrations) {
393393
Assert.notNull(resourceRegsitrations, "Resource handlers list must not be null");
394394
for (ResourceRegistration resource : resourceRegsitrations) {
395-
this.resources.put(resource.resource().name(), resource);
395+
this.resources.put(resource.resource().uri(), resource);
396396
}
397397
return this;
398398
}
@@ -417,7 +417,7 @@ public Builder resources(List<ResourceRegistration> resourceRegsitrations) {
417417
public Builder resources(ResourceRegistration... resourceRegistrations) {
418418
Assert.notNull(resourceRegistrations, "Resource handlers list must not be null");
419419
for (ResourceRegistration resource : resourceRegistrations) {
420-
this.resources.put(resource.resource().name(), resource);
420+
this.resources.put(resource.resource().uri(), resource);
421421
}
422422
return this;
423423
}

0 commit comments

Comments
 (0)