Skip to content

Commit

Permalink
refactor: improve resource management in StdioServerTransport
Browse files Browse the repository at this point in the history
- Use try-with-resources for proper reader cleanup
- Remove unnecessary reader/writer instance variables
- Access process streams directly when needed
  • Loading branch information
tzolov committed Dec 12, 2024
1 parent 8397849 commit 8a0cd4e
Showing 1 changed file with 8 additions and 18 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -17,8 +17,8 @@
package org.springframework.ai.mcp.client.stdio;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
Expand Down Expand Up @@ -51,12 +51,6 @@ public class StdioServerTransport extends AbstractMcpTransport {

private Process process;

private BufferedReader processErrorReader;

private BufferedReader processReader;

private BufferedWriter processWriter;

private Scheduler inboundScheduler;

private Scheduler outboundScheduler;
Expand Down Expand Up @@ -116,11 +110,6 @@ public void start() {
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
this.isRunning = true;
startInboundProcessing();
Expand All @@ -143,7 +132,8 @@ public void awaitForExit() {

private void startErrorProcessing() {
this.errorScheduler.schedule(() -> {
try {
try (BufferedReader processErrorReader = new BufferedReader(
new InputStreamReader(process.getErrorStream()))) {
String line;
while (isRunning && processErrorReader != null && (line = processErrorReader.readLine()) != null) {
try {
Expand All @@ -169,9 +159,9 @@ private void startErrorProcessing() {

private void startInboundProcessing() {
this.inboundScheduler.schedule(() -> {
try {
try (BufferedReader processReader = new BufferedReader(new InputStreamReader(process.getInputStream()))) {
String line;
while (this.isRunning && this.processReader != null && (line = this.processReader.readLine()) != null) {
while (this.isRunning && processReader != null && (line = processReader.readLine()) != null) {
try {
JSONRPCMessage message = deserializeJsonRpcMessage(line);
if (!this.getInboundSink().tryEmitNext(message).isSuccess()) {
Expand Down Expand Up @@ -222,9 +212,9 @@ private void startOutboundProcessing() {
.handle((message, s) -> {
if (message != null) {
try {
this.processWriter.write(objectMapper.writeValueAsString(message));
this.processWriter.newLine();
this.processWriter.flush();
this.process.outputWriter().write(objectMapper.writeValueAsString(message));
this.process.outputWriter().newLine();
this.process.outputWriter().flush();
s.next(message);
}
catch (IOException e) {
Expand Down

0 comments on commit 8a0cd4e

Please sign in to comment.