Add context support to find_comments and position ranges to search_replace

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
 	}

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)
 								}
 							}
 						}

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) {