package com.zy.ai.mcp.service;
|
|
import com.alibaba.fastjson.JSON;
|
import com.alibaba.fastjson.JSONObject;
|
import io.modelcontextprotocol.client.McpClient;
|
import io.modelcontextprotocol.client.McpSyncClient;
|
import io.modelcontextprotocol.client.transport.HttpClientSseClientTransport;
|
import io.modelcontextprotocol.spec.McpSchema;
|
import lombok.extern.slf4j.Slf4j;
|
import org.springframework.ai.mcp.SyncMcpToolCallbackProvider;
|
import org.springframework.ai.tool.ToolCallback;
|
import org.springframework.ai.tool.ToolCallbackProvider;
|
import org.springframework.ai.tool.definition.ToolDefinition;
|
import org.springframework.beans.factory.annotation.Value;
|
import org.springframework.stereotype.Service;
|
|
import jakarta.annotation.PreDestroy;
|
import java.time.Duration;
|
import java.util.ArrayList;
|
import java.util.Collections;
|
import java.util.Comparator;
|
import java.util.LinkedHashMap;
|
import java.util.List;
|
import java.util.Map;
|
|
@Slf4j
|
@Service
|
public class SpringAiMcpToolManager {
|
|
private static final McpSchema.Implementation CLIENT_INFO =
|
new McpSchema.Implementation("wcs-ai-assistant", "1.0.0");
|
|
private final Object clientMonitor = new Object();
|
|
@Value("${server.port:9090}")
|
private Integer serverPort;
|
|
@Value("${server.servlet.context-path:}")
|
private String contextPath;
|
|
@Value("${spring.ai.mcp.server.sse-endpoint:/ai/mcp/sse}")
|
private String sseEndpoint;
|
|
@Value("${spring.ai.mcp.server.request-timeout:20s}")
|
private Duration requestTimeout;
|
|
@Value("${app.ai.mcp.client.base-url:}")
|
private String configuredBaseUrl;
|
|
private volatile ClientSession clientSession;
|
|
public List<Map<String, Object>> listTools() {
|
List<Map<String, Object>> tools = new ArrayList<Map<String, Object>>();
|
for (ToolCallback callback : getToolCallbacks()) {
|
if (callback == null || callback.getToolDefinition() == null) {
|
continue;
|
}
|
ToolDefinition definition = callback.getToolDefinition();
|
Map<String, Object> item = new LinkedHashMap<String, Object>();
|
item.put("name", definition.name());
|
item.put("description", definition.description());
|
item.put("inputSchema", parseSchema(definition.inputSchema()));
|
tools.add(item);
|
}
|
tools.sort(new Comparator<Map<String, Object>>() {
|
@Override
|
public int compare(Map<String, Object> left, Map<String, Object> right) {
|
return String.valueOf(left.get("name")).compareTo(String.valueOf(right.get("name")));
|
}
|
});
|
return tools;
|
}
|
|
public List<Object> buildOpenAiTools() {
|
List<Object> tools = new ArrayList<Object>();
|
for (Map<String, Object> item : listTools()) {
|
Object name = item.get("name");
|
if (name == null) {
|
continue;
|
}
|
|
Map<String, Object> function = new LinkedHashMap<String, Object>();
|
function.put("name", String.valueOf(name));
|
Object description = item.get("description");
|
if (description != null) {
|
function.put("description", String.valueOf(description));
|
}
|
Object inputSchema = item.get("inputSchema");
|
function.put("parameters", inputSchema == null ? new LinkedHashMap<String, Object>() : inputSchema);
|
|
Map<String, Object> tool = new LinkedHashMap<String, Object>();
|
tool.put("type", "function");
|
tool.put("function", function);
|
tools.add(tool);
|
}
|
return tools;
|
}
|
|
public Object callTool(String toolName, JSONObject arguments) {
|
if (toolName == null || toolName.trim().isEmpty()) {
|
throw new IllegalArgumentException("missing tool name");
|
}
|
|
ToolCallback callback = findCallback(toolName);
|
if (callback == null) {
|
throw new IllegalArgumentException("tool not found: " + toolName);
|
}
|
|
String rawResult = callback.call(arguments == null ? "{}" : arguments.toJSONString());
|
return parseToolResult(rawResult);
|
}
|
|
private ToolCallback findCallback(String toolName) {
|
for (ToolCallback callback : getToolCallbacks()) {
|
if (callback == null || callback.getToolDefinition() == null) {
|
continue;
|
}
|
if (toolName.equals(callback.getToolDefinition().name())) {
|
return callback;
|
}
|
}
|
return null;
|
}
|
|
private ToolCallback[] getToolCallbacks() {
|
try {
|
ToolCallback[] callbacks = ensureToolCallbackProvider().getToolCallbacks();
|
return callbacks == null ? new ToolCallback[0] : callbacks;
|
} catch (Exception e) {
|
log.warn("Failed to load MCP tools through SSE client, baseUrl={}, sseEndpoint={}",
|
resolveBaseUrl(), resolveClientSseEndpoint(), e);
|
resetClientSession();
|
return new ToolCallback[0];
|
}
|
}
|
|
private Object parseToolResult(String rawResult) {
|
if (rawResult == null || rawResult.trim().isEmpty()) {
|
return rawResult;
|
}
|
try {
|
return JSON.parse(rawResult);
|
} catch (Exception ignore) {
|
return rawResult;
|
}
|
}
|
|
@SuppressWarnings("unchecked")
|
private Map<String, Object> parseSchema(String inputSchema) {
|
if (inputSchema == null || inputSchema.trim().isEmpty()) {
|
return Collections.emptyMap();
|
}
|
try {
|
Object parsed = JSON.parse(inputSchema);
|
if (parsed instanceof Map) {
|
return new LinkedHashMap<String, Object>((Map<String, Object>) parsed);
|
}
|
} catch (Exception e) {
|
log.warn("Failed to parse MCP tool schema: {}", inputSchema, e);
|
}
|
Map<String, Object> fallback = new LinkedHashMap<String, Object>();
|
fallback.put("type", "object");
|
return fallback;
|
}
|
|
private ToolCallbackProvider ensureToolCallbackProvider() {
|
return ensureClientSession().toolCallbackProvider;
|
}
|
|
private ClientSession ensureClientSession() {
|
ClientSession current = clientSession;
|
if (current != null) {
|
return current;
|
}
|
|
synchronized (clientMonitor) {
|
current = clientSession;
|
if (current != null) {
|
return current;
|
}
|
|
String baseUrl = resolveBaseUrl();
|
String clientSseEndpoint = resolveClientSseEndpoint();
|
HttpClientSseClientTransport transport = HttpClientSseClientTransport.builder(baseUrl)
|
.sseEndpoint(clientSseEndpoint)
|
.connectTimeout(requestTimeout)
|
.build();
|
|
McpSyncClient syncClient = McpClient.sync(transport)
|
.clientInfo(CLIENT_INFO)
|
.requestTimeout(requestTimeout)
|
.initializationTimeout(requestTimeout)
|
.build();
|
syncClient.initialize();
|
|
SyncMcpToolCallbackProvider callbackProvider = new SyncMcpToolCallbackProvider(syncClient);
|
current = new ClientSession(syncClient, callbackProvider, baseUrl);
|
clientSession = current;
|
log.info("Spring AI MCP SSE client initialized, baseUrl={}, sseEndpoint={}, tools={}",
|
baseUrl, clientSseEndpoint, current.toolCallbackProvider.getToolCallbacks().length);
|
return current;
|
}
|
}
|
|
private void resetClientSession() {
|
synchronized (clientMonitor) {
|
ClientSession current = clientSession;
|
clientSession = null;
|
if (current != null) {
|
current.close();
|
}
|
}
|
}
|
|
private String resolveBaseUrl() {
|
if (configuredBaseUrl != null && !configuredBaseUrl.trim().isEmpty()) {
|
return trimTrailingSlash(configuredBaseUrl.trim());
|
}
|
StringBuilder url = new StringBuilder("http://127.0.0.1:");
|
url.append(serverPort == null ? 9090 : serverPort);
|
return trimTrailingSlash(url.toString());
|
}
|
|
private String resolveClientSseEndpoint() {
|
String endpoint = normalizePath(sseEndpoint);
|
if (configuredBaseUrl != null && !configuredBaseUrl.trim().isEmpty()) {
|
return endpoint;
|
}
|
String context = normalizeContextPath(contextPath);
|
if (context.isEmpty()) {
|
return endpoint;
|
}
|
return context + endpoint;
|
}
|
|
private String normalizeContextPath(String path) {
|
if (path == null || path.trim().isEmpty() || "/".equals(path.trim())) {
|
return "";
|
}
|
String value = path.trim();
|
if (!value.startsWith("/")) {
|
value = "/" + value;
|
}
|
return trimTrailingSlash(value);
|
}
|
|
private String normalizePath(String path) {
|
if (path == null || path.trim().isEmpty()) {
|
return "/";
|
}
|
String value = path.trim();
|
if (!value.startsWith("/")) {
|
value = "/" + value;
|
}
|
return value;
|
}
|
|
private String trimTrailingSlash(String value) {
|
if (value == null || value.isEmpty()) {
|
return "";
|
}
|
return value.endsWith("/") && value.length() > 1 ? value.substring(0, value.length() - 1) : value;
|
}
|
|
@PreDestroy
|
public void destroy() {
|
resetClientSession();
|
}
|
|
private static final class ClientSession implements AutoCloseable {
|
|
private final McpSyncClient syncClient;
|
private final ToolCallbackProvider toolCallbackProvider;
|
private final String baseUrl;
|
|
private ClientSession(McpSyncClient syncClient, ToolCallbackProvider toolCallbackProvider, String baseUrl) {
|
this.syncClient = syncClient;
|
this.toolCallbackProvider = toolCallbackProvider;
|
this.baseUrl = baseUrl;
|
}
|
|
@Override
|
public void close() {
|
try {
|
syncClient.closeGracefully();
|
} catch (Exception e) {
|
log.debug("Close MCP SSE client gracefully failed, baseUrl={}", baseUrl, e);
|
}
|
try {
|
syncClient.close();
|
} catch (Exception e) {
|
log.debug("Close MCP SSE client failed, baseUrl={}", baseUrl, e);
|
}
|
}
|
}
|
}
|
