sift/internal/api/pdf_test.go

package api

// P5a task 06: PDFs through /v1/fetch and /v2/scrape — extraction in the child, the pool, the
// document cache, the map, pages, grep, and the errors.
//
// The child is the test binary itself (TestMain below), as in internal/pdftext's tests.
// SIFT_PDFTEXT_RUNLOG makes every child start append a line to a file, so a test can count
// extractions; SIFT_PDFTEXT_SLEEP slows a child down. Both hooks live only in this file.

import (
	"encoding/json"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"net/netip"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"git.wntrmute.dev/kyle/sift/internal/auth"
	"git.wntrmute.dev/kyle/sift/internal/dcache"
	"git.wntrmute.dev/kyle/sift/internal/digest"
	"git.wntrmute.dev/kyle/sift/internal/fetch"
	"git.wntrmute.dev/kyle/sift/internal/guard"
	"git.wntrmute.dev/kyle/sift/internal/pdfpool"
	"git.wntrmute.dev/kyle/sift/internal/pdftest"
	"git.wntrmute.dev/kyle/sift/internal/pdftext"
)

func TestMain(m *testing.M) {
	if os.Getenv("SIFT_PDFTEXT_CHILD") == "1" {
		if p := os.Getenv("SIFT_PDFTEXT_RUNLOG"); p != "" {
			f, err := os.OpenFile(p, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0o600)
			if err == nil {
				_, _ = f.WriteString("run\n")
				_ = f.Close()
			}
		}
		if d, err := time.ParseDuration(os.Getenv("SIFT_PDFTEXT_SLEEP")); err == nil {
			time.Sleep(d)
		}
		os.Exit(pdftext.ChildMain(os.Args[1:], os.Stdin, os.Stdout))
	}
	os.Exit(m.Run())
}

func threePDF() []byte {
	return pdftest.Build(pdftest.Doc{
		Title:   "Three Pages",
		Pages:   []string{"Chapter one begins here.", "The quick brown fox is on page two.", "Last page."},
		Outline: []pdftest.Outline{{Title: "1 Intro", Page: 1, Level: 1}, {Title: "2 Fox", Page: 2, Level: 1}},
	})
}

func longPDF(n int) []byte {
	var pages []string
	for p := 1; p <= n; p++ {
		pages = append(pages, pdftest.Lorem(p, 40))
	}
	return pdftest.Build(pdftest.Doc{Title: "Long Thesis", Pages: pages, Outline: []pdftest.Outline{{Title: "1 Introduction", Page: 1, Level: 1}, {Title: "2 Method", Page: n / 2, Level: 1}}})
}

// origin serves one PDF and counts GETs. It honours If-None-Match against its current ETag.
type origin struct {
	mu   sync.Mutex
	ct   string
	body []byte
	etag string
	gets atomic.Int32
	srv  *httptest.Server
}

func newOrigin(t *testing.T, ct string, body []byte) *origin {
	o := &origin{ct: ct, body: body, etag: `"v1"`}
	o.srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		o.gets.Add(1)
		o.mu.Lock()
		defer o.mu.Unlock()
		if r.Header.Get("If-None-Match") == o.etag {
			w.WriteHeader(http.StatusNotModified)
			return
		}
		w.Header().Set("ETag", o.etag)
		if o.ct != "" {
			w.Header().Set("Content-Type", o.ct)
		}
		_, _ = w.Write(o.body)
	}))
	t.Cleanup(o.srv.Close)
	return o
}

func (o *origin) replace(body []byte, etag string) {
	o.mu.Lock()
	o.body, o.etag = body, etag
	o.mu.Unlock()
}

type pdfRig struct {
	s      *Server
	fleet  string // bearer key of caller "test" (hermes-fleet)
	other  string // bearer key of caller "other" (hermes-fleet)
	gs     string // bearer key of restricted caller "straylight-gs"
	runlog string
	logs   *recHandler
}

func newPDFRig(t *testing.T, childEnv ...string) *pdfRig {
	t.Helper()
	dir := t.TempDir()
	keys := filepath.Join(dir, "keys.toml")
	key := func(name, ws string) string {
		out := filepath.Join(dir, name)
		if err := auth.Add(keys, name, ws, out); err != nil {
			t.Fatal(err)
		}
		b, err := os.ReadFile(out)
		if err != nil {
			t.Fatal(err)
		}
		return strings.TrimSpace(string(b))
	}
	r := &pdfRig{fleet: key("test", "hermes-fleet"), other: key("other", "hermes-fleet"), gs: key("straylight-gs", "guardian-star"), runlog: filepath.Join(dir, "runs"), logs: &recHandler{}}
	store, err := auth.Load(keys)
	if err != nil {
		t.Fatal(err)
	}
	g := &guard.Guard{AllowPrefixes: []netip.Prefix{netip.MustParsePrefix("127.0.0.0/8")}}
	s := New("test", g, &fetch.Fetcher{Guard: g, UA: "sift-test"}, store, slog.New(r.logs))
	s.PermitRestricted = []string{"straylight-gs"}
	db, err := dcache.Open(filepath.Join(dir, "cache.db"), 0, nil)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = db.Close() })
	s.Digest = &digest.Engine{Cache: db} // the cache lives on the digest engine; no model is configured
	s.PDFRunner = &pdftext.Runner{Exe: os.Args[0], Env: append([]string{"SIFT_PDFTEXT_CHILD=1", "SIFT_PDFTEXT_RUNLOG=" + r.runlog}, childEnv...), Deadline: 20 * time.Second}
	r.s = s
	return r
}

func (r *pdfRig) runs() int {
	b, _ := os.ReadFile(r.runlog)
	return strings.Count(string(b), "run\n")
}

func (r *pdfRig) fetch(t *testing.T, key string, body map[string]any) (int, map[string]any, http.Header) {
	t.Helper()
	rec := postJSON(t, r.s, "/v1/fetch", key, body)
	var out map[string]any
	_ = json.Unmarshal(rec.Body.Bytes(), &out)
	return rec.Code, out, rec.Header()
}

func pdfInfo(t *testing.T, out map[string]any) map[string]any {
	t.Helper()
	p, ok := out["pdf"].(map[string]any)
	if !ok {
		t.Fatalf("envelope has no pdf object: %v", out)
	}
	return p
}

func TestPDFCleanFetch(t *testing.T) {
	for _, ct := range []string{"application/pdf", "application/octet-stream", ""} {
		r := newPDFRig(t)
		o := newOrigin(t, ct, threePDF())
		code, out, _ := r.fetch(t, r.fleet, map[string]any{"url": o.srv.URL})
		if code != 200 {
			t.Fatalf("%q: %d %v", ct, code, out)
		}
		content, _ := out["content"].(string)
		if !strings.Contains(content, "--- page 2 ---") || !strings.Contains(content, "quick brown fox") || out["mode"] != "clean" {
			t.Errorf("%q: content %q mode %v", ct, content, out["mode"])
		}
		p := pdfInfo(t, out)
		if p["pages"] != 3.0 || p["title"] != "Three Pages" || p["map"] != false || p["first"] != 1.0 || p["last"] != 3.0 {
			t.Errorf("%q: pdf = %v", ct, p)
		}
		if ol, _ := p["outline"].([]any); len(ol) != 2 {
			t.Errorf("%q: outline = %v", ct, p["outline"])
		}
	}
}

func TestPDFMap(t *testing.T) {
	r := newPDFRig(t)
	r.s.OutCap = 20000
	o := newOrigin(t, "application/pdf", longPDF(40)) // ~40 × 2,700 characters
	code, out, _ := r.fetch(t, r.fleet, map[string]any{"url": o.srv.URL})
	if code != 200 {
		t.Fatalf("%d %v", code, out)
	}
	content, _ := out["content"].(string)
	p := pdfInfo(t, out)
	if p["map"] != true || !strings.Contains(content, "## Outline") || !strings.Contains(content, "- 2 Method (p. 20)") || len([]rune(content)) > 20000 {
		t.Errorf("map: pdf %v, %d chars:\n%.400s", p, len([]rune(content)), content)
	}
}

func TestPDFPagesAndGrep(t *testing.T) {
	r := newPDFRig(t)
	o := newOrigin(t, "application/pdf", threePDF())
	code, out, _ := r.fetch(t, r.fleet, map[string]any{"url": o.srv.URL, "pages": "2-3"})
	content, _ := out["content"].(string)
	if code != 200 || strings.Contains(content, "--- page 1 ---") || !strings.Contains(content, "--- page 3 ---") {
		t.Errorf("pages 2-3: %d %q", code, content)
	}
	if p := pdfInfo(t, out); p["first"] != 2.0 || p["last"] != 3.0 {
		t.Errorf("pages 2-3: pdf %v", p)
	}
	for _, bad := range []map[string]any{{"pages": "0-1"}, {"pages": "9"}, {"grep": strings.Repeat("a", 201)}, {"grep": "(unclosed"}} {
		bad["url"] = o.srv.URL
		if code, out, _ := r.fetch(t, r.fleet, bad); code != 400 || out["success"] != false {
			t.Errorf("%v: %d %v, want 400", bad, code, out)
		}
	}
	code, out, _ = r.fetch(t, r.fleet, map[string]any{"url": o.srv.URL, "grep": "quick brown"})
	content, _ = out["content"].(string)
	if code != 200 || !strings.Contains(content, "p. 2: The quick brown fox is on page two.") || pdfInfo(t, out)["grep_hits"] != 1.0 {
		t.Errorf("grep: %d %q %v", code, content, out["pdf"])
	}
	// pages and grep on a non-PDF are refused rather than ignored.
	html := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.Header().Set("Content-Type", "text/html")
		_, _ = io.WriteString(w, "<html><body><p>hi</p></body></html>")
	}))
	defer html.Close()
	if code, out, _ := r.fetch(t, r.fleet, map[string]any{"url": html.URL, "pages": "1"}); code != 400 {
		t.Errorf("pages on html: %d %v, want 400", code, out)
	}
}

func TestPDFCache(t *testing.T) {
	r := newPDFRig(t)
	o := newOrigin(t, "application/pdf", threePDF())
	get := func(extra map[string]any) string {
		body := map[string]any{"url": o.srv.URL}
		for k, v := range extra {
			body[k] = v
		}
		code, out, _ := r.fetch(t, r.fleet, body)
		if code != 200 {
			t.Fatalf("%v: %d %v", extra, code, out)
		}
		c, _ := out["content"].(string)
		return c
	}
	get(nil)
	if o.gets.Load() != 1 || r.runs() != 1 {
		t.Fatalf("first fetch: %d GETs, %d extractions; want 1 and 1", o.gets.Load(), r.runs())
	}
	if c := get(map[string]any{"pages": "2"}); !strings.Contains(c, "quick brown fox") {
		t.Errorf("cached pages=2: %q", c)
	}
	if o.gets.Load() != 1 || r.runs() != 1 {
		t.Errorf("fresh cache hit: %d GETs, %d extractions; want still 1 and 1", o.gets.Load(), r.runs())
	}
	// Stale: a conditional GET, answered 304, re-uses the text.
	r.s.PDFTTL = time.Nanosecond
	time.Sleep(2 * time.Millisecond)
	get(map[string]any{"grep": "Chapter"})
	if o.gets.Load() != 2 || r.runs() != 1 {
		t.Errorf("after 304: %d GETs, %d extractions; want 2 and 1", o.gets.Load(), r.runs())
	}
	// The document changed: a full GET and a new extraction.
	o.replace(pdftest.Build(pdftest.Doc{Pages: []string{"Brand new text."}}), `"v2"`)
	if c := get(nil); !strings.Contains(c, "Brand new text.") {
		t.Errorf("changed document: %q", c)
	}
	if o.gets.Load() != 3 || r.runs() != 2 {
		t.Errorf("after change: %d GETs, %d extractions; want 3 and 2", o.gets.Load(), r.runs())
	}
	if DefaultPDFTTL != 24*time.Hour {
		t.Errorf("DefaultPDFTTL = %v, want 24h", DefaultPDFTTL)
	}
}

func TestPDFErrors(t *testing.T) {
	r := newPDFRig(t)
	enc, err := os.ReadFile(filepath.Join("..", "..", "tests", "testdata", "pdf", "encrypted-user.pdf"))
	if err != nil {
		t.Fatal(err)
	}
	for _, tc := range []struct {
		name string
		body []byte
		want string
		hint string
	}{
		{"scanned", pdftest.Build(pdftest.Doc{Pages: []string{"", ""}, Scanned: []int{1, 2}}), "no text layer", "ocr"},
		{"encrypted", enc, "encrypted", ""},
		{"garbage", []byte("%PDF-1.4\nthis is not really a pdf"), "unreadable", ""},
	} {
		o := newOrigin(t, "application/pdf", tc.body)
		code, out, _ := r.fetch(t, r.fleet, map[string]any{"url": o.srv.URL})
		msg, _ := out["error"].(string)
		if code != 422 || !strings.Contains(msg, tc.want) || (tc.hint != "" && out["hint"] != tc.hint) {
			t.Errorf("%s: %d %v; want 422 %q hint %q", tc.name, code, out, tc.want, tc.hint)
		}
	}
}

func TestPDFTimeout(t *testing.T) {
	r := newPDFRig(t, "SIFT_PDFTEXT_SLEEP=10s")
	r.s.PDFRunner.Deadline = 300 * time.Millisecond
	o := newOrigin(t, "application/pdf", threePDF())
	code, out, _ := r.fetch(t, r.fleet, map[string]any{"url": o.srv.URL})
	if msg, _ := out["error"].(string); code != 422 || !strings.Contains(msg, "timed out") {
		t.Errorf("slow child: %d %v; want 422 … timed out", code, out)
	}
}

func TestPDFBusy(t *testing.T) {
	r := newPDFRig(t, "SIFT_PDFTEXT_SLEEP=1500ms")
	r.s.PDFPool = pdfpool.New(1, 0, 100*time.Millisecond)
	a := newOrigin(t, "application/pdf", threePDF())
	b := newOrigin(t, "application/pdf", longPDF(2))
	done := make(chan int, 1)
	go func() {
		code, _, _ := r.fetch(t, r.fleet, map[string]any{"url": a.srv.URL})
		done <- code
	}()
	time.Sleep(400 * time.Millisecond) // the first extraction holds the only slot
	code, out, hdr := r.fetch(t, r.other, map[string]any{"url": b.srv.URL})
	if code != 503 || hdr.Get("Retry-After") != "10" || !strings.Contains(fmt.Sprint(out["error"]), "busy") {
		t.Errorf("second caller with the pool full: %d %v Retry-After %q; want 503, busy, 10", code, out, hdr.Get("Retry-After"))
	}
	code, out, _ = r.fetch(t, r.fleet, map[string]any{"url": b.srv.URL})
	if code != 429 {
		t.Errorf("same caller, second document while its first is running: %d %v; want 429", code, out)
	}
	if c := <-done; c != 200 {
		t.Errorf("first fetch: %d", c)
	}
}

func TestPDFRestrictedCaller(t *testing.T) {
	r := newPDFRig(t)
	o := newOrigin(t, "application/pdf", threePDF())
	if code, out, _ := r.fetch(t, r.gs, map[string]any{"url": o.srv.URL}); code != 200 {
		t.Fatalf("restricted fetch: %d %v", code, out)
	}
	r.logs.mu.Lock()
	for _, rec := range r.logs.recs {
		var b strings.Builder
		b.WriteString(rec.Message)
		rec.Attrs(func(a slog.Attr) bool { fmt.Fprintf(&b, " %s=%v", a.Key, a.Value); return true })
		line := b.String()
		if strings.Contains(line, "127.0.0.1") || strings.Contains(line, "Three Pages") || strings.Contains(line, "quick brown") {
			t.Errorf("a restricted caller's PDF fetch logged %q", line)
		}
	}
	r.logs.mu.Unlock()
	// The fleet does not see the restricted caller's cached text: it extracts its own copy.
	if code, _, _ := r.fetch(t, r.fleet, map[string]any{"url": o.srv.URL}); code != 200 || r.runs() != 2 {
		t.Errorf("fleet after restricted: %d, %d extractions; want 200 and 2 (separate cache partitions)", code, r.runs())
	}
}

func TestPDFScrape(t *testing.T) {
	r := newPDFRig(t)
	o := newOrigin(t, "application/pdf", threePDF())
	rec := postJSON(t, r.s, "/v2/scrape", r.fleet, map[string]any{"url": o.srv.URL, "formats": []string{"markdown"}})
	var out struct {
		Success bool `json:"success"`
		Data    struct {
			Markdown string         `json:"markdown"`
			Metadata map[string]any `json:"metadata"`
		} `json:"data"`
	}
	if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil || rec.Code != 200 || !out.Success {
		t.Fatalf("%d %s", rec.Code, rec.Body.String())
	}
	if !strings.Contains(out.Data.Markdown, "--- page 2 ---") || out.Data.Metadata["numPages"] != 3.0 || out.Data.Metadata["title"] != "Three Pages" {
		t.Errorf("scrape: markdown %q metadata %v", out.Data.Markdown, out.Data.Metadata)
	}
	if ct, _ := out.Data.Metadata["contentType"].(string); !strings.HasPrefix(ct, "application/pdf") {
		t.Errorf("contentType = %q", ct)
	}
}

// A question on a document longer than PDFMaxQuestionChunks is refused with a pointer to grep
// and pages, before any model call.
func TestPDFQuestionTooLong(t *testing.T) {
	r := newPDFRig(t)
	r.s.PDFMaxQuestionChunks = 2
	o := newOrigin(t, "application/pdf", longPDF(12))
	code, out, _ := r.fetch(t, r.fleet, map[string]any{"url": o.srv.URL, "question": "what is the method?"})
	if msg, _ := out["error"].(string); code != 422 || !strings.Contains(msg, "too long") || !strings.Contains(msg, "grep") {
		t.Errorf("long question: %d %v; want 422 … too long … grep", code, out)
	}
	// P5b: long documents are answered by retrieval now; the limit is only a safety cap.
	if DefaultPDFMaxQuestionChunks != 5000 {
		t.Errorf("DefaultPDFMaxQuestionChunks = %d, want 5000", DefaultPDFMaxQuestionChunks)
	}
}