gopls/internal/mcp/context.go - tools.git - Git at Google

 // Copyright 2025 The Go Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.

 package mcp

 // This file defines the "context" operation, which returns a summary of the
 // specified package.

 import (
 	"bytes"
 	"context"
 	"fmt"
 	"go/ast"
 	"go/token"
 	"slices"
 	"strings"

 	"golang.org/x/tools/gopls/internal/cache"
 	"golang.org/x/tools/gopls/internal/cache/metadata"
 	"golang.org/x/tools/gopls/internal/cache/parsego"
 	"golang.org/x/tools/gopls/internal/file"
 	"golang.org/x/tools/gopls/internal/golang"
 	"golang.org/x/tools/gopls/internal/protocol"
 	"golang.org/x/tools/gopls/internal/util/astutil"
 	"golang.org/x/tools/internal/analysisinternal"
 	"golang.org/x/tools/internal/mcp"
 )

 type ContextParams struct {
 	File string `json:"file"`
 }

 func (h *handler) contextTool() *mcp.ServerTool {
 	return mcp.NewServerTool(
 		"go_context",
 		"Provide context for a region within a Go file",
 		h.contextHandler,
 		mcp.Input(
 			mcp.Property("file", mcp.Description("the absolute path to the file")),
 		),
 	)
 }

 func (h *handler) contextHandler(ctx context.Context, _ *mcp.ServerSession, params *mcp.CallToolParamsFor[ContextParams]) (*mcp.CallToolResultFor[any], error) {
 	fh, snapshot, release, err := h.fileOf(ctx, params.Arguments.File)
 	if err != nil {
 		return nil, err
 	}
 	defer release()

 	// TODO(hxjiang): support context for GoMod.
 	if snapshot.FileKind(fh) != file.Go {
 		return nil, fmt.Errorf("can't provide context for non-Go file")
 	}

 	pkg, pgf, err := golang.NarrowestPackageForFile(ctx, snapshot, fh.URI())
 	if err != nil {
 		return nil, err
 	}

 	var result strings.Builder

 	fmt.Fprintf(&result, "Current package %q (package %s):\n\n", pkg.Metadata().PkgPath, pkg.Metadata().Name)
 	// Write context of the current file.
 	{
 		fmt.Fprintf(&result, "%s (current file):\n", pgf.URI.Base())
 		result.WriteString("```go\n")
 		if err := writeFileSummary(ctx, snapshot, pgf.URI, &result, false, nil); err != nil {
 			return nil, err
 		}
 		result.WriteString("```\n\n")
 	}

 	// Write context of the rest of the files in the current package.
 	{
 		for _, file := range pkg.CompiledGoFiles() {
 			if file.URI == pgf.URI {
 				continue
 			}

 			fmt.Fprintf(&result, "%s:\n", file.URI.Base())
 			result.WriteString("```go\n")
 			if err := writeFileSummary(ctx, snapshot, file.URI, &result, false, nil); err != nil {
 				return nil, err
 			}
 			result.WriteString("```\n\n")
 		}
 	}

 	// Write dependencies context of current file.
 	if len(pgf.File.Imports) > 0 {
 		// Write import decls of the current file.
 		{
 			fmt.Fprintf(&result, "Current file %q contains this import declaration:\n", pgf.URI.Base())
 			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 {
 				genDecl, ok := decl.(*ast.GenDecl)
 				if !ok || genDecl.Tok != token.IMPORT {
 					continue
 				}

 				text, err := pgf.NodeText(genDecl)
 				if err != nil {
 					return nil, err
 				}

 				result.Write(text)
 				result.WriteString("\n")
 			}
 			result.WriteString("```\n\n")
 		}

 		var toSummarize []*ast.ImportSpec
 		for _, spec := range pgf.File.Imports {
 			// Skip the standard library to reduce token usage, operating on
 			// the assumption that the LLM is already familiar with its
 			// symbols and documentation.
 			if analysisinternal.IsStdPackage(spec.Path.Value) {
 				continue
 			}
 			toSummarize = append(toSummarize, spec)
 		}

 		// Write summaries from imported packages.
 		if len(toSummarize) > 0 {
 			result.WriteString("The imported packages declare the following symbols:\n\n")
 			for _, spec := range toSummarize {
 				path := metadata.UnquoteImportPath(spec)
 				id := pkg.Metadata().DepsByImpPath[path]
 				if id == "" {
 					continue // ignore error
 				}
 				md := snapshot.Metadata(id)
 				if md == nil {
 					continue // ignore error
 				}
 				if summary := summarizePackage(ctx, snapshot, md); summary != "" {
 					result.WriteString(summary)
 				}
 			}
 		}
 	}

 	return textResult(result.String()), nil
 }

 func summarizePackage(ctx context.Context, snapshot *cache.Snapshot, md *metadata.Package) string {
 	var buf strings.Builder
 	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("```go\n")
 		if err := writeFileSummary(ctx, snapshot, f, &buf, true, nil); err != nil {
 			return "" // ignore error
 		}
 		buf.WriteString("```\n\n")
 	}
 	return buf.String()
 }

 // writeFileSummary writes the file summary to the string builder based on
 // the input file URI.
 func writeFileSummary(ctx context.Context, snapshot *cache.Snapshot, f protocol.DocumentURI, out *strings.Builder, onlyExported bool, declsToSummarize map[string]bool) error {
 	fh, err := snapshot.ReadFile(ctx, f)
 	if err != nil {
 		return err
 	}
 	pgf, err := snapshot.ParseGo(ctx, fh, parsego.Full)
 	if err != nil {
 		return err
 	}

 	// If we're summarizing specific declarations, we don't need to copy the header.
 	if declsToSummarize == nil {
 		// Copy everything before the first non-import declaration:
 		// package decl, imports decl(s), and all comments (excluding copyright).
 		{
 			endPos := pgf.File.FileEnd

 		outerloop:
 			for _, decl := range pgf.File.Decls {
 				switch decl := decl.(type) {
 				case *ast.FuncDecl:
 					if decl.Doc != nil {
 						endPos = decl.Doc.Pos()
 					} else {
 						endPos = decl.Pos()
 					}
 					break outerloop
 				case *ast.GenDecl:
 					if decl.Tok == token.IMPORT {
 						continue
 					}
 					if decl.Doc != nil {
 						endPos = decl.Doc.Pos()
 					} else {
 						endPos = decl.Pos()
 					}
 					break outerloop
 				}
 			}

 			startPos := pgf.File.FileStart
 			if copyright := golang.CopyrightComment(pgf.File); copyright != nil {
 				startPos = copyright.End()
 			}

 			text, err := pgf.PosText(startPos, endPos)
 			if err != nil {
 				return err
 			}

 			out.Write(bytes.TrimSpace(text))
 			out.WriteString("\n\n")
 		}
 	}

 	// Write func decl and gen decl.
 	for _, decl := range pgf.File.Decls {
 		switch decl := decl.(type) {
 		case *ast.FuncDecl:
 			if declsToSummarize != nil {
 				if _, ok := declsToSummarize[decl.Name.Name]; !ok {
 					continue
 				}
 			}
 			if onlyExported {
 				if !decl.Name.IsExported() {
 					continue
 				}

 				if decl.Recv != nil && len(decl.Recv.List) > 0 {
 					_, rname, _ := astutil.UnpackRecv(decl.Recv.List[0].Type)
 					if !rname.IsExported() {
 						continue
 					}
 				}
 			}

 			// Write doc comment and func signature.
 			startPos := decl.Pos()
 			if decl.Doc != nil {
 				startPos = decl.Doc.Pos()
 			}

 			text, err := pgf.PosText(startPos, decl.Type.End())
 			if err != nil {
 				return err
 			}

 			out.Write(text)
 			out.WriteString("\n\n")

 		case *ast.GenDecl:
 			if decl.Tok == token.IMPORT {
 				continue
 			}

 			// If we are summarizing specific decls, check if any of them are in this GenDecl.
 			if declsToSummarize != nil {
 				found := false
 				for _, spec := range decl.Specs {
 					switch spec := spec.(type) {
 					case *ast.TypeSpec:
 						if _, ok := declsToSummarize[spec.Name.Name]; ok {
 							found = true
 						}
 					case *ast.ValueSpec:
 						for _, name := range spec.Names {
 							if _, ok := declsToSummarize[name.Name]; ok {
 								found = true
 							}
 						}
 					}
 				}
 				if !found {
 					continue
 				}
 			}

 			// Dump the entire GenDecl (exported or unexported)
 			// including doc comment without any filtering to the output.
 			if !onlyExported {
 				startPos := decl.Pos()
 				if decl.Doc != nil {
 					startPos = decl.Doc.Pos()
 				}
 				text, err := pgf.PosText(startPos, decl.End())
 				if err != nil {
 					return err
 				}

 				out.Write(text)
 				out.WriteString("\n")
 				continue
 			}

 			// Write only the GenDecl with exported identifier to the output.
 			var buf bytes.Buffer
 			if decl.Doc != nil {
 				text, err := pgf.NodeText(decl.Doc)
 				if err != nil {
 					return err
 				}
 				buf.Write(text)
 				buf.WriteString("\n")
 			}

 			buf.WriteString(decl.Tok.String() + " ")
 			if decl.Lparen.IsValid() {
 				buf.WriteString("(\n")
 			}

 			var anyExported bool
 			for _, spec := range decl.Specs {
 				// Captures the full byte range of the spec, including
 				// its associated doc comments and line comments.
 				// This range also covers any floating comments as these
 				// can be valuable for context. Like
 				// ```
 				// type foo struct { // floating comment.
 				// 		// floating comment.
 				//
 				// 		x int
 				// }
 				// ```
 				var startPos, endPos token.Pos

 				switch spec := spec.(type) {
 				case *ast.TypeSpec:
 					if declsToSummarize != nil {
 						if _, ok := declsToSummarize[spec.Name.Name]; !ok {
 							continue
 						}
 					}
 					// TODO(hxjiang): only keep the exported field of
 					// struct spec and exported method of interface spec.
 					if !spec.Name.IsExported() {
 						continue
 					}
 					anyExported = true

 					// Include preceding doc comment, if any.
 					if spec.Doc == nil {
 						startPos = spec.Pos()
 					} else {
 						startPos = spec.Doc.Pos()
 					}

 					// Include trailing line comment, if any.
 					if spec.Comment == nil {
 						endPos = spec.End()
 					} else {
 						endPos = spec.Comment.End()
 					}

 				case *ast.ValueSpec:
 					if declsToSummarize != nil {
 						found := false
 						for _, name := range spec.Names {
 							if _, ok := declsToSummarize[name.Name]; ok {
 								found = true
 							}
 						}
 						if !found {
 							continue
 						}
 					}
 					// TODO(hxjiang): only keep the exported identifier.
 					if !slices.ContainsFunc(spec.Names, (*ast.Ident).IsExported) {
 						continue
 					}
 					anyExported = true

 					if spec.Doc == nil {
 						startPos = spec.Pos()
 					} else {
 						startPos = spec.Doc.Pos()
 					}

 					if spec.Comment == nil {
 						endPos = spec.End()
 					} else {
 						endPos = spec.Comment.End()
 					}
 				}

 				indent, err := pgf.Indentation(startPos)
 				if err != nil {
 					return err
 				}

 				buf.WriteString(indent)

 				text, err := pgf.PosText(startPos, endPos)
 				if err != nil {
 					return err
 				}

 				buf.Write(text)
 				buf.WriteString("\n")
 			}

 			if decl.Lparen.IsValid() {
 				buf.WriteString(")\n")
 			}

 			// Only write the summary of the genDecl if there is
 			// any exported spec.
 			if anyExported {
 				out.Write(buf.Bytes())
 				out.WriteString("\n")
 			}
 		}
 	}
 	return nil
 }
	// Copyright 2025 The Go Authors. All rights reserved.
	// Use of this source code is governed by a BSD-style
	// license that can be found in the LICENSE file.

	package mcp

	// This file defines the "context" operation, which returns a summary of the
	// specified package.

	import (
	"bytes"
	"context"
	"fmt"
	"go/ast"
	"go/token"
	"slices"
	"strings"

	"golang.org/x/tools/gopls/internal/cache"
	"golang.org/x/tools/gopls/internal/cache/metadata"
	"golang.org/x/tools/gopls/internal/cache/parsego"
	"golang.org/x/tools/gopls/internal/file"
	"golang.org/x/tools/gopls/internal/golang"
	"golang.org/x/tools/gopls/internal/protocol"
	"golang.org/x/tools/gopls/internal/util/astutil"
	"golang.org/x/tools/internal/analysisinternal"
	"golang.org/x/tools/internal/mcp"
	)

	type ContextParams struct {
	File string `json:"file"`
	}

	func (h handler) contextTool() mcp.ServerTool {
	return mcp.NewServerTool(
	"go_context",
	"Provide context for a region within a Go file",
	h.contextHandler,
	mcp.Input(
	mcp.Property("file", mcp.Description("the absolute path to the file")),
	),
	)
	}

	func (h handler) contextHandler(ctx context.Context, _ mcp.ServerSession, params mcp.CallToolParamsFor[ContextParams]) (mcp.CallToolResultFor[any], error) {
	fh, snapshot, release, err := h.fileOf(ctx, params.Arguments.File)
	if err != nil {
	return nil, err
	}
	defer release()

	// TODO(hxjiang): support context for GoMod.
	if snapshot.FileKind(fh) != file.Go {
	return nil, fmt.Errorf("can't provide context for non-Go file")
	}

	pkg, pgf, err := golang.NarrowestPackageForFile(ctx, snapshot, fh.URI())
	if err != nil {
	return nil, err
	}

	var result strings.Builder

	fmt.Fprintf(&result, "Current package %q (package %s):\n\n", pkg.Metadata().PkgPath, pkg.Metadata().Name)
	// Write context of the current file.
	{
	fmt.Fprintf(&result, "%s (current file):\n", pgf.URI.Base())
	result.WriteString("```go\n")
	if err := writeFileSummary(ctx, snapshot, pgf.URI, &result, false, nil); err != nil {
	return nil, err
	}
	result.WriteString("```\n\n")
	}

	// Write context of the rest of the files in the current package.
	{
	for _, file := range pkg.CompiledGoFiles() {
	if file.URI == pgf.URI {
	continue
	}

	fmt.Fprintf(&result, "%s:\n", file.URI.Base())
	result.WriteString("```go\n")
	if err := writeFileSummary(ctx, snapshot, file.URI, &result, false, nil); err != nil {
	return nil, err
	}
	result.WriteString("```\n\n")
	}
	}

	// Write dependencies context of current file.
	if len(pgf.File.Imports) > 0 {
	// Write import decls of the current file.
	{
	fmt.Fprintf(&result, "Current file %q contains this import declaration:\n", pgf.URI.Base())
	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 {
	genDecl, ok := decl.(*ast.GenDecl)
	if !ok \|\| genDecl.Tok != token.IMPORT {
	continue
	}

	text, err := pgf.NodeText(genDecl)
	if err != nil {
	return nil, err
	}

	result.Write(text)
	result.WriteString("\n")
	}
	result.WriteString("```\n\n")
	}

	var toSummarize []*ast.ImportSpec
	for _, spec := range pgf.File.Imports {
	// Skip the standard library to reduce token usage, operating on
	// the assumption that the LLM is already familiar with its
	// symbols and documentation.
	if analysisinternal.IsStdPackage(spec.Path.Value) {
	continue
	}
	toSummarize = append(toSummarize, spec)
	}

	// Write summaries from imported packages.
	if len(toSummarize) > 0 {
	result.WriteString("The imported packages declare the following symbols:\n\n")
	for _, spec := range toSummarize {
	path := metadata.UnquoteImportPath(spec)
	id := pkg.Metadata().DepsByImpPath[path]
	if id == "" {
	continue // ignore error
	}
	md := snapshot.Metadata(id)
	if md == nil {
	continue // ignore error
	}
	if summary := summarizePackage(ctx, snapshot, md); summary != "" {
	result.WriteString(summary)
	}
	}
	}
	}

	return textResult(result.String()), nil
	}

	func summarizePackage(ctx context.Context, snapshot cache.Snapshot, md metadata.Package) string {
	var buf strings.Builder
	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("```go\n")
	if err := writeFileSummary(ctx, snapshot, f, &buf, true, nil); err != nil {
	return "" // ignore error
	}
	buf.WriteString("```\n\n")
	}
	return buf.String()
	}

	// writeFileSummary writes the file summary to the string builder based on
	// the input file URI.
	func writeFileSummary(ctx context.Context, snapshot cache.Snapshot, f protocol.DocumentURI, out strings.Builder, onlyExported bool, declsToSummarize map[string]bool) error {
	fh, err := snapshot.ReadFile(ctx, f)
	if err != nil {
	return err
	}
	pgf, err := snapshot.ParseGo(ctx, fh, parsego.Full)
	if err != nil {
	return err
	}

	// If we're summarizing specific declarations, we don't need to copy the header.
	if declsToSummarize == nil {
	// Copy everything before the first non-import declaration:
	// package decl, imports decl(s), and all comments (excluding copyright).
	{
	endPos := pgf.File.FileEnd

	outerloop:
	for _, decl := range pgf.File.Decls {
	switch decl := decl.(type) {
	case *ast.FuncDecl:
	if decl.Doc != nil {
	endPos = decl.Doc.Pos()
	} else {
	endPos = decl.Pos()
	}
	break outerloop
	case *ast.GenDecl:
	if decl.Tok == token.IMPORT {
	continue
	}
	if decl.Doc != nil {
	endPos = decl.Doc.Pos()
	} else {
	endPos = decl.Pos()
	}
	break outerloop
	}
	}

	startPos := pgf.File.FileStart
	if copyright := golang.CopyrightComment(pgf.File); copyright != nil {
	startPos = copyright.End()
	}

	text, err := pgf.PosText(startPos, endPos)
	if err != nil {
	return err
	}

	out.Write(bytes.TrimSpace(text))
	out.WriteString("\n\n")
	}
	}

	// Write func decl and gen decl.
	for _, decl := range pgf.File.Decls {
	switch decl := decl.(type) {
	case *ast.FuncDecl:
	if declsToSummarize != nil {
	if _, ok := declsToSummarize[decl.Name.Name]; !ok {
	continue
	}
	}
	if onlyExported {
	if !decl.Name.IsExported() {
	continue
	}

	if decl.Recv != nil && len(decl.Recv.List) > 0 {
	_, rname, _ := astutil.UnpackRecv(decl.Recv.List[0].Type)
	if !rname.IsExported() {
	continue
	}
	}
	}

	// Write doc comment and func signature.
	startPos := decl.Pos()
	if decl.Doc != nil {
	startPos = decl.Doc.Pos()
	}

	text, err := pgf.PosText(startPos, decl.Type.End())
	if err != nil {
	return err
	}

	out.Write(text)
	out.WriteString("\n\n")

	case *ast.GenDecl:
	if decl.Tok == token.IMPORT {
	continue
	}

	// If we are summarizing specific decls, check if any of them are in this GenDecl.
	if declsToSummarize != nil {
	found := false
	for _, spec := range decl.Specs {
	switch spec := spec.(type) {
	case *ast.TypeSpec:
	if _, ok := declsToSummarize[spec.Name.Name]; ok {
	found = true
	}
	case *ast.ValueSpec:
	for _, name := range spec.Names {
	if _, ok := declsToSummarize[name.Name]; ok {
	found = true
	}
	}
	}
	}
	if !found {
	continue
	}
	}

	// Dump the entire GenDecl (exported or unexported)
	// including doc comment without any filtering to the output.
	if !onlyExported {
	startPos := decl.Pos()
	if decl.Doc != nil {
	startPos = decl.Doc.Pos()
	}
	text, err := pgf.PosText(startPos, decl.End())
	if err != nil {
	return err
	}

	out.Write(text)
	out.WriteString("\n")
	continue
	}

	// Write only the GenDecl with exported identifier to the output.
	var buf bytes.Buffer
	if decl.Doc != nil {
	text, err := pgf.NodeText(decl.Doc)
	if err != nil {
	return err
	}
	buf.Write(text)
	buf.WriteString("\n")
	}

	buf.WriteString(decl.Tok.String() + " ")
	if decl.Lparen.IsValid() {
	buf.WriteString("(\n")
	}

	var anyExported bool
	for _, spec := range decl.Specs {
	// Captures the full byte range of the spec, including
	// its associated doc comments and line comments.
	// This range also covers any floating comments as these
	// can be valuable for context. Like
	// ```
	// type foo struct { // floating comment.
	// // floating comment.
	//
	// x int
	// }
	// ```
	var startPos, endPos token.Pos

	switch spec := spec.(type) {
	case *ast.TypeSpec:
	if declsToSummarize != nil {
	if _, ok := declsToSummarize[spec.Name.Name]; !ok {
	continue
	}
	}
	// TODO(hxjiang): only keep the exported field of
	// struct spec and exported method of interface spec.
	if !spec.Name.IsExported() {
	continue
	}
	anyExported = true

	// Include preceding doc comment, if any.
	if spec.Doc == nil {
	startPos = spec.Pos()
	} else {
	startPos = spec.Doc.Pos()
	}

	// Include trailing line comment, if any.
	if spec.Comment == nil {
	endPos = spec.End()
	} else {
	endPos = spec.Comment.End()
	}

	case *ast.ValueSpec:
	if declsToSummarize != nil {
	found := false
	for _, name := range spec.Names {
	if _, ok := declsToSummarize[name.Name]; ok {
	found = true
	}
	}
	if !found {
	continue
	}
	}
	// TODO(hxjiang): only keep the exported identifier.
	if !slices.ContainsFunc(spec.Names, (*ast.Ident).IsExported) {
	continue
	}
	anyExported = true

	if spec.Doc == nil {
	startPos = spec.Pos()
	} else {
	startPos = spec.Doc.Pos()
	}

	if spec.Comment == nil {
	endPos = spec.End()
	} else {
	endPos = spec.Comment.End()
	}
	}

	indent, err := pgf.Indentation(startPos)
	if err != nil {
	return err
	}

	buf.WriteString(indent)

	text, err := pgf.PosText(startPos, endPos)
	if err != nil {
	return err
	}

	buf.Write(text)
	buf.WriteString("\n")
	}

	if decl.Lparen.IsValid() {
	buf.WriteString(")\n")
	}

	// Only write the summary of the genDecl if there is
	// any exported spec.
	if anyExported {
	out.Write(buf.Bytes())
	out.WriteString("\n")
	}
	}
	}
	return nil
	}