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, "
hi
") })) 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) } }