gopls/internal/mcp: tune tools for editing

Make some changes to the gopls MCP server a result of experimentation
with our latest set of tools:

- Add instructions in instructions.md, and embed them in the MCP server.
- Prefer the 'tuned' version of references, context (outline), and
  diagnostics. Leave the original versions for now, for further
  experimentation. In some cases, refactor to share more code with the
  originals.
- Update 'go_outline' to just accept an array of package paths. This is how
  Gemini wanted to use it.
- More description for the 'go_search' tool.
- Accept an optional array of files for workspace_diagnostics, rather
  than a single file.
- Fix empty output when there were no diagnostics.
- Use fenced code blocks in more places.
- Allow configuring the set of tools in tests.
- Other minor test improvements.

For golang/go#73580

Change-Id: I47baf7f3c06904f0d016ac2d9bd5be68c5178b34
Reviewed-on: https://go-review.googlesource.com/c/tools/+/686755
Reviewed-by: Madeline Kalil <[email protected]>
Reviewed-by: Alan Donovan <[email protected]>
LUCI-TryBot-Result: Go LUCI <[email protected]>
Auto-Submit: Robert Findley <[email protected]>

Author: findleyr

gopls/internal/cache/metadata/graph.go

@@ -5,8 +5,10 @@
 package metadata
 
 import (
+	"cmp"
 	"iter"
 	"maps"
+	"slices"
 	"sort"
 	"strings"
 
@@ -23,6 +25,14 @@ type Graph struct {
 	// ImportedBy maps package IDs to the list of packages that import them.
 	ImportedBy map[PackageID][]PackageID
 
+	// ByPackagePath maps package by their package path to their package ID.
+	// Non-test packages appear before test packages, and within each of those
+	// categories, packages with fewer CompiledGoFiles appear first.
+	//
+	// TODO(rfindley): there's no reason for ImportedBy and IDs to not also hold
+	// pointers, rather than IDs, and pointers are more convenient.
+	ByPackagePath map[PackagePath][]*Package
+
 	// IDs maps file URIs to package IDs, sorted by (!valid, cli, packageID).
 	// A single file may belong to multiple packages due to tests packages.
 	//
@@ -84,10 +94,12 @@ func (g *Graph) Update(updates map[PackageID]*Package) *Graph {
 func newGraph(pkgs map[PackageID]*Package) *Graph {
 	// Build the import graph.
 	importedBy := make(map[PackageID][]PackageID)
+	byPackagePath := make(map[PackagePath][]*Package)
 	for id, mp := range pkgs {
 		for _, depID := range mp.DepsByPkgPath {
 			importedBy[depID] = append(importedBy[depID], id)
 		}
+		byPackagePath[mp.PkgPath] = append(byPackagePath[mp.PkgPath], mp)
 	}
 
 	// Collect file associations.
@@ -139,10 +151,26 @@ func newGraph(pkgs map[PackageID]*Package) *Graph {
 		}
 	}
 
+	for _, mps := range byPackagePath {
+		slices.SortFunc(mps, func(a, b *Package) int {
+			if (a.ForTest == "") != (b.ForTest == "") {
+				if a.ForTest == "" {
+					return -1
+				}
+				return 1
+			}
+			if c := cmp.Compare(len(a.CompiledGoFiles), len(b.CompiledGoFiles)); c != 0 {
+				return c
+			}
+			return cmp.Compare(a.ID, b.ID)
+		})
+	}
+
 	return &Graph{
-		Packages:   pkgs,
-		ImportedBy: importedBy,
-		IDs:        uriIDs,
+		Packages:      pkgs,
+		ImportedBy:    importedBy,
+		ByPackagePath: byPackagePath,
+		IDs:           uriIDs,
 	}
 }
 

gopls/internal/cache/snapshot.go

@@ -38,6 +38,7 @@ import (
 	"golang.org/x/tools/gopls/internal/util/bug"
 	"golang.org/x/tools/gopls/internal/util/constraints"
 	"golang.org/x/tools/gopls/internal/util/immutable"
+	"golang.org/x/tools/gopls/internal/util/moremaps"
 	"golang.org/x/tools/gopls/internal/util/pathutil"
 	"golang.org/x/tools/gopls/internal/util/persistent"
 	"golang.org/x/tools/gopls/internal/vulncheck"
@@ -1001,19 +1002,13 @@ func (s *Snapshot) IsWorkspacePackage(id PackageID) bool {
 // It may also contain ad-hoc packages for standalone files.
 // It includes all test variants.
 //
-// TODO(rfindley): Replace this with s.MetadataGraph().
+// TODO(rfindley): Replace usage of function this with s.LoadMetadataGraph().
 func (s *Snapshot) AllMetadata(ctx context.Context) ([]*metadata.Package, error) {
-	if err := s.awaitLoaded(ctx); err != nil {
+	g, err := s.LoadMetadataGraph(ctx)
+	if err != nil {
 		return nil, err
 	}
-
-	g := s.MetadataGraph()
-
-	meta := make([]*metadata.Package, 0, len(g.Packages))
-	for _, mp := range g.Packages {
-		meta = append(meta, mp)
-	}
-	return meta, nil
+	return moremaps.ValueSlice(g.Packages), nil
 }
 
 // GoModForFile returns the URI of the go.mod file for the given URI.
@@ -1171,6 +1166,14 @@ func (s *Snapshot) MetadataGraph() *metadata.Graph {
 	return s.meta
 }
 
+// LoadMetadataGraph is like [Snapshot.MetadataGraph], but awaits snapshot reloading.
+func (s *Snapshot) LoadMetadataGraph(ctx context.Context) (*metadata.Graph, error) {
+	if err := s.awaitLoaded(ctx); err != nil {
+		return nil, err
+	}
+	return s.MetadataGraph(), nil
+}
+
 // InitializationError returns the last error from initialization.
 func (s *Snapshot) InitializationError() *InitializationError {
 	s.mu.Lock()

gopls/internal/cmd/mcp_test.go

@@ -57,7 +57,7 @@ const B = 2
 	}()
 	var (
 		tool = "go_diagnostics"
-		args = map[string]any{"file": filepath.Join(tree, "a.go")}
+		args = map[string]any{"files": []string{filepath.Join(tree, "a.go")}}
 	)
 	// On the first diagnostics call, there should be no diagnostics.
 	{
@@ -157,15 +157,15 @@ func MyFun() {}
 	}()
 
 	var (
-		tool = "go_context"
+		tool = "go_file_metadata"
 		args = map[string]any{"file": filepath.Join(tree, "a.go")}
 	)
 	res, err := mcpSession.CallTool(ctx, &mcp.CallToolParams{Name: tool, Arguments: args})
 	if err != nil {
 		t.Fatal(err)
 	}
 	got := resultText(t, res)
-	want := "The imported packages declare the following symbols"
+	want := "example.com"
 	if !strings.Contains(got, want) {
 		t.Errorf("CallTool(%s, %v) = %+v, want containing %q", tool, args, got, want)
 	}

gopls/internal/mcp/context.go

@@ -93,7 +93,7 @@ func (h *handler) contextHandler(ctx context.Context, _ *mcp.ServerSession, para
 		// Write import decls of the current file.
 		{
 			fmt.Fprintf(&result, "Current file %q contains this import declaration:\n", pgf.URI.Base())
-			result.WriteString("--->\n")
+			result.WriteString("```go\n")
 			// Add all import decl to output including all floating comment by
 			// using GenDecl's start and end position.
 			for _, decl := range pgf.File.Decls {
@@ -110,7 +110,7 @@ func (h *handler) contextHandler(ctx context.Context, _ *mcp.ServerSession, para
 				result.Write(text)
 				result.WriteString("\n")
 			}
-			result.WriteString("<---\n\n")
+			result.WriteString("```\n\n")
 		}
 
 		var toSummarize []*ast.ImportSpec
@@ -137,7 +137,7 @@ func (h *handler) contextHandler(ctx context.Context, _ *mcp.ServerSession, para
 				if md == nil {
 					continue // ignore error
 				}
-				if summary := summarizePackage(ctx, snapshot, path, md); summary != "" {
+				if summary := summarizePackage(ctx, snapshot, md); summary != "" {
 					result.WriteString(summary)
 				}
 			}
@@ -147,16 +147,16 @@ func (h *handler) contextHandler(ctx context.Context, _ *mcp.ServerSession, para
 	return textResult(result.String()), nil
 }
 
-func summarizePackage(ctx context.Context, snapshot *cache.Snapshot, path metadata.ImportPath, md *metadata.Package) string {
+func summarizePackage(ctx context.Context, snapshot *cache.Snapshot, md *metadata.Package) string {
 	var buf strings.Builder
-	fmt.Fprintf(&buf, "%q (package %s)\n", path, md.Name)
+	fmt.Fprintf(&buf, "%q (package %s)\n", md.PkgPath, md.Name)
 	for _, f := range md.CompiledGoFiles {
 		fmt.Fprintf(&buf, "%s:\n", f.Base())
-		buf.WriteString("--->\n")
+		buf.WriteString("```go\n")
 		if err := writeFileSummary(ctx, snapshot, f, &buf, true); err != nil {
 			return "" // ignore error
 		}
-		buf.WriteString("<---\n\n")
+		buf.WriteString("```\n\n")
 	}
 	return buf.String()
 }

gopls/internal/mcp/diagnostics.go renamed to gopls/internal/mcp/file_diagnostics.go

@@ -27,76 +27,89 @@ type diagnosticsParams struct {
 	File string `json:"file"`
 }
 
-func (h *handler) diagnosticsTool() *mcp.ServerTool {
+func (h *handler) fileDiagnosticsTool() *mcp.ServerTool {
 	return mcp.NewServerTool(
-		"go_diagnostics",
+		"go_file_diagnostics",
 		"Provides diagnostics for a Go file",
-		h.diagnoseFileHandler,
+		h.fileDiagnosticsHandler,
 		mcp.Input(
 			mcp.Property("file", mcp.Description("the absolute path to the file to diagnose")),
 		),
 	)
 }
 
-func (h *handler) diagnoseFileHandler(ctx context.Context, _ *mcp.ServerSession, params *mcp.CallToolParamsFor[diagnosticsParams]) (*mcp.CallToolResultFor[any], error) {
+func (h *handler) fileDiagnosticsHandler(ctx context.Context, _ *mcp.ServerSession, params *mcp.CallToolParamsFor[diagnosticsParams]) (*mcp.CallToolResultFor[any], error) {
 	fh, snapshot, release, err := h.fileOf(ctx, params.Arguments.File)
 	if err != nil {
 		return nil, err
 	}
 	defer release()
 
-	diagnostics, err := golang.DiagnoseFile(ctx, snapshot, fh.URI())
+	diagnostics, fixes, err := h.diagnoseFile(ctx, snapshot, fh.URI())
 	if err != nil {
 		return nil, err
 	}
 
 	var builder strings.Builder
 	if len(diagnostics) == 0 {
-		builder.WriteString("No diagnostics")
-	} else {
-		// LSP [protocol.Diagnostic]s do not carry code edits directly.
-		// Instead, gopls provides associated [protocol.CodeAction]s with their
-		// diagnostics field populated.
-		// Ignore errors. It is still valuable to provide only the diagnostic
-		// without any text edits.
-		// TODO(hxjiang): support code actions that returns call back command.
-		actions, _ := h.lspServer.CodeAction(ctx, &protocol.CodeActionParams{
-			TextDocument: protocol.TextDocumentIdentifier{
-				URI: fh.URI(),
-			},
-			Context: protocol.CodeActionContext{
-				Only:        []protocol.CodeActionKind{protocol.QuickFix},
-				Diagnostics: cache.ToProtocolDiagnostics(diagnostics...),
-			},
-		})
-
-		type key struct {
-			Message string
-			Range   protocol.Range
-		}
+		return textResult("No diagnostics"), nil
+	}
 
-		fixes := make(map[key]*protocol.CodeAction)
-		for _, action := range actions {
-			for _, d := range action.Diagnostics {
-				k := key{d.Message, d.Range}
-				if alt, ok := fixes[k]; !ok || !alt.IsPreferred && action.IsPreferred {
-					fixes[k] = &action
-				}
-			}
-		}
+	if err := summarizeDiagnostics(ctx, snapshot, &builder, diagnostics, fixes); err != nil {
+		return nil, err
+	}
+
+	return textResult(builder.String()), nil
+}
+
+// diagnoseFile diagnoses a single file, including go/analysis and quick fixes.
+func (h *handler) diagnoseFile(ctx context.Context, snapshot *cache.Snapshot, uri protocol.DocumentURI) ([]*cache.Diagnostic, map[*cache.Diagnostic]*protocol.CodeAction, error) {
+	diagnostics, err := golang.DiagnoseFile(ctx, snapshot, uri)
+	if err != nil {
+		return nil, nil, err
+	}
+	if len(diagnostics) == 0 {
+		return nil, nil, nil
+	}
+
+	// LSP [protocol.Diagnostic]s do not carry code edits directly.
+	// Instead, gopls provides associated [protocol.CodeAction]s with their
+	// diagnostics field populated.
+	// Ignore errors. It is still valuable to provide only the diagnostic
+	// without any text edits.
+	// TODO(hxjiang): support code actions that returns call back command.
+	actions, _ := h.lspServer.CodeAction(ctx, &protocol.CodeActionParams{
+		TextDocument: protocol.TextDocumentIdentifier{
+			URI: uri,
+		},
+		Context: protocol.CodeActionContext{
+			Only:        []protocol.CodeActionKind{protocol.QuickFix},
+			Diagnostics: cache.ToProtocolDiagnostics(diagnostics...),
+		},
+	})
+
+	type key struct {
+		Message string
+		Range   protocol.Range
+	}
 
-		fixesByDiagnostic := make(map[*cache.Diagnostic]*protocol.CodeAction)
-		for _, d := range diagnostics {
-			if fix, ok := fixes[key{d.Message, d.Range}]; ok {
-				fixesByDiagnostic[d] = fix
+	actionMap := make(map[key]*protocol.CodeAction)
+	for _, action := range actions {
+		for _, d := range action.Diagnostics {
+			k := key{d.Message, d.Range}
+			if alt, ok := actionMap[k]; !ok || !alt.IsPreferred && action.IsPreferred {
+				actionMap[k] = &action
 			}
 		}
-		if err := summarizeDiagnostics(ctx, snapshot, &builder, diagnostics, fixesByDiagnostic); err != nil {
-			return nil, err
-		}
 	}
 
-	return textResult(builder.String()), nil
+	fixes := make(map[*cache.Diagnostic]*protocol.CodeAction)
+	for _, d := range diagnostics {
+		if fix, ok := actionMap[key{d.Message, d.Range}]; ok {
+			fixes[d] = fix
+		}
+	}
+	return diagnostics, fixes, nil
 }
 
 func summarizeDiagnostics(ctx context.Context, snapshot *cache.Snapshot, w io.Writer, diagnostics []*cache.Diagnostic, fixes map[*cache.Diagnostic]*protocol.CodeAction) error {

gopls/internal/mcp/instructions.md

@@ -0,0 +1,51 @@
+# The gopls MCP server
+
+These instructions describe how to efficiently work in the Go programming
+language using the gopls MCP server. They are intended to be provided as context
+for an interactive session using the gopls MCP tool: you can load this file
+directly into a session where the gopls MCP server is connected.
+
+## Detecting a Go workspace
+
+Use the `go_workspace` tool to learn about the Go workspace. These instructions
+apply whenever that tool indicates that the user is in a Go workspace.
+
+## Go Programming Guidelines
+
+These guidelines MUST be followed whenever working in a Go workspace. There
+are two workflows described below: the 'Read Workflow' must be followed when
+the user asks a question about a Go workspace. The 'Edit Workflow' must be
+followed when the user edits a Go workspace.
+
+You may re-do parts of each workflow as necessary to recover from errors.
+However, you cannot skip any steps.
+
+### Read Workflow
+
+1. **Search the workspace:** When the user asks about a symbol, use
+   `go_search` to search for the symbol in question. If you find no matches,
+   search for a substring of the user's referenced symbol. If `go_search`
+   fails, you may fall back to regular textual search.
+2. **Read files:** Read the relevant file(s). Use the `go_file_metadata` tool
+   to get package information for the file.
+3. **Understand packages:** If the user is asking about the use of one or more Go
+   package, use the `go_package_outline` command to summarize their API.
+
+### Editing Workflow
+
+1. **Read first:** Before making any edits, follow the Read Workflow to
+   understand the user's request.
+2. **Find references:** Before modifying the definition of any symbol, use the
+   `go_symbol_references` tool to find references to that identifier. These
+   references may need to be updated after editing the symbol. Read files
+   containing references to evaluate if any further edits are required.
+3. **Make edits:** Make the primary edit, as well as any edits to references.
+4. **Run diagnostics:** Every time, after making edits to one or more files,
+   you must call the `go_diagnostics` tool, passing the paths to the edited
+   files, to verify that the build is not broken. Apply edits to fix any
+   relevant diagnostics, and re-run the `go_diagnostics` tool to verify the
+   fixes. It is OK to ignore 'hint' or 'info' diagnostics if they are not
+   relevant.
+5. **Run tests** run `go test` for any packages that were edited. Invoke `go
+   test` with the package paths returned from `go_file_metadata`. Fix any test
+   failures.