From 833e898d201100951399295e2fc9204ad85c7889 Mon Sep 17 00:00:00 2001 From: Lukasz Raczylo Date: Fri, 19 Jun 2026 18:06:54 +0100 Subject: [PATCH] fix(server): skip oversized files in AST walk, fix line count Two AST-summary/query path fixes plus colocated QF1012 style fixups: - runASTQueryWalk now checks info.Size() > MaxFileSize before os.ReadFile, avoiding a full read of files it will discard. - generateASTSummary used len(splitLines(...)) for the '(N lines)' header, over-counting newline-terminated files by one (phantom trailing element). New countLines() helper counts via bytes.Count. Adds TestCountLines. --- internal/server/handlers_ast.go | 12 +++++++----- internal/server/handlers_file.go | 22 +++++++++++++++++++--- internal/server/server_test.go | 26 ++++++++++++++++++++++++++ 3 files changed, 52 insertions(+), 8 deletions(-) diff --git a/internal/server/handlers_ast.go b/internal/server/handlers_ast.go index c630b98..e1b1d85 100644 --- a/internal/server/handlers_ast.go +++ b/internal/server/handlers_ast.go @@ -137,11 +137,13 @@ func (s *Server) runASTQueryWalk(ctx context.Context, p *astQueryParams, exts [] return nil } - content, err := os.ReadFile(path) - if err != nil { + // Skip oversized files before reading them off disk. + if info.Size() > s.cfg.MaxFileSize { return nil } - if int64(len(content)) > s.cfg.MaxFileSize { + + content, err := os.ReadFile(path) + if err != nil { return nil } @@ -260,12 +262,12 @@ func (s *Server) generateASTSummary(ctx context.Context, path string, content [] } } - sb.WriteString(fmt.Sprintf("**%s** (%d lines, %s)\n\n", relPath, len(splitLines(string(content))), lang)) + fmt.Fprintf(&sb, "**%s** (%d lines, %s)\n\n", relPath, countLines(content), lang) sb.WriteString("Symbols:\n") for _, sym := range symbols { kindStr := symbolKindIcon(sym.Kind) - sb.WriteString(fmt.Sprintf(" %s %s L%d\n", kindStr, sym.Name, sym.Location.Line)) + fmt.Fprintf(&sb, " %s %s L%d\n", kindStr, sym.Name, sym.Location.Line) } return sb.String() diff --git a/internal/server/handlers_file.go b/internal/server/handlers_file.go index 7c19ed8..374f054 100644 --- a/internal/server/handlers_file.go +++ b/internal/server/handlers_file.go @@ -3,6 +3,7 @@ package server import ( "bufio" + "bytes" "context" "fmt" "os" @@ -184,7 +185,7 @@ func (s *Server) handleFileRead(ctx context.Context, request mcp.CallToolRequest } result, err := s.readOneFile(ctx, request, p) if err != nil { - output.WriteString(fmt.Sprintf("--- %s ---\n[error: %s]\n", p, errors.SanitizeError(err))) + fmt.Fprintf(&output, "--- %s ---\n[error: %s]\n", p, errors.SanitizeError(err)) continue } // Extract etag from result footer for dedup check. @@ -192,12 +193,12 @@ func (s *Server) handleFileRead(ctx context.Context, request mcp.CallToolRequest if etag != "" { if firstPath, seen := seenEtag[etag]; seen { // Duplicate content: emit pointer instead of full content. - output.WriteString(fmt.Sprintf("--- %s ---\n[duplicate of %s, etag: %s]\n", p, firstPath, etag)) + fmt.Fprintf(&output, "--- %s ---\n[duplicate of %s, etag: %s]\n", p, firstPath, etag) continue } seenEtag[etag] = p } - output.WriteString(fmt.Sprintf("--- %s ---\n%s", p, result)) + fmt.Fprintf(&output, "--- %s ---\n%s", p, result) } return mcp.NewToolResultText(output.String()), nil } @@ -571,6 +572,21 @@ func extractEtag(result string) string { // splitLines splits a string into lines. // For large files (> 1MB), uses bufio.Scanner which is more memory efficient. // For smaller files, uses simple string split which is faster. +// countLines returns the number of source lines in content. Unlike +// len(splitLines(...)), it does not over-count newline-terminated files +// (a trailing "\n" does not introduce a phantom final line) and avoids +// materializing the line slice. +func countLines(content []byte) int { + if len(content) == 0 { + return 0 + } + n := bytes.Count(content, []byte{'\n'}) + if content[len(content)-1] != '\n' { + n++ + } + return n +} + func splitLines(s string) []string { const largeSizeThreshold = 1024 * 1024 // 1MB diff --git a/internal/server/server_test.go b/internal/server/server_test.go index fa4a94c..1a8d435 100644 --- a/internal/server/server_test.go +++ b/internal/server/server_test.go @@ -415,6 +415,32 @@ func TestSplitLinesLargeFileNoTrailingNewline(t *testing.T) { } } +// TestCountLines verifies countLines reports source line counts without the +// phantom trailing line that len(splitLines(...)) produces for newline-terminated +// files. +func TestCountLines(t *testing.T) { + tests := []struct { + name string + content string + want int + }{ + {name: "empty", content: "", want: 0}, + {name: "single line no newline", content: "a", want: 1}, + {name: "single line trailing newline", content: "a\n", want: 1}, + {name: "two lines trailing newline", content: "a\nb\n", want: 2}, + {name: "two lines no trailing newline", content: "a\nb", want: 2}, + {name: "blank lines", content: "\n\n\n", want: 3}, + {name: "crlf trailing", content: "a\r\nb\r\n", want: 2}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := countLines([]byte(tt.content)); got != tt.want { + t.Errorf("countLines(%q) = %d, want %d", tt.content, got, tt.want) + } + }) + } +} + // TestSplitLinesLongLine verifies the scanner gracefully handles very long lines // (up to the 1MB buffer limit set in C-05 fix). func TestSplitLinesLongLine(t *testing.T) {