Skip to content

Commit b91380e

Browse files
chemicLchemicL
committed
Cleanup, javadoc, refactor
1 parent f7ea062 commit b91380e

File tree

13 files changed

+347
-178
lines changed

13 files changed

+347
-178
lines changed

‎mcp-spring/mcp-spring-webflux/pom.xml‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,15 @@
8888
<version>${byte-buddy.version}</version>
8989
<scope>test</scope>
9090
</dependency>
91+
<dependency>
92+
<groupId>org.springframework</groupId>
93+
<artifactId>spring-context</artifactId>
94+
<version>6.2.6</version>
95+
</dependency>
96+
<dependency>
97+
<groupId>io.projectreactor.netty</groupId>
98+
<artifactId>reactor-netty-http</artifactId>
99+
</dependency>
91100
<dependency>
92101
<groupId>io.projectreactor</groupId>
93102
<artifactId>reactor-test</artifactId>
@@ -117,7 +126,7 @@
117126
<groupId>ch.qos.logback</groupId>
118127
<artifactId>logback-classic</artifactId>
119128
<version>${logback.version}</version>
120-
<scope>test</scope>
129+
<!-- <scope>test</scope>-->
121130
</dependency>
122131

123132
<dependency>

‎mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/client/transport/Main.java‎

Lines changed: 48 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -3,15 +3,29 @@
33
import com.fasterxml.jackson.databind.ObjectMapper;
44
import io.modelcontextprotocol.client.McpAsyncClient;
55
import io.modelcontextprotocol.client.McpClient;
6+
import io.modelcontextprotocol.client.McpSyncClient;
67
import io.modelcontextprotocol.spec.McpSchema;
8+
import io.netty.channel.socket.nio.NioChannelOption;
9+
import jdk.net.ExtendedSocketOptions;
10+
import org.springframework.http.client.reactive.ReactorClientHttpConnector;
711
import org.springframework.web.reactive.function.client.WebClient;
12+
import reactor.netty.http.client.HttpClient;
13+
14+
import java.util.List;
15+
import java.util.Map;
16+
import java.util.Scanner;
17+
import java.util.concurrent.atomic.AtomicReference;
818

919
public class Main {
1020

11-
public static void main(String[] args) {
12-
McpAsyncClient client = McpClient
13-
.async(new WebClientStreamableHttpTransport(new ObjectMapper(),
14-
WebClient.builder().baseUrl("http://localhost:3001"), "/mcp", true, false))
21+
public static void main(String[] args) throws InterruptedException {
22+
McpSyncClient client = McpClient
23+
.sync(new WebClientStreamableHttpTransport(new ObjectMapper(),
24+
WebClient.builder()
25+
.clientConnector(new ReactorClientHttpConnector(
26+
HttpClient.create().option(NioChannelOption.of(ExtendedSocketOptions.TCP_KEEPIDLE), 5)))
27+
.baseUrl("http://localhost:3001"),
28+
"/mcp", true, false))
1529
.build();
1630

1731
/*
@@ -34,11 +48,36 @@ public static void main(String[] args) {
3448
* tools 6. -> 2xx response
3549
*/
3650

37-
client.initialize()
38-
.flatMap(r -> client.listTools())
39-
.map(McpSchema.ListToolsResult::tools)
40-
.doOnNext(System.out::println)
41-
.block();
51+
List<McpSchema.Tool> tools = null;
52+
while (tools == null) {
53+
try {
54+
client.initialize();
55+
tools = client.listTools().tools();
56+
}
57+
catch (Exception e) {
58+
System.out.println("Got exception. Retrying in 5s. " + e);
59+
Thread.sleep(5000);
60+
}
61+
}
62+
63+
Scanner scanner = new Scanner(System.in);
64+
while (scanner.hasNext()) {
65+
String text = scanner.nextLine();
66+
if (text == null || text.isEmpty()) {
67+
System.out.println("Done");
68+
break;
69+
}
70+
try {
71+
McpSchema.CallToolResult result = client
72+
.callTool(new McpSchema.CallToolRequest(tools.get(0).name(), Map.of("message", text)));
73+
System.out.println("Tool call result: " + result);
74+
}
75+
catch (Exception e) {
76+
System.out.println("Error calling tool " + e);
77+
}
78+
}
79+
80+
client.closeGracefully();
4281
}
4382

4483
}

‎mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/client/transport/WebClientStreamableHttpTransport.java‎

Lines changed: 28 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import com.fasterxml.jackson.core.type.TypeReference;
44
import com.fasterxml.jackson.databind.ObjectMapper;
5+
import io.modelcontextprotocol.spec.DefaultMcpTransportSession;
56
import io.modelcontextprotocol.spec.McpClientTransport;
67
import io.modelcontextprotocol.spec.McpError;
78
import io.modelcontextprotocol.spec.McpSchema;
@@ -51,20 +52,21 @@ public class WebClientStreamableHttpTransport implements McpClientTransport {
5152

5253
private final boolean resumableStreams;
5354

54-
private final AtomicReference<McpTransportSession> activeSession = new AtomicReference<>();
55+
private final AtomicReference<DefaultMcpTransportSession> activeSession = new AtomicReference<>();
5556

5657
private final AtomicReference<Function<Mono<McpSchema.JSONRPCMessage>, Mono<McpSchema.JSONRPCMessage>>> handler = new AtomicReference<>();
5758

5859
private final AtomicReference<Consumer<Throwable>> exceptionHandler = new AtomicReference<>();
5960

61+
// TODO: builder
6062
public WebClientStreamableHttpTransport(ObjectMapper objectMapper, WebClient.Builder webClientBuilder,
6163
String endpoint, boolean resumableStreams, boolean openConnectionOnStartup) {
6264
this.objectMapper = objectMapper;
6365
this.webClient = webClientBuilder.build();
6466
this.endpoint = endpoint;
6567
this.resumableStreams = resumableStreams;
6668
this.openConnectionOnStartup = openConnectionOnStartup;
67-
this.activeSession.set(new McpTransportSession());
69+
this.activeSession.set(new DefaultMcpTransportSession());
6870
}
6971

7072
@Override
@@ -80,15 +82,15 @@ public Mono<Void> connect(Function<Mono<McpSchema.JSONRPCMessage>, Mono<McpSchem
8082
}
8183

8284
@Override
83-
public void registerExceptionHandler(Consumer<Throwable> handler) {
85+
public void setExceptionHandler(Consumer<Throwable> handler) {
8486
logger.debug("Exception handler registered");
8587
this.exceptionHandler.set(handler);
8688
}
8789

8890
private void handleException(Throwable t) {
89-
logger.debug("Handling exception for session {}", activeSession.get().sessionId(), t);
91+
logger.debug("Handling exception for session {}", sessionIdRepresentation(this.activeSession.get()), t);
9092
if (t instanceof McpSessionNotFoundException) {
91-
McpTransportSession invalidSession = this.activeSession.getAndSet(new McpTransportSession());
93+
McpTransportSession<?> invalidSession = this.activeSession.getAndSet(new DefaultMcpTransportSession());
9294
logger.warn("Server does not recognize session {}. Invalidating.", invalidSession.sessionId());
9395
invalidSession.close();
9496
}
@@ -102,7 +104,7 @@ private void handleException(Throwable t) {
102104
public Mono<Void> closeGracefully() {
103105
return Mono.defer(() -> {
104106
logger.debug("Graceful close triggered");
105-
McpTransportSession currentSession = this.activeSession.get();
107+
DefaultMcpTransportSession currentSession = this.activeSession.get();
106108
if (currentSession != null) {
107109
return currentSession.closeGracefully();
108110
}
@@ -125,16 +127,14 @@ private void reconnect(McpStream stream, ContextView ctx) {
125127
// listen for messages.
126128
// If it doesn't, nothing actually happens here, that's just the way it is...
127129
final AtomicReference<Disposable> disposableRef = new AtomicReference<>();
128-
final McpTransportSession transportSession = this.activeSession.get();
130+
final McpTransportSession<Disposable> transportSession = this.activeSession.get();
129131
Disposable connection = webClient.get()
130132
.uri(this.endpoint)
131133
.accept(MediaType.TEXT_EVENT_STREAM)
132134
.headers(httpHeaders -> {
133-
if (transportSession.sessionId() != null) {
134-
httpHeaders.add("mcp-session-id", transportSession.sessionId());
135-
}
136-
if (stream != null && stream.lastId() != null) {
137-
httpHeaders.add("last-event-id", stream.lastId());
135+
transportSession.sessionId().ifPresent(id -> httpHeaders.add("mcp-session-id", id));
136+
if (stream != null) {
137+
stream.lastId().ifPresent(id -> httpHeaders.add("last-event-id", id));
138138
}
139139
})
140140
.exchangeToFlux(response -> {
@@ -161,7 +161,7 @@ else if (response.statusCode().isSameCodeAs(HttpStatus.NOT_FOUND)) {
161161
logger.warn("Session {} was not found on the MCP server", transportSession.sessionId());
162162

163163
McpSessionNotFoundException notFoundException = new McpSessionNotFoundException(
164-
transportSession.sessionId());
164+
sessionIdRepresentation(transportSession));
165165
// inform the stream/connection subscriber
166166
return Flux.error(notFoundException);
167167
}
@@ -187,6 +187,10 @@ else if (response.statusCode().isSameCodeAs(HttpStatus.NOT_FOUND)) {
187187
transportSession.addConnection(connection);
188188
}
189189

190+
private static String sessionIdRepresentation(McpTransportSession<?> transportSession) {
191+
return transportSession.sessionId().orElse("[missing_session_id]");
192+
}
193+
190194
@Override
191195
public Mono<Void> sendMessage(McpSchema.JSONRPCMessage message) {
192196
return Mono.create(sink -> {
@@ -197,15 +201,13 @@ public Mono<Void> sendMessage(McpSchema.JSONRPCMessage message) {
197201
// listen for messages.
198202
// If it doesn't, nothing actually happens here, that's just the way it is...
199203
final AtomicReference<Disposable> disposableRef = new AtomicReference<>();
200-
final McpTransportSession transportSession = this.activeSession.get();
204+
final McpTransportSession<Disposable> transportSession = this.activeSession.get();
201205

202206
Disposable connection = webClient.post()
203207
.uri(this.endpoint)
204208
.accept(MediaType.TEXT_EVENT_STREAM, MediaType.APPLICATION_JSON)
205209
.headers(httpHeaders -> {
206-
if (transportSession.sessionId() != null) {
207-
httpHeaders.add("mcp-session-id", transportSession.sessionId());
208-
}
210+
transportSession.sessionId().ifPresent(id -> httpHeaders.add("mcp-session-id", id));
209211
})
210212
.bodyValue(message)
211213
.exchangeToFlux(response -> {
@@ -287,7 +289,7 @@ else if (contentType.isCompatibleWith(MediaType.APPLICATION_JSON)) {
287289
logger.warn("Session {} was not found on the MCP server", transportSession.sessionId());
288290

289291
McpSessionNotFoundException notFoundException = new McpSessionNotFoundException(
290-
transportSession.sessionId());
292+
sessionIdRepresentation(transportSession));
291293
// inform the stream/connection subscriber
292294
return Flux.error(notFoundException);
293295
}
@@ -313,8 +315,8 @@ else if (contentType.isCompatibleWith(MediaType.APPLICATION_JSON)) {
313315
// invalidate the session
314316
// https://github.com/modelcontextprotocol/typescript-sdk/issues/389
315317
if (responseException.getStatusCode().isSameCodeAs(HttpStatus.BAD_REQUEST)) {
316-
return Mono.error(new McpSessionNotFoundException(this.activeSession.get().sessionId(),
317-
toPropagate));
318+
return Mono.error(new McpSessionNotFoundException(
319+
sessionIdRepresentation(this.activeSession.get()), toPropagate));
318320
}
319321
return Mono.empty();
320322
}).flux();
@@ -381,8 +383,8 @@ private class McpStream {
381383
this.resumable = resumable;
382384
}
383385

384-
String lastId() {
385-
return this.lastId.get();
386+
Optional<String> lastId() {
387+
return Optional.ofNullable(this.lastId.get());
386388
}
387389

388390
long streamId() {
@@ -395,9 +397,10 @@ Flux<McpSchema.JSONRPCMessage> consumeSseStream(
395397
if (resumable && !(e instanceof McpSessionNotFoundException)) {
396398
reconnect(this, ctx);
397399
}
398-
})
399-
.doOnNext(idAndMessage -> idAndMessage.getT1().ifPresent(this.lastId::set))
400-
.flatMapIterable(Tuple2::getT2));
400+
}).doOnNext(idAndMessage -> idAndMessage.getT1().ifPresent(id -> {
401+
String previousId = this.lastId.getAndSet(id);
402+
logger.debug("Updating last id {} -> {} for stream {}", previousId, id, this.streamId);
403+
})).flatMapIterable(Tuple2::getT2));
401404
}
402405

403406
}

‎mcp-spring/mcp-spring-webflux/src/main/resources/logback.xml‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
<?xml version="1.0" encoding="UTF-8"?>
2+
<!DOCTYPE configuration>
3+
4+
<configuration>
5+
<appender name="CONSOLE" class="ch.qos.logback.core.ConsoleAppender">
6+
<encoder>
7+
<pattern>%d{yyyy-MM-dd HH:mm:ss} [%thread] %-5level %logger{36} - %msg%n</pattern>
8+
</encoder>
9+
</appender>
10+
11+
<!-- Main MCP package -->
12+
<logger name="io.modelcontextprotocol" level="DEBUG"/>
13+
14+
<!-- Root logger -->
15+
<root level="INFO">
16+
<appender-ref ref="CONSOLE"/>
17+
</root>
18+
</configuration>

‎mcp-test/src/main/java/io/modelcontextprotocol/MockMcpTransport.java‎

Lines changed: 0 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,6 @@
77
import java.util.ArrayList;
88
import java.util.List;
99
import java.util.function.BiConsumer;
10-
import java.util.function.Consumer;
1110
import java.util.function.Function;
1211

1312
import com.fasterxml.jackson.core.type.TypeReference;
@@ -71,18 +70,6 @@ public McpSchema.JSONRPCMessage getLastSentMessage() {
7170

7271
private volatile boolean connected = false;
7372

74-
// @Override
75-
// public Mono<Void> connect(Consumer<McpSchema.JSONRPCMessage> consumer) {
76-
// if (connected) {
77-
// return Mono.error(new IllegalStateException("Already connected"));
78-
// }
79-
// connected = true;
80-
// return inbound.asFlux()
81-
// .doOnNext(consumer)
82-
// .doFinally(signal -> connected = false)
83-
// .then();
84-
// }
85-
8673
@Override
8774
public Mono<Void> connect(Function<Mono<McpSchema.JSONRPCMessage>, Mono<McpSchema.JSONRPCMessage>> handler) {
8875
if (connected) {

‎mcp-test/src/main/java/io/modelcontextprotocol/client/AbstractMcpAsyncClientResiliencyTests.java‎

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,8 @@
1616

1717
import java.io.IOException;
1818
import java.time.Duration;
19+
import java.util.List;
20+
import java.util.Map;
1921
import java.util.concurrent.atomic.AtomicReference;
2022
import java.util.function.Consumer;
2123
import java.util.function.Function;
@@ -60,7 +62,6 @@ public abstract class AbstractMcpAsyncClientResiliencyTests {
6062
final String ipAddressViaToxiproxy = toxiproxy.getHost();
6163
final int portViaToxiproxy = toxiproxy.getMappedPort(3000);
6264

63-
// int port = container.getMappedPort(3001);
6465
host = "http://" + ipAddressViaToxiproxy + ":" + portViaToxiproxy;
6566
}
6667

@@ -172,4 +173,26 @@ void testSessionInvalidation() {
172173
});
173174
}
174175

176+
@Test
177+
void testCallTool() {
178+
withClient(createMcpTransport(), mcpAsyncClient -> {
179+
AtomicReference<List<McpSchema.Tool>> tools = new AtomicReference<>();
180+
StepVerifier.create(mcpAsyncClient.initialize()).expectNextCount(1).verifyComplete();
181+
StepVerifier.create(mcpAsyncClient.listTools())
182+
.consumeNextWith(list -> tools.set(list.tools()))
183+
.verifyComplete();
184+
185+
disconnect();
186+
187+
String name = tools.get().get(0).name();
188+
// Assuming this is the echo tool
189+
McpSchema.CallToolRequest request = new McpSchema.CallToolRequest(name, Map.of("message", "hello"));
190+
StepVerifier.create(mcpAsyncClient.callTool(request)).expectError().verify();
191+
192+
reconnect();
193+
194+
StepVerifier.create(mcpAsyncClient.callTool(request)).expectNextCount(1).verifyComplete();
195+
});
196+
}
197+
175198
}

0 commit comments

Comments
 (0)