diff --git a/mcp/cmd_test.go b/mcp/cmd_test.go index bc149f4..82a35a8 100644 --- a/mcp/cmd_test.go +++ b/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 } diff --git a/mcp/transport.go b/mcp/transport.go index 5175f6f..dccc920 100644 --- a/mcp/transport.go +++ b/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, - 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) {