From 56ed3872f3c8a7a62ba9f985dc0a25d323a8768b Mon Sep 17 00:00:00 2001 From: Rob Findley Date: Tue, 29 Apr 2025 15:08:52 +0000 Subject: [PATCH] internal/mcp: add a command transport for connecting to a sidecar Add a missing CommandTransport that communicates with a sidecar process over newline-delimited stdin/stdout, the counterpart to the existing StdIOTransport. Add a test using the usual TestMain trick. Change-Id: I170ba78faf6b6d7fca7da1af11bf92e919b11ed5 Reviewed-on: https://go-review.googlesource.com/c/tools/+/668895 LUCI-TryBot-Result: Go LUCI Reviewed-by: Jonathan Amsterdam --- mcp/cmd.go | 106 +++++++++++++++++++++++++++++++++++++++++++++++ mcp/cmd_test.go | 67 ++++++++++++++++++++++++++++++ mcp/transport.go | 6 +-- 3 files changed, 175 insertions(+), 4 deletions(-) create mode 100644 mcp/cmd.go create mode 100644 mcp/cmd_test.go diff --git a/mcp/cmd.go b/mcp/cmd.go new file mode 100644 index 0000000..42af63b --- /dev/null +++ b/mcp/cmd.go @@ -0,0 +1,106 @@ +// Copyright 2025 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by an MIT-style +// license that can be found in the LICENSE file. + +package mcp + +import ( + "context" + "fmt" + "io" + "os/exec" + "syscall" + "time" +) + +// A CommandTransport is a [Transport] that runs a command and communicates +// with it over stdin/stdout, using newline-delimited JSON. +type CommandTransport struct { + cmd *exec.Cmd +} + +// NewCommandTransport returns a [CommandTransport] that runs the given command +// and communicates with it over stdin/stdout. +// +// The resulting transport takes ownership of the command, starting it during +// [CommandTransport.Connect], and stopping it when the connection is closed. +func NewCommandTransport(cmd *exec.Cmd) *CommandTransport { + return &CommandTransport{cmd} +} + +// Connect starts the command, and connects to it over stdin/stdout. +func (t *CommandTransport) Connect(ctx context.Context) (Stream, error) { + stdout, err := t.cmd.StdoutPipe() + if err != nil { + return nil, err + } + stdout = io.NopCloser(stdout) // close the connection by closing stdin, not stdout + stdin, err := t.cmd.StdinPipe() + if err != nil { + return nil, err + } + if err := t.cmd.Start(); err != nil { + return nil, err + } + return newIOStream(&pipeRWC{t.cmd, stdout, stdin}), nil +} + +// A pipeRWC is an io.ReadWriteCloser that communicates with a subprocess over +// stdin/stdout pipes. +type pipeRWC struct { + cmd *exec.Cmd + stdout io.ReadCloser + stdin io.WriteCloser +} + +func (s *pipeRWC) Read(p []byte) (n int, err error) { + return s.stdout.Read(p) +} + +func (s *pipeRWC) Write(p []byte) (n int, err error) { + return s.stdin.Write(p) +} + +// Close closes the input stream to the child process, and awaits normal +// termination of the command. If the command does not exit, it is signalled to +// terminate, and then eventually killed. +func (s *pipeRWC) Close() error { + // Spec: + // "For the stdio transport, the client SHOULD initiate shutdown by:... + + // "...First, closing the input stream to the child process (the server)" + if err := s.stdin.Close(); err != nil { + return fmt.Errorf("closing stdin: %v", err) + } + resChan := make(chan error, 1) + go func() { + resChan <- s.cmd.Wait() + }() + // "...Waiting for the server to exit, or sending SIGTERM if the server does not exit within a reasonable time" + wait := func() (error, bool) { + select { + case err := <-resChan: + return err, true + case <-time.After(5 * time.Second): + } + return nil, false + } + if err, ok := wait(); ok { + return err + } + // Note the condition here: if sending SIGTERM fails, don't wait and just + // move on to SIGKILL. + if err := s.cmd.Process.Signal(syscall.SIGTERM); err == nil { + if err, ok := wait(); ok { + return err + } + } + // "...Sending SIGKILL if the server does not exit within a reasonable time after SIGTERM" + if err := s.cmd.Process.Kill(); err != nil { + return err + } + if err, ok := wait(); ok { + return err + } + return fmt.Errorf("unresponsive subprocess") +} diff --git a/mcp/cmd_test.go b/mcp/cmd_test.go new file mode 100644 index 0000000..d989eb1 --- /dev/null +++ b/mcp/cmd_test.go @@ -0,0 +1,67 @@ +// Copyright 2025 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by an MIT-style +// license that can be found in the LICENSE file. + +package mcp_test + +import ( + "context" + "log" + "os" + "os/exec" + "testing" + + "github.com/google/go-cmp/cmp" + "github.com/modelcontextprotocol/go-sdk/mcp" +) + +const runAsServer = "_MCP_RUN_AS_SERVER" + +func TestMain(m *testing.M) { + if os.Getenv(runAsServer) != "" { + os.Unsetenv(runAsServer) + runServer() + return + } + os.Exit(m.Run()) +} + +func runServer() { + ctx := context.Background() + + server := mcp.NewServer("greeter", "v0.0.1", nil) + server.AddTools(mcp.MakeTool("greet", "say hi", SayHi)) + + if err := server.Run(ctx, mcp.NewStdIOTransport(), nil); err != nil { + log.Fatal(err) + } +} + +func TestCmdTransport(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + exe, err := os.Executable() + if err != nil { + t.Fatal(err) + } + cmd := exec.Command(exe) + cmd.Env = append(os.Environ(), runAsServer+"=true") + + client := mcp.NewClient("client", "v0.0.1", nil) + serverConn, err := client.Connect(ctx, mcp.NewCommandTransport(cmd), nil) + if err != nil { + log.Fatal(err) + } + got, err := serverConn.CallTool(ctx, "greet", SayHiParams{Name: "user"}) + if err != nil { + log.Fatal(err) + } + want := []mcp.Content{mcp.TextContent{Text: "Hi user"}} + if diff := cmp.Diff(want, got); diff != "" { + t.Errorf("greet returned unexpected content (-want +got):\n%s", diff) + } + if err := serverConn.Close(); err != nil { + t.Fatalf("closing server: %v", err) + } +} diff --git a/mcp/transport.go b/mcp/transport.go index 2a59c81..d06bcb1 100644 --- a/mcp/transport.go +++ b/mcp/transport.go @@ -227,13 +227,11 @@ type rwc struct { } func (r rwc) Read(p []byte) (n int, err error) { - n, err = r.rc.Read(p) - return n, err + return r.rc.Read(p) } func (r rwc) Write(p []byte) (n int, err error) { - n, err = r.wc.Write(p) - return n, err + return r.wc.Write(p) } func (r rwc) Close() error {