背景

智能助手项目需要接入 MCP 功能。一方面,由于多租户特性需求,系统需在用户提问时根据用户ID、租户或地区等维度动态实时加载 Mcp Server 及 Tools。另一方面因业务需要进行Tools入参进行处理,并记录原始第三方工具返回值。当前 Spring AI 1.1.2 版本仅支持通过配置文件方式加载 Mcp Server,且启动时即完成加载,无法实现基于用户ID或地区的动态加载需求。对于工具的入参与返回值处理官方框架不能直接支持。因此需要自主开发实现该功能。

说明:Mcp Server Tools是对OpenAPI JSON的包装实现通用快捷的方式定义工具,故自定义逻辑不可在此添加。

实现

1.依赖

<dependency>
    <groupId>org.springframework.ai</groupId>
    <artifactId>spring-ai-starter-mcp-client-webflux</artifactId>
</dependency>

2. 建立自定義的 MCP 工具包裝器 (ToolCallback Wrapper)



import java.util.Map;

public class CustomMcpToolWrapper implements ToolCallback {

    private final ToolCallback delegate;
    private final ChatService chatService; 

    public CustomMcpToolWrapper(ToolCallback delegate, chatService chatService) {
        this.delegate = delegate;
        this.chatService = chatService;
    }

    @Override
    public ToolDefinition getToolDefinition() {
        return delegate.getToolDefinition();
    }

    @Override
    public String call(String toolInput, ToolContext toolContext) {
		
        String toolName = delegate.getToolDefinition().name();
        String modifiedInput = toolInput;


        Map<String, Object> ctx = toolContext.getContext();
        String id = String.valueOf(ctx.get("id"));
        String empNo = String.valueOf(ctx.get("uid"));
        String agentCode = String.valueOf(ctx.get("agentCode"));

		//工具入參限定處理(例如:只能查自己的打卡記錄)  配置化管理 特定工具
        AiMcpToolT toolCfg=chatService.getMcpToolBy(toolName);
        if(ObjectUtils.isNotEmpty(toolCfg)){
            JSONObject inputJson = JSON.parseObject(toolInput);
            inputJson.put("uid", uid);
            modifiedInput = inputJson.toString();
        }


        
        String rawResponse;
        try {
            rawResponse = delegate.call(modifiedInput, toolContext);
            JSONArray array = JSON.parseArray(rawResponse);
            JSONObject firstObj = array.getJSONObject(0);
            String textValue = firstObj.getString("text");
            rawResponse=textValue;

        } catch (Exception e) {
            updateChatRecordStatus(id, "A"); 
            throw e;
        }

        //工具返回值限定處理(需要特定方式顯示)  配置化管理 特定工具
        if (ObjectUtils.isNotEmpty(toolCfg)) {
            handleOpenApiResponse(id, agentCode, rawResponse);
        }


        return rawResponse;
    }


    private void handleOpenApiResponse(String id, String agentCode, String resultJson) {
    	
	
    	String status="D";
        try {
            JSONObject rv = JSON.parseObject(resultJson);

            if(rv.getIntValue("code") == 1) {
                JSONObject arrays = rv.getJSONObject("answerData");
                if (arrays != null && !arrays.isEmpty()) {
                    AiChatRecordExtT ext = new AiChatRecordExtT();
                    ext.setRecordId(Long.valueOf(id));
                    ext.setAnswerData(arrays.toString());
                    ext.setAgentCode(agentCode);
                    chatService.addChatRecordExt(ext);
                    status= StringUtils.isNotEmpty(rv.getString("answerType"))?rv.getString("answerType"):"D";
                } else {
                	status="A";
                }
            } else {
            	status="A";
            }
        
        } catch (Exception e) {
          	status="A";
        }
        updateChatStatus(id, status);
        
    }

    private void updateChatStatus(String id, String status) {
    	
    	
        AiChatRecordT tp = new AiChatRecordT();
        tp.setId(Long.valueOf(id));
        tp.setAnswerType(status);
        chatService.updateChatRecordById(tp);
    }

	@Override
	public String call(String toolInput) {
		String rawResponse = delegate.call(toolInput);
		return rawResponse;
	}
}

3.构建上下文存储Mcp Client

public final class TenantMcpContext implements AutoCloseable {

 private final String key;
 private final List<McpSyncClient> clients;
 private final SyncMcpToolCallbackProvider provider;

 public TenantMcpContext(String key,
                         List<McpSyncClient> clients,
                         SyncMcpToolCallbackProvider provider) {
     this.key = key;
     this.clients = List.copyOf(clients);
     this.provider = provider;
 }

 public String key() {
     return key;
 }

 public ToolCallback[] toolCallbacks() {
     return provider.getToolCallbacks();
 }

 @Override
 public void close() {
     for (McpSyncClient client : clients) {
         try {
             client.closeGracefully();
         } catch (Exception ignored) {
         }
     }
 }
}

4.动态构建连接到McpServer 的McpClient,并优化无可用McpServer配置时,McpServer不可用时代码逻辑

@Service
public class TenantMcpRegistry {
	private static final Logger log = LoggerFactory.getLogger(TenantMcpRegistry.class);
	private final WebClient.Builder webClientBuilder;
	private final ObjectMapper objectMapper;
	private final AiMcpConfigTService aiMcpConfigTService;

	private final ConcurrentMap<String, TenantMcpContext> cache = new ConcurrentHashMap<>();
	private final ConcurrentMap<String, ReentrantLock> locks = new ConcurrentHashMap<>();

	public TenantMcpRegistry(WebClient.Builder webClientBuilder, ObjectMapper objectMapper,
			AiMcpConfigTService aiMcpConfigTService) {
		this.webClientBuilder = webClientBuilder;
		this.objectMapper = objectMapper;
		this.aiMcpConfigTService = aiMcpConfigTService;
	}

	public ToolCallback[] resolveToolCallbacks(String flag, String agentCode, String factory) {
		return getOrLoad(flag, agentCode, factory).toolCallbacks();
	}

	public TenantMcpContext refresh(String flag, String agentCode, String factory) {
		String key = buildKey(flag, agentCode, factory);
		ReentrantLock lock = locks.computeIfAbsent(key, k -> new ReentrantLock());
		lock.lock();
		try {
			TenantMcpContext fresh = buildContext(flag, agentCode, factory);
			TenantMcpContext old = cache.put(key, fresh);
			closeQuietly(old);
			return fresh;
		} finally {
			lock.unlock();
		}
	}

	public void evict(String flag, String agentCode, String factory) {
		String key = buildKey(flag, agentCode, factory);
		ReentrantLock lock = locks.computeIfAbsent(key, k -> new ReentrantLock());
		lock.lock();
		try {
			TenantMcpContext old = cache.remove(key);
			closeQuietly(old);
		} finally {
			lock.unlock();
		}
	}

	private TenantMcpContext getOrLoad(String flag, String agentCode, String factory) {
		String key = buildKey(flag, agentCode, factory);

		TenantMcpContext cached = cache.get(key);
		if (cached != null) {
			return cached;
		}

		ReentrantLock lock = locks.computeIfAbsent(key, k -> new ReentrantLock());
		lock.lock();
		try {
			TenantMcpContext again = cache.get(key);
			if (again != null) {
				return again;
			}
			TenantMcpContext created = buildContext(flag, agentCode, factory);
			cache.put(key, created);
			return created;
		} finally {
			lock.unlock();
		}
	}

	private TenantMcpContext buildContext(String flag, String agentCode, String factory) {

		List<AiMcpConfigT> dbConfigs = aiMcpConfigTService.getMcpConfigByFlag(flag, agentCode, factory);

		if (dbConfigs == null || dbConfigs.isEmpty()) {
			log.info("數據庫中無有效的MCP Server配置。");
			return createEmptyContext(flag, agentCode, factory);
		}
		
		List<McpSyncClient> clients = new ArrayList<>();

		for (AiMcpConfigT config : dbConfigs) {
			McpSyncClient client = null;
			try {
				HttpClient directHttpClient = HttpClient.newBuilder().connectTimeout(Duration.ofSeconds(60)).build();

				WebClient.Builder builderClone = webClientBuilder.clone()
						.baseUrl(config.getBaseUrl())
						.clientConnector(new org.springframework.http.client.reactive.JdkClientHttpConnector(directHttpClient));

				var transport = io.modelcontextprotocol.client.transport.WebClientStreamableHttpTransport
						.builder(builderClone).endpoint(config.getEndPoint())
						.jsonMapper(new io.modelcontextprotocol.json.jackson.JacksonMcpJsonMapper(objectMapper))
						.build();

				client = McpClient.sync(transport).requestTimeout(Duration.ofSeconds(60)).build();
				
				client.initialize();
				
				clients.add(client);
				log.info("成功載入並初始化 MCP Server: {} (URL: {})", config.getName(), config.getBaseUrl());
				
			} catch (Exception ex) {
				log.warn("MCP Server 連接或初始化失敗,將跳過此配置: name={}, url={}. 錯誤原因: {}", 
						config.getName(), config.getBaseUrl(), ex.getMessage());
				
				if (client != null) {
					try {
						client.closeGracefully();
					} catch (Exception ignored) {
					}
				}
			}
		}

		if (clients.isEmpty()) {
			return createEmptyContext(flag, agentCode, factory);
		}

		try {
			SyncMcpToolCallbackProvider provider = SyncMcpToolCallbackProvider.builder().mcpClients(clients).build();
			return new TenantMcpContext(buildKey(flag, agentCode, factory), clients, provider);
		} catch (Exception ex) {
			closeClients(clients);
			log.error("建立 MCP Provider 失敗, factory=" + factory, ex);
			throw new IllegalStateException("構建 MCP 執行環境失敗, factory=" + factory, ex);
		}
	}


	private String buildKey(String flag, String agentCode, String factory) {
		return flag + "|" + agentCode + "|" + factory;
	}

	private void closeQuietly(TenantMcpContext ctx) {
		if (ctx != null) {
			try {
				ctx.close();
			} catch (Exception ignored) {
			}
		}
	}

	private void closeClients(List<McpSyncClient> clients) {
		for (McpSyncClient client : clients) {
			try {
				client.closeGracefully();
			} catch (Exception ignored) {
			}
		}
	}

	@PreDestroy
	public void destroy() {
		for (TenantMcpContext ctx : cache.values()) {
			closeQuietly(ctx);
		}
		cache.clear();
	}
	

    private TenantMcpContext createEmptyContext(String flag, String agentCode, String factory) {
        String key = buildKey(flag, agentCode, factory);
        SyncMcpToolCallbackProvider emptyProvider = SyncMcpToolCallbackProvider.builder()
                .mcpClients(Collections.emptyList())
                .build();
        return new TenantMcpContext(key, Collections.emptyList(), emptyProvider);
    }
	
	
}

5.McpServer配置的表结构(STREAM_HTTP)

McpServerA  https://sm.mcp.com  /mi5Y/mcp  Y TenantA Guizhou
 McpServerb   http://a.adc.com  /mcp  Y TenantB Xian

6.实际使用

ToolCallback[] toolCallbacks = tenantMcpRegistry.resolveToolCallbacks("Y", agentCode, "NA");

List<ToolCallback> wrappedTools = Arrays.stream(toolCallbacks)

.map(tool -> new CustomMcpToolWrapper(tool, chatRecordService))

.collect(Collectors.toList());


 

Logo

AtomGit 是由开放原子开源基金会联合 CSDN 等生态伙伴共同推出的新一代开源与人工智能协作平台。平台坚持“开放、中立、公益”的理念,把代码托管、模型共享、数据集托管、智能体开发体验和算力服务整合在一起,为开发者提供从开发、训练到部署的一站式体验。

更多推荐