fix(output): make rendering cancellable mid-stream

cmd/grg/main.go

@@ -202,12 +202,13 @@ func emitResults(ctx context.Context, cfg *model.Config, results []*search.BlobR
 		return noMatchError{}
 	}
 
-	if err := ctx.Err(); err != nil {
-		return cancelError{err: err, quiet: cfg.Quiet}
-	}
-
+	// Rendering is cancellable: Format observes ctx at coarse boundaries, so the
+	// pre-flight check is folded into the render itself.
 	formatter := output.NewFormatter(cfg)
-	if err := formatter.Format(stdout, aggregated); err != nil {
+	if err := formatter.Format(ctx, stdout, aggregated); err != nil {
+		if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) || ctx.Err() != nil {
+			return cancelError{err: err, quiet: cfg.Quiet}
+		}
 		return err
 	}
 

internal/output/count.go

@@ -1,6 +1,7 @@
 package output
 
 import (
+	"context"
 	"fmt"
 	"io"
 	"strconv"
@@ -24,14 +25,20 @@ func NewCountFormatter(cfg *model.Config) *CountFormatter {
 	}
 }
 
-// Format writes <commit-short>:<path>:<count> lines to w.
-func (f *CountFormatter) Format(w io.Writer, results *aggregator.AggregatedResults) error {
+// Format writes <commit-short>:<path>:<count> lines to w, aborting with
+// ctx.Err() if ctx is cancelled mid-render.
+func (f *CountFormatter) Format(ctx context.Context, w io.Writer, results *aggregator.AggregatedResults) error {
 	if results == nil || len(results.Files) == 0 {
 		return nil
 	}
 
 	c := f.color
+	guard := newCancelGuard(ctx)
 	for _, file := range results.Files {
+		if err := guard.boundary(); err != nil {
+			return err
+		}
+
 		for _, commit := range file.Commits {
 			count := len(commit.Matches)
 			if commit.IsBinary {
@@ -57,6 +64,9 @@ func (f *CountFormatter) Format(w io.Writer, results *aggregator.AggregatedResul
 				formattedCount); err != nil {
 				return err
 			}
+			if err := guard.lines(1); err != nil {
+				return err
+			}
 		}
 	}
 

internal/output/files_with_matches.go

@@ -1,6 +1,7 @@
 package output
 
 import (
+	"context"
 	"fmt"
 	"io"
 
@@ -23,16 +24,22 @@ func NewFilesWithMatchesFormatter(cfg *model.Config) *FilesWithMatchesFormatter
 	}
 }
 
-// Format writes distinct <commit-short>:<path> lines to w.
-func (f *FilesWithMatchesFormatter) Format(w io.Writer, results *aggregator.AggregatedResults) error {
+// Format writes distinct <commit-short>:<path> lines to w, aborting with
+// ctx.Err() if ctx is cancelled mid-render.
+func (f *FilesWithMatchesFormatter) Format(ctx context.Context, w io.Writer, results *aggregator.AggregatedResults) error {
 	if results == nil || len(results.Files) == 0 {
 		return nil
 	}
 
 	c := f.color
+	guard := newCancelGuard(ctx)
 	seen := make(map[string]bool)
 
 	for _, file := range results.Files {
+		if err := guard.boundary(); err != nil {
+			return err
+		}
+
 		for _, commit := range file.Commits {
 			if len(commit.Matches) == 0 && !commit.IsBinary {
 				continue
@@ -56,6 +63,9 @@ func (f *FilesWithMatchesFormatter) Format(w io.Writer, results *aggregator.Aggr
 			if _, err := fmt.Fprintf(w, "%s%s%s\n", formattedCommit, formattedSep, formattedPath); err != nil {
 				return err
 			}
+			if err := guard.lines(1); err != nil {
+				return err
+			}
 		}
 	}
 

internal/output/grouped.go

@@ -1,6 +1,7 @@
 package output
 
 import (
+	"context"
 	"fmt"
 	"io"
 	"strconv"
@@ -24,16 +25,22 @@ func NewGroupedFormatter(cfg *model.Config) *GroupedFormatter {
 	}
 }
 
-// Format writes the grouped results to w.
-func (g *GroupedFormatter) Format(w io.Writer, results *aggregator.AggregatedResults) error {
+// Format writes the grouped results to w, aborting with ctx.Err() if ctx is
+// cancelled mid-render.
+func (g *GroupedFormatter) Format(ctx context.Context, w io.Writer, results *aggregator.AggregatedResults) error {
 	if results == nil || len(results.Files) == 0 {
 		return nil
 	}
 
 	c := g.color
+	guard := newCancelGuard(ctx)
 	firstFile := true
 
 	for _, file := range results.Files {
+		if err := guard.boundary(); err != nil {
+			return err
+		}
+
 		if !firstFile {
 			if _, err := fmt.Fprintln(w); err != nil {
 				return err
@@ -71,13 +78,19 @@ func (g *GroupedFormatter) Format(w io.Writer, results *aggregator.AggregatedRes
 						if err := g.writeLine(w, line.LineNum, line.LineText, line.IsMatch, line.Submatches); err != nil {
 							return err
 						}
+						if err := guard.lines(1); err != nil {
+							return err
+						}
 					}
 				}
 			} else {
 				for _, match := range commit.Matches {
 					if err := g.writeLine(w, match.LineNum, match.LineText, true, match.Submatches); err != nil {
 						return err
 					}
+					if err := guard.lines(1); err != nil {
+						return err
+					}
 				}
 			}
 		}

internal/output/output_test.go

@@ -2,6 +2,9 @@ package output
 
 import (
 	"bytes"
+	"context"
+	"errors"
+	"fmt"
 	"strings"
 	"testing"
 	"time"
@@ -60,7 +63,7 @@ func TestGroupedFormatter(t *testing.T) {
 
 	var buf bytes.Buffer
 	results := makeTestAggregatedResults()
-	if err := fmtter.Format(&buf, results); err != nil {
+	if err := fmtter.Format(context.Background(), &buf, results); err != nil {
 		t.Fatalf("unexpected error: %v", err)
 	}
 
@@ -84,7 +87,7 @@ func TestGroupedFormatter_NoLineNumber(t *testing.T) {
 
 	var buf bytes.Buffer
 	results := makeTestAggregatedResults()
-	if err := fmtter.Format(&buf, results); err != nil {
+	if err := fmtter.Format(context.Background(), &buf, results); err != nil {
 		t.Fatalf("unexpected error: %v", err)
 	}
 
@@ -140,7 +143,7 @@ func TestGroupedFormatter_ContextLines(t *testing.T) {
 	fmtter := NewFormatter(cfg)
 
 	var buf bytes.Buffer
-	if err := fmtter.Format(&buf, results); err != nil {
+	if err := fmtter.Format(context.Background(), &buf, results); err != nil {
 		t.Fatalf("unexpected error: %v", err)
 	}
 
@@ -167,7 +170,7 @@ func TestSingleLineFormatter(t *testing.T) {
 
 	var buf bytes.Buffer
 	results := makeTestAggregatedResults()
-	if err := fmtter.Format(&buf, results); err != nil {
+	if err := fmtter.Format(context.Background(), &buf, results); err != nil {
 		t.Fatalf("unexpected error: %v", err)
 	}
 
@@ -189,7 +192,7 @@ func TestSingleLineFormatter_NoLineNumber(t *testing.T) {
 
 	var buf bytes.Buffer
 	results := makeTestAggregatedResults()
-	if err := fmtter.Format(&buf, results); err != nil {
+	if err := fmtter.Format(context.Background(), &buf, results); err != nil {
 		t.Fatalf("unexpected error: %v", err)
 	}
 
@@ -210,7 +213,7 @@ func TestFilesWithMatchesFormatter(t *testing.T) {
 
 	var buf bytes.Buffer
 	results := makeTestAggregatedResults()
-	if err := fmtter.Format(&buf, results); err != nil {
+	if err := fmtter.Format(context.Background(), &buf, results); err != nil {
 		t.Fatalf("unexpected error: %v", err)
 	}
 
@@ -229,7 +232,7 @@ func TestCountFormatter(t *testing.T) {
 
 	var buf bytes.Buffer
 	results := makeTestAggregatedResults()
-	if err := fmtter.Format(&buf, results); err != nil {
+	if err := fmtter.Format(context.Background(), &buf, results); err != nil {
 		t.Fatalf("unexpected error: %v", err)
 	}
 
@@ -270,7 +273,7 @@ func TestBinaryFormatter(t *testing.T) {
 		Color:   model.ColorNever,
 	}
 	var buf bytes.Buffer
-	if err := NewFormatter(cfg).Format(&buf, results); err != nil {
+	if err := NewFormatter(cfg).Format(context.Background(), &buf, results); err != nil {
 		t.Fatal(err)
 	}
 
@@ -320,3 +323,133 @@ func TestColorHighlighting(t *testing.T) {
 		t.Errorf("expected %q, got %q", expected, highlighted)
 	}
 }
+
+// makeLargeAggregatedResults builds a result set whose rendering spans several
+// cancelCheckInterval windows, so that a mid-render cancellation is observable.
+func makeLargeAggregatedResults(files, matchesPerFile int) *aggregator.AggregatedResults {
+	date := time.Date(2026, 9, 4, 12, 0, 0, 0, time.UTC)
+	res := &aggregator.AggregatedResults{TotalFiles: files}
+	for f := range files {
+		matches := make([]model.SearchMatch, matchesPerFile)
+		for m := range matches {
+			matches[m] = model.SearchMatch{
+				LineNum:  m + 1,
+				LineText: fmt.Sprintf("needle in file %d line %d", f, m+1),
+			}
+		}
+		res.Files = append(res.Files, aggregator.FileMatches{
+			Path: fmt.Sprintf("pkg/file%04d.go", f),
+			Commits: []aggregator.CommitMatches{{
+				CommitSHA:  fmt.Sprintf("%040x", f),
+				ShortSHA:   fmt.Sprintf("%07x", f),
+				CommitDate: date,
+				Author:     "Alice",
+				AuthorName: "Alice",
+				Summary:    "bulk commit",
+				Matches:    matches,
+			}},
+		})
+		res.TotalMatches += matchesPerFile
+	}
+	return res
+}
+
+// cancelAfterWriter cancels the render's context once after lines have been
+// written, and keeps recording whatever the formatter emits afterwards so the
+// test can prove the remainder was never produced.
+type cancelAfterWriter struct {
+	buf    bytes.Buffer
+	cancel context.CancelFunc
+	after  int
+	lines  int
+}
+
+func (w *cancelAfterWriter) Write(p []byte) (int, error) {
+	n, err := w.buf.Write(p)
+	w.lines += bytes.Count(p, []byte("\n"))
+	if w.cancel != nil && w.lines >= w.after {
+		w.cancel()
+		w.cancel = nil
+	}
+	return n, err
+}
+
+func TestFormatterCancelledMidRender(t *testing.T) {
+	tests := []struct {
+		name    string
+		cfg     *model.Config
+		results *aggregator.AggregatedResults
+	}{
+		{
+			name:    "grouped",
+			cfg:     &model.Config{Heading: true, LineNumber: true, Color: model.ColorNever},
+			results: makeLargeAggregatedResults(8, 512),
+		},
+		{
+			name:    "single",
+			cfg:     &model.Config{Heading: false, LineNumber: true, Color: model.ColorNever},
+			results: makeLargeAggregatedResults(8, 512),
+		},
+		{
+			name:    "count",
+			cfg:     &model.Config{Count: true, Color: model.ColorNever},
+			results: makeLargeAggregatedResults(3000, 1),
+		},
+		{
+			name:    "files-with-matches",
+			cfg:     &model.Config{FilesWithMatches: true, Color: model.ColorNever},
+			results: makeLargeAggregatedResults(3000, 1),
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			var full bytes.Buffer
+			if err := NewFormatter(tc.cfg).Format(context.Background(), &full, tc.results); err != nil {
+				t.Fatalf("uncancelled render failed: %v", err)
+			}
+			want := full.String()
+			if lines := strings.Count(want, "\n"); lines <= 2*cancelCheckInterval {
+				t.Fatalf("fixture renders %d lines, too few to span several check intervals", lines)
+			}
+
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+			w := &cancelAfterWriter{cancel: cancel, after: 10}
+
+			err := NewFormatter(tc.cfg).Format(ctx, w, tc.results)
+			if !errors.Is(err, context.Canceled) {
+				t.Fatalf("expected context.Canceled, got %v", err)
+			}
+
+			got := w.buf.String()
+			if got == "" {
+				t.Fatal("expected the output written before cancellation to be retained")
+			}
+			if len(got) >= len(want) {
+				t.Fatalf("cancellation emitted %d of %d bytes: the remainder was not skipped", len(got), len(want))
+			}
+			if !strings.HasPrefix(want, got) {
+				t.Fatalf("output after cancellation is not a prefix of the full render (%d bytes written)", len(got))
+			}
+			if !strings.HasSuffix(got, "\n") {
+				t.Fatalf("cancellation truncated mid-line: %q", got[max(0, len(got)-64):])
+			}
+		})
+	}
+}
+
+func TestFormatterAlreadyCancelledEmitsNothing(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	cfg := &model.Config{Heading: true, LineNumber: true, Color: model.ColorNever}
+	var buf bytes.Buffer
+	err := NewFormatter(cfg).Format(ctx, &buf, makeTestAggregatedResults())
+	if !errors.Is(err, context.Canceled) {
+		t.Fatalf("expected context.Canceled, got %v", err)
+	}
+	if buf.Len() != 0 {
+		t.Fatalf("expected no output for an already-cancelled context, got %q", buf.String())
+	}
+}