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" {
-		// 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
 }