From 8c92d2849802421fbb266fe9244ee5a5bc5f1a97 Mon Sep 17 00:00:00 2001 From: JATMN Date: Tue, 18 Aug 2026 12:21:37 -0700 Subject: [PATCH 01/18] fix(daemon): synchronize TestPoolDrainKillsStraggler on active workers The test waited on pool.QueueDepth() == 1, which reflects slot-channel occupancy set immediately after Run acquires a slot. Drain(), however, reads len(p.active), and the worker handle is only added to p.active in runOnce after Launcher returns. On loaded CI runners the test goroutine could call Drain() in the window between slot acquisition and worker tracking, causing Drain() to take the early 'all workers drained' return and never force-kill the straggler. Synchronize on the state Drain() actually reads by waiting for the worker to appear in WorkerStats() (which is built from p.active). Also add a comment explaining why WorkerStats is the right signal here. Fixes Gitlawb/zero#919. --- internal/daemon/pool_test.go | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/internal/daemon/pool_test.go b/internal/daemon/pool_test.go index 7ced2bef4..54750fe75 100644 --- a/internal/daemon/pool_test.go +++ b/internal/daemon/pool_test.go @@ -224,7 +224,8 @@ func TestPoolDrainKillsStraggler(t *testing.T) { _, _ = pool.Run(context.Background(), WorkerSpec{Session: "a"}, &collectSink{}) close(runDone) }() - waitFor(t, func() bool { return pool.QueueDepth() == 1 }) + // Wait until the worker is tracked; Drain reads the active set, not slot occupancy. + waitFor(t, func() bool { return len(pool.WorkerStats()) == 1 }) pool.Drain() // KillTimeout elapses, straggler is force-killed if atomic.LoadInt32(&straggler.killed) != 1 { From dd0049c1cb8e786b1217eb138bd5ddb634bb11c7 Mon Sep 17 00:00:00 2001 From: jatmn Date: Tue, 18 Aug 2026 12:44:30 -0700 Subject: [PATCH 02/18] fix(daemon): close launch and drain race --- internal/daemon/pool.go | 31 +++++++++++++++++++++++++------ internal/daemon/pool_test.go | 34 ++++++++++++++++++++++++++++++++++ 2 files changed, 59 insertions(+), 6 deletions(-) diff --git a/internal/daemon/pool.go b/internal/daemon/pool.go index 6de38315b..de9a9607c 100644 --- a/internal/daemon/pool.go +++ b/internal/daemon/pool.go @@ -92,10 +92,11 @@ type Pool struct { opts PoolOptions slots chan struct{} - mu sync.Mutex - draining bool - active map[int]WorkerHandle // worker id -> handle, for drain/kill + status - nextID int + mu sync.Mutex + draining bool + active map[int]WorkerHandle // worker id -> handle, for drain/kill + status + launching int // launchers in progress; Drain must not mistake these for idle + nextID int drainOnce sync.Once drained chan struct{} @@ -239,11 +240,29 @@ func (p *Pool) Run(ctx context.Context, spec WorkerSpec, sink Sink) (int, error) // runOnce launches a single worker, pumps its output to sink, and returns its // exit code. The worker handle is tracked so Drain can kill it. func (p *Pool) runOnce(ctx context.Context, id int, spec WorkerSpec, sink Sink) (int, error) { + p.mu.Lock() + if p.draining { + p.mu.Unlock() + return 0, ErrPoolDraining + } + p.launching++ + p.mu.Unlock() handle, err := p.opts.Launcher(ctx, spec) + p.mu.Lock() + p.launching-- + draining := p.draining + if err == nil && !draining { + p.active[id] = handle + } + p.mu.Unlock() if err != nil { return 0, err } - p.track(id, handle) + if draining { + _ = handle.Kill() + _, _ = handle.Wait() + return 0, ErrPoolDraining + } defer p.untrack(id) // Pump stdout lines until the stream ends. @@ -333,7 +352,7 @@ func (p *Pool) Drain() { deadline := time.Now().Add(p.opts.KillTimeout) for time.Now().Before(deadline) { p.mu.Lock() - n := len(p.active) + n := len(p.active) + p.launching p.mu.Unlock() if n == 0 { return // all workers drained gracefully diff --git a/internal/daemon/pool_test.go b/internal/daemon/pool_test.go index 54750fe75..2a7976809 100644 --- a/internal/daemon/pool_test.go +++ b/internal/daemon/pool_test.go @@ -243,6 +243,40 @@ func TestPoolDrainKillsStraggler(t *testing.T) { } } +func TestPoolDrainKillsWorkerLaunchedAfterDrainStarts(t *testing.T) { + launchStarted := make(chan struct{}) + releaseLaunch := make(chan struct{}) + straggler := &fakeWorker{pid: 1, waitCh: make(chan struct{})} + pool, _ := NewPool(PoolOptions{Size: 1, KillTimeout: 10 * time.Millisecond, Launcher: func(context.Context, WorkerSpec) (WorkerHandle, error) { + close(launchStarted) + <-releaseLaunch + return straggler, nil + }}) + runDone := make(chan struct{}) + go func() { + _, _ = pool.Run(context.Background(), WorkerSpec{Session: "a"}, &collectSink{}) + close(runDone) + }() + <-launchStarted + drained := make(chan struct{}) + go func() { pool.Drain(); close(drained) }() + time.Sleep(20 * time.Millisecond) + close(releaseLaunch) + select { + case <-drained: + case <-time.After(2 * time.Second): + t.Fatal("Drain did not finish") + } + select { + case <-runDone: + case <-time.After(2 * time.Second): + t.Fatal("Run did not finish after Drain") + } + if atomic.LoadInt32(&straggler.killed) != 1 { + t.Fatal("Drain must kill a worker whose launch completed after draining began") + } +} + func waitFor(t *testing.T, cond func() bool) { t.Helper() deadline := time.Now().Add(2 * time.Second) From ea7894d0dfa4827df70a1e991d92da3029595c15 Mon Sep 17 00:00:00 2001 From: jatmn Date: Tue, 18 Aug 2026 12:57:00 -0700 Subject: [PATCH 03/18] fix(daemon): wait for late launch cleanup during drain --- internal/daemon/pool.go | 8 +++++++- internal/daemon/pool_test.go | 7 ++++++- 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/internal/daemon/pool.go b/internal/daemon/pool.go index de9a9607c..8db615108 100644 --- a/internal/daemon/pool.go +++ b/internal/daemon/pool.go @@ -249,18 +249,24 @@ func (p *Pool) runOnce(ctx context.Context, id int, spec WorkerSpec, sink Sink) p.mu.Unlock() handle, err := p.opts.Launcher(ctx, spec) p.mu.Lock() - p.launching-- draining := p.draining if err == nil && !draining { p.active[id] = handle + p.launching-- } p.mu.Unlock() if err != nil { + p.mu.Lock() + p.launching-- + p.mu.Unlock() return 0, err } if draining { _ = handle.Kill() _, _ = handle.Wait() + p.mu.Lock() + p.launching-- + p.mu.Unlock() return 0, ErrPoolDraining } defer p.untrack(id) diff --git a/internal/daemon/pool_test.go b/internal/daemon/pool_test.go index 2a7976809..407d4c6a0 100644 --- a/internal/daemon/pool_test.go +++ b/internal/daemon/pool_test.go @@ -260,7 +260,12 @@ func TestPoolDrainKillsWorkerLaunchedAfterDrainStarts(t *testing.T) { <-launchStarted drained := make(chan struct{}) go func() { pool.Drain(); close(drained) }() - time.Sleep(20 * time.Millisecond) + waitFor(t, pool.isDraining) + select { + case <-drained: + t.Fatal("Drain returned while a launcher was still in progress") + default: + } close(releaseLaunch) select { case <-drained: From 184ada9ac24ff22f72cc34681f89b946caae3423 Mon Sep 17 00:00:00 2001 From: jatmn Date: Tue, 18 Aug 2026 13:00:49 -0700 Subject: [PATCH 04/18] fix(imageinput): remove vulnerable PDF parser --- go.mod | 1 - go.sum | 2 - internal/imageinput/pdf.go | 99 ++++++++++----------------------- internal/imageinput/pdf_test.go | 21 +++---- 4 files changed, 36 insertions(+), 87 deletions(-) diff --git a/go.mod b/go.mod index 8d06eefc8..943ab15a2 100644 --- a/go.mod +++ b/go.mod @@ -15,7 +15,6 @@ require ( github.com/charmbracelet/x/ansi v0.11.7 github.com/charmbracelet/x/term v0.2.2 github.com/coder/websocket v1.8.15 - github.com/ledongthuc/pdf v0.0.0-20250511090121-5959a4027728 golang.org/x/image v0.45.0 golang.org/x/sys v0.47.0 mvdan.cc/sh/v3 v3.13.1 diff --git a/go.sum b/go.sum index d4601485b..d2190e20c 100644 --- a/go.sum +++ b/go.sum @@ -48,8 +48,6 @@ github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= -github.com/ledongthuc/pdf v0.0.0-20250511090121-5959a4027728 h1:QwWKgMY28TAXaDl+ExRDqGQltzXqN/xypdKP86niVn8= -github.com/ledongthuc/pdf v0.0.0-20250511090121-5959a4027728/go.mod h1:1fEHWurg7pvf5SG6XNE5Q8UZmOwex51Mkx3SLhrW5B4= github.com/lucasb-eyer/go-colorful v1.4.0 h1:UtrWVfLdarDgc44HcS7pYloGHJUjHV/4FwW4TvVgFr4= github.com/lucasb-eyer/go-colorful v1.4.0/go.mod h1:R4dSotOR9KMtayYi1e77YzuveK+i7ruzyGqttikkLy0= github.com/mattn/go-runewidth v0.0.24 h1:cpokDiIn0MGnhdHwuWnJBITySJ20QyNGnY2kR/ay2DU= diff --git a/internal/imageinput/pdf.go b/internal/imageinput/pdf.go index 93ac81023..c9f0803e6 100644 --- a/internal/imageinput/pdf.go +++ b/internal/imageinput/pdf.go @@ -14,13 +14,10 @@ import ( "time" "github.com/Gitlawb/zero/internal/zeroruntime" - "github.com/ledongthuc/pdf" ) -// Dependency posture (see stage 12): the DEFAULT build extracts a PDF's text -// layer in pure Go via github.com/ledongthuc/pdf (BSD-licensed, no CGO, no -// transitive deps), so ZERO stays a single static cross-compilable binary with -// no runtime dependencies. Rasterizing pages to images for vision models needs +// Dependency posture: PDF processing uses Poppler's bounded tools. Rasterizing +// pages to images for vision models needs // real font/graphics rendering, which no maintained pure-Go library does well; // that path is OPTIONAL and uses the poppler tools (pdftotext / pdftoppm) only // when they are already on PATH -- the same "external tool the user may have" @@ -154,29 +151,13 @@ func LoadDocument(path string, workspaceRoot string, opts DocumentOptions) (Docu } } - // Text path. Prefer poppler's pdftotext when present (it handles more font - // encodings); otherwise use the pure-Go extractor. Either way, absence of the - // external tool is not an error. + // Text path uses Poppler's pdftotext, which handles PDF parsing outside the + // process and is bounded by popplerTimeout. text, pages := "", 0 if useExternal { if t, ok := extractTextWithPoppler(data); ok { text = t - // pdftotext does not report a page count, so derive it from the pure-Go - // reader (cheap structural read, no text extraction) to keep - // Document.Pages correct regardless of which text path wins. - pages = pdfPageCount(data) - } - } - if strings.TrimSpace(text) == "" { - t, p, terr := extractTextPureGo(data) - if terr != nil { - // Only surface the pure-Go error when we have nothing else (no poppler - // text and no rasterized pages) to offer. - if len(images) == 0 { - return Document{}, terr - } - } else { - text, pages = t, p + pages = pdfPageCountWithPoppler(data) } } @@ -232,53 +213,6 @@ func readDocumentBytes(path string, workspaceRoot string) ([]byte, error) { return data, nil } -// extractTextPureGo extracts the full text layer with the pure-Go parser. The -// ledongthuc/pdf parser panics (not errors) on some malformed structures, so the -// whole call is wrapped in a recover: a bad PDF becomes a clean error, never a -// crash that escapes the package. It returns the joined text and the page count. -func extractTextPureGo(data []byte) (text string, pages int, err error) { - defer func() { - if rec := recover(); rec != nil { - text, pages = "", 0 - err = fmt.Errorf("could not parse PDF (malformed or unsupported): %v", rec) - } - }() - - reader, rerr := pdf.NewReader(bytes.NewReader(data), int64(len(data))) - if rerr != nil { - return "", 0, fmt.Errorf("could not parse PDF: %w", rerr) - } - pages = reader.NumPage() - - var buf strings.Builder - plain, perr := reader.GetPlainText() - if perr != nil { - return "", pages, fmt.Errorf("could not extract PDF text: %w", perr) - } - if _, cerr := io.Copy(&buf, plain); cerr != nil { - return "", pages, fmt.Errorf("could not read PDF text: %w", cerr) - } - return strings.TrimSpace(buf.String()), pages, nil -} - -// pdfPageCount returns the page count via the pure-Go reader without extracting -// any text. It backs Document.Pages on the poppler text path (pdftotext does not -// report a count). Like extractTextPureGo it recovers from the parser's panics on -// malformed input and reports 0 rather than crashing -- the page count is -// informational, so an unreadable structure simply yields 0. -func pdfPageCount(data []byte) (pages int) { - defer func() { - if recover() != nil { - pages = 0 - } - }() - reader, err := pdf.NewReader(bytes.NewReader(data), int64(len(data))) - if err != nil { - return 0 - } - return reader.NumPage() -} - // capDocumentText truncates text to MaxDocumentTextBytes on a UTF-8 rune // boundary and appends documentTruncatedMarker when it had to cut. The second // return reports whether truncation happened. The marker is counted against the @@ -346,6 +280,29 @@ func extractTextWithPoppler(data []byte) (string, bool) { return strings.TrimSpace(stdout.String()), true } +func pdfPageCountWithPoppler(data []byte) int { + if !popplerAvailable("pdfinfo") { + return 0 + } + ctx, cancel := context.WithTimeout(context.Background(), popplerTimeout) + defer cancel() + cmd := exec.CommandContext(ctx, "pdfinfo", "-") + cmd.Stdin = bytes.NewReader(data) + out, err := cmd.Output() + if err != nil { + return 0 + } + for _, line := range strings.Split(string(out), "\n") { + if value, ok := strings.CutPrefix(strings.TrimSpace(line), "Pages:"); ok { + var pages int + if _, err := fmt.Sscan(value, &pages); err == nil { + return pages + } + } + } + return 0 +} + // rasterizeWithPoppler renders the first maxPages pages to PNG via pdftoppm and // returns them as normalized ImageBlocks (reusing the image allow-list, sniff, // and per-image cap). It returns an error when pdftoppm is absent or rendering diff --git a/internal/imageinput/pdf_test.go b/internal/imageinput/pdf_test.go index 614250ab3..cdc2751e3 100644 --- a/internal/imageinput/pdf_test.go +++ b/internal/imageinput/pdf_test.go @@ -243,7 +243,7 @@ func TestLoadDocumentNoTextNoRaster(t *testing.T) { } // Force the pure-Go path with no external rasterizer so the no-text branch is // deterministic regardless of what is installed on the test host. - _, err := LoadDocument("scan.pdf", root, DocumentOptions{disableExternalTools: true}) + _, err := LoadDocument("scan.pdf", root, DocumentOptions{}) if err == nil { t.Fatal("expected an error for a PDF with no extractable text and no raster") } @@ -292,7 +292,7 @@ func TestLoadDocumentMalformedDoesNotPanic(t *testing.T) { if err := os.WriteFile(filepath.Join(root, "bad.pdf"), bad, 0o644); err != nil { t.Fatalf("write bad: %v", err) } - _, err := LoadDocument("bad.pdf", root, DocumentOptions{disableExternalTools: true}) + _, err := LoadDocument("bad.pdf", root, DocumentOptions{}) if err == nil { t.Fatal("expected an error for malformed PDF bytes") } @@ -306,7 +306,7 @@ func TestLoadDocumentFallsBackToPureGo(t *testing.T) { if err := os.WriteFile(filepath.Join(root, "doc.pdf"), buildMinimalPDF(want), 0o644); err != nil { t.Fatalf("write pdf: %v", err) } - doc, err := LoadDocument("doc.pdf", root, DocumentOptions{disableExternalTools: true}) + doc, err := LoadDocument("doc.pdf", root, DocumentOptions{}) if err != nil { t.Fatalf("LoadDocument (pure-Go): %v", err) } @@ -315,23 +315,18 @@ func TestLoadDocumentFallsBackToPureGo(t *testing.T) { } } -// Vision-mode extraction without an available rasterizer must not error: it -// degrades to the text layer (a vision model can still read the text block). -func TestLoadDocumentVisionWithoutRasterizerUsesText(t *testing.T) { +func TestLoadDocumentVisionUsesText(t *testing.T) { root := t.TempDir() want := "Vision degrade to text" if err := os.WriteFile(filepath.Join(root, "doc.pdf"), buildMinimalPDF(want), 0o644); err != nil { t.Fatalf("write pdf: %v", err) } - doc, err := LoadDocument("doc.pdf", root, DocumentOptions{Vision: true, disableExternalTools: true}) + doc, err := LoadDocument("doc.pdf", root, DocumentOptions{Vision: true}) if err != nil { t.Fatalf("LoadDocument (vision, no raster): %v", err) } - if len(doc.Images) != 0 { - t.Fatalf("no rasterizer available, expected 0 images, got %d", len(doc.Images)) - } if !strings.Contains(doc.Text, want) { - t.Fatalf("vision-without-raster should keep text, got %q", doc.Text) + t.Fatalf("vision extraction should keep text, got %q", doc.Text) } } @@ -367,10 +362,10 @@ func TestCapDocumentTextRespectsCap(t *testing.T) { // backs Document.Pages on the poppler text path, where pdftotext gives no count) // and must return 0 -- not panic -- on garbage. func TestPDFPageCount(t *testing.T) { - if got := pdfPageCount(buildMinimalPDF("one page")); got != 1 { + if got := pdfPageCountWithPoppler(buildMinimalPDF("one page")); got != 1 { t.Fatalf("pdfPageCount = %d, want 1", got) } - if got := pdfPageCount([]byte("not a pdf at all")); got != 0 { + if got := pdfPageCountWithPoppler([]byte("not a pdf at all")); got != 0 { t.Fatalf("pdfPageCount on garbage = %d, want 0", got) } } From 13255089c551d8c049b48217748c5d1a613524c3 Mon Sep 17 00:00:00 2001 From: jatmn Date: Tue, 18 Aug 2026 13:55:57 -0700 Subject: [PATCH 05/18] fix: restore in-process PDF fallback and make drain errors terminal Replace the GO-2026-6115 parser with Detective-XH/gopdf so text extraction and page counts still work without Poppler, keep page counting independent of pdftotext, and stop Run from retrying ErrPoolDraining so late-launch drain cleanup cannot hang or wrap as ErrPermanent. --- go.mod | 2 + go.sum | 4 ++ internal/daemon/pool.go | 18 ++++++ internal/daemon/pool_test.go | 2 +- internal/imageinput/pdf.go | 97 +++++++++++++++++++++++++++++---- internal/imageinput/pdf_test.go | 67 ++++++++++++++++++++--- 6 files changed, 170 insertions(+), 20 deletions(-) diff --git a/go.mod b/go.mod index 943ab15a2..e5dfadd2b 100644 --- a/go.mod +++ b/go.mod @@ -6,6 +6,7 @@ require ( charm.land/bubbles/v2 v2.1.1 charm.land/bubbletea/v2 v2.0.9 charm.land/lipgloss/v2 v2.0.5 + github.com/Detective-XH/gopdf v0.8.7 github.com/Microsoft/go-winio v0.6.2 github.com/alecthomas/chroma/v2 v2.27.0 github.com/atotto/clipboard v0.1.4 @@ -36,4 +37,5 @@ require ( github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e // indirect golang.org/x/exp v0.0.0-20260611194520-c48552f49976 // indirect golang.org/x/sync v0.22.0 // indirect + golang.org/x/text v0.41.0 // indirect ) diff --git a/go.sum b/go.sum index d2190e20c..2683c86de 100644 --- a/go.sum +++ b/go.sum @@ -4,6 +4,8 @@ charm.land/bubbletea/v2 v2.0.9 h1:DpJCMWKgzQK8SJv4zbKKFHAI10ymWy/evClPFk0k0f8= charm.land/bubbletea/v2 v2.0.9/go.mod h1:2SkdgoTXluXJHOUwAoRlRXF/28vklb1rFl6GcgV1/ss= charm.land/lipgloss/v2 v2.0.5 h1:kbNxgeeUOYv5J0YdpxFjfvf3dFvqH8Aci4zB6xqFtrY= charm.land/lipgloss/v2 v2.0.5/go.mod h1:9oqhxt4yxIMe6q5A4kHr44DremZk7J9UNh74GlWa5nc= +github.com/Detective-XH/gopdf v0.8.7 h1:ISH9pBjXlgeo2Jlmai1eW4eTNsAUwvhpOcricQVvgYg= +github.com/Detective-XH/gopdf v0.8.7/go.mod h1:zVhltv/ba8FRM+lTVRqhd28DE2E8VgT3zdcPjXrveYQ= github.com/Microsoft/go-winio v0.6.2 h1:F2VQgta7ecxGYO8k3ZZz3RS8fVIXVxONVUPlNERoyfY= github.com/Microsoft/go-winio v0.6.2/go.mod h1:yd8OoFMLzJbo9gZq8j5qaps8bJ9aShtEA8Ipt1oGCvU= github.com/alecthomas/assert/v2 v2.11.0 h1:2Q9r3ki8+JYXvGsDyBXwH3LcJ+WK5D0gc5E8vS6K3D0= @@ -68,5 +70,7 @@ golang.org/x/sync v0.22.0 h1:SZjpbeLmrCk4xhRSZFNZW5gFUeCeFgjekvI/+gfScek= golang.org/x/sync v0.22.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= golang.org/x/sys v0.47.0 h1:o7XGOvZQCADBQQ4Y7VNq2dRWQR7JmOUW8Kxx4ZsNgWs= golang.org/x/sys v0.47.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= +golang.org/x/text v0.41.0 h1:vz/seA0lnX87Othu2f/0L24RcgrXD9/YFTSuGjj3rH8= +golang.org/x/text v0.41.0/go.mod h1:jvf1O8ajNzZqhSrQBPbutR/EB83Cc0CFrezNQIwbb5M= mvdan.cc/sh/v3 v3.13.1 h1:DP3TfgZhDkT7lerUdnp6PTGKyxxzz6T+cOlY/xEvfWk= mvdan.cc/sh/v3 v3.13.1/go.mod h1:lXJ8SexMvEVcHCoDvAGLZgFJ9Wsm2sulmoNEXGhYZD0= diff --git a/internal/daemon/pool.go b/internal/daemon/pool.go index 8db615108..4a4d08cea 100644 --- a/internal/daemon/pool.go +++ b/internal/daemon/pool.go @@ -204,6 +204,11 @@ func (p *Pool) Run(ctx context.Context, spec WorkerSpec, sink Sink) (int, error) if ctx.Err() != nil { return 0, ctx.Err() } + // Drain is terminal: do not backoff/retry, and do not wrap the + // shutdown error as ErrPermanent when attempts are exhausted. + if errors.Is(err, ErrPoolDraining) { + return 0, ErrPoolDraining + } p.logf("worker %d launch/run error: %v", stat.id, err) case code == 0: return 0, nil // clean success @@ -377,6 +382,19 @@ func (p *Pool) Drain() { p.logf("drain: killing straggler worker pid=%d", h.Pid()) _ = h.Kill() } + + // Force-kill only covers handles already in active. A launcher that is + // still inside Launcher has no handle yet; keep Drain blocked until that + // late-launch path finishes its own kill+wait and decrements launching. + for { + p.mu.Lock() + n := p.launching + p.mu.Unlock() + if n == 0 { + return + } + time.Sleep(5 * time.Millisecond) + } }) } diff --git a/internal/daemon/pool_test.go b/internal/daemon/pool_test.go index 407d4c6a0..e00e64f2e 100644 --- a/internal/daemon/pool_test.go +++ b/internal/daemon/pool_test.go @@ -247,7 +247,7 @@ func TestPoolDrainKillsWorkerLaunchedAfterDrainStarts(t *testing.T) { launchStarted := make(chan struct{}) releaseLaunch := make(chan struct{}) straggler := &fakeWorker{pid: 1, waitCh: make(chan struct{})} - pool, _ := NewPool(PoolOptions{Size: 1, KillTimeout: 10 * time.Millisecond, Launcher: func(context.Context, WorkerSpec) (WorkerHandle, error) { + pool, _ := NewPool(PoolOptions{Size: 1, MaxAttempts: 1, KillTimeout: 10 * time.Millisecond, Backoff: func(int) time.Duration { return 0 }, Launcher: func(context.Context, WorkerSpec) (WorkerHandle, error) { close(launchStarted) <-releaseLaunch return straggler, nil diff --git a/internal/imageinput/pdf.go b/internal/imageinput/pdf.go index c9f0803e6..6e1704711 100644 --- a/internal/imageinput/pdf.go +++ b/internal/imageinput/pdf.go @@ -13,16 +13,21 @@ import ( "strings" "time" + pdf "github.com/Detective-XH/gopdf" "github.com/Gitlawb/zero/internal/zeroruntime" ) -// Dependency posture: PDF processing uses Poppler's bounded tools. Rasterizing -// pages to images for vision models needs -// real font/graphics rendering, which no maintained pure-Go library does well; -// that path is OPTIONAL and uses the poppler tools (pdftotext / pdftoppm) only -// when they are already on PATH -- the same "external tool the user may have" -// posture as the LSP language servers. When poppler is absent, extraction -// silently degrades to the pure-Go text layer; absence is never an error. +// Dependency posture: the DEFAULT build extracts a PDF's text layer in pure Go +// via github.com/Detective-XH/gopdf (maintained replacement for the unfixed +// github.com/ledongthuc/pdf lineage named by GO-2026-6115; BSD-licensed, no +// CGO), so ZERO stays a single static cross-compilable binary with no runtime +// tool dependencies for text extraction. Rasterizing pages to images for vision +// models needs real font/graphics rendering, which no maintained pure-Go +// library does well; that path is OPTIONAL and uses the poppler tools +// (pdftotext / pdftoppm) only when they are already on PATH -- the same +// "external tool the user may have" posture as the LSP language servers. When +// poppler is absent, extraction silently degrades to the pure-Go text layer; +// absence is never an error. // MaxDocumentBytes is the per-document raw-file cap (32 MiB). PDFs are routinely // larger than the image cap, but we still bound the file before it is read into @@ -151,14 +156,35 @@ func LoadDocument(path string, workspaceRoot string, opts DocumentOptions) (Docu } } - // Text path uses Poppler's pdftotext, which handles PDF parsing outside the - // process and is bounded by popplerTimeout. + // Text path. Prefer poppler's pdftotext when present (it handles more font + // encodings); otherwise use the pure-Go extractor. Either way, absence of the + // external tool is not an error. Page counting is independent of text + // extraction so a missing/failing pdftotext still reports Pages when pdfinfo + // or the in-process reader can. text, pages := "", 0 if useExternal { if t, ok := extractTextWithPoppler(data); ok { text = t - pages = pdfPageCountWithPoppler(data) } + pages = pdfPageCountWithPoppler(data) + } + if strings.TrimSpace(text) == "" { + t, p, terr := extractTextPureGo(data) + if terr != nil { + // Only surface the pure-Go error when we have nothing else (no poppler + // text and no rasterized pages) to offer. + if len(images) == 0 { + return Document{}, terr + } + } else { + text = t + if pages == 0 { + pages = p + } + } + } + if pages == 0 { + pages = pdfPageCount(data) } text, truncated := capDocumentText(text) @@ -213,6 +239,55 @@ func readDocumentBytes(path string, workspaceRoot string) ([]byte, error) { return data, nil } +// extractTextPureGo extracts the full text layer with the pure-Go parser. The +// underlying reader can still panic on some malformed structures, so the whole +// call is wrapped in a recover: a bad PDF becomes a clean error, never a crash +// that escapes the package. It returns the joined text and the page count. +func extractTextPureGo(data []byte) (text string, pages int, err error) { + defer func() { + if rec := recover(); rec != nil { + text, pages = "", 0 + err = fmt.Errorf("could not parse PDF (malformed or unsupported): %v", rec) + } + }() + + reader, rerr := pdf.NewReader(bytes.NewReader(data), int64(len(data))) + if rerr != nil { + return "", 0, fmt.Errorf("could not parse PDF: %w", rerr) + } + pages = reader.NumPage() + + var buf strings.Builder + ctx, cancel := context.WithTimeout(context.Background(), popplerTimeout) + defer cancel() + plain, perr := reader.GetPlainText(ctx) + if perr != nil { + return "", pages, fmt.Errorf("could not extract PDF text: %w", perr) + } + if _, cerr := io.Copy(&buf, plain); cerr != nil { + return "", pages, fmt.Errorf("could not read PDF text: %w", cerr) + } + return strings.TrimSpace(buf.String()), pages, nil +} + +// pdfPageCount returns the page count via the pure-Go reader without extracting +// any text. It backs Document.Pages when pdfinfo is absent (pdftotext does not +// report a count). Like extractTextPureGo it recovers from parser panics on +// malformed input and reports 0 rather than crashing -- the page count is +// informational, so an unreadable structure simply yields 0. +func pdfPageCount(data []byte) (pages int) { + defer func() { + if recover() != nil { + pages = 0 + } + }() + reader, err := pdf.NewReader(bytes.NewReader(data), int64(len(data))) + if err != nil { + return 0 + } + return reader.NumPage() +} + // capDocumentText truncates text to MaxDocumentTextBytes on a UTF-8 rune // boundary and appends documentTruncatedMarker when it had to cut. The second // return reports whether truncation happened. The marker is counted against the @@ -259,7 +334,7 @@ func popplerAvailable(name string) bool { // extractTextWithPoppler runs `pdftotext - -` (read stdin, write stdout) when // pdftotext is on PATH. The bool is false when the tool is absent or failed, so -// the caller can fall back to the pure-Go extractor. Absence is never an error. +// the caller falls back to the pure-Go extractor. Absence is never an error. func extractTextWithPoppler(data []byte) (string, bool) { if !popplerAvailable("pdftotext") { return "", false diff --git a/internal/imageinput/pdf_test.go b/internal/imageinput/pdf_test.go index cdc2751e3..9951ae2d8 100644 --- a/internal/imageinput/pdf_test.go +++ b/internal/imageinput/pdf_test.go @@ -8,13 +8,14 @@ import ( "strconv" "strings" "testing" + "time" ) const minimalPDFTextChunkSize = 80 // buildMinimalPDF assembles a tiny, single-page PDF whose content stream draws // the given text. It computes a real cross-reference table and trailer so a -// pure-Go PDF parser (ledongthuc/pdf) accepts it. Generating the fixture in-test +// pure-Go PDF parser (Detective-XH/gopdf) accepts it. Generating the fixture in-test // keeps the repo free of opaque binary blobs while still exercising the real // text-extraction path on real PDF bytes. func buildMinimalPDF(text string) []byte { @@ -243,7 +244,7 @@ func TestLoadDocumentNoTextNoRaster(t *testing.T) { } // Force the pure-Go path with no external rasterizer so the no-text branch is // deterministic regardless of what is installed on the test host. - _, err := LoadDocument("scan.pdf", root, DocumentOptions{}) + _, err := LoadDocument("scan.pdf", root, DocumentOptions{disableExternalTools: true}) if err == nil { t.Fatal("expected an error for a PDF with no extractable text and no raster") } @@ -292,7 +293,7 @@ func TestLoadDocumentMalformedDoesNotPanic(t *testing.T) { if err := os.WriteFile(filepath.Join(root, "bad.pdf"), bad, 0o644); err != nil { t.Fatalf("write bad: %v", err) } - _, err := LoadDocument("bad.pdf", root, DocumentOptions{}) + _, err := LoadDocument("bad.pdf", root, DocumentOptions{disableExternalTools: true}) if err == nil { t.Fatal("expected an error for malformed PDF bytes") } @@ -306,7 +307,7 @@ func TestLoadDocumentFallsBackToPureGo(t *testing.T) { if err := os.WriteFile(filepath.Join(root, "doc.pdf"), buildMinimalPDF(want), 0o644); err != nil { t.Fatalf("write pdf: %v", err) } - doc, err := LoadDocument("doc.pdf", root, DocumentOptions{}) + doc, err := LoadDocument("doc.pdf", root, DocumentOptions{disableExternalTools: true}) if err != nil { t.Fatalf("LoadDocument (pure-Go): %v", err) } @@ -321,12 +322,15 @@ func TestLoadDocumentVisionUsesText(t *testing.T) { if err := os.WriteFile(filepath.Join(root, "doc.pdf"), buildMinimalPDF(want), 0o644); err != nil { t.Fatalf("write pdf: %v", err) } - doc, err := LoadDocument("doc.pdf", root, DocumentOptions{Vision: true}) + doc, err := LoadDocument("doc.pdf", root, DocumentOptions{Vision: true, disableExternalTools: true}) if err != nil { t.Fatalf("LoadDocument (vision, no raster): %v", err) } + if len(doc.Images) != 0 { + t.Fatalf("no rasterizer available, expected 0 images, got %d", len(doc.Images)) + } if !strings.Contains(doc.Text, want) { - t.Fatalf("vision extraction should keep text, got %q", doc.Text) + t.Fatalf("vision-without-raster should keep text, got %q", doc.Text) } } @@ -362,14 +366,61 @@ func TestCapDocumentTextRespectsCap(t *testing.T) { // backs Document.Pages on the poppler text path, where pdftotext gives no count) // and must return 0 -- not panic -- on garbage. func TestPDFPageCount(t *testing.T) { - if got := pdfPageCountWithPoppler(buildMinimalPDF("one page")); got != 1 { + if got := pdfPageCount(buildMinimalPDF("one page")); got != 1 { t.Fatalf("pdfPageCount = %d, want 1", got) } - if got := pdfPageCountWithPoppler([]byte("not a pdf at all")); got != 0 { + if got := pdfPageCount([]byte("not a pdf at all")); got != 0 { t.Fatalf("pdfPageCount on garbage = %d, want 0", got) } } +func TestPDFPageCountIndependentOfTextExtraction(t *testing.T) { + root := t.TempDir() + if err := os.WriteFile(filepath.Join(root, "doc.pdf"), buildMinimalPDF("pages without pdftotext"), 0o644); err != nil { + t.Fatalf("write pdf: %v", err) + } + doc, err := LoadDocument("doc.pdf", root, DocumentOptions{disableExternalTools: true}) + if err != nil { + t.Fatalf("LoadDocument: %v", err) + } + if doc.Pages != 1 { + t.Fatalf("Pages = %d, want 1 when text extraction uses the in-process reader", doc.Pages) + } +} + +func TestLoadDocumentHostilePDFStaysBounded(t *testing.T) { + root := t.TempDir() + cases := map[string][]byte{ + "cycle.pdf": []byte("%PDF-1.4\n1 0 obj\n<< /Type /Catalog /Pages 1 0 R /Parent 1 0 R /Kids [1 0 R] /Count 999999999 /First 1 0 R /Next 1 0 R >>\nendobj\ntrailer\n<< /Root 1 0 R /Size 999999999 >>\nstartxref\n9\n%%EOF\n"), + "hex.pdf": []byte("%PDF-1.4\n1 0 obj\n<< /Type /Catalog /Pages 2 0 R >>\nendobj\nstream\n<" + strings.Repeat("A", 4096) + "\nendstream\n%%EOF\n"), + } + done := make(chan error, 1) + go func() { + var first error + for name, body := range cases { + path := filepath.Join(root, name) + if err := os.WriteFile(path, body, 0o644); err != nil { + first = err + break + } + _, err := LoadDocument(name, root, DocumentOptions{disableExternalTools: true}) + if err == nil { + first = fmt.Errorf("%s: expected error for hostile PDF", name) + break + } + } + done <- first + }() + select { + case err := <-done: + if err != nil { + t.Fatal(err) + } + case <-time.After(2 * time.Second): + t.Fatal("hostile PDF parsing exceeded the resource bound") + } +} + // LooksLikeDocumentFile sniffs PDF content by magic bytes, so a real PDF with no // ".pdf" extension is still recognized while a non-PDF (even named .pdf) is not. func TestLooksLikeDocumentFile(t *testing.T) { From be75eb2c1c59b5441b1f09fcc5f927a570c78986 Mon Sep 17 00:00:00 2001 From: jatmn Date: Tue, 18 Aug 2026 13:56:28 -0700 Subject: [PATCH 06/18] fix(daemon): always decrement in-flight launch count A panic or early return from Launcher left launching elevated, so Drain could wait forever after the grace window. Decrement via defer unless the worker was already moved into the active set. --- internal/daemon/pool.go | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/internal/daemon/pool.go b/internal/daemon/pool.go index 4a4d08cea..4bbddac32 100644 --- a/internal/daemon/pool.go +++ b/internal/daemon/pool.go @@ -252,26 +252,29 @@ func (p *Pool) runOnce(ctx context.Context, id int, spec WorkerSpec, sink Sink) } p.launching++ p.mu.Unlock() + inLaunch := true + defer func() { + if inLaunch { + p.mu.Lock() + p.launching-- + p.mu.Unlock() + } + }() handle, err := p.opts.Launcher(ctx, spec) p.mu.Lock() draining := p.draining if err == nil && !draining { p.active[id] = handle p.launching-- + inLaunch = false } p.mu.Unlock() if err != nil { - p.mu.Lock() - p.launching-- - p.mu.Unlock() return 0, err } if draining { _ = handle.Kill() _, _ = handle.Wait() - p.mu.Lock() - p.launching-- - p.mu.Unlock() return 0, ErrPoolDraining } defer p.untrack(id) From 196062ac5048cfb02038c485b9dc78ad2e0dcc69 Mon Sep 17 00:00:00 2001 From: jatmn Date: Tue, 18 Aug 2026 14:03:39 -0700 Subject: [PATCH 07/18] fix(daemon,imageinput): bound late-launch drain wait and restore page counts Cap the post-kill launching wait to KillTimeout so a stuck Launcher cannot wedge shutdown, surface in-flight launch errors as ErrPoolDraining, and keep successful pdftotext page counts on the in-process reader. --- internal/daemon/pool.go | 13 +++++++++---- internal/daemon/pool_test.go | 2 +- internal/imageinput/pdf.go | 9 ++++++++- 3 files changed, 18 insertions(+), 6 deletions(-) diff --git a/internal/daemon/pool.go b/internal/daemon/pool.go index 4bbddac32..f9a62e1df 100644 --- a/internal/daemon/pool.go +++ b/internal/daemon/pool.go @@ -270,6 +270,9 @@ func (p *Pool) runOnce(ctx context.Context, id int, spec WorkerSpec, sink Sink) } p.mu.Unlock() if err != nil { + if draining { + return 0, ErrPoolDraining + } return 0, err } if draining { @@ -386,10 +389,12 @@ func (p *Pool) Drain() { _ = h.Kill() } - // Force-kill only covers handles already in active. A launcher that is - // still inside Launcher has no handle yet; keep Drain blocked until that - // late-launch path finishes its own kill+wait and decrements launching. - for { + // Force-kill only covers handles already in active. A launcher still + // inside Launcher has no handle yet; wait one more KillTimeout for that + // late-launch path to finish kill+wait. Do not wait forever: a Launcher + // that ignores ctx would otherwise wedge shutdown. + deadline = time.Now().Add(p.opts.KillTimeout) + for time.Now().Before(deadline) { p.mu.Lock() n := p.launching p.mu.Unlock() diff --git a/internal/daemon/pool_test.go b/internal/daemon/pool_test.go index e00e64f2e..28063e965 100644 --- a/internal/daemon/pool_test.go +++ b/internal/daemon/pool_test.go @@ -247,7 +247,7 @@ func TestPoolDrainKillsWorkerLaunchedAfterDrainStarts(t *testing.T) { launchStarted := make(chan struct{}) releaseLaunch := make(chan struct{}) straggler := &fakeWorker{pid: 1, waitCh: make(chan struct{})} - pool, _ := NewPool(PoolOptions{Size: 1, MaxAttempts: 1, KillTimeout: 10 * time.Millisecond, Backoff: func(int) time.Duration { return 0 }, Launcher: func(context.Context, WorkerSpec) (WorkerHandle, error) { + pool, _ := NewPool(PoolOptions{Size: 1, MaxAttempts: 1, KillTimeout: 2 * time.Second, Backoff: func(int) time.Duration { return 0 }, Launcher: func(context.Context, WorkerSpec) (WorkerHandle, error) { close(launchStarted) <-releaseLaunch return straggler, nil diff --git a/internal/imageinput/pdf.go b/internal/imageinput/pdf.go index 6e1704711..9a309cafd 100644 --- a/internal/imageinput/pdf.go +++ b/internal/imageinput/pdf.go @@ -165,8 +165,15 @@ func LoadDocument(path string, workspaceRoot string, opts DocumentOptions) (Docu if useExternal { if t, ok := extractTextWithPoppler(data); ok { text = t + // pdftotext does not report a page count; use the in-process reader + // (cheap structural read) so Document.Pages stays correct on the + // poppler text path. + pages = pdfPageCount(data) + } else { + // pdftotext missing/failed: still take a page count from pdfinfo when + // that tool can succeed independently. + pages = pdfPageCountWithPoppler(data) } - pages = pdfPageCountWithPoppler(data) } if strings.TrimSpace(text) == "" { t, p, terr := extractTextPureGo(data) From 862315a1030cd9608db346ac67fc6d220b94e885 Mon Sep 17 00:00:00 2001 From: jatmn Date: Tue, 18 Aug 2026 14:36:53 -0700 Subject: [PATCH 08/18] fix(imageinput): fall back to pdfinfo when in-process page count is zero Page counting was still tied to which text extractor won, so a successful pdftotext path never asked pdfinfo. Count independently, and document that Poppler is preferred when present. --- internal/imageinput/pdf.go | 64 +++++++++++++++++++-------------- internal/imageinput/pdf_test.go | 27 ++++++++++++++ 2 files changed, 65 insertions(+), 26 deletions(-) diff --git a/internal/imageinput/pdf.go b/internal/imageinput/pdf.go index 9a309cafd..863864ec9 100644 --- a/internal/imageinput/pdf.go +++ b/internal/imageinput/pdf.go @@ -17,17 +17,17 @@ import ( "github.com/Gitlawb/zero/internal/zeroruntime" ) -// Dependency posture: the DEFAULT build extracts a PDF's text layer in pure Go -// via github.com/Detective-XH/gopdf (maintained replacement for the unfixed -// github.com/ledongthuc/pdf lineage named by GO-2026-6115; BSD-licensed, no -// CGO), so ZERO stays a single static cross-compilable binary with no runtime -// tool dependencies for text extraction. Rasterizing pages to images for vision -// models needs real font/graphics rendering, which no maintained pure-Go -// library does well; that path is OPTIONAL and uses the poppler tools -// (pdftotext / pdftoppm) only when they are already on PATH -- the same -// "external tool the user may have" posture as the LSP language servers. When -// poppler is absent, extraction silently degrades to the pure-Go text layer; -// absence is never an error. +// Dependency posture: LoadDocument prefers Poppler's pdftotext when it is on +// PATH and disableExternalTools is false (it handles more font encodings). +// github.com/Detective-XH/gopdf is the in-process fallback used when Poppler is +// missing, fails, or tests disable external tools -- a maintained replacement +// for the unfixed github.com/ledongthuc/pdf lineage named by GO-2026-6115 +// (BSD-licensed, no CGO). Rasterizing pages to images for vision models needs +// real font/graphics rendering, which no maintained pure-Go library does well; +// that path is OPTIONAL and uses pdftoppm only when it is already on PATH -- +// the same "external tool the user may have" posture as the LSP language +// servers. Absence of Poppler is never an error: text extraction degrades to +// the in-process reader. // MaxDocumentBytes is the per-document raw-file cap (32 MiB). PDFs are routinely // larger than the image cap, but we still bound the file before it is read into @@ -165,14 +165,6 @@ func LoadDocument(path string, workspaceRoot string, opts DocumentOptions) (Docu if useExternal { if t, ok := extractTextWithPoppler(data); ok { text = t - // pdftotext does not report a page count; use the in-process reader - // (cheap structural read) so Document.Pages stays correct on the - // poppler text path. - pages = pdfPageCount(data) - } else { - // pdftotext missing/failed: still take a page count from pdfinfo when - // that tool can succeed independently. - pages = pdfPageCountWithPoppler(data) } } if strings.TrimSpace(text) == "" { @@ -184,15 +176,13 @@ func LoadDocument(path string, workspaceRoot string, opts DocumentOptions) (Docu return Document{}, terr } } else { - text = t - if pages == 0 { - pages = p - } + text, pages = t, p } } - if pages == 0 { - pages = pdfPageCount(data) - } + // Page count is independent of which text extractor won: try the in-process + // reader first, then pdfinfo when external tools are enabled, so a PDF that + // pdftotext can read still reports Pages when gopdf cannot count it. + pages = resolvePageCount(data, useExternal, pages) text, truncated := capDocumentText(text) @@ -324,6 +314,28 @@ func utf8RuneStart(b byte) bool { return b&0xC0 != 0x80 } +// resolvePageCount fills Document.Pages from any available counter. already is +// a count captured during text extraction (0 means unknown). The in-process +// reader is tried first; pdfinfo is the fallback when external tools are on and +// the in-process count is still zero. +func resolvePageCount(data []byte, useExternal bool, already int) int { + if already > 0 { + return already + } + if pages := pageCountInProcess(data); pages > 0 { + return pages + } + if useExternal { + return pageCountPoppler(data) + } + return 0 +} + +var ( + pageCountInProcess = pdfPageCount + pageCountPoppler = pdfPageCountWithPoppler +) + func (o DocumentOptions) maxPages() int { if o.MaxPages > 0 { return o.MaxPages diff --git a/internal/imageinput/pdf_test.go b/internal/imageinput/pdf_test.go index 9951ae2d8..bf78d375d 100644 --- a/internal/imageinput/pdf_test.go +++ b/internal/imageinput/pdf_test.go @@ -388,6 +388,33 @@ func TestPDFPageCountIndependentOfTextExtraction(t *testing.T) { } } +// pdftotext success does not record a page count, so Pages must still fall +// through to pdfinfo when the in-process reader reports 0. +func TestResolvePageCountFallsBackToPopplerWhenInProcessIsZero(t *testing.T) { + origIn, origPop := pageCountInProcess, pageCountPoppler + t.Cleanup(func() { + pageCountInProcess, pageCountPoppler = origIn, origPop + }) + + pageCountInProcess = func([]byte) int { return 0 } + pageCountPoppler = func([]byte) int { return 7 } + + if got := resolvePageCount(nil, true, 0); got != 7 { + t.Fatalf("pdftotext-ok + in-process 0 + pdfinfo 7: Pages = %d, want 7", got) + } + if got := resolvePageCount(nil, false, 0); got != 0 { + t.Fatalf("external tools disabled: Pages = %d, want 0", got) + } + if got := resolvePageCount(nil, true, 3); got != 3 { + t.Fatalf("already-known count: Pages = %d, want 3", got) + } + + pageCountInProcess = func([]byte) int { return 2 } + if got := resolvePageCount(nil, true, 0); got != 2 { + t.Fatalf("in-process count wins over pdfinfo: Pages = %d, want 2", got) + } +} + func TestLoadDocumentHostilePDFStaysBounded(t *testing.T) { root := t.TempDir() cases := map[string][]byte{ From 1387c200d2ef7bbb635ac5e6673e20d7ead47b04 Mon Sep 17 00:00:00 2001 From: jatmn Date: Wed, 19 Aug 2026 11:35:24 -0700 Subject: [PATCH 09/18] chore: split PDF changes into separate PR --- go.mod | 3 +- go.sum | 8 +-- internal/imageinput/pdf.go | 93 ++++++++------------------------- internal/imageinput/pdf_test.go | 81 ++-------------------------- 4 files changed, 30 insertions(+), 155 deletions(-) diff --git a/go.mod b/go.mod index e5dfadd2b..8d06eefc8 100644 --- a/go.mod +++ b/go.mod @@ -6,7 +6,6 @@ require ( charm.land/bubbles/v2 v2.1.1 charm.land/bubbletea/v2 v2.0.9 charm.land/lipgloss/v2 v2.0.5 - github.com/Detective-XH/gopdf v0.8.7 github.com/Microsoft/go-winio v0.6.2 github.com/alecthomas/chroma/v2 v2.27.0 github.com/atotto/clipboard v0.1.4 @@ -16,6 +15,7 @@ require ( github.com/charmbracelet/x/ansi v0.11.7 github.com/charmbracelet/x/term v0.2.2 github.com/coder/websocket v1.8.15 + github.com/ledongthuc/pdf v0.0.0-20250511090121-5959a4027728 golang.org/x/image v0.45.0 golang.org/x/sys v0.47.0 mvdan.cc/sh/v3 v3.13.1 @@ -37,5 +37,4 @@ require ( github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e // indirect golang.org/x/exp v0.0.0-20260611194520-c48552f49976 // indirect golang.org/x/sync v0.22.0 // indirect - golang.org/x/text v0.41.0 // indirect ) diff --git a/go.sum b/go.sum index 2683c86de..232404cf1 100644 --- a/go.sum +++ b/go.sum @@ -4,8 +4,6 @@ charm.land/bubbletea/v2 v2.0.9 h1:DpJCMWKgzQK8SJv4zbKKFHAI10ymWy/evClPFk0k0f8= charm.land/bubbletea/v2 v2.0.9/go.mod h1:2SkdgoTXluXJHOUwAoRlRXF/28vklb1rFl6GcgV1/ss= charm.land/lipgloss/v2 v2.0.5 h1:kbNxgeeUOYv5J0YdpxFjfvf3dFvqH8Aci4zB6xqFtrY= charm.land/lipgloss/v2 v2.0.5/go.mod h1:9oqhxt4yxIMe6q5A4kHr44DremZk7J9UNh74GlWa5nc= -github.com/Detective-XH/gopdf v0.8.7 h1:ISH9pBjXlgeo2Jlmai1eW4eTNsAUwvhpOcricQVvgYg= -github.com/Detective-XH/gopdf v0.8.7/go.mod h1:zVhltv/ba8FRM+lTVRqhd28DE2E8VgT3zdcPjXrveYQ= github.com/Microsoft/go-winio v0.6.2 h1:F2VQgta7ecxGYO8k3ZZz3RS8fVIXVxONVUPlNERoyfY= github.com/Microsoft/go-winio v0.6.2/go.mod h1:yd8OoFMLzJbo9gZq8j5qaps8bJ9aShtEA8Ipt1oGCvU= github.com/alecthomas/assert/v2 v2.11.0 h1:2Q9r3ki8+JYXvGsDyBXwH3LcJ+WK5D0gc5E8vS6K3D0= @@ -50,6 +48,8 @@ github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= +github.com/ledongthuc/pdf v0.0.0-20250511090121-5959a4027728 h1:QwWKgMY28TAXaDl+ExRDqGQltzXqN/xypdKP86niVn8= +github.com/ledongthuc/pdf v0.0.0-20250511090121-5959a4027728/go.mod h1:1fEHWurg7pvf5SG6XNE5Q8UZmOwex51Mkx3SLhrW5B4= github.com/lucasb-eyer/go-colorful v1.4.0 h1:UtrWVfLdarDgc44HcS7pYloGHJUjHV/4FwW4TvVgFr4= github.com/lucasb-eyer/go-colorful v1.4.0/go.mod h1:R4dSotOR9KMtayYi1e77YzuveK+i7ruzyGqttikkLy0= github.com/mattn/go-runewidth v0.0.24 h1:cpokDiIn0MGnhdHwuWnJBITySJ20QyNGnY2kR/ay2DU= @@ -64,13 +64,13 @@ github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e h1:JVG44RsyaB9T2KIHavM github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e/go.mod h1:RbqR21r5mrJuqunuUZ/Dhy/avygyECGrLceyNeo4LiM= golang.org/x/exp v0.0.0-20260611194520-c48552f49976 h1:X8Hz2ImujgbmetVuW+w2YkyZChE3cBpZi2P158rTG9M= golang.org/x/exp v0.0.0-20260611194520-c48552f49976/go.mod h1:vnf4pv9iKZXY58sQE1L86zmNWJ4159e1RkcWiLCkeEY= +golang.org/x/image v0.44.0 h1:+tDekMZED9+LrtB3G5xzRggpVh9CARjZqROla3R3R+I= +golang.org/x/image v0.44.0/go.mod h1:V8K3KE9KKKE+pLpQDOeN18w9oacNSvy1tDOirTu4xtY= golang.org/x/image v0.45.0 h1:FMb1nTbH5H9vF55SriQHgFw5GnNL9Jg6L25BwXKzhB0= golang.org/x/image v0.45.0/go.mod h1:n62x/7RqlwXDvGsSU4u6IUTUf6KghUZ9Bt7cG/T9Fx4= golang.org/x/sync v0.22.0 h1:SZjpbeLmrCk4xhRSZFNZW5gFUeCeFgjekvI/+gfScek= golang.org/x/sync v0.22.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= golang.org/x/sys v0.47.0 h1:o7XGOvZQCADBQQ4Y7VNq2dRWQR7JmOUW8Kxx4ZsNgWs= golang.org/x/sys v0.47.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= -golang.org/x/text v0.41.0 h1:vz/seA0lnX87Othu2f/0L24RcgrXD9/YFTSuGjj3rH8= -golang.org/x/text v0.41.0/go.mod h1:jvf1O8ajNzZqhSrQBPbutR/EB83Cc0CFrezNQIwbb5M= mvdan.cc/sh/v3 v3.13.1 h1:DP3TfgZhDkT7lerUdnp6PTGKyxxzz6T+cOlY/xEvfWk= mvdan.cc/sh/v3 v3.13.1/go.mod h1:lXJ8SexMvEVcHCoDvAGLZgFJ9Wsm2sulmoNEXGhYZD0= diff --git a/internal/imageinput/pdf.go b/internal/imageinput/pdf.go index 863864ec9..93ac81023 100644 --- a/internal/imageinput/pdf.go +++ b/internal/imageinput/pdf.go @@ -13,21 +13,19 @@ import ( "strings" "time" - pdf "github.com/Detective-XH/gopdf" "github.com/Gitlawb/zero/internal/zeroruntime" + "github.com/ledongthuc/pdf" ) -// Dependency posture: LoadDocument prefers Poppler's pdftotext when it is on -// PATH and disableExternalTools is false (it handles more font encodings). -// github.com/Detective-XH/gopdf is the in-process fallback used when Poppler is -// missing, fails, or tests disable external tools -- a maintained replacement -// for the unfixed github.com/ledongthuc/pdf lineage named by GO-2026-6115 -// (BSD-licensed, no CGO). Rasterizing pages to images for vision models needs +// Dependency posture (see stage 12): the DEFAULT build extracts a PDF's text +// layer in pure Go via github.com/ledongthuc/pdf (BSD-licensed, no CGO, no +// transitive deps), so ZERO stays a single static cross-compilable binary with +// no runtime dependencies. Rasterizing pages to images for vision models needs // real font/graphics rendering, which no maintained pure-Go library does well; -// that path is OPTIONAL and uses pdftoppm only when it is already on PATH -- -// the same "external tool the user may have" posture as the LSP language -// servers. Absence of Poppler is never an error: text extraction degrades to -// the in-process reader. +// that path is OPTIONAL and uses the poppler tools (pdftotext / pdftoppm) only +// when they are already on PATH -- the same "external tool the user may have" +// posture as the LSP language servers. When poppler is absent, extraction +// silently degrades to the pure-Go text layer; absence is never an error. // MaxDocumentBytes is the per-document raw-file cap (32 MiB). PDFs are routinely // larger than the image cap, but we still bound the file before it is read into @@ -158,13 +156,15 @@ func LoadDocument(path string, workspaceRoot string, opts DocumentOptions) (Docu // Text path. Prefer poppler's pdftotext when present (it handles more font // encodings); otherwise use the pure-Go extractor. Either way, absence of the - // external tool is not an error. Page counting is independent of text - // extraction so a missing/failing pdftotext still reports Pages when pdfinfo - // or the in-process reader can. + // external tool is not an error. text, pages := "", 0 if useExternal { if t, ok := extractTextWithPoppler(data); ok { text = t + // pdftotext does not report a page count, so derive it from the pure-Go + // reader (cheap structural read, no text extraction) to keep + // Document.Pages correct regardless of which text path wins. + pages = pdfPageCount(data) } } if strings.TrimSpace(text) == "" { @@ -179,10 +179,6 @@ func LoadDocument(path string, workspaceRoot string, opts DocumentOptions) (Docu text, pages = t, p } } - // Page count is independent of which text extractor won: try the in-process - // reader first, then pdfinfo when external tools are enabled, so a PDF that - // pdftotext can read still reports Pages when gopdf cannot count it. - pages = resolvePageCount(data, useExternal, pages) text, truncated := capDocumentText(text) @@ -237,9 +233,9 @@ func readDocumentBytes(path string, workspaceRoot string) ([]byte, error) { } // extractTextPureGo extracts the full text layer with the pure-Go parser. The -// underlying reader can still panic on some malformed structures, so the whole -// call is wrapped in a recover: a bad PDF becomes a clean error, never a crash -// that escapes the package. It returns the joined text and the page count. +// ledongthuc/pdf parser panics (not errors) on some malformed structures, so the +// whole call is wrapped in a recover: a bad PDF becomes a clean error, never a +// crash that escapes the package. It returns the joined text and the page count. func extractTextPureGo(data []byte) (text string, pages int, err error) { defer func() { if rec := recover(); rec != nil { @@ -255,9 +251,7 @@ func extractTextPureGo(data []byte) (text string, pages int, err error) { pages = reader.NumPage() var buf strings.Builder - ctx, cancel := context.WithTimeout(context.Background(), popplerTimeout) - defer cancel() - plain, perr := reader.GetPlainText(ctx) + plain, perr := reader.GetPlainText() if perr != nil { return "", pages, fmt.Errorf("could not extract PDF text: %w", perr) } @@ -268,8 +262,8 @@ func extractTextPureGo(data []byte) (text string, pages int, err error) { } // pdfPageCount returns the page count via the pure-Go reader without extracting -// any text. It backs Document.Pages when pdfinfo is absent (pdftotext does not -// report a count). Like extractTextPureGo it recovers from parser panics on +// any text. It backs Document.Pages on the poppler text path (pdftotext does not +// report a count). Like extractTextPureGo it recovers from the parser's panics on // malformed input and reports 0 rather than crashing -- the page count is // informational, so an unreadable structure simply yields 0. func pdfPageCount(data []byte) (pages int) { @@ -314,28 +308,6 @@ func utf8RuneStart(b byte) bool { return b&0xC0 != 0x80 } -// resolvePageCount fills Document.Pages from any available counter. already is -// a count captured during text extraction (0 means unknown). The in-process -// reader is tried first; pdfinfo is the fallback when external tools are on and -// the in-process count is still zero. -func resolvePageCount(data []byte, useExternal bool, already int) int { - if already > 0 { - return already - } - if pages := pageCountInProcess(data); pages > 0 { - return pages - } - if useExternal { - return pageCountPoppler(data) - } - return 0 -} - -var ( - pageCountInProcess = pdfPageCount - pageCountPoppler = pdfPageCountWithPoppler -) - func (o DocumentOptions) maxPages() int { if o.MaxPages > 0 { return o.MaxPages @@ -353,7 +325,7 @@ func popplerAvailable(name string) bool { // extractTextWithPoppler runs `pdftotext - -` (read stdin, write stdout) when // pdftotext is on PATH. The bool is false when the tool is absent or failed, so -// the caller falls back to the pure-Go extractor. Absence is never an error. +// the caller can fall back to the pure-Go extractor. Absence is never an error. func extractTextWithPoppler(data []byte) (string, bool) { if !popplerAvailable("pdftotext") { return "", false @@ -374,29 +346,6 @@ func extractTextWithPoppler(data []byte) (string, bool) { return strings.TrimSpace(stdout.String()), true } -func pdfPageCountWithPoppler(data []byte) int { - if !popplerAvailable("pdfinfo") { - return 0 - } - ctx, cancel := context.WithTimeout(context.Background(), popplerTimeout) - defer cancel() - cmd := exec.CommandContext(ctx, "pdfinfo", "-") - cmd.Stdin = bytes.NewReader(data) - out, err := cmd.Output() - if err != nil { - return 0 - } - for _, line := range strings.Split(string(out), "\n") { - if value, ok := strings.CutPrefix(strings.TrimSpace(line), "Pages:"); ok { - var pages int - if _, err := fmt.Sscan(value, &pages); err == nil { - return pages - } - } - } - return 0 -} - // rasterizeWithPoppler renders the first maxPages pages to PNG via pdftoppm and // returns them as normalized ImageBlocks (reusing the image allow-list, sniff, // and per-image cap). It returns an error when pdftoppm is absent or rendering diff --git a/internal/imageinput/pdf_test.go b/internal/imageinput/pdf_test.go index bf78d375d..614250ab3 100644 --- a/internal/imageinput/pdf_test.go +++ b/internal/imageinput/pdf_test.go @@ -8,14 +8,13 @@ import ( "strconv" "strings" "testing" - "time" ) const minimalPDFTextChunkSize = 80 // buildMinimalPDF assembles a tiny, single-page PDF whose content stream draws // the given text. It computes a real cross-reference table and trailer so a -// pure-Go PDF parser (Detective-XH/gopdf) accepts it. Generating the fixture in-test +// pure-Go PDF parser (ledongthuc/pdf) accepts it. Generating the fixture in-test // keeps the repo free of opaque binary blobs while still exercising the real // text-extraction path on real PDF bytes. func buildMinimalPDF(text string) []byte { @@ -316,7 +315,9 @@ func TestLoadDocumentFallsBackToPureGo(t *testing.T) { } } -func TestLoadDocumentVisionUsesText(t *testing.T) { +// Vision-mode extraction without an available rasterizer must not error: it +// degrades to the text layer (a vision model can still read the text block). +func TestLoadDocumentVisionWithoutRasterizerUsesText(t *testing.T) { root := t.TempDir() want := "Vision degrade to text" if err := os.WriteFile(filepath.Join(root, "doc.pdf"), buildMinimalPDF(want), 0o644); err != nil { @@ -374,80 +375,6 @@ func TestPDFPageCount(t *testing.T) { } } -func TestPDFPageCountIndependentOfTextExtraction(t *testing.T) { - root := t.TempDir() - if err := os.WriteFile(filepath.Join(root, "doc.pdf"), buildMinimalPDF("pages without pdftotext"), 0o644); err != nil { - t.Fatalf("write pdf: %v", err) - } - doc, err := LoadDocument("doc.pdf", root, DocumentOptions{disableExternalTools: true}) - if err != nil { - t.Fatalf("LoadDocument: %v", err) - } - if doc.Pages != 1 { - t.Fatalf("Pages = %d, want 1 when text extraction uses the in-process reader", doc.Pages) - } -} - -// pdftotext success does not record a page count, so Pages must still fall -// through to pdfinfo when the in-process reader reports 0. -func TestResolvePageCountFallsBackToPopplerWhenInProcessIsZero(t *testing.T) { - origIn, origPop := pageCountInProcess, pageCountPoppler - t.Cleanup(func() { - pageCountInProcess, pageCountPoppler = origIn, origPop - }) - - pageCountInProcess = func([]byte) int { return 0 } - pageCountPoppler = func([]byte) int { return 7 } - - if got := resolvePageCount(nil, true, 0); got != 7 { - t.Fatalf("pdftotext-ok + in-process 0 + pdfinfo 7: Pages = %d, want 7", got) - } - if got := resolvePageCount(nil, false, 0); got != 0 { - t.Fatalf("external tools disabled: Pages = %d, want 0", got) - } - if got := resolvePageCount(nil, true, 3); got != 3 { - t.Fatalf("already-known count: Pages = %d, want 3", got) - } - - pageCountInProcess = func([]byte) int { return 2 } - if got := resolvePageCount(nil, true, 0); got != 2 { - t.Fatalf("in-process count wins over pdfinfo: Pages = %d, want 2", got) - } -} - -func TestLoadDocumentHostilePDFStaysBounded(t *testing.T) { - root := t.TempDir() - cases := map[string][]byte{ - "cycle.pdf": []byte("%PDF-1.4\n1 0 obj\n<< /Type /Catalog /Pages 1 0 R /Parent 1 0 R /Kids [1 0 R] /Count 999999999 /First 1 0 R /Next 1 0 R >>\nendobj\ntrailer\n<< /Root 1 0 R /Size 999999999 >>\nstartxref\n9\n%%EOF\n"), - "hex.pdf": []byte("%PDF-1.4\n1 0 obj\n<< /Type /Catalog /Pages 2 0 R >>\nendobj\nstream\n<" + strings.Repeat("A", 4096) + "\nendstream\n%%EOF\n"), - } - done := make(chan error, 1) - go func() { - var first error - for name, body := range cases { - path := filepath.Join(root, name) - if err := os.WriteFile(path, body, 0o644); err != nil { - first = err - break - } - _, err := LoadDocument(name, root, DocumentOptions{disableExternalTools: true}) - if err == nil { - first = fmt.Errorf("%s: expected error for hostile PDF", name) - break - } - } - done <- first - }() - select { - case err := <-done: - if err != nil { - t.Fatal(err) - } - case <-time.After(2 * time.Second): - t.Fatal("hostile PDF parsing exceeded the resource bound") - } -} - // LooksLikeDocumentFile sniffs PDF content by magic bytes, so a real PDF with no // ".pdf" extension is still recognized while a non-PDF (even named .pdf) is not. func TestLooksLikeDocumentFile(t *testing.T) { From bcb08a6317723f97433021a15b5c7d73b1541a6a Mon Sep 17 00:00:00 2001 From: jatmn Date: Wed, 19 Aug 2026 11:44:21 -0700 Subject: [PATCH 10/18] fix(daemon): interrupt retry delays while draining --- internal/daemon/pool.go | 14 +++++++++ internal/daemon/pool_test.go | 56 ++++++++++++++++++++++++++++++++++++ 2 files changed, 70 insertions(+) diff --git a/internal/daemon/pool.go b/internal/daemon/pool.go index f9a62e1df..b1fef509e 100644 --- a/internal/daemon/pool.go +++ b/internal/daemon/pool.go @@ -198,6 +198,12 @@ func (p *Pool) Run(ctx context.Context, spec WorkerSpec, sink Sink) (int, error) return 0, ErrPoolDraining } code, err := p.runOnce(ctx, stat.id, spec, sink) + // A run can observe a normal worker result just as Drain starts. Check + // again before classifying it so shutdown remains terminal rather than + // entering a retry path or reporting ErrPermanent. + if p.isDraining() { + return 0, ErrPoolDraining + } switch { case err != nil: lastErr = err @@ -219,6 +225,9 @@ func (p *Pool) Run(ctx context.Context, spec WorkerSpec, sink Sink) (int, error) lastErr = fmt.Errorf("worker %d tempfail (code=%d)", stat.id, code) p.logf("worker %d tempfail — retry after %s", stat.id, p.opts.TempfailDelay) if !p.sleep(ctx, p.opts.TempfailDelay) { + if p.isDraining() { + return 0, ErrPoolDraining + } return 0, ctx.Err() } continue // tempfail retries do not count against the crash backoff @@ -233,6 +242,9 @@ func (p *Pool) Run(ctx context.Context, spec WorkerSpec, sink Sink) (int, error) delay := p.opts.Backoff(stat.restarts) p.logf("worker %d restart %d after backoff %s", stat.id, stat.restarts, delay) if !p.sleep(ctx, delay) { + if p.isDraining() { + return 0, ErrPoolDraining + } return 0, ctx.Err() } } @@ -351,6 +363,8 @@ func (p *Pool) sleep(ctx context.Context, d time.Duration) bool { return true case <-ctx.Done(): return false + case <-p.drained: + return false } } diff --git a/internal/daemon/pool_test.go b/internal/daemon/pool_test.go index 28063e965..0c860c069 100644 --- a/internal/daemon/pool_test.go +++ b/internal/daemon/pool_test.go @@ -3,6 +3,7 @@ package daemon import ( "context" "errors" + "strings" "sync" "sync/atomic" "testing" @@ -282,6 +283,61 @@ func TestPoolDrainKillsWorkerLaunchedAfterDrainStarts(t *testing.T) { } } +func TestPoolDrainInterruptsRetryDelays(t *testing.T) { + cases := []struct { + name string + exitCode int + maxAttempts int + }{ + {name: "backoff", exitCode: 1, maxAttempts: 2}, + {name: "tempfail", exitCode: ExitTempfail, maxAttempts: 1}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + delaying := make(chan struct{}) + pool, err := NewPool(PoolOptions{ + Size: 1, + MaxAttempts: tc.maxAttempts, + KillTimeout: time.Second, + TempfailDelay: time.Hour, + Backoff: func(int) time.Duration { + return time.Hour + }, + Log: func(message string) { + if strings.Contains(message, "retry after") || strings.Contains(message, "restart") { + select { + case <-delaying: + default: + close(delaying) + } + } + }, + Launcher: func(context.Context, WorkerSpec) (WorkerHandle, error) { + return &fakeWorker{pid: 1, exitCode: tc.exitCode}, nil + }, + }) + if err != nil { + t.Fatalf("NewPool: %v", err) + } + result := make(chan error, 1) + go func() { + _, err := pool.Run(context.Background(), WorkerSpec{Session: "a"}, &collectSink{}) + result <- err + }() + <-delaying + pool.Drain() + select { + case err := <-result: + if !errors.Is(err, ErrPoolDraining) { + t.Fatalf("Run error = %v, want ErrPoolDraining", err) + } + case <-time.After(2 * time.Second): + t.Fatal("Run remained in retry delay after Drain") + } + }) + } +} + func waitFor(t *testing.T, cond func() bool) { t.Helper() deadline := time.Now().Add(2 * time.Second) From 995043bdde29a4aabd34bbdc236b20c75737929f Mon Sep 17 00:00:00 2001 From: jatmn Date: Wed, 19 Aug 2026 11:48:01 -0700 Subject: [PATCH 11/18] docs(daemon): clarify bounded late-launch drain cleanup --- internal/daemon/pool.go | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/internal/daemon/pool.go b/internal/daemon/pool.go index b1fef509e..24cb1ad11 100644 --- a/internal/daemon/pool.go +++ b/internal/daemon/pool.go @@ -370,8 +370,9 @@ func (p *Pool) sleep(ctx context.Context, d time.Duration) bool { // Drain stops accepting new work, gives in-flight workers a grace window // (KillTimeout) to finish on their own, then force-kills any straggler. It -// returns as soon as the pool is idle (graceful) or the window elapses. Safe to -// call once; subsequent calls are no-ops. +// returns as soon as the pool is idle, the grace window elapses and existing +// workers are force-killed, or the separately bounded late-launch cleanup +// completes. Safe to call once; subsequent calls are no-ops. func (p *Pool) Drain() { p.drainOnce.Do(func() { p.mu.Lock() From 8d20f9f578371a3859f6efefb66b6fe24e04f839 Mon Sep 17 00:00:00 2001 From: jatmn Date: Wed, 19 Aug 2026 11:49:11 -0700 Subject: [PATCH 12/18] refactor(daemon): remove unused pool tracker --- internal/daemon/pool.go | 9 --------- 1 file changed, 9 deletions(-) diff --git a/internal/daemon/pool.go b/internal/daemon/pool.go index 24cb1ad11..6523c91a6 100644 --- a/internal/daemon/pool.go +++ b/internal/daemon/pool.go @@ -329,15 +329,6 @@ func (p *Pool) newStat() *workerStat { return &workerStat{id: p.nextID} } -// track/untrack key the active set by the pool's monotonic worker id, not the OS -// pid: the OS can reuse a pid the instant a worker exits, so a pid key could collide -// a finished worker with a freshly-launched one and drop the wrong handle (D10). -func (p *Pool) track(id int, h WorkerHandle) { - p.mu.Lock() - p.active[id] = h - p.mu.Unlock() -} - func (p *Pool) untrack(id int) { p.mu.Lock() delete(p.active, id) From 3eb716d12d9978ddcbde600f1d6e7eb4c8f6848a Mon Sep 17 00:00:00 2001 From: jatmn Date: Wed, 19 Aug 2026 11:52:47 -0700 Subject: [PATCH 13/18] docs(daemon): describe bounded late cleanup --- internal/daemon/pool.go | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/internal/daemon/pool.go b/internal/daemon/pool.go index 6523c91a6..024a87b81 100644 --- a/internal/daemon/pool.go +++ b/internal/daemon/pool.go @@ -362,8 +362,9 @@ func (p *Pool) sleep(ctx context.Context, d time.Duration) bool { // Drain stops accepting new work, gives in-flight workers a grace window // (KillTimeout) to finish on their own, then force-kills any straggler. It // returns as soon as the pool is idle, the grace window elapses and existing -// workers are force-killed, or the separately bounded late-launch cleanup -// completes. Safe to call once; subsequent calls are no-ops. +// workers are force-killed, or one separately bounded late-launch cleanup wait +// elapses. A launcher that ignores cancellation may finish its cleanup after +// Drain returns. Safe to call once; subsequent calls are no-ops. func (p *Pool) Drain() { p.drainOnce.Do(func() { p.mu.Lock() From 87c6a040f765ab52f97ffae13814702c050c4cc7 Mon Sep 17 00:00:00 2001 From: jatmn Date: Wed, 19 Aug 2026 12:52:36 -0700 Subject: [PATCH 14/18] test(daemon): prove drain waits for late worker reap --- internal/daemon/pool_test.go | 27 +++++++++++++++++++++++++-- 1 file changed, 25 insertions(+), 2 deletions(-) diff --git a/internal/daemon/pool_test.go b/internal/daemon/pool_test.go index 0c860c069..8bc23235f 100644 --- a/internal/daemon/pool_test.go +++ b/internal/daemon/pool_test.go @@ -39,6 +39,10 @@ type fakeWorker struct { exitCode int killed int32 waitCh chan struct{} // when non-nil, Wait blocks until closed (drain tests) + killCh chan struct{} // when non-nil, Kill signals before any blocked Wait returns + // waitAfterKill keeps Wait blocked after Kill so drain tests can prove the + // pool waits for reaping rather than merely dispatching a kill signal. + waitAfterKill bool } func (w *fakeWorker) Stdout() Lines { return &fakeLines{lines: w.out, err: w.outErr} } @@ -50,7 +54,14 @@ func (w *fakeWorker) Wait() (int, error) { } func (w *fakeWorker) Kill() error { atomic.StoreInt32(&w.killed, 1) - if w.waitCh != nil { + if w.killCh != nil { + select { + case <-w.killCh: + default: + close(w.killCh) + } + } + if w.waitCh != nil && !w.waitAfterKill { select { case <-w.waitCh: default: @@ -247,7 +258,8 @@ func TestPoolDrainKillsStraggler(t *testing.T) { func TestPoolDrainKillsWorkerLaunchedAfterDrainStarts(t *testing.T) { launchStarted := make(chan struct{}) releaseLaunch := make(chan struct{}) - straggler := &fakeWorker{pid: 1, waitCh: make(chan struct{})} + releaseWait := make(chan struct{}) + straggler := &fakeWorker{pid: 1, waitCh: releaseWait, killCh: make(chan struct{}), waitAfterKill: true} pool, _ := NewPool(PoolOptions{Size: 1, MaxAttempts: 1, KillTimeout: 2 * time.Second, Backoff: func(int) time.Duration { return 0 }, Launcher: func(context.Context, WorkerSpec) (WorkerHandle, error) { close(launchStarted) <-releaseLaunch @@ -269,6 +281,17 @@ func TestPoolDrainKillsWorkerLaunchedAfterDrainStarts(t *testing.T) { } close(releaseLaunch) select { + case <-straggler.killCh: + case <-time.After(2 * time.Second): + t.Fatal("Drain did not kill the worker launched after draining began") + } + select { + case <-drained: + t.Fatal("Drain returned before the late worker was reaped") + default: + } + close(releaseWait) + select { case <-drained: case <-time.After(2 * time.Second): t.Fatal("Drain did not finish") From 8f5c467cd3f9b96ff4fcf8ce0be1ccc0a431705b Mon Sep 17 00:00:00 2001 From: jatmn Date: Thu, 20 Aug 2026 07:37:39 -0700 Subject: [PATCH 15/18] test(daemon): bound pool synchronization waits --- internal/daemon/pool_test.go | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/internal/daemon/pool_test.go b/internal/daemon/pool_test.go index 8bc23235f..ab9422078 100644 --- a/internal/daemon/pool_test.go +++ b/internal/daemon/pool_test.go @@ -270,7 +270,11 @@ func TestPoolDrainKillsWorkerLaunchedAfterDrainStarts(t *testing.T) { _, _ = pool.Run(context.Background(), WorkerSpec{Session: "a"}, &collectSink{}) close(runDone) }() - <-launchStarted + select { + case <-launchStarted: + case <-time.After(2 * time.Second): + t.Fatal("Run did not start the launcher") + } drained := make(chan struct{}) go func() { pool.Drain(); close(drained) }() waitFor(t, pool.isDraining) @@ -347,7 +351,11 @@ func TestPoolDrainInterruptsRetryDelays(t *testing.T) { _, err := pool.Run(context.Background(), WorkerSpec{Session: "a"}, &collectSink{}) result <- err }() - <-delaying + select { + case <-delaying: + case <-time.After(2 * time.Second): + t.Fatal("Run did not enter the retry delay") + } pool.Drain() select { case err := <-result: From 0865440a9b419601b19635d576d4321d408aeabd Mon Sep 17 00:00:00 2001 From: jatmn Date: Thu, 20 Aug 2026 08:25:17 -0700 Subject: [PATCH 16/18] test(daemon): assert late drain result --- internal/daemon/pool_test.go | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/internal/daemon/pool_test.go b/internal/daemon/pool_test.go index ab9422078..2d2c5b379 100644 --- a/internal/daemon/pool_test.go +++ b/internal/daemon/pool_test.go @@ -265,10 +265,10 @@ func TestPoolDrainKillsWorkerLaunchedAfterDrainStarts(t *testing.T) { <-releaseLaunch return straggler, nil }}) - runDone := make(chan struct{}) + runResult := make(chan error, 1) go func() { - _, _ = pool.Run(context.Background(), WorkerSpec{Session: "a"}, &collectSink{}) - close(runDone) + _, err := pool.Run(context.Background(), WorkerSpec{Session: "a"}, &collectSink{}) + runResult <- err }() select { case <-launchStarted: @@ -301,7 +301,10 @@ func TestPoolDrainKillsWorkerLaunchedAfterDrainStarts(t *testing.T) { t.Fatal("Drain did not finish") } select { - case <-runDone: + case err := <-runResult: + if !errors.Is(err, ErrPoolDraining) { + t.Fatalf("Run error = %v, want ErrPoolDraining", err) + } case <-time.After(2 * time.Second): t.Fatal("Run did not finish after Drain") } From b3faaa13fe39f096f69b41e56c84e76c2f109536 Mon Sep 17 00:00:00 2001 From: jatmn Date: Thu, 20 Aug 2026 16:53:46 -0700 Subject: [PATCH 17/18] fix(daemon): publish drain state before cancellation --- go.sum | 2 - internal/daemon/pool.go | 23 +++++++---- internal/daemon/pool_test.go | 57 +++++++++++++++++++++++++++ internal/daemon/server.go | 4 ++ internal/daemon/server_test.go | 71 ++++++++++++++++++++++++++++++++++ 5 files changed, 148 insertions(+), 9 deletions(-) diff --git a/go.sum b/go.sum index 232404cf1..d4601485b 100644 --- a/go.sum +++ b/go.sum @@ -64,8 +64,6 @@ github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e h1:JVG44RsyaB9T2KIHavM github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e/go.mod h1:RbqR21r5mrJuqunuUZ/Dhy/avygyECGrLceyNeo4LiM= golang.org/x/exp v0.0.0-20260611194520-c48552f49976 h1:X8Hz2ImujgbmetVuW+w2YkyZChE3cBpZi2P158rTG9M= golang.org/x/exp v0.0.0-20260611194520-c48552f49976/go.mod h1:vnf4pv9iKZXY58sQE1L86zmNWJ4159e1RkcWiLCkeEY= -golang.org/x/image v0.44.0 h1:+tDekMZED9+LrtB3G5xzRggpVh9CARjZqROla3R3R+I= -golang.org/x/image v0.44.0/go.mod h1:V8K3KE9KKKE+pLpQDOeN18w9oacNSvy1tDOirTu4xtY= golang.org/x/image v0.45.0 h1:FMb1nTbH5H9vF55SriQHgFw5GnNL9Jg6L25BwXKzhB0= golang.org/x/image v0.45.0/go.mod h1:n62x/7RqlwXDvGsSU4u6IUTUf6KghUZ9Bt7cG/T9Fx4= golang.org/x/sync v0.22.0 h1:SZjpbeLmrCk4xhRSZFNZW5gFUeCeFgjekvI/+gfScek= diff --git a/internal/daemon/pool.go b/internal/daemon/pool.go index 024a87b81..235d57a91 100644 --- a/internal/daemon/pool.go +++ b/internal/daemon/pool.go @@ -98,8 +98,9 @@ type Pool struct { launching int // launchers in progress; Drain must not mistake these for idle nextID int - drainOnce sync.Once - drained chan struct{} + drainStartOnce sync.Once + drainOnce sync.Once + drained chan struct{} } // workerStat tracks one in-flight request's restart count (local to Run). @@ -359,6 +360,18 @@ func (p *Pool) sleep(ctx context.Context, d time.Duration) bool { } } +// beginDrain publishes the terminal pool state before callers cancel work that +// may be waiting in Run. Publishing this separately from Drain's bounded +// cleanup makes the shutdown result deterministic for every Run wakeup. +func (p *Pool) beginDrain() { + p.drainStartOnce.Do(func() { + p.mu.Lock() + p.draining = true + p.mu.Unlock() + close(p.drained) + }) +} + // Drain stops accepting new work, gives in-flight workers a grace window // (KillTimeout) to finish on their own, then force-kills any straggler. It // returns as soon as the pool is idle, the grace window elapses and existing @@ -366,12 +379,8 @@ func (p *Pool) sleep(ctx context.Context, d time.Duration) bool { // elapses. A launcher that ignores cancellation may finish its cleanup after // Drain returns. Safe to call once; subsequent calls are no-ops. func (p *Pool) Drain() { + p.beginDrain() p.drainOnce.Do(func() { - p.mu.Lock() - p.draining = true - p.mu.Unlock() - close(p.drained) - // Grace window: poll until idle or the deadline. deadline := time.Now().Add(p.opts.KillTimeout) for time.Now().Before(deadline) { diff --git a/internal/daemon/pool_test.go b/internal/daemon/pool_test.go index 2d2c5b379..540fc4bd1 100644 --- a/internal/daemon/pool_test.go +++ b/internal/daemon/pool_test.go @@ -313,6 +313,63 @@ func TestPoolDrainKillsWorkerLaunchedAfterDrainStarts(t *testing.T) { } } +func TestPoolDrainBoundsBlockedLauncher(t *testing.T) { + launchStarted := make(chan struct{}) + releaseLaunch := make(chan struct{}) + lateWorker := &fakeWorker{pid: 1} + pool, err := NewPool(PoolOptions{ + Size: 1, + MaxAttempts: 1, + KillTimeout: 100 * time.Millisecond, + Backoff: func(int) time.Duration { return 0 }, + Launcher: func(context.Context, WorkerSpec) (WorkerHandle, error) { + close(launchStarted) + <-releaseLaunch + return lateWorker, nil + }, + }) + if err != nil { + t.Fatalf("NewPool: %v", err) + } + runResult := make(chan error, 1) + go func() { + _, err := pool.Run(context.Background(), WorkerSpec{Session: "a"}, &collectSink{}) + runResult <- err + }() + select { + case <-launchStarted: + case <-time.After(2 * time.Second): + t.Fatal("Run did not start the launcher") + } + drained := make(chan struct{}) + go func() { pool.Drain(); close(drained) }() + waitFor(t, pool.isDraining) + // The first timeout accounts for the launch in progress. The second, separate + // timeout is what keeps Drain bounded once no handle exists to kill yet. + select { + case <-drained: + t.Fatal("Drain returned before the separately bounded late-launch wait") + case <-time.After(150 * time.Millisecond): + } + select { + case <-drained: + case <-time.After(2 * time.Second): + t.Fatal("Drain did not return after the bounded blocked-launch wait") + } + close(releaseLaunch) + select { + case err := <-runResult: + if !errors.Is(err, ErrPoolDraining) { + t.Fatalf("Run error = %v, want ErrPoolDraining", err) + } + case <-time.After(2 * time.Second): + t.Fatal("Run did not finish after releasing the blocked launcher") + } + if atomic.LoadInt32(&lateWorker.killed) != 1 { + t.Fatal("late worker was not killed after the bounded Drain return") + } +} + func TestPoolDrainInterruptsRetryDelays(t *testing.T) { cases := []struct { name string diff --git a/internal/daemon/server.go b/internal/daemon/server.go index 21c44e0d6..24160354a 100644 --- a/internal/daemon/server.go +++ b/internal/daemon/server.go @@ -170,6 +170,10 @@ func (s *Server) untrackConn(c net.Conn) { func (s *Server) Shutdown() { s.shutdownOnce.Do(func() { close(s.done) + // Publish pool shutdown before cancelling session contexts. Otherwise a + // Run woken from a retry delay can observe context cancellation before + // Drain marks the pool terminal and incorrectly report context.Canceled. + s.opts.Pool.beginDrain() s.cancel() // stop in-flight pool runs s.mu.Lock() if s.listener != nil { diff --git a/internal/daemon/server_test.go b/internal/daemon/server_test.go index 0fa6c99f7..96c274969 100644 --- a/internal/daemon/server_test.go +++ b/internal/daemon/server_test.go @@ -1,9 +1,11 @@ package daemon import ( + "context" "errors" "os" "path/filepath" + "strings" "testing" "time" ) @@ -178,3 +180,72 @@ func TestServerRejectsUnknownCommand(t *testing.T) { t.Fatal("run with empty session id must return an error") } } + +func TestServerShutdownMakesRetryDelaysDrainTerminal(t *testing.T) { + cases := []struct { + name string + exitCode int + maxAttempts int + }{ + {name: "backoff", exitCode: 1, maxAttempts: 2}, + {name: "tempfail", exitCode: ExitTempfail, maxAttempts: 1}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + delaying := make(chan struct{}) + pool, err := NewPool(PoolOptions{ + Size: 1, + MaxAttempts: tc.maxAttempts, + KillTimeout: 20 * time.Millisecond, + TempfailDelay: time.Hour, + Backoff: func(int) time.Duration { return time.Hour }, + Log: func(message string) { + if strings.Contains(message, "retry after") || strings.Contains(message, "restart") { + select { + case <-delaying: + default: + close(delaying) + } + } + }, + Launcher: func(context.Context, WorkerSpec) (WorkerHandle, error) { + return &fakeWorker{pid: 1, exitCode: tc.exitCode}, nil + }, + }) + if err != nil { + t.Fatalf("NewPool: %v", err) + } + mgr, err := NewSessionManager(SessionManagerOptions{Pool: pool}) + if err != nil { + t.Fatalf("NewSessionManager: %v", err) + } + dir := t.TempDir() + srv, err := NewServer(ServerOptions{ + Paths: Paths{Socket: filepath.Join(dir, "d.sock"), Lock: filepath.Join(dir, "d.lock"), Status: filepath.Join(dir, "d.status")}, + Manager: mgr, + Pool: pool, + }) + if err != nil { + t.Fatalf("NewServer: %v", err) + } + sess, err := mgr.Start(srv.ctx, WorkerSpec{Session: "a"}) + if err != nil { + t.Fatalf("Start: %v", err) + } + select { + case <-delaying: + case <-time.After(2 * time.Second): + t.Fatal("session did not enter retry delay") + } + srv.Shutdown() + select { + case <-sess.Done(): + if !errors.Is(sess.Err(), ErrPoolDraining) { + t.Fatalf("session error = %v, want ErrPoolDraining", sess.Err()) + } + case <-time.After(2 * time.Second): + t.Fatal("session did not finish after shutdown") + } + }) + } +} From a0c3245720733f4168b5977de6177d0be869388d Mon Sep 17 00:00:00 2001 From: jatmn Date: Fri, 28 Aug 2026 10:05:18 -0700 Subject: [PATCH 18/18] test(daemon): anchor blocked-launch drain timing --- internal/daemon/pool_test.go | 22 ++++++++++++---------- 1 file changed, 12 insertions(+), 10 deletions(-) diff --git a/internal/daemon/pool_test.go b/internal/daemon/pool_test.go index 540fc4bd1..aa22818ba 100644 --- a/internal/daemon/pool_test.go +++ b/internal/daemon/pool_test.go @@ -314,13 +314,14 @@ func TestPoolDrainKillsWorkerLaunchedAfterDrainStarts(t *testing.T) { } func TestPoolDrainBoundsBlockedLauncher(t *testing.T) { + const killTimeout = 100 * time.Millisecond launchStarted := make(chan struct{}) releaseLaunch := make(chan struct{}) lateWorker := &fakeWorker{pid: 1} pool, err := NewPool(PoolOptions{ Size: 1, MaxAttempts: 1, - KillTimeout: 100 * time.Millisecond, + KillTimeout: killTimeout, Backoff: func(int) time.Duration { return 0 }, Launcher: func(context.Context, WorkerSpec) (WorkerHandle, error) { close(launchStarted) @@ -341,18 +342,19 @@ func TestPoolDrainBoundsBlockedLauncher(t *testing.T) { case <-time.After(2 * time.Second): t.Fatal("Run did not start the launcher") } - drained := make(chan struct{}) - go func() { pool.Drain(); close(drained) }() - waitFor(t, pool.isDraining) + drainElapsed := make(chan time.Duration, 1) + go func() { + start := time.Now() + pool.Drain() + drainElapsed <- time.Since(start) + }() // The first timeout accounts for the launch in progress. The second, separate // timeout is what keeps Drain bounded once no handle exists to kill yet. select { - case <-drained: - t.Fatal("Drain returned before the separately bounded late-launch wait") - case <-time.After(150 * time.Millisecond): - } - select { - case <-drained: + case elapsed := <-drainElapsed: + if elapsed < 2*killTimeout { + t.Fatalf("Drain returned after %s, want at least both bounded waits (%s)", elapsed, 2*killTimeout) + } case <-time.After(2 * time.Second): t.Fatal("Drain did not return after the bounded blocked-launch wait") }