gopls/internal/mcp: add -logfile and -rpc.trace to the headless server

For debugging interactions with gopls' headless MCP server, add a
-logfile and -rpc.trace flag. The names are chosen to align with gopls
serve, though the RPCs in question are now MCP, not LSP.

For golang/go#73580

Change-Id: I93525e8bd7ea3fe7c35445027961d430a6e75819
Reviewed-on: https://go-review.googlesource.com/c/tools/+/686815
Reviewed-by: Madeline Kalil <[email protected]>
LUCI-TryBot-Result: Go LUCI <[email protected]>

Author: findleyr

gopls/internal/cmd/mcp.go

@@ -8,6 +8,7 @@ import (
 	"context"
 	"flag"
 	"fmt"
+	"io"
 	"log"
 	"os"
 	"time"
@@ -21,7 +22,9 @@ import (
 type headlessMCP struct {
 	app *Application
 
-	Address string `flag:"listen" help:"the address on which to run the mcp server"`
+	Address  string `flag:"listen" help:"the address on which to run the mcp server"`
+	Logfile  string `flag:"logfile" help:"filename to log to; if unset, logs to stderr"`
+	RPCTrace bool   `flag:"rpc.trace" help:"print MCP rpc traces; cannot be used with -listen"`
 }
 
 func (m *headlessMCP) Name() string      { return "mcp" }
@@ -42,6 +45,21 @@ Examples:
 }
 
 func (m *headlessMCP) Run(ctx context.Context, args ...string) error {
+	if m.Address != "" && m.RPCTrace {
+		// There's currently no way to plumb logging instrumentation into the SSE
+		// transport that is created on connections to the HTTP handler, so we must
+		// disallow the -rpc.trace flag when using -listen.
+		return fmt.Errorf("-listen is incompatible with -rpc.trace")
+	}
+	if m.Logfile != "" {
+		f, err := os.Create(m.Logfile)
+		if err != nil {
+			return fmt.Errorf("opening logfile: %v", err)
+		}
+		log.SetOutput(f)
+		defer f.Close()
+	}
+
 	// Start a new in-process gopls session and create a fake client
 	// to connect to it.
 	cli, sess, err := m.app.connect(ctx)
@@ -112,6 +130,11 @@ func (m *headlessMCP) Run(ctx context.Context, args ...string) error {
 		return mcp.Serve(ctx, m.Address, eventChan, false)
 	} else {
 		countHeadlessMCPStdIO.Inc()
-		return mcp.StartStdIO(ctx, sess, cli.server)
+		var rpcLog io.Writer
+		if m.RPCTrace {
+			rpcLog = log.Writer() // possibly redirected by -logfile above
+		}
+		log.Printf("Listening for MCP messages on stdin...")
+		return mcp.StartStdIO(ctx, sess, cli.server, rpcLog)
 	}
 }

gopls/internal/cmd/mcp_test.go

@@ -22,7 +22,6 @@ import (
 
 func TestMCPCommandStdio(t *testing.T) {
 	// Test that the headless MCP subcommand works, and recognizes file changes.
-
 	testenv.NeedsExec(t) // stdio transport uses execve(2)
 	tree := writeTree(t, `
 -- go.mod --
@@ -93,6 +92,67 @@ const B = 2
 	}
 }
 
+func TestMCPCommandLogging(t *testing.T) {
+	// Test that logging flags for headless MCP subcommand work as intended.
+	testenv.NeedsExec(t) // stdio transport uses execve(2)
+
+	tests := []struct {
+		logFile  string // also the subtest name
+		trace    bool
+		want     string
+		dontWant string
+	}{
+		{"notrace.log", false, "stdin", "initialized"},
+		{"trace.log", true, "initialized", ""},
+	}
+
+	dir := t.TempDir()
+	for _, test := range tests {
+		t.Run(test.logFile, func(t *testing.T) {
+			tree := writeTree(t, `
+-- go.mod --
+module example.com
+go 1.18
+
+-- a.go --
+package p
+`)
+
+			logFile := filepath.Join(dir, test.logFile)
+			args := []string{"mcp", "-logfile", logFile}
+			if test.trace {
+				args = append(args, "-rpc.trace")
+			}
+			goplsCmd := exec.Command(os.Args[0], args...)
+			goplsCmd.Env = append(os.Environ(), "ENTRYPOINT=goplsMain")
+			goplsCmd.Dir = tree
+
+			ctx := t.Context()
+			client := mcp.NewClient("client", "v0.0.1", nil)
+			mcpSession, err := client.Connect(ctx, mcp.NewCommandTransport(goplsCmd))
+			if err != nil {
+				t.Fatal(err)
+			}
+			if err := mcpSession.Close(); err != nil {
+				t.Errorf("closing MCP connection: %v", err)
+			}
+			logs, err := os.ReadFile(logFile)
+			if err != nil {
+				t.Fatal(err)
+			}
+			if test.want != "" && !bytes.Contains(logs, []byte(test.want)) {
+				t.Errorf("logs do not contain expected %q", test.want)
+			}
+			if test.dontWant != "" && bytes.Contains(logs, []byte(test.dontWant)) {
+				t.Errorf("logs contain unexpected %q", test.dontWant)
+			}
+			if t.Failed() {
+				t.Logf("Logs:\n%s", string(logs))
+			}
+		})
+	}
+}
+
 func TestMCPCommandHTTP(t *testing.T) {
 	testenv.NeedsExec(t)
 	tree := writeTree(t, `

gopls/internal/cmd/usage/mcp.hlp

@@ -11,3 +11,7 @@ Examples:
   $ gopls mcp  //start over stdio
   -listen=string
     	the address on which to run the mcp server
+  -logfile=string
+    	filename to log to; if unset, logs to stderr
+  -rpc.trace
+    	print MCP rpc traces; cannot be used with -listen

gopls/internal/mcp/mcp.go

@@ -8,6 +8,7 @@ import (
 	"context"
 	_ "embed"
 	"fmt"
+	"io"
 	"log"
 	"net"
 	"net/http"
@@ -67,8 +68,12 @@ func Serve(ctx context.Context, address string, eventChan <-chan lsprpc.SessionE
 }
 
 // StartStdIO starts an MCP server over stdio.
-func StartStdIO(ctx context.Context, session *cache.Session, server protocol.Server) error {
-	t := mcp.NewLoggingTransport(mcp.NewStdioTransport(), os.Stderr)
+func StartStdIO(ctx context.Context, session *cache.Session, server protocol.Server, rpcLog io.Writer) error {
+	transport := mcp.NewStdioTransport()
+	var t mcp.Transport = transport
+	if rpcLog != nil {
+		t = mcp.NewLoggingTransport(transport, rpcLog)
+	}
 	s := newServer(session, server)
 	return s.Run(ctx, t)
 }