From 5f0bb90e283168ef9b4be982e9e9d877b01dad5a Mon Sep 17 00:00:00 2001
From: Lukasz Raczylo <lukasz@raczylo.com>
Date: Sat, 20 Jun 2026 01:44:47 +0100
Subject: [PATCH] fix(ast_query): reliable first-page pagination cursor
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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.
---
 internal/server/features_test.go | 30 ++++++++++++++++++++++++++++++
 internal/server/handlers_ast.go  | 17 +++++++----------
 2 files changed, 37 insertions(+), 10 deletions(-)

diff --git a/internal/server/features_test.go b/internal/server/features_test.go
index 0248fa0..e5bb585 100644
--- a/internal/server/features_test.go
+++ b/internal/server/features_test.go
@@ -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()
diff --git a/internal/server/handlers_ast.go b/internal/server/handlers_ast.go
index f629454..f7a460f 100644
--- a/internal/server/handlers_ast.go
+++ b/internal/server/handlers_ast.go
@@ -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 {