sift/internal/api/pdf.go

package api

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"sync/atomic"
	"time"

	"git.wntrmute.dev/kyle/sift/internal/auth"
	"git.wntrmute.dev/kyle/sift/internal/dcache"
	"git.wntrmute.dev/kyle/sift/internal/fetch"
	"git.wntrmute.dev/kyle/sift/internal/pdfdoc"
	"git.wntrmute.dev/kyle/sift/internal/pdfpool"
	"git.wntrmute.dev/kyle/sift/internal/pdftext"
)

const (
	// DefaultPDFTTL is how long a cached PDF is served before it is revalidated.
	DefaultPDFTTL = 24 * time.Hour
	// DefaultPDFMaxQuestionChunks is the safety cap on how long a document a question is
	// answered over; beyond it the question is refused in favour of grep or pages. Longer
	// documents are handled by retrieval (BM25 first, embedding their chunks in the
	// background), so this only bounds the un-answerable.
	DefaultPDFMaxQuestionChunks = 5000
)

// runner returns the configured extraction child, or a zero Runner that runs
// /proc/self/exe pdftext.
func (s *Server) runner() *pdftext.Runner {
	if s.PDFRunner != nil {
		return s.PDFRunner
	}
	return &pdftext.Runner{}
}

// pool returns the extraction pool, built once and only when none is configured.
func (s *Server) pool() *pdfpool.Pool {
	// Always through the Once: reading PDFPool outside it would race with the first writer.
	s.pdfPoolOnce.Do(func() {
		if s.PDFPool == nil {
			s.PDFPool = pdfpool.New(2, 4, 20*time.Second)
		}
	})
	return s.PDFPool
}

// extractPDF runs one extraction under the pool, coalesced by document hash and
// bounded to one per caller.
func (s *Server) extractPDF(ctx context.Context, c *auth.Caller, res *fetch.Resource) (*pdftext.Doc, error) {
	var execNano atomic.Int64
	fn := func(jobCtx context.Context) (any, error) {
		start := time.Now()
		v, err := s.runner().Run(jobCtx, res.Body)
		execNano.Store(int64(time.Since(start)))
		return v, err
	}
	start := time.Now()
	caller := c.Name
	if isBatchFetch(ctx) {
		caller = c.Name + "/batch"
	}
	v, err := s.pool().Do(ctx, s.callerClass(c)+"|"+res.SHA256(), caller, fn)
	total := time.Since(start)
	// The job's own extraction time is what this call measured it to be; the
	// rest of Do is time spent waiting for a slot.
	if exec := execNano.Load(); exec > 0 {
		wait := total - time.Duration(exec)
		if wait < 0 {
			wait = 0
		}
		s.Prom.ObservePDFWait(wait)
	}
	pages := 0
	if err == nil {
		if doc, ok := v.(*pdftext.Doc); ok {
			pages = doc.Pages
		}
	}
	s.Prom.ObservePDF(pdfOutcome(err), total, pages)
	if err != nil {
		return nil, pdfError(err)
	}
	doc, ok := v.(*pdftext.Doc)
	if !ok {
		return nil, fmt.Errorf("pdf: bad result %T", v)
	}
	return doc, nil
}

// pdfOutcome classifies an extraction error into the single value the
// sift_pdf_extract_total label carries; a nil error is a success and anything
// unrecognised is "other" so raw text never reaches a label.
func pdfOutcome(err error) string {
	switch {
	case err == nil:
		return "ok"
	case errors.Is(err, pdftext.ErrNoText):
		return "no_text"
	case errors.Is(err, pdftext.ErrEncrypted):
		return "encrypted"
	case errors.Is(err, pdftext.ErrUnreadable):
		return "unreadable"
	case errors.Is(err, pdftext.ErrTimeout):
		return "timeout"
	case errors.Is(err, pdftext.ErrTooMuchText):
		return "too_large"
	case errors.Is(err, pdfpool.ErrBusy), errors.Is(err, pdfpool.ErrCallerBusy):
		return "busy"
	default:
		return "other"
	}
}

// pdfError maps every PDF-path error to the status sift returns to its caller;
// a PDF never produces a 500.
func pdfError(err error) error {
	switch {
	case errors.Is(err, pdftext.ErrNoText):
		return &fetch.Error{Status: 422, Msg: "no text layer (scanned PDF?)", Hint: "ocr"}
	case errors.Is(err, pdftext.ErrEncrypted):
		return &fetch.Error{Status: 422, Msg: "encrypted PDF (password required)"}
	case errors.Is(err, pdftext.ErrTimeout):
		return &fetch.Error{Status: 422, Msg: "pdf extraction timed out"}
	case errors.Is(err, pdftext.ErrTooMuchText):
		return &fetch.Error{Status: 422, Msg: "pdf text over the output limit"}
	case errors.Is(err, pdfpool.ErrBusy):
		return &fetch.Error{Status: 503, Msg: "pdf workers busy", RetryAfter: 10}
	case errors.Is(err, pdfpool.ErrCallerBusy):
		return &fetch.Error{Status: 429, Msg: "busy"}
	case errors.Is(err, pdftext.ErrUnreadable):
		return &fetch.Error{Status: 422, Msg: "unreadable PDF"}
	default:
		return &fetch.Error{Status: 422, Msg: "unreadable PDF"}
	}
}

// renderPDF draws the document for an agent: the whole text when it fits the
// output cap, otherwise the map, a page range, or grep results.
func (s *Server) renderPDF(doc *pdftext.Doc, opts pdfdoc.Options) (pdfdoc.Rendered, error) {
	return pdfdoc.Render(doc, opts)
}

// pdfFromCache answers a PDF from the document cache before any fetch: a fresh
// entry directly, a stale one after a conditional GET (a 304 reuses the stored
// text), or a miss. A 200 from the revalidation is returned with ok=false so the
// caller extracts it; a nil cache is always a miss.
func (s *Server) pdfFromCache(ctx context.Context, c *auth.Caller, rawURL string) (*fetch.Resource, *pdftext.Doc, bool, error) {
	cache := s.cache()
	if cache == nil {
		return nil, nil, false, nil
	}
	key := dcache.PDFKey(rawURL, s.callerClass(c))
	e, ok := cache.GetPDF(key)
	if !ok {
		return nil, nil, false, nil
	}
	ttl := s.PDFTTL
	if ttl == 0 {
		ttl = DefaultPDFTTL
	}
	if time.Since(e.FetchedAt) < ttl {
		s.Prom.PDFCacheHit()
		res, doc := s.pdfFromEntry(rawURL, e)
		return res, doc, true, nil
	}
	// A revalidation is an upstream request like any other: it counts against the caller's and
	// the host's limits.
	if err := s.acquire(c.Name, rawURL); err != nil {
		return nil, nil, false, err
	}
	res, err := s.Fetch.FetchIf(ctx, rawURL, e.ETag, e.LastModified)
	s.release(c.Name, rawURL)
	if err != nil {
		return nil, nil, false, err
	}
	if res.Status == http.StatusNotModified {
		cache.TouchPDF(key, time.Now().UTC())
		s.Prom.PDFCacheHit()
		res, doc := s.pdfFromEntry(rawURL, e)
		return res, doc, true, nil
	}
	return res, nil, false, nil
}

// pdfFromEntry rebuilds the synthetic response and doc for a cached entry: the
// request's URL, a 200 PDF with the stored fetch time and hash, and no body.
func (s *Server) pdfFromEntry(rawURL string, e dcache.PDFEntry) (*fetch.Resource, *pdftext.Doc) {
	var doc pdftext.Doc
	_ = json.Unmarshal(e.Meta, &doc.Meta)
	doc.Text = e.Pages
	res := &fetch.Resource{
		URL:         rawURL,
		FinalURL:    rawURL,
		Status:      200,
		ContentType: "application/pdf",
		FetchedAt:   e.FetchedAt,
		Hash:        e.ContentHash,
	}
	return res, &doc
}

// storePDF writes an extracted PDF into the document cache for its TTL: its
// hash, the response's validators, the metadata as JSON and the per-page text.
// A nil cache is a no-op.
func (s *Server) storePDF(key string, res *fetch.Resource, doc *pdftext.Doc) {
	cache := s.cache()
	if cache == nil {
		return
	}
	meta, _ := json.Marshal(doc.Meta)
	cache.PutPDF(key, dcache.PDFEntry{
		ContentHash:  res.SHA256(),
		ETag:         res.Header.Get("ETag"),
		LastModified: res.Header.Get("Last-Modified"),
		FetchedAt:    time.Now().UTC(),
		Meta:         meta,
		Pages:        doc.Text,
	})
}