Skip to content

Commit baabb97

Browse files
predic8christiangoerdes
predic8
and
christiangoerdes
authored
MCPServer for Membrane (#2930)
* Refactor ROADMAP.md for clarity and remove outdated sections * feat: add JSON-RPC request, response, and MCP initialization support * refactor: streamline MCP request processing and response generation * refactor: remove redundant IOException declaration in toolsList method * Enhance MCP server error handling with detailed response generation * Refactor MCP server to replace MCPResponse with JSONRPCResponse for consistent handling * Add support for fetching Membrane runtime statistics in MCP server * Limit MCP server exchange retrieval to the last 100 entries * Refactor JSON-RPC request handling to distinguish between missing and null IDs, improving MCP server error handling and response generation. * Add ResponseKind to JSONRPCResponse for explicit result/error distinction and customize serialization * Improve JSON-RPC validation by rejecting notifications in MCPRequest parsing * Rename "lastExchanges" capability to "listChanged" in MCP server * Replace `Objects.requireNonNull` with static import `requireNonNull` for cleanup and consistency * Remove unnecessary `IOException` declarations from MCP server methods for cleanup * Add response details (status, body, headers) to MCP server exchange serialization * Reorder case blocks in `MembraneMCPServer` for consistency * Refactor error response handling in `MembraneMCPServer` to improve clarity and logging for notifications * Refactor JSON-RPC utilities to improve validation, normalization, and serialization handling * Add unit tests for `JSONRPCRequest` and `JSONRPCResponse` to validate parsing, normalization, and error handling * Introduce MCP tool registration and utilities, including `McpToolRegistry`, `McpToolDefinition`, and `McpToolHandler`. Minor cleanup and refactoring in MCP initialization and request handling. * Add `McpPayloadSanitizer` and `McpSessionContext` utilities for MCP payload sanitation and session state management * Remove unused imports in `McpPayloadSanitizer` for cleanup * Refactor and streamline MCP server request handling: add tool registration, session management, and improved error handling. * Refactor `MembraneMCPServer`: improve HTTP 405 response handling and clean up imports * Add unit tests for `MembraneMCPServer` and `MCPInitialize` to validate lifecycle handling, error scenarios, and MCP tool behaviors * Refactor `MembraneMCPServer` error handling: introduce `InvalidToolArgumentsException`, improve protocol version validation, and enhance test coverage. * Update `toString` implementation in `MCPToolsCall` to display argument keys instead of full argument details. * Extract utility methods from `MembraneMCPServer` to new `MCPUtil` class for improved modularity and code reuse. * Introduce session management in `MembraneMCPServer`: add `SESSION_HEADER` handling, session-based request processing, and enhanced error handling. Update corresponding unit tests for validation. --------- Co-authored-by: Christian Gördes <christian.goerdes@outlook.de> Co-authored-by: Christian Gördes <118011644+christiangoerdes@users.noreply.github.com>
1 parent a8fa588 commit baabb97

28 files changed

Lines changed: 3605 additions & 31 deletions

‎.gitignore‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -129,3 +129,4 @@ maven-plugin/target/surefire/
129129
/core/derby.log
130130
/distribution/conf/apis.yaml
131131
/distribution/conf/membrane.log
132+
/.codex

‎core/src/main/java/com/predic8/membrane/core/exchangestore/ClientStatistics.java‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -16,14 +16,14 @@
1616

1717
public interface ClientStatistics {
1818

19-
public int getCount();
19+
int getCount();
2020

21-
public String getClient();
21+
String getClient();
2222

23-
public long getMinDuration();
23+
long getMinDuration();
2424

25-
public long getMaxDuration();
25+
long getMaxDuration();
2626

27-
public long getAvgDuration();
27+
long getAvgDuration();
2828

2929
}

‎core/src/main/java/com/predic8/membrane/core/interceptor/mcp/MCPUtil.java‎

Lines changed: 114 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,114 @@
1+
package com.predic8.membrane.core.interceptor.mcp;
2+
3+
import com.predic8.membrane.core.exchange.AbstractExchange;
4+
import com.predic8.membrane.core.mcp.MCPToolsCall;
5+
import com.predic8.membrane.core.openapi.serviceproxy.APIProxy;
6+
import com.predic8.membrane.core.proxies.Proxy;
7+
import com.predic8.membrane.core.proxies.SOAPProxy;
8+
import com.predic8.membrane.core.proxies.ServiceProxy;
9+
import org.jetbrains.annotations.Nullable;
10+
11+
import java.util.LinkedHashMap;
12+
import java.util.Map;
13+
import java.util.Set;
14+
15+
public final class MCPUtil {
16+
17+
private MCPUtil() {
18+
}
19+
20+
public static Map<String, Object> describeProxy(Proxy proxy, Object statistics) {
21+
var description = new LinkedHashMap<String, Object>();
22+
description.put("name", proxy.getName());
23+
24+
String type;
25+
switch (proxy) {
26+
case APIProxy ignored -> type = "API";
27+
case SOAPProxy soapProxy -> {
28+
type = "soapProxy";
29+
description.put("wsdl", soapProxy.getWsdl());
30+
description.put("serviceName", soapProxy.getServiceName());
31+
}
32+
case ServiceProxy ignored -> type = "serviceProxy";
33+
default -> type = "unknown";
34+
}
35+
36+
description.put("type", type);
37+
description.put("rule", proxy.getKey().toString());
38+
description.put("interceptors", proxy.getFlow().stream()
39+
.map(interceptor -> Map.of("name", interceptor.getDisplayName()))
40+
.toList());
41+
description.put("statistics", statistics);
42+
return description;
43+
}
44+
45+
public static @Nullable Map<String, Object> describeExchange(AbstractExchange exchange, boolean includeBodies, McpPayloadSanitizer payloadSanitizer) {
46+
if (exchange.getResponse() == null) {
47+
return null;
48+
}
49+
50+
var description = new LinkedHashMap<String, Object>();
51+
description.put("id", exchange.getId());
52+
53+
var request = new LinkedHashMap<String, Object>();
54+
request.put("method", exchange.getRequest().getMethod());
55+
request.put("path", exchange.getRequest().getUri());
56+
request.put("headers", payloadSanitizer.sanitizeHeaders(exchange.getRequest().getHeader()));
57+
if (includeBodies) {
58+
request.put("body", payloadSanitizer.sanitizeBody(exchange.getRequest()));
59+
}
60+
61+
var response = new LinkedHashMap<String, Object>();
62+
response.put("status", exchange.getResponse().getStatusCode());
63+
response.put("headers", payloadSanitizer.sanitizeHeaders(exchange.getResponse().getHeader()));
64+
if (includeBodies) {
65+
response.put("body", payloadSanitizer.sanitizeBody(exchange.getResponse()));
66+
}
67+
68+
description.put("request", request);
69+
description.put("response", response);
70+
return description;
71+
}
72+
73+
public static int getOptionalIntArgument(MCPToolsCall call, String name, int defaultValue, int minimum, int maximum) {
74+
Object value = call.getArgument(name);
75+
if (value == null) {
76+
return defaultValue;
77+
}
78+
if (!(value instanceof Number number) || number.doubleValue() != Math.rint(number.doubleValue())) {
79+
throw new InvalidToolArgumentsException("Tool argument '" + name + "' must be an integer");
80+
}
81+
int parsed = number.intValue();
82+
if (parsed < minimum || parsed > maximum) {
83+
throw new InvalidToolArgumentsException(
84+
"Tool argument '" + name + "' must be between " + minimum + " and " + maximum
85+
);
86+
}
87+
return parsed;
88+
}
89+
90+
public static boolean getOptionalBooleanArgument(MCPToolsCall call, String name, boolean defaultValue) {
91+
Object value = call.getArgument(name);
92+
if (value == null) {
93+
return defaultValue;
94+
}
95+
if (value instanceof Boolean bool) {
96+
return bool;
97+
}
98+
throw new InvalidToolArgumentsException("Tool argument '" + name + "' must be a boolean");
99+
}
100+
101+
public static void rejectUnexpectedArguments(MCPToolsCall call, Set<String> allowed) {
102+
for (String argumentName : call.getArguments().keySet()) {
103+
if (!allowed.contains(argumentName)) {
104+
throw new InvalidToolArgumentsException("Unexpected tool argument: " + argumentName);
105+
}
106+
}
107+
}
108+
109+
public static final class InvalidToolArgumentsException extends IllegalArgumentException {
110+
private InvalidToolArgumentsException(String message) {
111+
super(message);
112+
}
113+
}
114+
}

‎core/src/main/java/com/predic8/membrane/core/interceptor/mcp/McpPayloadSanitizer.java‎

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,73 @@
1+
package com.predic8.membrane.core.interceptor.mcp;
2+
3+
import com.predic8.membrane.core.http.Header;
4+
import com.predic8.membrane.core.http.HeaderField;
5+
import com.predic8.membrane.core.http.Message;
6+
import com.predic8.membrane.core.http.MimeType;
7+
8+
import java.util.LinkedHashMap;
9+
import java.util.Map;
10+
import java.util.Set;
11+
12+
import static com.predic8.membrane.core.http.Header.*;
13+
import static java.util.Locale.ROOT;
14+
15+
public final class McpPayloadSanitizer {
16+
17+
private static final Set<String> SENSITIVE_HEADERS = Set.of(
18+
AUTHORIZATION.toLowerCase(ROOT),
19+
COOKIE.toLowerCase(ROOT),
20+
SET_COOKIE.toLowerCase(ROOT),
21+
PROXY_AUTHORIZATION.toLowerCase(ROOT)
22+
);
23+
24+
private static final String REDACTED = "<redacted>";
25+
private static final String BINARY_BODY_OMITTED = "<binary body omitted>";
26+
private static final String BODY_UNAVAILABLE = "<body unavailable>";
27+
private static final int MAX_BODY_LENGTH = 8 * 1024;
28+
29+
public Map<String, Object> sanitizeHeaders(Header header) {
30+
var sanitized = new LinkedHashMap<String, Object>();
31+
if (header == null) {
32+
return sanitized;
33+
}
34+
35+
for (HeaderField field : header.getAllHeaderFields()) {
36+
String name = field.getHeaderName().toString();
37+
38+
sanitized.merge(name, redactIfSensitive(field, name), (previous, current) -> previous + ", " + current);
39+
}
40+
41+
return sanitized;
42+
}
43+
44+
private static String redactIfSensitive(HeaderField field, String name) {
45+
return SENSITIVE_HEADERS.contains(name.toLowerCase(ROOT)) ? REDACTED : field.getValue();
46+
}
47+
48+
public String sanitizeBody(Message message) {
49+
if (message == null) {
50+
return BODY_UNAVAILABLE;
51+
}
52+
53+
try {
54+
if (message.isBodyEmpty()) {
55+
return "";
56+
}
57+
58+
// TODO: keep this?
59+
String contentType = message.getHeader().getContentType();
60+
if (contentType != null && !(MimeType.isText(contentType) || MimeType.isJson(contentType) || MimeType.isXML(contentType))) {
61+
return BINARY_BODY_OMITTED;
62+
}
63+
64+
String body = message.getBodyAsStringDecoded();
65+
if (body.length() <= MAX_BODY_LENGTH) {
66+
return body;
67+
}
68+
return body.substring(0, MAX_BODY_LENGTH) + "... <truncated>";
69+
} catch (Exception e) {
70+
return BODY_UNAVAILABLE;
71+
}
72+
}
73+
}

‎core/src/main/java/com/predic8/membrane/core/interceptor/mcp/McpSessionContext.java‎

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
package com.predic8.membrane.core.interceptor.mcp;
2+
3+
import com.predic8.membrane.core.mcp.MCPInitialize.ClientInfo;
4+
5+
import static com.predic8.membrane.core.interceptor.mcp.McpSessionContext.McpSessionState.*;
6+
7+
public final class McpSessionContext {
8+
9+
private McpSessionState state = NEW;
10+
11+
// TODO: only if we want to support multiple versions (maybe in the future?)
12+
private String negotiatedProtocolVersion;
13+
private ClientInfo clientInfo;
14+
15+
public synchronized McpSessionState getState() {
16+
return state;
17+
}
18+
19+
public synchronized boolean initialize(String protocolVersion, ClientInfo clientInfo) {
20+
if (state != NEW) {
21+
return false;
22+
}
23+
negotiatedProtocolVersion = protocolVersion;
24+
this.clientInfo = clientInfo;
25+
state = INITIALIZED;
26+
return true;
27+
}
28+
29+
public synchronized boolean markReady() {
30+
if (state != INITIALIZED) {
31+
return false;
32+
}
33+
state = READY;
34+
return true;
35+
}
36+
37+
@SuppressWarnings("BooleanMethodIsAlwaysInverted")
38+
public synchronized boolean isIn(McpSessionState... states) {
39+
for (McpSessionState candidate : states) {
40+
if (state == candidate) {
41+
return true;
42+
}
43+
}
44+
return false;
45+
}
46+
47+
public synchronized String getNegotiatedProtocolVersion() {
48+
return negotiatedProtocolVersion;
49+
}
50+
51+
public synchronized ClientInfo getClientInfo() {
52+
return clientInfo;
53+
}
54+
55+
public enum McpSessionState {
56+
NEW,
57+
INITIALIZED,
58+
READY,
59+
CLOSED
60+
}
61+
62+
}

‎core/src/main/java/com/predic8/membrane/core/interceptor/mcp/McpToolDefinition.java‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
package com.predic8.membrane.core.interceptor.mcp;
2+
3+
import com.predic8.membrane.core.mcp.MCPToolsListResponse;
4+
5+
import java.util.Map;
6+
7+
public record McpToolDefinition(
8+
String name,
9+
String description,
10+
Map<String, Object> inputSchema,
11+
McpToolHandler handler
12+
) {
13+
public MCPToolsListResponse.Tool toTool() {
14+
return new MCPToolsListResponse.Tool(name, description, inputSchema);
15+
}
16+
}

‎core/src/main/java/com/predic8/membrane/core/interceptor/mcp/McpToolHandler.java‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
package com.predic8.membrane.core.interceptor.mcp;
2+
3+
import com.predic8.membrane.core.exchange.Exchange;
4+
import com.predic8.membrane.core.mcp.MCPToolsCall;
5+
import com.predic8.membrane.core.mcp.MCPToolsCallResponse;
6+
7+
@FunctionalInterface
8+
public interface McpToolHandler {
9+
10+
MCPToolsCallResponse handle(MCPToolsCall call, Exchange exc) throws Exception;
11+
12+
}

‎core/src/main/java/com/predic8/membrane/core/interceptor/mcp/McpToolRegistry.java‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
package com.predic8.membrane.core.interceptor.mcp;
2+
3+
import java.util.Collection;
4+
import java.util.LinkedHashMap;
5+
import java.util.Map;
6+
7+
public final class McpToolRegistry {
8+
9+
private final Map<String, McpToolDefinition> tools = new LinkedHashMap<>();
10+
11+
public McpToolRegistry register(McpToolDefinition definition) {
12+
if (tools.putIfAbsent(definition.name(), definition) != null) {
13+
throw new IllegalArgumentException("Duplicate MCP tool registration: " + definition.name());
14+
}
15+
return this;
16+
}
17+
18+
public McpToolDefinition find(String name) {
19+
return tools.get(name);
20+
}
21+
22+
public Collection<McpToolDefinition> list() {
23+
return tools.values();
24+
}
25+
}

0 commit comments

Comments
 (0)