Skip to content

Commit e0135a6

Browse files
tzolovtzolov
committed
improve HttpClientStreamableHttpTransportErrorHandlingTest
1 parent 675b13a commit e0135a6

File tree

2 files changed

+80
-74
lines changed

2 files changed

+80
-74
lines changed

‎mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/client/transport/WebClientStreamableHttpTransportErrorHandlingTest.java‎

Lines changed: 23 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -52,9 +52,9 @@ public class WebClientStreamableHttpTransportErrorHandlingTest {
5252

5353
private HttpServer server;
5454

55-
private AtomicReference<Integer> responseStatus = new AtomicReference<>(200);
55+
private AtomicReference<Integer> serverResponseStatus = new AtomicReference<>(200);
5656

57-
private AtomicReference<String> sessionId = new AtomicReference<>(null);
57+
private AtomicReference<String> currentServerSessionId = new AtomicReference<>(null);
5858

5959
private AtomicReference<String> lastReceivedSessionId = new AtomicReference<>(null);
6060

@@ -93,7 +93,7 @@ void startServer() throws IOException {
9393
String requestSessionId = exchange.getRequestHeaders().getFirst(HttpHeaders.MCP_SESSION_ID);
9494
lastReceivedSessionId.set(requestSessionId);
9595

96-
int status = responseStatus.get();
96+
int status = serverResponseStatus.get();
9797

9898
// Track which request this is
9999
if (firstRequestLatch.getCount() > 0) {
@@ -109,7 +109,7 @@ else if (secondRequestLatch.getCount() > 0) {
109109

110110
// Don't include session ID in 404 and 400 responses - the implementation
111111
// checks if the transport has a session stored locally
112-
String responseSessionId = sessionId.get();
112+
String responseSessionId = currentServerSessionId.get();
113113
if (responseSessionId != null && status == 200) {
114114
exchange.getResponseHeaders().set(HttpHeaders.MCP_SESSION_ID, responseSessionId);
115115
}
@@ -144,8 +144,8 @@ void stopServer() {
144144
*/
145145
@Test
146146
void test404WithoutSessionId() {
147-
responseStatus.set(404);
148-
sessionId.set(null); // No session ID in response
147+
serverResponseStatus.set(404);
148+
currentServerSessionId.set(null); // No session ID in response
149149

150150
var testMessage = createTestMessage();
151151

@@ -163,8 +163,8 @@ void test404WithoutSessionId() {
163163
@Test
164164
void test404WithSessionId() throws InterruptedException {
165165
// First establish a session
166-
responseStatus.set(200);
167-
sessionId.set("test-session-123");
166+
serverResponseStatus.set(200);
167+
currentServerSessionId.set("test-session-123");
168168

169169
// Set up exception handler to verify session invalidation
170170
@SuppressWarnings("unchecked")
@@ -187,7 +187,7 @@ void test404WithSessionId() throws InterruptedException {
187187
assertThat(getRequestLatch.await(5, TimeUnit.SECONDS)).isTrue();
188188

189189
// Now return 404 for next request
190-
responseStatus.set(404);
190+
serverResponseStatus.set(404);
191191

192192
// Use delaySubscription to ensure session is fully processed before next
193193
// request
@@ -212,8 +212,8 @@ void test404WithSessionId() throws InterruptedException {
212212
*/
213213
@Test
214214
void test400WithoutSessionId() {
215-
responseStatus.set(400);
216-
sessionId.set(null); // No session ID
215+
serverResponseStatus.set(400);
216+
currentServerSessionId.set(null); // No session ID
217217

218218
var testMessage = createTestMessage();
219219

@@ -232,8 +232,8 @@ void test400WithoutSessionId() {
232232
void test400WithSessionId() throws InterruptedException {
233233

234234
// First establish a session
235-
responseStatus.set(200);
236-
sessionId.set("test-session-456");
235+
serverResponseStatus.set(200);
236+
currentServerSessionId.set("test-session-456");
237237

238238
// Set up exception handler
239239
@SuppressWarnings("unchecked")
@@ -258,7 +258,7 @@ void test400WithSessionId() throws InterruptedException {
258258
assertThat(getCompleted).isTrue();
259259

260260
// Now return 400 for next request (simulating unknown session ID)
261-
responseStatus.set(400);
261+
serverResponseStatus.set(400);
262262

263263
// Use delaySubscription to ensure session is fully processed before next
264264
// request
@@ -284,24 +284,24 @@ void test400WithSessionId() throws InterruptedException {
284284
@Test
285285
void testSessionRecoveryAfter404() {
286286
// First establish a session
287-
responseStatus.set(200);
288-
sessionId.set("session-1");
287+
serverResponseStatus.set(200);
288+
currentServerSessionId.set("session-1");
289289

290290
// Send initial message to establish session
291291
var testMessage = createTestMessage();
292292

293293
// Use Mono.defer to ensure proper sequencing
294294
Mono<Void> establishSession = transport.sendMessage(testMessage).then(Mono.defer(() -> {
295295
// Simulate session loss - return 404
296-
responseStatus.set(404);
296+
serverResponseStatus.set(404);
297297
return transport.sendMessage(testMessage).onErrorResume(McpTransportSessionNotFoundException.class, e -> {
298298
// Expected error, continue with recovery
299299
return Mono.empty();
300300
});
301301
})).then(Mono.defer(() -> {
302302
// Now server is back with new session
303-
responseStatus.set(200);
304-
sessionId.set("session-2");
303+
serverResponseStatus.set(200);
304+
currentServerSessionId.set("session-2");
305305
lastReceivedSessionId.set(null); // Reset to verify new session
306306

307307
// Should be able to establish new session
@@ -336,7 +336,7 @@ void testReconnectErrorHandling() throws InterruptedException {
336336

337337
if ("GET".equals(method)) {
338338
sseConnectionLatch.countDown();
339-
int status = responseStatus.get();
339+
int status = serverResponseStatus.get();
340340

341341
if (status == 404 && requestSessionId != null) {
342342
// 404 with session ID - should trigger SessionNotFoundException
@@ -358,7 +358,7 @@ else if (status == 404) {
358358
else {
359359
// POST request handling
360360
exchange.getResponseHeaders().set("Content-Type", "application/json");
361-
String responseSessionId = sessionId.get();
361+
String responseSessionId = currentServerSessionId.get();
362362
if (responseSessionId != null) {
363363
exchange.getResponseHeaders().set(HttpHeaders.MCP_SESSION_ID, responseSessionId);
364364
}
@@ -370,8 +370,8 @@ else if (status == 404) {
370370
});
371371

372372
// Test with session ID - should get SessionNotFoundException
373-
responseStatus.set(200);
374-
sessionId.set("sse-session-1");
373+
serverResponseStatus.set(200);
374+
currentServerSessionId.set("sse-session-1");
375375

376376
var transport = WebClientStreamableHttpTransport.builder(WebClient.builder().baseUrl(HOST))
377377
.endpoint("/mcp-sse")

‎mcp/src/test/java/io/modelcontextprotocol/client/transport/HttpClientStreamableHttpTransportErrorHandlingTest.java‎

Lines changed: 57 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@
1616

1717
import org.junit.jupiter.api.AfterEach;
1818
import org.junit.jupiter.api.BeforeEach;
19-
import org.junit.jupiter.api.Disabled;
2019
import org.junit.jupiter.api.Test;
2120
import org.junit.jupiter.api.Timeout;
2221

@@ -47,9 +46,9 @@ public class HttpClientStreamableHttpTransportErrorHandlingTest {
4746

4847
private HttpServer server;
4948

50-
private AtomicReference<Integer> responseStatus = new AtomicReference<>(200);
49+
private AtomicReference<Integer> serverResponseStatus = new AtomicReference<>(200);
5150

52-
private AtomicReference<String> sessionId = new AtomicReference<>(null);
51+
private AtomicReference<String> currentServerSessionId = new AtomicReference<>(null);
5352

5453
private AtomicReference<String> lastReceivedSessionId = new AtomicReference<>(null);
5554

@@ -60,32 +59,37 @@ void startServer() throws IOException {
6059
server = HttpServer.create(new InetSocketAddress(PORT), 0);
6160

6261
// Configure the /mcp endpoint with dynamic response
63-
server.createContext("/mcp", exchange -> {
64-
// Capture session ID from request if present
65-
String requestSessionId = exchange.getRequestHeaders().getFirst(HttpHeaders.MCP_SESSION_ID);
66-
lastReceivedSessionId.set(requestSessionId);
62+
server.createContext("/mcp", httpExchange -> {
63+
if ("DELETE".equals(httpExchange.getRequestMethod())) {
64+
httpExchange.sendResponseHeaders(200, 0);
65+
}
66+
else {
67+
// Capture session ID from request if present
68+
String requestSessionId = httpExchange.getRequestHeaders().getFirst(HttpHeaders.MCP_SESSION_ID);
69+
lastReceivedSessionId.set(requestSessionId);
6770

68-
int status = responseStatus.get();
71+
int status = serverResponseStatus.get();
6972

70-
// Set response headers
71-
exchange.getResponseHeaders().set("Content-Type", "application/json");
73+
// Set response headers
74+
httpExchange.getResponseHeaders().set("Content-Type", "application/json");
7275

73-
// Add session ID to response if configured
74-
String responseSessionId = sessionId.get();
75-
if (responseSessionId != null) {
76-
exchange.getResponseHeaders().set(HttpHeaders.MCP_SESSION_ID, responseSessionId);
77-
}
76+
// Add session ID to response if configured
77+
String responseSessionId = currentServerSessionId.get();
78+
if (responseSessionId != null) {
79+
httpExchange.getResponseHeaders().set(HttpHeaders.MCP_SESSION_ID, responseSessionId);
80+
}
7881

79-
// Send response based on configured status
80-
if (status == 200) {
81-
String response = "{\"jsonrpc\":\"2.0\",\"result\":{},\"id\":\"test-id\"}";
82-
exchange.sendResponseHeaders(200, response.length());
83-
exchange.getResponseBody().write(response.getBytes());
84-
}
85-
else {
86-
exchange.sendResponseHeaders(status, 0);
82+
// Send response based on configured status
83+
if (status == 200) {
84+
String response = "{\"jsonrpc\":\"2.0\",\"result\":{},\"id\":\"test-id\"}";
85+
httpExchange.sendResponseHeaders(200, response.length());
86+
httpExchange.getResponseBody().write(response.getBytes());
87+
}
88+
else {
89+
httpExchange.sendResponseHeaders(status, 0);
90+
}
8791
}
88-
exchange.close();
92+
httpExchange.close();
8993
});
9094

9195
server.setExecutor(null);
@@ -107,10 +111,10 @@ void stopServer() {
107111
*/
108112
@Test
109113
void test404WithoutSessionId() {
110-
responseStatus.set(404);
111-
sessionId.set(null); // No session ID in response
114+
serverResponseStatus.set(404);
115+
currentServerSessionId.set(null); // No session ID in response
112116

113-
var testMessage = createTestMessage();
117+
var testMessage = createTestRequestMessage();
114118

115119
StepVerifier.create(transport.sendMessage(testMessage))
116120
.expectErrorMatches(throwable -> throwable instanceof McpTransportException
@@ -127,8 +131,8 @@ void test404WithoutSessionId() {
127131
@Test
128132
void test404WithSessionId() {
129133
// First establish a session
130-
responseStatus.set(200);
131-
sessionId.set("test-session-123");
134+
serverResponseStatus.set(200);
135+
currentServerSessionId.set("test-session-123");
132136

133137
// Set up exception handler to verify session invalidation
134138
@SuppressWarnings("unchecked")
@@ -139,12 +143,12 @@ void test404WithSessionId() {
139143
StepVerifier.create(transport.connect(msg -> msg)).verifyComplete();
140144

141145
// Send initial message to establish session
142-
var testMessage = createTestMessage();
146+
var testMessage = createTestRequestMessage();
143147
StepVerifier.create(transport.sendMessage(testMessage)).verifyComplete();
144148

145149
// The session should now be established, next request will include session ID
146150
// Now return 404 for next request
147-
responseStatus.set(404);
151+
serverResponseStatus.set(404);
148152

149153
// Send another message - should get SessionNotFoundException
150154
StepVerifier.create(transport.sendMessage(testMessage))
@@ -163,10 +167,10 @@ void test404WithSessionId() {
163167
*/
164168
@Test
165169
void test400WithoutSessionId() {
166-
responseStatus.set(400);
167-
sessionId.set(null); // No session ID
170+
serverResponseStatus.set(400);
171+
currentServerSessionId.set(null); // No session ID
168172

169-
var testMessage = createTestMessage();
173+
var testMessage = createTestRequestMessage();
170174

171175
StepVerifier.create(transport.sendMessage(testMessage))
172176
.expectErrorMatches(throwable -> throwable instanceof McpTransportException
@@ -185,8 +189,8 @@ void test400WithoutSessionId() {
185189
@Test
186190
void test400WithSessionId() {
187191
// First establish a session
188-
responseStatus.set(200);
189-
sessionId.set("test-session-456");
192+
serverResponseStatus.set(200);
193+
currentServerSessionId.set("test-session-456");
190194

191195
// Set up exception handler
192196
@SuppressWarnings("unchecked")
@@ -197,12 +201,12 @@ void test400WithSessionId() {
197201
StepVerifier.create(transport.connect(msg -> msg)).verifyComplete();
198202

199203
// Send initial message to establish session
200-
var testMessage = createTestMessage();
204+
var testMessage = createTestRequestMessage();
201205
StepVerifier.create(transport.sendMessage(testMessage)).verifyComplete();
202206

203207
// The session should now be established, next request will include session ID
204208
// Now return 400 for next request (simulating unknown session ID)
205-
responseStatus.set(400);
209+
serverResponseStatus.set(400);
206210

207211
// Send another message - should get SessionNotFoundException
208212
StepVerifier.create(transport.sendMessage(testMessage))
@@ -222,25 +226,27 @@ void test400WithSessionId() {
222226
@Test
223227
void testSessionRecoveryAfter404() {
224228
// First establish a session
225-
responseStatus.set(200);
226-
sessionId.set("session-1");
229+
serverResponseStatus.set(200);
230+
currentServerSessionId.set("session-1");
227231

228232
// Send initial message to establish session
229-
var testMessage = createTestMessage();
233+
var testMessage = createTestRequestMessage();
230234
StepVerifier.create(transport.sendMessage(testMessage)).verifyComplete();
231235

236+
assertThat(lastReceivedSessionId.get()).isNull();
237+
232238
// The session should now be established
233239
// Simulate session loss - return 404
234-
responseStatus.set(404);
240+
serverResponseStatus.set(404);
235241

236242
// This should fail with SessionNotFoundException
237243
StepVerifier.create(transport.sendMessage(testMessage))
238244
.expectError(McpTransportSessionNotFoundException.class)
239245
.verify();
240246

241247
// Now server is back with new session
242-
responseStatus.set(200);
243-
sessionId.set("session-2");
248+
serverResponseStatus.set(200);
249+
currentServerSessionId.set("session-2");
244250
lastReceivedSessionId.set(null); // Reset to verify new session
245251

246252
// Should be able to establish new session
@@ -270,7 +276,7 @@ void testReconnectErrorHandling() {
270276
String requestSessionId = exchange.getRequestHeaders().getFirst(HttpHeaders.MCP_SESSION_ID);
271277

272278
if ("GET".equals(method)) {
273-
int status = responseStatus.get();
279+
int status = serverResponseStatus.get();
274280

275281
if (status == 404 && requestSessionId != null) {
276282
// 404 with session ID - should trigger SessionNotFoundException
@@ -292,7 +298,7 @@ else if (status == 404) {
292298
else {
293299
// POST request handling
294300
exchange.getResponseHeaders().set("Content-Type", "application/json");
295-
String responseSessionId = sessionId.get();
301+
String responseSessionId = currentServerSessionId.get();
296302
if (responseSessionId != null) {
297303
exchange.getResponseHeaders().set(HttpHeaders.MCP_SESSION_ID, responseSessionId);
298304
}
@@ -304,8 +310,8 @@ else if (status == 404) {
304310
});
305311

306312
// Test with session ID - should get SessionNotFoundException
307-
responseStatus.set(200);
308-
sessionId.set("sse-session-1");
313+
serverResponseStatus.set(200);
314+
currentServerSessionId.set("sse-session-1");
309315

310316
var transport = HttpClientStreamableHttpTransport.builder(HOST)
311317
.endpoint("/mcp-sse")
@@ -316,19 +322,19 @@ else if (status == 404) {
316322
StepVerifier.create(transport.connect(msg -> msg)).verifyComplete();
317323

318324
// Send message to establish session
319-
var testMessage = createTestMessage();
325+
var testMessage = createTestRequestMessage();
320326
StepVerifier.create(transport.sendMessage(testMessage)).verifyComplete();
321327

322328
// Now simulate server returning 404 on reconnect
323-
responseStatus.set(404);
329+
serverResponseStatus.set(404);
324330

325331
// This should trigger reconnect which will fail
326332
// The error should be handled internally and passed to exception handler
327333

328334
StepVerifier.create(transport.closeGracefully()).verifyComplete();
329335
}
330336

331-
private McpSchema.JSONRPCRequest createTestMessage() {
337+
private McpSchema.JSONRPCRequest createTestRequestMessage() {
332338
var initializeRequest = new McpSchema.InitializeRequest(ProtocolVersions.MCP_2025_03_26,
333339
McpSchema.ClientCapabilities.builder().roots(true).build(),
334340
new McpSchema.Implementation("Test Client", "1.0.0"));

0 commit comments

Comments
 (0)