Skip to content

Commit 6678c46

Browse files
habiblawal1habiblawal1
authored
McpRequestId type to store incoming request ids (#32794)
* feat: create a McpRequestId class to store incoming request ids * fix: correct the way MessageParsingTest parses requests * chore: make RequestId make our new class * docs: add javadoc to McpRequestId class * chore: convert RequestId class to ExecutionRequestId * test: update MessageParsingTest with McpRequestId deserialiser * feat: add serializer for McpRequestId * chore: handle null request ids better * test: add Cancellation test for numeric ids * test: use map to specify countdown latch to use for cancellation test --------- Co-authored-by: Habib Lawal <[email protected]>
1 parent 35ce3c8 commit 6678c46

File tree

20 files changed

+563
-99
lines changed

20 files changed

+563
-99
lines changed

‎dev/io.openliberty.mcp.internal/src/io/openliberty/mcp/internal/McpConnectionTracker.java‎

Lines changed: 13 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -9,45 +9,46 @@
99
*******************************************************************************/
1010
package io.openliberty.mcp.internal;
1111

12+
import static io.openliberty.mcp.internal.exceptions.jsonrpc.JSONRPCErrorCode.INVALID_PARAMS;
13+
1214
import java.util.concurrent.ConcurrentHashMap;
1315
import java.util.concurrent.ConcurrentMap;
1416

1517
import io.openliberty.mcp.internal.exceptions.jsonrpc.JSONRPCException;
16-
import io.openliberty.mcp.internal.requests.RequestId;
18+
import io.openliberty.mcp.internal.requests.ExecutionRequestId;
19+
import io.openliberty.mcp.internal.requests.McpRequestId;
1720
import io.openliberty.mcp.messaging.Cancellation;
1821
import jakarta.enterprise.context.ApplicationScoped;
1922

20-
import static io.openliberty.mcp.internal.exceptions.jsonrpc.JSONRPCErrorCode.INVALID_PARAMS;
21-
2223
/**
2324
* This is a connection tracker bean. It keeps track of ongoing tool call requests
2425
*/
2526

2627
@ApplicationScoped
2728
public class McpConnectionTracker {
2829

29-
private final ConcurrentMap<String, Cancellation> ongoingRequests;
30+
private final ConcurrentMap<McpRequestId, Cancellation> ongoingRequests;
3031

3132
public McpConnectionTracker() {
3233
this.ongoingRequests = new ConcurrentHashMap<>();
3334
}
3435

35-
public void deregisterOngoingRequest(RequestId id) {
36-
ongoingRequests.remove(id.getUniqueId());
36+
public void deregisterOngoingRequest(ExecutionRequestId id) {
37+
ongoingRequests.remove(id.id());
3738
}
3839

39-
public void registerOngoingRequest(RequestId id, Cancellation cancellation) {
40-
Cancellation previous = ongoingRequests.putIfAbsent(id.getUniqueId(), cancellation);
40+
public void registerOngoingRequest(ExecutionRequestId id, Cancellation cancellation) {
41+
Cancellation previous = ongoingRequests.putIfAbsent(id.id(), cancellation);
4142
if (previous != null) {
4243
throw new JSONRPCException(INVALID_PARAMS, "A request with id " + id.id() + " is already in progress");
4344
}
4445
}
4546

46-
public boolean isOngoingRequest(RequestId id) {
47-
return ongoingRequests.containsKey(id.getUniqueId());
47+
public boolean isOngoingRequest(ExecutionRequestId id) {
48+
return ongoingRequests.containsKey(id.id());
4849
}
4950

50-
public Cancellation getOngoingRequestCancellation(RequestId id) {
51-
return ongoingRequests.get(id.getUniqueId());
51+
public Cancellation getOngoingRequestCancellation(ExecutionRequestId id) {
52+
return ongoingRequests.get(id.id());
5253
}
5354
}

‎dev/io.openliberty.mcp.internal/src/io/openliberty/mcp/internal/McpServlet.java‎

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -28,10 +28,11 @@
2828
import io.openliberty.mcp.internal.exceptions.jsonrpc.JSONRPCErrorCode;
2929
import io.openliberty.mcp.internal.exceptions.jsonrpc.JSONRPCException;
3030
import io.openliberty.mcp.internal.requests.CancellationImpl;
31+
import io.openliberty.mcp.internal.requests.ExecutionRequestId;
3132
import io.openliberty.mcp.internal.requests.McpInitializeParams;
3233
import io.openliberty.mcp.internal.requests.McpNotificationParams;
34+
import io.openliberty.mcp.internal.requests.McpRequestId;
3335
import io.openliberty.mcp.internal.requests.McpToolCallParams;
34-
import io.openliberty.mcp.internal.requests.RequestId;
3536
import io.openliberty.mcp.internal.responses.McpInitializeResult;
3637
import io.openliberty.mcp.internal.responses.McpInitializeResult.ServerInfo;
3738
import io.openliberty.mcp.messaging.Cancellation;
@@ -122,7 +123,7 @@ protected void callRequest(McpTransport transport)
122123

123124
@FFDCIgnore({ JSONRPCException.class, InvocationTargetException.class, IllegalAccessException.class, IllegalArgumentException.class })
124125
private void callTool(McpTransport transport) {
125-
RequestId requestId = createOngoingRequestId(transport);
126+
ExecutionRequestId requestId = createOngoingRequestId(transport);
126127
McpToolCallParams params = transport.getParams(McpToolCallParams.class);
127128

128129
if (params.getMetadata() == null) {
@@ -223,7 +224,7 @@ private ToolResponse toErrorResponse(Throwable t) {
223224
* @param requestId the ongoing request Id
224225
* @param toolMetadata the tool metadata
225226
*/
226-
private void addSpecialArguments(Object[] argumentsArray, RequestId requestId, ToolMetadata toolMetadata) {
227+
private void addSpecialArguments(Object[] argumentsArray, ExecutionRequestId requestId, ToolMetadata toolMetadata) {
227228
for (SpecialArgumentMetadata argMetadata : toolMetadata.specialArguments()) {
228229
switch (argMetadata.typeResolution().specialArgsType()) {
229230
case CANCELLATION -> {
@@ -303,7 +304,8 @@ private void ping(McpTransport transport) {
303304

304305
private void cancelRequest(McpTransport transport) {
305306
McpNotificationParams notificationParams = transport.getMcpRequest().getParams(McpNotificationParams.class, jsonb);
306-
RequestId requestId = new RequestId(notificationParams.getRequestId(), transport.getRequestIpAddress());
307+
McpRequestId mcpRedId = notificationParams.getRequestId();
308+
ExecutionRequestId requestId = new ExecutionRequestId(mcpRedId, transport.getRequestIpAddress());
307309
Optional<String> reason = Optional.ofNullable(notificationParams.getReason());
308310

309311
if (TraceComponent.isAnyTracingEnabled() && tc.isEventEnabled()) {
@@ -320,8 +322,8 @@ private void cancelRequest(McpTransport transport) {
320322
transport.sendEmptyResponse();
321323
}
322324

323-
private RequestId createOngoingRequestId(McpTransport transport) {
324-
return new RequestId(transport.getMcpRequest().id().toString(),
325-
transport.getRequestIpAddress());
325+
private ExecutionRequestId createOngoingRequestId(McpTransport transport) {
326+
return new ExecutionRequestId(transport.getMcpRequest().id(),
327+
transport.getRequestIpAddress());
326328
}
327329
}

‎dev/io.openliberty.mcp.internal/src/io/openliberty/mcp/internal/McpTransport.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
import io.openliberty.mcp.internal.exceptions.jsonrpc.JSONRPCErrorCode;
2626
import io.openliberty.mcp.internal.exceptions.jsonrpc.JSONRPCException;
2727
import io.openliberty.mcp.internal.requests.McpRequest;
28+
import io.openliberty.mcp.internal.requests.McpRequestId;
2829
import io.openliberty.mcp.internal.responses.McpErrorResponse;
2930
import io.openliberty.mcp.internal.responses.McpResponse;
3031
import io.openliberty.mcp.internal.responses.McpResultResponse;
@@ -191,7 +192,7 @@ public void sendError(Exception e) throws IOException {
191192
* @param e The JSONRPCException to be included in the error response.
192193
*/
193194
public void sendJsonRpcException(JSONRPCException e) {
194-
McpResponse mcpResponse = new McpErrorResponse(mcpRequest == null ? "" : mcpRequest.id(), e);
195+
McpResponse mcpResponse = new McpErrorResponse(mcpRequest == null ? new McpRequestId("") : mcpRequest.id(), e);
195196
jsonb.toJson(mcpResponse, writer);
196197
}
197198

‎dev/io.openliberty.mcp.internal/src/io/openliberty/mcp/internal/requests/CancellationImpl.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919
*/
2020
public class CancellationImpl implements Cancellation {
2121

22-
private RequestId requestId;
22+
private ExecutionRequestId requestId;
2323
private volatile Optional<String> reason = null;
2424

2525
/**
@@ -36,7 +36,7 @@ public Result check() {
3636
return new Result(true, reason);
3737
}
3838

39-
public void setRequestId(RequestId requestId) {
39+
public void setRequestId(ExecutionRequestId requestId) {
4040
this.requestId = requestId;
4141
}
4242

dev/io.openliberty.mcp.internal/src/io/openliberty/mcp/internal/requests/RequestId.java renamed to dev/io.openliberty.mcp.internal/src/io/openliberty/mcp/internal/requests/ExecutionRequestId.java

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -9,8 +9,4 @@
99
*******************************************************************************/
1010
package io.openliberty.mcp.internal.requests;
1111

12-
public record RequestId(String id, String sourceIp) {
13-
public String getUniqueId() {
14-
return id + sourceIp;
15-
}
16-
}
12+
public record ExecutionRequestId(McpRequestId id, String sourceIp) {}

‎dev/io.openliberty.mcp.internal/src/io/openliberty/mcp/internal/requests/McpNotificationParams.java‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15,14 +15,14 @@
1515
public class McpNotificationParams {
1616

1717
//Cancelled Notification params
18-
private String requestId;
18+
private McpRequestId requestId;
1919
private String reason;
2020

21-
public String getRequestId() {
21+
public McpRequestId getRequestId() {
2222
return requestId;
2323
}
2424

25-
public void setRequestId(String requestId) {
25+
public void setRequestId(McpRequestId requestId) {
2626
this.requestId = requestId;
2727
}
2828

‎dev/io.openliberty.mcp.internal/src/io/openliberty/mcp/internal/requests/McpRequest.java‎

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@
2424
import jakarta.json.bind.Jsonb;
2525

2626
public record McpRequest(String jsonrpc,
27-
Object id,
27+
McpRequestId id,
2828
String method,
2929
JsonObject params) {
3030

@@ -72,7 +72,7 @@ public static McpRequest createValidMCPRequest(Reader reader) throws JsonExcepti
7272
return createMCPNotificationRequest(jsonRpc, method, params);
7373
}
7474

75-
Object idObj = parseAndValidateId(id, errors);
75+
McpRequestId idObj = parseAndValidateId(id, errors);
7676

7777
if (!errors.isEmpty()) {
7878
throw new MCPRequestValidationException(errors);
@@ -98,17 +98,16 @@ private static void validateMethod(String method, List<String> errors) {
9898
}
9999
}
100100

101-
private static Object parseAndValidateId(JsonValue id, List<String> errors) {
102-
101+
private static McpRequestId parseAndValidateId(JsonValue id, List<String> errors) {
103102
return switch (id.getValueType()) {
104-
case NUMBER -> ((JsonNumber) id).numberValue();
103+
case NUMBER -> new McpRequestId(((JsonNumber) id).bigDecimalValue());
105104
case STRING -> {
106105
String idString = ((JsonString) id).getString();
107106
if (idString.isBlank()) {
108107
errors.add("id must not be empty");
109108
yield null;
110109
}
111-
yield idString;
110+
yield new McpRequestId(idString);
112111
}
113112
default -> {
114113
errors.add("id must be a string or number");

‎dev/io.openliberty.mcp.internal/src/io/openliberty/mcp/internal/requests/McpRequestId.java‎

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,87 @@
1+
/*******************************************************************************
2+
* Copyright (c) 2025 IBM Corporation and others.
3+
* All rights reserved. This program and the accompanying materials
4+
* are made available under the terms of the Eclipse Public License 2.0
5+
* which accompanies this distribution, and is available at
6+
* http://www.eclipse.org/legal/epl-2.0/
7+
*
8+
* SPDX-License-Identifier: EPL-2.0
9+
*******************************************************************************/
10+
package io.openliberty.mcp.internal.requests;
11+
12+
import java.math.BigDecimal;
13+
import java.util.Objects;
14+
15+
import jakarta.json.bind.annotation.JsonbTypeDeserializer;
16+
import jakarta.json.bind.annotation.JsonbTypeSerializer;
17+
18+
/**
19+
* Stores the id of an MCP Request, which can be represented as a String or Number.
20+
*/
21+
@JsonbTypeSerializer(McpRequestIdSerializer.class)
22+
@JsonbTypeDeserializer(McpRequestIdDeserializer.class)
23+
public final class McpRequestId {
24+
25+
private final String strVal;
26+
private final BigDecimal numVal;
27+
28+
public McpRequestId(String value) {
29+
this.strVal = value;
30+
this.numVal = null;
31+
}
32+
33+
public McpRequestId(BigDecimal value) {
34+
this.numVal = value;
35+
this.strVal = null;
36+
}
37+
38+
/**
39+
* @return the strVal
40+
*/
41+
public String getStrVal() {
42+
return strVal;
43+
}
44+
45+
/**
46+
* @return the numVal
47+
*/
48+
public BigDecimal getNumVal() {
49+
return numVal;
50+
}
51+
52+
/**
53+
* Overrides the equals method to compare if two MCP Request IDs are equal
54+
*
55+
* @param obj The McpRequestId object to compare.
56+
* @return True if the MCP Request IDs are equal, false otherwise.
57+
*/
58+
@Override
59+
public boolean equals(Object obj) {
60+
if (this == obj)
61+
return true;
62+
if (obj == null)
63+
return false;
64+
if (getClass() != obj.getClass())
65+
return false;
66+
McpRequestId other = (McpRequestId) obj;
67+
return Objects.equals(numVal, other.numVal) && Objects.equals(strVal, other.strVal);
68+
}
69+
70+
/**
71+
* Overrides the hashCode method to generate a hash code based on the stored id value.
72+
*
73+
* @return The hash code for this McpRequestId object based on the stored id value.
74+
*/
75+
@Override
76+
public int hashCode() {
77+
return Objects.hash(numVal, strVal);
78+
}
79+
80+
@Override
81+
public String toString() {
82+
if (getStrVal() != null)
83+
return getStrVal();
84+
return getNumVal().toString();
85+
}
86+
87+
}

‎dev/io.openliberty.mcp.internal/src/io/openliberty/mcp/internal/requests/McpRequestIdDeserializer.java‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
/*******************************************************************************
2+
* Copyright (c) 2025 IBM Corporation and others.
3+
* All rights reserved. This program and the accompanying materials
4+
* are made available under the terms of the Eclipse Public License 2.0
5+
* which accompanies this distribution, and is available at
6+
* http://www.eclipse.org/legal/epl-2.0/
7+
*
8+
* SPDX-License-Identifier: EPL-2.0
9+
*******************************************************************************/
10+
package io.openliberty.mcp.internal.requests;
11+
12+
import java.lang.reflect.Type;
13+
14+
import io.openliberty.mcp.internal.exceptions.jsonrpc.JSONRPCErrorCode;
15+
import io.openliberty.mcp.internal.exceptions.jsonrpc.JSONRPCException;
16+
import jakarta.json.JsonNumber;
17+
import jakarta.json.JsonString;
18+
import jakarta.json.JsonValue;
19+
import jakarta.json.bind.serializer.DeserializationContext;
20+
import jakarta.json.bind.serializer.JsonbDeserializer;
21+
import jakarta.json.stream.JsonParser;
22+
23+
/**
24+
* Instructions for how Jsonb should deserialize JSON values into a {@link McpRequestId} type
25+
*/
26+
public class McpRequestIdDeserializer implements JsonbDeserializer<McpRequestId> {
27+
28+
@Override
29+
public McpRequestId deserialize(JsonParser parser, DeserializationContext ctx, Type rtType) {
30+
JsonValue jsonVal = parser.getValue();
31+
switch (jsonVal.getValueType()) {
32+
case STRING:
33+
String strVal = ((JsonString) jsonVal).getString();
34+
if (strVal.isBlank())
35+
throw new JSONRPCException(JSONRPCErrorCode.PARSE_ERROR, "MCP Request Id must not be blank");
36+
return new McpRequestId(strVal);
37+
case NUMBER:
38+
return new McpRequestId(((JsonNumber) jsonVal).bigDecimalValue());
39+
default:
40+
throw new JSONRPCException(JSONRPCErrorCode.PARSE_ERROR, "MCP Request Id must be a string or number");
41+
42+
}
43+
}
44+
45+
}

‎dev/io.openliberty.mcp.internal/src/io/openliberty/mcp/internal/requests/McpRequestIdSerializer.java‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
/*******************************************************************************
2+
* Copyright (c) 2025 IBM Corporation and others.
3+
* All rights reserved. This program and the accompanying materials
4+
* are made available under the terms of the Eclipse Public License 2.0
5+
* which accompanies this distribution, and is available at
6+
* http://www.eclipse.org/legal/epl-2.0/
7+
*
8+
* SPDX-License-Identifier: EPL-2.0
9+
*******************************************************************************/
10+
package io.openliberty.mcp.internal.requests;
11+
12+
import jakarta.json.bind.serializer.JsonbSerializer;
13+
import jakarta.json.bind.serializer.SerializationContext;
14+
import jakarta.json.stream.JsonGenerator;
15+
16+
/**
17+
* Instructions for how Jsonb should serialize {@link McpRequestId} types into JSON
18+
*/
19+
public class McpRequestIdSerializer implements JsonbSerializer<McpRequestId> {
20+
21+
@Override
22+
public void serialize(McpRequestId id, JsonGenerator generator, SerializationContext ctx) {
23+
if (id.getStrVal() != null && !id.getStrVal().isEmpty())
24+
generator.write(id.getStrVal());
25+
else if (id.getNumVal() != null)
26+
generator.write(id.getNumVal());
27+
else
28+
generator.writeNull();
29+
}
30+
31+
}

0 commit comments

Comments
 (0)