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 {