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

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

mcp/cmd_test.go

@@ -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" {
+	cmd.Process.Signal(os.Interrupt)
-		// Windows does not support os.Interrupt
-		session.Close()
-	} else {
-		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
 }

mcp/transport.go

@@ -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,
+		rwc:      rwc,
-		in:  json.NewDecoder(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)
+	// As a matter of principle, enforce that reads on a closed context return an
-}
+	// error.
-
-func (t *ioConn) read(ctx context.Context, in *json.Decoder) (jsonrpc.Message, 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 {
+	select {
-		return nil, err
+	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) {