From 46714df60478ba09e2ad442d3f670bc1de26fd9c Mon Sep 17 00:00:00 2001 From: ageerle Date: Wed, 23 Sep 2026 17:35:58 +0800 Subject: [PATCH] feat: support bearer key in remote MCP tool config --- .../org/ruoyi/domain/entity/mcp/McpTool.java | 2 +- .../LangChain4jMcpToolProviderService.java | 18 ++- .../service/core/RemoteMcpBearerKeyTest.java | 127 ++++++++++++++++++ 3 files changed, 142 insertions(+), 5 deletions(-) create mode 100644 ruoyi-modules/ruoyi-chat/src/test/java/org/ruoyi/mcp/service/core/RemoteMcpBearerKeyTest.java diff --git a/ruoyi-modules/ruoyi-chat/src/main/java/org/ruoyi/domain/entity/mcp/McpTool.java b/ruoyi-modules/ruoyi-chat/src/main/java/org/ruoyi/domain/entity/mcp/McpTool.java index 8ce47880..b0415611 100644 --- a/ruoyi-modules/ruoyi-chat/src/main/java/org/ruoyi/domain/entity/mcp/McpTool.java +++ b/ruoyi-modules/ruoyi-chat/src/main/java/org/ruoyi/domain/entity/mcp/McpTool.java @@ -47,7 +47,7 @@ public class McpTool extends TenantEntity { /** * 配置信息（JSON格式） * LOCAL: {"command": "npx", "args": ["-y", "@example/mcp-server"], "env": {...}} - * REMOTE: {"baseUrl": "http://localhost:8080/mcp"} + * REMOTE: {"baseUrl": "https://example.com/mcp", "bearerKey": "optional-key"} * BUILTIN: null */ private String configJson; diff --git a/ruoyi-modules/ruoyi-chat/src/main/java/org/ruoyi/mcp/service/core/LangChain4jMcpToolProviderService.java b/ruoyi-modules/ruoyi-chat/src/main/java/org/ruoyi/mcp/service/core/LangChain4jMcpToolProviderService.java index d6fb504d..bb47a12c 100644 --- a/ruoyi-modules/ruoyi-chat/src/main/java/org/ruoyi/mcp/service/core/LangChain4jMcpToolProviderService.java +++ b/ruoyi-modules/ruoyi-chat/src/main/java/org/ruoyi/mcp/service/core/LangChain4jMcpToolProviderService.java @@ -525,11 +525,21 @@ public class LangChain4jMcpToolProviderService { String baseUrl = configNode.get("baseUrl").asText(); log.info("mcp_client transport=HTTP status=CREATING"); - // 创建 HTTP/SSE 传输层 - McpTransport transport = StreamableHttpMcpTransport.builder() + // bearerKey is optional and belongs to this REMOTE tool's client only. + StreamableHttpMcpTransport.Builder transportBuilder = StreamableHttpMcpTransport.builder() .url(baseUrl) - .logRequests(TRAFFIC_LOGGING_ENABLED) - .build(); + .logRequests(TRAFFIC_LOGGING_ENABLED); + JsonNode bearerKeyNode = configNode.get("bearerKey"); + if (bearerKeyNode != null && !bearerKeyNode.isNull()) { + if (!bearerKeyNode.isTextual()) { + throw new IllegalArgumentException("bearerKey must be a string"); + } + String bearerKey = bearerKeyNode.asText().trim(); + if (!bearerKey.isEmpty()) { + transportBuilder.customHeaders(Map.of("Authorization", "Bearer " + bearerKey)); + } + } + McpTransport transport = transportBuilder.build(); // 创建客户端 return new DefaultMcpClient.Builder() diff --git a/ruoyi-modules/ruoyi-chat/src/test/java/org/ruoyi/mcp/service/core/RemoteMcpBearerKeyTest.java b/ruoyi-modules/ruoyi-chat/src/test/java/org/ruoyi/mcp/service/core/RemoteMcpBearerKeyTest.java new file mode 100644 index 00000000..1be1a2a8 --- /dev/null +++ b/ruoyi-modules/ruoyi-chat/src/test/java/org/ruoyi/mcp/service/core/RemoteMcpBearerKeyTest.java @@ -0,0 +1,127 @@ +package org.ruoyi.mcp.service.core; + +import com.baomidou.mybatisplus.core.conditions.Wrapper; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ObjectNode; +import com.sun.net.httpserver.HttpServer; +import dev.langchain4j.agent.tool.ToolExecutionRequest; +import dev.langchain4j.agentic.AgenticServices; +import dev.langchain4j.data.message.AiMessage; +import dev.langchain4j.data.message.ToolExecutionResultMessage; +import dev.langchain4j.model.chat.ChatModel; +import dev.langchain4j.model.chat.request.ChatRequest; +import dev.langchain4j.model.chat.response.ChatResponse; +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; +import org.ruoyi.agent.WebSearchAgent; +import org.ruoyi.domain.entity.mcp.McpTool; +import org.ruoyi.mapper.mcp.McpToolMapper; + +import java.net.InetSocketAddress; +import java.nio.charset.StandardCharsets; +import java.util.List; +import java.util.concurrent.CopyOnWriteArrayList; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +@Tag("dev") +class RemoteMcpBearerKeyTest { + + private final ObjectMapper json = new ObjectMapper(); + + @Test + void selectedRemoteToolSendsBearerKeyDuringDiscoveryAndCall() throws Exception { + List requests = runAgentWithRemoteTool("synthetic-key"); + assertTrue(requests.stream().anyMatch(r -> "initialize".equals(r.method()))); + assertTrue(requests.stream().anyMatch(r -> "tools/list".equals(r.method()))); + assertTrue(requests.stream().anyMatch(r -> "tools/call".equals(r.method()))); + assertTrue(requests.stream().allMatch(r -> "Bearer synthetic-key".equals(r.authorization()))); + } + + @Test + void remoteToolWithoutBearerKeyRemainsAnonymous() throws Exception { + List requests = runAgentWithRemoteTool(null); + assertTrue(requests.stream().anyMatch(r -> "tools/call".equals(r.method()))); + assertTrue(requests.stream().allMatch(r -> r.authorization() == null)); + } + + private List runAgentWithRemoteTool(String bearerKey) throws Exception { + List requests = new CopyOnWriteArrayList<>(); + HttpServer server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0); + server.createContext("/mcp", exchange -> { + if (!"POST".equals(exchange.getRequestMethod())) { + exchange.sendResponseHeaders(204, -1); + exchange.close(); + return; + } + var request = json.readTree(exchange.getRequestBody()); + requests.add(new RequestHeader(request.path("method").asText(), + exchange.getRequestHeaders().getFirst("Authorization"))); + if (!request.has("id")) { + exchange.sendResponseHeaders(202, -1); + exchange.close(); + return; + } + String result = switch (request.path("method").asText()) { + case "initialize" -> "{\"protocolVersion\":\"2025-06-18\",\"capabilities\":{\"tools\":{}},\"serverInfo\":{\"name\":\"fixture\",\"version\":\"1\"}}"; + case "tools/list" -> "{\"tools\":[{\"name\":\"fixture_search\",\"inputSchema\":{\"type\":\"object\",\"properties\":{}}}]}"; + case "tools/call" -> "{\"content\":[{\"type\":\"text\",\"text\":\"fixture result\"}],\"isError\":false}"; + default -> "{}"; + }; + byte[] body = ("{\"jsonrpc\":\"2.0\",\"id\":" + request.get("id") + ",\"result\":" + result + "}") + .getBytes(StandardCharsets.UTF_8); + exchange.getResponseHeaders().set("Content-Type", "application/json"); + exchange.sendResponseHeaders(200, body.length); + exchange.getResponseBody().write(body); + exchange.close(); + }); + server.start(); + + McpTool tool = new McpTool(); + tool.setId(41L); + tool.setName("Fixture remote search"); + tool.setType("REMOTE"); + tool.setStatus("ENABLED"); + ObjectNode config = json.createObjectNode() + .put("baseUrl", "http://127.0.0.1:" + server.getAddress().getPort() + "/mcp"); + if (bearerKey != null) { + config.put("bearerKey", bearerKey); + } + tool.setConfigJson(json.writeValueAsString(config)); + McpToolMapper mapper = mock(McpToolMapper.class); + when(mapper.selectList(any(Wrapper.class))).thenReturn(List.of(tool)); + var service = new LangChain4jMcpToolProviderService(mapper, json, mock(BuiltinToolRegistry.class)); + try { + ChatModel router = new ChatModel() { + @Override + public ChatResponse doChat(ChatRequest request) { + var last = request.messages().get(request.messages().size() - 1); + if (last instanceof ToolExecutionResultMessage result) { + return ChatResponse.builder().aiMessage(AiMessage.from(result.text())).build(); + } + assertTrue(request.toolSpecifications().stream() + .anyMatch(spec -> "fixture_search".equals(spec.name()))); + return ChatResponse.builder().aiMessage(AiMessage.from(ToolExecutionRequest.builder() + .id("call").name("fixture_search").arguments("{}").build())).build(); + } + }; + String answer = AgenticServices.agentBuilder(WebSearchAgent.class) + .chatModel(router) + .toolProvider(service.getToolProvider(List.of(41L))) + .build() + .search("Use the selected search tool"); + assertTrue(answer.contains("fixture result")); + } finally { + service.cleanup(); + server.stop(0); + } + return requests.stream().filter(r -> !r.method().startsWith("notifications/")).toList(); + } + + private record RequestHeader(String method, String authorization) { + } +}