Skip to content

Commit 476f9db

Browse files
tzolovtzolov
authored
feat: enhance error handling with custom error code preservation (#653)
- Improve McpClientSession error handling to preserve custom error codes and data from McpError instances. - Add aggregated exception messages to error data field for better debugging. - Include test coverage for various McpClientSession error scenarios. Signed-off-by: Christian Tzolov <christian.tzolov@broadcom.com>
1 parent 6e9af40 commit 476f9db

File tree

2 files changed

+160
-39
lines changed

2 files changed

+160
-39
lines changed

‎mcp-core/src/main/java/io/modelcontextprotocol/spec/McpClientSession.java‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -166,9 +166,15 @@ private void handle(McpSchema.JSONRPCMessage message) {
166166
else if (message instanceof McpSchema.JSONRPCRequest request) {
167167
logger.debug("Received request: {}", request);
168168
handleIncomingRequest(request).onErrorResume(error -> {
169+
170+
McpSchema.JSONRPCResponse.JSONRPCError jsonRpcError = (error instanceof McpError mcpError
171+
&& mcpError.getJsonRpcError() != null) ? mcpError.getJsonRpcError()
172+
// TODO: add error message through the data field
173+
: new McpSchema.JSONRPCResponse.JSONRPCError(McpSchema.ErrorCodes.INTERNAL_ERROR,
174+
error.getMessage(), McpError.aggregateExceptionMessages(error));
175+
169176
var errorResponse = new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), null,
170-
new McpSchema.JSONRPCResponse.JSONRPCError(McpSchema.ErrorCodes.INTERNAL_ERROR,
171-
error.getMessage(), null));
177+
jsonRpcError);
172178
return Mono.just(errorResponse);
173179
}).flatMap(this.transport::sendMessage).onErrorComplete(t -> {
174180
logger.warn("Issue sending response to the client, ", t);

‎mcp-core/src/test/java/io/modelcontextprotocol/spec/McpClientSessionTests.java‎

Lines changed: 152 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -6,11 +6,10 @@
66

77
import java.time.Duration;
88
import java.util.Map;
9+
import java.util.function.Function;
910

1011
import io.modelcontextprotocol.MockMcpClientTransport;
1112
import io.modelcontextprotocol.json.TypeRef;
12-
import org.junit.jupiter.api.AfterEach;
13-
import org.junit.jupiter.api.BeforeEach;
1413
import org.junit.jupiter.api.Test;
1514
import org.slf4j.Logger;
1615
import org.slf4j.LoggerFactory;
@@ -19,7 +18,6 @@
1918
import reactor.test.StepVerifier;
2019

2120
import static org.assertj.core.api.Assertions.assertThat;
22-
import static org.assertj.core.api.Assertions.assertThatThrownBy;
2321

2422
/**
2523
* Test suite for {@link McpClientSession} that verifies its JSON-RPC message handling,
@@ -39,35 +37,6 @@ class McpClientSessionTests {
3937

4038
private static final String ECHO_METHOD = "echo";
4139

42-
private McpClientSession session;
43-
44-
private MockMcpClientTransport transport;
45-
46-
@BeforeEach
47-
void setUp() {
48-
transport = new MockMcpClientTransport();
49-
session = new McpClientSession(TIMEOUT, transport, Map.of(),
50-
Map.of(TEST_NOTIFICATION, params -> Mono.fromRunnable(() -> logger.info("Status update: {}", params))));
51-
}
52-
53-
@AfterEach
54-
void tearDown() {
55-
if (session != null) {
56-
session.close();
57-
}
58-
}
59-
60-
@Test
61-
void testConstructorWithInvalidArguments() {
62-
assertThatThrownBy(() -> new McpClientSession(null, transport, Map.of(), Map.of()))
63-
.isInstanceOf(IllegalArgumentException.class)
64-
.hasMessageContaining("The requestTimeout can not be null");
65-
66-
assertThatThrownBy(() -> new McpClientSession(TIMEOUT, null, Map.of(), Map.of()))
67-
.isInstanceOf(IllegalArgumentException.class)
68-
.hasMessageContaining("transport can not be null");
69-
}
70-
7140
TypeRef<String> responseType = new TypeRef<>() {
7241
};
7342

@@ -76,6 +45,11 @@ void testSendRequest() {
7645
String testParam = "test parameter";
7746
String responseData = "test response";
7847

48+
var transport = new MockMcpClientTransport();
49+
var session = new McpClientSession(TIMEOUT, transport, Map.of(),
50+
Map.of(TEST_NOTIFICATION, params -> Mono.fromRunnable(() -> logger.info("Status update: {}", params))),
51+
Function.identity());
52+
7953
// Create a Mono that will emit the response after the request is sent
8054
Mono<String> responseMono = session.sendRequest(TEST_METHOD, testParam, responseType);
8155
// Verify response handling
@@ -92,10 +66,17 @@ void testSendRequest() {
9266
assertThat(request.params()).isEqualTo(testParam);
9367
assertThat(response).isEqualTo(responseData);
9468
}).verifyComplete();
69+
70+
session.close();
9571
}
9672

9773
@Test
9874
void testSendRequestWithError() {
75+
var transport = new MockMcpClientTransport();
76+
var session = new McpClientSession(TIMEOUT, transport, Map.of(),
77+
Map.of(TEST_NOTIFICATION, params -> Mono.fromRunnable(() -> logger.info("Status update: {}", params))),
78+
Function.identity());
79+
9980
Mono<String> responseMono = session.sendRequest(TEST_METHOD, "test", responseType);
10081

10182
// Verify error handling
@@ -107,20 +88,34 @@ void testSendRequestWithError() {
10788
transport.simulateIncomingMessage(
10889
new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), null, error));
10990
}).expectError(McpError.class).verify();
91+
92+
session.close();
11093
}
11194

11295
@Test
11396
void testRequestTimeout() {
97+
var transport = new MockMcpClientTransport();
98+
var session = new McpClientSession(TIMEOUT, transport, Map.of(),
99+
Map.of(TEST_NOTIFICATION, params -> Mono.fromRunnable(() -> logger.info("Status update: {}", params))),
100+
Function.identity());
101+
114102
Mono<String> responseMono = session.sendRequest(TEST_METHOD, "test", responseType);
115103

116104
// Verify timeout
117105
StepVerifier.create(responseMono)
118106
.expectError(java.util.concurrent.TimeoutException.class)
119107
.verify(TIMEOUT.plusSeconds(1));
108+
109+
session.close();
120110
}
121111

122112
@Test
123113
void testSendNotification() {
114+
var transport = new MockMcpClientTransport();
115+
var session = new McpClientSession(TIMEOUT, transport, Map.of(),
116+
Map.of(TEST_NOTIFICATION, params -> Mono.fromRunnable(() -> logger.info("Status update: {}", params))),
117+
Function.identity());
118+
124119
Map<String, Object> params = Map.of("key", "value");
125120
Mono<Void> notificationMono = session.sendNotification(TEST_NOTIFICATION, params);
126121

@@ -132,15 +127,17 @@ void testSendNotification() {
132127
assertThat(notification.method()).isEqualTo(TEST_NOTIFICATION);
133128
assertThat(notification.params()).isEqualTo(params);
134129
}).verifyComplete();
130+
131+
session.close();
135132
}
136133

137134
@Test
138135
void testRequestHandling() {
139136
String echoMessage = "Hello MCP!";
140137
Map<String, McpClientSession.RequestHandler<?>> requestHandlers = Map.of(ECHO_METHOD,
141138
params -> Mono.just(params));
142-
transport = new MockMcpClientTransport();
143-
session = new McpClientSession(TIMEOUT, transport, requestHandlers, Map.of());
139+
var transport = new MockMcpClientTransport();
140+
var session = new McpClientSession(TIMEOUT, transport, requestHandlers, Map.of(), Function.identity());
144141

145142
// Simulate incoming request
146143
McpSchema.JSONRPCRequest request = new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, ECHO_METHOD,
@@ -153,15 +150,18 @@ void testRequestHandling() {
153150
McpSchema.JSONRPCResponse response = (McpSchema.JSONRPCResponse) sentMessage;
154151
assertThat(response.result()).isEqualTo(echoMessage);
155152
assertThat(response.error()).isNull();
153+
154+
session.close();
156155
}
157156

158157
@Test
159158
void testNotificationHandling() {
160159
Sinks.One<Object> receivedParams = Sinks.one();
161160

162-
transport = new MockMcpClientTransport();
163-
session = new McpClientSession(TIMEOUT, transport, Map.of(),
164-
Map.of(TEST_NOTIFICATION, params -> Mono.fromRunnable(() -> receivedParams.tryEmitValue(params))));
161+
var transport = new MockMcpClientTransport();
162+
var session = new McpClientSession(TIMEOUT, transport, Map.of(),
163+
Map.of(TEST_NOTIFICATION, params -> Mono.fromRunnable(() -> receivedParams.tryEmitValue(params))),
164+
Function.identity());
165165

166166
// Simulate incoming notification from the server
167167
Map<String, Object> notificationParams = Map.of("status", "ready");
@@ -173,10 +173,18 @@ void testNotificationHandling() {
173173

174174
// Verify handler was called
175175
assertThat(receivedParams.asMono().block(Duration.ofSeconds(1))).isEqualTo(notificationParams);
176+
177+
session.close();
176178
}
177179

178180
@Test
179181
void testUnknownMethodHandling() {
182+
183+
var transport = new MockMcpClientTransport();
184+
var session = new McpClientSession(TIMEOUT, transport, Map.of(),
185+
Map.of(TEST_NOTIFICATION, params -> Mono.fromRunnable(() -> logger.info("Status update: {}", params))),
186+
Function.identity());
187+
180188
// Simulate incoming request for unknown method
181189
McpSchema.JSONRPCRequest request = new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, "unknown.method",
182190
"test-id", null);
@@ -188,10 +196,117 @@ void testUnknownMethodHandling() {
188196
McpSchema.JSONRPCResponse response = (McpSchema.JSONRPCResponse) sentMessage;
189197
assertThat(response.error()).isNotNull();
190198
assertThat(response.error().code()).isEqualTo(McpSchema.ErrorCodes.METHOD_NOT_FOUND);
199+
200+
session.close();
201+
}
202+
203+
@Test
204+
void testRequestHandlerThrowsMcpErrorWithJsonRpcError() {
205+
// Setup: Create a request handler that throws McpError with custom error code and
206+
// data
207+
String testMethod = "test.customError";
208+
Map<String, Object> errorData = Map.of("customField", "customValue");
209+
McpClientSession.RequestHandler<?> failingHandler = params -> Mono
210+
.error(McpError.builder(123).message("Custom error message").data(errorData).build());
211+
212+
var transport = new MockMcpClientTransport();
213+
var session = new McpClientSession(TIMEOUT, transport, Map.of(testMethod, failingHandler), Map.of(),
214+
Function.identity());
215+
216+
// Simulate incoming request that will trigger the error
217+
McpSchema.JSONRPCRequest request = new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, testMethod,
218+
"test-id", null);
219+
transport.simulateIncomingMessage(request);
220+
221+
// Verify: The response should contain the custom error from McpError
222+
McpSchema.JSONRPCMessage sentMessage = transport.getLastSentMessage();
223+
assertThat(sentMessage).isInstanceOf(McpSchema.JSONRPCResponse.class);
224+
McpSchema.JSONRPCResponse response = (McpSchema.JSONRPCResponse) sentMessage;
225+
assertThat(response.error()).isNotNull();
226+
assertThat(response.error().code()).isEqualTo(123);
227+
assertThat(response.error().message()).isEqualTo("Custom error message");
228+
assertThat(response.error().data()).isEqualTo(errorData);
229+
230+
session.close();
231+
}
232+
233+
@Test
234+
void testRequestHandlerThrowsGenericException() {
235+
// Setup: Create a request handler that throws a generic RuntimeException
236+
String testMethod = "test.genericError";
237+
RuntimeException exception = new RuntimeException("Something went wrong");
238+
McpClientSession.RequestHandler<?> failingHandler = params -> Mono.error(exception);
239+
240+
var transport = new MockMcpClientTransport();
241+
var session = new McpClientSession(TIMEOUT, transport, Map.of(testMethod, failingHandler), Map.of(),
242+
Function.identity());
243+
244+
// Simulate incoming request that will trigger the error
245+
McpSchema.JSONRPCRequest request = new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, testMethod,
246+
"test-id", null);
247+
transport.simulateIncomingMessage(request);
248+
249+
// Verify: The response should contain INTERNAL_ERROR with aggregated exception
250+
// messages in data field
251+
McpSchema.JSONRPCMessage sentMessage = transport.getLastSentMessage();
252+
assertThat(sentMessage).isInstanceOf(McpSchema.JSONRPCResponse.class);
253+
McpSchema.JSONRPCResponse response = (McpSchema.JSONRPCResponse) sentMessage;
254+
assertThat(response.error()).isNotNull();
255+
assertThat(response.error().code()).isEqualTo(McpSchema.ErrorCodes.INTERNAL_ERROR);
256+
assertThat(response.error().message()).isEqualTo("Something went wrong");
257+
// Verify data field contains aggregated exception messages
258+
assertThat(response.error().data()).isNotNull();
259+
assertThat(response.error().data().toString()).contains("RuntimeException");
260+
assertThat(response.error().data().toString()).contains("Something went wrong");
261+
262+
session.close();
263+
}
264+
265+
@Test
266+
void testRequestHandlerThrowsExceptionWithCause() {
267+
// Setup: Create a request handler that throws an exception with a cause chain
268+
String testMethod = "test.chainedError";
269+
RuntimeException rootCause = new IllegalArgumentException("Root cause message");
270+
RuntimeException middleCause = new IllegalStateException("Middle cause message", rootCause);
271+
RuntimeException topException = new RuntimeException("Top level message", middleCause);
272+
McpClientSession.RequestHandler<?> failingHandler = params -> Mono.error(topException);
273+
274+
var transport = new MockMcpClientTransport();
275+
var session = new McpClientSession(TIMEOUT, transport, Map.of(testMethod, failingHandler), Map.of(),
276+
Function.identity());
277+
278+
// Simulate incoming request that will trigger the error
279+
McpSchema.JSONRPCRequest request = new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, testMethod,
280+
"test-id", null);
281+
transport.simulateIncomingMessage(request);
282+
283+
// Verify: The response should contain INTERNAL_ERROR with full exception chain
284+
// in data field
285+
McpSchema.JSONRPCMessage sentMessage = transport.getLastSentMessage();
286+
assertThat(sentMessage).isInstanceOf(McpSchema.JSONRPCResponse.class);
287+
McpSchema.JSONRPCResponse response = (McpSchema.JSONRPCResponse) sentMessage;
288+
assertThat(response.error()).isNotNull();
289+
assertThat(response.error().code()).isEqualTo(McpSchema.ErrorCodes.INTERNAL_ERROR);
290+
assertThat(response.error().message()).isEqualTo("Top level message");
291+
// Verify data field contains the full exception chain
292+
String dataString = response.error().data().toString();
293+
assertThat(dataString).contains("RuntimeException");
294+
assertThat(dataString).contains("Top level message");
295+
assertThat(dataString).contains("IllegalStateException");
296+
assertThat(dataString).contains("Middle cause message");
297+
assertThat(dataString).contains("IllegalArgumentException");
298+
assertThat(dataString).contains("Root cause message");
299+
300+
session.close();
191301
}
192302

193303
@Test
194304
void testGracefulShutdown() {
305+
var transport = new MockMcpClientTransport();
306+
var session = new McpClientSession(TIMEOUT, transport, Map.of(),
307+
Map.of(TEST_NOTIFICATION, params -> Mono.fromRunnable(() -> logger.info("Status update: {}", params))),
308+
Function.identity());
309+
195310
StepVerifier.create(session.closeGracefully()).verifyComplete();
196311
}
197312

0 commit comments

Comments
 (0)