mirror of
https://github.com/lukaszraczylo/filepuff-mcp.git
synced 2026-10-20 05:57:51 +00:00
V2/token optimization (#11)
* v2.0: token-optimization overhaul Additive (backward-compatible flags): - file_read: skeleton mode, strip (imports/license/block_comments), compact_line_numbers, 8-char etag with prefix-match compat - ast_query: format=verbose|compact|location, pagination cursor - file_search: cluster mode, pagination cursor - lsp_query (references): compact output Breaking (v2): - Preambles removed; opt-in verbose=true restores - edit_apply: response=count|diff|none, default count - ping tool removed - symbol_at/find_definition/find_references merged into lsp_query - Tool descriptions trimmed -83%, help moved to filepuff://help/<tool> - Batch file_read dedups by etag Protocol: - ResourceLink returned for file_read >64 KiB (force_inline override) - OnAfterInitialize hook reads capabilities.experimental.filepuff for session defaults (default_format, default_max_results, default_cluster, compact_refs, line_numbers, resource_link_threshold) * fix: drop --max-total-count from ripgrep args The flag does not exist in stable ripgrep (confirmed up to 15.1.0 -- "unrecognized flag --max-total-count, similar flags that are available: --max-count"). Every file_search call failed on hosts with stock rg. --max-count is per-file, not a drop-in replacement, so rely on the in-process truncation in parseOutput that was already the documented safety net.
This commit is contained in:
+198
-105
@@ -8,12 +8,193 @@ import (
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
"github.com/lukaszraczylo/mcp-filepuff/internal/cursor"
|
||||
"github.com/lukaszraczylo/mcp-filepuff/internal/parser"
|
||||
"github.com/lukaszraczylo/mcp-filepuff/internal/query"
|
||||
"github.com/lukaszraczylo/mcp-filepuff/pkg/protocol"
|
||||
"github.com/mark3labs/mcp-go/mcp"
|
||||
)
|
||||
|
||||
// astQueryParams holds resolved parameters for an AST query invocation.
|
||||
type astQueryParams struct {
|
||||
pattern string
|
||||
language string
|
||||
nameMatches string
|
||||
nameExact string
|
||||
kindIn []string
|
||||
paths []string
|
||||
maxResults int
|
||||
format string
|
||||
offset int
|
||||
queryHash string
|
||||
verbose bool
|
||||
}
|
||||
|
||||
// resolveASTQueryParams parses the request and resolves session-pref defaults.
|
||||
// Returns (params, errorResult, error); when errorResult is non-nil, the caller
|
||||
// should return it directly.
|
||||
func (s *Server) resolveASTQueryParams(request mcp.CallToolRequest) (*astQueryParams, *mcp.CallToolResult) {
|
||||
pattern, err := request.RequireString("pattern")
|
||||
if err != nil {
|
||||
return nil, mcp.NewToolResultError("pattern is required")
|
||||
}
|
||||
language, err := request.RequireString("language")
|
||||
if err != nil {
|
||||
return nil, mcp.NewToolResultError("language is required")
|
||||
}
|
||||
|
||||
p := &astQueryParams{
|
||||
pattern: pattern,
|
||||
language: language,
|
||||
nameMatches: request.GetString("name_matches", ""),
|
||||
nameExact: request.GetString("name_exact", ""),
|
||||
kindIn: request.GetStringSlice("kind_in", nil),
|
||||
paths: request.GetStringSlice("paths", nil),
|
||||
verbose: request.GetBool("verbose", false),
|
||||
}
|
||||
|
||||
sp := s.sessionPrefs.Load()
|
||||
var prefsMaxResults int
|
||||
var prefsFormat string
|
||||
if sp != nil {
|
||||
prefsMaxResults = sp.DefaultMaxResults
|
||||
prefsFormat = sp.ASTQueryFormat
|
||||
}
|
||||
p.maxResults = effectiveInt(request, "max_results", prefsMaxResults, 100)
|
||||
|
||||
p.format = request.GetString("format", "")
|
||||
if p.format == "" {
|
||||
if prefsFormat != "" {
|
||||
p.format = prefsFormat
|
||||
} else {
|
||||
p.format = "verbose"
|
||||
}
|
||||
}
|
||||
|
||||
p.queryHash = cursor.HashParams(map[string]string{
|
||||
"pattern": p.pattern,
|
||||
"language": p.language,
|
||||
"name_matches": p.nameMatches,
|
||||
"name_exact": p.nameExact,
|
||||
"kind_in": strings.Join(p.kindIn, ","),
|
||||
"paths": strings.Join(p.paths, ","),
|
||||
})
|
||||
|
||||
if cursorStr := request.GetString("cursor", ""); cursorStr != "" {
|
||||
off, hash, decErr := cursor.Decode(cursorStr)
|
||||
if decErr != nil {
|
||||
return nil, mcp.NewToolResultError(fmt.Sprintf("invalid cursor: %s", decErr))
|
||||
}
|
||||
if hash != p.queryHash {
|
||||
return nil, mcp.NewToolResultError("cursor is for a different query, re-run without cursor")
|
||||
}
|
||||
p.offset = off
|
||||
}
|
||||
|
||||
return p, nil
|
||||
}
|
||||
|
||||
// runASTQueryWalk walks the configured paths and collects matches.
|
||||
func (s *Server) runASTQueryWalk(ctx context.Context, p *astQueryParams, exts []string) []query.MatchResult {
|
||||
astQuery := &query.ASTQuery{
|
||||
Pattern: p.pattern,
|
||||
Language: p.language,
|
||||
Filters: query.QueryFilters{
|
||||
NameMatches: p.nameMatches,
|
||||
NameExact: p.nameExact,
|
||||
KindIn: p.kindIn,
|
||||
},
|
||||
}
|
||||
|
||||
// Collect limit for early-exit: when paginating we need all results first.
|
||||
collectLimit := p.maxResults
|
||||
if p.offset > 0 {
|
||||
collectLimit = 0
|
||||
}
|
||||
|
||||
var allResults []query.MatchResult
|
||||
for _, searchPath := range p.paths {
|
||||
if !s.cfg.IsPathAllowed(searchPath) {
|
||||
continue
|
||||
}
|
||||
|
||||
walkErr := filepath.Walk(searchPath, func(path string, info os.FileInfo, err error) error {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
default:
|
||||
}
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
if info.IsDir() {
|
||||
if strings.HasPrefix(info.Name(), ".") {
|
||||
return filepath.SkipDir
|
||||
}
|
||||
return nil
|
||||
}
|
||||
if !hasAnySuffix(path, exts) {
|
||||
return nil
|
||||
}
|
||||
|
||||
content, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
if int64(len(content)) > s.cfg.MaxFileSize {
|
||||
return nil
|
||||
}
|
||||
|
||||
result, err := s.parser.Parse(ctx, path, content)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
matches, err := s.matcher.Match(ctx, astQuery, result.Tree, content, path)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
allResults = append(allResults, matches...)
|
||||
|
||||
if collectLimit > 0 && len(allResults) >= collectLimit {
|
||||
return filepath.SkipAll
|
||||
}
|
||||
return nil
|
||||
})
|
||||
if walkErr != nil {
|
||||
s.logger.Warn("error walking path", "path", searchPath, "error", walkErr)
|
||||
}
|
||||
}
|
||||
return allResults
|
||||
}
|
||||
|
||||
// hasAnySuffix reports whether path ends with any of the given suffixes.
|
||||
func hasAnySuffix(path string, suffixes []string) bool {
|
||||
for _, ext := range suffixes {
|
||||
if strings.HasSuffix(path, ext) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// buildASTCursorFooter computes the cursor footer line for truncated results.
|
||||
func buildASTCursorFooter(total, offset, maxResults int, queryHash string) string {
|
||||
if offset > 0 && offset < total {
|
||||
totalAfterOffset := total - offset
|
||||
if maxResults > 0 && totalAfterOffset > maxResults {
|
||||
remaining := totalAfterOffset - maxResults
|
||||
nextOffset := offset + maxResults
|
||||
nextCursor := cursor.Encode(nextOffset, queryHash)
|
||||
return fmt.Sprintf("[cursor: %s, remaining: %d]", nextCursor, remaining)
|
||||
}
|
||||
} else if offset == 0 && maxResults > 0 && total > maxResults {
|
||||
remaining := total - maxResults
|
||||
nextCursor := cursor.Encode(maxResults, queryHash)
|
||||
return fmt.Sprintf("[cursor: %s, remaining: %d]", nextCursor, remaining)
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// handleASTQuery handles the ast_query tool.
|
||||
func (s *Server) handleASTQuery(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
||||
// Acquire semaphore to limit concurrent queries (prevents CPU exhaustion)
|
||||
@@ -24,121 +205,33 @@ func (s *Server) handleASTQuery(ctx context.Context, request mcp.CallToolRequest
|
||||
return mcp.NewToolResultError("request cancelled"), nil
|
||||
}
|
||||
|
||||
pattern, err := request.RequireString("pattern")
|
||||
if err != nil {
|
||||
return mcp.NewToolResultError("pattern is required"), nil
|
||||
p, errResult := s.resolveASTQueryParams(request)
|
||||
if errResult != nil {
|
||||
return errResult, nil
|
||||
}
|
||||
|
||||
language, err := request.RequireString("language")
|
||||
if err != nil {
|
||||
return mcp.NewToolResultError("language is required"), nil
|
||||
if len(p.paths) == 0 {
|
||||
p.paths = []string{s.cfg.WorkspaceRoot}
|
||||
}
|
||||
|
||||
// Build query
|
||||
astQuery := &query.ASTQuery{
|
||||
Pattern: pattern,
|
||||
Language: language,
|
||||
Filters: query.QueryFilters{
|
||||
NameMatches: request.GetString("name_matches", ""),
|
||||
NameExact: request.GetString("name_exact", ""),
|
||||
KindIn: request.GetStringSlice("kind_in", nil),
|
||||
},
|
||||
}
|
||||
|
||||
maxResults := request.GetInt("max_results", 100)
|
||||
paths := request.GetStringSlice("paths", nil)
|
||||
|
||||
// Default to workspace root if no paths specified
|
||||
if len(paths) == 0 {
|
||||
paths = []string{s.cfg.WorkspaceRoot}
|
||||
}
|
||||
|
||||
// Find files to search based on language
|
||||
exts := languageToExtensions(language)
|
||||
exts := languageToExtensions(p.language)
|
||||
if len(exts) == 0 {
|
||||
return mcp.NewToolResultError(fmt.Sprintf("unsupported language: %s (supported: go, typescript, javascript, python, c, cpp, html, vue, elixir, rust)", language)), nil
|
||||
return mcp.NewToolResultError(fmt.Sprintf("unsupported language: %s (supported: go, typescript, javascript, python, c, cpp, html, vue, elixir, rust)", p.language)), nil
|
||||
}
|
||||
|
||||
var allResults []query.MatchResult
|
||||
allResults := s.runASTQueryWalk(ctx, p, exts)
|
||||
cursorFooter := buildASTCursorFooter(len(allResults), p.offset, p.maxResults, p.queryHash)
|
||||
|
||||
// Walk through paths and find matching files
|
||||
for _, searchPath := range paths {
|
||||
// Validate path is within workspace
|
||||
if !s.cfg.IsPathAllowed(searchPath) {
|
||||
continue
|
||||
}
|
||||
|
||||
err := filepath.Walk(searchPath, func(path string, info os.FileInfo, err error) error {
|
||||
// Check for context cancellation
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
default:
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
return nil // Skip files with errors
|
||||
}
|
||||
|
||||
if info.IsDir() {
|
||||
// Skip hidden directories
|
||||
if strings.HasPrefix(info.Name(), ".") {
|
||||
return filepath.SkipDir
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Check file extension matches language
|
||||
matched := false
|
||||
for _, ext := range exts {
|
||||
if strings.HasSuffix(path, ext) {
|
||||
matched = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !matched {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Read and parse file
|
||||
content, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return nil // Skip unreadable files
|
||||
}
|
||||
|
||||
// Check file size
|
||||
if int64(len(content)) > s.cfg.MaxFileSize {
|
||||
return nil // Skip large files
|
||||
}
|
||||
|
||||
// Parse file
|
||||
result, err := s.parser.Parse(ctx, path, content)
|
||||
if err != nil {
|
||||
return nil // Skip unparseable files
|
||||
}
|
||||
|
||||
// Run query
|
||||
matches, err := s.matcher.Match(ctx, astQuery, result.Tree, content, path)
|
||||
if err != nil {
|
||||
return nil // Skip on error
|
||||
}
|
||||
|
||||
allResults = append(allResults, matches...)
|
||||
|
||||
// Stop if we have enough results
|
||||
if maxResults > 0 && len(allResults) >= maxResults {
|
||||
return filepath.SkipAll
|
||||
}
|
||||
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
s.logger.Warn("error walking path", "path", searchPath, "error", err)
|
||||
output := query.FormatResultsWithOptions(allResults, p.maxResults, p.format, p.offset, p.verbose)
|
||||
if cursorFooter != "" {
|
||||
// Replace the [remaining: N] placeholder emitted by FormatResultsWithOptions
|
||||
// with the full [cursor: ..., remaining: N] line.
|
||||
output = strings.ReplaceAll(output, fmt.Sprintf("[remaining: %d]\n", len(allResults)-p.offset-p.maxResults), cursorFooter+"\n")
|
||||
// Fallback: if placeholder wasn't present (e.g. format=location), append footer.
|
||||
if !strings.Contains(output, cursorFooter) {
|
||||
output = strings.TrimRight(output, "\n") + "\n" + cursorFooter + "\n"
|
||||
}
|
||||
}
|
||||
|
||||
// Format and return results
|
||||
output := query.FormatResults(allResults, maxResults)
|
||||
return mcp.NewToolResultText(output), nil
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user