package com.zy.ai.mcp.service;
|
|
import com.alibaba.fastjson.JSON;
|
import com.alibaba.fastjson.JSONObject;
|
import com.zy.ai.entity.AiMcpMount;
|
import com.zy.ai.enums.AiMcpTransportType;
|
import com.zy.ai.service.AiMcpMountService;
|
import io.modelcontextprotocol.client.McpClient;
|
import io.modelcontextprotocol.client.McpSyncClient;
|
import io.modelcontextprotocol.client.transport.HttpClientSseClientTransport;
|
import io.modelcontextprotocol.client.transport.HttpClientStreamableHttpTransport;
|
import io.modelcontextprotocol.spec.McpSchema;
|
import lombok.RequiredArgsConstructor;
|
import lombok.extern.slf4j.Slf4j;
|
import org.springframework.ai.mcp.SyncMcpToolCallbackProvider;
|
import org.springframework.ai.tool.ToolCallback;
|
import org.springframework.ai.tool.definition.DefaultToolDefinition;
|
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.net.URI;
|
import java.time.Duration;
|
import java.util.ArrayList;
|
import java.util.Collections;
|
import java.util.Comparator;
|
import java.util.LinkedHashMap;
|
import java.util.LinkedHashSet;
|
import java.util.List;
|
import java.util.Map;
|
|
@Slf4j
|
@Service
|
@RequiredArgsConstructor
|
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();
|
private final AiMcpMountService aiMcpMountService;
|
|
@Value("${spring.ai.mcp.server.request-timeout:20s}")
|
private Duration defaultRequestTimeout;
|
|
private volatile ClientRegistry clientRegistry;
|
|
public List<Map<String, Object>> listTools() {
|
return ensureClientRegistry().toolList;
|
}
|
|
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 (isBlank(toolName)) {
|
throw new IllegalArgumentException("missing tool name");
|
}
|
|
MountedTool mountedTool = ensureClientRegistry().toolMap.get(toolName);
|
if (mountedTool == null) {
|
throw new IllegalArgumentException("tool not found: " + toolName);
|
}
|
|
try {
|
String rawResult = mountedTool.callback.call(arguments == null ? "{}" : arguments.toJSONString());
|
return parseToolResult(rawResult);
|
} catch (Exception e) {
|
evictCache();
|
throw e;
|
}
|
}
|
|
public Map<String, Object> testMount(AiMcpMount mount) {
|
if (mount == null) {
|
throw new IllegalArgumentException("参数不能为空");
|
}
|
MountSession session = null;
|
long start = System.currentTimeMillis();
|
try {
|
session = openSession(mount);
|
LinkedHashMap<String, Object> result = new LinkedHashMap<String, Object>();
|
result.put("ok", true);
|
result.put("message", "连接成功，已发现 " + session.callbacks.length + " 个工具");
|
result.put("latencyMs", System.currentTimeMillis() - start);
|
result.put("url", session.url);
|
result.put("transportType", mount.getTransportType());
|
result.put("toolCount", session.callbacks.length);
|
result.put("toolNames", collectToolNames(session.callbacks, mount, reservedToolNames(mount.getMountCode())));
|
return result;
|
} catch (Exception e) {
|
TransportTarget target = resolveTransportTarget(mount);
|
LinkedHashMap<String, Object> result = new LinkedHashMap<String, Object>();
|
result.put("ok", false);
|
result.put("message", safeMessage(e));
|
result.put("latencyMs", System.currentTimeMillis() - start);
|
result.put("url", target.url);
|
result.put("transportType", mount.getTransportType());
|
result.put("toolCount", 0);
|
result.put("toolNames", new ArrayList<String>());
|
return result;
|
} finally {
|
if (session != null) {
|
session.close();
|
}
|
}
|
}
|
|
public void evictCache() {
|
resetClientRegistry();
|
}
|
|
private ClientRegistry ensureClientRegistry() {
|
ClientRegistry current = clientRegistry;
|
if (current != null) {
|
return current;
|
}
|
|
synchronized (clientMonitor) {
|
current = clientRegistry;
|
if (current != null) {
|
return current;
|
}
|
|
current = buildClientRegistry();
|
clientRegistry = current;
|
return current;
|
}
|
}
|
|
private ClientRegistry buildClientRegistry() {
|
List<AiMcpMount> mounts = loadEnabledMounts();
|
List<MountSession> sessions = new ArrayList<MountSession>();
|
LinkedHashMap<String, MountedTool> toolMap = new LinkedHashMap<String, MountedTool>();
|
List<Map<String, Object>> toolList = new ArrayList<Map<String, Object>>();
|
|
for (AiMcpMount mount : mounts) {
|
if (mount == null) {
|
continue;
|
}
|
MountSession session = null;
|
try {
|
session = openSession(mount);
|
sessions.add(session);
|
for (ToolCallback callback : session.callbacks) {
|
MountedTool mountedTool = buildMountedTool(mount, callback, toolMap.keySet());
|
toolMap.put(mountedTool.toolName, mountedTool);
|
toolList.add(toToolDescriptor(mountedTool));
|
}
|
log.info("MCP mount loaded, mountCode={}, transport={}, url={}, tools={}",
|
mount.getMountCode(), mount.getTransportType(), session.url, session.callbacks.length);
|
} catch (Exception e) {
|
log.warn("Failed to load MCP mount, mountCode={}, transport={}, url={}",
|
mount.getMountCode(), mount.getTransportType(), resolveUrl(mount), e);
|
if (session != null) {
|
session.close();
|
}
|
}
|
}
|
|
toolList.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 new ClientRegistry(sessions, toolMap, toolList);
|
}
|
|
private List<AiMcpMount> loadEnabledMounts() {
|
try {
|
List<AiMcpMount> mounts = aiMcpMountService.listEnabledOrdered();
|
if (mounts == null || mounts.isEmpty()) {
|
aiMcpMountService.initDefaultsIfMissing();
|
mounts = aiMcpMountService.listEnabledOrdered();
|
}
|
return mounts == null ? Collections.<AiMcpMount>emptyList() : mounts;
|
} catch (Exception e) {
|
log.warn("Failed to query MCP mount configuration", e);
|
return Collections.emptyList();
|
}
|
}
|
|
private MountedTool buildMountedTool(AiMcpMount mount, ToolCallback callback, java.util.Set<String> usedNames) {
|
ToolDefinition definition = callback == null ? null : callback.getToolDefinition();
|
if (definition == null || isBlank(definition.name())) {
|
throw new IllegalArgumentException("invalid tool definition");
|
}
|
|
String originalName = definition.name();
|
String preferredName = mount.getMountCode() + "_" + originalName;
|
String finalName = ensureUniqueToolName(preferredName, mount, originalName, usedNames);
|
ToolCallback effectiveCallback = originalName.equals(finalName)
|
? callback
|
: new MountedToolCallback(finalName, callback);
|
|
return new MountedTool(mount, originalName, finalName, effectiveCallback);
|
}
|
|
private String ensureUniqueToolName(String preferredName,
|
AiMcpMount mount,
|
String originalName,
|
java.util.Set<String> usedNames) {
|
if (!usedNames.contains(preferredName)) {
|
return preferredName;
|
}
|
|
String fallbackBase = mount.getMountCode() + "_" + originalName;
|
if (!usedNames.contains(fallbackBase)) {
|
log.warn("Duplicate MCP tool name detected, fallback rename applied, mountCode={}, originalName={}, finalName={}",
|
mount.getMountCode(), originalName, fallbackBase);
|
return fallbackBase;
|
}
|
|
int index = 2;
|
String candidate = fallbackBase + "_" + index;
|
while (usedNames.contains(candidate)) {
|
index++;
|
candidate = fallbackBase + "_" + index;
|
}
|
log.warn("Duplicate MCP tool name detected, numbered rename applied, mountCode={}, originalName={}, finalName={}",
|
mount.getMountCode(), originalName, candidate);
|
return candidate;
|
}
|
|
private Map<String, Object> toToolDescriptor(MountedTool mountedTool) {
|
ToolDefinition definition = mountedTool.callback.getToolDefinition();
|
Map<String, Object> item = new LinkedHashMap<String, Object>();
|
item.put("name", definition.name());
|
item.put("originalName", mountedTool.originalName);
|
item.put("mountCode", mountedTool.mount.getMountCode());
|
item.put("mountName", mountedTool.mount.getName());
|
item.put("transportType", mountedTool.mount.getTransportType());
|
item.put("description", definition.description());
|
item.put("inputSchema", parseSchema(definition.inputSchema()));
|
return item;
|
}
|
|
private List<String> collectToolNames(ToolCallback[] callbacks, AiMcpMount mount, java.util.Set<String> reservedNames) {
|
List<String> names = new ArrayList<String>();
|
java.util.Set<String> used = new LinkedHashSet<String>();
|
if (reservedNames != null && !reservedNames.isEmpty()) {
|
used.addAll(reservedNames);
|
}
|
if (callbacks == null) {
|
return names;
|
}
|
for (ToolCallback callback : callbacks) {
|
ToolDefinition definition = callback == null ? null : callback.getToolDefinition();
|
if (definition == null || isBlank(definition.name())) {
|
continue;
|
}
|
String preferred = mount.getMountCode() + "_" + definition.name();
|
String finalName = ensureUniqueToolName(preferred, mount, definition.name(), used);
|
used.add(finalName);
|
names.add(finalName);
|
}
|
return names;
|
}
|
|
private java.util.Set<String> reservedToolNames(String mountCode) {
|
LinkedHashSet<String> reserved = new LinkedHashSet<String>();
|
ClientRegistry current = clientRegistry;
|
if (current == null || current.toolMap == null || current.toolMap.isEmpty()) {
|
return reserved;
|
}
|
for (MountedTool mountedTool : current.toolMap.values()) {
|
if (mountedTool == null || mountedTool.mount == null) {
|
continue;
|
}
|
if (mountCode != null && mountCode.equals(mountedTool.mount.getMountCode())) {
|
continue;
|
}
|
reserved.add(mountedTool.toolName);
|
}
|
return reserved;
|
}
|
|
private MountSession openSession(AiMcpMount mount) {
|
Duration timeout = resolveTimeout(mount);
|
TransportTarget target = resolveTransportTarget(mount);
|
String baseUrl = target.baseUrl;
|
String endpoint = target.endpoint;
|
AiMcpTransportType transportType = AiMcpTransportType.ofCode(mount.getTransportType());
|
McpSyncClient syncClient;
|
|
if (transportType == AiMcpTransportType.STREAMABLE_HTTP) {
|
HttpClientStreamableHttpTransport transport = HttpClientStreamableHttpTransport.builder(baseUrl)
|
.endpoint(endpoint)
|
.connectTimeout(timeout)
|
.build();
|
syncClient = McpClient.sync(transport)
|
.clientInfo(CLIENT_INFO)
|
.requestTimeout(timeout)
|
.initializationTimeout(timeout)
|
.build();
|
} else {
|
HttpClientSseClientTransport transport = HttpClientSseClientTransport.builder(baseUrl)
|
.sseEndpoint(endpoint)
|
.connectTimeout(timeout)
|
.build();
|
syncClient = McpClient.sync(transport)
|
.clientInfo(CLIENT_INFO)
|
.requestTimeout(timeout)
|
.initializationTimeout(timeout)
|
.build();
|
}
|
|
syncClient.initialize();
|
SyncMcpToolCallbackProvider callbackProvider = new SyncMcpToolCallbackProvider(syncClient);
|
ToolCallback[] callbacks = callbackProvider.getToolCallbacks();
|
return new MountSession(mount, syncClient, target.url, callbacks == null ? new ToolCallback[0] : callbacks);
|
}
|
|
private Duration resolveTimeout(AiMcpMount mount) {
|
if (mount != null && mount.getRequestTimeoutMs() != null && mount.getRequestTimeoutMs() > 0) {
|
return Duration.ofMillis(mount.getRequestTimeoutMs());
|
}
|
return defaultRequestTimeout == null ? Duration.ofSeconds(20) : defaultRequestTimeout;
|
}
|
|
private String resolveUrl(AiMcpMount mount) {
|
return trim(mount == null ? null : mount.getUrl());
|
}
|
|
private TransportTarget resolveTransportTarget(AiMcpMount mount) {
|
String rawUrl = resolveUrl(mount);
|
if (isBlank(rawUrl)) {
|
throw new IllegalArgumentException("missing url");
|
}
|
try {
|
URI finalUri = URI.create(rawUrl);
|
String authority = finalUri.getRawAuthority();
|
String scheme = finalUri.getScheme();
|
if (isBlank(scheme) || isBlank(authority)) {
|
throw new IllegalArgumentException("invalid MCP url");
|
}
|
String baseUrl = scheme + "://" + authority;
|
String endpoint = finalUri.getRawPath();
|
if (isBlank(endpoint)) {
|
throw new IllegalArgumentException("missing MCP path");
|
}
|
if (finalUri.getRawQuery() != null && !finalUri.getRawQuery().isEmpty()) {
|
endpoint = endpoint + "?" + finalUri.getRawQuery();
|
}
|
return new TransportTarget(rawUrl, baseUrl, endpoint);
|
} catch (IllegalArgumentException e) {
|
throw e;
|
} catch (Exception e) {
|
throw new IllegalArgumentException("invalid url: " + rawUrl, e);
|
}
|
}
|
|
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 String normalizePath(String path) {
|
if (isBlank(path)) {
|
return "/";
|
}
|
String value = path.trim();
|
if (!value.startsWith("/")) {
|
value = "/" + value;
|
}
|
return value;
|
}
|
|
private String safeMessage(Throwable throwable) {
|
if (throwable == null) {
|
return "unknown error";
|
}
|
if (!isBlank(throwable.getMessage())) {
|
return throwable.getMessage();
|
}
|
return throwable.getClass().getSimpleName();
|
}
|
|
private String trim(String text) {
|
return text == null ? null : text.trim();
|
}
|
|
private boolean isBlank(String text) {
|
return text == null || text.trim().isEmpty();
|
}
|
|
private void resetClientRegistry() {
|
synchronized (clientMonitor) {
|
ClientRegistry current = clientRegistry;
|
clientRegistry = null;
|
if (current != null) {
|
current.close();
|
}
|
}
|
}
|
|
@PreDestroy
|
public void destroy() {
|
resetClientRegistry();
|
}
|
|
private static final class ClientRegistry implements AutoCloseable {
|
|
private final List<MountSession> sessions;
|
private final Map<String, MountedTool> toolMap;
|
private final List<Map<String, Object>> toolList;
|
|
private ClientRegistry(List<MountSession> sessions,
|
Map<String, MountedTool> toolMap,
|
List<Map<String, Object>> toolList) {
|
this.sessions = sessions;
|
this.toolMap = toolMap;
|
this.toolList = toolList;
|
}
|
|
@Override
|
public void close() {
|
if (sessions == null) {
|
return;
|
}
|
for (MountSession session : sessions) {
|
if (session != null) {
|
session.close();
|
}
|
}
|
}
|
}
|
|
private static final class MountSession implements AutoCloseable {
|
|
private final AiMcpMount mount;
|
private final McpSyncClient syncClient;
|
private final String url;
|
private final ToolCallback[] callbacks;
|
|
private MountSession(AiMcpMount mount,
|
McpSyncClient syncClient,
|
String url,
|
ToolCallback[] callbacks) {
|
this.mount = mount;
|
this.syncClient = syncClient;
|
this.url = url;
|
this.callbacks = callbacks;
|
}
|
|
@Override
|
public void close() {
|
try {
|
syncClient.closeGracefully();
|
} catch (Exception e) {
|
log.debug("Close MCP client gracefully failed, mountCode={}", mount == null ? null : mount.getMountCode(), e);
|
}
|
try {
|
syncClient.close();
|
} catch (Exception e) {
|
log.debug("Close MCP client failed, mountCode={}", mount == null ? null : mount.getMountCode(), e);
|
}
|
}
|
}
|
|
private static final class MountedTool {
|
|
private final AiMcpMount mount;
|
private final String originalName;
|
private final String toolName;
|
private final ToolCallback callback;
|
|
private MountedTool(AiMcpMount mount, String originalName, String toolName, ToolCallback callback) {
|
this.mount = mount;
|
this.originalName = originalName;
|
this.toolName = toolName;
|
this.callback = callback;
|
}
|
}
|
|
private static final class MountedToolCallback implements ToolCallback {
|
|
private final ToolCallback delegate;
|
private final ToolDefinition definition;
|
|
private MountedToolCallback(String name, ToolCallback delegate) {
|
this.delegate = delegate;
|
ToolDefinition source = delegate.getToolDefinition();
|
this.definition = DefaultToolDefinition.builder()
|
.name(name)
|
.description(source == null ? "" : source.description())
|
.inputSchema(source == null ? "{}" : source.inputSchema())
|
.build();
|
}
|
|
@Override
|
public ToolDefinition getToolDefinition() {
|
return definition;
|
}
|
|
@Override
|
public org.springframework.ai.tool.metadata.ToolMetadata getToolMetadata() {
|
return delegate.getToolMetadata();
|
}
|
|
@Override
|
public String call(String toolInput) {
|
return delegate.call(toolInput);
|
}
|
|
@Override
|
public String call(String toolInput, org.springframework.ai.chat.model.ToolContext toolContext) {
|
return delegate.call(toolInput, toolContext);
|
}
|
}
|
|
private static final class TransportTarget {
|
|
private final String url;
|
private final String baseUrl;
|
private final String endpoint;
|
|
private TransportTarget(String url, String baseUrl, String endpoint) {
|
this.url = url;
|
this.baseUrl = baseUrl;
|
this.endpoint = endpoint;
|
}
|
}
|
}
|
