refactor: improve resource management in StdioServerTransport

- Use try-with-resources for proper reader cleanup
- Remove unnecessary reader/writer instance variables
- Access process streams directly when needed
This commit is contained in:
Christian Tzolov 2024-12-12 23:18:19 +01:00
commit 8a0cd4ef06

View file

@ -17,8 +17,8 @@
package org.springframework.ai.mcp.client.stdio; package org.springframework.ai.mcp.client.stdio;
import java.io.BufferedReader; import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.IOException; import java.io.IOException;
import java.io.InputStreamReader;
import java.util.ArrayList; import java.util.ArrayList;
import java.util.HashMap; import java.util.HashMap;
import java.util.List; import java.util.List;
@ -51,12 +51,6 @@ public class StdioServerTransport extends AbstractMcpTransport {
private Process process; private Process process;
private BufferedReader processErrorReader;
private BufferedReader processReader;
private BufferedWriter processWriter;
private Scheduler inboundScheduler; private Scheduler inboundScheduler;
private Scheduler outboundScheduler; private Scheduler outboundScheduler;
@ -116,11 +110,6 @@ public class StdioServerTransport extends AbstractMcpTransport {
throw new RuntimeException("Process input or output stream is null"); throw new RuntimeException("Process input or output stream is null");
} }
// Initialize readers and writers
this.processErrorReader = this.process.errorReader();
this.processReader = this.process.inputReader();
this.processWriter = this.process.outputWriter();
// Start threads // Start threads
this.isRunning = true; this.isRunning = true;
startInboundProcessing(); startInboundProcessing();
@ -143,7 +132,8 @@ public class StdioServerTransport extends AbstractMcpTransport {
private void startErrorProcessing() { private void startErrorProcessing() {
this.errorScheduler.schedule(() -> { this.errorScheduler.schedule(() -> {
try { try (BufferedReader processErrorReader = new BufferedReader(
new InputStreamReader(process.getErrorStream()))) {
String line; String line;
while (isRunning && processErrorReader != null && (line = processErrorReader.readLine()) != null) { while (isRunning && processErrorReader != null && (line = processErrorReader.readLine()) != null) {
try { try {
@ -169,9 +159,9 @@ public class StdioServerTransport extends AbstractMcpTransport {
private void startInboundProcessing() { private void startInboundProcessing() {
this.inboundScheduler.schedule(() -> { this.inboundScheduler.schedule(() -> {
try { try (BufferedReader processReader = new BufferedReader(new InputStreamReader(process.getInputStream()))) {
String line; String line;
while (this.isRunning && this.processReader != null && (line = this.processReader.readLine()) != null) { while (this.isRunning && processReader != null && (line = processReader.readLine()) != null) {
try { try {
JSONRPCMessage message = deserializeJsonRpcMessage(line); JSONRPCMessage message = deserializeJsonRpcMessage(line);
if (!this.getInboundSink().tryEmitNext(message).isSuccess()) { if (!this.getInboundSink().tryEmitNext(message).isSuccess()) {
@ -222,9 +212,9 @@ public class StdioServerTransport extends AbstractMcpTransport {
.handle((message, s) -> { .handle((message, s) -> {
if (message != null) { if (message != null) {
try { try {
this.processWriter.write(objectMapper.writeValueAsString(message)); this.process.outputWriter().write(objectMapper.writeValueAsString(message));
this.processWriter.newLine(); this.process.outputWriter().newLine();
this.processWriter.flush(); this.process.outputWriter().flush();
s.next(message); s.next(message);
} }
catch (IOException e) { catch (IOException e) {