modelcontextprotocol/go-sdk
Watch
1
0
Fork
You've already forked go-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

mcp: unblock reads of ioConn as soon as Close is called

Browse source 
Use a read loop and incoming channel so that we can select on incoming
messages and unblock calls of ioConn.Read as soon as Close is called.
This avoids the scenario of #224, where a close of the StdioTransport
does not gracefully shut down the JSON-RPC connection.

Fixes #224
This commit is contained in:
Rob Findley 2025-08-04 23:20:55 +00:00 • committed by Robert Findley
commit a02c2ffd0e
2 changed files with 142 additions and 24 deletions
Download patch file Download diff file Expand all files Collapse all files

97
mcp/cmd_test.go
View file

@ -10,7 +10,9 @@ import (
"log"
"os"
"os/exec"
"os/signal"
"runtime"
"syscall"
"testing"
"time"
@ -21,14 +23,27 @@ import (
const runAsServer = "_MCP_RUN_AS_SERVER"
func TestMain(m *testing.M) {
if os.Getenv(runAsServer) != "" {
// If the runAsServer variable is set, execute the relevant serverFunc
// instead of running tests (aka the fork and exec trick).
if name := os.Getenv(runAsServer); name != "" {
run := serverFuncs[name]
if run == nil {
log.Fatalf("Unknown server %q", name)
}
os.Unsetenv(runAsServer)
runServer()
run()
return
}
os.Exit(m.Run())
}
// serverFuncs defines server functions that may be run as subprocesses via
// [TestMain].
var serverFuncs = map[string]func(){
"default": runServer,
"cancelContext": runCancelContextServer,
}
func runServer() {
ctx := context.Background()
@ -39,6 +54,16 @@ func runServer() {
}
}
func runCancelContextServer() {
ctx, done := signal.NotifyContext(context.Background(), syscall.SIGINT)
defer done()
server := mcp.NewServer(testImpl, nil)
if err := server.Run(ctx, mcp.NewStdioTransport()); err != nil {
log.Fatal(err)
}
}
func TestServerRunContextCancel(t *testing.T) {
server := mcp.NewServer(&mcp.Implementation{Name: "greeter", Version: "v0.0.1"}, nil)
mcp.AddTool(server, &mcp.Tool{Name: "greet", Description: "say hi"}, SayHi)
@ -80,15 +105,18 @@ func TestServerRunContextCancel(t *testing.T) {
}
func TestServerInterrupt(t *testing.T) {
if runtime.GOOS == "windows" {
t.Skip("requires POSIX signals")
}
requireExec(t)
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
cmd := createServerCommand(t)
cmd := createServerCommand(t, "default")
client := mcp.NewClient(testImpl, nil)
session, err := client.Connect(ctx, mcp.NewCommandTransport(cmd))
_, err := client.Connect(ctx, mcp.NewCommandTransport(cmd))
if err != nil {
t.Fatal(err)
}
@ -101,29 +129,64 @@ func TestServerInterrupt(t *testing.T) {
}()
// send a signal to the server process to terminate it
if runtime.GOOS == "windows" {
// Windows does not support os.Interrupt
session.Close()
} else {
cmd.Process.Signal(os.Interrupt)
}
cmd.Process.Signal(os.Interrupt)
// wait for the server to exit
// TODO: use synctest when availble
// TODO: use synctest when available
select {
case <-time.After(5 * time.Second):
t.Fatal("server did not exit after SIGTERM")
t.Fatal("server did not exit after SIGINT")
case <-onExit:
}
}
func TestStdioContextCancellation(t *testing.T) {
if runtime.GOOS == "windows" {
t.Skip("requires POSIX signals")
}
requireExec(t)
// This test is a variant of TestServerInterrupt reproducing the conditions
// of #224, where interrupt failed to shut down the server because reads of
// Stdin were not unblocked.
cmd := createServerCommand(t, "cancelContext")
// Creating a stdin pipe causes os.Stdin.Close to not immediately unblock
// pending reads.
_, _ = cmd.StdinPipe()
// Just Start the command, rather than connecting to the server, because we
// don't want the client connection to indirectly flush stdin through writes.
if err := cmd.Start(); err != nil {
t.Fatalf("starting command: %v", err)
}
// Sleep to make it more likely that the server is blocked in the read loop.
time.Sleep(100 * time.Millisecond)
onExit := make(chan struct{})
go func() {
cmd.Process.Wait()
close(onExit)
}()
cmd.Process.Signal(os.Interrupt)
select {
case <-time.After(5 * time.Second):
t.Fatal("server did not exit after SIGINT")
case <-onExit:
t.Logf("done.")
}
}
func TestCmdTransport(t *testing.T) {
requireExec(t)
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
cmd := createServerCommand(t)
cmd := createServerCommand(t, "default")
client := mcp.NewClient(&mcp.Implementation{Name: "client", Version: "v0.0.1"}, nil)
session, err := client.Connect(ctx, mcp.NewCommandTransport(cmd))
@ -150,7 +213,11 @@ func TestCmdTransport(t *testing.T) {
}
}
func createServerCommand(t *testing.T) *exec.Cmd {
// createServerCommand creates a command to fork and exec the test binary as an
// MCP server.
//
// serverName must refer to an entry in the [serverFuncs] map.
func createServerCommand(t *testing.T, serverName string) *exec.Cmd {
t.Helper()
exe, err := os.Executable()
@ -158,7 +225,7 @@ func createServerCommand(t *testing.T) *exec.Cmd {
t.Fatal(err)
}
cmd := exec.Command(exe)
cmd.Env = append(os.Environ(), runAsServer+"=true")
cmd.Env = append(os.Environ(), runAsServer+"="+serverName)
return cmd
}

71
mcp/transport.go
View file

@ -265,7 +265,9 @@ func (r rwc) Close() error {
// See [msgBatch] for more discussion of message batching.
type ioConn struct {
rwc io.ReadWriteCloser // the underlying stream
in *json.Decoder // a decoder bound to rwc
// incoming receives messages from the read loop started in [newIOConn].
incoming <-chan msgOrErr
// If outgoiBatch has a positive capacity, it will be used to batch requests
// and notifications before sending.
@ -279,12 +281,47 @@ type ioConn struct {
// Since writes may be concurrent to reads, we need to guard this with a mutex.
batchMu sync.Mutex
batches map[jsonrpc2.ID]*msgBatch // lazily allocated
closeOnce sync.Once
closed chan struct{}
closeErr error
}
type msgOrErr struct {
msg json.RawMessage
err error
}
func newIOConn(rwc io.ReadWriteCloser) *ioConn {
var (
incoming = make(chan msgOrErr)
closed = make(chan struct{})
)
// Start a goroutine for reads, so that we can select on the incoming channel
// in [ioConn.Read] and unblock the read as soon as Close is called (see #224).
//
// This leaks a goroutine, but that is unavoidable since AFAIK there is no
// (easy and portable) way to guarantee that reads of stdin are unblocked
// when closed.
go func() {
dec := json.NewDecoder(rwc)
for {
var raw json.RawMessage
err := dec.Decode(&raw)
select {
case incoming <- msgOrErr{msg: raw, err: err}:
case <-closed:
return
}
if err != nil {
return
}
}
}()
return &ioConn{
rwc: rwc,
in: json.NewDecoder(rwc),
rwc: rwc,
incoming: incoming,
closed: closed,
}
}
@ -356,10 +393,8 @@ type msgBatch struct {
}
func (t *ioConn) Read(ctx context.Context) (jsonrpc.Message, error) {
return t.read(ctx, t.in)
}
func (t *ioConn) read(ctx context.Context, in *json.Decoder) (jsonrpc.Message, error) {
// As a matter of principle, enforce that reads on a closed context return an
// error.
select {
case <-ctx.Done():
return nil, ctx.Err()
@ -372,9 +407,20 @@ func (t *ioConn) read(ctx context.Context, in *json.Decoder) (jsonrpc.Message, e
}
var raw json.RawMessage
if err := in.Decode(&raw); err != nil {
return nil, err
select {
case <-ctx.Done():
return nil, ctx.Err()
case v := <-t.incoming:
if v.err != nil {
return nil, v.err
}
raw = v.msg
case <-t.closed:
return nil, io.EOF
}
msgs, batch, err := readBatch(raw)
if err != nil {
return nil, err
@ -431,6 +477,7 @@ func readBatch(data []byte) (msgs []jsonrpc.Message, isBatch bool, _ error) {
}
func (t *ioConn) Write(ctx context.Context, msg jsonrpc.Message) error {
// As in [ioConn.Read], enforce that Writes on a closed context are an error.
select {
case <-ctx.Done():
return ctx.Err()
@ -478,7 +525,11 @@ func (t *ioConn) Write(ctx context.Context, msg jsonrpc.Message) error {
}
func (t *ioConn) Close() error {
return t.rwc.Close()
t.closeOnce.Do(func() {
t.closeErr = t.rwc.Close()
close(t.closed)
})
return t.closeErr
}
func marshalMessages[T jsonrpc.Message](msgs []T) ([]byte, error) {