fix(ast_query): reliable first-page pagination cursor

The page-1 walk was capped at max_results via SkipAll, but buildASTCursorFooter
derives the remaining count from the total number of matches seen. Capping
truncated that total, so on page 1 the cursor fired only when the last parsed
file happened to overshoot the cap — nondeterministic — and reported a wrong
remaining count when it did. Same failure class as the file_search first-page
cursor bug.

Always walk the full tree (pages 2+ already did). The common case where matches
fit in one page already visited everything (SkipAll never triggered), and parse
results are LRU-cached, so the extra cost falls only on multi-page queries and is
mostly file I/O.

Adds a strict regression test asserting page 1 emits a cursor with an accurate
remaining count; the existing pagination test silently skipped when no cursor.
This commit is contained in:
2026-06-20 01:44:47 +01:00
parent 62971c9f8c
commit 5f0bb90e28
2 changed files with 37 additions and 10 deletions
+30
View File
@@ -215,6 +215,36 @@ func TestASTQuerySkipsDependencyDirs(t *testing.T) {
// ---- Feature 3 (ast_query): pagination cursor ----
// TestASTQueryFirstPageCursor is a strict regression guard: page 1 must emit a
// cursor with an accurate remaining count when more matches exist. Previously
// the page-1 walk was capped at max_results, so the total was truncated and the
// cursor fired nondeterministically (and with a wrong remaining count).
func TestASTQueryFirstPageCursor(t *testing.T) {
srv, tmpDir := newFeaturesServer(t)
ctx := context.Background()
// Fixture has 5 functions (Alpha/Beta/Gamma/Delta/Epsilon); cap at 2.
req := mcp.CallToolRequest{}
req.Params.Arguments = map[string]interface{}{
"pattern": "func $NAME() $RET",
"language": "go",
"paths": []interface{}{tmpDir},
"max_results": float64(2),
}
res, err := srv.handleASTQuery(ctx, req)
if err != nil {
t.Fatalf("page1 error: %v", err)
}
text := res.Content[0].(mcp.TextContent).Text
if !strings.Contains(text, "[cursor:") {
t.Fatalf("page1 with 5 matches capped at 2 must emit a cursor, got:\n%s", text)
}
// 5 total - 2 shown = 3 remaining.
if !strings.Contains(text, "remaining: 3") {
t.Errorf("expected accurate 'remaining: 3', got:\n%s", text)
}
}
func TestASTQueryPaginationCursor(t *testing.T) {
srv, tmpDir := newFeaturesServer(t)
ctx := context.Background()
+7 -10
View File
@@ -109,12 +109,13 @@ func (s *Server) runASTQueryWalk(ctx context.Context, p *astQueryParams, exts []
},
}
// Collect limit for early-exit: when paginating we need all results first.
collectLimit := p.maxResults
if p.offset > 0 {
collectLimit = 0
}
// A full walk is required on every page: the cursor's "remaining" count is
// derived from the total match count (see buildASTCursorFooter), so the page-1
// response can only emit a correct cursor if it has seen every match. Capping
// page 1 at max_results made the cursor fire nondeterministically (only when
// the last parsed file overshot the cap) with a wrong remaining count. Parse
// results are LRU-cached, so re-walking on later pages is mostly file I/O, and
// matches below max_results already visit the whole tree regardless.
var allResults []query.MatchResult
for _, searchPath := range p.paths {
if !s.cfg.IsPathAllowed(searchPath) {
@@ -162,10 +163,6 @@ func (s *Server) runASTQueryWalk(ctx context.Context, p *astQueryParams, exts []
return nil
}
allResults = append(allResults, matches...)
if collectLimit > 0 && len(allResults) >= collectLimit {
return filepath.SkipAll
}
return nil
})
if walkErr != nil {