From 125ee726d6a0bd1e4f7321f70cc4bdb67389c345 Mon Sep 17 00:00:00 2001 From: Ian Gulliver Date: Fri, 27 Jun 2025 22:20:01 -0700 Subject: [PATCH] Add context support to find_comments and position ranges to search_replace --- main.go | 6 ++++- tool_find_comments.go | 55 +++++++++++++++++++++++++++++++++++------- tool_search_replace.go | 44 ++++++++++++++++++++++++++++++--- 3 files changed, 92 insertions(+), 13 deletions(-) diff --git a/main.go b/main.go index d435db0..c7dcd47 100644 --- a/main.go +++ b/main.go @@ -169,6 +169,9 @@ func main() { mcp.WithString("filter", mcp.Description("Optional regex to filter comments (applies to 'todo' and 'all' types)"), ), + mcp.WithBool("include_context", + mcp.Description("Include surrounding lines of code as context (default: false)"), + ), ) mcpServer.AddTool(findCommentsTool, findCommentsHandler) @@ -682,8 +685,9 @@ func findCommentsHandler(ctx context.Context, request mcp.CallToolRequest) (*mcp dir := request.GetString("dir", "./") commentType := request.GetString("type", "all") filter := request.GetString("filter", "") + includeContext := request.GetBool("include_context", false) - comments, err := findComments(dir, commentType, filter) + comments, err := findComments(dir, commentType, filter, includeContext) if err != nil { return mcp.NewToolResultError(fmt.Sprintf("failed to find comments: %v", err)), nil } diff --git a/tool_find_comments.go b/tool_find_comments.go index 112429f..46dd0bb 100644 --- a/tool_find_comments.go +++ b/tool_find_comments.go @@ -4,6 +4,7 @@ import ( "go/ast" "go/token" "regexp" + "strings" ) // Comment analysis types @@ -18,9 +19,29 @@ type CommentItem struct { Comment string `json:"comment,omitempty"` Type string `json:"type"` Position Position `json:"position"` + Context []string `json:"context,omitempty"` } -func findComments(dir string, commentType string, filter string) ([]CommentInfo, error) { +func getContext(src []byte, pos token.Position, contextLines int) []string { + lines := strings.Split(string(src), "\n") + startLine := pos.Line - contextLines - 1 + endLine := pos.Line + contextLines - 1 + + if startLine < 0 { + startLine = 0 + } + if endLine >= len(lines) { + endLine = len(lines) - 1 + } + + var context []string + for i := startLine; i <= endLine; i++ { + context = append(context, lines[i]) + } + return context +} + +func findComments(dir string, commentType string, filter string, includeContext bool) ([]CommentInfo, error) { var comments []CommentInfo err := walkGoFiles(dir, func(path string, src []byte, file *ast.File, fset *token.FileSet) error { @@ -47,11 +68,15 @@ func findComments(dir string, commentType string, filter string) ([]CommentInfo, // If no filter or filter matches, include the comment if filterRegex == nil || filterRegex.MatchString(c.Text) { pos := fset.Position(c.Pos()) - info.TODOs = append(info.TODOs, CommentItem{ + item := CommentItem{ Comment: c.Text, Type: "comment", Position: newPosition(pos), - }) + } + if includeContext { + item.Context = getContext(src, pos, 3) + } + info.TODOs = append(info.TODOs, item) } } } @@ -64,11 +89,15 @@ func findComments(dir string, commentType string, filter string) ([]CommentInfo, case *ast.FuncDecl: if ast.IsExported(x.Name.Name) && x.Doc == nil { pos := fset.Position(x.Pos()) - info.Undocumented = append(info.Undocumented, CommentItem{ + item := CommentItem{ Name: x.Name.Name, Type: "function", Position: newPosition(pos), - }) + } + if includeContext { + item.Context = getContext(src, pos, 3) + } + info.Undocumented = append(info.Undocumented, item) } case *ast.GenDecl: for _, spec := range x.Specs { @@ -76,21 +105,29 @@ func findComments(dir string, commentType string, filter string) ([]CommentInfo, case *ast.TypeSpec: if ast.IsExported(s.Name.Name) && x.Doc == nil && s.Doc == nil { pos := fset.Position(s.Pos()) - info.Undocumented = append(info.Undocumented, CommentItem{ + item := CommentItem{ Name: s.Name.Name, Type: "type", Position: newPosition(pos), - }) + } + if includeContext { + item.Context = getContext(src, pos, 3) + } + info.Undocumented = append(info.Undocumented, item) } case *ast.ValueSpec: for _, name := range s.Names { if ast.IsExported(name.Name) && x.Doc == nil && s.Doc == nil { pos := fset.Position(name.Pos()) - info.Undocumented = append(info.Undocumented, CommentItem{ + item := CommentItem{ Name: name.Name, Type: "value", Position: newPosition(pos), - }) + } + if includeContext { + item.Context = getContext(src, pos, 3) + } + info.Undocumented = append(info.Undocumented, item) } } } diff --git a/tool_search_replace.go b/tool_search_replace.go index 7ee49d6..b38420c 100644 --- a/tool_search_replace.go +++ b/tool_search_replace.go @@ -27,6 +27,12 @@ type SearchMatch struct { Column int `json:"column"` Text string `json:"text"` Context string `json:"context,omitempty"` + StartLine int `json:"start_line"` + StartCol int `json:"start_col"` + EndLine int `json:"end_line"` + EndCol int `json:"end_col"` + StartByte int `json:"start_byte"` + EndByte int `json:"end_byte"` } func searchReplace(paths []string, pattern string, replacement *string, useRegex, caseInsensitive bool, includeContext bool, beforePattern, afterPattern string) (*SearchReplaceResult, error) { @@ -247,10 +253,42 @@ func processFile(path string, searchFunc func(string) [][]int, replaceFunc func( } column := match[0] - lineStart + 1 + // Calculate end position + endLineNum := lineNum + endColumn := column + len(content[match[0]:match[1]]) + + // Check if match spans multiple lines + matchText := content[match[0]:match[1]] + newlineCount := strings.Count(matchText, "\n") + if newlineCount > 0 { + // Find the end line + for i := lineNum; i < len(lineStarts); i++ { + if match[1] <= lineStarts[i] { + endLineNum = i + break + } else if i == len(lineStarts)-1 { + endLineNum = len(lines) + } + } + + // Calculate end column on the last line + endLineStart := 0 + if endLineNum > 0 && endLineNum <= len(lineStarts) { + endLineStart = lineStarts[endLineNum-1] + } + endColumn = match[1] - endLineStart + 1 + } + searchMatch := SearchMatch{ - Line: lineNum, - Column: column, - Text: content[match[0]:match[1]], + Line: lineNum, + Column: column, + Text: matchText, + StartLine: lineNum, + StartCol: column, + EndLine: endLineNum, + EndCol: endColumn, + StartByte: match[0], + EndByte: match[1], } if includeContext && lineNum > 0 && lineNum <= len(lines) {